record_video.py

import argparse
import os

import gymnasium as gym
import imageio.v2 as imageio

from enjoy import load_policy_from_checkpoint


def record(weights_path, output_path, seed=42, max_steps=1000):
    encoder, actor, cfg = load_policy_from_checkpoint(weights_path, seed=seed)
    env_kwargs = dict(cfg.get("env_kwargs", {}))
    env_kwargs["render_mode"] = "rgb_array"
    env = gym.make("Ant-v5", **env_kwargs)
    frames = []
    obs, _ = env.reset(seed=seed)
    for _ in range(max_steps):
        Hr, Hi = encoder.encode(obs)
        action = actor.greedy_action_np(Hr, Hi)
        obs, reward, term, trunc, _ = env.step(action)
        frames.append(env.render())
        if term or trunc:
            break
    env.close()
    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
    imageio.mimsave(output_path, frames, fps=30)
    return len(frames)


def main():
    parser = argparse.ArgumentParser(description="Record Ant-v5 replay video")
    parser.add_argument("--weights", default="hdppo-Ant-v5/weights.npz")
    parser.add_argument("--output", default="replay.mp4")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--max-steps", type=int, default=1000)
    args = parser.parse_args()
    n = record(args.weights, args.output, seed=args.seed, max_steps=args.max_steps)
    print(f"saved {n} frames to {args.output}")


if __name__ == "__main__":
    main()