From a058fd8e45f6ef4c804a5f75b2f22a023cd4126b Mon Sep 17 00:00:00 2001 From: Billy Okal Date: Tue, 27 Oct 2026 13:54:04 -0500 Subject: [PATCH 1/6] feat: record metric inputs and run metadata during generation Generation, annotation and retargeting now record what dataset metrics need, without computing any metric: - Every producer records step_success from the task's success term. Generation also records the per-step EEF and commanded target poses, but not object poses, which annotation keeps for the source pool. - Each exported dataset gets an autodata entry in env_args: run id, exported count, pose-action layout, robot joint names with simulator limits, and optional datasheet hardware_limits from the embodiment. - --result_file adds a run block (run id, seed, revision, input digests, hardware, setup and generation time) and an attempt log, one line per attempt with its failure cause and algorithm extras. - GenerationAlgorithm.attempt_extras() lets an algorithm report per-attempt facts; SkillGen reports its motion-planning time. --- autodata_core/algorithms.py | 17 ++- autodata_core/data_generator.py | 25 ++++- .../embodiments/franka_ik_rel.yaml | 13 +++ .../embodiments/franka_ik_rel_skillgen.yaml | 13 +++ .../absolute_pose_embodiment_adapter.py | 9 ++ .../bimanual_embodiment_adapter.py | 8 ++ .../embodiments/embodiment_adapter.py | 9 ++ autodata_interfaces/embodiments/factory.py | 3 +- .../single_arm_embodiment_adapter.py | 6 ++ autodata_interfaces/env/__init__.py | 2 + .../env/datagen_info_recorder.py | 79 ++++++++++++++ .../env/isaaclab_env_interface.py | 37 ++++++- autodata_retargeting/runner.py | 9 +- autodata_tests/core/test_data_generator.py | 17 ++- autodata_tests/e2e/test_annotation.py | 4 +- .../e2e/test_mimicgen_data_generation.py | 6 +- .../e2e/test_skillgen_bin_data_generation.py | 5 +- autodata_tests/utils/utils.py | 41 +++++++ autodata_utils/dataset_metadata.py | 38 +++++++ autodata_utils/generation_result.py | 51 +++++++++ scripts/annotate_demos.py | 61 +++++------ scripts/generate_dataset.py | 101 +++++++++++++++++- .../references/cli-reference.md | 2 +- 23 files changed, 507 insertions(+), 49 deletions(-) create mode 100644 autodata_interfaces/env/datagen_info_recorder.py create mode 100644 autodata_utils/dataset_metadata.py diff --git a/autodata_core/algorithms.py b/autodata_core/algorithms.py index 65d6daf..8c969be 100644 --- a/autodata_core/algorithms.py +++ b/autodata_core/algorithms.py @@ -15,7 +15,8 @@ from __future__ import annotations -from collections.abc import Iterator +import time +from collections.abc import Iterable, Iterator from typing import TYPE_CHECKING, Any if TYPE_CHECKING: @@ -71,6 +72,14 @@ def validate_setup(self, datastream: Datastream) -> None: or ``datastream.get_env()`` (escape hatch) to enforce algorithm-specific invariants. """ + def attempt_extras(self, eef_states: Iterable[_EEFGenerationState]) -> dict[str, Any]: + """Algorithm-specific facts about one finished attempt, written to its attempt-log line. Default: none. + + Args: + eef_states: Per-EEF generation state of the attempt. + """ + return {} + def plan_subtask_trajectory( self, *, @@ -190,6 +199,10 @@ def __init__(self, motion_planners: dict[int, Any]) -> None: assert motion_planners, "SkillGen requires at least one motion planner" self.motion_planners = motion_planners + def attempt_extras(self, eef_states: Iterable[_EEFGenerationState]) -> dict[str, Any]: + """Report the attempt's wall-clock motion-planning time [s] as ``planning_time_s``.""" + return {"planning_time_s": sum(state.planning_time_s for state in eef_states)} + def plan_subtask_trajectory( self, *, @@ -238,6 +251,7 @@ def plan_subtask_trajectory( ) planner = self.motion_planners[env_id] + start_time = time.perf_counter() planning_success = planner.update_world_and_plan_motion( target_pose=target_pose, expected_attached_object=expected_attached_object, @@ -245,6 +259,7 @@ def plan_subtask_trajectory( step_size=getattr(planner, "step_size", None), enable_retiming=(hasattr(planner, "step_size") and planner.step_size is not None), ) + eef_state.planning_time_s += time.perf_counter() - start_time if not planning_success: return None diff --git a/autodata_core/data_generator.py b/autodata_core/data_generator.py index 59481cb..59d1e84 100644 --- a/autodata_core/data_generator.py +++ b/autodata_core/data_generator.py @@ -15,7 +15,7 @@ import torch from copy import deepcopy from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal from isaaclab.managers import TerminationTermCfg @@ -47,10 +47,21 @@ async def _optional_lock(lock: asyncio.Lock | None): @dataclass(frozen=True) class GenerationResult: - """Output of one :meth:`DataGenerator.generate` invocation.""" + """Output of one :meth:`DataGenerator.generate` invocation. + + Attributes: + initial_state: Scene state of every env right after this env's reset. + success: Whether the task succeeded at any step of the attempt. + failure_cause: ``"planner"`` when motion planning aborted the attempt, otherwise ``"execution"`` + for a failed attempt and None on success. + extras: Algorithm-specific facts about the attempt, from + :meth:`~autodata_core.algorithms.GenerationAlgorithm.attempt_extras`. + """ initial_state: dict success: bool + failure_cause: Literal["planner", "execution"] | None = None + extras: dict[str, Any] = field(default_factory=dict) @dataclass @@ -76,6 +87,7 @@ class _EEFGenerationState: # subtask trajectory here so the follow-up call splices it in. pending_subtask_trajectory: WaypointTrajectory | None = None is_in_motion_plan: bool = False + planning_time_s: float = 0.0 class DataGenerator: @@ -569,7 +581,12 @@ async def generate( ) if result is None: # Planning failure (e.g. SkillGen motion planner). Abort with no success. - return GenerationResult(initial_state=initial_state, success=False) + return GenerationResult( + initial_state=initial_state, + success=False, + failure_cause="planner", + extras=self.algorithm.attempt_extras(eef_states.values()), + ) next_trajectory, _is_motion_plan = result eef_state.current_trajectory = next_trajectory eef_state.subtask_step_index = 0 @@ -620,6 +637,8 @@ async def generate( return GenerationResult( initial_state=initial_state, success=buffers.success, + failure_cause=None if buffers.success else "execution", + extras=self.algorithm.attempt_extras(eef_states.values()), ) # ------------------------------------------------------------------ diff --git a/autodata_examples/embodiments/franka_ik_rel.yaml b/autodata_examples/embodiments/franka_ik_rel.yaml index 2ad1e97..3194ae2 100644 --- a/autodata_examples/embodiments/franka_ik_rel.yaml +++ b/autodata_examples/embodiments/franka_ik_rel.yaml @@ -31,3 +31,16 @@ action_layout: # that same frame, so no shift is needed. SkillGen's cube-stack dataset is annotated in the # panda_hand frame, so it needs [0, 0, 0.1034] — use franka_cube_stack_embodiment_skillgen.yaml. eef_offset: [0.0, 0.0, 0.0] + +# Franka FCI datasheet joint limits, keyed by joint-name pattern. Read only by the dataset metrics. +hardware_limits: + velocity: # [rad/s] + panda_joint[1-4]: 2.175 + panda_joint[5-7]: 2.61 + acceleration: # [rad/s^2] + panda_joint1: 15.0 + panda_joint2: 7.5 + panda_joint3: 10.0 + panda_joint4: 12.5 + panda_joint5: 15.0 + panda_joint[6-7]: 20.0 diff --git a/autodata_examples/embodiments/franka_ik_rel_skillgen.yaml b/autodata_examples/embodiments/franka_ik_rel_skillgen.yaml index 5e27381..618a461 100644 --- a/autodata_examples/embodiments/franka_ik_rel_skillgen.yaml +++ b/autodata_examples/embodiments/franka_ik_rel_skillgen.yaml @@ -35,3 +35,16 @@ action_layout: # approach axis) while the SkillGen dataset is annotated in the panda_hand frame, so the adapter # shifts reports back to panda_hand to keep planned transits and recorded skill segments aligned. eef_offset: [0.0, 0.0, 0.1034] + +# Franka FCI datasheet joint limits, keyed by joint-name pattern. Read only by the dataset metrics. +hardware_limits: + velocity: # [rad/s] + panda_joint[1-4]: 2.175 + panda_joint[5-7]: 2.61 + acceleration: # [rad/s^2] + panda_joint1: 15.0 + panda_joint2: 7.5 + panda_joint3: 10.0 + panda_joint4: 12.5 + panda_joint5: 15.0 + panda_joint[6-7]: 20.0 diff --git a/autodata_interfaces/embodiments/absolute_pose_embodiment_adapter.py b/autodata_interfaces/embodiments/absolute_pose_embodiment_adapter.py index c1d6c09..33cc04c 100644 --- a/autodata_interfaces/embodiments/absolute_pose_embodiment_adapter.py +++ b/autodata_interfaces/embodiments/absolute_pose_embodiment_adapter.py @@ -77,6 +77,15 @@ def from_dict(cls, data: dict[str, Any]) -> AbsolutePoseEmbodimentAdapter: """Build the adapter from a validated descriptor.""" return cls(**data) + def metric_context(self) -> dict[str, Any]: + """Adds the pose-action columns of every EEF that commands a pose; they are not clipped.""" + pose_action = { + name: {"columns": list(spec["pose_action_slice"]), "bound": None} + for name, spec in self.eefs.items() + if spec.get("pose_action_slice") is not None + } + return {**super().metric_context(), "pose_action": pose_action} + def get_eef_names(self) -> tuple[str, ...]: return tuple(self.eefs) diff --git a/autodata_interfaces/embodiments/bimanual_embodiment_adapter.py b/autodata_interfaces/embodiments/bimanual_embodiment_adapter.py index 206cc86..ad698e4 100644 --- a/autodata_interfaces/embodiments/bimanual_embodiment_adapter.py +++ b/autodata_interfaces/embodiments/bimanual_embodiment_adapter.py @@ -191,6 +191,14 @@ def action_dim(self) -> int: channel_ends = [channel_slice[1] for channel_slice in self.passthrough_channels.values()] return max([self.hand_joints_slice[1], *channel_ends]) + def metric_context(self) -> dict[str, Any]: + """Adds the pose-action columns of both hands; they are not clipped.""" + pose_action = { + self.left.name: {"columns": list(self.left_pose_slice), "bound": None}, + self.right.name: {"columns": list(self.right_pose_slice), "bound": None}, + } + return {**super().metric_context(), "pose_action": pose_action} + def action_to_target_eef_pose(self, action: torch.Tensor) -> dict[str, torch.Tensor]: """Convert env action to target EEF poses for both arms. diff --git a/autodata_interfaces/embodiments/embodiment_adapter.py b/autodata_interfaces/embodiments/embodiment_adapter.py index f230bfe..a219c6b 100644 --- a/autodata_interfaces/embodiments/embodiment_adapter.py +++ b/autodata_interfaces/embodiments/embodiment_adapter.py @@ -50,6 +50,15 @@ class EmbodimentAdapter(ABC): robot_asset_name: str = "robot" raw_config: dict[str, Any] | None = None + def metric_context(self) -> dict[str, Any]: + """Facts about this embodiment that the dataset metrics read, stored with every dataset it records. + + Subclasses add their pose-action layout. The base adds the optional ``hardware_limits`` section of the + embodiment config (datasheet joint limits, keyed by joint-name pattern). + """ + hardware_limits = (self.raw_config or {}).get("hardware_limits") + return {"hardware_limits": hardware_limits} if hardware_limits else {} + def get_navigation_eef_name(self) -> str | None: """Return the configured navigation channel without requiring a live environment.""" diff --git a/autodata_interfaces/embodiments/factory.py b/autodata_interfaces/embodiments/factory.py index fb1cc98..6267e59 100644 --- a/autodata_interfaces/embodiments/factory.py +++ b/autodata_interfaces/embodiments/factory.py @@ -70,7 +70,8 @@ def embodiment_adapter_from_dict(data: dict[str, Any], env: Any = None) -> Embod type_key in EMBODIMENT_TYPE_REGISTRY ), f"Unknown embodiment type {type_key!r}. Registered: {sorted(EMBODIMENT_TYPE_REGISTRY)}" cls = EMBODIMENT_TYPE_REGISTRY[type_key] - payload = {k: v for k, v in data.items() if k != "type"} + # ``hardware_limits`` is read from ``raw_config`` by :meth:`EmbodimentAdapter.metric_context`. + payload = {k: v for k, v in data.items() if k not in ("type", "hardware_limits")} adapter = cls.from_dict(payload) # type: ignore[attr-defined] adapter.raw_config = data # the full parsed config, for sections no typed adapter field models if env is not None: diff --git a/autodata_interfaces/embodiments/single_arm_embodiment_adapter.py b/autodata_interfaces/embodiments/single_arm_embodiment_adapter.py index 2c332a0..3b25dec 100644 --- a/autodata_interfaces/embodiments/single_arm_embodiment_adapter.py +++ b/autodata_interfaces/embodiments/single_arm_embodiment_adapter.py @@ -133,6 +133,12 @@ class DeltaPoseIKSingleArmAdapter(SingleArmEmbodimentAdapter): clip_pose_action_to_unit: bool = True + def metric_context(self) -> dict[str, Any]: + """Adds the pose-action columns and their clip bound.""" + bound = 1.0 if self.clip_pose_action_to_unit else None + pose_action = {self.eef_name: {"columns": [0, _DELTA_POSE_ACTION_DIM], "bound": bound}} + return {**super().metric_context(), "pose_action": pose_action} + @property def action_dim(self) -> int: """Total width of the env action vector.""" diff --git a/autodata_interfaces/env/__init__.py b/autodata_interfaces/env/__init__.py index a673b88..b75c594 100644 --- a/autodata_interfaces/env/__init__.py +++ b/autodata_interfaces/env/__init__.py @@ -6,6 +6,7 @@ from autodata_interfaces.env.env_profile import EnvironmentProfile from autodata_interfaces.env.isaaclab_env_interface import ( apply_env_profile, + collect_autodata_entry, env_loop, get_env_name_from_dataset, setup_env_config, @@ -18,5 +19,6 @@ "get_env_name_from_dataset", "setup_output_paths", "env_loop", + "collect_autodata_entry", "setup_env_config", ] diff --git a/autodata_interfaces/env/datagen_info_recorder.py b/autodata_interfaces/env/datagen_info_recorder.py new file mode 100644 index 0000000..15bb80b --- /dev/null +++ b/autodata_interfaces/env/datagen_info_recorder.py @@ -0,0 +1,79 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Recorder terms that store the per-step data annotation, generation and the dataset metrics read.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg +from isaaclab.managers import RecorderTerm, RecorderTermCfg, TerminationTermCfg +from isaaclab.utils import configclass + +if TYPE_CHECKING: + from autodata_interfaces.datastream import Datastream + + +class PreStepDatagenInfoRecorder(RecorderTerm): + """Records end-effector and commanded target poses, and optionally object poses, before each step. + + Attributes: + datastream: Datastream to read. Set it on the class once the Datastream exists, since Isaac Lab + copies recorder configurations when building the env. + """ + + datastream: Datastream | None = None + + def record_pre_step(self): + assert self.datastream is not None, "Set PreStepDatagenInfoRecorder.datastream before stepping the env." + datagen_info = { + "eef_pose": self.datastream.embodiment_adapter.get_eef_poses(env_ids=None), + "target_eef_pose": self.datastream.action_to_target_eef_pose(self._env.action_manager.action), + } + if self.cfg.record_object_poses: + datagen_info["object_pose"] = self.datastream.get_object_poses() + return "obs/datagen_info", datagen_info + + +@configclass +class PreStepDatagenInfoRecorderCfg(RecorderTermCfg): + """Configuration for :class:`PreStepDatagenInfoRecorder`.""" + + class_type: type[RecorderTerm] = PreStepDatagenInfoRecorder + record_object_poses: bool = True + """Annotation needs object poses for the source pool; generation leaves them out, as no metric reads them.""" + + +class PostStepSuccessRecorder(RecorderTerm): + """Records whether the task's success term holds after each step, as ``step_success``.""" + + def record_post_step(self): + success_term = self.cfg.success_term + return "step_success", success_term.func(self._env, **success_term.params) + + +@configclass +class PostStepSuccessRecorderCfg(RecorderTermCfg): + """Configuration for :class:`PostStepSuccessRecorder`.""" + + class_type: type[RecorderTerm] = PostStepSuccessRecorder + success_term: TerminationTermCfg | None = None + """The task's success term; the env itself runs without termination terms.""" + + +@configclass +class DatagenInfoRecorderManagerCfg(ActionStateRecorderManagerCfg): + """Action/state recorder plus the per-step ``obs/datagen_info`` term.""" + + record_pre_step_datagen_info = PreStepDatagenInfoRecorderCfg() + + +def add_step_success_recorder(recorder_cfg: object, success_term: TerminationTermCfg) -> None: + """Add the ``step_success`` term to any recorder manager configuration. + + Args: + recorder_cfg: Recorder manager configuration the env will be built with. + success_term: The task's success term. + """ + recorder_cfg.record_post_step_success = PostStepSuccessRecorderCfg(success_term=success_term) diff --git a/autodata_interfaces/env/isaaclab_env_interface.py b/autodata_interfaces/env/isaaclab_env_interface.py index 3a8f916..80754f0 100644 --- a/autodata_interfaces/env/isaaclab_env_interface.py +++ b/autodata_interfaces/env/isaaclab_env_interface.py @@ -18,7 +18,7 @@ import contextlib import os import torch -from typing import Any +from typing import TYPE_CHECKING, Any import isaaclab.sim as sim_utils from isaaclab.assets import RigidObjectCfg @@ -32,9 +32,43 @@ from isaaclab_physx.sim.schemas import PhysxRigidBodyPropertiesCfg from isaaclab_tasks.utils.parse_cfg import parse_env_cfg +from autodata_interfaces.env.datagen_info_recorder import add_step_success_recorder from autodata_interfaces.env.env_profile import EnvironmentProfile, convert_event_params from autodata_interfaces.env.pink_ik_action import PinkInverseKinematicsActionSharedUrdf from autodata_interfaces.tasks.generation_policy_spec import GenerationPolicy +from autodata_utils.dataset_metadata import AUTODATA_ENTRY_VERSION +from autodata_utils.tensor_utils import as_torch + +if TYPE_CHECKING: + from autodata_interfaces.embodiments.embodiment_adapter import EmbodimentAdapter + + +def collect_autodata_entry(env: Any, embodiment_adapter: EmbodimentAdapter, run_id: str) -> dict[str, Any]: + """The ``autodata`` entry of the datasets this env exports, read while the env is alive. + + It holds the run id, the embodiment's metric context, and the robot's joint names with env 0's simulator + position limits [rad or m] and velocity limits [rad/s or m/s]. Unbounded limits are stored as None. + :func:`~autodata_utils.dataset_metadata.merge_autodata_entry` writes it once the env has closed the file. + + Args: + env: Live environment. + embodiment_adapter: Adapter of the embodiment the env drives. + run_id: Identifier shared by everything this run writes. + """ + entry: dict[str, Any] = {"version": AUTODATA_ENTRY_VERSION, "run_id": run_id, **embodiment_adapter.metric_context()} + robot = env.scene.articulations.get(embodiment_adapter.robot_asset_name) + if robot is not None: + entry["joints"] = { + "names": list(robot.joint_names), + "position_limits": _finite_or_none(as_torch(robot.data.joint_pos_limits)[0].cpu()), + "velocity_limits": _finite_or_none(as_torch(robot.data.joint_vel_limits)[0].cpu()), + } + return entry + + +def _finite_or_none(values: torch.Tensor) -> list: + """Nested lists of ``values`` with non-finite entries as None, so the JSON stays standard.""" + return [_finite_or_none(row) if row.dim() else (row.item() if torch.isfinite(row) else None) for row in values] def setup_output_paths(output_file_path: str) -> tuple[str, str]: @@ -231,6 +265,7 @@ def setup_env_config( env_cfg.recorders = ActionStateRecorderManagerCfg() else: env_cfg.recorders = recorder_cfg + add_step_success_recorder(env_cfg.recorders, success_term) env_cfg.recorders.dataset_export_dir_path = output_dir env_cfg.recorders.dataset_filename = output_file_name diff --git a/autodata_retargeting/runner.py b/autodata_retargeting/runner.py index f1e9ddb..e4a05d2 100644 --- a/autodata_retargeting/runner.py +++ b/autodata_retargeting/runner.py @@ -7,13 +7,15 @@ import numpy as np import random import torch +import uuid from copy import deepcopy from isaaclab.utils.datasets import HDF5DatasetFileHandler from autodata_interfaces.embodiments import embodiment_adapter_from_yaml -from autodata_interfaces.env import setup_env_config, setup_output_paths +from autodata_interfaces.env import collect_autodata_entry, setup_env_config, setup_output_paths from autodata_interfaces.tasks.generation_policy_spec import GenerationPolicy +from autodata_utils.dataset_metadata import exported_dataset_files, merge_autodata_entry from .config import RetargetConfig from .eef_offset import build_eef_offsets, compose_retarget_eef_offsets_from_adapters @@ -117,8 +119,10 @@ def run(args, simulation_app) -> int: torch.manual_seed(generation_policy.seed) dataset_file_handler = None + autodata_entry = None try: target_adapter.bind_env(env) + autodata_entry = collect_autodata_entry(env, target_adapter, uuid.uuid4().hex) env.reset() # Snapshot the target's home scene state once; each replay restores its robot from here # (deepcopy so later sim steps cannot mutate the cached tensors). @@ -180,3 +184,6 @@ def run(args, simulation_app) -> int: if dataset_file_handler is not None: dataset_file_handler.close() env.close() + if autodata_entry is not None: + for dataset_file in exported_dataset_files(output_dir, output_file_name): + merge_autodata_entry(dataset_file, autodata_entry) diff --git a/autodata_tests/core/test_data_generator.py b/autodata_tests/core/test_data_generator.py index 806c0e9..6667ca1 100644 --- a/autodata_tests/core/test_data_generator.py +++ b/autodata_tests/core/test_data_generator.py @@ -196,6 +196,21 @@ def test_navigation_does_not_skip_playback_after_cancellation_or_while_paused(pa assert all(state.subtask_step_index == 4 for state in states.values()) +def test_planning_failure_is_attributed_to_the_planner(monkeypatch): + generator = _navigation_generator(active=False, reached=True) + generator.datastream.env = SimpleNamespace(recorder_manager=Mock()) + monkeypatch.setattr(generator.algorithm, "plan_subtask_trajectory", lambda **kwargs: None) + monkeypatch.setattr(generator, "_reset_environment_for_generation", AsyncMock(return_value=(torch.tensor([0]), {}))) + + result = asyncio.run( + generator.generate( + env_id=0, success_term=None, env_reset_queue=asyncio.Queue(), env_action_queue=asyncio.Queue() + ) + ) + + assert result.success is False and result.failure_cause == "planner" + + def test_unresponsive_navigation_fails_and_finalizes_the_attempt(monkeypatch): generator = _navigation_generator(active=True, reached=False) generator.datastream.generation_policy.max_navigation_hold_steps = 3 @@ -219,7 +234,7 @@ def test_unresponsive_navigation_fails_and_finalizes_the_attempt(monkeypatch): ) ) - assert result.success is False + assert result.success is False and result.failure_cause == "execution" assert execute.await_count == 3 assert recorder.set_success_to_episodes.call_args.args[1].item() is False recorder.export_episodes.assert_called_once() diff --git a/autodata_tests/e2e/test_annotation.py b/autodata_tests/e2e/test_annotation.py index f5383d3..bedc99c 100644 --- a/autodata_tests/e2e/test_annotation.py +++ b/autodata_tests/e2e/test_annotation.py @@ -12,7 +12,7 @@ from autodata_tests.utils.constants import TestPaths from autodata_tests.utils.subprocess import run_subprocess -from autodata_tests.utils.utils import assert_valid_dataset +from autodata_tests.utils.utils import assert_recorded_episodes, assert_valid_dataset @pytest.mark.with_subprocess @@ -38,6 +38,8 @@ def test_automatic_annotation_from_dataset(tmp_path): str(output_file), ]) assert_valid_dataset(str(output_file), min_num_demos=1) + entry = assert_recorded_episodes(str(output_file)) + assert entry["annotation"]["annotated_episodes"] == entry["exported_episodes"] with h5py.File(input_file, "r") as source, h5py.File(output_file, "r") as annotated: assert len(annotated["data"]) == len(source["data"]) assert json.loads(annotated["data"].attrs["env_args"])["env_name"] == "IsaacContrib-Stack-Cube-Franka-IK-Rel" diff --git a/autodata_tests/e2e/test_mimicgen_data_generation.py b/autodata_tests/e2e/test_mimicgen_data_generation.py index cfb02b9..a97271e 100644 --- a/autodata_tests/e2e/test_mimicgen_data_generation.py +++ b/autodata_tests/e2e/test_mimicgen_data_generation.py @@ -8,7 +8,7 @@ from autodata_tests.utils.constants import TestPaths from autodata_tests.utils.subprocess import run_subprocess -from autodata_tests.utils.utils import assert_valid_dataset +from autodata_tests.utils.utils import assert_recorded_run, assert_valid_dataset HEADLESS = True GENERATION_NUM_TRIALS = 1 @@ -19,6 +19,7 @@ def _run_franka_cube_stack_mimicgen(num_envs: int, device: str) -> None: with tempfile.TemporaryDirectory() as temp_dir: output_file = os.path.join(temp_dir, "generated.hdf5") + result_file = os.path.join(temp_dir, "generation_result.json") args = [ TestPaths.python_path, @@ -35,6 +36,8 @@ def _run_franka_cube_stack_mimicgen(num_envs: int, device: str) -> None: os.path.join(TestPaths.test_data_dir, "annotated_dataset_franka_stack_mimicgen.hdf5"), "--output_file", output_file, + "--result_file", + result_file, "--generation_num_trials", str(GENERATION_NUM_TRIALS), "--num_envs", @@ -47,6 +50,7 @@ def _run_franka_cube_stack_mimicgen(num_envs: int, device: str) -> None: run_subprocess(args) assert_valid_dataset(output_file, min_num_demos=GENERATION_NUM_TRIALS) + assert_recorded_run(output_file, result_file) @pytest.mark.with_subprocess diff --git a/autodata_tests/e2e/test_skillgen_bin_data_generation.py b/autodata_tests/e2e/test_skillgen_bin_data_generation.py index 0627f18..90b5a17 100644 --- a/autodata_tests/e2e/test_skillgen_bin_data_generation.py +++ b/autodata_tests/e2e/test_skillgen_bin_data_generation.py @@ -17,7 +17,7 @@ from autodata_tests.utils.constants import TestPaths from autodata_tests.utils.subprocess import run_subprocess -from autodata_tests.utils.utils import assert_valid_dataset +from autodata_tests.utils.utils import assert_recorded_run, assert_valid_dataset HEADLESS = True GENERATION_NUM_TRIALS = 1 @@ -62,6 +62,9 @@ def _run_franka_bin_stack_skillgen(num_envs: int, device: str) -> None: run_subprocess(args) assert_valid_dataset(output_file, min_num_demos=GENERATION_NUM_TRIALS) + attempts = assert_recorded_run(output_file, result_file) + # Every SkillGen attempt plans at least one transit, and the algorithm reports how long it took. + assert all(attempt["extras"]["planning_time_s"] > 0 for attempt in attempts) with open(result_file, encoding="utf-8") as result_handle: result = json.load(result_handle) diff --git a/autodata_tests/utils/utils.py b/autodata_tests/utils/utils.py index f9cbde4..dd528a3 100644 --- a/autodata_tests/utils/utils.py +++ b/autodata_tests/utils/utils.py @@ -41,3 +41,44 @@ def read_generation_result(result_file: str) -> tuple[int, int]: assert num_attempts > 0, f"Generation result has invalid num_attempts: {num_attempts}." assert num_success + num_failures == num_attempts, "Generation result counts do not add up." return num_success, num_attempts + + +def assert_recorded_episodes(dataset_file: str) -> dict: + """Assert the per-step recording and the ``autodata`` entry agree with the dataset's episodes. + + Every episode has one ``step_success`` value per action, its success attribute equals whether + ``step_success`` ever held, and the entry counts every exported episode. + + Returns: + The dataset's ``autodata`` entry. + """ + + with h5py.File(dataset_file, "r") as h5_file: + data = h5_file["data"] + entry = json.loads(data.attrs["env_args"])["autodata"] + assert entry["exported_episodes"] == len(data), f"The entry counts {entry['exported_episodes']} episodes." + for name, episode in data.items(): + step_success = episode["step_success"][:] + assert len(step_success) == len(episode["actions"]), f"{name}: step_success and actions differ in length." + assert bool(step_success.any()) == bool(episode.attrs["success"]), f"{name}: step_success disagrees." + return entry + + +def assert_recorded_run(dataset_file: str, result_file: str) -> list[dict]: + """Assert a generation run's dataset, result JSON and attempt log describe the same run. + + Returns: + The attempt-log lines. + """ + + entry = assert_recorded_episodes(dataset_file) + with open(result_file, encoding="utf-8") as result_handle: + result = json.load(result_handle) + attempts_file = os.path.join(os.path.dirname(result_file), result["run"]["attempts_file"]) + with open(attempts_file, encoding="utf-8") as attempts_handle: + attempts = [json.loads(line) for line in attempts_handle if line.strip()] + assert entry["exported_episodes"] == result["num_success"], "The dataset and result count different successes." + assert len(attempts) == result["num_attempts"], "The attempt log and result count different attempts." + run_ids = {entry["run_id"], result["run"]["run_id"], *(attempt["run_id"] for attempt in attempts)} + assert len(run_ids) == 1, f"The dataset, result and attempt log carry different run ids: {run_ids}" + return attempts diff --git a/autodata_utils/dataset_metadata.py b/autodata_utils/dataset_metadata.py new file mode 100644 index 0000000..745077b --- /dev/null +++ b/autodata_utils/dataset_metadata.py @@ -0,0 +1,38 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The ``autodata`` entry of a dataset's ``env_args``: what the dataset metrics need to know about its making.""" + +from __future__ import annotations + +import h5py +import json +import os +from collections.abc import Mapping +from typing import Any + +AUTODATA_ENTRY_VERSION = 1 + + +def exported_dataset_files(output_dir: str, output_file_name: str) -> list[str]: + """The HDF5 files a recorder exported: the dataset and, when failed episodes were kept, its ``_failed`` file.""" + candidates = [os.path.join(output_dir, f"{output_file_name}{suffix}.hdf5") for suffix in ("", "_failed")] + return [path for path in candidates if os.path.exists(path)] + + +def merge_autodata_entry(dataset_path: str, entry: Mapping[str, Any]) -> None: + """Store ``entry`` under ``autodata`` in the dataset's ``env_args``, with the number of exported episodes. + + Call it after the env has closed the file. The episode count lets a reader detect a dataset merged from + several runs, since merging keeps the first file's ``env_args``. + + Args: + dataset_path: Exported HDF5 dataset. + entry: The entry, as built by :func:`~autodata_interfaces.env.collect_autodata_entry`. + """ + with h5py.File(dataset_path, "r+") as h5_file: + data = h5_file["data"] + raw_env_args = data.attrs.get("env_args", "{}") + env_args = json.loads(raw_env_args.decode() if isinstance(raw_env_args, bytes) else raw_env_args) + env_args["autodata"] = {**entry, "exported_episodes": len(data)} + data.attrs["env_args"] = json.dumps(env_args) diff --git a/autodata_utils/generation_result.py b/autodata_utils/generation_result.py index c38213d..090d4db 100644 --- a/autodata_utils/generation_result.py +++ b/autodata_utils/generation_result.py @@ -6,11 +6,15 @@ from __future__ import annotations import contextlib +import hashlib import json import os +import shutil +import subprocess import tempfile from collections.abc import Mapping from pathlib import Path +from typing import Any def write_generation_result( @@ -19,6 +23,7 @@ def write_generation_result( requested_trials: int, stats: Mapping[str, int], env_profile: Mapping[str, str | None] | None = None, + run_metadata: Mapping[str, Any] | None = None, ) -> None: """Atomically write the final outcome of a completed data-generation run. @@ -31,6 +36,7 @@ def write_generation_result( during generation, or None when the run used the unmodified base task. Recorded here because the generated dataset stores only the base env id and would otherwise not reveal that the scene was modified. + run_metadata: How the run was configured and how long it took, stored under ``run``. """ required_stats = ("num_success", "num_failures", "num_attempts") for stat_name in required_stats: @@ -51,6 +57,8 @@ def write_generation_result( } if env_profile is not None: result["env_profile"] = dict(env_profile) + if run_metadata is not None: + result["run"] = dict(run_metadata) file_descriptor, temporary_file_name = tempfile.mkstemp( prefix=f".{result_path.name}.", suffix=".tmp", dir=result_path.parent, text=True @@ -65,3 +73,46 @@ def write_generation_result( finally: with contextlib.suppress(FileNotFoundError): os.unlink(temporary_file_name) + + +def attempts_file_for(result_file: str) -> str: + """Path of the attempt log written next to ``result_file``.""" + return str(Path(result_file).with_suffix(".attempts.jsonl")) + + +class AttemptLog: + """Appends one JSON line per completed generation attempt, stamped with the run id.""" + + def __init__(self, path: str, run_id: str) -> None: + self._path = Path(path) + self._run_id = run_id + self._path.parent.mkdir(parents=True, exist_ok=True) + self._path.write_text("", encoding="utf-8") + + def write(self, record: Mapping[str, Any]) -> None: + with self._path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps({"run_id": self._run_id, **record}, sort_keys=True) + "\n") + + +def input_record(path: str | None) -> dict[str, str] | None: + """Path and SHA-256 of a run input, or None when the input was not given.""" + if path is None: + return None + digest = hashlib.sha256() + with open(path, "rb") as handle: + for chunk in iter(lambda: handle.read(1 << 20), b""): + digest.update(chunk) + return {"path": path, "sha256": digest.hexdigest()} + + +def code_revision() -> str | None: + """Git commit of the AutoData checkout, or None without git or outside a checkout.""" + if shutil.which("git") is None: + return None + result = subprocess.run( + ["git", "-C", str(Path(__file__).resolve().parent), "rev-parse", "HEAD"], + capture_output=True, + text=True, + check=False, + ) + return result.stdout.strip() if result.returncode == 0 else None diff --git a/scripts/annotate_demos.py b/scripts/annotate_demos.py index b0d6b2b..d27984c 100644 --- a/scripts/annotate_demos.py +++ b/scripts/annotate_demos.py @@ -104,23 +104,33 @@ import math # noqa: E402 import os # noqa: E402 import sys # noqa: E402 +import time # noqa: E402 import torch # noqa: E402 import traceback # noqa: E402 +import uuid # noqa: E402 from collections.abc import Callable # noqa: E402 import isaaclab_tasks # noqa: F401, E402 (registers gym envs) from isaaclab.devices import Se3Keyboard, Se3KeyboardCfg # noqa: E402 -from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg # noqa: E402 -from isaaclab.managers import DatasetExportMode, RecorderTerm, RecorderTermCfg, TerminationTermCfg # noqa: E402 -from isaaclab.utils import configclass # noqa: E402 +from isaaclab.managers import DatasetExportMode, TerminationTermCfg # noqa: E402 from isaaclab.utils.datasets import EpisodeData, HDF5DatasetFileHandler # noqa: E402 from autodata_core.pool import DataGenInfoPool # noqa: E402 from autodata_interfaces.datastream import Datastream # noqa: E402 from autodata_interfaces.embodiments import embodiment_adapter_from_yaml # noqa: E402 -from autodata_interfaces.env import get_env_name_from_dataset, setup_env_config, setup_output_paths # noqa: E402 +from autodata_interfaces.env import ( # noqa: E402 + collect_autodata_entry, + get_env_name_from_dataset, + setup_env_config, + setup_output_paths, +) +from autodata_interfaces.env.datagen_info_recorder import ( # noqa: E402 + DatagenInfoRecorderManagerCfg, + PreStepDatagenInfoRecorder, +) from autodata_interfaces.env.registration import register_external_environment # noqa: E402 from autodata_interfaces.tasks.task_descriptor import TaskDescriptor # noqa: E402 +from autodata_utils.dataset_metadata import exported_dataset_files, merge_autodata_entry # noqa: E402 is_paused = False current_action_index = 0 @@ -151,34 +161,6 @@ def mark_subtask_cb() -> None: print(f"Marked a subtask signal at action index: {current_action_index}") -# Recorder term that writes per-step datagen info read through the Datastream. -class PreStepDatagenInfoRecorder(RecorderTerm): - """Recorder term that records ``obs/datagen_info`` from the Datastream on each step.""" - - def record_pre_step(self): - assert _datastream is not None, "Datastream must be initialized before recording." - datagen_info = { - "object_pose": _datastream.get_object_poses(), - "eef_pose": _datastream.embodiment_adapter.get_eef_poses(env_ids=None), - "target_eef_pose": _datastream.action_to_target_eef_pose(self._env.action_manager.action), - } - return "obs/datagen_info", datagen_info - - -@configclass -class PreStepDatagenInfoRecorderCfg(RecorderTermCfg): - """Configuration for :class:`PreStepDatagenInfoRecorder`.""" - - class_type: type[RecorderTerm] = PreStepDatagenInfoRecorder - - -@configclass -class AnnotationRecorderManagerCfg(ActionStateRecorderManagerCfg): - """Action/state recorder plus the datagen-info term used during annotation.""" - - record_pre_step_datagen_info = PreStepDatagenInfoRecorderCfg() - - def _build_signal_names( task_descriptor: TaskDescriptor, annotate_start_signals: bool, @@ -264,13 +246,14 @@ def main() -> int: num_envs=1, device=args_cli.device, generation_policy_params=generation_policy, - recorder_cfg=AnnotationRecorderManagerCfg(), + recorder_cfg=DatagenInfoRecorderManagerCfg(), enable_cameras=args_cli.record_images, ) # Only export episodes we explicitly mark successful (i.e. fully annotated). env_cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_SUCCEEDED_ONLY env = gym.make(env_name, cfg=env_cfg).unwrapped + autodata_entry = None try: # Create the Datastream embodiment_adapter = embodiment_adapter_from_yaml(args_cli.embodiment) @@ -287,6 +270,8 @@ def main() -> int: source_pool=empty_pool, uses_start_signals=annotate_start_signals, ) + PreStepDatagenInfoRecorder.datastream = _datastream + autodata_entry = collect_autodata_entry(env, embodiment_adapter, uuid.uuid4().hex) env.reset() @@ -311,6 +296,7 @@ def main() -> int: exported_episode_count = 0 processed_episode_count = 0 + annotation_start = time.monotonic() with contextlib.suppress(KeyboardInterrupt), torch.inference_mode(): for episode_index, episode_name in enumerate(episode_names): if not simulation_app.is_running() or simulation_app.is_exiting(): @@ -342,10 +328,19 @@ def main() -> int: f"\nExported {exported_episode_count} (out of {processed_episode_count}) annotated" f" episode{'s' if exported_episode_count != 1 else ''} to {args_cli.output_file}." ) + autodata_entry["annotation"] = { + "input_episodes": len(episode_names), + "processed_episodes": processed_episode_count, + "annotated_episodes": exported_episode_count, + "time_s": time.monotonic() - annotation_start, + } print("Exiting the app.") return exported_episode_count finally: env.close() + if autodata_entry is not None: + for dataset_file in exported_dataset_files(output_dir, output_file_name): + merge_autodata_entry(dataset_file, autodata_entry) def replay_episode( diff --git a/scripts/generate_dataset.py b/scripts/generate_dataset.py index d6f5167..5a3bd48 100755 --- a/scripts/generate_dataset.py +++ b/scripts/generate_dataset.py @@ -33,6 +33,7 @@ """Launch Isaac Sim Simulator first.""" import argparse +import time from isaaclab.app import AppLauncher @@ -42,6 +43,8 @@ from autodata_interfaces.tasks.task_descriptor_utils import resolve_task_algorithm from autodata_utils.cli import callable_path +PROCESS_START = time.monotonic() + # Hardcoded to keep argparse importable without pulling in the heavy core package. # Add new algorithms here when registering them in autodata_core.algorithms. _ALG_CHOICES = ["mimicgen", "dexmimicgen", "skillgen"] @@ -97,7 +100,10 @@ "--result_file", type=str, default=None, - help="Optional JSON path for the final structured generation result.", + help=( + "Optional JSON path for the final structured generation result. An attempt log with one JSON " + "line per attempt is written next to it (.attempts.jsonl)." + ), ) parser.add_argument( "--pause_subtask", @@ -137,24 +143,41 @@ import asyncio # noqa: E402 import gymnasium as gym # noqa: E402 import numpy as np # noqa: E402 +import platform # noqa: E402 import random # noqa: E402 import sys # noqa: E402 import torch # noqa: E402 import traceback # noqa: E402 +import uuid # noqa: E402 +from pathlib import Path # noqa: E402 from typing import Any # noqa: E402 from autodata_core import DataGenerator, get_algorithm # noqa: E402 from autodata_core.algorithms import REGISTERED_ALGORITHMS # noqa: E402 +from autodata_core.data_generator import GenerationResult # noqa: E402 from autodata_interfaces.datastream import Datastream # noqa: E402 from autodata_interfaces.env import ( # noqa: E402 EnvironmentProfile, + collect_autodata_entry, env_loop, get_env_name_from_dataset, setup_env_config, setup_output_paths, ) +from autodata_interfaces.env.datagen_info_recorder import ( # noqa: E402 + DatagenInfoRecorderManagerCfg, + PreStepDatagenInfoRecorder, +) from autodata_interfaces.env.registration import load_callable, register_external_environment # noqa: E402 -from autodata_utils.generation_result import write_generation_result # noqa: E402 +from autodata_utils.dataset_metadata import exported_dataset_files, merge_autodata_entry # noqa: E402 +from autodata_utils.generation_result import ( # noqa: E402 + AttemptLog, + attempts_file_for, + code_revision, + input_record, + write_generation_result, +) +from autodata_utils.tensor_utils import as_torch # noqa: E402 async def run_data_generator( @@ -166,9 +189,11 @@ async def run_data_generator( success_term, pause_subtask: bool, stats: dict, + attempt_log: AttemptLog | None, ) -> None: - """Repeatedly call ``data_generator.generate`` and tally outcomes into ``stats``.""" + """Repeatedly call ``data_generator.generate``, tally outcomes into ``stats`` and log each attempt.""" while True: + attempt_start = time.monotonic() try: result = await data_generator.generate( env_id=env_id, @@ -187,6 +212,45 @@ async def run_data_generator( else: stats["num_failures"] += 1 stats["num_attempts"] += 1 + if attempt_log is not None: + attempt_log.write(_attempt_record(env_id, result, attempt_start, time.monotonic())) + + +def _attempt_record(env_id: int, result: GenerationResult, start_time: float, end_time: float) -> dict: + """One attempt-log line; times [s] are monotonic-clock readings.""" + return { + "env_id": env_id, + "success": result.success, + "failure_cause": result.failure_cause, + "start_time_s": start_time, + "end_time_s": end_time, + "initial_object_positions": { + name: as_torch(state["root_pose"])[env_id, :2].tolist() + for name, state in result.initial_state.get("rigid_object", {}).items() + }, + "extras": result.extras, + } + + +def _hardware(device: str) -> dict[str, str | None]: + """Model names of the GPU and CPU the run used; physics runs on the CPU when ``device`` is ``cpu``.""" + gpu = None + if torch.cuda.is_available(): + index = torch.device(device).index if device.startswith("cuda") else None + gpu = torch.cuda.get_device_name(torch.cuda.current_device() if index is None else index) + cpu = None + cpuinfo = Path("/proc/cpuinfo") + if cpuinfo.exists(): + model_lines = (line for line in cpuinfo.read_text().splitlines() if line.startswith("model name")) + cpu = next((line.split(":", 1)[1].strip() for line in model_lines), None) + return {"gpu": gpu, "num_gpus": 1 if gpu else 0, "cpu": cpu or platform.processor() or None} + + +def _generation_recorder_cfg() -> DatagenInfoRecorderManagerCfg: + """Actions and states, plus the per-step poses the dataset metrics read; object poses are left out.""" + recorder_cfg = DatagenInfoRecorderManagerCfg() + recorder_cfg.record_pre_step_datagen_info.record_object_poses = False + return recorder_cfg def build_datastream( @@ -216,6 +280,7 @@ def setup_async_generation( success_term, algorithm, pause_subtask: bool, + attempt_log: AttemptLog | None = None, ) -> dict: """Build a single :class:`DataGenerator` over ``datastream`` and ``num_envs`` async tasks.""" asyncio_event_loop = asyncio.get_event_loop() @@ -238,6 +303,7 @@ def setup_async_generation( success_term=success_term, pause_subtask=pause_subtask, stats=stats, + attempt_log=attempt_log, ) ) tasks.append(task) @@ -327,7 +393,7 @@ def main() -> None: num_envs=args_cli.num_envs, device=args_cli.device, generation_policy_params=generation_policy_params, - recorder_cfg=load_callable(args_cli.recorder_cfg)() if args_cli.recorder_cfg else None, + recorder_cfg=load_callable(args_cli.recorder_cfg)() if args_cli.recorder_cfg else _generation_recorder_cfg(), env_profile=env_profile, enable_cameras=args_cli.record_images, ) @@ -347,6 +413,9 @@ def main() -> None: task_descriptor=task_descriptor, embodiment_adapter=embodiment_adapter, ) + PreStepDatagenInfoRecorder.datastream = datastream + run_id = uuid.uuid4().hex + autodata_entry = collect_autodata_entry(env, embodiment_adapter, run_id) # SkillGen needs one curobo planner per env. Mimic/DexMimic take no kwargs. motion_planners: dict | None = None @@ -362,6 +431,8 @@ def main() -> None: alg_kwargs["motion_planners"] = motion_planners algorithm = get_algorithm(args_cli.alg, **alg_kwargs) + attempts_file = attempts_file_for(args_cli.result_file) if args_cli.result_file else None + attempt_log = AttemptLog(attempts_file, run_id) if attempts_file else None try: async_components = setup_async_generation( datastream=datastream, @@ -369,10 +440,12 @@ def main() -> None: success_term=success_term, algorithm=algorithm, pause_subtask=args_cli.pause_subtask, + attempt_log=attempt_log, ) data_gen_tasks = asyncio.ensure_future(asyncio.gather(*async_components["tasks"])) generation_completed = False + generation_start = time.monotonic() try: generation_completed = env_loop( env, @@ -383,6 +456,7 @@ def main() -> None: stats=async_components["stats"], data_gen_tasks=data_gen_tasks, ) + generation_end = time.monotonic() except asyncio.CancelledError: print("Async tasks cancelled.") finally: @@ -406,9 +480,28 @@ def main() -> None: if env_profile else None ), + run_metadata={ + "run_id": run_id, + "seed": generation_policy_params.seed, + "revision": code_revision(), + "inputs": { + "source": input_record(args_cli.input_file), + "embodiment": input_record(args_cli.embodiment), + "task_descriptor": input_record(args_cli.task_descriptor), + "env_profile": input_record(args_cli.env_profile), + }, + "device": args_cli.device, + **_hardware(args_cli.device), + "num_envs": args_cli.num_envs, + "setup_time_s": generation_start - PROCESS_START, + "generation_time_s": generation_end - generation_start, + "attempts_file": Path(attempts_file).name, + }, ) finally: env.close() + for dataset_file in exported_dataset_files(output_dir, output_file_name): + merge_autodata_entry(dataset_file, autodata_entry) if __name__ == "__main__": diff --git a/skills/autodata-generate-datasets/references/cli-reference.md b/skills/autodata-generate-datasets/references/cli-reference.md index 19a5e3b..cfb4183 100644 --- a/skills/autodata-generate-datasets/references/cli-reference.md +++ b/skills/autodata-generate-datasets/references/cli-reference.md @@ -25,7 +25,7 @@ YAML, live env, and HDF5 source dataset, then runs the selected generation algor | `--output_file ` | no | `./datasets/output_dataset.hdf5` | Destination for generated demos. | | `--generation_num_trials ` | no | `None` | Number of demos to generate. | | `--num_envs ` | no | `1` | Number of parallel simulation environments. | -| `--result_file ` | no | — | Where to write the generation result summary. | +| `--result_file ` | no | — | Where to write the generation result summary. An attempt log (`.attempts.jsonl`) is written next to it. | | `--pause_subtask` | no | off | Pause at subtask boundaries (debugging). | | `--visualize_plan` | no | off | Visualize the motion plan (SkillGen debugging). | From 96a15bde9133fcb641cde03bf28b1f57bb518f5f Mon Sep 17 00:00:00 2001 From: Billy Okal Date: Tue, 27 Oct 2026 15:36:40 -0500 Subject: [PATCH 2/6] docs: describe what generation records per step, attempt and run --- docs/pages/concepts/data_generator.rst | 35 +++++++++++++++++++++++--- 1 file changed, 32 insertions(+), 3 deletions(-) diff --git a/docs/pages/concepts/data_generator.rst b/docs/pages/concepts/data_generator.rst index b182383..514111b 100644 --- a/docs/pages/concepts/data_generator.rst +++ b/docs/pages/concepts/data_generator.rst @@ -121,9 +121,15 @@ report progress and decide when to stop: Results and Statistics ---------------------- -Three layers report what happened: - -* **Per attempt** — ``GenerationResult``: ``success`` and ``initial_state``. +Four layers report what happened: + +* **Per step** — every exported episode records ``step_success``, whether the task's success term + held after each step. Generation also records the end-effector pose and its commanded target under + ``obs/datagen_info``; object poses are recorded there only during annotation, for the source pool. +* **Per attempt** — ``GenerationResult``: ``success``, ``failure_cause`` (``planner`` or + ``execution``), ``extras`` from the algorithm's ``attempt_extras()``, and ``initial_state``. With + ``--result_file``, each completed attempt is also written as one line of + ``.attempts.jsonl``. * **Per run, in memory** — the ``stats`` counters: ``num_success``, ``num_failures``, ``num_attempts`` (``num_success + num_failures = num_attempts``). * **Per run, on disk** — with ``--result_file``, a JSON sidecar is written atomically at the @@ -141,6 +147,20 @@ Three layers report what happened: "name": "franka_bin_stack", "path": "autodata_examples/env_profiles/franka_bin_stack.yaml", "planner": "franka_stack_cube_bin" + }, + "run": { + "run_id": "3f2a9c...", + "seed": 1, + "revision": "a058fd8...", + "inputs": { + "source": {"path": "datasets/source.hdf5", "sha256": "9b1c..."}, + "embodiment": {"path": "autodata_examples/embodiments/franka_ik_rel.yaml", "sha256": "4e07..."}, + "task_descriptor": {"path": "autodata_examples/tasks/franka_cube_stack.yaml", "sha256": "c2d8..."}, + "env_profile": null + }, + "device": "cuda:0", "gpu": "...", "num_gpus": 1, "cpu": "...", "num_envs": 1, + "setup_time_s": 22.3, "generation_time_s": 228.5, + "attempts_file": "generation_result.attempts.jsonl" } } @@ -148,6 +168,12 @@ Three layers report what happened: applied — the generated dataset itself records only the base env id, so the sidecar is what documents that the scene was modified. +Every exported dataset, whether from generation, annotation or retargeting, also carries an ``autodata`` +entry in its ``env_args`` attribute. The entry holds the ``run_id`` and the number of exported +episodes, the embodiment's pose-action layout, the robot's joint names with their simulator limits, +and the embodiment's optional datasheet ``hardware_limits``. The ``run_id`` ties the dataset to its +result file and attempt log. + The GenerationAlgorithm Plug-In ------------------------------- @@ -176,6 +202,9 @@ follow. Five class attributes declare how the algorithm differs from a vanilla M The shipped algorithms — ``mimicgen``, ``dexmimicgen``, ``skillgen`` — are described in :doc:`algorithms`. +An algorithm can also override ``attempt_extras()`` to report facts about each attempt in its +attempt-log line; SkillGen reports its motion-planning time as ``planning_time_s``. + The behavioral hook: ``plan_subtask_trajectory`` ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ From 5303c8fd8d6c2cfb48232214c74aaaf2836eeebd Mon Sep 17 00:00:00 2001 From: Billy Okal Date: Tue, 27 Oct 2026 21:12:36 -0500 Subject: [PATCH 3/6] fix: keep earlier runs' records and label only this run's datasets - Each run writes its own attempt log, ..attempts.jsonl, so rerunning with the same --result_file keeps earlier logs, and an interrupted rerun leaves the earlier result and its log consistent. - A _failed dataset gets this run's autodata entry only when this run kept failed episodes, so one left by an earlier run is not relabelled. - The e2e helper checks the logged successes against the result. - Docs: only generated datasets have a result file and attempt log. --- autodata_retargeting/runner.py | 2 +- autodata_tests/utils/test_dataset_metadata.py | 14 ++++++++++++++ autodata_tests/utils/test_generation_result.py | 13 ++++++++++++- autodata_tests/utils/utils.py | 3 +++ autodata_utils/dataset_metadata.py | 16 +++++++++++++--- autodata_utils/generation_result.py | 6 +++--- docs/pages/concepts/data_generator.rst | 8 ++++---- scripts/annotate_demos.py | 2 +- scripts/generate_dataset.py | 6 +++--- .../references/cli-reference.md | 2 +- 10 files changed, 55 insertions(+), 17 deletions(-) create mode 100644 autodata_tests/utils/test_dataset_metadata.py diff --git a/autodata_retargeting/runner.py b/autodata_retargeting/runner.py index e4a05d2..cb82c9b 100644 --- a/autodata_retargeting/runner.py +++ b/autodata_retargeting/runner.py @@ -185,5 +185,5 @@ def run(args, simulation_app) -> int: dataset_file_handler.close() env.close() if autodata_entry is not None: - for dataset_file in exported_dataset_files(output_dir, output_file_name): + for dataset_file in exported_dataset_files(output_dir, output_file_name, generation_policy.keep_failed): merge_autodata_entry(dataset_file, autodata_entry) diff --git a/autodata_tests/utils/test_dataset_metadata.py b/autodata_tests/utils/test_dataset_metadata.py new file mode 100644 index 0000000..99fe7cf --- /dev/null +++ b/autodata_tests/utils/test_dataset_metadata.py @@ -0,0 +1,14 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path + +from autodata_utils.dataset_metadata import exported_dataset_files + + +def test_failed_file_left_by_an_earlier_run_is_not_exported_again(tmp_path: Path) -> None: + for name in ("generated.hdf5", "generated_failed.hdf5"): + (tmp_path / name).touch() + + assert exported_dataset_files(str(tmp_path), "generated", kept_failed=False) == [str(tmp_path / "generated.hdf5")] + assert len(exported_dataset_files(str(tmp_path), "generated", kept_failed=True)) == 2 diff --git a/autodata_tests/utils/test_generation_result.py b/autodata_tests/utils/test_generation_result.py index 31bb6b9..37fe93b 100644 --- a/autodata_tests/utils/test_generation_result.py +++ b/autodata_tests/utils/test_generation_result.py @@ -4,8 +4,9 @@ """Tests for structured data-generation result files.""" import json +from pathlib import Path -from autodata_utils.generation_result import write_generation_result +from autodata_utils.generation_result import AttemptLog, attempts_file_for, write_generation_result def test_write_generation_result_writes_completed_stats_atomically(tmp_path): @@ -52,3 +53,13 @@ def test_write_generation_result_records_env_profile_when_given(tmp_path): "path": "autodata_examples/env_profiles/franka_bin_stack.yaml", "planner": "franka_stack_cube_bin", } + + +def test_rerun_with_the_same_result_file_keeps_the_earlier_attempt_log(tmp_path: Path) -> None: + result_file = str(tmp_path / "result.json") + AttemptLog(attempts_file_for(result_file, "run-a"), "run-a").write({"success": True}) + + AttemptLog(attempts_file_for(result_file, "run-b"), "run-b") + + earlier = Path(attempts_file_for(result_file, "run-a")).read_text().splitlines() + assert [json.loads(line)["run_id"] for line in earlier] == ["run-a"] diff --git a/autodata_tests/utils/utils.py b/autodata_tests/utils/utils.py index dd528a3..c300e97 100644 --- a/autodata_tests/utils/utils.py +++ b/autodata_tests/utils/utils.py @@ -79,6 +79,9 @@ def assert_recorded_run(dataset_file: str, result_file: str) -> list[dict]: attempts = [json.loads(line) for line in attempts_handle if line.strip()] assert entry["exported_episodes"] == result["num_success"], "The dataset and result count different successes." assert len(attempts) == result["num_attempts"], "The attempt log and result count different attempts." + assert ( + sum(attempt["success"] for attempt in attempts) == result["num_success"] + ), "The attempt log and result count different successes." run_ids = {entry["run_id"], result["run"]["run_id"], *(attempt["run_id"] for attempt in attempts)} assert len(run_ids) == 1, f"The dataset, result and attempt log carry different run ids: {run_ids}" return attempts diff --git a/autodata_utils/dataset_metadata.py b/autodata_utils/dataset_metadata.py index 745077b..35ac6e9 100644 --- a/autodata_utils/dataset_metadata.py +++ b/autodata_utils/dataset_metadata.py @@ -14,9 +14,19 @@ AUTODATA_ENTRY_VERSION = 1 -def exported_dataset_files(output_dir: str, output_file_name: str) -> list[str]: - """The HDF5 files a recorder exported: the dataset and, when failed episodes were kept, its ``_failed`` file.""" - candidates = [os.path.join(output_dir, f"{output_file_name}{suffix}.hdf5") for suffix in ("", "_failed")] +def exported_dataset_files(output_dir: str, output_file_name: str, kept_failed: bool) -> list[str]: + """The HDF5 files this run's recorder exported. + + A ``_failed`` file counts only when this run kept failed episodes; otherwise one left by an earlier run would be + labelled with this run's id. + + Args: + output_dir: Export directory. + output_file_name: Dataset file name without the extension. + kept_failed: Whether the run exported failed episodes to a separate ``_failed`` file. + """ + suffixes = ("", "_failed") if kept_failed else ("",) + candidates = [os.path.join(output_dir, f"{output_file_name}{suffix}.hdf5") for suffix in suffixes] return [path for path in candidates if os.path.exists(path)] diff --git a/autodata_utils/generation_result.py b/autodata_utils/generation_result.py index 090d4db..fc5bc4e 100644 --- a/autodata_utils/generation_result.py +++ b/autodata_utils/generation_result.py @@ -75,9 +75,9 @@ def write_generation_result( os.unlink(temporary_file_name) -def attempts_file_for(result_file: str) -> str: - """Path of the attempt log written next to ``result_file``.""" - return str(Path(result_file).with_suffix(".attempts.jsonl")) +def attempts_file_for(result_file: str, run_id: str) -> str: + """Path of the run's attempt log, next to ``result_file``; each run writes its own, so a rerun keeps earlier logs.""" + return str(Path(result_file).with_suffix(f".{run_id}.attempts.jsonl")) class AttemptLog: diff --git a/docs/pages/concepts/data_generator.rst b/docs/pages/concepts/data_generator.rst index 514111b..787d88d 100644 --- a/docs/pages/concepts/data_generator.rst +++ b/docs/pages/concepts/data_generator.rst @@ -129,7 +129,7 @@ Four layers report what happened: * **Per attempt** — ``GenerationResult``: ``success``, ``failure_cause`` (``planner`` or ``execution``), ``extras`` from the algorithm's ``attempt_extras()``, and ``initial_state``. With ``--result_file``, each completed attempt is also written as one line of - ``.attempts.jsonl``. + ``..attempts.jsonl``. * **Per run, in memory** — the ``stats`` counters: ``num_success``, ``num_failures``, ``num_attempts`` (``num_success + num_failures = num_attempts``). * **Per run, on disk** — with ``--result_file``, a JSON sidecar is written atomically at the @@ -160,7 +160,7 @@ Four layers report what happened: }, "device": "cuda:0", "gpu": "...", "num_gpus": 1, "cpu": "...", "num_envs": 1, "setup_time_s": 22.3, "generation_time_s": 228.5, - "attempts_file": "generation_result.attempts.jsonl" + "attempts_file": "generation_result.3f2a9c.attempts.jsonl" } } @@ -171,8 +171,8 @@ Four layers report what happened: Every exported dataset, whether from generation, annotation or retargeting, also carries an ``autodata`` entry in its ``env_args`` attribute. The entry holds the ``run_id`` and the number of exported episodes, the embodiment's pose-action layout, the robot's joint names with their simulator limits, -and the embodiment's optional datasheet ``hardware_limits``. The ``run_id`` ties the dataset to its -result file and attempt log. +and the embodiment's optional datasheet ``hardware_limits``. For a generated dataset, the ``run_id`` +ties it to its result file and attempt log. The GenerationAlgorithm Plug-In ------------------------------- diff --git a/scripts/annotate_demos.py b/scripts/annotate_demos.py index d27984c..95d370f 100644 --- a/scripts/annotate_demos.py +++ b/scripts/annotate_demos.py @@ -339,7 +339,7 @@ def main() -> int: finally: env.close() if autodata_entry is not None: - for dataset_file in exported_dataset_files(output_dir, output_file_name): + for dataset_file in exported_dataset_files(output_dir, output_file_name, kept_failed=False): merge_autodata_entry(dataset_file, autodata_entry) diff --git a/scripts/generate_dataset.py b/scripts/generate_dataset.py index 5a3bd48..19791a5 100755 --- a/scripts/generate_dataset.py +++ b/scripts/generate_dataset.py @@ -102,7 +102,7 @@ default=None, help=( "Optional JSON path for the final structured generation result. An attempt log with one JSON " - "line per attempt is written next to it (.attempts.jsonl)." + "line per attempt is written next to it (..attempts.jsonl)." ), ) parser.add_argument( @@ -431,7 +431,7 @@ def main() -> None: alg_kwargs["motion_planners"] = motion_planners algorithm = get_algorithm(args_cli.alg, **alg_kwargs) - attempts_file = attempts_file_for(args_cli.result_file) if args_cli.result_file else None + attempts_file = attempts_file_for(args_cli.result_file, run_id) if args_cli.result_file else None attempt_log = AttemptLog(attempts_file, run_id) if attempts_file else None try: async_components = setup_async_generation( @@ -500,7 +500,7 @@ def main() -> None: ) finally: env.close() - for dataset_file in exported_dataset_files(output_dir, output_file_name): + for dataset_file in exported_dataset_files(output_dir, output_file_name, generation_policy_params.keep_failed): merge_autodata_entry(dataset_file, autodata_entry) diff --git a/skills/autodata-generate-datasets/references/cli-reference.md b/skills/autodata-generate-datasets/references/cli-reference.md index cfb4183..9062e4a 100644 --- a/skills/autodata-generate-datasets/references/cli-reference.md +++ b/skills/autodata-generate-datasets/references/cli-reference.md @@ -25,7 +25,7 @@ YAML, live env, and HDF5 source dataset, then runs the selected generation algor | `--output_file ` | no | `./datasets/output_dataset.hdf5` | Destination for generated demos. | | `--generation_num_trials ` | no | `None` | Number of demos to generate. | | `--num_envs ` | no | `1` | Number of parallel simulation environments. | -| `--result_file ` | no | — | Where to write the generation result summary. An attempt log (`.attempts.jsonl`) is written next to it. | +| `--result_file ` | no | — | Where to write the generation result summary. An attempt log (`..attempts.jsonl`) is written next to it. | | `--pause_subtask` | no | off | Pause at subtask boundaries (debugging). | | `--visualize_plan` | no | off | Visualize the motion plan (SkillGen debugging). | From 5948fc060560f3674a14203602098a52fd4328f9 Mon Sep 17 00:00:00 2001 From: Billy Okal Date: Tue, 27 Oct 2026 21:33:20 -0500 Subject: [PATCH 4/6] docs: say the result file and attempt log come with --result_file --- docs/pages/concepts/data_generator.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/pages/concepts/data_generator.rst b/docs/pages/concepts/data_generator.rst index 787d88d..6fa9454 100644 --- a/docs/pages/concepts/data_generator.rst +++ b/docs/pages/concepts/data_generator.rst @@ -171,8 +171,8 @@ Four layers report what happened: Every exported dataset, whether from generation, annotation or retargeting, also carries an ``autodata`` entry in its ``env_args`` attribute. The entry holds the ``run_id`` and the number of exported episodes, the embodiment's pose-action layout, the robot's joint names with their simulator limits, -and the embodiment's optional datasheet ``hardware_limits``. For a generated dataset, the ``run_id`` -ties it to its result file and attempt log. +and the embodiment's optional datasheet ``hardware_limits``. For a dataset generated with +``--result_file``, the ``run_id`` ties it to that run's result file and attempt log. The GenerationAlgorithm Plug-In ------------------------------- From 77b511de37d24a0993432a7707738ea6fd4947c5 Mon Sep 17 00:00:00 2001 From: Billy Okal Date: Wed, 28 Oct 2026 13:51:26 -0500 Subject: [PATCH 5/6] feat: make metric recording opt-in and keep algorithm diagnostics in extras - Metric recording is off by default, so datasets and result files are unchanged unless --record_metric_inputs is given to generation, annotation or retargeting. It records the per-step poses and task success, the autodata entry and, with --result_file, the run block and attempt log. - GenerationResult no longer has failure_cause; SkillGen reports planning_failed in its attempt extras, next to planning_time_s. - metric_context() docstrings say what they record; the pose-action entry's columns key is now action_slice. - Docs: the run linkage holds once the run completes. --- autodata_core/algorithms.py | 9 ++- autodata_core/data_generator.py | 8 +-- .../absolute_pose_embodiment_adapter.py | 7 +- .../bimanual_embodiment_adapter.py | 9 ++- .../embodiments/embodiment_adapter.py | 8 ++- .../single_arm_embodiment_adapter.py | 7 +- .../env/isaaclab_env_interface.py | 6 +- autodata_retargeting/runner.py | 4 +- autodata_tests/core/test_algorithms.py | 24 +++++++ autodata_tests/core/test_data_generator.py | 17 +---- autodata_tests/e2e/test_annotation.py | 1 + .../e2e/test_dexmimicgen_data_generation.py | 3 +- .../e2e/test_mimicgen_data_generation.py | 1 + .../e2e/test_skillgen_bin_data_generation.py | 1 + autodata_tests/utils/utils.py | 11 +++ docs/pages/concepts/data_generator.rst | 28 ++++---- scripts/annotate_demos.py | 22 ++++-- scripts/generate_dataset.py | 69 ++++++++++++------- scripts/retarget_dataset.py | 5 ++ .../references/cli-reference.md | 3 +- 20 files changed, 160 insertions(+), 83 deletions(-) diff --git a/autodata_core/algorithms.py b/autodata_core/algorithms.py index 8c969be..e0df4f8 100644 --- a/autodata_core/algorithms.py +++ b/autodata_core/algorithms.py @@ -200,8 +200,12 @@ def __init__(self, motion_planners: dict[int, Any]) -> None: self.motion_planners = motion_planners def attempt_extras(self, eef_states: Iterable[_EEFGenerationState]) -> dict[str, Any]: - """Report the attempt's wall-clock motion-planning time [s] as ``planning_time_s``.""" - return {"planning_time_s": sum(state.planning_time_s for state in eef_states)} + """Report the attempt's wall-clock motion-planning time [s] and whether a motion plan failed.""" + states = list(eef_states) + return { + "planning_time_s": sum(state.planning_time_s for state in states), + "planning_failed": any(state.planning_failed for state in states), + } def plan_subtask_trajectory( self, @@ -261,6 +265,7 @@ def plan_subtask_trajectory( ) eef_state.planning_time_s += time.perf_counter() - start_time if not planning_success: + eef_state.planning_failed = True return None mp_waypoints = self._convert_planned_trajectory_to_waypoints(planner, target_passthrough_action) diff --git a/autodata_core/data_generator.py b/autodata_core/data_generator.py index 59d1e84..3f20fa5 100644 --- a/autodata_core/data_generator.py +++ b/autodata_core/data_generator.py @@ -15,7 +15,7 @@ import torch from copy import deepcopy from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any from isaaclab.managers import TerminationTermCfg @@ -52,15 +52,12 @@ class GenerationResult: Attributes: initial_state: Scene state of every env right after this env's reset. success: Whether the task succeeded at any step of the attempt. - failure_cause: ``"planner"`` when motion planning aborted the attempt, otherwise ``"execution"`` - for a failed attempt and None on success. extras: Algorithm-specific facts about the attempt, from :meth:`~autodata_core.algorithms.GenerationAlgorithm.attempt_extras`. """ initial_state: dict success: bool - failure_cause: Literal["planner", "execution"] | None = None extras: dict[str, Any] = field(default_factory=dict) @@ -88,6 +85,7 @@ class _EEFGenerationState: pending_subtask_trajectory: WaypointTrajectory | None = None is_in_motion_plan: bool = False planning_time_s: float = 0.0 + planning_failed: bool = False class DataGenerator: @@ -584,7 +582,6 @@ async def generate( return GenerationResult( initial_state=initial_state, success=False, - failure_cause="planner", extras=self.algorithm.attempt_extras(eef_states.values()), ) next_trajectory, _is_motion_plan = result @@ -637,7 +634,6 @@ async def generate( return GenerationResult( initial_state=initial_state, success=buffers.success, - failure_cause=None if buffers.success else "execution", extras=self.algorithm.attempt_extras(eef_states.values()), ) diff --git a/autodata_interfaces/embodiments/absolute_pose_embodiment_adapter.py b/autodata_interfaces/embodiments/absolute_pose_embodiment_adapter.py index 33cc04c..08016f4 100644 --- a/autodata_interfaces/embodiments/absolute_pose_embodiment_adapter.py +++ b/autodata_interfaces/embodiments/absolute_pose_embodiment_adapter.py @@ -78,9 +78,12 @@ def from_dict(cls, data: dict[str, Any]) -> AbsolutePoseEmbodimentAdapter: return cls(**data) def metric_context(self) -> dict[str, Any]: - """Adds the pose-action columns of every EEF that commands a pose; they are not clipped.""" + """Each end effector with a ``pose_action_slice`` commands its target pose there: position, then quaternion. + + These entries are not clipped. + """ pose_action = { - name: {"columns": list(spec["pose_action_slice"]), "bound": None} + name: {"action_slice": list(spec["pose_action_slice"]), "bound": None} for name, spec in self.eefs.items() if spec.get("pose_action_slice") is not None } diff --git a/autodata_interfaces/embodiments/bimanual_embodiment_adapter.py b/autodata_interfaces/embodiments/bimanual_embodiment_adapter.py index ad698e4..e948e3b 100644 --- a/autodata_interfaces/embodiments/bimanual_embodiment_adapter.py +++ b/autodata_interfaces/embodiments/bimanual_embodiment_adapter.py @@ -192,10 +192,13 @@ def action_dim(self) -> int: return max([self.hand_joints_slice[1], *channel_ends]) def metric_context(self) -> dict[str, Any]: - """Adds the pose-action columns of both hands; they are not clipped.""" + """``left_pose_slice`` and ``right_pose_slice`` of the action vector hold the hands' target poses. + + Each holds a position, then a quaternion. These entries are not clipped. + """ pose_action = { - self.left.name: {"columns": list(self.left_pose_slice), "bound": None}, - self.right.name: {"columns": list(self.right_pose_slice), "bound": None}, + self.left.name: {"action_slice": list(self.left_pose_slice), "bound": None}, + self.right.name: {"action_slice": list(self.right_pose_slice), "bound": None}, } return {**super().metric_context(), "pose_action": pose_action} diff --git a/autodata_interfaces/embodiments/embodiment_adapter.py b/autodata_interfaces/embodiments/embodiment_adapter.py index a219c6b..d0b9859 100644 --- a/autodata_interfaces/embodiments/embodiment_adapter.py +++ b/autodata_interfaces/embodiments/embodiment_adapter.py @@ -51,10 +51,12 @@ class EmbodimentAdapter(ABC): raw_config: dict[str, Any] | None = None def metric_context(self) -> dict[str, Any]: - """Facts about this embodiment that the dataset metrics read, stored with every dataset it records. + """Describe this embodiment for the dataset metrics; the description is saved with every dataset it records. - Subclasses add their pose-action layout. The base adds the optional ``hardware_limits`` section of the - embodiment config (datasheet joint limits, keyed by joint-name pattern). + The base returns the optional ``hardware_limits`` section of the embodiment config: datasheet joint limits + keyed by joint-name pattern. Subclasses add ``pose_action``, which says for each end effector which part of + the action vector commands its pose (``action_slice``, a half-open ``[start, end)`` index range) and the + value the adapter clips those entries to (``bound``, or None when they are not clipped). """ hardware_limits = (self.raw_config or {}).get("hardware_limits") return {"hardware_limits": hardware_limits} if hardware_limits else {} diff --git a/autodata_interfaces/embodiments/single_arm_embodiment_adapter.py b/autodata_interfaces/embodiments/single_arm_embodiment_adapter.py index 3b25dec..6cd8df0 100644 --- a/autodata_interfaces/embodiments/single_arm_embodiment_adapter.py +++ b/autodata_interfaces/embodiments/single_arm_embodiment_adapter.py @@ -134,9 +134,12 @@ class DeltaPoseIKSingleArmAdapter(SingleArmEmbodimentAdapter): clip_pose_action_to_unit: bool = True def metric_context(self) -> dict[str, Any]: - """Adds the pose-action columns and their clip bound.""" + """Action entries 0-5 command the end effector's pose change: delta position, then axis-angle rotation. + + They are clipped to ±1 when ``clip_pose_action_to_unit`` is set. + """ bound = 1.0 if self.clip_pose_action_to_unit else None - pose_action = {self.eef_name: {"columns": [0, _DELTA_POSE_ACTION_DIM], "bound": bound}} + pose_action = {self.eef_name: {"action_slice": [0, _DELTA_POSE_ACTION_DIM], "bound": bound}} return {**super().metric_context(), "pose_action": pose_action} @property diff --git a/autodata_interfaces/env/isaaclab_env_interface.py b/autodata_interfaces/env/isaaclab_env_interface.py index 80754f0..9e4ebf6 100644 --- a/autodata_interfaces/env/isaaclab_env_interface.py +++ b/autodata_interfaces/env/isaaclab_env_interface.py @@ -200,6 +200,7 @@ def setup_env_config( recorder_cfg: RecorderManagerBaseCfg | None = None, env_profile: EnvironmentProfile | None = None, enable_cameras: bool = False, + record_step_success: bool = False, ) -> tuple[Any, Any]: """Configure the environment for data generation. @@ -220,6 +221,8 @@ def setup_env_config( reset-event changes) before the generation-specific adjustments below. enable_cameras: Record the task's image observations. Defaults to low-dimensional observations, excluding the cameras used for teleoperation previews. + record_step_success: Also record whether the task's success term holds after each step, as + ``step_success``, for the dataset metrics. Returns: A tuple of the environment configuration and the success termination condition. @@ -265,7 +268,8 @@ def setup_env_config( env_cfg.recorders = ActionStateRecorderManagerCfg() else: env_cfg.recorders = recorder_cfg - add_step_success_recorder(env_cfg.recorders, success_term) + if record_step_success: + add_step_success_recorder(env_cfg.recorders, success_term) env_cfg.recorders.dataset_export_dir_path = output_dir env_cfg.recorders.dataset_filename = output_file_name diff --git a/autodata_retargeting/runner.py b/autodata_retargeting/runner.py index cb82c9b..7621d1c 100644 --- a/autodata_retargeting/runner.py +++ b/autodata_retargeting/runner.py @@ -98,6 +98,7 @@ def run(args, simulation_app) -> int: num_envs=num_envs, device=args.device, generation_policy_params=generation_policy, + record_step_success=getattr(args, "record_metric_inputs", False), ) # Compensate the target env's (auto-detected) arm-action scale *in the encoded action* rather than @@ -122,7 +123,8 @@ def run(args, simulation_app) -> int: autodata_entry = None try: target_adapter.bind_env(env) - autodata_entry = collect_autodata_entry(env, target_adapter, uuid.uuid4().hex) + if getattr(args, "record_metric_inputs", False): + autodata_entry = collect_autodata_entry(env, target_adapter, uuid.uuid4().hex) env.reset() # Snapshot the target's home scene state once; each replay restores its robot from here # (deepcopy so later sim steps cannot mutate the cached tensors). diff --git a/autodata_tests/core/test_algorithms.py b/autodata_tests/core/test_algorithms.py index b3b0867..ae90f59 100644 --- a/autodata_tests/core/test_algorithms.py +++ b/autodata_tests/core/test_algorithms.py @@ -3,6 +3,9 @@ """Unit tests for :mod:`autodata_core.algorithms`.""" +from types import SimpleNamespace +from unittest.mock import Mock + import pytest from autodata_core.algorithms import ( @@ -14,6 +17,7 @@ get_algorithm, iter_algorithms, ) +from autodata_core.data_generator import _EEFGenerationState def test_registry_contents(): @@ -81,3 +85,23 @@ def test_get_algorithm_skillgen_needs_planners(): assert isinstance( get_algorithm("skillgen", motion_planners={0: object()}), SkillGen ) # base implementation must not raise + + +def test_skillgen_reports_a_failed_motion_plan(): + planner = Mock(update_world_and_plan_motion=Mock(return_value=False)) + algorithm = SkillGen(motion_planners={0: planner}) + data_generator = Mock() + data_generator.generate_eef_subtask_trajectory.return_value = [SimpleNamespace(pose=None, passthrough_action=None)] + eef_state = _EEFGenerationState() + + result = algorithm.plan_subtask_trajectory( + data_generator=data_generator, + env_id=0, + eef_name="franka", + eef_state=eef_state, + all_randomized_subtask_boundaries={}, + runtime_subtask_constraints_dict={}, + selected_src_demo_inds={}, + ) + + assert result is None and algorithm.attempt_extras([eef_state])["planning_failed"] is True diff --git a/autodata_tests/core/test_data_generator.py b/autodata_tests/core/test_data_generator.py index 6667ca1..806c0e9 100644 --- a/autodata_tests/core/test_data_generator.py +++ b/autodata_tests/core/test_data_generator.py @@ -196,21 +196,6 @@ def test_navigation_does_not_skip_playback_after_cancellation_or_while_paused(pa assert all(state.subtask_step_index == 4 for state in states.values()) -def test_planning_failure_is_attributed_to_the_planner(monkeypatch): - generator = _navigation_generator(active=False, reached=True) - generator.datastream.env = SimpleNamespace(recorder_manager=Mock()) - monkeypatch.setattr(generator.algorithm, "plan_subtask_trajectory", lambda **kwargs: None) - monkeypatch.setattr(generator, "_reset_environment_for_generation", AsyncMock(return_value=(torch.tensor([0]), {}))) - - result = asyncio.run( - generator.generate( - env_id=0, success_term=None, env_reset_queue=asyncio.Queue(), env_action_queue=asyncio.Queue() - ) - ) - - assert result.success is False and result.failure_cause == "planner" - - def test_unresponsive_navigation_fails_and_finalizes_the_attempt(monkeypatch): generator = _navigation_generator(active=True, reached=False) generator.datastream.generation_policy.max_navigation_hold_steps = 3 @@ -234,7 +219,7 @@ def test_unresponsive_navigation_fails_and_finalizes_the_attempt(monkeypatch): ) ) - assert result.success is False and result.failure_cause == "execution" + assert result.success is False assert execute.await_count == 3 assert recorder.set_success_to_episodes.call_args.args[1].item() is False recorder.export_episodes.assert_called_once() diff --git a/autodata_tests/e2e/test_annotation.py b/autodata_tests/e2e/test_annotation.py index bedc99c..fa93c2e 100644 --- a/autodata_tests/e2e/test_annotation.py +++ b/autodata_tests/e2e/test_annotation.py @@ -36,6 +36,7 @@ def test_automatic_annotation_from_dataset(tmp_path): str(input_file), "--output_file", str(output_file), + "--record_metric_inputs", ]) assert_valid_dataset(str(output_file), min_num_demos=1) entry = assert_recorded_episodes(str(output_file)) diff --git a/autodata_tests/e2e/test_dexmimicgen_data_generation.py b/autodata_tests/e2e/test_dexmimicgen_data_generation.py index fa73a6c..602c954 100644 --- a/autodata_tests/e2e/test_dexmimicgen_data_generation.py +++ b/autodata_tests/e2e/test_dexmimicgen_data_generation.py @@ -8,7 +8,7 @@ from autodata_tests.utils.constants import TestPaths from autodata_tests.utils.subprocess import run_subprocess -from autodata_tests.utils.utils import assert_valid_dataset +from autodata_tests.utils.utils import assert_no_metric_inputs, assert_valid_dataset HEADLESS = True GENERATION_NUM_TRIALS = 1 @@ -47,6 +47,7 @@ def _run_gr1_pick_place_dexmimicgen(num_envs: int, device: str) -> None: run_subprocess(args) assert_valid_dataset(output_file, min_num_demos=GENERATION_NUM_TRIALS) + assert_no_metric_inputs(output_file) def _run_g1_pick_place_dexmimicgen(num_envs: int, device: str) -> None: diff --git a/autodata_tests/e2e/test_mimicgen_data_generation.py b/autodata_tests/e2e/test_mimicgen_data_generation.py index a97271e..0552279 100644 --- a/autodata_tests/e2e/test_mimicgen_data_generation.py +++ b/autodata_tests/e2e/test_mimicgen_data_generation.py @@ -38,6 +38,7 @@ def _run_franka_cube_stack_mimicgen(num_envs: int, device: str) -> None: output_file, "--result_file", result_file, + "--record_metric_inputs", "--generation_num_trials", str(GENERATION_NUM_TRIALS), "--num_envs", diff --git a/autodata_tests/e2e/test_skillgen_bin_data_generation.py b/autodata_tests/e2e/test_skillgen_bin_data_generation.py index 90b5a17..580d086 100644 --- a/autodata_tests/e2e/test_skillgen_bin_data_generation.py +++ b/autodata_tests/e2e/test_skillgen_bin_data_generation.py @@ -50,6 +50,7 @@ def _run_franka_bin_stack_skillgen(num_envs: int, device: str) -> None: output_file, "--result_file", result_file, + "--record_metric_inputs", "--generation_num_trials", str(GENERATION_NUM_TRIALS), "--num_envs", diff --git a/autodata_tests/utils/utils.py b/autodata_tests/utils/utils.py index c300e97..4792836 100644 --- a/autodata_tests/utils/utils.py +++ b/autodata_tests/utils/utils.py @@ -85,3 +85,14 @@ def assert_recorded_run(dataset_file: str, result_file: str) -> list[dict]: run_ids = {entry["run_id"], result["run"]["run_id"], *(attempt["run_id"] for attempt in attempts)} assert len(run_ids) == 1, f"The dataset, result and attempt log carry different run ids: {run_ids}" return attempts + + +def assert_no_metric_inputs(dataset_file: str) -> None: + """Assert a dataset recorded without ``--record_metric_inputs`` has no metric inputs, as before they existed.""" + + with h5py.File(dataset_file, "r") as h5_file: + data = h5_file["data"] + assert "autodata" not in json.loads(data.attrs["env_args"]), "The dataset has an autodata entry." + for name, episode in data.items(): + assert "step_success" not in episode, f"{name} records step_success." + assert "datagen_info" not in episode["obs"], f"{name} records datagen_info." diff --git a/docs/pages/concepts/data_generator.rst b/docs/pages/concepts/data_generator.rst index 6fa9454..244451b 100644 --- a/docs/pages/concepts/data_generator.rst +++ b/docs/pages/concepts/data_generator.rst @@ -123,13 +123,13 @@ Results and Statistics Four layers report what happened: -* **Per step** — every exported episode records ``step_success``, whether the task's success term - held after each step. Generation also records the end-effector pose and its commanded target under - ``obs/datagen_info``; object poses are recorded there only during annotation, for the source pool. -* **Per attempt** — ``GenerationResult``: ``success``, ``failure_cause`` (``planner`` or - ``execution``), ``extras`` from the algorithm's ``attempt_extras()``, and ``initial_state``. With - ``--result_file``, each completed attempt is also written as one line of - ``..attempts.jsonl``. +* **Per step** — annotation records ``obs/datagen_info`` (object, end-effector and commanded + target poses) for the source pool. With ``--record_metric_inputs``, every exported episode also + records ``step_success``, whether the task's success term held after each step, and generation + records the end-effector pose and its commanded target. +* **Per attempt** — ``GenerationResult``: ``success``, ``extras`` from the algorithm's + ``attempt_extras()``, and ``initial_state``. With ``--record_metric_inputs`` and ``--result_file``, + each completed attempt is also written as one line of ``..attempts.jsonl``. * **Per run, in memory** — the ``stats`` counters: ``num_success``, ``num_failures``, ``num_attempts`` (``num_success + num_failures = num_attempts``). * **Per run, on disk** — with ``--result_file``, a JSON sidecar is written atomically at the @@ -168,11 +168,12 @@ Four layers report what happened: applied — the generated dataset itself records only the base env id, so the sidecar is what documents that the scene was modified. -Every exported dataset, whether from generation, annotation or retargeting, also carries an ``autodata`` -entry in its ``env_args`` attribute. The entry holds the ``run_id`` and the number of exported -episodes, the embodiment's pose-action layout, the robot's joint names with their simulator limits, -and the embodiment's optional datasheet ``hardware_limits``. For a dataset generated with -``--result_file``, the ``run_id`` ties it to that run's result file and attempt log. +The ``run`` block is written only with ``--record_metric_inputs``. With that option, every exported +dataset, whether from generation, annotation or retargeting, also carries an ``autodata`` entry in +its ``env_args`` attribute. The entry holds the ``run_id`` and the number of exported episodes, the +embodiment's pose-action layout, the robot's joint names with their simulator limits, and the +embodiment's optional datasheet ``hardware_limits``. For a dataset generated with ``--result_file``, +the ``run_id`` ties it to that run's attempt log and, once the run completes, its result file. The GenerationAlgorithm Plug-In ------------------------------- @@ -203,7 +204,8 @@ The shipped algorithms — ``mimicgen``, ``dexmimicgen``, ``skillgen`` — are d :doc:`algorithms`. An algorithm can also override ``attempt_extras()`` to report facts about each attempt in its -attempt-log line; SkillGen reports its motion-planning time as ``planning_time_s``. +attempt-log line. SkillGen reports its motion-planning time as ``planning_time_s`` and whether a +motion plan failed as ``planning_failed``. The behavioral hook: ``plan_subtask_trajectory`` ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/scripts/annotate_demos.py b/scripts/annotate_demos.py index 95d370f..ee5434f 100644 --- a/scripts/annotate_demos.py +++ b/scripts/annotate_demos.py @@ -88,6 +88,11 @@ action="store_true", help="Render and record the task's image observations.", ) +parser.add_argument( + "--record_metric_inputs", + action="store_true", + help="Record what the dataset metrics read: per-step task success and run metadata in the dataset. Off by default.", +) AppLauncher.add_app_launcher_args(parser) parser.add_argument("--external_callback", type=callable_path, help="Callable registering environments after startup.") args_cli, environment_args = parser.parse_known_args() @@ -248,6 +253,7 @@ def main() -> int: generation_policy_params=generation_policy, recorder_cfg=DatagenInfoRecorderManagerCfg(), enable_cameras=args_cli.record_images, + record_step_success=args_cli.record_metric_inputs, ) # Only export episodes we explicitly mark successful (i.e. fully annotated). env_cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_SUCCEEDED_ONLY @@ -271,7 +277,8 @@ def main() -> int: uses_start_signals=annotate_start_signals, ) PreStepDatagenInfoRecorder.datastream = _datastream - autodata_entry = collect_autodata_entry(env, embodiment_adapter, uuid.uuid4().hex) + if args_cli.record_metric_inputs: + autodata_entry = collect_autodata_entry(env, embodiment_adapter, uuid.uuid4().hex) env.reset() @@ -328,12 +335,13 @@ def main() -> int: f"\nExported {exported_episode_count} (out of {processed_episode_count}) annotated" f" episode{'s' if exported_episode_count != 1 else ''} to {args_cli.output_file}." ) - autodata_entry["annotation"] = { - "input_episodes": len(episode_names), - "processed_episodes": processed_episode_count, - "annotated_episodes": exported_episode_count, - "time_s": time.monotonic() - annotation_start, - } + if autodata_entry is not None: + autodata_entry["annotation"] = { + "input_episodes": len(episode_names), + "processed_episodes": processed_episode_count, + "annotated_episodes": exported_episode_count, + "time_s": time.monotonic() - annotation_start, + } print("Exiting the app.") return exported_episode_count finally: diff --git a/scripts/generate_dataset.py b/scripts/generate_dataset.py index 19791a5..81ebb62 100755 --- a/scripts/generate_dataset.py +++ b/scripts/generate_dataset.py @@ -101,8 +101,16 @@ type=str, default=None, help=( - "Optional JSON path for the final structured generation result. An attempt log with one JSON " - "line per attempt is written next to it (..attempts.jsonl)." + "Optional JSON path for the final structured generation result. With --record_metric_inputs, an attempt " + "log with one JSON line per attempt is written next to it (..attempts.jsonl)." + ), +) +parser.add_argument( + "--record_metric_inputs", + action="store_true", + help=( + "Record what the dataset metrics read: per-step end-effector and target poses and task success, run " + "metadata in the dataset and, with --result_file, a run block and an attempt log. Off by default." ), ) parser.add_argument( @@ -221,7 +229,6 @@ def _attempt_record(env_id: int, result: GenerationResult, start_time: float, en return { "env_id": env_id, "success": result.success, - "failure_cause": result.failure_cause, "start_time_s": start_time, "end_time_s": end_time, "initial_object_positions": { @@ -393,7 +400,12 @@ def main() -> None: num_envs=args_cli.num_envs, device=args_cli.device, generation_policy_params=generation_policy_params, - recorder_cfg=load_callable(args_cli.recorder_cfg)() if args_cli.recorder_cfg else _generation_recorder_cfg(), + recorder_cfg=( + load_callable(args_cli.recorder_cfg)() + if args_cli.recorder_cfg + else _generation_recorder_cfg() if args_cli.record_metric_inputs else None + ), + record_step_success=args_cli.record_metric_inputs, env_profile=env_profile, enable_cameras=args_cli.record_images, ) @@ -415,7 +427,7 @@ def main() -> None: ) PreStepDatagenInfoRecorder.datastream = datastream run_id = uuid.uuid4().hex - autodata_entry = collect_autodata_entry(env, embodiment_adapter, run_id) + autodata_entry = collect_autodata_entry(env, embodiment_adapter, run_id) if args_cli.record_metric_inputs else None # SkillGen needs one curobo planner per env. Mimic/DexMimic take no kwargs. motion_planners: dict | None = None @@ -431,7 +443,7 @@ def main() -> None: alg_kwargs["motion_planners"] = motion_planners algorithm = get_algorithm(args_cli.alg, **alg_kwargs) - attempts_file = attempts_file_for(args_cli.result_file, run_id) if args_cli.result_file else None + attempts_file = attempts_file_for(args_cli.result_file, run_id) if autodata_entry and args_cli.result_file else None attempt_log = AttemptLog(attempts_file, run_id) if attempts_file else None try: async_components = setup_async_generation( @@ -480,28 +492,35 @@ def main() -> None: if env_profile else None ), - run_metadata={ - "run_id": run_id, - "seed": generation_policy_params.seed, - "revision": code_revision(), - "inputs": { - "source": input_record(args_cli.input_file), - "embodiment": input_record(args_cli.embodiment), - "task_descriptor": input_record(args_cli.task_descriptor), - "env_profile": input_record(args_cli.env_profile), - }, - "device": args_cli.device, - **_hardware(args_cli.device), - "num_envs": args_cli.num_envs, - "setup_time_s": generation_start - PROCESS_START, - "generation_time_s": generation_end - generation_start, - "attempts_file": Path(attempts_file).name, - }, + run_metadata=( + None + if autodata_entry is None + else { + "run_id": run_id, + "seed": generation_policy_params.seed, + "revision": code_revision(), + "inputs": { + "source": input_record(args_cli.input_file), + "embodiment": input_record(args_cli.embodiment), + "task_descriptor": input_record(args_cli.task_descriptor), + "env_profile": input_record(args_cli.env_profile), + }, + "device": args_cli.device, + **_hardware(args_cli.device), + "num_envs": args_cli.num_envs, + "setup_time_s": generation_start - PROCESS_START, + "generation_time_s": generation_end - generation_start, + "attempts_file": Path(attempts_file).name, + } + ), ) finally: env.close() - for dataset_file in exported_dataset_files(output_dir, output_file_name, generation_policy_params.keep_failed): - merge_autodata_entry(dataset_file, autodata_entry) + if autodata_entry is not None: + for dataset_file in exported_dataset_files( + output_dir, output_file_name, generation_policy_params.keep_failed + ): + merge_autodata_entry(dataset_file, autodata_entry) if __name__ == "__main__": diff --git a/scripts/retarget_dataset.py b/scripts/retarget_dataset.py index d2906ad..af10815 100755 --- a/scripts/retarget_dataset.py +++ b/scripts/retarget_dataset.py @@ -138,6 +138,11 @@ "expected to be low." ), ) +parser.add_argument( + "--record_metric_inputs", + action="store_true", + help="Record what the dataset metrics read: per-step task success and run metadata in the dataset. Off by default.", +) AppLauncher.add_app_launcher_args(parser) args_cli = parser.parse_args() diff --git a/skills/autodata-generate-datasets/references/cli-reference.md b/skills/autodata-generate-datasets/references/cli-reference.md index 9062e4a..39d00a8 100644 --- a/skills/autodata-generate-datasets/references/cli-reference.md +++ b/skills/autodata-generate-datasets/references/cli-reference.md @@ -25,7 +25,8 @@ YAML, live env, and HDF5 source dataset, then runs the selected generation algor | `--output_file ` | no | `./datasets/output_dataset.hdf5` | Destination for generated demos. | | `--generation_num_trials ` | no | `None` | Number of demos to generate. | | `--num_envs ` | no | `1` | Number of parallel simulation environments. | -| `--result_file ` | no | — | Where to write the generation result summary. An attempt log (`..attempts.jsonl`) is written next to it. | +| `--result_file ` | no | — | Where to write the generation result summary. With `--record_metric_inputs`, an attempt log (`..attempts.jsonl`) is written next to it. | +| `--record_metric_inputs` | no | off | Record what the dataset metrics read: per-step poses and task success, run metadata, and the attempt log. | | `--pause_subtask` | no | off | Pause at subtask boundaries (debugging). | | `--visualize_plan` | no | off | Visualize the motion plan (SkillGen debugging). | From 9ea6a76ae42092d56bccacdb6648a3e4b796816b Mon Sep 17 00:00:00 2001 From: Billy Okal Date: Wed, 28 Oct 2026 16:37:03 -0500 Subject: [PATCH 6/6] fix: record metric inputs with a custom recorder and without a success term - With --record_metric_inputs, the per-step pose term is added to a custom --recorder_cfg too, unless it already records datagen info. - Without a success term, step_success is not recorded instead of failing on the first step. --- .../env/datagen_info_recorder.py | 22 +++++++++++++-- .../env/test_datagen_info_recorder.py | 28 +++++++++++++++++++ scripts/generate_dataset.py | 24 ++++++---------- 3 files changed, 55 insertions(+), 19 deletions(-) create mode 100644 autodata_tests/interfaces/env/test_datagen_info_recorder.py diff --git a/autodata_interfaces/env/datagen_info_recorder.py b/autodata_interfaces/env/datagen_info_recorder.py index 15bb80b..81d5fa6 100644 --- a/autodata_interfaces/env/datagen_info_recorder.py +++ b/autodata_interfaces/env/datagen_info_recorder.py @@ -69,11 +69,27 @@ class DatagenInfoRecorderManagerCfg(ActionStateRecorderManagerCfg): record_pre_step_datagen_info = PreStepDatagenInfoRecorderCfg() -def add_step_success_recorder(recorder_cfg: object, success_term: TerminationTermCfg) -> None: +def add_step_success_recorder(recorder_cfg: object, success_term: TerminationTermCfg | None) -> None: """Add the ``step_success`` term to any recorder manager configuration. + Nothing is added without a success term, since there is then no success to record. + + Args: + recorder_cfg: Recorder manager configuration the env will be built with. + success_term: The task's success term, or None when the task has none. + """ + if success_term is not None: + recorder_cfg.record_post_step_success = PostStepSuccessRecorderCfg(success_term=success_term) + + +def add_pose_recorder(recorder_cfg: object) -> None: + """Add the per-step end-effector and commanded target pose term to any recorder manager configuration. + + Object poses are left out, as no metric reads them. A configuration that already has the + ``record_pre_step_datagen_info`` term, such as :class:`DatagenInfoRecorderManagerCfg`, keeps its own. + Args: recorder_cfg: Recorder manager configuration the env will be built with. - success_term: The task's success term. """ - recorder_cfg.record_post_step_success = PostStepSuccessRecorderCfg(success_term=success_term) + if not hasattr(recorder_cfg, "record_pre_step_datagen_info"): + recorder_cfg.record_pre_step_datagen_info = PreStepDatagenInfoRecorderCfg(record_object_poses=False) diff --git a/autodata_tests/interfaces/env/test_datagen_info_recorder.py b/autodata_tests/interfaces/env/test_datagen_info_recorder.py new file mode 100644 index 0000000..2929133 --- /dev/null +++ b/autodata_tests/interfaces/env/test_datagen_info_recorder.py @@ -0,0 +1,28 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg + +from autodata_interfaces.env.datagen_info_recorder import ( + DatagenInfoRecorderManagerCfg, + add_pose_recorder, + add_step_success_recorder, +) + + +def test_metric_recorders_extend_a_custom_configuration(): + custom = ActionStateRecorderManagerCfg() + + add_pose_recorder(custom) + add_step_success_recorder(custom, success_term=None) + + assert custom.record_pre_step_datagen_info.record_object_poses is False + assert not hasattr(custom, "record_post_step_success") + + +def test_pose_recorder_keeps_an_existing_datagen_info_term(): + annotation = DatagenInfoRecorderManagerCfg() + + add_pose_recorder(annotation) + + assert annotation.record_pre_step_datagen_info.record_object_poses is True diff --git a/scripts/generate_dataset.py b/scripts/generate_dataset.py index 81ebb62..2aa3604 100755 --- a/scripts/generate_dataset.py +++ b/scripts/generate_dataset.py @@ -160,6 +160,8 @@ from pathlib import Path # noqa: E402 from typing import Any # noqa: E402 +from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg # noqa: E402 + from autodata_core import DataGenerator, get_algorithm # noqa: E402 from autodata_core.algorithms import REGISTERED_ALGORITHMS # noqa: E402 from autodata_core.data_generator import GenerationResult # noqa: E402 @@ -172,10 +174,7 @@ setup_env_config, setup_output_paths, ) -from autodata_interfaces.env.datagen_info_recorder import ( # noqa: E402 - DatagenInfoRecorderManagerCfg, - PreStepDatagenInfoRecorder, -) +from autodata_interfaces.env.datagen_info_recorder import PreStepDatagenInfoRecorder, add_pose_recorder # noqa: E402 from autodata_interfaces.env.registration import load_callable, register_external_environment # noqa: E402 from autodata_utils.dataset_metadata import exported_dataset_files, merge_autodata_entry # noqa: E402 from autodata_utils.generation_result import ( # noqa: E402 @@ -253,13 +252,6 @@ def _hardware(device: str) -> dict[str, str | None]: return {"gpu": gpu, "num_gpus": 1 if gpu else 0, "cpu": cpu or platform.processor() or None} -def _generation_recorder_cfg() -> DatagenInfoRecorderManagerCfg: - """Actions and states, plus the per-step poses the dataset metrics read; object poses are left out.""" - recorder_cfg = DatagenInfoRecorderManagerCfg() - recorder_cfg.record_pre_step_datagen_info.record_object_poses = False - return recorder_cfg - - def build_datastream( env: Any, input_file: str, @@ -393,6 +385,10 @@ def main() -> None: # names the motion-planner profile tuned for the resulting scene. env_profile = EnvironmentProfile.from_yaml(args_cli.env_profile) if args_cli.env_profile else None + recorder_cfg = load_callable(args_cli.recorder_cfg)() if args_cli.recorder_cfg else None + if args_cli.record_metric_inputs: + recorder_cfg = recorder_cfg or ActionStateRecorderManagerCfg() + add_pose_recorder(recorder_cfg) env_cfg, success_term = setup_env_config( env_name=env_name, output_dir=output_dir, @@ -400,11 +396,7 @@ def main() -> None: num_envs=args_cli.num_envs, device=args_cli.device, generation_policy_params=generation_policy_params, - recorder_cfg=( - load_callable(args_cli.recorder_cfg)() - if args_cli.recorder_cfg - else _generation_recorder_cfg() if args_cli.record_metric_inputs else None - ), + recorder_cfg=recorder_cfg, record_step_success=args_cli.record_metric_inputs, env_profile=env_profile, enable_cameras=args_cli.record_images,