Dtingdahl/datagen collector port

isaaclab_arena/environments/isaaclab_arena_manager_based_env_cfg.py

@@ -86,6 +86,12 @@ class IsaacLabArenaManagerBasedRLEnvCfg(ManagerBasedRLEnvCfg):
     decimation: int = 8
     wait_for_textures: bool = False
 
+    # Force extra RTX sensor refreshes after every reset. IsaacLab's own default (0)
+    # leaves camera buffers stale on the first frame of every episode after the
+    # first, so the previous episode's final rendered frame leaks in (corrupting
+    # RGB/depth/flow during datagen). See IsaacLab-Arena #339.
+    num_rerenders_on_reset: int = 5
+
 
 def apply_arena_global_settings() -> None:
     """Apply Arena's process-global RTX and physics settings before environment construction."""

isaaclab_arena/evaluation/arena_run.py

@@ -65,6 +65,12 @@ class ArenaRunCfg:
     variations: dict[str, Any] = field(default_factory=dict)
     """Variation values applied when the environment is compiled."""
 
+    datagen: dict[str, Any] | None = field(default=None)
+    """Per-run datagen collection config (output_dir, cameras, ...), or None to disable
+    collection for this run. Consumed by a datagen_collector_factory injected into
+    execute_experiment/experiment_runner.main -- Arena itself does not interpret its
+    contents beyond passing it to that factory."""
+
     def __post_init__(self) -> None:
         assert self.name, "run name must not be empty"
         assert self.num_rebuilds > 0, "num_rebuilds must be greater than zero"

isaaclab_arena/evaluation/datagen_collector.py

@@ -0,0 +1,75 @@
+# 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
+"""Interface for datagen data collectors, and their CallbackRecorderTerm wiring.
+
+DatagenCollectorBase is implemented by the collecting package (e.g. nvblox_next's
+datagen.arena_data_collector.DatagenCollector). build_datagen_callback_handlers adapts
+one to the generic CallbackRecorderTermHandlers shape recording.callback_recorder_term
+expects, keeping that module free of any datagen-specific knowledge.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from isaaclab_arena.evaluation.episode_outcome import EpisodeOutcome, classify_outcome
+from isaaclab_arena.recording.callback_recorder_term import CallbackRecorderTermHandlers
+
+
+class DatagenCollectorBase(ABC):
+    """Interface a datagen data collector implements, driven via CallbackRecorderTerm.
+
+    Implementations record per-step data during a policy rollout. on_step fires after
+    every env.step (for every env, before any reset that step); on_episode_end fires
+    once per env_id right before that env's reset, while its terminal state is still
+    intact -- also the place to prepare that env's cameras for the next episode, since
+    the reset immediately following flushes any re-aimed poses via IsaacLab's
+    num_rerenders_on_reset. finalize/close run at rollout/job teardown.
+    """
+
+    @abstractmethod
+    def on_step(self, env: Any) -> None:
+        """Record one frame for every env.
+
+        Reads env.obs_buf / env.action_manager / env.episode_length_buf directly (no
+        args beyond env: state lives on it).
+        """
+
+    @abstractmethod
+    def on_episode_end(self, env: Any, env_id: int, outcome: EpisodeOutcome = "timeout") -> None:
+        """Flush env_id's in-progress episode and prepare its cameras for the next one."""
+
+    @abstractmethod
+    def finalize(self, env: Any | None = None) -> None:
+        """Flush any in-progress episodes and stop recording. Idempotent."""
+
+    @abstractmethod
+    def close(self, env: Any | None = None) -> None:
+        """Finalize, then release resources such as spawned cameras. Idempotent."""
+
+
+def build_datagen_callback_handlers(
+    collector: DatagenCollectorBase, env: Any | None = None
+) -> CallbackRecorderTermHandlers:
+    """Adapt a DatagenCollectorBase to the CallbackRecorderTerm handler shape.
+
+    Args:
+        collector: The collector to drive.
+        env: If given, on_close calls collector.close(env) with this env instead of
+            None (CallbackRecorderTerm's on_close only ever receives a file_path, not
+            an env, so callers that need close(env) must bind env here at
+            build_handlers time -- see run_execution.py's usage).
+    """
+
+    def on_pre_reset(pre_reset_env: Any, env_ids) -> None:
+        for env_id in env_ids:
+            collector.on_episode_end(pre_reset_env, int(env_id), outcome=classify_outcome(pre_reset_env, int(env_id)))
+
+    return CallbackRecorderTermHandlers(
+        on_post_step=collector.on_step,
+        on_pre_reset=on_pre_reset,
+        on_close=lambda _file_path: collector.close(env),
+    )

isaaclab_arena/evaluation/episode_outcome.py

@@ -0,0 +1,33 @@
+# Copyright (c) 2025-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
+"""Classify how a datagen episode ended, from the env's termination state."""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+EpisodeOutcome = Literal["success", "failure", "timeout"]
+
+
+def classify_outcome(env: Any, env_id: int) -> EpisodeOutcome:
+    """Classify env_id's just-finished episode from its active termination terms.
+
+    Mirrors record_core_episode_results in isaaclab_arena/recording/common_terms.py,
+    which reads the same "success" termination term for its own per-episode record.
+
+    Args:
+        env: IsaacLab environment instance (must have a termination_manager).
+        env_id: Index of the env whose episode just ended.
+
+    Returns:
+        "success" if the success termination term fired, "timeout" if the time_out
+        term fired, otherwise "failure".
+    """
+    active_terms = env.termination_manager.active_terms
+    if "success" in active_terms and bool(env.termination_manager.get_term("success")[env_id]):
+        return "success"
+    if "time_out" in active_terms and bool(env.termination_manager.get_term("time_out")[env_id]):
+        return "timeout"
+    return "failure"

isaaclab_arena/evaluation/experiment_runner.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from pathlib import Path
 from typing import TYPE_CHECKING
 
@@ -23,8 +24,11 @@
 from isaaclab_arena.video.video_recording import timestamped_run_dir
 
 if TYPE_CHECKING:
+    import gymnasium as gym
+
     from isaaclab_arena.evaluation.arena_experiment import ArenaExperimentCfg
-    from isaaclab_arena.evaluation.arena_run import ArenaRunResult
+    from isaaclab_arena.evaluation.arena_run import ArenaRunCfg, ArenaRunResult
+    from isaaclab_arena.evaluation.datagen_collector import DatagenCollectorBase
 
 
 # TODO(cvolk): Move experiment-level variation inspection out of this CLI entry point.
@@ -95,7 +99,15 @@ def _write_arena_experiment_result(
     return ArenaExperimentResult(experiment_output_directory, run_metadata_by_name).write()
 
 
-def main():
+def main(datagen_collector_factory: Callable[[ArenaRunCfg, gym.Env], DatagenCollectorBase] | None = None) -> None:
+    """Run an Arena Experiment (one or more typed or legacy-JSON Runs).
+
+    Args:
+        datagen_collector_factory: Optional Callable[[ArenaRunCfg, gym.Env], DatagenCollectorBase].
+            When given, each Run whose ArenaRunCfg.datagen is not None gets a collector
+            built from it, driven via a CallbackRecorderTerm. When None, no datagen
+            collection runs and behavior matches the plain evaluation path.
+    """
     args_cli, experiment_overrides = parse_experiment_runner_args()
     experiment_config_path = validate_experiment_config_path(args_cli.experiment_config)
     legacy_experiment_config = load_legacy_json_experiment_config(
@@ -173,6 +185,7 @@ def main():
             record_viewport_video=args_cli.record_viewport_video,
             record_camera_video=args_cli.record_camera_video,
             continue_on_error=args_cli.continue_on_error,
+            datagen_collector_factory=datagen_collector_factory,
         )
         for run_result in run_results:
             if run_result.metrics is not None:

isaaclab_arena/evaluation/legacy_eval_config.py

@@ -84,6 +84,7 @@ def _run_cfg_from_legacy_job(
         ),
         num_rebuilds=job_config.get("num_rebuilds", 1),
         variations=variations,
+        datagen=job_config.get("datagen"),
     )
 
 

isaaclab_arena/evaluation/run_execution.py

@@ -9,21 +9,27 @@
 
 import os
 import traceback
+from collections.abc import Callable
 from copy import deepcopy
 from dataclasses import fields, replace
 from pathlib import Path
 from typing import TYPE_CHECKING
 
+from isaaclab.managers.recorder_manager import RecorderManagerBaseCfg
+
 from isaaclab_arena.assets.registries import EnvironmentRegistry, PolicyRegistry
 from isaaclab_arena.evaluation.arena_experiment import ArenaExperimentCfg
 from isaaclab_arena.evaluation.arena_run import ArenaRunCfg, ArenaRunResult, RunStatus
+from isaaclab_arena.evaluation.datagen_collector import DatagenCollectorBase, build_datagen_callback_handlers
 from isaaclab_arena.evaluation.legacy_graph_environment_cli import (
     LegacyGraphEnvironmentCfg,
     build_arena_builder_from_legacy_graph,
 )
 from isaaclab_arena.evaluation.policy_runner import rollout_policy
 from isaaclab_arena.evaluation.resource_cleanup import close_run_resources
 from isaaclab_arena.metrics.aggregate_metrics import aggregate_metrics
+from isaaclab_arena.recording.callback_recorder_term import CallbackRecorderTermCfg, CallbackRecorderTermHandlers
+from isaaclab_arena.utils.configclass import combine_configclass_instances, make_configclass
 from isaaclab_arena.variations.variations_hydra import overrides_from_dict
 from isaaclab_arena.video.video_recording import VideoRecordingCfg, wrap_env_for_video
 
@@ -41,6 +47,7 @@ def execute_experiment(
     record_viewport_video: bool = False,
     record_camera_video: bool = False,
     continue_on_error: bool = False,
+    datagen_collector_factory: Callable[[ArenaRunCfg, gym.Env], DatagenCollectorBase] | None = None,
 ) -> list[ArenaRunResult]:
     """Execute an experiment's runs in order and return their results.
 
@@ -50,6 +57,8 @@ def execute_experiment(
         record_viewport_video: Whether to record the viewport for each run.
         record_camera_video: Whether to record observation cameras for each run.
         continue_on_error: Whether to continue with later runs after one fails.
+        datagen_collector_factory: Optional Callable[[ArenaRunCfg, gym.Env], DatagenCollectorBase]
+            forwarded to build_and_run for each run.
 
     Returns:
         One result per attempted run, in execution order.
@@ -67,6 +76,7 @@ def execute_experiment(
                     record_camera_video=record_camera_video,
                     video_base_dir=str(run_output_dir),
                 ),
+                datagen_collector_factory=datagen_collector_factory,
             )
         except Exception as error:
             results.append(ArenaRunResult(run_name=run_cfg.name, status=RunStatus.FAILED))
@@ -84,8 +94,14 @@ def build_and_run(
     cfg: ArenaRunCfg,
     output_dir: str | Path,
     video_cfg: VideoRecordingCfg | None = None,
+    datagen_collector_factory: Callable[[ArenaRunCfg, gym.Env], DatagenCollectorBase] | None = None,
 ) -> ArenaRunResult:
-    """Build and execute one typed Arena run, then return its result."""
+    """Build and execute one typed Arena run, then return its result.
+
+    Args:
+        datagen_collector_factory: Optional Callable[[ArenaRunCfg, gym.Env], DatagenCollectorBase]
+            forwarded to _build_environment_from_cfg for each rebuild.
+    """
     metrics_per_rebuild: list[MetricsDataCollection] = []
     output_dir = str(output_dir)
     video_cfg = video_cfg or VideoRecordingCfg(video_base_dir=output_dir)
@@ -105,7 +121,9 @@ 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)
+            env = _build_environment_from_cfg(
+                rebuild_cfg, rebuild_video_cfg.render_mode, datagen_collector_factory=datagen_collector_factory
+            )
             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)
@@ -137,15 +155,55 @@ def _seed_cfg_for_rebuild(cfg: ArenaRunCfg, rebuild_index: int) -> ArenaRunCfg:
     return cfg
 
 
+def _with_datagen_recorder_term(
+    recorders_cfg: RecorderManagerBaseCfg | None,
+    build_handlers: Callable[[gym.Env], CallbackRecorderTermHandlers],
+) -> RecorderManagerBaseCfg:
+    """Merge a CallbackRecorderTerm using build_handlers into recorders_cfg.
+
+    recorders_cfg may be None (no other recorder terms configured for this run).
+    """
+    datagen_recorders_cfg = make_configclass(
+        "DatagenRecorderManagerCfg",
+        [("datagen_callback", CallbackRecorderTermCfg, CallbackRecorderTermCfg(build_handlers=build_handlers))],
+        bases=(RecorderManagerBaseCfg,),
+    )()
+    # datagen_recorders_cfg is passed last so recorders_cfg's already-configured values (not
+    # datagen_recorders_cfg's inherited base-class defaults) win on any field both share.
+    return combine_configclass_instances(
+        "RecorderManagerCfg", datagen_recorders_cfg, recorders_cfg, bases=(RecorderManagerBaseCfg,)
+    )
+
+
 def _build_environment_from_cfg(
     cfg: ArenaRunCfg,
     render_mode: str | None,
+    datagen_collector_factory: Callable[[ArenaRunCfg, gym.Env], DatagenCollectorBase] | None = None,
 ) -> gym.Env:
-    """Compile and instantiate a run's environment."""
+    """Compile and instantiate a run's environment.
+
+    Args:
+        datagen_collector_factory: Optional Callable[[ArenaRunCfg, gym.Env], DatagenCollectorBase].
+            When given and cfg.datagen is not None, its collector is driven via a
+            CallbackRecorderTerm merged into the env's recorders config.
+    """
     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}"
+    if datagen_collector_factory is not None and cfg.datagen is not None:
+        # ArenaEnvBuilder only sets env_cfg.recorders when mimic is disabled, so a
+        # requested datagen collector would silently never be invoked in mimic mode.
+        assert not cfg.environment_builder.mimic, (
+            f"Run '{cfg.name}' requests datagen collection but mimic mode never sets env_cfg.recorders,"
+            " so the collector would silently never be invoked"
+        )
+
+        def build_handlers(env: gym.Env, run_cfg: ArenaRunCfg = cfg) -> CallbackRecorderTermHandlers:
+            collector: DatagenCollectorBase = datagen_collector_factory(run_cfg, env)
+            return build_datagen_callback_handlers(collector, env=env)
+
+        env_cfg.recorders = _with_datagen_recorder_term(env_cfg.recorders, build_handlers)
     return arena_builder.make_registered(env_cfg, env_kwargs, render_mode=render_mode)