Episode Recorder (#810)

## Summary
Introduce an `EpisodeRecorderManager` that records per-episode data.

## Detailed description
- Adds a per-episode data recorder to our `env` that accumulates
per-episode data and optionally writes it to disk.
- Users can register custom `EpisodeRecorderTermCfgs` to record their
custom data.
- Implemented 2 basic recorders.
  - core metadata recorder (env_id, episode_idx, success, etc.)
  - variations recorder (records samples for the variations)
- Added the recorded meta data into the evaluation report.

## Note
- Getting at the per environment success flag requires overriding
internal an internal method of `IsaacLabManagerBasedRLEnv`

## Not done
- Gracefully handle rebuilds and `run_in_chunks`. Will be handled in
later MRs.

---------

Signed-off-by: alex <amillane@nvidia.com>
Co-authored-by: aiguldzh-nvidia <adzhumamurat@nvidia.com>

Author: alexmillane

docs/pages/quickstart/first_experiments/running_a_real_policy/gr00t.rst

@@ -153,7 +153,7 @@ episode; the runner writes an ``index.html`` which is then served over HTTP.
    python isaaclab_arena/evaluation/eval_runner.py \
      --viz kit \
      --eval_jobs_config isaaclab_arena_environments/eval_jobs_configs/droid_pnp_srl_gr00t_jobs_config.json \
-     --video_base_dir ./output \
+     --output_base_dir ./output \
      --record_camera_video --serve_evaluation_report
 
 You can also (re)build and serve a report later by pointing the standalone tool at the output

docs/pages/quickstart/first_experiments/running_a_real_policy/openpi.rst

@@ -162,7 +162,7 @@ episode; the runner writes an ``index.html`` which is then served over HTTP.
    python isaaclab_arena/evaluation/eval_runner.py \
      --viz kit \
      --eval_jobs_config isaaclab_arena_environments/eval_jobs_configs/droid_pnp_srl_openpi_jobs_config.json \
-     --video_base_dir ./output \
+     --output_base_dir ./output \
      --record_camera_video --serve_evaluation_report
 
 You can also (re)build and serve a report later by pointing the standalone tool at the output

isaaclab_arena/environments/arena_env_builder.py

@@ -34,6 +34,8 @@
     make_progress_tracking_events_cfg,
     make_progress_tracking_recorder_cfg,
 )
+from isaaclab_arena.recording.common_terms import CoreEpisodeRecorderTermCfg, VariationEpisodeRecorderTermCfg
+from isaaclab_arena.recording.episode_recorder_manager import EpisodeRecorderTermCfg
 from isaaclab_arena.relations.placement_events import PLACEMENT_RESET_EVENT_NAME
 from isaaclab_arena.tasks.no_task import NoTask
 from isaaclab_arena.utils.configclass import combine_configclass_instances, make_configclass
@@ -152,14 +154,30 @@ def _modify_recorder_cfg_dataset_filename(self, recorder_cfg: RecorderManagerBas
         )
         return recorder_cfg
 
-    @staticmethod
-    def _metrics_to_metrics_cfg(metrics: list[MetricBase] | None) -> object | None:
+    def _compose_metrics_cfg(self, metrics: list[MetricBase] | None) -> object | None:
         """Build a configclass container with one ``MetricTermCfg`` field per metric."""
         if not metrics:
             return None
         fields = [(m.name, MetricTermCfg, m.get_metric_term_cfg()) for m in metrics]
         return make_configclass("MetricsCfg", fields)()
 
+    def _compose_episode_recorders_cfg(self, extra_terms: dict[str, EpisodeRecorderTermCfg] | None = None) -> object:
+        """Build a configclass container with one EpisodeRecorderTermCfg field per episode recorder term.
+
+        Note that this function automatically adds the core and variations terms.
+        """
+        fields = [
+            ("core", EpisodeRecorderTermCfg, CoreEpisodeRecorderTermCfg()),
+            ("variations", EpisodeRecorderTermCfg, VariationEpisodeRecorderTermCfg()),
+        ]
+        for name, term_cfg in (extra_terms or {}).items():
+            assert name not in (
+                "core",
+                "variations",
+            ), f"Episode recorder term name '{name}' collides with a built-in term."
+            fields.append((name, EpisodeRecorderTermCfg, term_cfg))
+        return make_configclass("EpisodeRecorderManagerCfg", fields)()
+
     def compose_manager_cfg(self) -> tuple[IsaacLabArenaManagerBasedRLEnvCfg, dict[str, Any]]:
         """Return the base ManagerBased cfg and the env kwargs (no registration).
 
@@ -240,7 +258,7 @@ def compose_manager_cfg(self) -> tuple[IsaacLabArenaManagerBasedRLEnvCfg, dict[s
             elif isinstance(device_cfg, DeviceCfg):
                 teleop_devices_cfg = DevicesCfg(devices={self.arena_env.teleop_device.name: device_cfg})
         metrics = task.get_metrics()
-        metrics_cfg = self._metrics_to_metrics_cfg(metrics)
+        metrics_cfg = self._compose_metrics_cfg(metrics)
         metrics_recorder_manager_cfg = metrics_to_recorder_manager_cfg(metrics)
         progress_tracking_recorder_cfg: Any = (
             make_progress_tracking_recorder_cfg(progress_objectives) if progress_objectives else None
@@ -278,11 +296,15 @@ def compose_manager_cfg(self) -> tuple[IsaacLabArenaManagerBasedRLEnvCfg, dict[s
             task.get_commands_cfg(),
         )
 
+        episode_recorders_cfg = self._compose_episode_recorders_cfg(self.arena_env.episode_recorder_terms)
+
         viewer_cfg = task.get_viewer_cfg()
 
         episode_length_s = task.get_episode_length_s()
 
-        task_description = task.get_task_description()
+        # Language instruction is optionally overridden on the CLI.
+        language_instruction = getattr(self.args, "language_instruction", None)
+        task_description = language_instruction or task.get_task_description()
 
         # Build the environment configuration
         if not self.args.mimic:
@@ -300,6 +322,7 @@ def compose_manager_cfg(self) -> tuple[IsaacLabArenaManagerBasedRLEnvCfg, dict[s
                 teleop_devices=teleop_devices_cfg,
                 recorders=recorder_manager_cfg,
                 metrics=metrics_cfg,
+                episode_recorders=episode_recorders_cfg,
                 task_description=task_description,
                 viewer=viewer_cfg,
             )

isaaclab_arena/environments/isaaclab_arena_environment.py

@@ -12,6 +12,7 @@
     from isaaclab_arena.assets.teleop_device_base import TeleopDeviceBase
     from isaaclab_arena.embodiments.embodiment_base import EmbodimentBase
     from isaaclab_arena.environments.isaaclab_arena_manager_based_env_cfg import IsaacLabArenaManagerBasedRLEnvCfg
+    from isaaclab_arena.recording.episode_recorder_manager import EpisodeRecorderTermCfg
     from isaaclab_arena.scene.scene import Scene
     from isaaclab_arena.tasks.task_base import TaskBase
 
@@ -29,6 +30,7 @@ def __init__(
         env_cfg_callback: Callable[IsaacLabArenaManagerBasedRLEnvCfg] | None = None,
         rl_framework_entry_point: str | None = None,
         rl_policy_cfg: str | None = None,
+        episode_recorder_terms: dict[str, EpisodeRecorderTermCfg] | None = None,
     ):
         """
         Args:
@@ -46,6 +48,8 @@ def __init__(
                 ``rl_policy_cfg`` is set.
             rl_policy_cfg: Import path to the RL policy config class, e.g.
                 ``"my_module:RLPolicyCfg"``.
+            episode_recorder_terms: Additional per-episode recorder terms to record alongside the
+                built-in ones, keyed by name.
         """
         self.name = name
         self.scene = scene
@@ -57,3 +61,4 @@ def __init__(
             raise ValueError("rl_framework_entry_point and rl_policy_cfg must both be set or both be None.")
         self.rl_framework_entry_point = rl_framework_entry_point
         self.rl_policy_cfg = rl_policy_cfg
+        self.episode_recorder_terms = episode_recorder_terms or {}

isaaclab_arena/environments/isaaclab_arena_manager_based_env.py

@@ -5,11 +5,14 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
+
 from isaaclab.envs import ManagerBasedRLEnv
 
 from isaaclab_arena.environments.isaaclab_arena_manager_based_env_cfg import IsaacLabArenaManagerBasedRLEnvCfg
 from isaaclab_arena.metrics.metric_data import MetricsDataCollection
 from isaaclab_arena.metrics.metrics_manager import MetricsManager
+from isaaclab_arena.recording.episode_recorder_manager import EpisodeRecorderManager
 from isaaclab_arena.variations.variation_recorder import VariationRecorder
 
 
@@ -26,6 +29,13 @@ def __init__(
         **kwargs,
     ):
         self._variation_recorder = variation_recorder
+        if variation_recorder is not None:
+            # Bind so run-time variation draws can be attributed to the current episode index.
+            variation_recorder.bind_env(self)
+        # Per-env count of completed episodes; advanced in ``_reset_idx``.
+        self._episode_counts: dict[int, int] = {}
+        # The initial reset touches every env before any episode has run; skip it.
+        self._first_reset = True
         super().__init__(cfg=cfg, render_mode=render_mode, **kwargs)
 
     @property
@@ -38,9 +48,41 @@ def variation_recorder(self) -> VariationRecorder | None:
             )
         return self._variation_recorder
 
+    @property
+    def episode_recorder(self) -> EpisodeRecorderManager:
+        """The per-episode recorder."""
+        return self.episode_recorder_manager
+
     def load_managers(self) -> None:
         super().load_managers()
         self.metrics_manager = MetricsManager(self.cfg.metrics, self)
+        self.episode_recorder_manager = EpisodeRecorderManager(self.cfg.episode_recorders, self)
+
+    def get_language_instruction(self) -> str | None:
+        """Return the language instruction that is passed to the policy."""
+        return self.cfg.task_description
+
+    def get_episode_index(self, env_id: int) -> int:
+        """Return the index of the current episode in ``env_id``."""
+        return self._episode_counts.get(env_id, 0)
+
+    def _advance_episode_indices(self, env_ids: Sequence[int]) -> None:
+        """Advance the per-env episode counter for each episode in ``env_ids``."""
+        for env_id in env_ids:
+            env_id = int(env_id)
+            self._episode_counts[env_id] = self._episode_counts.get(env_id, 0) + 1
+
+    def _reset_idx(self, env_ids: Sequence[int]) -> None:
+        # The initial reset touches every env before any episode has run; nothing to record or count.
+        if self._first_reset:
+            self._first_reset = False
+            super()._reset_idx(env_ids)
+            return
+        # Runs recorder before super() so the just-finished episode is still intact.
+        self.episode_recorder_manager.record_pre_reset(env_ids)
+        # Advance before super() so reset-mode variation draws are tagged with the episode they begin.
+        self._advance_episode_indices(env_ids)
+        super()._reset_idx(env_ids)
 
     def compute_metrics(self) -> MetricsDataCollection:
         """Compute all registered metrics.

isaaclab_arena/environments/isaaclab_arena_manager_based_env_cfg.py

@@ -64,6 +64,8 @@ class IsaacLabArenaManagerBasedRLEnvCfg(ManagerBasedRLEnvCfg):
 
     metrics: object | None = None
 
+    episode_recorders: object | None = None
+
     # Task language description
     task_description: str | None = None
 

isaaclab_arena/evaluation/eval_runner.py

@@ -38,17 +38,20 @@ def load_env(
     job_name: str,
     variations: list[str] | None = None,
     render_mode: str | None = None,
+    language_instruction: str | None = None,
 ):
 
     args_parser = get_isaaclab_arena_environments_cli_parser()
 
     arena_env_args_cli = args_parser.parse_args(arena_env_args)
+    # Optionally override the language instruction.
+    arena_env_args_cli.language_instruction = language_instruction
     arena_builder = get_arena_builder_from_cli(arena_env_args_cli, hydra_overrides=variations)
 
-    env_name, env_cfg, env_kwargs = arena_builder.build_registered()
+    _, env_cfg, env_kwargs = arena_builder.build_registered()
 
     # Set unique dataset filename for this job to avoid file locking conflicts
-    if hasattr(env_cfg, "recorders") and env_cfg.recorders is not None:
+    if env_cfg.recorders is not None:
         env_cfg.recorders.dataset_filename = f"dataset_{job_name}"
 
     env = arena_builder.make_registered(env_cfg, env_kwargs, render_mode=render_mode)
@@ -262,11 +265,11 @@ def main():
         # Always dated so every run produces its own report dir, recording or not.
         # TODO(alexmillane): Currently each chunk produces its own output directory.
         # We should use the same output directory for all chunks in the future.
-        run_video_dir = timestamped_run_dir(args_cli.video_base_dir)
+        run_output_dir = timestamped_run_dir(args_cli.output_base_dir)
 
         if args_cli.record_viewport_video:
-            os.makedirs(run_video_dir, exist_ok=True)
-            print(f"[INFO] Video recording enabled. Videos will be saved to: {run_video_dir}")
+            os.makedirs(run_output_dir, exist_ok=True)
+            print(f"[INFO] Video recording enabled. Videos will be saved to: {run_output_dir}")
 
         for job in job_manager:
             if job is None:
@@ -283,17 +286,30 @@ def main():
             # aggregate the metrics across rebuilds into a single result.
             for rebuild_idx in range(job.num_rebuilds):
                 try:
+                    job_output_dir = os.path.join(run_output_dir, job.name)
+
                     # Per-job video output directory; cameras are tagged with the rebuild index.
                     video_cfg = VideoRecordingCfg(
                         record_viewport_video=args_cli.record_viewport_video,
                         record_camera_video=args_cli.record_camera_video,
-                        video_base_dir=os.path.join(run_video_dir, job.name),
+                        video_base_dir=job_output_dir,
                         camera_name_prefix=f"robot-cam-rebuild{rebuild_idx}",
                     )
                     env = load_env(
-                        job.arena_env_args, job.name, variations=job.variations, render_mode=video_cfg.render_mode
+                        job.arena_env_args,
+                        job.name,
+                        variations=job.variations,
+                        render_mode=video_cfg.render_mode,
+                        language_instruction=job.language_instruction,
                     )
 
+                    # Write per-episode results to disk.
+                    # TODO: Aggregate the per-episode records across rebuilds into a single file,
+                    # as is done for the metrics below.
+                    results_path = os.path.join(job_output_dir, f"episode_results_rebuild{rebuild_idx}.jsonl")
+                    env.unwrapped.episode_recorder.set_job_name(job.name)
+                    env.unwrapped.episode_recorder.set_output_path(results_path)
+
                     policy = get_policy_from_job(job)
 
                     # Episodes allotted to this rebuild (None when the job is length-driven by steps).
@@ -314,7 +330,6 @@ def main():
                         policy,
                         num_steps=job.num_steps,
                         num_episodes=num_episodes_this_rebuild,
-                        language_instruction=job.language_instruction,
                     )
 
                     job_manager.complete_job(job, metrics=metrics, status=Status.COMPLETED)
@@ -347,7 +362,7 @@ def main():
         metrics_logger.print_metrics()
 
         # Write HTML report.
-        report_path = build_report(run_video_dir)
+        report_path = build_report(run_output_dir)
         if args_cli.serve_evaluation_report:
             serve_until_ctrl_c(report_path.parent, args_cli.evaluation_report_port, report_path.name)
 

isaaclab_arena/evaluation/eval_runner_cli.py

@@ -27,10 +27,13 @@ def add_eval_runner_arguments(parser: argparse.ArgumentParser) -> None:
         help="Record one mp4 per (env, camera, episode) from obs['camera_obs'] for each eval job.",
     )
     parser.add_argument(
-        "--video_base_dir",
+        "--output_base_dir",
         type=str,
-        default="/eval/videos",
-        help="Base directory for recorded videos; a reverse-dated run subdirectory and per-job subdirectory are added.",
+        default="/eval/output",
+        help=(
+            "Base directory for evaluation outputs (videos, per-episode results, report); a"
+            " reverse-dated run subdirectory and per-job subdirectory are added."
+        ),
     )
     parser.add_argument(
         "--serve_evaluation_report",

isaaclab_arena/evaluation/policy_runner.py

@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 import argparse
+import os
 import torch
 import tqdm
 from importlib import import_module
@@ -64,7 +65,6 @@ def rollout_policy(
     policy: PolicyBase,
     num_steps: int | None,
     num_episodes: int | None,
-    language_instruction: str | None = None,
 ) -> MetricsDataCollection | None:
     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"
@@ -73,10 +73,7 @@ def rollout_policy(
     try:
         obs, _ = env.reset()
         policy.reset()
-        # Determine language instruction: CLI/job-level override takes precedence over the task's own
-        # description. Use unwrapped to reach the base env through any gym wrappers (e.g. OrderEnforcing).
-        task_description = language_instruction or env.unwrapped.cfg.task_description
-        policy.set_task_description(task_description)
+        policy.set_task_description(env.unwrapped.get_language_instruction())
 
         # Setup progress bar based on num_steps or num_episodes
         if num_steps is not None:
@@ -192,12 +189,18 @@ def main():
             print(arena_builder.get_variations_catalogue_as_string())
             return
 
+        output_dir = timestamped_run_dir(args_cli.output_base_dir)
         video_cfg = VideoRecordingCfg(
             record_viewport_video=args_cli.record_viewport_video,
             record_camera_video=args_cli.record_camera_video,
-            video_base_dir=timestamped_run_dir(args_cli.video_base_dir),
+            video_base_dir=output_dir,
         )
-        env, cfg = arena_builder.make_registered_and_return_cfg(render_mode=video_cfg.render_mode)
+        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 from the arguments
         policy = policy_cls.from_args(args_cli)
@@ -223,7 +226,7 @@ def main():
 
         steps_str = f"{num_steps} steps" if num_steps is not None else f"{num_episodes} episodes"
         print(f"[Rank {local_rank}/{world_size}] Starting rollout ({steps_str})")
-        metrics = rollout_policy(env, policy, num_steps, num_episodes, args_cli.language_instruction)
+        metrics = rollout_policy(env, policy, num_steps, num_episodes)
 
         if metrics is not None:
             print(f"[Rank {local_rank}/{world_size}] Metrics: {metrics_to_plain_python_types(metrics)}")
@@ -241,7 +244,7 @@ def main():
         # Write and serve the evaluation report.
         # Only the local rank 0 writes/serves it, to avoid races on a shared output dir.
         if get_local_rank() == 0:
-            report_path = build_report(video_cfg.video_base_dir)
+            report_path = build_report(output_dir)
             if args_cli.serve_evaluation_report:
                 serve_until_ctrl_c(report_path.parent, args_cli.evaluation_report_port, report_path.name)
 

isaaclab_arena/evaluation/policy_runner_cli.py

@@ -39,12 +39,12 @@ def add_policy_runner_arguments(parser: argparse.ArgumentParser) -> None:
         help="Record an mp4 video of the rollout viewport (uses gymnasium.wrappers.RecordVideo).",
     )
     parser.add_argument(
-        "--video_base_dir",
+        "--output_base_dir",
         type=str,
-        default="/eval/videos",
+        default="/eval/output",
         help=(
-            "Base directory for recorded videos; a reverse-dated run subdirectory is added per run."
-            " Used with --record_viewport_video and/or --record_camera_video."
+            "Base directory for evaluation outputs (videos, per-episode results, report); a"
+            " reverse-dated run subdirectory is added per run."
         ),
     )
     parser.add_argument(