diff --git a/script/main.py b/script/main.py
index 7f9f9aa0..4963f322 100644
--- a/script/main.py
+++ b/script/main.py
@@ -1,10 +1,11 @@
 from metadrive.envs.metadrive_env import MetaDriveEnv
 import gymnasium as gym
+from metadrive.envs.gym_wrapper import createGymWrapper # import the wrapper
 
-env = MetaDriveEnv(config={"use_render": True})
-obs, info = env.reset()
+env = createGymWrapper(MetaDriveEnv)(config={"use_render": True}) # wrap the environment
+obs = env.reset()
 for i in range(1000):
-    obs, reward, terminated, truncated, info = env.step(env.action_space.sample())
-    if terminated or truncated:
+    obs, reward, done, info = env.step(env.action_space.sample()) # the return value contains no truncate
+    if done:
         env.reset()
 env.close()
\ No newline at end of file
diff --git a/script/requirements.txt b/script/requirements.txt
index f84ab97d..c6cd7a39 100644
--- a/script/requirements.txt
+++ b/script/requirements.txt
@@ -1,2 +1,3 @@
 metadrive-simulator
-gymnasium
\ No newline at end of file
+gymnasium
+gym
\ No newline at end of file