Feature/nvblox next datagen slim

isaaclab_arena/environments/isaaclab_arena_manager_based_env_cfg.py

@@ -83,6 +83,11 @@ class IsaacLabArenaManagerBasedRLEnvCfg(ManagerBasedRLEnvCfg):
     decimation: int = 4
     episode_length_s: float = 50.0
     wait_for_textures: bool = False
+    # Force at least one RTX sensor refresh after reset. The IsaacLab 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
 
 
 @configclass

isaaclab_arena/evaluation/datagen_collector.py

@@ -0,0 +1,130 @@
+# 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
+"""Interfaces for datagen data collection driven by the evaluation runners.
+
+DatagenRunManagerBase is the single object injected into eval_runner.main. It drives one
+DatagenCollectorBase per job and keeps whatever run-level records it needs. The runners
+drive both objects only through these interfaces. NoOpDatagenRunManager is the default
+when none is injected, so the runners never special-case the no-collection path.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Literal
+
+# How an episode ended.
+EpisodeOutcome = Literal["success", "failure", "timeout"]
+
+
+class DatagenCollectorBase(ABC):
+    """Interface the evaluation runners use to drive a datagen data collector.
+
+    Implementations record per-step data during a policy rollout (see rollout_policy in
+    policy_runner.py). The runner calls on_step after every env.step, on_episode_end when the
+    env resets between episodes, finalize when the rollout finishes, and close when the job is
+    done. How a collector is constructed and configured is up to the implementing package. The
+    runners depend only on this interface.
+    """
+
+    @abstractmethod
+    def on_step(self, env: Any, obs: Any, actions: Any, step_idx: int) -> None:
+        """Record one frame after an env.step."""
+
+    @abstractmethod
+    def on_episode_end(self, env: Any, outcome: EpisodeOutcome = "timeout") -> None:
+        """Flush the in-progress episode with the outcome that ended it.
+
+        Also the place to prepare the next episode's cameras (e.g. re-randomize their
+        placement). The runner calls this once the env has auto-reset, so any re-aimed poses
+        take effect over the next episode's leading frames, which the collector is expected
+        to drop.
+        """
+
+    @abstractmethod
+    def finalize(self, env: Any | None = None) -> None:
+        """Flush any in-progress episode and stop recording. Idempotent."""
+
+    @abstractmethod
+    def close(self, env: Any | None = None) -> None:
+        """Finalize, then release resources such as spawned cameras. Idempotent."""
+
+
+class DatagenRunManagerBase(ABC):
+    """Interface for the run-level datagen object injected into eval_runner.main.
+
+    on_job_finished runs while the job's env is still alive, so an implementation can release
+    per-job resources such as spawned cameras before the stage teardown and record whatever
+    run-level bookkeeping it needs.
+    """
+
+    @abstractmethod
+    def start_run(self, eval_jobs_config: dict, description: str | None, device: str) -> None:
+        """Called once before any job runs."""
+
+    @abstractmethod
+    def create_collector(self, job_name: str, datagen_job_dict: dict, env: Any) -> DatagenCollectorBase:
+        """Build, retain, and return the collector for one job's live env."""
+
+    @abstractmethod
+    def on_job_finished(self, job: Any, env: Any) -> None:
+        """Called once per collecting job after its rollout, while the env is still alive."""
+
+    @abstractmethod
+    def finish_run(self) -> None:
+        """Called once after all jobs ran."""
+
+    def needs_cameras(self) -> bool:
+        """Whether collection requires the sim to start with camera support enabled."""
+        return True
+
+    def prepare_env_cfg(self, env_cfg: Any) -> None:
+        """Adjust an env cfg before its env is built. Called once per job.
+
+        The default drops the env's metrics and their recorder terms so the collector's own
+        dataset is the only one written.
+        """
+        if hasattr(env_cfg, "metrics"):
+            env_cfg.metrics = None
+        if hasattr(env_cfg, "recorders"):
+            env_cfg.recorders = None
+
+
+class NoOpDatagenCollector(DatagenCollectorBase):
+    """Collector that records nothing, used when no datagen collection is configured."""
+
+    def on_step(self, env: Any, obs: Any, actions: Any, step_idx: int) -> None:
+        pass
+
+    def on_episode_end(self, env: Any, outcome: EpisodeOutcome = "timeout") -> None:
+        pass
+
+    def finalize(self, env: Any | None = None) -> None:
+        pass
+
+    def close(self, env: Any | None = None) -> None:
+        pass
+
+
+class NoOpDatagenRunManager(DatagenRunManagerBase):
+    """Run manager that collects nothing. eval_runner defaults to it when none is injected."""
+
+    def start_run(self, eval_jobs_config: dict, description: str | None, device: str) -> None:
+        pass
+
+    def create_collector(self, job_name: str, datagen_job_dict: dict, env: Any) -> DatagenCollectorBase:
+        return NoOpDatagenCollector()
+
+    def on_job_finished(self, job: Any, env: Any) -> None:
+        pass
+
+    def finish_run(self) -> None:
+        pass
+
+    def needs_cameras(self) -> bool:
+        return False
+
+    def prepare_env_cfg(self, env_cfg: Any) -> None:
+        pass

isaaclab_arena/evaluation/eval_runner.py

@@ -17,6 +17,7 @@
 
 from isaaclab_arena.assets.registries import PolicyRegistry
 from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser
+from isaaclab_arena.evaluation.datagen_collector import DatagenRunManagerBase, NoOpDatagenRunManager
 from isaaclab_arena.evaluation.eval_runner_cli import add_eval_runner_arguments
 from isaaclab_arena.evaluation.job_manager import Job, JobManager, Status
 from isaaclab_arena.evaluation.policy_runner import get_policy_cls, rollout_policy
@@ -42,6 +43,7 @@ def load_env(
     variations: list[str] | None = None,
     render_mode: str | None = None,
     language_instruction: str | None = None,
+    datagen_run_manager: DatagenRunManagerBase | None = None,
 ):
 
     args_parser = get_isaaclab_arena_environments_cli_parser()
@@ -57,6 +59,11 @@ def load_env(
     if env_cfg.recorders is not None:
         env_cfg.recorders.dataset_filename = f"dataset_{job_name}"
 
+    # Datagen: let the run manager adjust the cfg before the env is built (e.g. drop the
+    # env's metric recorders so the collector's own dataset is the only one written).
+    if datagen_run_manager is not None:
+        datagen_run_manager.prepare_env_cfg(env_cfg)
+
     env = arena_builder.make_registered(env_cfg, env_kwargs, 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
@@ -202,7 +209,17 @@ def _run_in_chunks(args_cli: argparse.Namespace, master_cfg: dict) -> None:
             sys.exit(returncode)
 
 
-def main():
+def main(datagen_run_manager: DatagenRunManagerBase | None = None):
+    """Run all jobs in an eval config.
+
+    Args:
+        datagen_run_manager: Optional datagen run manager. It drives one data collector
+            per job (implementations live in the collecting package). When None, the no-op
+            default runs and behavior matches the plain evaluation path.
+    """
+    if datagen_run_manager is None:
+        datagen_run_manager = NoOpDatagenRunManager()
+
     args_parser = get_isaaclab_arena_cli_parser()
     args_cli, unknown = args_parser.parse_known_args()
 
@@ -236,10 +253,22 @@ def main():
         _run_in_chunks(args_cli, eval_jobs_config)
         return
 
+    # Optional top-level datagen defaults, applied to every job (a per-job "datagen"
+    # block overrides these). When present, the injected run manager collects data
+    # alongside each rollout.
+    datagen_defaults = eval_jobs_config.get("datagen")
+
+    # start_run asserts the config actually configures collection (a datagen block with an
+    # output_dir). The no-op default manager does nothing.
+    datagen_run_manager.start_run(eval_jobs_config, args_cli.datagen_description, args_cli.device)
+
     # Check if any job requires cameras and enable them if needed before starting simulation
     if args_cli.record_camera_video:
         args_cli.enable_cameras = True
     enable_cameras_if_required(eval_jobs_config, args_cli)
+    # Datagen collection renders dedicated cameras, which requires camera support.
+    if datagen_run_manager.needs_cameras():
+        args_cli.enable_cameras = True
 
     with SimulationAppContext(args_cli):
         job_manager = JobManager(eval_jobs_config["jobs"])
@@ -262,12 +291,16 @@ def main():
                 continue
             env = None
             policy = None
+            collector = None
 
             metrics_per_run: list[MetricsDataCollection] = []
 
             # num_episodes is the total across rebuilds, so split it over the rebuilds.
             num_episodes_per_rebuild = _split_episodes_across_rebuilds(job.num_episodes, job.num_rebuilds, job.name)
 
+            # Top-level datagen defaults overridden by the job's own datagen block.
+            datagen_job_dict = {**(datagen_defaults or {}), **(job.datagen or {})}
+
             # Rebuild the environment and re-run the rollout job.num_rebuilds times, then
             # aggregate the metrics across rebuilds into a single result.
             for rebuild_idx in range(job.num_rebuilds):
@@ -287,6 +320,7 @@ def main():
                         variations=job.variations,
                         render_mode=video_cfg.render_mode,
                         language_instruction=job.language_instruction,
+                        datagen_run_manager=datagen_run_manager,
                     )
 
                     # Write per-episode results to disk.
@@ -311,11 +345,15 @@ def main():
 
                     env = wrap_env_for_video(env, video_cfg, job.num_steps, num_episodes_this_rebuild)
 
+                    collector = datagen_run_manager.create_collector(job.name, datagen_job_dict, env)
+
                     metrics = rollout_policy(
                         env,
                         policy,
                         num_steps=job.num_steps,
                         num_episodes=num_episodes_this_rebuild,
+                        language_instruction=job.language_instruction,
+                        collector=collector,
                     )
 
                     job_manager.complete_job(job, metrics=metrics, status=Status.COMPLETED)
@@ -333,17 +371,27 @@ def main():
 
                 finally:
                     try:
-                        _close_job_resources(policy, env)
+                        # Run the per-job hook while the env is still alive and before the
+                        # stage teardown, so an implementation can release its resources.
+                        datagen_run_manager.on_job_finished(job, env)
+                    except Exception as exc:  # never let datagen bookkeeping fail a run
+                        print(f"[datagen] Warning: failed to record job '{job.name}': {exc}")
                     finally:
-                        policy = None
-                        env = None
-                        collect_garbage_and_clear_cuda_cache()
-
-            # Aggregate the metrics from the different experiments into a single view.
+                        try:
+                            _close_job_resources(policy, env)
+                        finally:
+                            policy = None
+                            env = None
+                            collector = None
+                            collect_garbage_and_clear_cuda_cache()
+
+            # Aggregate the metrics from the different rebuilds into a single view.
             if metrics_per_run:
                 aggregated_metrics = aggregate_metrics(metrics_per_run)
                 metrics_logger.append_job_metrics(job.name, aggregated_metrics)
 
+        datagen_run_manager.finish_run()
+
         job_manager.print_jobs_info()
         metrics_logger.print_metrics()
 

isaaclab_arena/evaluation/eval_runner_cli.py

@@ -53,6 +53,15 @@ def add_eval_runner_arguments(parser: argparse.ArgumentParser) -> None:
         default=False,
         help="Continue evaluation with remaining jobs when a job fails instead of stopping immediately.",
     )
+    parser.add_argument(
+        "--datagen-description",
+        type=str,
+        default=None,
+        help=(
+            "Free-text reason this datagen dataset was generated; recorded in manifest.json "
+            "(overrides the eval config's datagen.description)."
+        ),
+    )
     parser.add_argument(
         "--chunk_size",
         type=int,

isaaclab_arena/evaluation/job_manager.py

@@ -30,6 +30,7 @@ def __init__(
         policy_config_dict: dict = None,
         status: Status = None,
         language_instruction: str = None,
+        datagen: dict = None,
         variations: list[str] = None,
     ):
         """Initialize a Job instance.
@@ -65,6 +66,9 @@ def __init__(
         self.policy_type = policy_type
         self.policy_config_dict = policy_config_dict if policy_config_dict is not None else {}
         self.language_instruction = language_instruction
+        # Optional per-job datagen collection settings (camera pose, output dir, ...).
+        # When None, datagen collection is disabled for this job.
+        self.datagen = datagen
         self.status = status if status is not None else Status.PENDING
         self.start_time = None
         self.end_time = None
@@ -122,6 +126,7 @@ def from_dict(cls, data: dict) -> "Job":
             policy_config_dict=data["policy_config_dict"],
             status=status,
             language_instruction=data.get("language_instruction"),
+            datagen=data.get("datagen"),
             variations=cls.convert_variations_dict_to_hydra_overrides(data.get("variations", {})),
         )