Merge pull request #37 from maximecb/patch-2

Update visualize.py

scripts/visualize.py

@@ -23,6 +23,9 @@
                     help="pause duration between two consequent actions of the agent (default: 0.1)")
 parser.add_argument("--gif", type=str, default=None,
                     help="store output as gif with the given filename")
+parser.add_argument("--episodes", type=int, default=1000000,
+                    help="number of episodes to visualize")
+
 args = parser.parse_args()
 
 # Set seed for all randomness sources
@@ -53,24 +56,28 @@
    from array2gif import write_gif
    frames = []
 
-done = True
+# Create a window to view the environment
+env.render('human')
 
-while True:
-    if done:
-        obs = env.reset()
+for episode in range(args.episodes):
+    obs = env.reset()
 
-    time.sleep(args.pause)
-    renderer = env.render()
-    if args.gif:
-        frames.append(numpy.moveaxis(env.render("rgb_array"), 2, 0))
+    while True:
+        env.render('human')
+        if args.gif:
+            frames.append(numpy.moveaxis(env.render("rgb_array"), 2, 0))
 
-    action = agent.get_action(obs)
-    obs, reward, done, _ = env.step(action)
-    agent.analyze_feedback(reward, done)
+        action = agent.get_action(obs)
+        obs, reward, done, _ = env.step(action)
+        agent.analyze_feedback(reward, done)
 
-    if renderer.window is None:
-        if args.gif:
-            print("Saving gif... ", end="")
-            write_gif(numpy.array(frames), args.gif+".gif", fps=1/args.pause)
-            print("Done.")
+        if done or env.window.closed:
+            break
+
+    if env.window.closed:
         break
+
+if args.gif:
+    print("Saving gif... ", end="")
+    write_gif(numpy.array(frames), args.gif+".gif", fps=1/args.pause)
+    print("Done.")