Merge pull request #441 from kjang96/visualizer_fix

Minor fixes
flow-project · Mar 7, 2019 · 32a0ecc · 32a0ecc
2 parents 2397a15 + fb17351
commit 32a0ecc
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/flow/visualize/visualizer_rllib.py b/flow/visualize/visualizer_rllib.py
@@ -160,7 +160,8 @@ def visualizer_rllib(args):
     checkpoint = checkpoint + '/checkpoint-' + args.checkpoint_num
     agent.restore(checkpoint)
 
-    if hasattr(agent, "local_evaluator") and os.environ["TEST_FLAG"] != 'True':
+    if hasattr(agent, "local_evaluator") and \
+            os.environ.get("TEST_FLAG") != 'True':
         env = agent.local_evaluator.env
     else:
         env = gym.make(env_name)
@@ -193,6 +194,9 @@ def visualizer_rllib(args):
     else:
         use_lstm = False
 
+    env.restart_simulation(
+        sim_params=sim_params, render=sim_params.render)
+
     final_outflows = []
     mean_speed = []
     for i in range(args.num_rollouts):