From 93bf5e362da30fc19bb0d608618ed8f04ebcf58e Mon Sep 17 00:00:00 2001 From: Qian Lin Date: Wed, 9 Sep 2026 15:49:48 +0800 Subject: [PATCH] Migrate viewport recording to Isaac Lab Signed-off-by: Qian Lin --- .../evaluation/experiment_runner.py | 6 +- isaaclab_arena/evaluation/policy_runner.py | 33 +++++---- isaaclab_arena/evaluation/run_execution.py | 19 +++--- isaaclab_arena/tests/test_run_execution.py | 8 +-- isaaclab_arena/tests/test_video_recording.py | 48 +++++++++++++ isaaclab_arena/video/video_recording.py | 67 +++++++++++-------- 6 files changed, 126 insertions(+), 55 deletions(-) create mode 100644 isaaclab_arena/tests/test_video_recording.py diff --git a/isaaclab_arena/evaluation/experiment_runner.py b/isaaclab_arena/evaluation/experiment_runner.py index 7ab6b3c9a8..a0c6e4e232 100644 --- a/isaaclab_arena/evaluation/experiment_runner.py +++ b/isaaclab_arena/evaluation/experiment_runner.py @@ -104,8 +104,10 @@ def main(): ) # AppLauncher must enable camera support before SimulationApp starts. Check if this is required. - if args_cli.record_camera_video or _experiment_requires_cameras( - experiment_config_path, legacy_experiment_config, experiment_overrides + if ( + args_cli.record_viewport_video + or args_cli.record_camera_video + or _experiment_requires_cameras(experiment_config_path, legacy_experiment_config, experiment_overrides) ): args_cli.enable_cameras = True diff --git a/isaaclab_arena/evaluation/policy_runner.py b/isaaclab_arena/evaluation/policy_runner.py index d8323b0da1..bf4913ca87 100644 --- a/isaaclab_arena/evaluation/policy_runner.py +++ b/isaaclab_arena/evaluation/policy_runner.py @@ -22,7 +22,12 @@ from isaaclab_arena.utils.hydra_overrides import assert_hydra_overrides from isaaclab_arena.utils.isaaclab_utils.simulation_app import SimulationAppContext from isaaclab_arena.utils.multiprocess import get_local_rank, get_world_size -from isaaclab_arena.video.video_recording import VideoRecordingCfg, timestamped_run_dir, wrap_env_for_video +from isaaclab_arena.video.video_recording import ( + VideoRecordingCfg, + configure_env_for_video, + timestamped_run_dir, + wrap_env_for_video, +) from isaaclab_arena.visualization.report import build_report, serve_until_ctrl_c from isaaclab_arena_environments.cli import get_arena_builder_from_cli, get_isaaclab_arena_environments_cli_parser @@ -161,8 +166,8 @@ def main(): args_cli.device = f"cuda:{local_rank}" print(f"[Rank {local_rank}/{world_size}] One Isaac Lab instance per process on cuda:{local_rank}") - # --record_camera_video requires cameras to be enabled at sim startup, before SimulationAppContext. - if "--record_camera_video" in unknown: + # Video capture through Kit requires cameras to be enabled before SimulationApp starts. + if "--record_camera_video" in unknown or "--record_viewport_video" in unknown: args_cli.enable_cameras = True with SimulationAppContext(args_cli): @@ -198,10 +203,10 @@ def main(): args_cli.seed += local_rank # Re-apply enable_cameras: the full parse resets it to default False. - if args_cli.record_camera_video: + if args_cli.record_camera_video or args_cli.record_viewport_video: args_cli.enable_cameras = True - # Build scene. Use rgb_array render mode when recording so RecordVideo can grab frames. + # Build the policy before the environment so its intrinsic length can configure native video recording. arena_builder = get_arena_builder_from_cli(args_cli, hydra_overrides=hydra_overrides) output_dir = timestamped_run_dir(args_cli.output_base_dir) @@ -210,12 +215,6 @@ def main(): record_camera_video=args_cli.record_camera_video, video_base_dir=output_dir, ) - env = arena_builder.make_registered(render_mode=video_cfg.render_mode) - - # Write per-episode results to disk. - results_path = os.path.join(output_dir, f"episode_results_rank{local_rank}.jsonl") - env.unwrapped.episode_recorder.set_job_name("policy_runner") - env.unwrapped.episode_recorder.set_output_path(results_path) # Create the policy through the typed config compatibility adapter. policy = build_policy_from_cli(policy_cls, args_cli) @@ -236,7 +235,17 @@ def main(): else: raise ValueError(f"[Rank {local_rank}/{world_size}] Either num_steps or num_episodes must be provided") - # Optionally wrap with the viewport/camera video recorders (both independent). + # Configure native viewport recording before the environment creates its video recorders. + _, env_cfg, env_kwargs = arena_builder.build_registered() + configure_env_for_video(env_cfg, video_cfg, num_steps, num_episodes) + env = arena_builder.make_registered(env_cfg, env_kwargs) + + # Write per-episode results to disk. + results_path = os.path.join(output_dir, f"episode_results_rank{local_rank}.jsonl") + env.unwrapped.episode_recorder.set_job_name("policy_runner") + env.unwrapped.episode_recorder.set_output_path(results_path) + + # Optionally wrap with the camera-observation recorder. env = wrap_env_for_video(env, video_cfg, num_steps, num_episodes) steps_str = f"{num_steps} steps" if num_steps is not None else f"{num_episodes} episodes" diff --git a/isaaclab_arena/evaluation/run_execution.py b/isaaclab_arena/evaluation/run_execution.py index eb8f4950ff..4ea475aeff 100644 --- a/isaaclab_arena/evaluation/run_execution.py +++ b/isaaclab_arena/evaluation/run_execution.py @@ -25,7 +25,7 @@ from isaaclab_arena.evaluation.resource_cleanup import close_run_resources from isaaclab_arena.metrics.aggregate_metrics import aggregate_metrics from isaaclab_arena.variations.variations_hydra import overrides_from_dict -from isaaclab_arena.video.video_recording import VideoRecordingCfg, wrap_env_for_video +from isaaclab_arena.video.video_recording import VideoRecordingCfg, configure_env_for_video, wrap_env_for_video if TYPE_CHECKING: import gymnasium as gym @@ -105,17 +105,17 @@ def build_and_run( camera_name_prefix=f"robot-cam-rebuild{rebuild_index}", ) rebuild_cfg = _seed_cfg_for_rebuild(cfg, rebuild_index) - env = _build_environment_from_cfg(rebuild_cfg, rebuild_video_cfg.render_mode) - results_path = os.path.join(output_dir, f"episode_results_rebuild{rebuild_index}.jsonl") - env.unwrapped.episode_recorder.set_job_name(cfg.name) - env.unwrapped.episode_recorder.set_output_path(results_path) - policy = _build_policy_from_cfg(rebuild_cfg) num_steps, num_episodes = _resolve_rollout_limit( cfg, policy, num_episodes, ) + env = _build_environment_from_cfg(rebuild_cfg, rebuild_video_cfg, num_steps, num_episodes) + results_path = os.path.join(output_dir, f"episode_results_rebuild{rebuild_index}.jsonl") + env.unwrapped.episode_recorder.set_job_name(cfg.name) + env.unwrapped.episode_recorder.set_output_path(results_path) + env = wrap_env_for_video(env, rebuild_video_cfg, num_steps, num_episodes) metrics = rollout_policy(env, policy, num_steps=num_steps, num_episodes=num_episodes) if metrics is not None: @@ -139,14 +139,17 @@ def _seed_cfg_for_rebuild(cfg: ArenaRunCfg, rebuild_index: int) -> ArenaRunCfg: def _build_environment_from_cfg( cfg: ArenaRunCfg, - render_mode: str | None, + video_cfg: VideoRecordingCfg, + num_steps: int | None, + num_episodes: int | None, ) -> gym.Env: """Compile and instantiate a run's environment.""" arena_builder = build_arena_builder_from_run_cfg(cfg) _, env_cfg, env_kwargs = arena_builder.build_registered() if env_cfg.recorders is not None: env_cfg.recorders.dataset_filename = f"dataset_{cfg.name}" - return arena_builder.make_registered(env_cfg, env_kwargs, render_mode=render_mode) + configure_env_for_video(env_cfg, video_cfg, num_steps, num_episodes) + return arena_builder.make_registered(env_cfg, env_kwargs) def build_arena_builder_from_run_cfg(cfg: ArenaRunCfg) -> ArenaEnvBuilder: diff --git a/isaaclab_arena/tests/test_run_execution.py b/isaaclab_arena/tests/test_run_execution.py index d0ffc87c48..bba296a866 100644 --- a/isaaclab_arena/tests/test_run_execution.py +++ b/isaaclab_arena/tests/test_run_execution.py @@ -64,7 +64,7 @@ def test_build_and_run_splits_episode_budget_without_mutating_config(monkeypatch rollout_limits = [] received_run_cfgs = [] - def make_environment(cfg, render_mode): + def make_environment(cfg, video_cfg, num_steps, num_episodes): received_run_cfgs.append(cfg) return _environment() @@ -117,7 +117,7 @@ def test_build_and_run_raises_and_closes_resources(monkeypatch, tmp_path): monkeypatch.setattr( run_execution, "_build_environment_from_cfg", - lambda cfg, render_mode: environment, + lambda cfg, video_cfg, num_steps, num_episodes: environment, ) monkeypatch.setattr(run_execution, "_build_policy_from_cfg", lambda cfg: policy) monkeypatch.setattr(run_execution, "wrap_env_for_video", lambda env, video_cfg, steps, episodes: env) @@ -149,7 +149,7 @@ def test_build_and_run_requires_a_limit_for_an_unbounded_policy(monkeypatch, tmp monkeypatch.setattr( run_execution, "_build_environment_from_cfg", - lambda cfg, render_mode: environment, + lambda cfg, video_cfg, num_steps, num_episodes: environment, ) monkeypatch.setattr(run_execution, "_build_policy_from_cfg", lambda cfg: policy) monkeypatch.setattr( @@ -164,7 +164,7 @@ def test_build_and_run_requires_a_limit_for_an_unbounded_policy(monkeypatch, tmp output_dir=tmp_path, ) - assert closed_resources == [(policy, environment)] + assert closed_resources == [(policy, None)] def test_execute_experiment_runs_in_declaration_order(monkeypatch, tmp_path): diff --git a/isaaclab_arena/tests/test_video_recording.py b/isaaclab_arena/tests/test_video_recording.py new file mode 100644 index 0000000000..4064a0d46f --- /dev/null +++ b/isaaclab_arena/tests/test_video_recording.py @@ -0,0 +1,48 @@ +# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +from isaaclab_arena.video.video_recording import VideoRecordingCfg, configure_env_for_video + + +def _env_cfg(): + return SimpleNamespace( + episode_length_s=10.0, + decimation=2, + sim=SimpleNamespace(dt=0.1), + video_recorders=[], + ) + + +def test_configure_env_for_viewport_video_uses_native_recorder(tmp_path): + env_cfg = _env_cfg() + + configure_env_for_video( + env_cfg, + VideoRecordingCfg(record_viewport_video=True, video_base_dir=str(tmp_path)), + num_steps=25, + num_episodes=None, + ) + + assert len(env_cfg.video_recorders) == 1 + recorder_cfg = env_cfg.video_recorders[0] + assert recorder_cfg.source == "visualizer:kit" + assert recorder_cfg.output_dir == str(tmp_path) + assert recorder_cfg.output_filename_prefix == "viewport" + assert recorder_cfg.video_length == 25 + + +def test_configure_env_for_viewport_video_sizes_episode_rollout(tmp_path): + env_cfg = _env_cfg() + + configure_env_for_video( + env_cfg, + VideoRecordingCfg(record_viewport_video=True, video_base_dir=str(tmp_path)), + num_steps=None, + num_episodes=3, + ) + + assert env_cfg.video_recorders[0].video_length == 150 diff --git a/isaaclab_arena/video/video_recording.py b/isaaclab_arena/video/video_recording.py index 2c51f359db..165d34d5b6 100644 --- a/isaaclab_arena/video/video_recording.py +++ b/isaaclab_arena/video/video_recording.py @@ -7,6 +7,7 @@ import dataclasses import datetime +import math import os @@ -31,11 +32,6 @@ def enabled(self) -> bool: """Whether any recorder is requested.""" return self.record_viewport_video or self.record_camera_video - @property - def render_mode(self) -> str | None: - """The ``render_mode`` the env must be built with to capture the viewport video.""" - return "rgb_array" if self.record_viewport_video else None - def timestamped_run_dir(base_dir: str) -> str: """Append a reverse-dated subdirectory to ``base_dir``, e.g. ``base_dir/2026-06-16_14-42-54``. @@ -47,15 +43,42 @@ def timestamped_run_dir(base_dir: str) -> str: return os.path.join(base_dir, timestamp) -def _resolve_video_length(env, num_steps: int | None, num_episodes: int | None) -> int: - """Number of env steps to record: the step budget, or one episode's worth per episode. +def _resolve_video_length(env_cfg, num_steps: int | None, num_episodes: int | None) -> int: + """Number of env steps to record from the rollout limit and environment config. - ``max_episode_length`` is in environment steps, which matches the rollout cadence. + ``episode_length_s / (sim.dt * decimation)`` is the configured maximum episode + length in environment steps. """ if num_steps is not None: return num_steps assert num_episodes is not None, "Cannot determine video length: both num_steps and num_episodes are None." - return num_episodes * env.unwrapped.max_episode_length + max_episode_length = math.ceil(env_cfg.episode_length_s / (env_cfg.sim.dt * env_cfg.decimation)) + return num_episodes * max_episode_length + + +def configure_env_for_video( + env_cfg, + video_cfg: VideoRecordingCfg, + num_steps: int | None, + num_episodes: int | None, +) -> None: + """Add requested native video recorders to an Isaac Lab environment config.""" + if not video_cfg.record_viewport_video: + return + + from isaaclab.envs.utils.video_recorder_cfg import VideoRecorderCfg + + os.makedirs(video_cfg.video_base_dir, exist_ok=True) + video_length = _resolve_video_length(env_cfg, num_steps, num_episodes) + env_cfg.video_recorders.append( + VideoRecorderCfg( + source="visualizer:kit", + output_dir=video_cfg.video_base_dir, + output_filename_prefix="viewport", + video_length=video_length, + ) + ) + print(f"Recording {video_length}-step viewport video to: {video_cfg.video_base_dir}") def wrap_env_for_video( @@ -64,36 +87,22 @@ def wrap_env_for_video( num_steps: int | None, num_episodes: int | None, ): - """Wrap ``env`` with the recorders enabled in ``video_cfg`` and return the wrapped env. + """Wrap ``env`` with the camera-observation recorder when requested. - Returns ``env`` unchanged when no recorder is requested. ``num_steps`` and ``num_episodes`` - are mutually exclusive and size the viewport video. + Viewport recording is configured natively before environment construction by + :func:`configure_env_for_video`. Args: env: The env to wrap. video_cfg: The video recording configuration struct. - num_steps: Step budget for the rollout, or ``None`` when episode-driven. - num_episodes: Episode budget for the rollout, or ``None`` when step-driven. + num_steps: Unused; retained for call-site compatibility. + num_episodes: Unused; retained for call-site compatibility. """ - if not video_cfg.enabled: + if not video_cfg.record_camera_video: return env os.makedirs(video_cfg.video_base_dir, exist_ok=True) - # Record the kit viewport (via env.render()). - if video_cfg.record_viewport_video: - from gymnasium.wrappers import RecordVideo - - video_length = _resolve_video_length(env, num_steps, num_episodes) - env = RecordVideo( - env, - video_folder=video_cfg.video_base_dir, - step_trigger=lambda step: step == 0, - video_length=video_length, - disable_logger=True, - ) - print(f"Recording {video_length}-step viewport video to: {video_cfg.video_base_dir}") - # Record the embodiment-mounted cameras (from obs["camera_obs"]), # flushed at each episode reset rather than after a fixed number of steps. if video_cfg.record_camera_video: