From 5afe2b1caef94705fc764963c76cbbb296296b7e Mon Sep 17 00:00:00 2001
From: David Tingdahl <dtingdahl@nvidia.com>
Date: Fri, 10 Jul 2026 09:43:39 +0200
Subject: [PATCH 1/4] Add datagen collection hooks to the evaluation runner

Inject nvblox_next datagen collection into eval_runner and policy_runner via an
optional collector_provider argument. The provider builds a per-job collector from
a job's datagen config block; rollouts then record per-step data, per-episode
outcomes, and a run-level provenance manifest. Adds episode outcome
classification, a --datagen-description CLI flag, and a max_recorded_frames cap.

Signed-off-by: David Tingdahl <dtingdahl@nvidia.com>
---
 .../isaaclab_arena_manager_based_env_cfg.py   |   5 +
 isaaclab_arena/evaluation/episode_outcome.py  |  25 ++
 isaaclab_arena/evaluation/eval_runner.py      |  77 ++++++-
 isaaclab_arena/evaluation/eval_runner_cli.py  |   9 +
 isaaclab_arena/evaluation/job_manager.py      |   5 +
 isaaclab_arena/evaluation/policy_runner.py    | 214 +++++++++++++++---
 isaaclab_arena/tests/test_episode_outcome.py  |  19 ++
 isaaclab_arena/tests/test_eval_runner.py      |   2 +-
 8 files changed, 309 insertions(+), 47 deletions(-)
 create mode 100644 isaaclab_arena/evaluation/episode_outcome.py
 create mode 100644 isaaclab_arena/tests/test_episode_outcome.py

diff --git a/isaaclab_arena/environments/isaaclab_arena_manager_based_env_cfg.py b/isaaclab_arena/environments/isaaclab_arena_manager_based_env_cfg.py
index ba9bfb620c..3927f90631 100644
--- a/isaaclab_arena/environments/isaaclab_arena_manager_based_env_cfg.py
+++ b/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
diff --git a/isaaclab_arena/evaluation/episode_outcome.py b/isaaclab_arena/evaluation/episode_outcome.py
new file mode 100644
index 0000000000..f9bd0e21a8
--- /dev/null
+++ b/isaaclab_arena/evaluation/episode_outcome.py
@@ -0,0 +1,25 @@
+# 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
+"""Map the termination term that ended a datagen episode to an outcome label."""
+
+from __future__ import annotations
+
+
+def classify_outcome(ended_by: str | None) -> str:
+    """Classify an episode outcome from the termination term that ended it.
+
+    Args:
+        ended_by: Name of the stashed termination term that fired, or ``None``
+            when the episode ended on the ``max_episode_length`` cap.
+
+    Returns:
+        ``"success"`` if the success term fired, ``"timeout"`` if nothing fired
+        (length cap), otherwise ``"failure"``.
+    """
+    if ended_by is None:
+        return "timeout"
+    if ended_by == "success":
+        return "success"
+    return "failure"
diff --git a/isaaclab_arena/evaluation/eval_runner.py b/isaaclab_arena/evaluation/eval_runner.py
index d2dfc0979e..d5ad593ca9 100644
--- a/isaaclab_arena/evaluation/eval_runner.py
+++ b/isaaclab_arena/evaluation/eval_runner.py
@@ -19,7 +19,7 @@
 from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser
 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
+from isaaclab_arena.evaluation.policy_runner import get_policy_cls, prepare_env_cfg_for_datagen, rollout_policy
 from isaaclab_arena.metrics.aggregate_metrics import aggregate_metrics
 from isaaclab_arena.metrics.metrics_logger import MetricsLogger
 from isaaclab_arena.utils.isaaclab_utils.simulation_app import (
@@ -42,6 +42,7 @@ def load_env(
     variations: list[str] | None = None,
     render_mode: str | None = None,
     language_instruction: str | None = None,
+    disable_auto_reset: bool = False,
 ):
 
     args_parser = get_isaaclab_arena_environments_cli_parser()
@@ -57,9 +58,13 @@ def load_env(
     if env_cfg.recorders is not None:
         env_cfg.recorders.dataset_filename = f"dataset_{job_name}"
 
+    # Datagen: disable the env's in-step() auto-reset (and metric recorders) so the rollout
+    # loop can drive episode boundaries and resets explicitly (see prepare_env_cfg_for_datagen).
+    reset_terms = prepare_env_cfg_for_datagen(env_cfg) if disable_auto_reset else None
+
     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
+    return env, reset_terms
 
 
 def list_variations(eval_jobs_config: dict) -> None:
@@ -202,7 +207,14 @@ def _run_in_chunks(args_cli: argparse.Namespace, master_cfg: dict) -> None:
             sys.exit(returncode)
 
 
-def main():
+def main(collector_provider=None):
+    """Run all jobs in an eval config.
+
+    Args:
+        collector_provider: Optional datagen provider (see Task 8's
+            ``nvblox_next`` implementation). When ``None``, no datagen
+            collection runs and behavior matches the plain evaluation path.
+    """
     args_parser = get_isaaclab_arena_cli_parser()
     args_cli, unknown = args_parser.parse_known_args()
 
@@ -236,10 +248,21 @@ 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, datagen data collection runs alongside
+    # each rollout, writing one HDF5 file per episode.
+    datagen_defaults = eval_jobs_config.get("datagen")
+
+    if collector_provider is not None and datagen_defaults is not None:
+        collector_provider.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_defaults or any(job_dict.get("datagen") for job_dict in eval_jobs_config["jobs"]):
+        args_cli.enable_cameras = True
 
     with SimulationAppContext(args_cli):
         job_manager = JobManager(eval_jobs_config["jobs"])
@@ -262,12 +285,18 @@ 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)
 
+            # Datagen is active for this job when the provider is enabled for it (top-level
+            # defaults overridden by the job's own datagen block).
+            job_datagen = {**(datagen_defaults or {}), **(job.datagen or {})}
+            is_datagen = collector_provider is not None and collector_provider.is_enabled(job_datagen)
+
             # 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):
@@ -281,12 +310,15 @@ def main():
                         video_base_dir=job_output_dir,
                         camera_name_prefix=f"robot-cam-rebuild{rebuild_idx}",
                     )
-                    env = load_env(
+                    # Datagen: disable the env's auto-reset so the rollout drives episode
+                    # boundaries and resets explicitly (see prepare_env_cfg_for_datagen).
+                    env, datagen_reset_terms = load_env(
                         job.arena_env_args,
                         job.name,
                         variations=job.variations,
                         render_mode=video_cfg.render_mode,
                         language_instruction=job.language_instruction,
+                        disable_auto_reset=is_datagen,
                     )
 
                     # Write per-episode results to disk.
@@ -311,11 +343,22 @@ def main():
 
                     env = wrap_env_for_video(env, video_cfg, job.num_steps, num_episodes_this_rebuild)
 
+                    collector = collector_provider.create(job.name, job_datagen, env) if is_datagen else None
+
+                    # Per-episode step cap for datagen, bounded by the env's own cap.
+                    max_episode_length = int(env.unwrapped.max_episode_length)
+                    if is_datagen:
+                        max_episode_length = collector_provider.max_episode_length_cap(job_datagen, max_episode_length)
+
                     metrics = rollout_policy(
                         env,
                         policy,
                         num_steps=job.num_steps,
                         num_episodes=num_episodes_this_rebuild,
+                        language_instruction=job.language_instruction,
+                        collector=collector,
+                        datagen_reset_terms=datagen_reset_terms,
+                        max_episode_length=max_episode_length,
                     )
 
                     job_manager.complete_job(job, metrics=metrics, status=Status.COMPLETED)
@@ -333,17 +376,31 @@ def main():
 
                 finally:
                     try:
-                        _close_job_resources(policy, env)
+                        # Release datagen cameras BEFORE tearing down the stage, so
+                        # their replicator annotators do not leak into the next job.
+                        if collector is not None:
+                            collector.close(env)
+                            try:
+                                collector_provider.on_job_finished(job, collector, env, job.status.value)
+                            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)
 
+        if collector_provider is not None and datagen_defaults is not None:
+            collector_provider.finish_run()
+
         job_manager.print_jobs_info()
         metrics_logger.print_metrics()
 
diff --git a/isaaclab_arena/evaluation/eval_runner_cli.py b/isaaclab_arena/evaluation/eval_runner_cli.py
index 235b4302f6..3c85d73f26 100644
--- a/isaaclab_arena/evaluation/eval_runner_cli.py
+++ b/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,
diff --git a/isaaclab_arena/evaluation/job_manager.py b/isaaclab_arena/evaluation/job_manager.py
index 25e8a9ed43..66a2565f90 100644
--- a/isaaclab_arena/evaluation/job_manager.py
+++ b/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", {})),
         )
 
diff --git a/isaaclab_arena/evaluation/policy_runner.py b/isaaclab_arena/evaluation/policy_runner.py
index 1f898e94d8..83c938852c 100644
--- a/isaaclab_arena/evaluation/policy_runner.py
+++ b/isaaclab_arena/evaluation/policy_runner.py
@@ -7,10 +7,11 @@
 
 import argparse
 import os
+import dataclasses
 import torch
 import tqdm
 from importlib import import_module
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from isaaclab_arena.assets.registries import PolicyRegistry
 from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser
@@ -19,6 +20,8 @@
     add_policy_runner_arguments,
     build_policy_from_cli,
 )
+from isaaclab_arena.evaluation.episode_outcome import classify_outcome
+from isaaclab_arena.evaluation.policy_runner_cli import add_policy_runner_arguments
 from isaaclab_arena.metrics.metrics_logger import metrics_to_plain_python_types
 from isaaclab_arena.utils.hydra_overrides import assert_hydra_overrides
 from isaaclab_arena.utils.isaaclab_utils.simulation_app import SimulationAppContext
@@ -63,14 +66,136 @@ def is_distributed(args_cli: argparse.Namespace) -> bool:
     )
 
 
+def prepare_env_cfg_for_datagen(env_cfg) -> list[tuple[str, Any]]:
+    """Prepare *env_cfg* for datagen collection. Call on the cfg *before* the env is built.
+
+    Two changes, both so the dedicated datagen cameras capture clean frames:
+
+    1. Remove the termination terms so the env never auto-resets inside ``step()``. Isaac
+       Lab resets a done env *within* ``step()`` and re-renders before the new scene is
+       flushed, so a dedicated camera reads back the previous episode's final frame as the
+       first frame of the next episode. The rollout loop instead evaluates the returned
+       terms manually and drives a clean, explicit ``env.reset()`` between episodes (which
+       flushes to the renderer before re-rendering). Mirrors
+       ``submodules/IsaacLab/scripts/tools/record_demos.py``.
+    2. Drop the metrics and their recorder terms. The datagen collector writes its own
+       per-episode HDF5, and the success-rate recorder asserts the (now removed) success
+       termination is active, so it must not run.
+
+    Returns:
+        The stashed non-timeout terms as ``(name, term)`` pairs (e.g. success, object_dropped)
+        for manual evaluation. Timeout terms are replaced by the env's ``max_episode_length``
+        cap in the loop.
+    """
+    # Datagen has its own writer; drop the env's metrics + recorder terms (the success
+    # recorder also depends on the success termination removed below). recorders=None makes
+    # the RecorderManager a no-op.
+    if hasattr(env_cfg, "metrics"):
+        env_cfg.metrics = None
+    if hasattr(env_cfg, "recorders"):
+        env_cfg.recorders = None
+
+    terminations = getattr(env_cfg, "terminations", None)
+    if terminations is None:
+        return []
+    stashed = []
+    for field in dataclasses.fields(terminations):
+        term = getattr(terminations, field.name)
+        # Skip unset/empty fields; real termination terms are TerminationTermCfg (have .func).
+        if term is None or not hasattr(term, "func"):
+            continue
+        setattr(terminations, field.name, None)
+        is_timeout = field.name == "time_out" or getattr(term, "time_out", False)
+        if not is_timeout:
+            stashed.append((field.name, term))
+    return stashed
+
+
+def _manual_episode_done(env, reset_terms: list) -> str | None:
+    """Return the name of the first stashed termination term that fires, else ``None``."""
+    base_env = env.unwrapped
+    for name, term in reset_terms:
+        result = term.func(base_env, **(term.params or {}))
+        if bool(torch.as_tensor(result).any()):
+            return name
+    return None
+
+
+def _run_datagen_rollout(
+    env, policy, collector, pbar, num_steps, num_episodes, reset_terms, max_episode_length, obs
+) -> None:
+    """Rollout loop for datagen collection with the env's auto-reset disabled.
+
+    Records every (settled) frame, decides episode end from the stashed termination terms
+    plus a max-length cap, then flushes the episode and performs a clean explicit
+    ``env.reset()`` before the next one so the first frame of each episode is correct.
+    """
+    assert max_episode_length is not None, "datagen rollout requires max_episode_length"
+    num_episodes_completed = 0
+    num_steps_completed = 0
+    steps_in_episode = 0
+    while True:
+        with torch.inference_mode():
+            actions = policy.get_action(env, obs)
+            obs, _, _, _, _ = env.step(actions)
+        steps_in_episode += 1
+        collector.on_step(env, obs, actions, num_steps_completed)
+        num_steps_completed += 1
+
+        if num_steps is not None:
+            pbar.update(1)
+            if num_steps_completed >= num_steps:
+                break
+
+        ended_by = _manual_episode_done(env, reset_terms)
+        hit_cap = steps_in_episode >= max_episode_length
+        if ended_by is not None or hit_cap:
+            collector.end_episode(env, outcome=classify_outcome(ended_by))
+            num_episodes_completed += 1
+            if num_episodes is not None:
+                pbar.update(1)
+                if num_episodes_completed >= num_episodes:
+                    break
+            # Re-aim the datagen cameras before the reset so the reset's RTX rerenders
+            # (num_rerenders_on_reset) flush the new poses; otherwise the next episode's
+            # first frame is rendered from the previous layout.
+            collector.resample_cameras()
+            obs, _ = env.reset()
+            policy.reset()
+            steps_in_episode = 0
+
+
 def rollout_policy(
     env,
     policy: PolicyBase,
     num_steps: int | None,
     num_episodes: int | None,
+    language_instruction: str | None = None,
+    collector: Any = None,
+    datagen_reset_terms: list | None = None,
+    max_episode_length: int | None = None,
 ) -> MetricsDataCollection | None:
+    """Roll out *policy* in *env*.
+
+    Args:
+        collector: Optional data collector. When provided, ``collector.on_step(env,
+            obs, actions, step_idx)`` is called after every environment step and
+            ``collector.finalize(env)`` once the rollout finishes. Duck-typed so
+            core does not depend on any specific datagen collection package.
+        datagen_reset_terms: Stashed termination terms returned by
+            :func:`prepare_env_cfg_for_datagen`. Required whenever a collector is
+            given: with the env's auto-reset disabled, the loop evaluates these terms
+            (plus ``max_episode_length``) to end episodes and resets explicitly, so no
+            frame is captured mid-reset.
+        max_episode_length: Per-episode step cap (replaces the removed timeout term).
+            Required when ``datagen_reset_terms`` is given.
+    """
     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"
+    assert collector is None or datagen_reset_terms is not None, (
+        "A datagen collector requires the env's auto-reset to be disabled; pass datagen_reset_terms"
+        " from prepare_env_cfg_for_datagen() (and max_episode_length)."
+    )
 
     pbar = None
     try:
@@ -84,41 +209,48 @@ def rollout_policy(
         else:
             pbar = tqdm.tqdm(total=num_episodes, desc="Episodes", unit="episode")
 
-        num_episodes_completed = 0
-        num_steps_completed = 0
-
-        while True:
-            with torch.inference_mode():
-                actions = policy.get_action(env, obs)
-                obs, _, terminated, truncated, _ = env.step(actions)
-
-                if terminated.any() or truncated.any():
-                    # Only reset policy for those envs that are terminated or truncated
-                    print(
-                        f"Resetting policy for terminated env_ids: {terminated.nonzero().flatten()}"
-                        f" and truncated env_ids: {truncated.nonzero().flatten()}"
-                    )
-                    env_ids = (terminated | truncated).nonzero().flatten()
-                    policy.reset(env_ids=env_ids)
-                    # Break if number of episodes is reached
-                    completed_episodes = env_ids.shape[0]
-                    num_episodes_completed += completed_episodes
-                    if hasattr(env.unwrapped.cfg, "metrics") and env.unwrapped.cfg.metrics is not None:
-                        metrics = env.unwrapped.compute_metrics()
-                        tqdm.tqdm.write(
-                            f"[Rank {get_local_rank()}/{get_world_size()}] Metrics:"
-                            f" {metrics_to_plain_python_types(metrics)}"
+        if collector is not None:
+            # Datagen path: env auto-reset is disabled, so we drive episode boundaries
+            # and resets explicitly (see _run_datagen_rollout / record_demos.py).
+            _run_datagen_rollout(
+                env, policy, collector, pbar, num_steps, num_episodes, datagen_reset_terms, max_episode_length, obs
+            )
+        else:
+            num_episodes_completed = 0
+            num_steps_completed = 0
+
+            while True:
+                with torch.inference_mode():
+                    actions = policy.get_action(env, obs)
+                    obs, _, terminated, truncated, _ = env.step(actions)
+
+                    if terminated.any() or truncated.any():
+                        # Only reset policy for those envs that are terminated or truncated
+                        print(
+                            f"Resetting policy for terminated env_ids: {terminated.nonzero().flatten()}"
+                            f" and truncated env_ids: {truncated.nonzero().flatten()}"
                         )
-                    if num_episodes is not None:
-                        pbar.update(completed_episodes)
-                        if num_episodes_completed >= num_episodes:
+                        env_ids = (terminated | truncated).nonzero().flatten()
+                        policy.reset(env_ids=env_ids)
+                        # Break if number of episodes is reached
+                        completed_episodes = env_ids.shape[0]
+                        num_episodes_completed += completed_episodes
+                        if hasattr(env.unwrapped.cfg, "metrics") and env.unwrapped.cfg.metrics is not None:
+                            metrics = env.unwrapped.compute_metrics()
+                            tqdm.tqdm.write(
+                                f"[Rank {get_local_rank()}/{get_world_size()}] Metrics:"
+                                f" {metrics_to_plain_python_types(metrics)}"
+                            )
+                        if num_episodes is not None:
+                            pbar.update(completed_episodes)
+                            if num_episodes_completed >= num_episodes:
+                                break
+                    # Break if number of steps is reached
+                    num_steps_completed += 1
+                    if num_steps is not None:
+                        pbar.update(1)
+                        if num_steps_completed >= num_steps:
                             break
-                # Break if number of steps is reached
-                num_steps_completed += 1
-                if num_steps is not None:
-                    pbar.update(1)
-                    if num_steps_completed >= num_steps:
-                        break
 
         pbar.close()
 
@@ -129,6 +261,10 @@ def rollout_policy(
 
     else:
 
+        # Persist and close any datagen dataset before returning.
+        if collector is not None:
+            collector.finalize(env)
+
         # Only compute metrics if env has non-None metrics.
         # Use unwrapped to reach the base env through any gym wrappers (e.g. OrderEnforcing)
         if hasattr(env.unwrapped.cfg, "metrics") and env.unwrapped.cfg.metrics is not None:
@@ -185,7 +321,7 @@ def main():
         if args_cli.record_camera_video:
             args_cli.enable_cameras = True
 
-        # Build scene. Use rgb_array render mode when recording so RecordVideo can grab frames.
+        # Build scene. Use rgb_array render mode when recording so the recorders can grab frames.
         arena_builder = get_arena_builder_from_cli(args_cli, hydra_overrides=hydra_overrides)
 
         if args_cli.list_variations:
@@ -229,7 +365,13 @@ 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)
+        metrics = rollout_policy(
+            env,
+            policy,
+            num_steps,
+            num_episodes,
+            args_cli.language_instruction,
+        )
 
         if metrics is not None:
             print(f"[Rank {local_rank}/{world_size}] Metrics: {metrics_to_plain_python_types(metrics)}")
diff --git a/isaaclab_arena/tests/test_episode_outcome.py b/isaaclab_arena/tests/test_episode_outcome.py
new file mode 100644
index 0000000000..b6281a625e
--- /dev/null
+++ b/isaaclab_arena/tests/test_episode_outcome.py
@@ -0,0 +1,19 @@
+# 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
+"""Unit tests for episode outcome classification (no Isaac Sim required)."""
+
+from isaaclab_arena.evaluation.episode_outcome import classify_outcome
+
+
+def test_success_term_maps_to_success():
+    assert classify_outcome("success") == "success"
+
+
+def test_no_term_means_timeout():
+    assert classify_outcome(None) == "timeout"
+
+
+def test_other_term_means_failure():
+    assert classify_outcome("object_dropped") == "failure"
diff --git a/isaaclab_arena/tests/test_eval_runner.py b/isaaclab_arena/tests/test_eval_runner.py
index 2451a3d2df..2034979791 100644
--- a/isaaclab_arena/tests/test_eval_runner.py
+++ b/isaaclab_arena/tests/test_eval_runner.py
@@ -264,7 +264,7 @@ def _test_eval_config_variation_lands_in_events_cfg(simulation_app):
         "variations": {"droid_abs_joint_pos": {f"camera_extrinsics_{camera_name}": {"enabled": True}}},
     })
 
-    env = load_env(job.arena_env_args, job.name, variations=job.variations)
+    env, _ = load_env(job.arena_env_args, job.name, variations=job.variations)
     try:
         env_cfg = env.unwrapped.cfg
         assert hasattr(env_cfg.events, event_name), (

From b363360898b73f54096bf12482ced7722fb747de Mon Sep 17 00:00:00 2001
From: David Tingdahl <dtingdahl@nvidia.com>
Date: Mon, 13 Jul 2026 10:48:42 +0200
Subject: [PATCH 2/4] Add DatagenCollectorBase interface for datagen collectors

Define the collector interface the evaluation runners drive (on_step,
on_episode_end, finalize, close) as an ABC in
evaluation/datagen_collector.py, and type rollout_policy's collector
argument against it. Implementations live in the collecting package;
the runners no longer rely on duck typing.

on_episode_end doubles as the pre-reset hook for preparing the next
episode's cameras, since the runner calls it before the explicit
env.reset() whose RTX rerenders flush new camera poses.

Signed-off-by: David Tingdahl <dtingdahl@nvidia.com>
---
 .../evaluation/datagen_collector.py           | 44 +++++++++++++++++++
 isaaclab_arena/evaluation/policy_runner.py    | 15 +++----
 2 files changed, 51 insertions(+), 8 deletions(-)
 create mode 100644 isaaclab_arena/evaluation/datagen_collector.py

diff --git a/isaaclab_arena/evaluation/datagen_collector.py b/isaaclab_arena/evaluation/datagen_collector.py
new file mode 100644
index 0000000000..449d861e0b
--- /dev/null
+++ b/isaaclab_arena/evaluation/datagen_collector.py
@@ -0,0 +1,44 @@
+# 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
+"""Interface for datagen data collectors driven by the evaluation runners."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+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 at each
+    episode boundary before its explicit reset, 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: str = "timeout") -> None:
+        """Flush the in-progress episode. outcome is "success", "failure", or "timeout".
+
+        This is also the place to prepare the next episode's cameras (e.g. re-randomize
+        their placement): the runner calls on_episode_end before the explicit env.reset()
+        that starts the next episode, and it is that reset's RTX rerenders that flush new
+        camera poses into the renderer. Re-aimed after the reset, the next episode's
+        first frame would still render from the previous layout.
+        """
+
+    @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."""
diff --git a/isaaclab_arena/evaluation/policy_runner.py b/isaaclab_arena/evaluation/policy_runner.py
index 83c938852c..fc127f1a01 100644
--- a/isaaclab_arena/evaluation/policy_runner.py
+++ b/isaaclab_arena/evaluation/policy_runner.py
@@ -20,6 +20,7 @@
     add_policy_runner_arguments,
     build_policy_from_cli,
 )
+from isaaclab_arena.evaluation.datagen_collector import DatagenCollectorBase
 from isaaclab_arena.evaluation.episode_outcome import classify_outcome
 from isaaclab_arena.evaluation.policy_runner_cli import add_policy_runner_arguments
 from isaaclab_arena.metrics.metrics_logger import metrics_to_plain_python_types
@@ -150,16 +151,14 @@ def _run_datagen_rollout(
         ended_by = _manual_episode_done(env, reset_terms)
         hit_cap = steps_in_episode >= max_episode_length
         if ended_by is not None or hit_cap:
-            collector.end_episode(env, outcome=classify_outcome(ended_by))
+            collector.on_episode_end(env, outcome=classify_outcome(ended_by))
             num_episodes_completed += 1
             if num_episodes is not None:
                 pbar.update(1)
                 if num_episodes_completed >= num_episodes:
                     break
-            # Re-aim the datagen cameras before the reset so the reset's RTX rerenders
-            # (num_rerenders_on_reset) flush the new poses; otherwise the next episode's
-            # first frame is rendered from the previous layout.
-            collector.resample_cameras()
+            # The reset must come after on_episode_end: its RTX rerenders (num_rerenders_on_reset)
+            # flush any camera poses the collector re-aimed for the next episode.
             obs, _ = env.reset()
             policy.reset()
             steps_in_episode = 0
@@ -171,7 +170,7 @@ def rollout_policy(
     num_steps: int | None,
     num_episodes: int | None,
     language_instruction: str | None = None,
-    collector: Any = None,
+    collector: DatagenCollectorBase | None = None,
     datagen_reset_terms: list | None = None,
     max_episode_length: int | None = None,
 ) -> MetricsDataCollection | None:
@@ -180,8 +179,8 @@ def rollout_policy(
     Args:
         collector: Optional data collector. When provided, ``collector.on_step(env,
             obs, actions, step_idx)`` is called after every environment step and
-            ``collector.finalize(env)`` once the rollout finishes. Duck-typed so
-            core does not depend on any specific datagen collection package.
+            ``collector.finalize(env)`` once the rollout finishes. Implementations
+            live in the collecting package (see :class:`DatagenCollectorBase`).
         datagen_reset_terms: Stashed termination terms returned by
             :func:`prepare_env_cfg_for_datagen`. Required whenever a collector is
             given: with the env's auto-reset disabled, the loop evaluates these terms

From ebfa9820c9cf17111f8f712301fdf57f598a5cb2 Mon Sep 17 00:00:00 2001
From: David Tingdahl <dtingdahl@nvidia.com>
Date: Wed, 15 Jul 2026 16:25:37 +0200
Subject: [PATCH 3/4] Datagen run manager owns its collectors; interface and
 hook cleanup

eval_runner now takes a datagen_run_manager typed against the new
DatagenRunManagerBase ABC. Injecting the manager enables collection for
every job (start_run asserts the config configures it), so the per-job
is_enabled hook is gone. The manager owns the collector it creates per
job: on_job_finished(job, env) closes and records it, so the runner no
longer manages collector lifetime or passes the collector back. The
per-episode recording cap moves onto the collector interface as
cap_episode_length with a no-cap default, and episode outcomes are typed
as the EpisodeOutcome literal that classify_outcome produces.

Signed-off-by: David Tingdahl <dtingdahl@nvidia.com>
---
 .../evaluation/datagen_collector.py           | 49 +++++++++++++++++--
 isaaclab_arena/evaluation/episode_outcome.py  |  4 +-
 isaaclab_arena/evaluation/eval_runner.py      | 47 ++++++++++--------
 3 files changed, 75 insertions(+), 25 deletions(-)

diff --git a/isaaclab_arena/evaluation/datagen_collector.py b/isaaclab_arena/evaluation/datagen_collector.py
index 449d861e0b..49f4171e44 100644
--- a/isaaclab_arena/evaluation/datagen_collector.py
+++ b/isaaclab_arena/evaluation/datagen_collector.py
@@ -2,12 +2,20 @@
 # All rights reserved.
 #
 # SPDX-License-Identifier: Apache-2.0
-"""Interface for datagen data collectors driven by the evaluation runners."""
+"""Interfaces for datagen data collection driven by the evaluation runners.
+
+DatagenRunManagerBase is the single object injected into eval_runner.main. It spawns
+and owns one DatagenCollectorBase per job and keeps whatever run-level records it
+needs. The runners drive both objects only through these interfaces.
+"""
 
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Any
+from typing import Any, Literal
+
+# How an episode ended, as classified by episode_outcome.classify_outcome.
+EpisodeOutcome = Literal["success", "failure", "timeout"]
 
 
 class DatagenCollectorBase(ABC):
@@ -25,8 +33,8 @@ 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: str = "timeout") -> None:
-        """Flush the in-progress episode. outcome is "success", "failure", or "timeout".
+    def on_episode_end(self, env: Any, outcome: EpisodeOutcome = "timeout") -> None:
+        """Flush the in-progress episode with the outcome that ended it.
 
         This is also the place to prepare the next episode's cameras (e.g. re-randomize
         their placement): the runner calls on_episode_end before the explicit env.reset()
@@ -42,3 +50,36 @@ def finalize(self, env: Any | None = None) -> None:
     @abstractmethod
     def close(self, env: Any | None = None) -> None:
         """Finalize, then release resources such as spawned cameras. Idempotent."""
+
+    def cap_episode_length(self, env_max_episode_length: int) -> int:
+        """Return the per-episode step cap the rollout should use.
+
+        Collectors with a recording budget return a smaller cap. The default keeps
+        the env's own limit.
+        """
+        return env_max_episode_length
+
+
+class DatagenRunManagerBase(ABC):
+    """Interface for the run-level datagen object injected into eval_runner.main.
+
+    The manager owns the collectors it creates: on_job_finished must close the job's
+    collector (its cameras must be released before the stage teardown) and may record
+    any bookkeeping it needs for the run.
+    """
+
+    @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."""
diff --git a/isaaclab_arena/evaluation/episode_outcome.py b/isaaclab_arena/evaluation/episode_outcome.py
index f9bd0e21a8..fb9ae41853 100644
--- a/isaaclab_arena/evaluation/episode_outcome.py
+++ b/isaaclab_arena/evaluation/episode_outcome.py
@@ -6,8 +6,10 @@
 
 from __future__ import annotations
 
+from isaaclab_arena.evaluation.datagen_collector import EpisodeOutcome
 
-def classify_outcome(ended_by: str | None) -> str:
+
+def classify_outcome(ended_by: str | None) -> EpisodeOutcome:
     """Classify an episode outcome from the termination term that ended it.
 
     Args:
diff --git a/isaaclab_arena/evaluation/eval_runner.py b/isaaclab_arena/evaluation/eval_runner.py
index d5ad593ca9..a7ba76659b 100644
--- a/isaaclab_arena/evaluation/eval_runner.py
+++ b/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
 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, prepare_env_cfg_for_datagen, rollout_policy
@@ -207,13 +208,14 @@ def _run_in_chunks(args_cli: argparse.Namespace, master_cfg: dict) -> None:
             sys.exit(returncode)
 
 
-def main(collector_provider=None):
+def main(datagen_run_manager: DatagenRunManagerBase | None = None):
     """Run all jobs in an eval config.
 
     Args:
-        collector_provider: Optional datagen provider (see Task 8's
-            ``nvblox_next`` implementation). When ``None``, no datagen
-            collection runs and behavior matches the plain evaluation path.
+        datagen_run_manager: Optional datagen run manager. It spawns and owns one
+            data collector per job (implementations live in the collecting package).
+            When ``None``, no datagen collection runs and behavior matches the plain
+            evaluation path.
     """
     args_parser = get_isaaclab_arena_cli_parser()
     args_cli, unknown = args_parser.parse_known_args()
@@ -253,8 +255,11 @@ def main(collector_provider=None):
     # each rollout, writing one HDF5 file per episode.
     datagen_defaults = eval_jobs_config.get("datagen")
 
-    if collector_provider is not None and datagen_defaults is not None:
-        collector_provider.start_run(eval_jobs_config, args_cli.datagen_description, args_cli.device)
+    # Injecting a run manager enables collection for every job; start_run asserts that
+    # the config actually configures collection (datagen block with an output_dir).
+    datagen_is_enabled = datagen_run_manager is not None
+    if datagen_is_enabled:
+        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:
@@ -292,10 +297,8 @@ def main(collector_provider=None):
             # 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)
 
-            # Datagen is active for this job when the provider is enabled for it (top-level
-            # defaults overridden by the job's own datagen block).
-            job_datagen = {**(datagen_defaults or {}), **(job.datagen or {})}
-            is_datagen = collector_provider is not None and collector_provider.is_enabled(job_datagen)
+            # 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.
@@ -318,7 +321,7 @@ def main(collector_provider=None):
                         variations=job.variations,
                         render_mode=video_cfg.render_mode,
                         language_instruction=job.language_instruction,
-                        disable_auto_reset=is_datagen,
+                        disable_auto_reset=datagen_is_enabled,
                     )
 
                     # Write per-episode results to disk.
@@ -343,12 +346,16 @@ def main(collector_provider=None):
 
                     env = wrap_env_for_video(env, video_cfg, job.num_steps, num_episodes_this_rebuild)
 
-                    collector = collector_provider.create(job.name, job_datagen, env) if is_datagen else None
+                    collector = (
+                        datagen_run_manager.create_collector(job.name, datagen_job_dict, env)
+                        if datagen_is_enabled
+                        else None
+                    )
 
                     # Per-episode step cap for datagen, bounded by the env's own cap.
                     max_episode_length = int(env.unwrapped.max_episode_length)
-                    if is_datagen:
-                        max_episode_length = collector_provider.max_episode_length_cap(job_datagen, max_episode_length)
+                    if collector is not None:
+                        max_episode_length = collector.cap_episode_length(max_episode_length)
 
                     metrics = rollout_policy(
                         env,
@@ -376,12 +383,12 @@ def main(collector_provider=None):
 
                 finally:
                     try:
-                        # Release datagen cameras BEFORE tearing down the stage, so
-                        # their replicator annotators do not leak into the next job.
+                        # on_job_finished closes the job's collector, releasing its cameras
+                        # BEFORE the stage teardown so their replicator annotators do not
+                        # leak into the next job.
                         if collector is not None:
-                            collector.close(env)
                             try:
-                                collector_provider.on_job_finished(job, collector, env, job.status.value)
+                                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:
@@ -398,8 +405,8 @@ def main(collector_provider=None):
                 aggregated_metrics = aggregate_metrics(metrics_per_run)
                 metrics_logger.append_job_metrics(job.name, aggregated_metrics)
 
-        if collector_provider is not None and datagen_defaults is not None:
-            collector_provider.finish_run()
+        if datagen_is_enabled:
+            datagen_run_manager.finish_run()
 
         job_manager.print_jobs_info()
         metrics_logger.print_metrics()

From 282e47d974af667b937e48611884acb348ad3da8 Mon Sep 17 00:00:00 2001
From: David Tingdahl <dtingdahl@nvidia.com>
Date: Sun, 19 Jul 2026 13:04:59 +0200
Subject: [PATCH 4/4] Datagen: use the env's native auto-reset and a no-op
 default collector

- Collapse the datagen rollout onto the standard auto-reset loop; drop
  _run_datagen_rollout, _manual_episode_done, and the termination stashing. On a
  done step the finished episode is closed before the post-reset obs is recorded,
  so the leaked reset frame lands in the next episode where the collector's
  skip_initial_frames drops it.
- Add NoOpDatagenCollector / NoOpDatagenRunManager and default to them, so the
  runners no longer special-case the no-collection path.
- Move env-cfg prep (drop metric recorders) and the camera requirement onto the
  run manager (prepare_env_cfg, needs_cameras); remove the unused
  cap_episode_length hook.
- Delete episode_outcome (classify_outcome unused after the rewrite).

Signed-off-by: David Tingdahl <dtingdahl@nvidia.com>
---
 .../evaluation/datagen_collector.py           |  93 ++++++--
 isaaclab_arena/evaluation/episode_outcome.py  |  27 ---
 isaaclab_arena/evaluation/eval_runner.py      |  74 +++---
 isaaclab_arena/evaluation/policy_runner.py    | 221 +++++-------------
 isaaclab_arena/tests/test_episode_outcome.py  |  19 --
 5 files changed, 161 insertions(+), 273 deletions(-)
 delete mode 100644 isaaclab_arena/evaluation/episode_outcome.py
 delete mode 100644 isaaclab_arena/tests/test_episode_outcome.py

diff --git a/isaaclab_arena/evaluation/datagen_collector.py b/isaaclab_arena/evaluation/datagen_collector.py
index 49f4171e44..5aaff76f17 100644
--- a/isaaclab_arena/evaluation/datagen_collector.py
+++ b/isaaclab_arena/evaluation/datagen_collector.py
@@ -4,9 +4,10 @@
 # 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 spawns
-and owns one DatagenCollectorBase per job and keeps whatever run-level records it
-needs. The runners drive both objects only through these interfaces.
+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
@@ -14,7 +15,7 @@
 from abc import ABC, abstractmethod
 from typing import Any, Literal
 
-# How an episode ended, as classified by episode_outcome.classify_outcome.
+# How an episode ended.
 EpisodeOutcome = Literal["success", "failure", "timeout"]
 
 
@@ -22,10 +23,10 @@ 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 at each
-    episode boundary before its explicit reset, 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.
+    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
@@ -36,11 +37,10 @@ def on_step(self, env: Any, obs: Any, actions: Any, step_idx: int) -> None:
     def on_episode_end(self, env: Any, outcome: EpisodeOutcome = "timeout") -> None:
         """Flush the in-progress episode with the outcome that ended it.
 
-        This is also the place to prepare the next episode's cameras (e.g. re-randomize
-        their placement): the runner calls on_episode_end before the explicit env.reset()
-        that starts the next episode, and it is that reset's RTX rerenders that flush new
-        camera poses into the renderer. Re-aimed after the reset, the next episode's
-        first frame would still render from the previous layout.
+        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
@@ -51,21 +51,13 @@ def finalize(self, env: Any | None = None) -> None:
     def close(self, env: Any | None = None) -> None:
         """Finalize, then release resources such as spawned cameras. Idempotent."""
 
-    def cap_episode_length(self, env_max_episode_length: int) -> int:
-        """Return the per-episode step cap the rollout should use.
-
-        Collectors with a recording budget return a smaller cap. The default keeps
-        the env's own limit.
-        """
-        return env_max_episode_length
-
 
 class DatagenRunManagerBase(ABC):
     """Interface for the run-level datagen object injected into eval_runner.main.
 
-    The manager owns the collectors it creates: on_job_finished must close the job's
-    collector (its cameras must be released before the stage teardown) and may record
-    any bookkeeping it needs for the run.
+    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
@@ -83,3 +75,56 @@ def on_job_finished(self, job: Any, env: Any) -> None:
     @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
diff --git a/isaaclab_arena/evaluation/episode_outcome.py b/isaaclab_arena/evaluation/episode_outcome.py
deleted file mode 100644
index fb9ae41853..0000000000
--- a/isaaclab_arena/evaluation/episode_outcome.py
+++ /dev/null
@@ -1,27 +0,0 @@
-# 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
-"""Map the termination term that ended a datagen episode to an outcome label."""
-
-from __future__ import annotations
-
-from isaaclab_arena.evaluation.datagen_collector import EpisodeOutcome
-
-
-def classify_outcome(ended_by: str | None) -> EpisodeOutcome:
-    """Classify an episode outcome from the termination term that ended it.
-
-    Args:
-        ended_by: Name of the stashed termination term that fired, or ``None``
-            when the episode ended on the ``max_episode_length`` cap.
-
-    Returns:
-        ``"success"`` if the success term fired, ``"timeout"`` if nothing fired
-        (length cap), otherwise ``"failure"``.
-    """
-    if ended_by is None:
-        return "timeout"
-    if ended_by == "success":
-        return "success"
-    return "failure"
diff --git a/isaaclab_arena/evaluation/eval_runner.py b/isaaclab_arena/evaluation/eval_runner.py
index a7ba76659b..46190ec3ed 100644
--- a/isaaclab_arena/evaluation/eval_runner.py
+++ b/isaaclab_arena/evaluation/eval_runner.py
@@ -17,10 +17,10 @@
 
 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
+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, prepare_env_cfg_for_datagen, rollout_policy
+from isaaclab_arena.evaluation.policy_runner import get_policy_cls, rollout_policy
 from isaaclab_arena.metrics.aggregate_metrics import aggregate_metrics
 from isaaclab_arena.metrics.metrics_logger import MetricsLogger
 from isaaclab_arena.utils.isaaclab_utils.simulation_app import (
@@ -43,7 +43,7 @@ def load_env(
     variations: list[str] | None = None,
     render_mode: str | None = None,
     language_instruction: str | None = None,
-    disable_auto_reset: bool = False,
+    datagen_run_manager: DatagenRunManagerBase | None = None,
 ):
 
     args_parser = get_isaaclab_arena_environments_cli_parser()
@@ -59,13 +59,14 @@ def load_env(
     if env_cfg.recorders is not None:
         env_cfg.recorders.dataset_filename = f"dataset_{job_name}"
 
-    # Datagen: disable the env's in-step() auto-reset (and metric recorders) so the rollout
-    # loop can drive episode boundaries and resets explicitly (see prepare_env_cfg_for_datagen).
-    reset_terms = prepare_env_cfg_for_datagen(env_cfg) if disable_auto_reset else None
+    # 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, reset_terms
+    return env
 
 
 def list_variations(eval_jobs_config: dict) -> None:
@@ -212,11 +213,13 @@ def main(datagen_run_manager: DatagenRunManagerBase | None = None):
     """Run all jobs in an eval config.
 
     Args:
-        datagen_run_manager: Optional datagen run manager. It spawns and owns one
-            data collector per job (implementations live in the collecting package).
-            When ``None``, no datagen collection runs and behavior matches the plain
-            evaluation path.
+        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()
 
@@ -251,22 +254,20 @@ def main(datagen_run_manager: DatagenRunManagerBase | None = None):
         return
 
     # Optional top-level datagen defaults, applied to every job (a per-job "datagen"
-    # block overrides these). When present, datagen data collection runs alongside
-    # each rollout, writing one HDF5 file per episode.
+    # block overrides these). When present, the injected run manager collects data
+    # alongside each rollout.
     datagen_defaults = eval_jobs_config.get("datagen")
 
-    # Injecting a run manager enables collection for every job; start_run asserts that
-    # the config actually configures collection (datagen block with an output_dir).
-    datagen_is_enabled = datagen_run_manager is not None
-    if datagen_is_enabled:
-        datagen_run_manager.start_run(eval_jobs_config, args_cli.datagen_description, args_cli.device)
+    # 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_defaults or any(job_dict.get("datagen") for job_dict in eval_jobs_config["jobs"]):
+    if datagen_run_manager.needs_cameras():
         args_cli.enable_cameras = True
 
     with SimulationAppContext(args_cli):
@@ -313,15 +314,13 @@ def main(datagen_run_manager: DatagenRunManagerBase | None = None):
                         video_base_dir=job_output_dir,
                         camera_name_prefix=f"robot-cam-rebuild{rebuild_idx}",
                     )
-                    # Datagen: disable the env's auto-reset so the rollout drives episode
-                    # boundaries and resets explicitly (see prepare_env_cfg_for_datagen).
-                    env, datagen_reset_terms = load_env(
+                    env = load_env(
                         job.arena_env_args,
                         job.name,
                         variations=job.variations,
                         render_mode=video_cfg.render_mode,
                         language_instruction=job.language_instruction,
-                        disable_auto_reset=datagen_is_enabled,
+                        datagen_run_manager=datagen_run_manager,
                     )
 
                     # Write per-episode results to disk.
@@ -346,16 +345,7 @@ def main(datagen_run_manager: DatagenRunManagerBase | None = None):
 
                     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)
-                        if datagen_is_enabled
-                        else None
-                    )
-
-                    # Per-episode step cap for datagen, bounded by the env's own cap.
-                    max_episode_length = int(env.unwrapped.max_episode_length)
-                    if collector is not None:
-                        max_episode_length = collector.cap_episode_length(max_episode_length)
+                    collector = datagen_run_manager.create_collector(job.name, datagen_job_dict, env)
 
                     metrics = rollout_policy(
                         env,
@@ -364,8 +354,6 @@ def main(datagen_run_manager: DatagenRunManagerBase | None = None):
                         num_episodes=num_episodes_this_rebuild,
                         language_instruction=job.language_instruction,
                         collector=collector,
-                        datagen_reset_terms=datagen_reset_terms,
-                        max_episode_length=max_episode_length,
                     )
 
                     job_manager.complete_job(job, metrics=metrics, status=Status.COMPLETED)
@@ -383,14 +371,11 @@ def main(datagen_run_manager: DatagenRunManagerBase | None = None):
 
                 finally:
                     try:
-                        # on_job_finished closes the job's collector, releasing its cameras
-                        # BEFORE the stage teardown so their replicator annotators do not
-                        # leak into the next job.
-                        if collector is not None:
-                            try:
-                                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}")
+                        # 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:
                         try:
                             _close_job_resources(policy, env)
@@ -405,8 +390,7 @@ def main(datagen_run_manager: DatagenRunManagerBase | None = None):
                 aggregated_metrics = aggregate_metrics(metrics_per_run)
                 metrics_logger.append_job_metrics(job.name, aggregated_metrics)
 
-        if datagen_is_enabled:
-            datagen_run_manager.finish_run()
+        datagen_run_manager.finish_run()
 
         job_manager.print_jobs_info()
         metrics_logger.print_metrics()
diff --git a/isaaclab_arena/evaluation/policy_runner.py b/isaaclab_arena/evaluation/policy_runner.py
index fc127f1a01..060dead825 100644
--- a/isaaclab_arena/evaluation/policy_runner.py
+++ b/isaaclab_arena/evaluation/policy_runner.py
@@ -7,11 +7,10 @@
 
 import argparse
 import os
-import dataclasses
 import torch
 import tqdm
 from importlib import import_module
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 from isaaclab_arena.assets.registries import PolicyRegistry
 from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser
@@ -20,8 +19,11 @@
     add_policy_runner_arguments,
     build_policy_from_cli,
 )
-from isaaclab_arena.evaluation.datagen_collector import DatagenCollectorBase
-from isaaclab_arena.evaluation.episode_outcome import classify_outcome
+from isaaclab_arena.evaluation.datagen_collector import (
+    DatagenCollectorBase,
+    EpisodeOutcome,
+    NoOpDatagenCollector,
+)
 from isaaclab_arena.evaluation.policy_runner_cli import add_policy_runner_arguments
 from isaaclab_arena.metrics.metrics_logger import metrics_to_plain_python_types
 from isaaclab_arena.utils.hydra_overrides import assert_hydra_overrides
@@ -67,101 +69,22 @@ def is_distributed(args_cli: argparse.Namespace) -> bool:
     )
 
 
-def prepare_env_cfg_for_datagen(env_cfg) -> list[tuple[str, Any]]:
-    """Prepare *env_cfg* for datagen collection. Call on the cfg *before* the env is built.
-
-    Two changes, both so the dedicated datagen cameras capture clean frames:
+def _datagen_episode_outcome(env, truncated) -> EpisodeOutcome:
+    """Best-effort outcome for an auto-reset datagen episode from the termination manager.
 
-    1. Remove the termination terms so the env never auto-resets inside ``step()``. Isaac
-       Lab resets a done env *within* ``step()`` and re-renders before the new scene is
-       flushed, so a dedicated camera reads back the previous episode's final frame as the
-       first frame of the next episode. The rollout loop instead evaluates the returned
-       terms manually and drives a clean, explicit ``env.reset()`` between episodes (which
-       flushes to the renderer before re-rendering). Mirrors
-       ``submodules/IsaacLab/scripts/tools/record_demos.py``.
-    2. Drop the metrics and their recorder terms. The datagen collector writes its own
-       per-episode HDF5, and the success-rate recorder asserts the (now removed) success
-       termination is active, so it must not run.
-
-    Returns:
-        The stashed non-timeout terms as ``(name, term)`` pairs (e.g. success, object_dropped)
-        for manual evaluation. Timeout terms are replaced by the env's ``max_episode_length``
-        cap in the loop.
-    """
-    # Datagen has its own writer; drop the env's metrics + recorder terms (the success
-    # recorder also depends on the success termination removed below). recorders=None makes
-    # the RecorderManager a no-op.
-    if hasattr(env_cfg, "metrics"):
-        env_cfg.metrics = None
-    if hasattr(env_cfg, "recorders"):
-        env_cfg.recorders = None
-
-    terminations = getattr(env_cfg, "terminations", None)
-    if terminations is None:
-        return []
-    stashed = []
-    for field in dataclasses.fields(terminations):
-        term = getattr(terminations, field.name)
-        # Skip unset/empty fields; real termination terms are TerminationTermCfg (have .func).
-        if term is None or not hasattr(term, "func"):
-            continue
-        setattr(terminations, field.name, None)
-        is_timeout = field.name == "time_out" or getattr(term, "time_out", False)
-        if not is_timeout:
-            stashed.append((field.name, term))
-    return stashed
-
-
-def _manual_episode_done(env, reset_terms: list) -> str | None:
-    """Return the name of the first stashed termination term that fires, else ``None``."""
-    base_env = env.unwrapped
-    for name, term in reset_terms:
-        result = term.func(base_env, **(term.params or {}))
-        if bool(torch.as_tensor(result).any()):
-            return name
-    return None
-
-
-def _run_datagen_rollout(
-    env, policy, collector, pbar, num_steps, num_episodes, reset_terms, max_episode_length, obs
-) -> None:
-    """Rollout loop for datagen collection with the env's auto-reset disabled.
-
-    Records every (settled) frame, decides episode end from the stashed termination terms
-    plus a max-length cap, then flushes the episode and performs a clean explicit
-    ``env.reset()`` before the next one so the first frame of each episode is correct.
+    Right after the auto-reset inside step(), the termination manager still holds the buffers
+    that triggered it: a fired success term maps to "success", a truncation to "timeout", and
+    anything else (e.g. object dropped) to "failure".
     """
-    assert max_episode_length is not None, "datagen rollout requires max_episode_length"
-    num_episodes_completed = 0
-    num_steps_completed = 0
-    steps_in_episode = 0
-    while True:
-        with torch.inference_mode():
-            actions = policy.get_action(env, obs)
-            obs, _, _, _, _ = env.step(actions)
-        steps_in_episode += 1
-        collector.on_step(env, obs, actions, num_steps_completed)
-        num_steps_completed += 1
-
-        if num_steps is not None:
-            pbar.update(1)
-            if num_steps_completed >= num_steps:
-                break
-
-        ended_by = _manual_episode_done(env, reset_terms)
-        hit_cap = steps_in_episode >= max_episode_length
-        if ended_by is not None or hit_cap:
-            collector.on_episode_end(env, outcome=classify_outcome(ended_by))
-            num_episodes_completed += 1
-            if num_episodes is not None:
-                pbar.update(1)
-                if num_episodes_completed >= num_episodes:
-                    break
-            # The reset must come after on_episode_end: its RTX rerenders (num_rerenders_on_reset)
-            # flush any camera poses the collector re-aimed for the next episode.
-            obs, _ = env.reset()
-            policy.reset()
-            steps_in_episode = 0
+    try:
+        termination_manager = env.unwrapped.termination_manager
+        if "success" in termination_manager.active_terms and bool(
+            termination_manager.get_term("success")[0]
+        ):
+            return "success"
+    except Exception:  # best-effort provenance: never fail a rollout on outcome labelling
+        pass
+    return "timeout" if bool(torch.as_tensor(truncated)[0]) else "failure"
 
 
 def rollout_policy(
@@ -171,30 +94,19 @@ def rollout_policy(
     num_episodes: int | None,
     language_instruction: str | None = None,
     collector: DatagenCollectorBase | None = None,
-    datagen_reset_terms: list | None = None,
-    max_episode_length: int | None = None,
 ) -> MetricsDataCollection | None:
     """Roll out *policy* in *env*.
 
     Args:
-        collector: Optional data collector. When provided, ``collector.on_step(env,
-            obs, actions, step_idx)`` is called after every environment step and
-            ``collector.finalize(env)`` once the rollout finishes. Implementations
-            live in the collecting package (see :class:`DatagenCollectorBase`).
-        datagen_reset_terms: Stashed termination terms returned by
-            :func:`prepare_env_cfg_for_datagen`. Required whenever a collector is
-            given: with the env's auto-reset disabled, the loop evaluates these terms
-            (plus ``max_episode_length``) to end episodes and resets explicitly, so no
-            frame is captured mid-reset.
-        max_episode_length: Per-episode step cap (replaces the removed timeout term).
-            Required when ``datagen_reset_terms`` is given.
+        collector: Data collector driven during the rollout, defaulting to a no-op. on_step
+            runs after every env step, on_episode_end when the env resets, and finalize when
+            the rollout finishes. On a done step the returned obs is already the next episode's
+            first frame, so the finished episode is closed before that obs is recorded. See
+            DatagenCollectorBase.
     """
     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"
-    assert collector is None or datagen_reset_terms is not None, (
-        "A datagen collector requires the env's auto-reset to be disabled; pass datagen_reset_terms"
-        " from prepare_env_cfg_for_datagen() (and max_episode_length)."
-    )
+    collector = collector or NoOpDatagenCollector()
 
     pbar = None
     try:
@@ -208,48 +120,42 @@ def rollout_policy(
         else:
             pbar = tqdm.tqdm(total=num_episodes, desc="Episodes", unit="episode")
 
-        if collector is not None:
-            # Datagen path: env auto-reset is disabled, so we drive episode boundaries
-            # and resets explicitly (see _run_datagen_rollout / record_demos.py).
-            _run_datagen_rollout(
-                env, policy, collector, pbar, num_steps, num_episodes, datagen_reset_terms, max_episode_length, obs
-            )
-        else:
-            num_episodes_completed = 0
-            num_steps_completed = 0
-
-            while True:
-                with torch.inference_mode():
-                    actions = policy.get_action(env, obs)
-                    obs, _, terminated, truncated, _ = env.step(actions)
-
-                    if terminated.any() or truncated.any():
-                        # Only reset policy for those envs that are terminated or truncated
-                        print(
-                            f"Resetting policy for terminated env_ids: {terminated.nonzero().flatten()}"
-                            f" and truncated env_ids: {truncated.nonzero().flatten()}"
-                        )
-                        env_ids = (terminated | truncated).nonzero().flatten()
-                        policy.reset(env_ids=env_ids)
-                        # Break if number of episodes is reached
-                        completed_episodes = env_ids.shape[0]
-                        num_episodes_completed += completed_episodes
-                        if hasattr(env.unwrapped.cfg, "metrics") and env.unwrapped.cfg.metrics is not None:
-                            metrics = env.unwrapped.compute_metrics()
-                            tqdm.tqdm.write(
-                                f"[Rank {get_local_rank()}/{get_world_size()}] Metrics:"
-                                f" {metrics_to_plain_python_types(metrics)}"
-                            )
-                        if num_episodes is not None:
-                            pbar.update(completed_episodes)
-                            if num_episodes_completed >= num_episodes:
-                                break
-                    # Break if number of steps is reached
-                    num_steps_completed += 1
-                    if num_steps is not None:
-                        pbar.update(1)
-                        if num_steps_completed >= num_steps:
-                            break
+        num_episodes_completed = 0
+        num_steps_completed = 0
+
+        while True:
+            with torch.inference_mode():
+                actions = policy.get_action(env, obs)
+                obs, _, terminated, truncated, _ = env.step(actions)
+
+            done = bool(terminated.any() or truncated.any())
+
+            if done:
+                # Close the finished episode before recording obs: the auto-reset already made
+                # it the next episode's first frame, so it must be attributed there, not here.
+                collector.on_episode_end(env, outcome=_datagen_episode_outcome(env, truncated))
+                env_ids = (terminated | truncated).nonzero().flatten()
+                policy.reset(env_ids=env_ids)
+                completed_episodes = env_ids.shape[0]
+                num_episodes_completed += completed_episodes
+                if hasattr(env.unwrapped.cfg, "metrics") and env.unwrapped.cfg.metrics is not None:
+                    metrics = env.unwrapped.compute_metrics()
+                    tqdm.tqdm.write(
+                        f"[Rank {get_local_rank()}/{get_world_size()}] Metrics:"
+                        f" {metrics_to_plain_python_types(metrics)}"
+                    )
+                if num_episodes is not None:
+                    pbar.update(completed_episodes)
+                    if num_episodes_completed >= num_episodes:
+                        break
+
+            collector.on_step(env, obs, actions, num_steps_completed)
+
+            num_steps_completed += 1
+            if num_steps is not None:
+                pbar.update(1)
+                if num_steps_completed >= num_steps:
+                    break
 
         pbar.close()
 
@@ -261,8 +167,7 @@ def rollout_policy(
     else:
 
         # Persist and close any datagen dataset before returning.
-        if collector is not None:
-            collector.finalize(env)
+        collector.finalize(env)
 
         # Only compute metrics if env has non-None metrics.
         # Use unwrapped to reach the base env through any gym wrappers (e.g. OrderEnforcing)
diff --git a/isaaclab_arena/tests/test_episode_outcome.py b/isaaclab_arena/tests/test_episode_outcome.py
deleted file mode 100644
index b6281a625e..0000000000
--- a/isaaclab_arena/tests/test_episode_outcome.py
+++ /dev/null
@@ -1,19 +0,0 @@
-# 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
-"""Unit tests for episode outcome classification (no Isaac Sim required)."""
-
-from isaaclab_arena.evaluation.episode_outcome import classify_outcome
-
-
-def test_success_term_maps_to_success():
-    assert classify_outcome("success") == "success"
-
-
-def test_no_term_means_timeout():
-    assert classify_outcome(None) == "timeout"
-
-
-def test_other_term_means_failure():
-    assert classify_outcome("object_dropped") == "failure"