Add video recording support to eval_runner.py

isaaclab_arena/environments/arena_env_builder.py

@@ -309,12 +309,14 @@ def build_registered(
         )
         return name, cfg
 
-    def make_registered(self, env_cfg: None | IsaacLabArenaManagerBasedRLEnvCfg = None) -> ManagerBasedEnv:
-        env, _ = self.make_registered_and_return_cfg(env_cfg)
+    def make_registered(
+        self, env_cfg: None | IsaacLabArenaManagerBasedRLEnvCfg = None, render_mode: str | None = None
+    ) -> ManagerBasedEnv:
+        env, _ = self.make_registered_and_return_cfg(env_cfg, render_mode=render_mode)
         return env
 
     def make_registered_and_return_cfg(
-        self, env_cfg: None | IsaacLabArenaManagerBasedRLEnvCfg = None
+        self, env_cfg: None | IsaacLabArenaManagerBasedRLEnvCfg = None, render_mode: str | None = None
     ) -> tuple[ManagerBasedEnv, IsaacLabArenaManagerBasedRLEnvCfg]:
         name, cfg = self.build_registered(env_cfg)
-        return gym.make(name, cfg=cfg), cfg
+        return gym.make(name, cfg=cfg, render_mode=render_mode), cfg

isaaclab_arena/evaluation/eval_runner.py

@@ -8,6 +8,7 @@
 import json
 import os
 import traceback
+from gymnasium.wrappers import RecordVideo
 from typing import TYPE_CHECKING
 
 from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser
@@ -23,7 +24,7 @@
     from isaaclab_arena.policy.policy_base import PolicyBase
 
 
-def load_env(arena_env_args: list[str], job_name: str):
+def load_env(arena_env_args: list[str], job_name: str, render_mode: str | None = None):
 
     reload_arena_modules()
 
@@ -38,7 +39,7 @@ def load_env(arena_env_args: list[str], job_name: str):
     if hasattr(env_cfg, "recorders") and env_cfg.recorders is not None:
         env_cfg.recorders.dataset_filename = f"dataset_{job_name}"
 
-    env = arena_builder.make_registered(env_cfg)
+    env = arena_builder.make_registered(env_cfg, render_mode=render_mode)
     # Don't reset here - rollout_policy() will reset the env. Every reset triggers a new episode, initializing recorder & creating a new hdf5 entry.
     return env
 
@@ -113,12 +114,16 @@ def main():
 
         job_manager.print_jobs_info()
 
+        if args_cli.video:
+            os.makedirs(args_cli.video_dir, exist_ok=True)
+            print(f"[INFO] Video recording enabled. Videos will be saved to: {args_cli.video_dir}")
+
         for job in job_manager:
             if job is not None:
                 env = None
                 try:
-                    # Modules reloading first, otherwise 2 instances of same class are created (e.g. Enum)
-                    env = load_env(job.arena_env_args, job.name)
+                    render_mode = "rgb_array" if args_cli.video else None
+                    env = load_env(job.arena_env_args, job.name, render_mode=render_mode)
 
                     policy = get_policy_from_job(job)
 
@@ -129,6 +134,21 @@ def main():
                             job.num_steps = policy.length()
                         else:
                             job.num_steps = args_cli.num_steps
+
+                    if args_cli.video:
+                        if job.num_steps is not None:
+                            video_length = job.num_steps
+                        else:
+                            video_length = job.num_episodes * env.unwrapped.max_episode_length
+                        video_kwargs = {
+                            "video_folder": os.path.join(args_cli.video_dir, job.name),
+                            "step_trigger": lambda step: step == 0,
+                            "video_length": video_length,
+                            "disable_logger": True,
+                        }
+                        print(f"[INFO] Recording video for job '{job.name}' -> {video_kwargs['video_folder']}")
+                        env = RecordVideo(env, **video_kwargs)
+
                     metrics = rollout_policy(env, policy, num_steps=job.num_steps, num_episodes=job.num_episodes)
 
                     job_manager.complete_job(job, metrics=metrics, status=Status.COMPLETED)

isaaclab_arena/evaluation/eval_runner_cli.py

@@ -14,6 +14,13 @@ def add_eval_runner_arguments(parser: argparse.ArgumentParser) -> None:
         default="isaaclab_arena_environments/eval_jobs_configs/zero_action_jobs_config.json",
         help="Path to the eval jobs config file.",
     )
+    parser.add_argument("--video", action="store_true", default=False, help="Record videos for each eval job.")
+    parser.add_argument(
+        "--video_dir",
+        type=str,
+        default="/eval/videos",
+        help="Root directory for recorded videos. Each job gets a subdirectory.",
+    )
     parser.add_argument(
         "--continue_on_error",
         action="store_true",