Add option to render expert rollout

scratch/etienne/intersimple/data/expert.py

@@ -74,7 +74,17 @@ class DummyVecEnvPolicy():
         states = [p[1] for p in predictions]
         return actions, states
 
-def demonstrations(expert='NormalizedIntersimpleExpert', env='NRasterizedRandomAgent', path=None, min_timesteps=25000, min_episodes=None, env_args={}, policy_args={}):
+def save_video(env, expert):
+    env.reset()
+    env.render()
+    done = False
+    while not done:
+        actions, _ = expert.predict()
+        _, _, done, _ = env.step(actions)
+        env.render()
+    env.close()
+
+def demonstrations(expert='NormalizedIntersimpleExpert', env='NRasterizedRandomAgent', path=None, min_timesteps=25000, min_episodes=None, video=False, env_args={}, policy_args={}):
     """Rollout and save expert demos.
     
     Usage:
@@ -91,6 +101,9 @@ def demonstrations(expert='NormalizedIntersimpleExpert', env='NRasterizedRandomA
     policy = Expert(env, **policy_args)
     venv_policy = DummyVecEnvPolicy([lambda: policy])
 
+    if video:
+        save_video(env, policy)
+    
     path = path or (policy.__class__.__name__ + '_' + env.__class__.__name__ + '.pkl')
 
     rollout.rollout_and_save(