fix visualize.py

lcswillems · Oct 5, 2022 · 7dcb2be · 7dcb2be
1 parent 5e309dc
commit 7dcb2be
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/scripts/visualize.py b/scripts/visualize.py
@@ -41,7 +41,7 @@
 
 # Load environment
 
-env = utils.make_env(args.env, args.seed)
+env = utils.make_env(args.env, args.seed, render_mode="human")
 for _ in range(args.shift):
     env.reset()
 print("Environment loaded\n")
@@ -61,15 +61,15 @@
     frames = []
 
 # Create a window to view the environment
-env.render("human")
+env.render()
 
 for episode in range(args.episodes):
     obs, _ = env.reset()
 
     while True:
-        env.render("human")
+        env.render()
         if args.gif:
-            frames.append(numpy.moveaxis(env.render("rgb_array"), 2, 0))
+            frames.append(numpy.moveaxis(env.get_frame(), 2, 0))
 
         action = agent.get_action(obs)
         obs, reward, done, _, _ = env.step(action)

diff --git a/utils/env.py b/utils/env.py
@@ -1,7 +1,10 @@
 import gymnasium as gym
 
 
-def make_env(env_key, seed=None):
-    env = gym.make(env_key)
+def make_env(env_key, seed=None, render_mode=None):
+    if render_mode is None:
+        env = gym.make(env_key)
+    else:
+        env = gym.make(env_key, render_mode=render_mode)
     env.reset(seed=seed)
     return env