ALE-Pacman-v5 / record_video.py
ledmands
Changed the video recorder directory
2690fb6
raw
history blame
1.87 kB
import gymnasium as gym
from stable_baselines3 import DQN
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.vec_env import VecVideoRecorder, DummyVecEnv, VecEnv
model_name = "agents/dqn_v2-7/ALE-Pacman-v5" # path to model, should be an argument
env_id = "ALE/Pacman-v5"
video_folder = "videos/"
video_length = 100000 #steps by hard coding this, I can almost ensure only one episode is recorded...
vec_env = DummyVecEnv([lambda: gym.make(env_id, render_mode="rgb_array")])
model = DQN.load(model_name)
# output: <stable_baselines3.common.vec_env.dummy_vec_env.DummyVecEnv object at 0x0000029974DC6550>
# vec_env = gym.make(env_id, render_mode="rgb_array")
# output <OrderEnforcing<PassiveEnvChecker<AtariEnv<ALE/Pacman-v5>>>>
# vec_env = Monitor(gym.make(env_id, render_mode="rgb_array"))
print("\n\n\n")
print(vec_env)
print("\n\n\n")
obs = vec_env.reset()
# Record the video starting at the first step
vec_env = VecVideoRecorder(vec_env,
video_folder,
record_video_trigger=lambda x: x == 0,
video_length=video_length,
name_prefix=f"video-{model_name}"
)
# Once I make the environment, now I need to walk through it...???
# I want to act according to the policy that has been trained
obs = vec_env.reset()
print(vec_env)
# for _ in range(video_length + 1):
# action, states = model.predict(obs)
# obs, _, _, _ = vec_env.step(action)
# Instead of using the specified steps in a for loop
# use a while loop to check if the episode has terminated
# Stop recording when the episode ends
end = True
while end == True:
action, states = model.predict(obs)
obs, _, done, _ = vec_env.step(action)
if done == True:
print("exiting loop")
end = False
# # Save the video
vec_env.close()