Recording of multiple Breakout games from a Gymnasium Vectorized Environment

78 views Asked by At

I want to record the interaction with a vectorized environment of two games of Atari's Breakout using Gymnasium. Unfortunately, when executed, my code only records the last game of the vector during an episode (or maybe it internally overwrites recordings until the last recording). I am looking for a solution to record all games of a vector for each episode. Here is my current code:

import gymnasium as gym
from gymnasium.wrappers.record_video import RecordVideo

from pyvirtualdisplay import Display
from IPython.display import display, Video

wrap_record_video=lambda x: RecordVideo(x,
                  video_folder="recordings/",
                  name_prefix="breakout_test",
                  episode_trigger=lambda x: True,
                  disable_logger=True)


env = gym.make_vec('Breakout-v4', 
                   render_mode='rgb_array', 
                   num_envs=2, 
                   vectorization_mode="async", 
                   wrappers=[wrap_record_video])

d = Display(visible=0, size=(800,600))
d.start()

for episode in range(0, 5):
  state = env.reset()
  done = False
  score = 0

  while not done:
    action = env.action_space.sample() 
    n_state, reward, terminated, truncated, info = env.step(action)
    done = True in terminated or True in truncated
    score += reward

  print("Episode:{} Score:{}".format(episode, score))
  display(Video("recordings/breakout_test-episode-{}.mp4".format(episode), embed=True))

env.close()
d.stop() 

I have seen examples that wrap the vectorized environment into a VecFrameStack environment of stable-baselines3. However, VecFrameStack requires the vector to be of class stable_baselines3.common.vec_env.VecEnv, while the output of gym.make_vec is of class gymnasium.vector.VectorEnv.

0

There are 0 answers