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/datagen_collector.py b/isaaclab_arena/evaluation/datagen_collector.py new file mode 100644 index 0000000000..5aaff76f17 --- /dev/null +++ b/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 diff --git a/isaaclab_arena/evaluation/eval_runner.py b/isaaclab_arena/evaluation/eval_runner.py index d2dfc0979e..46190ec3ed 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, 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() 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..060dead825 100644 --- a/isaaclab_arena/evaluation/policy_runner.py +++ b/isaaclab_arena/evaluation/policy_runner.py @@ -19,6 +19,12 @@ add_policy_runner_arguments, build_policy_from_cli, ) +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 from isaaclab_arena.utils.isaaclab_utils.simulation_app import SimulationAppContext @@ -63,14 +69,44 @@ def is_distributed(args_cli: argparse.Namespace) -> bool: ) +def _datagen_episode_outcome(env, truncated) -> EpisodeOutcome: + """Best-effort outcome for an auto-reset datagen episode from the termination manager. + + 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". + """ + 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( env, policy: PolicyBase, num_steps: int | None, num_episodes: int | None, + language_instruction: str | None = None, + collector: DatagenCollectorBase | None = None, ) -> MetricsDataCollection | None: + """Roll out *policy* in *env*. + + Args: + 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" + collector = collector or NoOpDatagenCollector() pbar = None try: @@ -92,34 +128,35 @@ def rollout_policy( 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()}" + 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)}" ) - 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: + 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() except Exception as e: @@ -129,6 +166,9 @@ def rollout_policy( else: + # Persist and close any datagen dataset before returning. + 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 +225,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 +269,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_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), (