File size: 1,330 Bytes
34ff700 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | import argparse
import os
import gymnasium as gym
import imageio.v2 as imageio
from enjoy import load_agent_from_checkpoint
def record(weights_path, output_path, seed=42, max_steps=1000):
agent = load_agent_from_checkpoint(weights_path, seed=seed)
env = gym.make("Pendulum-v1", render_mode="rgb_array")
frames = []
obs, _ = env.reset(seed=seed)
for _ in range(max_steps):
Hr, Hi = agent.encoder.encode(obs)
action = agent.actor.greedy(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 Pendulum-v1 replay video")
parser.add_argument("--weights", default="hdppo-Pendulum-v1/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()
|