From 7dcb2be5c1b17eb7f951b2a3c7b3d79323227dff Mon Sep 17 00:00:00 2001
From: Salem <salemlahlou9@gmail.com>
Date: Wed, 5 Oct 2022 12:31:47 -0400
Subject: [PATCH] fix visualize.py

---
 scripts/visualize.py | 8 ++++----
 utils/env.py         | 7 +++++--
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/scripts/visualize.py b/scripts/visualize.py
index cd9ceac4..664c3131 100644
--- 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
index 2327d3db..a866a6a0 100644
--- 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