Add video recording support to eval_runner.py

isaaclab_arena/environments/arena_env_builder.py

@@ -289,12 +289,17 @@ 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).unwrapped, cfg
+        env = gym.make(name, cfg=cfg, render_mode=render_mode)
+        if render_mode is None:
+            env = env.unwrapped
+        return env, 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,3 +14,10 @@ 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.",
+    )

isaaclab_arena/evaluation/policy_runner.py

@@ -63,11 +63,13 @@ def rollout_policy(
     assert num_steps is not None or num_episodes is not None, "Either num_steps or num_episodes must be provided"
     assert num_steps is None or num_episodes is None, "Only one of num_steps or num_episodes must be provided"
 
+    pbar = None
     try:
         obs, _ = env.reset()
         policy.reset()
         # set task description (could be None) from the task being evaluated
-        policy.set_task_description(env.cfg.isaaclab_arena_env.task.get_task_description())
+        # Use unwrapped to reach the base env through any gym wrappers (e.g. RecordVideo)
+        policy.set_task_description(env.unwrapped.cfg.isaaclab_arena_env.task.get_task_description())
 
         # Setup progress bar based on num_steps or num_episodes
         if num_steps is not None:
@@ -107,16 +109,17 @@ def rollout_policy(
         pbar.close()
 
     except Exception as e:
-        pbar.close()
+        if pbar is not None:
+            pbar.close()
         raise RuntimeError(f"Error rolling out policy: {e}")
 
     else:
         # only compute metrics if env has metrics registered
-        if hasattr(env.cfg, "metrics"):
+        if hasattr(env.unwrapped.cfg, "metrics"):
             # NOTE(xinjieyao, 2025-10-07): lazy import to prevent app stalling caused by omni.kit
             from isaaclab_arena.metrics.metrics import compute_metrics
 
-            metrics = compute_metrics(env)
+            metrics = compute_metrics(env.unwrapped)
             return metrics
         return None