OpenAI Gym Video Recorder

gym-recorder-demo.py

# This works for gym 0.26.2. It might not work on previous versions.

import gym
import cv2

fps = 30

# Create environment
env = gym.make('CartPole-v1', render_mode='rgb_array')

# Get initial observation and frame
obs = env.reset()
frame = env.render()
print("Frame shape:", frame.shape)
size = frame.shape

# Create recorder
out = cv2.VideoWriter('output.mp4', cv2.VideoWriter_fourcc(*'mp4v'), fps, (size[1], size[0]))
out.write(frame)

# Run episode
obs = env.reset()
for _ in range(500):
    frame = env.render()
    out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
    action = env.action_space.sample()
    obs, reward, terminated, truncated, info = env.step(action)
    if terminated:
        env.reset()
        continue
        
# Close out the environment
env.close()
out.release()
print("Done!")