From 4ee960b1047f25f2a53bf8c406be3d2b883e88ff Mon Sep 17 00:00:00 2001
From: ebuehrle <43623224+ebuehrle@users.noreply.github.com>
Date: Fri, 5 Nov 2021 08:35:14 +0100
Subject: [PATCH] Fix rendering

---
 scratch/etienne/intersimple/gail/envs.py                   | 7 +++++--
 scratch/etienne/intersimple/gail/options2.py               | 2 +-
 .../intersimple/gail_options_image_random_location.py      | 6 ++++--
 3 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/scratch/etienne/intersimple/gail/envs.py b/scratch/etienne/intersimple/gail/envs.py
index e290f9d..6f3b3fd 100644
--- a/scratch/etienne/intersimple/gail/envs.py
+++ b/scratch/etienne/intersimple/gail/envs.py
@@ -34,10 +34,13 @@ class NRasterizedRouteSpeedRandomAgentLocation(RandomLocation, RandomAgent, Rewa
         NormalizedActionSpace, ActionVisualization, InteractionSimulatorMarkerViz, ImitationCompat, Intersimple):
     pass
 
-class NoBSTimeLimit(TimeLimit):
+class TransparentTimeLimit(TimeLimit):
 
     def __getattr__(self, name):
         return getattr(self.env, name)
+    
+    def close(self, *args, **kwargs):
+        return self.env.close(*args, **kwargs)
 
 def TLNRasterizedRouteRandomAgentLocation(max_episode_steps, *args, **kwargs):
-    return NoBSTimeLimit(NRasterizedRouteRandomAgentLocation(*args, **kwargs), max_episode_steps=max_episode_steps)
+    return TransparentTimeLimit(NRasterizedRouteRandomAgentLocation(*args, **kwargs), max_episode_steps=max_episode_steps)
diff --git a/scratch/etienne/intersimple/gail/options2.py b/scratch/etienne/intersimple/gail/options2.py
index 69514a5..4745157 100644
--- a/scratch/etienne/intersimple/gail/options2.py
+++ b/scratch/etienne/intersimple/gail/options2.py
@@ -91,7 +91,7 @@ class RenderOptions(OptionsEnv):
 
     def _ll_step(self, action):
         out = super()._ll_step(action)
-        self.env.render()
+        self.env.render(mode='post')
         return out
     
     def close(self, *args, **kwargs):
diff --git a/scratch/etienne/intersimple/gail_options_image_random_location.py b/scratch/etienne/intersimple/gail_options_image_random_location.py
index 9131ba3..41ce439 100644
--- a/scratch/etienne/intersimple/gail_options_image_random_location.py
+++ b/scratch/etienne/intersimple/gail_options_image_random_location.py
@@ -19,7 +19,6 @@ from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv
 
 model_name = 'gail_options_image_random_location'
 env_settings = {'width': 70, 'height': 70, 'm_per_px': 1, 'mu': 0.001, 'random_skip': True, 'max_episode_steps': 50}
-env = TLNRasterizedRouteRandomAgentLocation(**env_settings)
 
 ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is safe fallback
 
@@ -33,6 +32,7 @@ def train(
         discount=0.99,
         epochs=100,
     ):
+    env = TLNRasterizedRouteRandomAgentLocation(**env_settings)
 
     tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
     tempdir_path = pathlib.Path(tempdir.name)
@@ -79,7 +79,6 @@ def train(
     return generator
 
 def video(model_name, env):
-    env = RenderOptions(env, options=ALL_OPTIONS)
     model = stable_baselines3.PPO.load(model_name)
 
     done = False
@@ -91,6 +90,9 @@ def video(model_name, env):
     env.close(filestr='render/'+model_name)
 
 def evaluate():
+    video_settings = { **env_settings, 'random_skip': False, 'max_episode_steps': 1000 }
+    env = TLNRasterizedRouteRandomAgentLocation(**video_settings)
+    env = RenderOptions(env, options=ALL_OPTIONS)
     video(
         model_name=model_name,
         env=env