Skip to content

Commit 92e05c6

Browse files
cvolkcvolkclaude
andauthored
Add video recording support to eval_runner.py (#445)
## Summary Adds the ability to record videos of evaluation jobs via `--video` flag. Each job produces an mp4 in a per-job subdirectory under `--video_dir` ``` python isaaclab_arena/evaluation/eval_runner.py \ --eval_jobs_config isaaclab_arena_environments/eval_jobs_configs/droid_pnp_srl_gr00t_jobs_config.json \ --video ``` --------- Signed-off-by: Clemens Volk <cvolk@nvidia.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 04a157f commit 92e05c6

File tree

3 files changed

+37
-8
lines changed

3 files changed

+37
-8
lines changed

isaaclab_arena/environments/arena_env_builder.py

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -309,12 +309,14 @@ def build_registered(
309309
)
310310
return name, cfg
311311

312-
def make_registered(self, env_cfg: None | IsaacLabArenaManagerBasedRLEnvCfg = None) -> ManagerBasedEnv:
313-
env, _ = self.make_registered_and_return_cfg(env_cfg)
312+
def make_registered(
313+
self, env_cfg: None | IsaacLabArenaManagerBasedRLEnvCfg = None, render_mode: str | None = None
314+
) -> ManagerBasedEnv:
315+
env, _ = self.make_registered_and_return_cfg(env_cfg, render_mode=render_mode)
314316
return env
315317

316318
def make_registered_and_return_cfg(
317-
self, env_cfg: None | IsaacLabArenaManagerBasedRLEnvCfg = None
319+
self, env_cfg: None | IsaacLabArenaManagerBasedRLEnvCfg = None, render_mode: str | None = None
318320
) -> tuple[ManagerBasedEnv, IsaacLabArenaManagerBasedRLEnvCfg]:
319321
name, cfg = self.build_registered(env_cfg)
320-
return gym.make(name, cfg=cfg), cfg
322+
return gym.make(name, cfg=cfg, render_mode=render_mode), cfg

isaaclab_arena/evaluation/eval_runner.py

Lines changed: 24 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import json
99
import os
1010
import traceback
11+
from gymnasium.wrappers import RecordVideo
1112
from typing import TYPE_CHECKING
1213

1314
from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser
@@ -23,7 +24,7 @@
2324
from isaaclab_arena.policy.policy_base import PolicyBase
2425

2526

26-
def load_env(arena_env_args: list[str], job_name: str):
27+
def load_env(arena_env_args: list[str], job_name: str, render_mode: str | None = None):
2728

2829
reload_arena_modules()
2930

@@ -38,7 +39,7 @@ def load_env(arena_env_args: list[str], job_name: str):
3839
if hasattr(env_cfg, "recorders") and env_cfg.recorders is not None:
3940
env_cfg.recorders.dataset_filename = f"dataset_{job_name}"
4041

41-
env = arena_builder.make_registered(env_cfg)
42+
env = arena_builder.make_registered(env_cfg, render_mode=render_mode)
4243
# Don't reset here - rollout_policy() will reset the env. Every reset triggers a new episode, initializing recorder & creating a new hdf5 entry.
4344
return env
4445

@@ -113,12 +114,16 @@ def main():
113114

114115
job_manager.print_jobs_info()
115116

117+
if args_cli.video:
118+
os.makedirs(args_cli.video_dir, exist_ok=True)
119+
print(f"[INFO] Video recording enabled. Videos will be saved to: {args_cli.video_dir}")
120+
116121
for job in job_manager:
117122
if job is not None:
118123
env = None
119124
try:
120-
# Modules reloading first, otherwise 2 instances of same class are created (e.g. Enum)
121-
env = load_env(job.arena_env_args, job.name)
125+
render_mode = "rgb_array" if args_cli.video else None
126+
env = load_env(job.arena_env_args, job.name, render_mode=render_mode)
122127

123128
policy = get_policy_from_job(job)
124129

@@ -129,6 +134,21 @@ def main():
129134
job.num_steps = policy.length()
130135
else:
131136
job.num_steps = args_cli.num_steps
137+
138+
if args_cli.video:
139+
if job.num_steps is not None:
140+
video_length = job.num_steps
141+
else:
142+
video_length = job.num_episodes * env.unwrapped.max_episode_length
143+
video_kwargs = {
144+
"video_folder": os.path.join(args_cli.video_dir, job.name),
145+
"step_trigger": lambda step: step == 0,
146+
"video_length": video_length,
147+
"disable_logger": True,
148+
}
149+
print(f"[INFO] Recording video for job '{job.name}' -> {video_kwargs['video_folder']}")
150+
env = RecordVideo(env, **video_kwargs)
151+
132152
metrics = rollout_policy(env, policy, num_steps=job.num_steps, num_episodes=job.num_episodes)
133153

134154
job_manager.complete_job(job, metrics=metrics, status=Status.COMPLETED)

isaaclab_arena/evaluation/eval_runner_cli.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,13 @@ def add_eval_runner_arguments(parser: argparse.ArgumentParser) -> None:
1414
default="isaaclab_arena_environments/eval_jobs_configs/zero_action_jobs_config.json",
1515
help="Path to the eval jobs config file.",
1616
)
17+
parser.add_argument("--video", action="store_true", default=False, help="Record videos for each eval job.")
18+
parser.add_argument(
19+
"--video_dir",
20+
type=str,
21+
default="/eval/videos",
22+
help="Root directory for recorded videos. Each job gets a subdirectory.",
23+
)
1724
parser.add_argument(
1825
"--continue_on_error",
1926
action="store_true",

0 commit comments

Comments
 (0)