record.py

import gymnasium as gym
import ale_py
import json
from stable_baselines3 import PPO
from stable_baselines3.common.evaluation import evaluate_policy
from stable_baselines3.common.atari_wrappers import AtariWrapper
from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.vec_env import VecFrameStack

import imageio

if __name__ == "__main__":

    gym.register_envs(ale_py)
    env = make_vec_env(env_id="ALE/Enduro-v5", wrapper_class=AtariWrapper, wrapper_kwargs={"frame_skip":1})
    env = VecFrameStack(env, n_stack=4)
    
    agent = PPO.load("model.zip")
    
    frames = []

    obs = env.reset()

    for _ in range(1500):
        action, _ = agent.predict(obs)
        obs, _, done, _ = env.step(action)
        frame = env.render()
        frames.append(frame)
        if done:
            break

    env.close()

    # Save as video
    imageio.mimsave("replay.mp4", frames, fps=30)