diff --git a/autodata_core/algorithms.py b/autodata_core/algorithms.py index 65d6daf..e0df4f8 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,14 @@ 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] 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, *, @@ -238,6 +255,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,7 +263,9 @@ 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: + 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 59481cb..3f20fa5 100644 --- a/autodata_core/data_generator.py +++ b/autodata_core/data_generator.py @@ -47,10 +47,18 @@ 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. + extras: Algorithm-specific facts about the attempt, from + :meth:`~autodata_core.algorithms.GenerationAlgorithm.attempt_extras`. + """ initial_state: dict success: bool + extras: dict[str, Any] = field(default_factory=dict) @dataclass @@ -76,6 +84,8 @@ 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 + planning_failed: bool = False class DataGenerator: @@ -569,7 +579,11 @@ 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, + 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 +634,7 @@ async def generate( return GenerationResult( initial_state=initial_state, success=buffers.success, + 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..08016f4 100644 --- a/autodata_interfaces/embodiments/absolute_pose_embodiment_adapter.py +++ b/autodata_interfaces/embodiments/absolute_pose_embodiment_adapter.py @@ -77,6 +77,18 @@ 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]: + """Each end effector with a ``pose_action_slice`` commands its target pose there: position, then quaternion. + + These entries are not clipped. + """ + pose_action = { + 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 + } + 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..e948e3b 100644 --- a/autodata_interfaces/embodiments/bimanual_embodiment_adapter.py +++ b/autodata_interfaces/embodiments/bimanual_embodiment_adapter.py @@ -191,6 +191,17 @@ 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]: + """``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: {"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} + 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..d0b9859 100644 --- a/autodata_interfaces/embodiments/embodiment_adapter.py +++ b/autodata_interfaces/embodiments/embodiment_adapter.py @@ -50,6 +50,17 @@ class EmbodimentAdapter(ABC): robot_asset_name: str = "robot" raw_config: dict[str, Any] | None = None + def metric_context(self) -> dict[str, Any]: + """Describe this embodiment for the dataset metrics; the description is saved with every dataset it records. + + 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 {} + 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..6cd8df0 100644 --- a/autodata_interfaces/embodiments/single_arm_embodiment_adapter.py +++ b/autodata_interfaces/embodiments/single_arm_embodiment_adapter.py @@ -133,6 +133,15 @@ class DeltaPoseIKSingleArmAdapter(SingleArmEmbodimentAdapter): clip_pose_action_to_unit: bool = True + def metric_context(self) -> dict[str, Any]: + """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: {"action_slice": [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..81d5fa6 --- /dev/null +++ b/autodata_interfaces/env/datagen_info_recorder.py @@ -0,0 +1,95 @@ +# 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) -> 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. + """ + 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_interfaces/env/isaaclab_env_interface.py b/autodata_interfaces/env/isaaclab_env_interface.py index 3a8f916..9e4ebf6 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]: @@ -166,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. @@ -186,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. @@ -231,6 +268,8 @@ def setup_env_config( env_cfg.recorders = ActionStateRecorderManagerCfg() else: env_cfg.recorders = recorder_cfg + 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 f1e9ddb..7621d1c 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 @@ -96,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 @@ -117,8 +120,11 @@ 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) + 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). @@ -180,3 +186,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, generation_policy.keep_failed): + merge_autodata_entry(dataset_file, autodata_entry) 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/e2e/test_annotation.py b/autodata_tests/e2e/test_annotation.py index f5383d3..fa93c2e 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 @@ -36,8 +36,11 @@ 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)) + 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_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 cfb02b9..0552279 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,9 @@ 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, + "--record_metric_inputs", "--generation_num_trials", str(GENERATION_NUM_TRIALS), "--num_envs", @@ -47,6 +51,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..580d086 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 @@ -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", @@ -62,6 +63,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/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/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 f9cbde4..4792836 100644 --- a/autodata_tests/utils/utils.py +++ b/autodata_tests/utils/utils.py @@ -41,3 +41,58 @@ 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." + 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 + + +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/autodata_utils/dataset_metadata.py b/autodata_utils/dataset_metadata.py new file mode 100644 index 0000000..35ac6e9 --- /dev/null +++ b/autodata_utils/dataset_metadata.py @@ -0,0 +1,48 @@ +# 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, 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)] + + +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..fc5bc4e 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, 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: + """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/docs/pages/concepts/data_generator.rst b/docs/pages/concepts/data_generator.rst index b182383..244451b 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** — 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 @@ -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.3f2a9c.attempts.jsonl" } } @@ -148,6 +168,13 @@ 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. +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 ------------------------------- @@ -176,6 +203,10 @@ 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`` 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 b0d6b2b..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() @@ -104,23 +109,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 +166,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 +251,15 @@ 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, + 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 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 +276,9 @@ def main() -> int: source_pool=empty_pool, uses_start_signals=annotate_start_signals, ) + PreStepDatagenInfoRecorder.datastream = _datastream + if args_cli.record_metric_inputs: + autodata_entry = collect_autodata_entry(env, embodiment_adapter, uuid.uuid4().hex) env.reset() @@ -311,6 +303,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 +335,20 @@ 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}." ) + 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: env.close() + if autodata_entry is not None: + for dataset_file in exported_dataset_files(output_dir, output_file_name, kept_failed=False): + merge_autodata_entry(dataset_file, autodata_entry) def replay_episode( diff --git a/scripts/generate_dataset.py b/scripts/generate_dataset.py index d6f5167..2aa3604 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,18 @@ "--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. 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( "--pause_subtask", @@ -137,24 +151,40 @@ 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 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 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 PreStepDatagenInfoRecorder, add_pose_recorder # noqa: E402 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 +196,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 +219,37 @@ 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, + "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 build_datastream( @@ -216,6 +279,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 +302,7 @@ def setup_async_generation( success_term=success_term, pause_subtask=pause_subtask, stats=stats, + attempt_log=attempt_log, ) ) tasks.append(task) @@ -320,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, @@ -327,7 +396,8 @@ 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=recorder_cfg, + record_step_success=args_cli.record_metric_inputs, env_profile=env_profile, enable_cameras=args_cli.record_images, ) @@ -347,6 +417,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) 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 @@ -362,6 +435,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, 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( datastream=datastream, @@ -369,10 +444,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 +460,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 +484,35 @@ def main() -> None: if env_profile else None ), + 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() + 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 19a5e3b..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. | +| `--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). |