Add hdppo-InvertedDoublePendulum-v5 package (weights, code, model card)

README.md

@@ -107,6 +107,7 @@ python train_hdppo.py
 |------|-------------|
 | `hdppo-InvertedDoublePendulum-v5/weights.npz` | Published actor (+ critic if HD) and FPE encoder (D=64) |
 | `hdppo-InvertedDoublePendulum-v5/weights_D512_teacher.npz` | Teacher checkpoint (D=512) |
+| `replay.mp4` | Sample rollout video from the published min-D checkpoint |
 | `results.json` | Evaluation summary for the published checkpoint |
 | `results_D512_teacher.json` | Evaluation summary for the teacher |
 | `config.yml` | Training hyperparameters |

record_video.py

@@ -1,44 +1,42 @@
-
-        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("InvertedDoublePendulum-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 InvertedDoublePendulum-v5 replay video")
-          parser.add_argument("--weights", default="hdppo-InvertedDoublePendulum-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()
+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("InvertedDoublePendulum-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 InvertedDoublePendulum-v5 replay video")
+    parser.add_argument("--weights", default="hdppo-InvertedDoublePendulum-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()