From 1f95ba0ef3abdd14bf5a05b97ffc2e021b472ef0 Mon Sep 17 00:00:00 2001 From: Billy Okal Date: Fri, 16 Oct 2026 12:09:02 -0500 Subject: [PATCH 1/4] feat: report dataset KPIs from generation records Generation now writes a run block in the result JSON, an attempt log next to it, and per-step intended and achieved end-effector poses. Channels whose action commands no pose record no intended pose. The cuRobo planner keeps a summary of each planning call, and annotate_demos.py can write a summary file. validate_dataset.py --report computes the KPIs from these files. --- autodata_core/algorithms.py | 8 +- autodata_core/data_generator.py | 78 +- autodata_core/pool.py | 5 + autodata_interfaces/datastream/datastream.py | 54 +- .../absolute_pose_embodiment_adapter.py | 3 + .../embodiments/embodiment_adapter.py | 8 + .../motion_planners/curobo/curobo_planner.py | 54 +- autodata_tests/core/test_data_generator.py | 68 +- autodata_tests/e2e/test_annotation.py | 6 + .../e2e/test_mimicgen_data_generation.py | 61 +- .../e2e/test_skillgen_bin_data_generation.py | 28 +- .../interfaces/datastream/test_datastream.py | 14 + .../test_absolute_pose_embodiment_adapter.py | 4 + autodata_tests/utils/test_dataset_metrics.py | 118 +++ autodata_tests/utils/test_validate_dataset.py | 31 + autodata_tests/utils/utils.py | 44 ++ autodata_utils/dataset_metrics.py | 711 ++++++++++++++++++ autodata_utils/generation_result.py | 33 + docs/index.rst | 1 + docs/pages/advanced/motion_planners.rst | 5 +- docs/pages/concepts/data_generator.rst | 9 +- docs/pages/references/kpis.rst | 142 ++++ scripts/annotate_demos.py | 38 +- scripts/generate_dataset.py | 115 ++- scripts/validate_dataset.py | 36 +- .../references/cli-reference.md | 11 +- 26 files changed, 1643 insertions(+), 42 deletions(-) create mode 100644 autodata_tests/utils/test_dataset_metrics.py create mode 100644 autodata_utils/dataset_metrics.py create mode 100644 docs/pages/references/kpis.rst diff --git a/autodata_core/algorithms.py b/autodata_core/algorithms.py index 65d6daf..bd36256 100644 --- a/autodata_core/algorithms.py +++ b/autodata_core/algorithms.py @@ -185,7 +185,8 @@ def __init__(self, motion_planners: dict[int, Any]) -> None: motion_planners: One planner per env_id. Planners must expose ``update_world_and_plan_motion(target_pose, expected_attached_object, env_id, ...)`` and ``get_planned_poses() -> list[torch.Tensor[4, 4]]``. ``config.motion_noise_scale`` - is read if present. + is read if present, and so is ``last_plan_summary``, a dict describing the latest + planning call, which is copied into the attempt's ``planner_calls``. """ assert motion_planners, "SkillGen requires at least one motion planner" self.motion_planners = motion_planners @@ -245,6 +246,11 @@ def plan_subtask_trajectory( step_size=getattr(planner, "step_size", None), enable_retiming=(hasattr(planner, "step_size") and planner.step_size is not None), ) + plan_summary = getattr(planner, "last_plan_summary", None) + if plan_summary is not None: + eef_state.planner_calls.append( + {"eef": eef_name, "subtask": eef_state.current_subtask_index, **plan_summary} + ) if not planning_success: return None diff --git a/autodata_core/data_generator.py b/autodata_core/data_generator.py index 59481cb..4b3c77b 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,26 @@ 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 planning aborted the attempt, ``"execution"`` when every + subtask ran but the task never succeeded, None on success. + failed_subtask: Per-EEF index of the subtask where the attempt failed. For a planner failure, + the subtask being planned; for an execution failure, the first subtask whose termination + signal never fired (the last subtask if all fired). None when unknown. + num_steps: Number of env steps executed. + planner_calls: One summary per motion-planner call made during the attempt. + """ initial_state: dict success: bool + failure_cause: Literal["planner", "execution"] | None = None + failed_subtask: dict[str, int] | None = None + num_steps: int = 0 + planner_calls: tuple[dict, ...] = () @dataclass @@ -76,6 +92,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 + planner_calls: list[dict] = field(default_factory=list) class DataGenerator: @@ -107,6 +124,8 @@ def __init__( self.algorithm = algorithm self.src_demo_datagen_info_pool = datastream.source_pool self.demo_keys = demo_keys + # Resolved on the first attempt, once the env has produced observations. + self._track_term_signals: bool | None = None self._validate_eef_count_for_algorithm() self._validate_coordination_for_algorithm() @@ -537,6 +556,10 @@ async def generate( eef_states = self._initialize_eef_states() selected_src_demo_inds: dict[str, int | None] = {name: None for name in self.datastream.get_eef_names()} buffers = _GenerationBuffers() + if self._track_term_signals is None: + self._track_term_signals = self.datastream.has_subtask_term_signals() + fired_term_signals: set[str] = set() + num_steps = 0 randomized_subtask_boundaries: dict[str, np.ndarray] | None = None prev_pool_size = 0 @@ -569,7 +592,14 @@ 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", + failed_subtask={eef_name: eef_state.current_subtask_index}, + num_steps=num_steps, + planner_calls=self._collect_planner_calls(eef_states), + ) next_trajectory, _is_motion_plan = result eef_state.current_trajectory = next_trajectory eef_state.subtask_step_index = 0 @@ -579,6 +609,7 @@ async def generate( eef_states=eef_states, ) multi_waypoint = MultiWaypoint(eef_waypoints) + self._record_step_poses(env_id, eef_waypoints) await asyncio.sleep(0) @@ -588,6 +619,11 @@ async def generate( env_id=env_id, env_action_queue=env_action_queue, ) + num_steps += 1 + if self._track_term_signals: + signals = self.datastream.get_subtask_term_signals(env_ids=[env_id]) + fired_now = torch.stack([value[0] for value in signals.values()]).tolist() + fired_term_signals.update(name for name, fired in zip(signals, fired_now) if fired) buffers.success = buffers.success or exec_success if generation_policy.use_navigation_controller: hold, was_navigating = self._synchronize_navigation(eef_states, env_id, was_navigating) @@ -617,9 +653,14 @@ async def generate( if export_demo: env.recorder_manager.export_episodes(env_id_tensor) + failed = not buffers.success return GenerationResult( initial_state=initial_state, success=buffers.success, + failure_cause="execution" if failed else None, + failed_subtask=self._failed_subtasks(fired_term_signals) if failed and self._track_term_signals else None, + num_steps=num_steps, + planner_calls=self._collect_planner_calls(eef_states), ) # ------------------------------------------------------------------ @@ -675,6 +716,37 @@ def _build_runtime_subtask_constraints(self) -> dict: def _initialize_eef_states(self) -> dict[str, _EEFGenerationState]: return {name: _EEFGenerationState() for name in self.datastream.get_eef_names()} + def _record_step_poses(self, env_id: int, eef_waypoints: dict[str, Waypoint]) -> None: + """Append the noise-free waypoint pose commanded at this step, and the EEF pose before it, to the episode. + + Tracking-only channels command no pose, so they record no intended pose. + """ + commanded = self.datastream.get_eef_names_with_pose_action() + intended = {name: waypoint.pose.unsqueeze(0) for name, waypoint in eef_waypoints.items() if name in commanded} + achieved = {name: self.datastream.get_robot_eef_pose(env_ids=[env_id], eef_name=name) for name in eef_waypoints} + self.datastream.get_env().recorder_manager.add_to_episodes( + "obs/datagen_info", {"intended_eef_pose": intended, "eef_pose": achieved}, env_ids=[env_id] + ) + + def _failed_subtasks(self, fired_term_signals: set[str]) -> dict[str, int]: + """Per EEF, the first subtask whose termination signal never fired, else the last subtask.""" + failed: dict[str, int] = {} + for eef_name in self.datastream.get_eef_names(): + subtasks = self.datastream.get_subtasks(eef_name) + failed[eef_name] = next( + ( + index + for index, subtask in enumerate(subtasks) + if subtask.subtask_term_signal and subtask.subtask_term_signal not in fired_term_signals + ), + len(subtasks) - 1, + ) + return failed + + @staticmethod + def _collect_planner_calls(eef_states: dict[str, _EEFGenerationState]) -> tuple[dict, ...]: + return tuple(call for eef_state in eef_states.values() for call in eef_state.planner_calls) + def _maybe_refresh_randomized_subtask_boundaries( self, randomized_subtask_boundaries: dict[str, np.ndarray] | None, diff --git a/autodata_core/pool.py b/autodata_core/pool.py index 102b8a2..0fca268 100644 --- a/autodata_core/pool.py +++ b/autodata_core/pool.py @@ -122,6 +122,11 @@ def _add_episode(self, episode: EpisodeData) -> None: if "datagen_info" not in ep_grp["obs"]: raise ValueError("Episode lacks 'datagen_info' obs annotations") + missing = {"eef_pose", "object_pose", "target_eef_pose", "subtask_term_signals"} - set( + ep_grp["obs"]["datagen_info"] + ) + if missing: + raise ValueError(f"Episode 'datagen_info' lacks {sorted(missing)}; pass an annotated source dataset") eef_pose = ep_grp["obs"]["datagen_info"]["eef_pose"] object_poses_dict = ep_grp["obs"]["datagen_info"]["object_pose"] diff --git a/autodata_interfaces/datastream/datastream.py b/autodata_interfaces/datastream/datastream.py index 9733fa7..497f24f 100644 --- a/autodata_interfaces/datastream/datastream.py +++ b/autodata_interfaces/datastream/datastream.py @@ -191,6 +191,11 @@ def get_robot_eef_pose(self, env_ids: Sequence[int] | None, eef_name: str) -> to return self.embodiment_adapter.get_eef_poses(env_ids=env_ids)[eef_name] + def get_eef_names_with_pose_action(self) -> tuple[str, ...]: + """Return the EEFs whose target pose the action commands, excluding tracking-only channels.""" + + return self.embodiment_adapter.get_eef_names_with_pose_action() + def get_robot_joint_positions(self, env_ids: Sequence[int] | None = None) -> torch.Tensor: """Read the robot's current joint positions [rad] via the embodiment adapter. @@ -271,16 +276,38 @@ def get_object_poses(self, env_ids: Sequence[int] | None = None) -> dict[str, to object_pose_matrix[obj_name] = reference_inverse @ object_pose_matrix[obj_name] return object_pose_matrix + def _boundary_signal_names(self) -> list[str]: + """Termination-signal names of every non-final subtask, across EEFs. + + The final subtask ends with the trajectory, so its signal never marks a boundary. + """ + + return [name for eef in self.get_eef_names() for name in self.get_term_signal_names(eef)[:-1] if name] + + def has_subtask_term_signals(self, obs_group: str = "subtask_terms") -> bool: + """Whether :meth:`get_subtask_term_signals` can read every subtask-boundary signal. + + True when ``env.obs_buf[obs_group]`` is a non-concatenated group holding one term per + boundary signal. Not every env publishes such a group. + """ + + obs_buf = getattr(self.env, "obs_buf", None) + group = obs_buf.get(obs_group) if isinstance(obs_buf, dict) else None + if not isinstance(group, dict): + return False + signal_names = self._boundary_signal_names() + return bool(signal_names) and all(name in group for name in signal_names) + def get_subtask_term_signals( self, env_ids: Sequence[int] | None = None, obs_group: str = "subtask_terms" ) -> dict[str, torch.Tensor]: - """Read the current per-subtask termination signals from the env observation buffer. + """Read the current subtask-boundary signals from the env observation buffer. - For every non-empty termination-signal name declared by the task descriptor (across all - EEFs), reads the identically named boolean observation term from - ``env.obs_buf[obs_group]``. The env must therefore publish one observation term per - term-signal name in a non-concatenated group. Auto annotation samples this each replay - step to locate subtask boundaries without a human in the loop. + For the termination signal of every non-final subtask (across all EEFs), reads the + identically named boolean observation term from ``env.obs_buf[obs_group]``. The env must + therefore publish one observation term per boundary signal in a non-concatenated group. + Auto annotation samples this each replay step to locate subtask boundaries without a human + in the loop. Args: env_ids: Environments to read; ``None`` reads all. @@ -302,15 +329,12 @@ def get_subtask_term_signals( "signal reads require the group to set 'concatenate_terms = False'." ) signals: dict[str, torch.Tensor] = {} - for eef_name in self.get_eef_names(): - for signal_name in self.get_term_signal_names(eef_name): - if not signal_name: - continue - assert signal_name in group, ( - f"Termination signal {signal_name!r} has no matching observation term in " - f"group {obs_group!r}. Available terms: {sorted(group.keys())}" - ) - signals[signal_name] = as_torch(group[signal_name])[index].reshape(-1).bool() + for signal_name in self._boundary_signal_names(): + assert signal_name in group, ( + f"Termination signal {signal_name!r} has no matching observation term in " + f"group {obs_group!r}. Available terms: {sorted(group.keys())}" + ) + signals[signal_name] = as_torch(group[signal_name])[index].reshape(-1).bool() return signals def get_robot_root_pose(self, env_ids: Sequence[int] | None = None) -> torch.Tensor: diff --git a/autodata_interfaces/embodiments/absolute_pose_embodiment_adapter.py b/autodata_interfaces/embodiments/absolute_pose_embodiment_adapter.py index c1d6c09..3e1409e 100644 --- a/autodata_interfaces/embodiments/absolute_pose_embodiment_adapter.py +++ b/autodata_interfaces/embodiments/absolute_pose_embodiment_adapter.py @@ -80,6 +80,9 @@ def from_dict(cls, data: dict[str, Any]) -> AbsolutePoseEmbodimentAdapter: def get_eef_names(self) -> tuple[str, ...]: return tuple(self.eefs) + def get_eef_names_with_pose_action(self) -> tuple[str, ...]: + return tuple(name for name, spec in self.eefs.items() if spec.get("pose_action_slice") is not None) + def get_eef_poses(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]: assert self.env is not None, "Call bind_env(env) before reading state." index = slice(None) if env_ids is None else env_ids diff --git a/autodata_interfaces/embodiments/embodiment_adapter.py b/autodata_interfaces/embodiments/embodiment_adapter.py index 4f63034..1694fd5 100644 --- a/autodata_interfaces/embodiments/embodiment_adapter.py +++ b/autodata_interfaces/embodiments/embodiment_adapter.py @@ -93,6 +93,14 @@ def get_eef_names(self) -> tuple[str, ...]: Tuple of EEF names. Stable across calls; never empty. """ + def get_eef_names_with_pose_action(self) -> tuple[str, ...]: + """Return the end-effectors whose target pose the action commands. + + Tracking-only channels, whose actions carry only passthrough commands, are excluded. + """ + + return self.get_eef_names() + @abstractmethod def get_eef_poses(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]: """Read current end-effector poses [m, m, m] from the live env. diff --git a/autodata_interfaces/motion_planners/curobo/curobo_planner.py b/autodata_interfaces/motion_planners/curobo/curobo_planner.py index 1a43077..2ac8e1e 100644 --- a/autodata_interfaces/motion_planners/curobo/curobo_planner.py +++ b/autodata_interfaces/motion_planners/curobo/curobo_planner.py @@ -5,6 +5,7 @@ import logging import numpy as np +import time import torch from dataclasses import dataclass from types import SimpleNamespace @@ -286,6 +287,9 @@ def __init__( # Initialize planning state self._current_plan: JointState | None = None self._plan_index: int = 0 + # One record per cuRobo plan_single call of the current update_world_and_plan_motion call. + self._plan_phase_results: list[dict[str, Any]] = [] + self.last_plan_summary: dict[str, Any] | None = None # Initialize visualization state self.frame_counter: int = 0 @@ -1307,8 +1311,9 @@ def _plan_to_contact_pose( ) planning_success = False + result: Any = None try: - result: Any = self.motion_gen.plan_single(start_state, goal_pose, self.plan_config) + result = self.motion_gen.plan_single(start_state, goal_pose, self.plan_config) if result.success.item(): if result.optimized_plan is not None and len(result.optimized_plan.position) != 0: @@ -1332,6 +1337,13 @@ def _plan_to_contact_pose( except Exception as e: self.logger.debug(f"Error during planning: {e}") + # Telemetry must never change the planning outcome or skip the sphere restore below. + try: + phase_summary = self._summarize_plan_phase(result, planning_success) + except Exception as e: + self.logger.debug(f"Error summarizing the planning phase: {e}") + phase_summary = {"success": planning_success, "status": "Success" if planning_success else "exception"} + self._plan_phase_results.append(phase_summary) # Always restore sphere state after planning, regardless of success if contact: @@ -1843,7 +1855,9 @@ def update_world_and_plan_motion( Provides a high-level interface that handles the complete planning workflow: world synchronization, object attachment/detachment, gripper configuration, - and motion planning. + and motion planning. Afterwards :attr:`last_plan_summary` holds the call's outcome, the + cuRobo status of the first failed phase, wall-clock time [s], summed cuRobo solve, IK and + trajectory-optimization times [s], and summed cuRobo attempts. Args: target_pose: Target end-effector pose as 4x4 transformation matrix @@ -1855,6 +1869,8 @@ def update_world_and_plan_motion( Returns: True if complete planning pipeline succeeded, False if any step failed """ + start_time = time.perf_counter() + self._plan_phase_results = [] # Always reset the plan before starting a new one to ensure a clean state self.reset_plan() @@ -1919,8 +1935,42 @@ def update_world_and_plan_motion( self._detach_objects() + failed_phase = next((phase for phase in self._plan_phase_results if not phase["success"]), None) + self.last_plan_summary = { + "success": plan_success, + "status": "Success" if plan_success else (failed_phase["status"] if failed_phase else "unknown"), + "wall_time_s": time.perf_counter() - start_time, + "num_phases": len(self._plan_phase_results), + **{ + key: sum(phase.get(key, 0) for phase in self._plan_phase_results) + for key in ("solve_time_s", "ik_time_s", "trajopt_time_s", "attempts") + }, + } return plan_success + @staticmethod + def _summarize_plan_phase(result: Any, success: bool) -> dict[str, Any]: + """Outcome, cuRobo status, timing [s] and attempt count of one planning phase. + + ``result`` is None when ``plan_single`` raised. A phase whose cuRobo result succeeded but whose + post-processing raised is reported with status ``exception``. + """ + if success: + status = "Success" + elif result is None or (result.success is not None and bool(result.success.item())): + status = "exception" + elif result.status is None: + status = "unknown" + else: + status = str(getattr(result.status, "value", result.status)) + summary: dict[str, Any] = {"success": success, "status": status} + if result is not None: + summary["solve_time_s"] = float(result.solve_time) + summary["ik_time_s"] = float(result.ik_time) + summary["trajopt_time_s"] = float(result.trajopt_time) + summary["attempts"] = int(result.attempts) + return summary + # ===================================================================================== # UTILITY METHODS # ===================================================================================== diff --git a/autodata_tests/core/test_data_generator.py b/autodata_tests/core/test_data_generator.py index 806c0e9..d6ff769 100644 --- a/autodata_tests/core/test_data_generator.py +++ b/autodata_tests/core/test_data_generator.py @@ -3,8 +3,8 @@ """Unit tests for :class:`autodata_core.data_generator.DataGenerator`. -Construction-time validation and navigation playback decisions are covered here. -Simulation integration is covered by the e2e suite. +Construction-time validation, navigation playback decisions and the planner-failure path of the +generation loop are covered here. Simulation integration is covered by the e2e suite. """ import asyncio @@ -15,7 +15,7 @@ import pytest -from autodata_core.algorithms import DexMimicGen, MimicGen, SkillGen +from autodata_core.algorithms import DexMimicGen, GenerationAlgorithm, MimicGen, SkillGen from autodata_core.data_generator import DataGenerator, GenerationResult, _EEFGenerationState from autodata_core.waypoint import MultiWaypoint, Waypoint from autodata_interfaces.tasks.generation_policy_spec import GenerationPolicy @@ -66,6 +66,15 @@ def get_generation_policy(self): def get_env(self): return self.env + def has_subtask_term_signals(self): + return False + + def get_eef_names_with_pose_action(self): + return tuple(self._subtasks) + + def get_robot_eef_pose(self, env_ids, eef_name): + return torch.eye(4).unsqueeze(0) + def _assert_constructed(generator: DataGenerator, datastream: _MockDatastream, algorithm) -> None: assert generator.algorithm is algorithm @@ -133,6 +142,59 @@ def test_skillgen_coordination_unsupported_error(): DataGenerator(datastream=datastream, algorithm=SkillGen(motion_planners={0: object()})) +class _PlanningFails(GenerationAlgorithm): + """Unregistered algorithm whose planner never finds a trajectory.""" + + def plan_subtask_trajectory(self, **kwargs): + return None + + +class _GenerationDatastream(_MockDatastream): + """Adds the reads ``DataGenerator.generate`` makes before the first plan.""" + + def __init__(self, subtasks_by_eef: dict) -> None: + super().__init__(subtasks_by_eef) + self.source_pool.subtask_boundaries = {eef: [[(0, 5), (5, 10)]] for eef in subtasks_by_eef} + self.env = SimpleNamespace(recorder_manager=SimpleNamespace(reset=lambda env_ids: None)) + + def get_scene_state(self, is_relative: bool = True): + return {"rigid_object": {}} + + +def test_planner_failure_records_cause_and_subtask(): + datastream = _GenerationDatastream({"franka": [_subtask("grasp_1"), _subtask("")]}) + generator = DataGenerator(datastream=datastream, algorithm=_PlanningFails()) + + async def run_attempt() -> GenerationResult: + reset_queue: asyncio.Queue = asyncio.Queue() + + async def serve_reset() -> None: + await reset_queue.get() + reset_queue.task_done() + + server = asyncio.create_task(serve_reset()) + result = await generator.generate( + env_id=0, success_term=None, env_reset_queue=reset_queue, env_action_queue=asyncio.Queue() + ) + await server + return result + + result = asyncio.run(run_attempt()) + + assert result.success is False + assert result.failure_cause == "planner" + assert result.failed_subtask == {"franka": 0} + assert result.num_steps == 0 + + +def test_execution_failure_is_attributed_to_the_first_unfinished_subtask(): + datastream = _MockDatastream({"franka": [_subtask("grasp_1"), _subtask("stack_1"), _subtask("")]}) + generator = DataGenerator(datastream=datastream, algorithm=MimicGen()) + + assert generator._failed_subtasks({"grasp_1"}) == {"franka": 1} + assert generator._failed_subtasks({"grasp_1", "stack_1"}) == {"franka": 2} + + def test_generation_result_fields_and_frozen(): result = GenerationResult(initial_state={"foo": 1}, success=True) assert result.initial_state == {"foo": 1} diff --git a/autodata_tests/e2e/test_annotation.py b/autodata_tests/e2e/test_annotation.py index f5383d3..aebe37c 100644 --- a/autodata_tests/e2e/test_annotation.py +++ b/autodata_tests/e2e/test_annotation.py @@ -20,6 +20,7 @@ def test_automatic_annotation_from_dataset(tmp_path): """Use the stored task name and export usable termination signals during replay.""" input_file = Path(TestPaths.test_data_dir) / "annotated_dataset_franka_stack_mimicgen.hdf5" output_file = tmp_path / "annotated.hdf5" + summary_file = tmp_path / "annotation_summary.json" run_subprocess([ TestPaths.python_path, str(Path(TestPaths.scripts_dir) / "annotate_demos.py"), @@ -36,10 +37,15 @@ def test_automatic_annotation_from_dataset(tmp_path): str(input_file), "--output_file", str(output_file), + "--summary_file", + str(summary_file), ]) assert_valid_dataset(str(output_file), min_num_demos=1) + summary = json.loads(summary_file.read_text(encoding="utf-8")) with h5py.File(input_file, "r") as source, h5py.File(output_file, "r") as annotated: assert len(annotated["data"]) == len(source["data"]) + assert summary["num_processed"] == summary["num_exported"] == len(source["data"]) + assert summary["rejections"] == {} assert json.loads(annotated["data"].attrs["env_args"])["env_name"] == "IsaacContrib-Stack-Cube-Franka-IK-Rel" for episode in annotated["data"].values(): signals = episode["obs/datagen_info/subtask_term_signals"] diff --git a/autodata_tests/e2e/test_mimicgen_data_generation.py b/autodata_tests/e2e/test_mimicgen_data_generation.py index cfb02b9..aa493d0 100644 --- a/autodata_tests/e2e/test_mimicgen_data_generation.py +++ b/autodata_tests/e2e/test_mimicgen_data_generation.py @@ -1,6 +1,9 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import h5py +import json +import numpy as np import os import tempfile @@ -8,17 +11,36 @@ 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_kpi_report, + assert_step_records_align, + assert_valid_dataset, + read_attempt_log, +) HEADLESS = True GENERATION_NUM_TRIALS = 1 +# KPIs that a one-demo MimicGen run with a source dataset and embodiment must produce. +APPLICABLE_KPIS = { + "success_rate", + "demos_per_hour", + "demos_per_human_demo", + "tracking_error", + "action_saturation", + "motion_smoothness", + "failure_share", +} def _run_franka_cube_stack_mimicgen(num_envs: int, device: str) -> None: - """Run MimicGen data generation for the Franka cube-stack task""" + """Run MimicGen data generation for the Franka cube-stack task, then report its KPIs.""" 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") + report_file = os.path.join(temp_dir, "kpis.json") + source_file = os.path.join(TestPaths.test_data_dir, "annotated_dataset_franka_stack_mimicgen.hdf5") + embodiment_file = os.path.join(TestPaths.embodiments_dir, "franka_ik_rel.yaml") args = [ TestPaths.python_path, @@ -30,11 +52,13 @@ def _run_franka_cube_stack_mimicgen(num_envs: int, device: str) -> None: "--task_descriptor", os.path.join(TestPaths.tasks_dir, "franka_cube_stack.yaml"), "--embodiment", - os.path.join(TestPaths.embodiments_dir, "franka_ik_rel.yaml"), + embodiment_file, "--input_file", - os.path.join(TestPaths.test_data_dir, "annotated_dataset_franka_stack_mimicgen.hdf5"), + source_file, "--output_file", output_file, + "--result_file", + result_file, "--generation_num_trials", str(GENERATION_NUM_TRIALS), "--num_envs", @@ -47,6 +71,35 @@ 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_step_records_align(output_file) + with h5py.File(output_file, "r") as dataset: + for episode in dataset["data"].values(): + # franka_ik_rel.yaml has no EEF offset, so the recorded pose must equal the pre-step observation. + recorded = episode["obs/datagen_info/eef_pose/franka"][:, :3, 3] + np.testing.assert_allclose(recorded, episode["obs/eef_pos"][:], atol=1e-6) + with open(result_file, encoding="utf-8") as result_handle: + result = json.load(result_handle) + assert len(read_attempt_log(result_file)) == result["num_attempts"] + + run_subprocess([ + TestPaths.python_path, + os.path.join(TestPaths.scripts_dir, "validate_dataset.py"), + output_file, + "--report", + "--result_file", + result_file, + "--source", + source_file, + "--embodiment", + embodiment_file, + "--report_file", + report_file, + ]) + kpis = assert_kpi_report(report_file, APPLICABLE_KPIS) + assert bool(kpis["failure_share"]["value"]) == (result["num_failures"] > 0) + assert any(term["func"].endswith(":randomize_object_pose") for term in result["run"]["reset_ranges"]) + for name in ("initial_state_coverage", "success_by_initial_configuration"): + assert "frame is wrong" not in json.dumps(kpis[name]), f"{name}: {kpis[name]}" @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..dc0fb91 100644 --- a/autodata_tests/e2e/test_skillgen_bin_data_generation.py +++ b/autodata_tests/e2e/test_skillgen_bin_data_generation.py @@ -17,10 +17,12 @@ 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_kpi_report, assert_valid_dataset, read_attempt_log HEADLESS = True GENERATION_NUM_TRIALS = 1 +# KPIs that a one-demo SkillGen run with a result file must produce. +APPLICABLE_KPIS = {"success_rate", "failure_share", "planner_failure_types", "planning_time"} def _run_franka_bin_stack_skillgen(num_envs: int, device: str) -> None: @@ -29,6 +31,7 @@ def _run_franka_bin_stack_skillgen(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") + report_file = os.path.join(temp_dir, "kpis.json") env_profile = os.path.join(TestPaths.env_profiles_dir, "franka_bin_stack.yaml") args = [ @@ -70,6 +73,29 @@ def _run_franka_bin_stack_skillgen(num_envs: int, device: str) -> None: "path": env_profile, "planner": "franka_stack_cube_bin", } + attempts = read_attempt_log(result_file) + assert len(attempts) == result["num_attempts"] + # Every SkillGen attempt plans at least its first transit. + assert all(attempt["planner_calls"] for attempt in attempts) + assert {"status", "wall_time_s", "solve_time_s"} <= set(attempts[0]["planner_calls"][0]) + for attempt in attempts: + last_call = attempt["planner_calls"][-1] + planner_failed = attempt["failure_cause"] == "planner" + assert last_call["success"] != planner_failed and (last_call["status"] == "Success") != planner_failed + + run_subprocess([ + TestPaths.python_path, + os.path.join(TestPaths.scripts_dir, "validate_dataset.py"), + output_file, + "--report", + "--result_file", + result_file, + "--report_file", + report_file, + ]) + kpis = assert_kpi_report(report_file, APPLICABLE_KPIS) + num_calls = sum(len(attempt["planner_calls"]) for attempt in attempts) + assert kpis["planner_failure_types"]["num_calls"] == num_calls @pytest.mark.with_subprocess diff --git a/autodata_tests/interfaces/datastream/test_datastream.py b/autodata_tests/interfaces/datastream/test_datastream.py index 56486fa..6ff07bd 100644 --- a/autodata_tests/interfaces/datastream/test_datastream.py +++ b/autodata_tests/interfaces/datastream/test_datastream.py @@ -284,3 +284,17 @@ def test_get_robot_root_pose(): def test_get_scene_state_passthrough(): datastream, _, _, _ = _datastream() assert datastream.get_scene_state(is_relative=True) is _SCENE_STATE_SENTINEL + + +def test_subtask_signals_ignore_the_final_subtask(): + task_dict = _task_dict() + task_dict["subtasks"]["franka"][-1]["subtask_term_signal"] = "stack_2" + task = TaskDescriptor.from_dict(task_dict) + adapter = _embodiment_adapter() + env = _env() + env.obs_buf["subtask_terms"] = {"grasp_1": torch.tensor([True]), "stack_1": torch.tensor([False])} + pool = DataGenInfoPool(task_descriptor=task, embodiment_adapter=adapter, device="cpu", uses_start_signals=False) + datastream = Datastream(env=env, task_descriptor=task, embodiment_adapter=adapter, source_pool=pool) + + assert datastream.has_subtask_term_signals() + assert set(datastream.get_subtask_term_signals(env_ids=[0])) == {"grasp_1", "stack_1"} diff --git a/autodata_tests/interfaces/embodiments/test_absolute_pose_embodiment_adapter.py b/autodata_tests/interfaces/embodiments/test_absolute_pose_embodiment_adapter.py index b58fff6..8e54941 100644 --- a/autodata_tests/interfaces/embodiments/test_absolute_pose_embodiment_adapter.py +++ b/autodata_tests/interfaces/embodiments/test_absolute_pose_embodiment_adapter.py @@ -40,6 +40,10 @@ def test_pose_conversion_preserves_scalar_grippers_and_body_commands(): assert all(torch.equal(noisy_channels[name], command[i]) for name, command in channels.items()) +def test_tracking_only_channel_commands_no_pose(): + assert _adapter().get_eef_names_with_pose_action() == ("left", "right") + + def test_controller_state_queries_select_requested_environment(): adapter = _adapter() state = torch.zeros(2, 2, 13) diff --git a/autodata_tests/utils/test_dataset_metrics.py b/autodata_tests/utils/test_dataset_metrics.py new file mode 100644 index 0000000..d92ac63 --- /dev/null +++ b/autodata_tests/utils/test_dataset_metrics.py @@ -0,0 +1,118 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Known-answer tests for the dataset KPI functions.""" + +import numpy as np + +import pytest + +from autodata_utils import dataset_metrics + + +def test_wilson_interval_matches_known_values(): + assert dataset_metrics.wilson_interval(0, 10) == pytest.approx((0.0, 0.2775), abs=1e-4) + assert dataset_metrics.wilson_interval(10, 10) == pytest.approx((0.7225, 1.0), abs=1e-4) + + +def test_eef_jerk_of_cubic_trajectory_is_six(): + t = np.arange(10) * 0.05 + positions = np.stack([t**3, np.zeros_like(t), np.zeros_like(t)], axis=1) + np.testing.assert_allclose(dataset_metrics.eef_jerk(positions, 0.05), 6.0, rtol=1e-6) + + +def test_saturated_fraction_counts_entries_at_the_limit(): + pose_actions = np.array([[1.0, -1.0, 0.5, 0.0, 0.2, -0.3], [0.9, 0.0, 0.0, 0.0, 0.0, 1.0]]) + assert dataset_metrics.saturated_fraction(pose_actions) == pytest.approx(3 / 12) + + +def test_path_deviation_ignores_lag_along_the_path(): + intended = np.stack([np.linspace(0.0, 1.0, 50), np.zeros(50), np.zeros(50)], axis=1) + lagging = intended[np.maximum(np.arange(50) - 3, 0)] + [0.0, 0.005, 0.0] + np.testing.assert_allclose(dataset_metrics.path_deviation(lagging, intended), 0.005, atol=1e-9) + + +def test_tracking_error_baseline_skips_channels_without_a_generated_intended_pose(): + path = np.linspace(0.0, 0.1, 30)[:, None] * np.array([1.0, 0.0, 0.0]) + beside = path + np.array([0.0, 0.002, 0.0]) + generated = [{"eef_pose": {"left": beside}, "intended_eef_pose": {"left": path}}] + # A tracking-only channel's source target is a placeholder far from its achieved pose. + source = [{"eef_pose": {"left": beside, "body": path}, "target_eef_pose": {"left": path, "body": path + 1.0}}] + + kpi = dataset_metrics._tracking_error(generated, source) + + assert kpi["source_baseline"]["p95"] == pytest.approx(2.0) + + +def _attempt(env_id: int, success: bool, start: float, duration: float) -> dict: + return {"env_id": env_id, "success": success, "start_time_s": start, "end_time_s": start + duration} + + +def test_attempts_before_stop_keeps_attempts_started_before_the_first_env_finished(): + long_successes = [_attempt(0, True, 4.0 * k, 4.0) for k in range(3)] # env 0 last finishes at 12 s + short_failures = [_attempt(1, False, 1.6 * k, 1.6) for k in range(10)] # one runs from 11.2 to 12.8 s + + kept = dataset_metrics.attempts_before_stop(long_successes + short_failures, num_envs=2) + + assert len(kept) == 11 and sum(attempt["success"] for attempt in kept) == 3 + assert dataset_metrics.attempts_before_stop(long_successes[:2] + short_failures, num_envs=2) is None + assert dataset_metrics.attempts_before_stop(short_failures, num_envs=2) is None + + +def test_coverage_ratio_drops_when_successes_leave_a_region_empty(): + rng = np.random.default_rng(1) + attempts = rng.uniform(0.0, 1.0, size=(400, 2)) + spread = attempts[rng.choice(400, 100, replace=False)] + left_half_only = attempts[attempts[:, 0] < 0.5][:100] + bounds = [(0.0, 1.0), (0.0, 1.0)] + + assert dataset_metrics.coverage_ratio(spread, attempts, bounds)["ratio"] == pytest.approx(1.0, abs=0.1) + assert dataset_metrics.coverage_ratio(left_half_only, attempts, bounds)["ratio"] == pytest.approx(0.5, abs=0.1) + + +def _attempts(x_positions: np.ndarray, successes: np.ndarray) -> list[dict]: + return [ + {"success": bool(success), "initial_object_poses": {"cube": [float(x), 0.0, 0.0]}} + for x, success in zip(x_positions, successes) + ] + + +def test_range_points_resolve_offset_and_absolute_ranges_and_reject_wrong_frames(): + attempts = _attempts(np.array([0.39, 0.41]), np.array([True, True])) + offset = { + "func": "isaaclab.envs.mdp.events:reset_root_state_uniform", + "pose_range": {"x": [-0.01, 0.01], "y": [0.0, 0.0]}, + "default_positions": {"cube": [0.4, 0.0, 0.0]}, + } + absolute = {**offset, "func": "isaaclab_tasks.contrib.stack.mdp.stack_events:randomize_object_pose"} + + points, bounds, axes = dataset_metrics._range_points(offset, "cube", attempts) + + np.testing.assert_allclose(points[:, 0], [-0.01, 0.01], atol=1e-9) + assert axes == ["x"] and bounds == [(-0.01, 0.01)] + assert "frame is wrong" in dataset_metrics._range_points(absolute, "cube", attempts) + + +def test_success_by_configuration_flags_only_a_truly_weak_region(): + rng = np.random.default_rng(2) + x_positions = rng.uniform(0.0, 0.3, size=600) + result = { + "run": { + "reset_ranges": [{ + "func": "m:randomize_object_pose", + "assets": ["cube"], + "pose_range": {"x": [0.0, 0.3]}, + "default_positions": {}, + }] + } + } + uniform = _attempts(x_positions, rng.uniform(size=600) < 0.7) + weak_left = _attempts(x_positions, (rng.uniform(size=600) < 0.7) & (x_positions > 0.1)) + # 9 of 11 at an overall rate of 65/67: unremarkable by an exact test, though a Wilson bound reads it as low. + near_perfect = _attempts( + np.array([0.05] * 11 + [0.15] * 28 + [0.25] * 28), np.array([True] * 9 + [False] * 2 + [True] * 56) + ) + + assert dataset_metrics._success_by_configuration(result, uniform, corrected=False)["value"] == 0 + assert dataset_metrics._success_by_configuration(result, weak_left, corrected=False)["value"] == 1 + assert dataset_metrics._success_by_configuration(result, near_perfect, corrected=False)["value"] == 0 diff --git a/autodata_tests/utils/test_validate_dataset.py b/autodata_tests/utils/test_validate_dataset.py index 3e1ea73..cf351d9 100644 --- a/autodata_tests/utils/test_validate_dataset.py +++ b/autodata_tests/utils/test_validate_dataset.py @@ -45,3 +45,34 @@ def test_main_returns_nonzero_for_missing_file(tmp_path: Path) -> None: """An unreadable input path should produce a failing exit status.""" assert validate_dataset.main([str(tmp_path / "missing.hdf5")]) == 1 + + +def test_report_covers_a_run_whose_every_attempt_failed(tmp_path: Path) -> None: + """A run without successes leaves an empty, invalid dataset whose outcome KPIs are still reported.""" + + dataset_path = tmp_path / "empty.hdf5" + with h5py.File(dataset_path, "w") as h5_file: + h5_file.create_group("data").attrs["env_args"] = json.dumps({"env_name": "Test-Env-v0"}) + result_path = tmp_path / "result.json" + result_path.write_text( + json.dumps({ + "num_success": 0, + "num_failures": 3, + "num_attempts": 3, + "run": {"num_envs": 1, "attempts_file": "result.attempts.jsonl"}, + }) + ) + attempt = {"env_id": 0, "success": False, "failure_cause": "execution", "failed_subtask": None, "planner_calls": []} + (tmp_path / "result.attempts.jsonl").write_text( + "".join(json.dumps({**attempt, "start_time_s": t, "end_time_s": t + 1.0}) + "\n" for t in range(3)) + ) + report_path = tmp_path / "kpis.json" + + exit_code = validate_dataset.main( + [str(dataset_path), "--report", "--result_file", str(result_path), "--report_file", str(report_path)] + ) + + kpis = json.loads(report_path.read_text())["kpis"] + assert exit_code == 1 + assert kpis["success_rate"]["value"] == 0.0 + assert kpis["failure_share"]["value"] == {"execution": 1.0} diff --git a/autodata_tests/utils/utils.py b/autodata_tests/utils/utils.py index f9cbde4..315c250 100644 --- a/autodata_tests/utils/utils.py +++ b/autodata_tests/utils/utils.py @@ -5,8 +5,11 @@ import h5py import json +import math import os +from autodata_utils.generation_result import attempts_file_for + def assert_valid_dataset(output_file: str, min_num_demos: int) -> None: """Assert the run wrote a dataset containing at least min_num_demos entries.""" @@ -41,3 +44,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 read_attempt_log(result_file: str) -> list[dict]: + """Read the attempt log written next to a generation result file.""" + + with open(attempts_file_for(result_file), encoding="utf-8") as attempts_handle: + return [json.loads(line) for line in attempts_handle if line.strip()] + + +def assert_kpi_report(report_file: str, applicable: set[str]) -> dict: + """Assert every KPI is reported or explained, and the ``applicable`` ones carry finite values. + + Returns: + The report's KPIs. + """ + + with open(report_file, encoding="utf-8") as report_handle: + kpis = json.load(report_handle)["kpis"] + for name, kpi in kpis.items(): + assert "value" in kpi or kpi.get("not_applicable"), f"KPI {name} has neither a value nor a reason: {kpi}" + for name in applicable: + assert "value" in kpis[name], f"KPI {name} is not applicable: {kpis[name].get('not_applicable')}" + value = kpis[name]["value"] + numbers = value.values() if isinstance(value, dict) else [value] + assert all(math.isfinite(number) for number in numbers), f"KPI {name} is not finite: {value}" + return kpis + + +def assert_step_records_align(dataset_file: str) -> None: + """Assert each generated episode records one intended and one achieved end-effector pose per action.""" + + with h5py.File(dataset_file, "r") as dataset: + for name, episode in dataset["data"].items(): + num_actions = len(episode["actions"]) + datagen_info = episode["obs/datagen_info"] + lengths = { + f"{group}/{eef}": len(poses) + for group in ("intended_eef_pose", "eef_pose") + for eef, poses in datagen_info[group].items() + } + assert set(lengths.values()) == {num_actions}, f"{name}: {num_actions} actions but {lengths}" diff --git a/autodata_utils/dataset_metrics.py b/autodata_utils/dataset_metrics.py new file mode 100644 index 0000000..45791b8 --- /dev/null +++ b/autodata_utils/dataset_metrics.py @@ -0,0 +1,711 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""KPIs of a generated AutoData dataset. + +The report reads the generated HDF5, the generation result JSON and its attempt log, the source +dataset, the embodiment YAML and an optional annotation summary. It needs only h5py, numpy and PyYAML, +so it runs without the simulator. Every KPI is either a value or ``not_applicable`` with a reason; +missing inputs never read as zero. +""" + +from __future__ import annotations + +import h5py +import json +import math +import numpy as np +import yaml +from collections import Counter +from pathlib import Path +from typing import Any + +from autodata_utils.generation_result import file_sha256 + +REPORT_SCHEMA_VERSION = 1 +Z_95 = 1.959963984540054 +"""Two-sided 95% standard-normal quantile.""" + +SPARSE_CELL_ATTEMPTS = 10 +"""Grid cells with fewer attempts are left out of the success-by-configuration headline.""" + +MIN_ATTEMPTS_PER_ENV_FOR_STOP_CORRECTION = 3 +"""Below this many completed attempts in some env, the correction for attempts cancelled at the stop is itself biased.""" + +SAMPLES_PER_COVERAGE_CELL = 5 +"""The coverage grid is sized so each cell holds about this many successful starts.""" + +PATH_DEVIATION_WINDOW = 20 +"""Achieved steps are matched to intended path segments at most this many steps away.""" + +RANGE_TOLERANCE = 1e-3 +"""Starts farther than this [m] outside their reset range mean the range frame was resolved wrongly.""" + +# Isaac Lab reset functions whose pose_range semantics are known. +_OFFSET_RANGE_FUNCS = ("reset_root_state_uniform",) +_ABSOLUTE_RANGE_FUNCS = ("randomize_object_pose",) + +# Whether each embodiment adapter type clips the pose part (first six entries) of its action. The report +# stays torch-free, so it cannot import the adapters themselves. +_CLIPPING_BY_EMBODIMENT_TYPE = { + "absolute_pose": False, + "delta_pose_ik_single_arm": True, + "absolute_pose_whole_body_bimanual": False, +} + + +# ---------------------------------------------------------------------------------------------- +# Statistics +# ---------------------------------------------------------------------------------------------- + + +def binomial_cdf(successes: int, trials: int, rate: float) -> float: + """Probability of at most ``successes`` successes in ``trials`` Bernoulli trials with success ``rate``.""" + if successes >= trials: + return 1.0 + assert 0.0 < rate < 1.0, "rate must lie strictly between 0 and 1 when some trial fails" + # Log space: binomial coefficients overflow a float beyond about a thousand trials. + log_terms = ( + math.lgamma(trials + 1) + - math.lgamma(k + 1) + - math.lgamma(trials - k + 1) + + k * math.log(rate) + + (trials - k) * math.log1p(-rate) + for k in range(successes + 1) + ) + return min(1.0, sum(math.exp(term) for term in log_terms)) + + +def wilson_interval(successes: int, trials: int, z: float = Z_95) -> tuple[float, float]: + """Wilson score interval for a binomial proportion.""" + assert trials > 0, "trials must be positive" + p = successes / trials + denominator = 1 + z**2 / trials + center = (p + z**2 / (2 * trials)) / denominator + half_width = z * math.sqrt(p * (1 - p) / trials + z**2 / (4 * trials**2)) / denominator + return max(0.0, center - half_width), min(1.0, center + half_width) + + +def eef_jerk(positions: np.ndarray, step_dt: float) -> np.ndarray: + """Linear jerk magnitude [m/s³] along an end-effector position trajectory. + + Args: + positions: ``[T, 3]`` positions [m] sampled every ``step_dt``. + step_dt: Control step [s]. + """ + if len(positions) < 4: + return np.empty(0) + return np.linalg.norm(np.diff(positions, n=3, axis=0), axis=1) / step_dt**3 + + +def saturated_fraction(pose_actions: np.ndarray) -> float: + """Share of pose-action entries clipped to ±1.""" + return float(np.mean(np.isclose(np.abs(pose_actions), 1.0, rtol=0.0, atol=1e-6))) + + +def path_deviation(achieved: np.ndarray, intended: np.ndarray, window: int = PATH_DEVIATION_WINDOW) -> np.ndarray: + """Distance [m] from each achieved position to the intended path near the same step. + + Args: + achieved: ``[T, 3]`` achieved end-effector positions [m]. + intended: ``[T, 3]`` intended end-effector positions [m], step-aligned with ``achieved``. + window: Steps on either side whose intended segments are candidates. + """ + assert len(achieved) == len(intended), "achieved and intended positions must be step-aligned" + if len(intended) < 2: + return np.linalg.norm(achieved - intended, axis=1) + starts, ends = intended[:-1], intended[1:] + deviations = np.empty(len(achieved)) + for step, point in enumerate(achieved): + lo, hi = max(0, step - window), min(len(starts), step + window + 1) + segment_start, segment = starts[lo:hi], ends[lo:hi] - starts[lo:hi] + length_sq = np.einsum("ij,ij->i", segment, segment) + t = np.einsum("ij,ij->i", point - segment_start, segment) / np.where(length_sq > 0, length_sq, 1.0) + closest = segment_start + np.clip(t, 0.0, 1.0)[:, None] * segment + deviations[step] = np.linalg.norm(closest - point, axis=1).min() + return deviations + + +def grid_cells(points: np.ndarray, bounds: list[tuple[float, float]], bins: int) -> np.ndarray: + """Flat cell index of each point on a ``bins``-per-dimension grid over ``bounds``.""" + index = np.zeros(len(points), dtype=int) + for dim, (lo, hi) in enumerate(bounds): + cell = np.clip(((points[:, dim] - lo) / (hi - lo) * bins).astype(int), 0, bins - 1) + index = index * bins + cell + return index + + +def min_coverage_successes(num_dims: int) -> int: + """Fewest successful starts that give the coverage grid at least two bins per dimension.""" + return SAMPLES_PER_COVERAGE_CELL * 2**num_dims + + +def coverage_ratio( + success_points: np.ndarray, attempt_points: np.ndarray, bounds: list[tuple[float, float]], seed: int = 0 +) -> dict[str, float]: + """Cells occupied by successful starts over cells occupied by as many attempted starts. + + The grid is sized to about :data:`SAMPLES_PER_COVERAGE_CELL` successes per cell, so raw sample count + does not drive the ratio. + + Args: + success_points: ``[S, D]`` starts of successful attempts, in range coordinates, with + ``S >= min_coverage_successes(D)``. + attempt_points: ``[A, D]`` starts of all attempts, ``A >= S``. + bounds: Per-dimension ``(lo, hi)`` of the reset range. + seed: Seed of the random attempt subsets. + """ + num_success = len(success_points) + assert num_success >= min_coverage_successes(len(bounds)), "too few successes for a two-bin grid" + rng = np.random.default_rng(seed) + bins = round((num_success / SAMPLES_PER_COVERAGE_CELL) ** (1 / len(bounds))) + + def occupied(points: np.ndarray) -> int: + return len(np.unique(grid_cells(points, bounds, bins))) + + attempt_occupancy = np.mean( + [occupied(attempt_points[rng.choice(len(attempt_points), num_success, replace=False)]) for _ in range(50)] + ) + return {"ratio": occupied(success_points) / attempt_occupancy, "bins_per_dim": bins} + + +# ---------------------------------------------------------------------------------------------- +# Inputs +# ---------------------------------------------------------------------------------------------- + + +def embodiment_clips_pose_actions(path: str) -> bool: + """Whether the embodiment's adapter clips the pose part (first six entries) of its action to ±1.""" + with open(path, encoding="utf-8") as handle: + data = yaml.safe_load(handle) + embodiment_type = data.get("type") + assert ( + embodiment_type in _CLIPPING_BY_EMBODIMENT_TYPE + ), f"Embodiment type {embodiment_type!r} is unknown to the KPI report; add it to _CLIPPING_BY_EMBODIMENT_TYPE." + return _CLIPPING_BY_EMBODIMENT_TYPE[embodiment_type] and bool( + data.get("action_layout", {}).get("clip_pose_action_to_unit", True) + ) + + +def load_episodes(path: str) -> tuple[list[dict], dict]: + """Per-episode arrays needed by the KPIs, plus the dataset's ``env_args``. + + End-effector positions (``eef_pos``, ``[T, E, 3]``) come from ``obs/datagen_info/eef_pose``, the control-link + pose that annotation and generation both record before each step, so source and generated data share a frame + and a step. + """ + episodes = [] + with h5py.File(path, "r") as h5_file: + data = h5_file["data"] + env_args_raw = data.attrs.get("env_args", "{}") + env_args = json.loads(env_args_raw.decode() if isinstance(env_args_raw, bytes) else env_args_raw) + for name in data: + demo = data[name] + episode: dict[str, Any] = {"actions": demo["actions"][:]} + datagen_info = demo["obs"].get("datagen_info") + if datagen_info is not None: + for group in ("intended_eef_pose", "target_eef_pose", "eef_pose"): + if group in datagen_info: + episode[group] = {eef: datagen_info[group][eef][:][:, :3, 3] for eef in datagen_info[group]} + if "eef_pose" in episode: + episode["eef_pos"] = np.stack([episode["eef_pose"][eef] for eef in sorted(episode["eef_pose"])], 1) + episodes.append(episode) + return episodes, env_args + + +def _step_dt_from_env_args(env_args: dict) -> float | None: + sim_args = env_args.get("sim_args") or {} + if "dt" in sim_args and "decimation" in sim_args: + return float(sim_args["dt"]) * int(sim_args["decimation"]) + return None + + +def _read_attempts(result_file: str, result: dict) -> list[dict] | None: + attempts_file = result.get("run", {}).get("attempts_file") + if attempts_file is None: + return None + path = Path(result_file).parent / attempts_file + assert path.exists(), f"{result_file} names the attempt log {path}, which is missing." + with path.open(encoding="utf-8") as handle: + return [json.loads(line) for line in handle if line.strip()] + + +# ---------------------------------------------------------------------------------------------- +# KPIs +# ---------------------------------------------------------------------------------------------- + + +def _kpi(value: Any, unit: str, n: int, **details: Any) -> dict[str, Any]: + return {"value": value, "unit": unit, "n": n, **details} + + +def _not_applicable(reason: str) -> dict[str, str]: + return {"not_applicable": reason} + + +def _percentiles(values: np.ndarray, scale: float = 1.0) -> dict[str, float]: + p50, p95 = np.percentile(values, [50, 95]) * scale + return {"p50": float(p50), "p95": float(p95)} + + +def attempts_before_stop(attempts: list[dict], num_envs: int) -> list[dict] | None: + """Attempts that started before the first env completed its last logged attempt. + + Attempts still running when generation stops are cancelled and never logged, and long attempts are the likelier + to be running, so the full log under-represents long outcomes. Each env's cancelled attempt started at or after + that env's last completion, so every attempt that started before the earliest such completion was logged. This + subset is selected by start time, not duration. None when some env completed fewer than + :data:`MIN_ATTEMPTS_PER_ENV_FOR_STOP_CORRECTION` attempts, where neither the subset nor the full log is unbiased. + """ + last_end: dict[int, float] = {} + completed: Counter[int] = Counter() + for attempt in attempts: + last_end[attempt["env_id"]] = max(last_end.get(attempt["env_id"], 0.0), attempt["end_time_s"]) + completed[attempt["env_id"]] += 1 + if len(completed) < num_envs or min(completed.values()) < MIN_ATTEMPTS_PER_ENV_FOR_STOP_CORRECTION: + return None + cutoff = min(last_end.values()) + return [attempt for attempt in attempts if attempt["start_time_s"] < cutoff] + + +_UNCORRECTED_NOTE = ( + "uncorrected: attempts cancelled at the stop can bias it (correcting needs an attempt log with " + f"{MIN_ATTEMPTS_PER_ENV_FOR_STOP_CORRECTION}+ completed attempts from every env)" +) +_CORRECTED_NOTE = "counts only attempts that started before the first env's last completion" + + +def _success_rate(result: dict | None, stop_attempts: list[dict] | None) -> dict: + if result is None: + return _not_applicable("needs --result_file") + successes, attempts = result["num_success"], result["num_attempts"] + if attempts == 0: + return _not_applicable("the run made no attempts") + raw_rate = successes / attempts + note = _UNCORRECTED_NOTE + if stop_attempts is not None: + successes, attempts = sum(attempt["success"] for attempt in stop_attempts), len(stop_attempts) + note = _CORRECTED_NOTE + lo, hi = wilson_interval(successes, attempts) + return _kpi( + successes / attempts, + "fraction", + attempts, + interval_95=[lo, hi], + raw_rate=raw_rate, + stop_rule=result.get("run", {}).get("stop_rule"), + note=note, + ) + + +def _demos_per_hour(result: dict | None) -> dict: + run = (result or {}).get("run", {}) + if "generation_time_s" not in run: + return _not_applicable("needs a --result_file with run timing") + return _kpi( + result["num_success"] / (run["generation_time_s"] / 3600), + "demos/h", + result["num_success"], + generation_time_s=run["generation_time_s"], + setup_time_s=run["setup_time_s"], + device=run["device"], + num_envs=run["num_envs"], + ) + + +def _demos_per_human_demo(result: dict | None, source_episodes: list[dict] | None) -> dict: + if result is None: + return _not_applicable("needs --result_file") + num_source = result.get("run", {}).get("source", {}).get("num_episodes") + if num_source is None and source_episodes is not None: + num_source = len(source_episodes) + if not num_source: + return _not_applicable("needs the source episode count (--result_file run metadata or --source)") + note = "input ratio: about requested demos / source demos when the run stops at a success count" + return _kpi(result["num_success"] / num_source, "ratio", num_source, note=note) + + +def _annotation_yield(annotation_summary: dict | None) -> dict: + if annotation_summary is None: + return _not_applicable("needs --annotation_summary") + processed = annotation_summary["num_processed"] + if processed == 0: + return _not_applicable("no episodes were processed") + return _kpi( + annotation_summary["num_exported"] / processed, + "fraction", + processed, + rejections=annotation_summary["rejections"], + manual_remarks=annotation_summary["manual_remarks"], + mode=annotation_summary["mode"], + ) + + +def _pooled_path_deviation(episodes: list[dict], intended_key: str, eefs: set[str] | None = None) -> np.ndarray: + deviations = [ + path_deviation(episode["eef_pose"][eef], episode[intended_key][eef]) + for episode in episodes + if intended_key in episode and "eef_pose" in episode + for eef in episode[intended_key] + if eefs is None or eef in eefs + ] + return np.concatenate(deviations) if deviations else np.empty(0) + + +def _tracking_error(episodes: list[dict], source_episodes: list[dict] | None) -> dict: + if not episodes: + return _not_applicable("the dataset has no episodes") + deviations = _pooled_path_deviation(episodes, "intended_eef_pose") + if not len(deviations): + return _not_applicable("dataset has no obs/datagen_info/intended_eef_pose") + baseline = None + if source_episodes is not None: + # Generation records intended poses only for commanded EEFs; the source's target poses include + # placeholders for tracking-only channels. + commanded = {eef for episode in episodes for eef in episode.get("intended_eef_pose", {})} + source_deviations = _pooled_path_deviation(source_episodes, "target_eef_pose", commanded) + baseline = _percentiles(source_deviations, 1000.0) if len(source_deviations) else None + return _kpi( + _percentiles(deviations, 1000.0), + "mm from the noise-free intended path (position)", + len(deviations), + source_baseline=baseline, + note="generated values include injected action noise; the source baseline is against executed commands", + ) + + +def _action_saturation(episodes: list[dict], source_episodes: list[dict] | None, clips: bool | None) -> dict: + if clips is None: + return _not_applicable("needs --embodiment") + if not clips: + return _not_applicable("the embodiment adapter does not clip pose actions") + if not episodes: + return _not_applicable("the dataset has no episodes") + pose_actions = np.concatenate([episode["actions"][:, :6] for episode in episodes]) + baseline = None + if source_episodes is not None: + baseline = saturated_fraction(np.concatenate([episode["actions"][:, :6] for episode in source_episodes])) + return _kpi(saturated_fraction(pose_actions), "fraction", pose_actions.size, source_baseline=baseline) + + +def _pooled_jerk(episodes: list[dict], step_dt: float) -> np.ndarray: + jerks = [ + eef_jerk(episode["eef_pos"][:, eef], step_dt) + for episode in episodes + if "eef_pos" in episode + for eef in range(episode["eef_pos"].shape[1]) + ] + return np.concatenate(jerks) if jerks else np.empty(0) + + +def _motion_smoothness( + episodes: list[dict], + step_dt: float | None, + source_episodes: list[dict] | None, + source_step_dt: float | None, +) -> dict: + if step_dt is None: + return _not_applicable("no step_dt in the result file or dataset sim_args") + if not episodes: + return _not_applicable("the dataset has no episodes") + jerk = _pooled_jerk(episodes, step_dt) + if not len(jerk): + return _not_applicable("dataset has no obs/datagen_info/eef_pose") + baseline = None + if source_episodes is not None: + source_jerk = _pooled_jerk(source_episodes, source_step_dt or step_dt) + baseline = _percentiles(source_jerk) if len(source_jerk) else None + return _kpi( + _percentiles(jerk), + "m/s^3 end-effector jerk", + len(jerk), + source_baseline=baseline, + note=( + None + if source_episodes is None or source_step_dt is not None + else "source has no sim_args; its baseline assumes the generated step_dt" + ), + ) + + +def _range_points(reset_range: dict, asset: str, attempts: list[dict]) -> tuple | str: + """Starts of ``asset`` in reset-range coordinates over the range's non-zero-width x/y axes.""" + func_name = reset_range["func"].rsplit(":", 1)[-1] + if func_name in _OFFSET_RANGE_FUNCS: + origin = reset_range["default_positions"].get(asset) + if origin is None: + return f"no default position recorded for {asset}" + elif func_name in _ABSOLUTE_RANGE_FUNCS: + origin = [0.0, 0.0, 0.0] + else: + return f"unknown pose_range semantics of {reset_range['func']}" + axes = [ + (index, axis) + for index, axis in enumerate(("x", "y")) + if axis in reset_range["pose_range"] and reset_range["pose_range"][axis][1] > reset_range["pose_range"][axis][0] + ] + if not axes: + return f"{asset} has no x/y variation" + starts = np.array([attempt["initial_object_poses"][asset][:2] for attempt in attempts]) - np.array(origin[:2]) + points = starts[:, [index for index, _ in axes]] + bounds = [tuple(reset_range["pose_range"][axis]) for _, axis in axes] + in_range = np.all( + [ + (points[:, dim] >= lo - RANGE_TOLERANCE) & (points[:, dim] <= hi + RANGE_TOLERANCE) + for dim, (lo, hi) in enumerate(bounds) + ], + axis=0, + ) + if not in_range.all(): + return f"only {in_range.mean():.0%} of {asset} starts lie inside its reset range, so its frame is wrong" + return points, bounds, [axis for _, axis in axes] + + +def _configuration_inputs(result: dict | None, attempts: list[dict] | None) -> tuple[list, list[str]] | str: + if result is None or attempts is None: + return "needs --result_file with an attempt log" + if not attempts: + return "the attempt log is empty" + reset_ranges = result.get("run", {}).get("reset_ranges", []) + entries, skipped = [], [] + for reset_range in reset_ranges: + for asset in reset_range["assets"]: + if not all(asset in attempt["initial_object_poses"] for attempt in attempts): + skipped.append(f"{asset} is not a rigid object") + continue + resolved = _range_points(reset_range, asset, attempts) + if isinstance(resolved, str): + skipped.append(resolved) + else: + entries.append((asset, *resolved)) + if not entries: + return "no reset range with x/y variation" + (f" ({'; '.join(skipped)})" if skipped else "") + return entries, skipped + + +def _success_by_configuration(result: dict | None, attempts: list[dict] | None, corrected: bool) -> dict: + resolved = _configuration_inputs(result, attempts) + if isinstance(resolved, str): + return _not_applicable(resolved) + entries, skipped = resolved + success = np.array([attempt["success"] for attempt in attempts]) + overall_rate = float(success.mean()) + if overall_rate == 0.0: + return _not_applicable("no attempt succeeded") + per_asset = {} + for asset, points, bounds, axes in entries: + cells = grid_cells(points, bounds, 3) + table = [] + for cell in range(3 ** len(bounds)): + in_cell = cells == cell + n = int(in_cell.sum()) + table.append({"cell": cell, "attempts": n, "successes": int(success[in_cell].sum())}) + per_asset[asset] = {"axes": axes, "range": [list(bound) for bound in bounds], "cells": table} + dense = [ + entry for grid in per_asset.values() for entry in grid["cells"] if entry["attempts"] >= SPARSE_CELL_ATTEMPTS + ] + if not dense: + reasons = [f"no grid cell has {SPARSE_CELL_ATTEMPTS} or more attempts", *skipped] + return _not_applicable("; ".join(reasons)) + for entry in dense: + entry["rate"] = entry["successes"] / entry["attempts"] + # One-sided exact test: how likely so few successes are if the cell had the overall rate. + entry["p_value"] = binomial_cdf(entry["successes"], entry["attempts"], overall_rate) + rate_label = "corrected" if corrected else "uncorrected" + return _kpi( + sum(entry["p_value"] < 0.025 / len(dense) for entry in dense), + "cells significantly below the overall success rate", + len(attempts), + note=( + f"of {len(dense)} cells with {SPARSE_CELL_ATTEMPTS}+ attempts; " + f"overall success rate {overall_rate:.3f} ({rate_label})" + ), + overall_rate=overall_rate, + per_asset=per_asset, + skipped=skipped, + ) + + +def _initial_state_coverage(result: dict | None, attempts: list[dict] | None) -> dict: + resolved = _configuration_inputs(result, attempts) + if isinstance(resolved, str): + return _not_applicable(resolved) + entries, skipped = resolved + success = np.array([attempt["success"] for attempt in attempts]) + per_asset = {} + for asset, points, bounds, axes in entries: + if success.sum() < min_coverage_successes(len(bounds)): + skipped.append(f"{asset} needs {min_coverage_successes(len(bounds))} successful attempts") + continue + per_asset[asset] = { + "axes": axes, + "range": [list(bound) for bound in bounds], + **coverage_ratio(points[success], points, bounds), + } + if not per_asset: + return _not_applicable("; ".join(skipped)) + worst = min(entry["ratio"] for entry in per_asset.values()) + return _kpi(worst, "worst per-object coverage ratio", int(success.sum()), per_asset=per_asset, skipped=skipped) + + +def _failure_share(attempts: list[dict] | None, stop_attempts: list[dict] | None) -> dict: + if attempts is None: + return _not_applicable("needs --result_file with an attempt log") + note = _UNCORRECTED_NOTE + if stop_attempts is not None: + attempts = stop_attempts + note = _CORRECTED_NOTE + failures = [attempt for attempt in attempts if not attempt["success"]] + if not failures: + return _kpi({}, "fraction of failures", 0, note=note) + by_cause = Counter(attempt["failure_cause"] or "unknown" for attempt in failures) + by_subtask: dict[str, dict[str, dict[str, int]]] = {} + for attempt in failures: + for eef, subtask in (attempt["failed_subtask"] or {}).items(): + cause_table = by_subtask.setdefault(attempt["failure_cause"] or "unknown", {}).setdefault(eef, {}) + cause_table[str(subtask)] = cause_table.get(str(subtask), 0) + 1 + return _kpi( + {cause: count / len(failures) for cause, count in by_cause.items()}, + "fraction of failures", + len(failures), + by_subtask=by_subtask, + note=note, + ) + + +def _planner_calls(attempts: list[dict] | None) -> list[dict]: + return [call for attempt in attempts or [] for call in attempt.get("planner_calls", [])] + + +def _planner_failure_types(attempts: list[dict] | None) -> dict: + calls = _planner_calls(attempts) + if not calls: + return _not_applicable("no motion-planner calls in the attempt log") + failed = [call for call in calls if not call["success"]] + return _kpi( + {status: count / len(failed) for status, count in Counter(call["status"] for call in failed).items()}, + "fraction of failed plans", + len(failed), + num_calls=len(calls), + ) + + +def _planning_time(result: dict | None, attempts: list[dict] | None) -> dict: + calls = _planner_calls(attempts) + if not calls: + return _not_applicable("no motion-planner calls in the attempt log") + per_attempt = np.array([ + sum(call["wall_time_s"] for call in attempt["planner_calls"]) + for attempt in attempts + if attempt["planner_calls"] + ]) + generation_time = (result or {}).get("run", {}).get("generation_time_s") + share = float(per_attempt.sum() / generation_time) if generation_time else None + return _kpi(_percentiles(per_attempt), "s per attempt", len(per_attempt), share_of_generation_time=share) + + +def build_report( + dataset: str, + result_file: str | None = None, + source: str | None = None, + embodiment: str | None = None, + annotation_summary: str | None = None, +) -> dict[str, Any]: + """Compute every KPI for a generated dataset. + + Args: + dataset: Generated HDF5 dataset holding the run's successful demos. + result_file: Generation result JSON (``--result_file`` of ``generate_dataset.py``). + source: Annotated source HDF5 dataset the generation started from; enables source baselines. + embodiment: Embodiment YAML used for generation; needed for action saturation. + annotation_summary: Annotation summary JSON (``--summary_file`` of ``annotate_demos.py``). + """ + episodes, env_args = load_episodes(dataset) + source_episodes, source_env_args = load_episodes(source) if source else (None, {}) + result = None + attempts = None + stop_attempts = None + if result_file: + with open(result_file, encoding="utf-8") as handle: + result = json.load(handle) + assert len(episodes) == result["num_success"], ( + f"{dataset} holds {len(episodes)} demos but {result_file} counts {result['num_success']} successes;" + " they come from different runs." + ) + attempts = _read_attempts(result_file, result) + assert attempts is None or len(attempts) == result["num_attempts"], ( + f"The attempt log has {len(attempts)} lines but {result_file} counts {result['num_attempts']} attempts;" + " they come from different runs." + ) + num_envs = result.get("run", {}).get("num_envs") + if attempts and num_envs: + stop_attempts = attempts_before_stop(attempts, num_envs) + recorded_source = result.get("run", {}).get("source", {}) + assert source is None or recorded_source.get("sha256") in ( + None, + file_sha256(source), + ), f"--source {source} is not the dataset this run generated from ({recorded_source.get('path')})." + annotation = None + if annotation_summary: + with open(annotation_summary, encoding="utf-8") as handle: + annotation = json.load(handle) + + step_dt = (result or {}).get("run", {}).get("step_dt") or _step_dt_from_env_args(env_args) + source_step_dt = _step_dt_from_env_args(source_env_args) + kpis = { + "success_rate": _success_rate(result, stop_attempts), + "demos_per_hour": _demos_per_hour(result), + "demos_per_human_demo": _demos_per_human_demo(result, source_episodes), + "annotation_yield": _annotation_yield(annotation), + "tracking_error": _tracking_error(episodes, source_episodes), + "action_saturation": _action_saturation( + episodes, source_episodes, embodiment_clips_pose_actions(embodiment) if embodiment else None + ), + "motion_smoothness": _motion_smoothness(episodes, step_dt, source_episodes, source_step_dt), + "initial_state_coverage": _initial_state_coverage(result, stop_attempts or attempts), + "success_by_initial_configuration": _success_by_configuration( + result, stop_attempts or attempts, corrected=stop_attempts is not None + ), + "failure_share": _failure_share(attempts, stop_attempts), + "planner_failure_types": _planner_failure_types(attempts), + "planning_time": _planning_time(result, attempts), + } + return { + "schema_version": REPORT_SCHEMA_VERSION, + "dataset": dataset, + "num_episodes": len(episodes), + "num_source_episodes": len(source_episodes) if source_episodes is not None else None, + "kpis": kpis, + } + + +def _format_value(value: Any) -> str: + if value is None: + return "-" + if isinstance(value, float): + return f"{value:.3g}" + if isinstance(value, dict): + return ", ".join(f"{key}={_format_value(item)}" for key, item in value.items()) or "none" + return str(value) + + +def format_report(report: dict[str, Any]) -> str: + """Render the report as an aligned text table.""" + rows = [] + for name, kpi in report["kpis"].items(): + if "not_applicable" in kpi: + rows.append((name, "n/a", "", "", kpi["not_applicable"])) + continue + details = [] + if "interval_95" in kpi: + details.append(f"95% CI {kpi['interval_95'][0]:.3f}-{kpi['interval_95'][1]:.3f}") + if kpi.get("source_baseline") is not None: + details.append(f"source {_format_value(kpi['source_baseline'])}") + if kpi.get("note"): + details.append(kpi["note"]) + rows.append((name, _format_value(kpi["value"]), kpi["unit"], str(kpi["n"]), "; ".join(details))) + headers = ("KPI", "Value", "Unit", "n", "Baseline / note") + widths = [max(len(row[col]) for row in [headers, *rows]) for col in range(len(headers))] + lines = [" ".join(cell.ljust(width) for cell, width in zip(row, widths)).rstrip() for row in [headers, *rows]] + lines.insert(1, " ".join("-" * width for width in widths)) + return "\n".join(lines) diff --git a/autodata_utils/generation_result.py b/autodata_utils/generation_result.py index c38213d..4ea98df 100644 --- a/autodata_utils/generation_result.py +++ b/autodata_utils/generation_result.py @@ -6,11 +6,13 @@ from __future__ import annotations import contextlib +import hashlib import json import os import tempfile from collections.abc import Mapping from pathlib import Path +from typing import Any def write_generation_result( @@ -19,6 +21,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 +34,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 +55,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 +71,30 @@ def write_generation_result( finally: with contextlib.suppress(FileNotFoundError): os.unlink(temporary_file_name) + + +def file_sha256(path: str) -> str: + """Hex SHA-256 digest of a file, recorded so a report can tell which source a run used.""" + digest = hashlib.sha256() + with open(path, "rb") as handle: + for chunk in iter(lambda: handle.read(1 << 20), b""): + digest.update(chunk) + return digest.hexdigest() + + +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.""" + + def __init__(self, path: str) -> None: + self._path = Path(path) + 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(record, sort_keys=True) + "\n") diff --git a/docs/index.rst b/docs/index.rst index 6d4c9bf..6c5694f 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -146,5 +146,6 @@ Table of Contents :maxdepth: 1 :caption: References + pages/references/kpis pages/references/troubleshooting pages/references/release_notes diff --git a/docs/pages/advanced/motion_planners.rst b/docs/pages/advanced/motion_planners.rst index c06038c..6ebe494 100644 --- a/docs/pages/advanced/motion_planners.rst +++ b/docs/pages/advanced/motion_planners.rst @@ -96,8 +96,9 @@ Visualization and Debugging The cuRobo backend can visualize its collision-sphere model and planned trajectories via `rerun `_ (``visualize_spheres`` / ``visualize_plan`` config flags). During multi-env generation these are enabled only for env 0 to keep the simulation -responsive. Planner diagnostics (success rates, timing) are available through -``get_planner_info()``. +responsive. After each ``update_world_and_plan_motion()`` call, the cuRobo backend's +``last_plan_summary`` holds the outcome, the cuRobo status, and wall-clock and solver times. +SkillGen copies it into each attempt's ``planner_calls`` (see :doc:`../references/kpis`). Adding a New Backend -------------------- diff --git a/docs/pages/concepts/data_generator.rst b/docs/pages/concepts/data_generator.rst index b182383..ce7bb5b 100644 --- a/docs/pages/concepts/data_generator.rst +++ b/docs/pages/concepts/data_generator.rst @@ -123,7 +123,10 @@ Results and Statistics Three layers report what happened: -* **Per attempt** — ``GenerationResult``: ``success`` and ``initial_state``. +* **Per attempt** — ``GenerationResult``: ``success``, ``failure_cause``, ``failed_subtask``, + ``num_steps``, ``planner_calls`` 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 @@ -146,7 +149,9 @@ Three layers report what happened: ``env_profile`` appears only when an :doc:`environment profile ` was applied — the generated dataset itself records only the base env id, so the sidecar is - what documents that the scene was modified. + what documents that the scene was modified. A ``run`` block records how the run was + configured and how long it took. :doc:`../references/kpis` describes it and the report + computed from these files. The GenerationAlgorithm Plug-In ------------------------------- diff --git a/docs/pages/references/kpis.rst b/docs/pages/references/kpis.rst new file mode 100644 index 0000000..5c42529 --- /dev/null +++ b/docs/pages/references/kpis.rst @@ -0,0 +1,142 @@ +KPIs +==== + +The KPI report describes what a generated dataset delivers. It answers four questions: + +- how many usable demonstrations a run yields, and at what cost; +- whether they make good training data; +- whether they cover the task's variation; +- why attempts fail. + +``scripts/validate_dataset.py --report`` computes the report from files that generation and +annotation already write. It needs no simulator. + +Producing the Inputs +-------------------- + +Every generated dataset carries two per-step entries under ``obs/datagen_info``: + +- ``intended_eef_pose``: the noise-free waypoint pose commanded at that step, for end effectors whose + action commands a pose (tracking-only channels have none); +- ``eef_pose``: the achieved control-link pose, read before that step. + +Pass ``--result_file`` to ``generate_dataset.py`` to also write: + +- The **result JSON**, with a ``run`` block: seed, stop rule, device, env count, control step + (``step_dt``), setup and generation time, source dataset path and digest, and the reset events + that sample object poses from a ``pose_range``. +- The **attempt log**, ``.attempts.jsonl``. It holds one JSON line per completed + attempt: env id, start and end time, outcome, ``failure_cause`` (``planner`` or ``execution``), the + failed subtask per end effector, step count, initial object poses, and one summary per + motion-planner call. + +Two inputs come from elsewhere: + +- Pass ``--summary_file`` to ``annotate_demos.py`` to record annotation yield. +- Pass the annotated source dataset as ``--source`` to enable the source baselines. Its + ``obs/datagen_info/eef_pose`` is in the same control-link frame as the generated data's. + +The embodiment YAML is needed only for ``action_saturation``. + +Running the Report +------------------ + +.. code-block:: bash + + python scripts/validate_dataset.py datasets/generated.hdf5 --report \ + --result_file datasets/generated.json \ + --source datasets/annotated_datasets/dataset_franka_annotated.hdf5 \ + --embodiment autodata_examples/embodiments/franka_ik_rel.yaml \ + --report_file datasets/generated_kpis.json + +The command validates the dataset, prints a KPI table, and optionally writes the full report as +JSON. A dataset left empty because every attempt failed is still reported. The report stops with an +error in three cases: + +- the dataset's demonstration count differs from the result file's success count; +- the attempt log's line count differs from the result file's attempt count; +- ``--source`` is not the dataset the run generated from. + +KPI Definitions +--------------- + +.. list-table:: + :widths: 26 52 22 + :header-rows: 1 + + * - KPI + - Definition + - Needs + * - ``success_rate`` + - Successful attempts over attempts, with a Wilson 95% interval. An attempt succeeds if the + task is successful at any step. ``raw_rate`` is the uncorrected rate from the result file. + - Result JSON, attempt log + * - ``demos_per_hour`` + - Successful demonstrations per hour of generation time. One run uses one GPU, so this is + also demonstrations per GPU-hour. Setup time is reported separately. + - Result JSON + * - ``demos_per_human_demo`` + - Successful demonstrations over source demonstrations. When the run stops at a success + count, this is the requested count divided by the source count. + - Result JSON + * - ``annotation_yield`` + - Exported over processed episodes, with rejection reasons and manual re-marks. + - Annotation summary + * - ``tracking_error`` + - p50 and p95 distance [mm] from each achieved end-effector position to the noise-free + intended path. Position only. + - Generated dataset + * - ``action_saturation`` + - Share of pose-action entries at ±1, for adapters that clip pose actions. + - Embodiment YAML + * - ``motion_smoothness`` + - p50 and p95 end-effector linear jerk [m/s³]. + - ``step_dt`` + * - ``initial_state_coverage`` + - Per object, grid cells occupied by successful starts over cells occupied by as many + attempted starts. The worst object is reported. + - Attempt log, reset ranges + * - ``success_by_initial_configuration`` + - Number of cells whose success rate is significantly below the overall success rate. Each + object's reset range has three cells per varying axis (3 for one axis, 9 for two). Only + cells with at least 10 attempts count. A cell is flagged when the exact binomial + probability of so few successes at the overall rate is below 0.025 divided by the number + of counted cells. ``per_asset`` holds every cell's counts. + - Attempt log, reset ranges + * - ``failure_share`` + - Share of failed attempts by cause, and counts by end effector and subtask. + - Attempt log + * - ``planner_failure_types`` + - Share of failed motion plans by cuRobo status. + - Attempt log (SkillGen) + * - ``planning_time`` + - p50 and p95 planning wall-clock time per attempt [s], and its share of generation time. + Planning blocks the event loop, so this share is time no env steps. + - Attempt log (SkillGen) + +Reading the Values +------------------ + +- **Source baselines.** Quality KPIs carry a ``source_baseline`` computed on the source + demonstrations. Compare against it rather than reading absolute values. +- **Tracking error includes noise.** Generated values include the injected action noise by design. + The source baseline is measured against the executed teleoperation commands. +- **Attempts cancelled at the stop.** When generation stops, attempts still running are cancelled + and never logged. Long attempts are the likelier to be running, so a raw rate under-counts long + outcomes. ``success_rate``, ``failure_share`` and the configuration KPIs count only attempts + that started before the first env finished its last logged attempt; all of those completed. The + correction needs at least three logged attempts from every env, since with fewer it is biased + too. Otherwise they use all attempts and mark the value uncorrected. +- **Not applicable.** A KPI that cannot be computed reports ``not_applicable`` with a reason. + Missing inputs never read as zero. +- **Coverage grid.** The grid is sized to about five successful starts per cell, with at least two + bins per varying axis. An object with fewer successful starts is skipped. +- **Range frames.** Ranges of ``randomize_object_pose`` are absolute in the env frame. Ranges of + ``reset_root_state_uniform`` are offsets from the asset's default position. The configuration + KPIs skip, and list the reason for, other reset functions, ranges without x or y variation, and + objects whose starts fall outside their range. +- **Failed subtasks.** Execution failures are attributed only when the env publishes the + termination signal of every non-final subtask in its ``subtask_terms`` observation group. + +The report's JSON layout is versioned by ``schema_version``. The functions live in +``autodata_utils/dataset_metrics.py``. diff --git a/scripts/annotate_demos.py b/scripts/annotate_demos.py index b0d6b2b..c972e98 100644 --- a/scripts/annotate_demos.py +++ b/scripts/annotate_demos.py @@ -25,6 +25,9 @@ Pass ``--auto`` to annotate subtask termination signals without keyboard input (``--viz none`` supported). ``--signal_obs_group`` selects the observation group holding the per-subtask boolean terms (default ``subtask_terms``). SkillGen start signals still require manual mode. + +Pass ``--summary_file`` to save the processed and exported episode counts, the rejection reasons and +the number of manual re-marks as JSON. """ """Launch Isaac Sim Simulator first.""" @@ -81,6 +84,12 @@ default="subtask_terms", help="Observation group holding the per-subtask boolean terms read in --auto mode.", ) +parser.add_argument( + "--summary_file", + type=str, + default=None, + help="Optional JSON path for the annotation summary (processed and exported counts, rejection reasons).", +) parser.add_argument( "--enable_cameras", @@ -101,11 +110,13 @@ import contextlib # noqa: E402 import gymnasium as gym # noqa: E402 +import json # noqa: E402 import math # noqa: E402 import os # noqa: E402 import sys # noqa: E402 import torch # noqa: E402 import traceback # noqa: E402 +from collections import Counter # noqa: E402 from collections.abc import Callable # noqa: E402 import isaaclab_tasks # noqa: F401, E402 (registers gym envs) @@ -126,6 +137,8 @@ current_action_index = 0 marked_subtask_action_indices: list[int] = [] skip_episode = False +rejection_counts: Counter[str] = Counter() +manual_remark_count = 0 _datastream: Datastream | None = None @@ -315,7 +328,6 @@ def main() -> int: for episode_index, episode_name in enumerate(episode_names): if not simulation_app.is_running() or simulation_app.is_exiting(): break - processed_episode_count += 1 print(f"\nAnnotating episode #{episode_index} ({episode_name})") episode = dataset_file_handler.load_episode(episode_name, env.device) @@ -327,6 +339,7 @@ def main() -> int: annotated = annotate_episode_in_manual_mode( env, episode, success_term, term_signal_names, start_signal_names, annotate_start_signals ) + processed_episode_count += 1 if annotated and not skip_episode: env.recorder_manager.set_success_to_episodes( @@ -342,6 +355,17 @@ 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 args_cli.summary_file: + summary = { + "mode": "auto" if args_cli.auto else "manual", + "num_processed": processed_episode_count, + "num_exported": exported_episode_count, + "rejections": dict(rejection_counts), + "manual_remarks": manual_remark_count, + } + with open(args_cli.summary_file, "w", encoding="utf-8") as summary_handle: + json.dump(summary, summary_handle, indent=2, sort_keys=True) + summary_handle.write("\n") print("Exiting the app.") return exported_episode_count finally: @@ -427,7 +451,7 @@ def annotate_episode_in_manual_mode( Returns: True if the episode was fully annotated, False if it was skipped or never succeeded. """ - global is_paused, marked_subtask_action_indices, skip_episode + global is_paused, marked_subtask_action_indices, skip_episode, manual_remark_count term_signal_action_indices: dict[str, int] = {} start_signal_action_indices: dict[str, int] = {} @@ -455,6 +479,7 @@ def annotate_episode_in_manual_mode( task_succeeded = replay_episode(env, episode, success_term) if skip_episode: print("\tSkipping the episode.") + rejection_counts["skipped"] += 1 return False print(f"\tSubtasks marked at action indices: {marked_subtask_action_indices}") @@ -480,12 +505,14 @@ def annotate_episode_in_manual_mode( if not task_succeeded: print("\tThe final task was not completed.") + rejection_counts["task_not_completed"] += 1 return False print( f"\tOnly {len(marked_subtask_action_indices)} out of {expected_count} subtask signals" f' for eef "{eef_name}" were annotated.' ) + manual_remark_count += 1 print(f'\tThe episode will be replayed again for re-marking subtask signals for eef "{eef_name}".\n') num_steps = len(episode.data["actions"]) @@ -580,6 +607,7 @@ def sample_signals(action_index: int) -> None: task_succeeded = replay_episode(env, episode, success_term, on_step=sample_signals) if not task_succeeded: print("\tThe final task was not completed.") + rejection_counts["task_not_completed"] += 1 return False num_steps = len(episode.data["actions"]) @@ -591,18 +619,22 @@ def sample_signals(action_index: int) -> None: signal = torch.stack([step[signal_name][0] for step in sampled_signals]) edge = _first_rising_edge(signal) if edge is None: - reason = "already active at episode start" if bool(signal[0]) else "never fired" + active_at_start = bool(signal[0]) + reason = "already active at episode start" if active_at_start else "never fired" print(f'\tSubtask signal "{signal_name}" (eef "{eef_name}") {reason}.') + rejection_counts["signal_active_at_start" if active_at_start else "signal_never_fired"] += 1 return False eef_edges.append(edge) if any(later <= earlier for earlier, later in zip(eef_edges, eef_edges[1:])): print(f'\tSubtask signals for eef "{eef_name}" fired out of order: {dict(zip(eef_term_names, eef_edges))}.') + rejection_counts["signals_out_of_order"] += 1 return False spacing_error = _validate_boundary_spacing(task_descriptor, eef_name, eef_edges, num_steps) if spacing_error is not None: print(f'\tRejected annotations for eef "{eef_name}": {spacing_error}.') + rejection_counts["boundaries_too_close"] += 1 return False print(f'\tDetected subtask signals for eef "{eef_name}": {dict(zip(eef_term_names, eef_edges))}') diff --git a/scripts/generate_dataset.py b/scripts/generate_dataset.py index d6f5167..6f90b49 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,9 @@ from autodata_interfaces.tasks.task_descriptor_utils import resolve_task_algorithm from autodata_utils.cli import callable_path +# Taken before the simulator launches, so the recorded setup time includes its startup. +PROCESS_START_TIME = 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 +101,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", @@ -141,10 +148,12 @@ import sys # noqa: E402 import torch # noqa: E402 import traceback # 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, @@ -154,7 +163,14 @@ setup_output_paths, ) 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 import pose_math # noqa: E402 +from autodata_utils.generation_result import ( # noqa: E402 + AttemptLog, + attempts_file_for, + file_sha256, + write_generation_result, +) +from autodata_utils.tensor_utils import as_torch # noqa: E402 async def run_data_generator( @@ -166,9 +182,14 @@ async def run_data_generator( success_term, pause_subtask: bool, stats: dict, + timing: 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() + if timing["generation_start"] is None: + timing["generation_start"] = attempt_start try: result = await data_generator.generate( env_id=env_id, @@ -187,6 +208,67 @@ 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 measured from process start.""" + return { + "env_id": env_id, + "success": result.success, + "failure_cause": result.failure_cause, + "failed_subtask": result.failed_subtask, + "num_steps": result.num_steps, + "start_time_s": start_time - PROCESS_START_TIME, + "end_time_s": end_time - PROCESS_START_TIME, + "initial_object_poses": _initial_object_poses(result.initial_state, env_id), + "planner_calls": list(result.planner_calls), + } + + +def _initial_object_poses(initial_state: dict, env_id: int) -> dict[str, list[float]]: + """``[x, y, yaw]`` [m, m, rad] of each rigid object in the env frame at the start of the attempt.""" + poses = {} + for name, object_state in initial_state.get("rigid_object", {}).items(): + root_pose = as_torch(object_state["root_pose"])[env_id] + rotation = pose_math.matrix_from_quat(root_pose[3:7]) + yaw = torch.atan2(rotation[1, 0], rotation[0, 0]) + poses[name] = [float(root_pose[0]), float(root_pose[1]), float(yaw)] + return poses + + +def _reset_ranges(env_cfg: Any) -> list[dict]: + """Reset event terms that sample object poses from a ``pose_range``. + + ``default_positions`` lets a reader resolve ranges that are offsets from an asset's default pose. + """ + ranges = [] + events = getattr(env_cfg, "events", None) + for term_name, term in (vars(events) if events is not None else {}).items(): + params = getattr(term, "params", None) + if getattr(term, "mode", None) != "reset" or not params or not isinstance(params.get("pose_range"), dict): + continue + func = ( + term.func + if isinstance(term.func, str) + else f"{getattr(term.func, '__module__', '')}:{getattr(term.func, '__name__', '')}" + ) + asset_cfgs = list(params.get("asset_cfgs", [])) + ([params["asset_cfg"]] if "asset_cfg" in params else []) + assets = [asset_cfg.name for asset_cfg in asset_cfgs] + default_positions = {} + for name in assets: + init_state = getattr(getattr(env_cfg.scene, name, None), "init_state", None) + if init_state is not None: + default_positions[name] = [float(value) for value in init_state.pos] + ranges.append({ + "term": term_name, + "func": func, + "assets": assets, + "pose_range": {axis: [float(bounds[0]), float(bounds[1])] for axis, bounds in params["pose_range"].items()}, + "default_positions": default_positions, + }) + return ranges def build_datastream( @@ -216,6 +298,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() @@ -225,6 +308,7 @@ def setup_async_generation( data_generator = DataGenerator(datastream=datastream, algorithm=algorithm) stats = {"num_success": 0, "num_failures": 0, "num_attempts": 0} + timing = {"generation_start": None} tasks = [] for env_id in range(num_envs): @@ -238,6 +322,8 @@ def setup_async_generation( success_term=success_term, pause_subtask=pause_subtask, stats=stats, + timing=timing, + attempt_log=attempt_log, ) ) tasks.append(task) @@ -249,6 +335,7 @@ def setup_async_generation( "action_queue": env_action_queue, "info_pool": datastream.source_pool, "stats": stats, + "timing": timing, } @@ -362,6 +449,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) if attempts_file else None try: async_components = setup_async_generation( datastream=datastream, @@ -369,10 +458,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_end = None try: generation_completed = env_loop( env, @@ -383,6 +474,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: @@ -396,6 +488,7 @@ def main() -> None: _close_motion_planners(motion_planners) if generation_completed and args_cli.result_file: + generation_start = async_components["timing"]["generation_start"] write_generation_result( result_file=args_cli.result_file, algorithm=args_cli.alg, @@ -406,6 +499,22 @@ def main() -> None: if env_profile else None ), + run_metadata={ + "seed": generation_policy_params.seed, + "stop_rule": "successes" if generation_policy_params.guarantee_success else "attempts", + "device": args_cli.device, + "num_envs": args_cli.num_envs, + "step_dt": float(env.step_dt), + "setup_time_s": generation_start - PROCESS_START_TIME, + "generation_time_s": generation_end - generation_start, + "source": { + "path": args_cli.input_file, + "sha256": file_sha256(args_cli.input_file), + "num_episodes": datastream.num_source_demos, + }, + "reset_ranges": _reset_ranges(env_cfg), + "attempts_file": Path(attempts_file).name, + }, ) finally: env.close() diff --git a/scripts/validate_dataset.py b/scripts/validate_dataset.py index a53db2d..253b016 100644 --- a/scripts/validate_dataset.py +++ b/scripts/validate_dataset.py @@ -1,14 +1,17 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Structurally validate AutoData HDF5 datasets. +"""Structurally validate AutoData HDF5 datasets and report the KPIs of a generated one. -Prints a summary table (episode count, env id, sim args) per file followed by any issues. +Prints a summary table (episode count, env id, sim args) per file followed by any issues. With +``--report``, also prints the KPIs of a single generated dataset. Usage:: python scripts/validate_dataset.py [ ...] python scripts/validate_dataset.py directory/*.hdf5 + python scripts/validate_dataset.py --report --result_file \\ + --source --embodiment [--report_file ] The command exits nonzero if any supplied file is invalid. """ @@ -21,6 +24,8 @@ import os import sys +from autodata_utils.dataset_metrics import build_report, format_report + REQUIRED_DEMO_FIELDS = { "actions": h5py.Dataset, "initial_state": h5py.Group, @@ -163,7 +168,15 @@ def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser(description="Validate AutoData demonstration HDF5 datasets.") parser.add_argument("files", nargs="+", help="HDF5 dataset file(s) to validate.") + parser.add_argument("--report", action="store_true", help="Also report the KPIs of a single generated dataset.") + parser.add_argument("--result_file", help="Generation result JSON written by generate_dataset.py --result_file.") + parser.add_argument("--source", help="Annotated source HDF5 the generation started from; enables baselines.") + parser.add_argument("--embodiment", help="Embodiment YAML used for generation; needed for action saturation.") + parser.add_argument("--annotation_summary", help="Summary JSON written by annotate_demos.py --summary_file.") + parser.add_argument("--report_file", help="Also write the KPI report to this JSON file.") args = parser.parse_args(argv) + if args.report and len(args.files) != 1: + parser.error("--report takes exactly one dataset") results = [validate_file(path) for path in args.files] print_results(results) @@ -173,6 +186,25 @@ def main(argv: list[str] | None = None) -> int: print() print(f"All {total} file(s) valid." if num_invalid == 0 else f"{num_invalid} of {total} file(s) invalid.") + # A run whose every attempt failed leaves an empty dataset; its outcome KPIs are still worth reporting. + only_empty = results[0].num_episodes == 0 and len(results[0].issues) == 1 + if args.report and not (num_invalid == 0 or only_empty): + print("\nKPI report skipped: the dataset is invalid.") + elif args.report: + report = build_report( + dataset=args.files[0], + result_file=args.result_file, + source=args.source, + embodiment=args.embodiment, + annotation_summary=args.annotation_summary, + ) + print() + print(format_report(report)) + if args.report_file: + with open(args.report_file, "w", encoding="utf-8") as report_handle: + json.dump(report, report_handle, indent=2) + report_handle.write("\n") + return 1 if num_invalid else 0 diff --git a/skills/autodata-generate-datasets/references/cli-reference.md b/skills/autodata-generate-datasets/references/cli-reference.md index 19a5e3b..a52ac13 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). | @@ -44,6 +44,7 @@ Annotate source demonstrations with subtask signals. | `--output_file ` | Annotated output. | | `--auto` | Annotate termination signals without keyboard input (`--viz none` supported). | | `--signal_obs_group ` | Observation group holding per-subtask boolean terms (default `subtask_terms`). | +| `--summary_file ` | Write processed/exported counts and rejection reasons as JSON. | Interactive controls: `N` play/resume, `B` pause, `S` mark subtask signal, `Q` skip episode. SkillGen start signals still require manual mode. @@ -60,3 +61,11 @@ python scripts/validate_dataset.py directory/*.hdf5 Required demo fields checked: `actions` (dataset), `initial_state` (group), `obs` (group). The command exits non-zero if any file is invalid, so workflows can gate on a clean exit code. + +Add `--report` to print the KPIs of one generated dataset (see the KPIs docs page): + +```bash +python scripts/validate_dataset.py --report --result_file \ + --source --embodiment [--annotation_summary ] \ + [--report_file ] +``` From 1f79f4a6ae9417183a03a8f0eb9cdf5bc827fb3a Mon Sep 17 00:00:00 2001 From: Billy Okal Date: Mon, 19 Oct 2026 13:47:16 -0500 Subject: [PATCH 2/4] refactor: trim KPI recording to what the KPIs read Label demos per hour with the GPU or CPU model. Drop recorded fields no KPI reads (seed, stop rule, setup time, step count, source digest, cuRobo solver timings, annotation mode and re-marks), simplify the cuRobo plan summary to outcome, status and wall time, and shorten the KPI docs. --- autodata_core/algorithms.py | 3 +- autodata_core/data_generator.py | 6 - autodata_core/pool.py | 5 - autodata_interfaces/datastream/datastream.py | 23 +-- .../motion_planners/curobo/curobo_planner.py | 51 +----- autodata_tests/core/test_data_generator.py | 1 - .../e2e/test_mimicgen_data_generation.py | 13 +- .../e2e/test_skillgen_bin_data_generation.py | 6 +- autodata_tests/utils/utils.py | 15 -- autodata_utils/dataset_metrics.py | 101 +++--------- autodata_utils/generation_result.py | 10 -- docs/pages/advanced/motion_planners.rst | 5 +- docs/pages/concepts/data_generator.rst | 10 +- docs/pages/references/kpis.rst | 148 ++++++------------ scripts/annotate_demos.py | 22 +-- scripts/generate_dataset.py | 56 +++---- scripts/validate_dataset.py | 8 +- .../references/cli-reference.md | 3 +- 18 files changed, 131 insertions(+), 355 deletions(-) diff --git a/autodata_core/algorithms.py b/autodata_core/algorithms.py index bd36256..cb03a15 100644 --- a/autodata_core/algorithms.py +++ b/autodata_core/algorithms.py @@ -185,8 +185,7 @@ def __init__(self, motion_planners: dict[int, Any]) -> None: motion_planners: One planner per env_id. Planners must expose ``update_world_and_plan_motion(target_pose, expected_attached_object, env_id, ...)`` and ``get_planned_poses() -> list[torch.Tensor[4, 4]]``. ``config.motion_noise_scale`` - is read if present, and so is ``last_plan_summary``, a dict describing the latest - planning call, which is copied into the attempt's ``planner_calls``. + and ``last_plan_summary`` (copied into the attempt's ``planner_calls``) are read if present. """ assert motion_planners, "SkillGen requires at least one motion planner" self.motion_planners = motion_planners diff --git a/autodata_core/data_generator.py b/autodata_core/data_generator.py index 4b3c77b..6456bd0 100644 --- a/autodata_core/data_generator.py +++ b/autodata_core/data_generator.py @@ -57,7 +57,6 @@ class GenerationResult: failed_subtask: Per-EEF index of the subtask where the attempt failed. For a planner failure, the subtask being planned; for an execution failure, the first subtask whose termination signal never fired (the last subtask if all fired). None when unknown. - num_steps: Number of env steps executed. planner_calls: One summary per motion-planner call made during the attempt. """ @@ -65,7 +64,6 @@ class GenerationResult: success: bool failure_cause: Literal["planner", "execution"] | None = None failed_subtask: dict[str, int] | None = None - num_steps: int = 0 planner_calls: tuple[dict, ...] = () @@ -559,7 +557,6 @@ async def generate( if self._track_term_signals is None: self._track_term_signals = self.datastream.has_subtask_term_signals() fired_term_signals: set[str] = set() - num_steps = 0 randomized_subtask_boundaries: dict[str, np.ndarray] | None = None prev_pool_size = 0 @@ -597,7 +594,6 @@ async def generate( success=False, failure_cause="planner", failed_subtask={eef_name: eef_state.current_subtask_index}, - num_steps=num_steps, planner_calls=self._collect_planner_calls(eef_states), ) next_trajectory, _is_motion_plan = result @@ -619,7 +615,6 @@ async def generate( env_id=env_id, env_action_queue=env_action_queue, ) - num_steps += 1 if self._track_term_signals: signals = self.datastream.get_subtask_term_signals(env_ids=[env_id]) fired_now = torch.stack([value[0] for value in signals.values()]).tolist() @@ -659,7 +654,6 @@ async def generate( success=buffers.success, failure_cause="execution" if failed else None, failed_subtask=self._failed_subtasks(fired_term_signals) if failed and self._track_term_signals else None, - num_steps=num_steps, planner_calls=self._collect_planner_calls(eef_states), ) diff --git a/autodata_core/pool.py b/autodata_core/pool.py index 0fca268..102b8a2 100644 --- a/autodata_core/pool.py +++ b/autodata_core/pool.py @@ -122,11 +122,6 @@ def _add_episode(self, episode: EpisodeData) -> None: if "datagen_info" not in ep_grp["obs"]: raise ValueError("Episode lacks 'datagen_info' obs annotations") - missing = {"eef_pose", "object_pose", "target_eef_pose", "subtask_term_signals"} - set( - ep_grp["obs"]["datagen_info"] - ) - if missing: - raise ValueError(f"Episode 'datagen_info' lacks {sorted(missing)}; pass an annotated source dataset") eef_pose = ep_grp["obs"]["datagen_info"]["eef_pose"] object_poses_dict = ep_grp["obs"]["datagen_info"]["object_pose"] diff --git a/autodata_interfaces/datastream/datastream.py b/autodata_interfaces/datastream/datastream.py index 497f24f..e133354 100644 --- a/autodata_interfaces/datastream/datastream.py +++ b/autodata_interfaces/datastream/datastream.py @@ -277,19 +277,12 @@ def get_object_poses(self, env_ids: Sequence[int] | None = None) -> dict[str, to return object_pose_matrix def _boundary_signal_names(self) -> list[str]: - """Termination-signal names of every non-final subtask, across EEFs. - - The final subtask ends with the trajectory, so its signal never marks a boundary. - """ + """Termination-signal names of every non-final subtask; the final one ends with the trajectory.""" return [name for eef in self.get_eef_names() for name in self.get_term_signal_names(eef)[:-1] if name] def has_subtask_term_signals(self, obs_group: str = "subtask_terms") -> bool: - """Whether :meth:`get_subtask_term_signals` can read every subtask-boundary signal. - - True when ``env.obs_buf[obs_group]`` is a non-concatenated group holding one term per - boundary signal. Not every env publishes such a group. - """ + """Whether the env publishes every signal :meth:`get_subtask_term_signals` reads; not every env does.""" obs_buf = getattr(self.env, "obs_buf", None) group = obs_buf.get(obs_group) if isinstance(obs_buf, dict) else None @@ -301,13 +294,13 @@ def has_subtask_term_signals(self, obs_group: str = "subtask_terms") -> bool: def get_subtask_term_signals( self, env_ids: Sequence[int] | None = None, obs_group: str = "subtask_terms" ) -> dict[str, torch.Tensor]: - """Read the current subtask-boundary signals from the env observation buffer. + """Read the current per-subtask termination signals from the env observation buffer. - For the termination signal of every non-final subtask (across all EEFs), reads the - identically named boolean observation term from ``env.obs_buf[obs_group]``. The env must - therefore publish one observation term per boundary signal in a non-concatenated group. - Auto annotation samples this each replay step to locate subtask boundaries without a human - in the loop. + For every non-final subtask's termination-signal name declared by the task descriptor (across all + EEFs), reads the identically named boolean observation term from + ``env.obs_buf[obs_group]``. The env must therefore publish one observation term per + term-signal name in a non-concatenated group. Auto annotation samples this each replay + step to locate subtask boundaries without a human in the loop. Args: env_ids: Environments to read; ``None`` reads all. diff --git a/autodata_interfaces/motion_planners/curobo/curobo_planner.py b/autodata_interfaces/motion_planners/curobo/curobo_planner.py index 2ac8e1e..f35595e 100644 --- a/autodata_interfaces/motion_planners/curobo/curobo_planner.py +++ b/autodata_interfaces/motion_planners/curobo/curobo_planner.py @@ -287,8 +287,7 @@ def __init__( # Initialize planning state self._current_plan: JointState | None = None self._plan_index: int = 0 - # One record per cuRobo plan_single call of the current update_world_and_plan_motion call. - self._plan_phase_results: list[dict[str, Any]] = [] + self._plan_failure_status: str | None = None self.last_plan_summary: dict[str, Any] | None = None # Initialize visualization state @@ -1311,9 +1310,8 @@ def _plan_to_contact_pose( ) planning_success = False - result: Any = None try: - result = self.motion_gen.plan_single(start_state, goal_pose, self.plan_config) + result: Any = self.motion_gen.plan_single(start_state, goal_pose, self.plan_config) if result.success.item(): if result.optimized_plan is not None and len(result.optimized_plan.position) != 0: @@ -1334,16 +1332,11 @@ def _plan_to_contact_pose( self.logger.debug(f"Contact planning succeeded with {len(self._current_plan.position)} waypoints") else: self.logger.debug(f"Contact planning failed: {result.status}") + self._plan_failure_status = str(getattr(result.status, "value", result.status)) except Exception as e: self.logger.debug(f"Error during planning: {e}") - # Telemetry must never change the planning outcome or skip the sphere restore below. - try: - phase_summary = self._summarize_plan_phase(result, planning_success) - except Exception as e: - self.logger.debug(f"Error summarizing the planning phase: {e}") - phase_summary = {"success": planning_success, "status": "Success" if planning_success else "exception"} - self._plan_phase_results.append(phase_summary) + self._plan_failure_status = "exception" # Always restore sphere state after planning, regardless of success if contact: @@ -1856,8 +1849,7 @@ def update_world_and_plan_motion( Provides a high-level interface that handles the complete planning workflow: world synchronization, object attachment/detachment, gripper configuration, and motion planning. Afterwards :attr:`last_plan_summary` holds the call's outcome, the - cuRobo status of the first failed phase, wall-clock time [s], summed cuRobo solve, IK and - trajectory-optimization times [s], and summed cuRobo attempts. + cuRobo status of the failed phase, and the wall-clock time [s]. Args: target_pose: Target end-effector pose as 4x4 transformation matrix @@ -1870,7 +1862,7 @@ def update_world_and_plan_motion( True if complete planning pipeline succeeded, False if any step failed """ start_time = time.perf_counter() - self._plan_phase_results = [] + self._plan_failure_status = None # Always reset the plan before starting a new one to ensure a clean state self.reset_plan() @@ -1935,42 +1927,13 @@ def update_world_and_plan_motion( self._detach_objects() - failed_phase = next((phase for phase in self._plan_phase_results if not phase["success"]), None) self.last_plan_summary = { "success": plan_success, - "status": "Success" if plan_success else (failed_phase["status"] if failed_phase else "unknown"), + "status": "Success" if plan_success else self._plan_failure_status or "unknown", "wall_time_s": time.perf_counter() - start_time, - "num_phases": len(self._plan_phase_results), - **{ - key: sum(phase.get(key, 0) for phase in self._plan_phase_results) - for key in ("solve_time_s", "ik_time_s", "trajopt_time_s", "attempts") - }, } return plan_success - @staticmethod - def _summarize_plan_phase(result: Any, success: bool) -> dict[str, Any]: - """Outcome, cuRobo status, timing [s] and attempt count of one planning phase. - - ``result`` is None when ``plan_single`` raised. A phase whose cuRobo result succeeded but whose - post-processing raised is reported with status ``exception``. - """ - if success: - status = "Success" - elif result is None or (result.success is not None and bool(result.success.item())): - status = "exception" - elif result.status is None: - status = "unknown" - else: - status = str(getattr(result.status, "value", result.status)) - summary: dict[str, Any] = {"success": success, "status": status} - if result is not None: - summary["solve_time_s"] = float(result.solve_time) - summary["ik_time_s"] = float(result.ik_time) - summary["trajopt_time_s"] = float(result.trajopt_time) - summary["attempts"] = int(result.attempts) - return summary - # ===================================================================================== # UTILITY METHODS # ===================================================================================== diff --git a/autodata_tests/core/test_data_generator.py b/autodata_tests/core/test_data_generator.py index d6ff769..2fde05b 100644 --- a/autodata_tests/core/test_data_generator.py +++ b/autodata_tests/core/test_data_generator.py @@ -184,7 +184,6 @@ async def serve_reset() -> None: assert result.success is False assert result.failure_cause == "planner" assert result.failed_subtask == {"franka": 0} - assert result.num_steps == 0 def test_execution_failure_is_attributed_to_the_first_unfinished_subtask(): diff --git a/autodata_tests/e2e/test_mimicgen_data_generation.py b/autodata_tests/e2e/test_mimicgen_data_generation.py index aa493d0..a8cdceb 100644 --- a/autodata_tests/e2e/test_mimicgen_data_generation.py +++ b/autodata_tests/e2e/test_mimicgen_data_generation.py @@ -11,12 +11,7 @@ from autodata_tests.utils.constants import TestPaths from autodata_tests.utils.subprocess import run_subprocess -from autodata_tests.utils.utils import ( - assert_kpi_report, - assert_step_records_align, - assert_valid_dataset, - read_attempt_log, -) +from autodata_tests.utils.utils import assert_kpi_report, assert_valid_dataset HEADLESS = True GENERATION_NUM_TRIALS = 1 @@ -71,15 +66,14 @@ 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_step_records_align(output_file) with h5py.File(output_file, "r") as dataset: for episode in dataset["data"].values(): + assert len(episode["obs/datagen_info/intended_eef_pose/franka"]) == len(episode["actions"]) # franka_ik_rel.yaml has no EEF offset, so the recorded pose must equal the pre-step observation. recorded = episode["obs/datagen_info/eef_pose/franka"][:, :3, 3] np.testing.assert_allclose(recorded, episode["obs/eef_pos"][:], atol=1e-6) with open(result_file, encoding="utf-8") as result_handle: result = json.load(result_handle) - assert len(read_attempt_log(result_file)) == result["num_attempts"] run_subprocess([ TestPaths.python_path, @@ -96,8 +90,9 @@ def _run_franka_cube_stack_mimicgen(num_envs: int, device: str) -> None: report_file, ]) kpis = assert_kpi_report(report_file, APPLICABLE_KPIS) - assert bool(kpis["failure_share"]["value"]) == (result["num_failures"] > 0) assert any(term["func"].endswith(":randomize_object_pose") for term in result["run"]["reset_ranges"]) + hardware = result["run"]["gpu"] if device == "cuda" else result["run"]["cpu"] + assert hardware and hardware in kpis["demos_per_hour"]["note"] for name in ("initial_state_coverage", "success_by_initial_configuration"): assert "frame is wrong" not in json.dumps(kpis[name]), f"{name}: {kpis[name]}" diff --git a/autodata_tests/e2e/test_skillgen_bin_data_generation.py b/autodata_tests/e2e/test_skillgen_bin_data_generation.py index dc0fb91..1088ce0 100644 --- a/autodata_tests/e2e/test_skillgen_bin_data_generation.py +++ b/autodata_tests/e2e/test_skillgen_bin_data_generation.py @@ -74,10 +74,8 @@ def _run_franka_bin_stack_skillgen(num_envs: int, device: str) -> None: "planner": "franka_stack_cube_bin", } attempts = read_attempt_log(result_file) - assert len(attempts) == result["num_attempts"] # Every SkillGen attempt plans at least its first transit. assert all(attempt["planner_calls"] for attempt in attempts) - assert {"status", "wall_time_s", "solve_time_s"} <= set(attempts[0]["planner_calls"][0]) for attempt in attempts: last_call = attempt["planner_calls"][-1] planner_failed = attempt["failure_cause"] == "planner" @@ -93,9 +91,7 @@ def _run_franka_bin_stack_skillgen(num_envs: int, device: str) -> None: "--report_file", report_file, ]) - kpis = assert_kpi_report(report_file, APPLICABLE_KPIS) - num_calls = sum(len(attempt["planner_calls"]) for attempt in attempts) - assert kpis["planner_failure_types"]["num_calls"] == num_calls + assert_kpi_report(report_file, APPLICABLE_KPIS) @pytest.mark.with_subprocess diff --git a/autodata_tests/utils/utils.py b/autodata_tests/utils/utils.py index 315c250..121580d 100644 --- a/autodata_tests/utils/utils.py +++ b/autodata_tests/utils/utils.py @@ -70,18 +70,3 @@ def assert_kpi_report(report_file: str, applicable: set[str]) -> dict: numbers = value.values() if isinstance(value, dict) else [value] assert all(math.isfinite(number) for number in numbers), f"KPI {name} is not finite: {value}" return kpis - - -def assert_step_records_align(dataset_file: str) -> None: - """Assert each generated episode records one intended and one achieved end-effector pose per action.""" - - with h5py.File(dataset_file, "r") as dataset: - for name, episode in dataset["data"].items(): - num_actions = len(episode["actions"]) - datagen_info = episode["obs/datagen_info"] - lengths = { - f"{group}/{eef}": len(poses) - for group in ("intended_eef_pose", "eef_pose") - for eef, poses in datagen_info[group].items() - } - assert set(lengths.values()) == {num_actions}, f"{name}: {num_actions} actions but {lengths}" diff --git a/autodata_utils/dataset_metrics.py b/autodata_utils/dataset_metrics.py index 45791b8..422972b 100644 --- a/autodata_utils/dataset_metrics.py +++ b/autodata_utils/dataset_metrics.py @@ -1,12 +1,9 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""KPIs of a generated AutoData dataset. +"""KPIs of a generated AutoData dataset, computed offline from generation and annotation outputs. -The report reads the generated HDF5, the generation result JSON and its attempt log, the source -dataset, the embodiment YAML and an optional annotation summary. It needs only h5py, numpy and PyYAML, -so it runs without the simulator. Every KPI is either a value or ``not_applicable`` with a reason; -missing inputs never read as zero. +Every KPI is either a value or ``not_applicable`` with a reason; missing inputs never read as zero. """ from __future__ import annotations @@ -20,24 +17,14 @@ from pathlib import Path from typing import Any -from autodata_utils.generation_result import file_sha256 - -REPORT_SCHEMA_VERSION = 1 Z_95 = 1.959963984540054 -"""Two-sided 95% standard-normal quantile.""" - SPARSE_CELL_ATTEMPTS = 10 """Grid cells with fewer attempts are left out of the success-by-configuration headline.""" - MIN_ATTEMPTS_PER_ENV_FOR_STOP_CORRECTION = 3 -"""Below this many completed attempts in some env, the correction for attempts cancelled at the stop is itself biased.""" - +"""Below this many completed attempts in some env, the stop correction is itself biased.""" SAMPLES_PER_COVERAGE_CELL = 5 -"""The coverage grid is sized so each cell holds about this many successful starts.""" - PATH_DEVIATION_WINDOW = 20 """Achieved steps are matched to intended path segments at most this many steps away.""" - RANGE_TOLERANCE = 1e-3 """Starts farther than this [m] outside their reset range mean the range frame was resolved wrongly.""" @@ -45,19 +32,6 @@ _OFFSET_RANGE_FUNCS = ("reset_root_state_uniform",) _ABSOLUTE_RANGE_FUNCS = ("randomize_object_pose",) -# Whether each embodiment adapter type clips the pose part (first six entries) of its action. The report -# stays torch-free, so it cannot import the adapters themselves. -_CLIPPING_BY_EMBODIMENT_TYPE = { - "absolute_pose": False, - "delta_pose_ik_single_arm": True, - "absolute_pose_whole_body_bimanual": False, -} - - -# ---------------------------------------------------------------------------------------------- -# Statistics -# ---------------------------------------------------------------------------------------------- - def binomial_cdf(successes: int, trials: int, rate: float) -> float: """Probability of at most ``successes`` successes in ``trials`` Bernoulli trials with success ``rate``.""" @@ -169,30 +143,22 @@ def occupied(points: np.ndarray) -> int: return {"ratio": occupied(success_points) / attempt_occupancy, "bins_per_dim": bins} -# ---------------------------------------------------------------------------------------------- -# Inputs -# ---------------------------------------------------------------------------------------------- - - def embodiment_clips_pose_actions(path: str) -> bool: - """Whether the embodiment's adapter clips the pose part (first six entries) of its action to ±1.""" + """Whether the embodiment's adapter clips the pose part (first six entries) of its action to ±1. + + Only the single-arm delta-pose IK adapter clips; the report stays torch-free, so it reads the YAML. + """ with open(path, encoding="utf-8") as handle: data = yaml.safe_load(handle) - embodiment_type = data.get("type") - assert ( - embodiment_type in _CLIPPING_BY_EMBODIMENT_TYPE - ), f"Embodiment type {embodiment_type!r} is unknown to the KPI report; add it to _CLIPPING_BY_EMBODIMENT_TYPE." - return _CLIPPING_BY_EMBODIMENT_TYPE[embodiment_type] and bool( + return data.get("type") == "delta_pose_ik_single_arm" and bool( data.get("action_layout", {}).get("clip_pose_action_to_unit", True) ) def load_episodes(path: str) -> tuple[list[dict], dict]: - """Per-episode arrays needed by the KPIs, plus the dataset's ``env_args``. + """Per-episode actions and ``obs/datagen_info`` positions, plus the dataset's ``env_args``. - End-effector positions (``eef_pos``, ``[T, E, 3]``) come from ``obs/datagen_info/eef_pose``, the control-link - pose that annotation and generation both record before each step, so source and generated data share a frame - and a step. + Annotation and generation both record ``eef_pose`` before each step, so source and generated data share a frame. """ episodes = [] with h5py.File(path, "r") as h5_file: @@ -230,11 +196,6 @@ def _read_attempts(result_file: str, result: dict) -> list[dict] | None: return [json.loads(line) for line in handle if line.strip()] -# ---------------------------------------------------------------------------------------------- -# KPIs -# ---------------------------------------------------------------------------------------------- - - def _kpi(value: Any, unit: str, n: int, **details: Any) -> dict[str, Any]: return {"value": value, "unit": unit, "n": n, **details} @@ -251,11 +212,9 @@ def _percentiles(values: np.ndarray, scale: float = 1.0) -> dict[str, float]: def attempts_before_stop(attempts: list[dict], num_envs: int) -> list[dict] | None: """Attempts that started before the first env completed its last logged attempt. - Attempts still running when generation stops are cancelled and never logged, and long attempts are the likelier - to be running, so the full log under-represents long outcomes. Each env's cancelled attempt started at or after - that env's last completion, so every attempt that started before the earliest such completion was logged. This - subset is selected by start time, not duration. None when some env completed fewer than - :data:`MIN_ATTEMPTS_PER_ENV_FOR_STOP_CORRECTION` attempts, where neither the subset nor the full log is unbiased. + Attempts running at the stop are cancelled unlogged, and long ones are likelier to be running. Each env's + cancelled attempt started after that env's last completion, so this start-time subset was fully logged. None + when some env completed fewer than :data:`MIN_ATTEMPTS_PER_ENV_FOR_STOP_CORRECTION` attempts. """ last_end: dict[int, float] = {} completed: Counter[int] = Counter() @@ -293,7 +252,6 @@ def _success_rate(result: dict | None, stop_attempts: list[dict] | None) -> dict attempts, interval_95=[lo, hi], raw_rate=raw_rate, - stop_rule=result.get("run", {}).get("stop_rule"), note=note, ) @@ -302,25 +260,25 @@ def _demos_per_hour(result: dict | None) -> dict: run = (result or {}).get("run", {}) if "generation_time_s" not in run: return _not_applicable("needs a --result_file with run timing") + gpu, cpu = run["gpu"] or "unknown GPU", run["cpu"] or "unknown CPU" + hardware = f"on {gpu}" if run["device"] != "cpu" else f"CPU physics on {cpu}, with {gpu}" return _kpi( result["num_success"] / (run["generation_time_s"] / 3600), "demos/h", result["num_success"], + note=f"{hardware}, {run['num_envs']} envs; compare only on the same hardware and env count", generation_time_s=run["generation_time_s"], - setup_time_s=run["setup_time_s"], device=run["device"], + gpu=run["gpu"], + cpu=run["cpu"], num_envs=run["num_envs"], ) -def _demos_per_human_demo(result: dict | None, source_episodes: list[dict] | None) -> dict: - if result is None: - return _not_applicable("needs --result_file") - num_source = result.get("run", {}).get("source", {}).get("num_episodes") - if num_source is None and source_episodes is not None: - num_source = len(source_episodes) +def _demos_per_human_demo(result: dict | None) -> dict: + num_source = (result or {}).get("run", {}).get("num_source_episodes") if not num_source: - return _not_applicable("needs the source episode count (--result_file run metadata or --source)") + return _not_applicable("needs a --result_file with the source episode count") note = "input ratio: about requested demos / source demos when the run stops at a success count" return _kpi(result["num_success"] / num_source, "ratio", num_source, note=note) @@ -336,8 +294,6 @@ def _annotation_yield(annotation_summary: dict | None) -> dict: "fraction", processed, rejections=annotation_summary["rejections"], - manual_remarks=annotation_summary["manual_remarks"], - mode=annotation_summary["mode"], ) @@ -640,11 +596,6 @@ def build_report( num_envs = result.get("run", {}).get("num_envs") if attempts and num_envs: stop_attempts = attempts_before_stop(attempts, num_envs) - recorded_source = result.get("run", {}).get("source", {}) - assert source is None or recorded_source.get("sha256") in ( - None, - file_sha256(source), - ), f"--source {source} is not the dataset this run generated from ({recorded_source.get('path')})." annotation = None if annotation_summary: with open(annotation_summary, encoding="utf-8") as handle: @@ -655,7 +606,7 @@ def build_report( kpis = { "success_rate": _success_rate(result, stop_attempts), "demos_per_hour": _demos_per_hour(result), - "demos_per_human_demo": _demos_per_human_demo(result, source_episodes), + "demos_per_human_demo": _demos_per_human_demo(result), "annotation_yield": _annotation_yield(annotation), "tracking_error": _tracking_error(episodes, source_episodes), "action_saturation": _action_saturation( @@ -670,13 +621,7 @@ def build_report( "planner_failure_types": _planner_failure_types(attempts), "planning_time": _planning_time(result, attempts), } - return { - "schema_version": REPORT_SCHEMA_VERSION, - "dataset": dataset, - "num_episodes": len(episodes), - "num_source_episodes": len(source_episodes) if source_episodes is not None else None, - "kpis": kpis, - } + return {"dataset": dataset, "num_episodes": len(episodes), "kpis": kpis} def _format_value(value: Any) -> str: diff --git a/autodata_utils/generation_result.py b/autodata_utils/generation_result.py index 4ea98df..65f545d 100644 --- a/autodata_utils/generation_result.py +++ b/autodata_utils/generation_result.py @@ -6,7 +6,6 @@ from __future__ import annotations import contextlib -import hashlib import json import os import tempfile @@ -73,15 +72,6 @@ def write_generation_result( os.unlink(temporary_file_name) -def file_sha256(path: str) -> str: - """Hex SHA-256 digest of a file, recorded so a report can tell which source a run used.""" - digest = hashlib.sha256() - with open(path, "rb") as handle: - for chunk in iter(lambda: handle.read(1 << 20), b""): - digest.update(chunk) - return digest.hexdigest() - - 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")) diff --git a/docs/pages/advanced/motion_planners.rst b/docs/pages/advanced/motion_planners.rst index 6ebe494..c06038c 100644 --- a/docs/pages/advanced/motion_planners.rst +++ b/docs/pages/advanced/motion_planners.rst @@ -96,9 +96,8 @@ Visualization and Debugging The cuRobo backend can visualize its collision-sphere model and planned trajectories via `rerun `_ (``visualize_spheres`` / ``visualize_plan`` config flags). During multi-env generation these are enabled only for env 0 to keep the simulation -responsive. After each ``update_world_and_plan_motion()`` call, the cuRobo backend's -``last_plan_summary`` holds the outcome, the cuRobo status, and wall-clock and solver times. -SkillGen copies it into each attempt's ``planner_calls`` (see :doc:`../references/kpis`). +responsive. Planner diagnostics (success rates, timing) are available through +``get_planner_info()``. Adding a New Backend -------------------- diff --git a/docs/pages/concepts/data_generator.rst b/docs/pages/concepts/data_generator.rst index ce7bb5b..9127a5f 100644 --- a/docs/pages/concepts/data_generator.rst +++ b/docs/pages/concepts/data_generator.rst @@ -124,9 +124,8 @@ Results and Statistics Three layers report what happened: * **Per attempt** — ``GenerationResult``: ``success``, ``failure_cause``, ``failed_subtask``, - ``num_steps``, ``planner_calls`` and ``initial_state``. With - ``--result_file``, each completed attempt is also written as one line of - ``.attempts.jsonl``. + ``planner_calls`` 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 @@ -149,9 +148,8 @@ Three layers report what happened: ``env_profile`` appears only when an :doc:`environment profile ` was applied — the generated dataset itself records only the base env id, so the sidecar is - what documents that the scene was modified. A ``run`` block records how the run was - configured and how long it took. :doc:`../references/kpis` describes it and the report - computed from these files. + what documents that the scene was modified. The ``run`` block is described in + :doc:`../references/kpis`. The GenerationAlgorithm Plug-In ------------------------------- diff --git a/docs/pages/references/kpis.rst b/docs/pages/references/kpis.rst index 5c42529..7f776f9 100644 --- a/docs/pages/references/kpis.rst +++ b/docs/pages/references/kpis.rst @@ -1,45 +1,24 @@ KPIs ==== -The KPI report describes what a generated dataset delivers. It answers four questions: - -- how many usable demonstrations a run yields, and at what cost; -- whether they make good training data; -- whether they cover the task's variation; -- why attempts fail. - -``scripts/validate_dataset.py --report`` computes the report from files that generation and -annotation already write. It needs no simulator. - -Producing the Inputs --------------------- - -Every generated dataset carries two per-step entries under ``obs/datagen_info``: - -- ``intended_eef_pose``: the noise-free waypoint pose commanded at that step, for end effectors whose - action commands a pose (tracking-only channels have none); -- ``eef_pose``: the achieved control-link pose, read before that step. - -Pass ``--result_file`` to ``generate_dataset.py`` to also write: - -- The **result JSON**, with a ``run`` block: seed, stop rule, device, env count, control step - (``step_dt``), setup and generation time, source dataset path and digest, and the reset events - that sample object poses from a ``pose_range``. -- The **attempt log**, ``.attempts.jsonl``. It holds one JSON line per completed - attempt: env id, start and end time, outcome, ``failure_cause`` (``planner`` or ``execution``), the - failed subtask per end effector, step count, initial object poses, and one summary per - motion-planner call. - -Two inputs come from elsewhere: - -- Pass ``--summary_file`` to ``annotate_demos.py`` to record annotation yield. -- Pass the annotated source dataset as ``--source`` to enable the source baselines. Its - ``obs/datagen_info/eef_pose`` is in the same control-link frame as the generated data's. - -The embodiment YAML is needed only for ``action_saturation``. - -Running the Report ------------------- +``scripts/validate_dataset.py --report`` computes KPIs of a generated dataset from files that +generation and annotation write. It needs no simulator. + +Inputs +------ + +- **Generated dataset.** Every episode records, per step under ``obs/datagen_info``, the + noise-free waypoint pose commanded (``intended_eef_pose``, only for end effectors whose action + commands a pose) and the control-link pose read before the step (``eef_pose``). +- **Result JSON and attempt log.** ``generate_dataset.py --result_file`` adds a ``run`` block to + the result JSON (device, GPU and CPU model, env count, control step, generation time, source + episode count, reset events that sample object poses from a ``pose_range``) and writes + ``.attempts.jsonl``, one line per completed attempt: env id, start and end + time, outcome, ``failure_cause`` (``planner`` or ``execution``), failed subtask per end + effector, initial object poses, and one summary per motion-planner call. +- **Source dataset** (``--source``): enables the source baselines. +- **Embodiment YAML** (``--embodiment``): needed for ``action_saturation``. +- **Annotation summary** (``annotate_demos.py --summary_file``): needed for ``annotation_yield``. .. code-block:: bash @@ -49,94 +28,59 @@ Running the Report --embodiment autodata_examples/embodiments/franka_ik_rel.yaml \ --report_file datasets/generated_kpis.json -The command validates the dataset, prints a KPI table, and optionally writes the full report as -JSON. A dataset left empty because every attempt failed is still reported. The report stops with an -error in three cases: - -- the dataset's demonstration count differs from the result file's success count; -- the attempt log's line count differs from the result file's attempt count; -- ``--source`` is not the dataset the run generated from. +A dataset left empty because every attempt failed is still reported. A KPI that cannot be +computed reports ``not_applicable`` with a reason. -KPI Definitions ---------------- +Definitions +----------- .. list-table:: - :widths: 26 52 22 + :widths: 26 74 :header-rows: 1 * - KPI - Definition - - Needs * - ``success_rate`` - Successful attempts over attempts, with a Wilson 95% interval. An attempt succeeds if the - task is successful at any step. ``raw_rate`` is the uncorrected rate from the result file. - - Result JSON, attempt log + task is successful at any step. * - ``demos_per_hour`` - - Successful demonstrations per hour of generation time. One run uses one GPU, so this is - also demonstrations per GPU-hour. Setup time is reported separately. - - Result JSON + - Successful demonstrations per hour of generation time, labelled with the hardware (the GPU, + or the CPU for ``--device cpu``) and env count. Compare only on the same hardware and env count. * - ``demos_per_human_demo`` - - Successful demonstrations over source demonstrations. When the run stops at a success - count, this is the requested count divided by the source count. - - Result JSON + - Successful demonstrations over source demonstrations. * - ``annotation_yield`` - - Exported over processed episodes, with rejection reasons and manual re-marks. - - Annotation summary + - Exported over processed episodes, with rejection reasons. * - ``tracking_error`` - - p50 and p95 distance [mm] from each achieved end-effector position to the noise-free - intended path. Position only. - - Generated dataset + - p50 and p95 distance [mm] from each achieved end-effector position to the intended path. + Generated values include the injected action noise. * - ``action_saturation`` - Share of pose-action entries at ±1, for adapters that clip pose actions. - - Embodiment YAML * - ``motion_smoothness`` - p50 and p95 end-effector linear jerk [m/s³]. - - ``step_dt`` * - ``initial_state_coverage`` - Per object, grid cells occupied by successful starts over cells occupied by as many - attempted starts. The worst object is reported. - - Attempt log, reset ranges + attempted starts, on a grid of about five successes per cell. The worst object is reported. * - ``success_by_initial_configuration`` - - Number of cells whose success rate is significantly below the overall success rate. Each - object's reset range has three cells per varying axis (3 for one axis, 9 for two). Only - cells with at least 10 attempts count. A cell is flagged when the exact binomial - probability of so few successes at the overall rate is below 0.025 divided by the number - of counted cells. ``per_asset`` holds every cell's counts. - - Attempt log, reset ranges + - Number of reset-range cells (three per varying axis, at least 10 attempts each) whose + success count is improbably low at the overall rate: one-sided exact binomial test at + 0.025 divided by the number of cells. * - ``failure_share`` - - Share of failed attempts by cause, and counts by end effector and subtask. - - Attempt log + - Share of failed attempts by cause, with counts by end effector and subtask. Execution + failures are attributed to a subtask only when the env publishes every non-final subtask's + termination signal in its ``subtask_terms`` observation group. * - ``planner_failure_types`` - Share of failed motion plans by cuRobo status. - - Attempt log (SkillGen) * - ``planning_time`` - p50 and p95 planning wall-clock time per attempt [s], and its share of generation time. - Planning blocks the event loop, so this share is time no env steps. - - Attempt log (SkillGen) -Reading the Values ------------------- +Quality KPIs carry a ``source_baseline`` computed the same way on the source demonstrations. -- **Source baselines.** Quality KPIs carry a ``source_baseline`` computed on the source - demonstrations. Compare against it rather than reading absolute values. -- **Tracking error includes noise.** Generated values include the injected action noise by design. - The source baseline is measured against the executed teleoperation commands. -- **Attempts cancelled at the stop.** When generation stops, attempts still running are cancelled - and never logged. Long attempts are the likelier to be running, so a raw rate under-counts long - outcomes. ``success_rate``, ``failure_share`` and the configuration KPIs count only attempts - that started before the first env finished its last logged attempt; all of those completed. The - correction needs at least three logged attempts from every env, since with fewer it is biased - too. Otherwise they use all attempts and mark the value uncorrected. -- **Not applicable.** A KPI that cannot be computed reports ``not_applicable`` with a reason. - Missing inputs never read as zero. -- **Coverage grid.** The grid is sized to about five successful starts per cell, with at least two - bins per varying axis. An object with fewer successful starts is skipped. -- **Range frames.** Ranges of ``randomize_object_pose`` are absolute in the env frame. Ranges of - ``reset_root_state_uniform`` are offsets from the asset's default position. The configuration - KPIs skip, and list the reason for, other reset functions, ranges without x or y variation, and - objects whose starts fall outside their range. -- **Failed subtasks.** Execution failures are attributed only when the env publishes the - termination signal of every non-final subtask in its ``subtask_terms`` observation group. +When generation stops, attempts still running are cancelled and never logged, and long attempts +are the likelier to be running. ``success_rate``, ``failure_share`` and the configuration KPIs +therefore count only attempts that started before the first env finished its last logged +attempt. This needs at least three logged attempts from every env; otherwise they use all attempts +and say the value is uncorrected. -The report's JSON layout is versioned by ``schema_version``. The functions live in -``autodata_utils/dataset_metrics.py``. +Ranges of ``randomize_object_pose`` are absolute in the env frame; ranges of +``reset_root_state_uniform`` are offsets from the asset's default position. The configuration KPIs +skip other reset functions and say why. diff --git a/scripts/annotate_demos.py b/scripts/annotate_demos.py index c972e98..06e2488 100644 --- a/scripts/annotate_demos.py +++ b/scripts/annotate_demos.py @@ -25,9 +25,6 @@ Pass ``--auto`` to annotate subtask termination signals without keyboard input (``--viz none`` supported). ``--signal_obs_group`` selects the observation group holding the per-subtask boolean terms (default ``subtask_terms``). SkillGen start signals still require manual mode. - -Pass ``--summary_file`` to save the processed and exported episode counts, the rejection reasons and -the number of manual re-marks as JSON. """ """Launch Isaac Sim Simulator first.""" @@ -138,7 +135,6 @@ marked_subtask_action_indices: list[int] = [] skip_episode = False rejection_counts: Counter[str] = Counter() -manual_remark_count = 0 _datastream: Datastream | None = None @@ -357,11 +353,9 @@ def main() -> int: ) if args_cli.summary_file: summary = { - "mode": "auto" if args_cli.auto else "manual", "num_processed": processed_episode_count, "num_exported": exported_episode_count, "rejections": dict(rejection_counts), - "manual_remarks": manual_remark_count, } with open(args_cli.summary_file, "w", encoding="utf-8") as summary_handle: json.dump(summary, summary_handle, indent=2, sort_keys=True) @@ -451,7 +445,7 @@ def annotate_episode_in_manual_mode( Returns: True if the episode was fully annotated, False if it was skipped or never succeeded. """ - global is_paused, marked_subtask_action_indices, skip_episode, manual_remark_count + global is_paused, marked_subtask_action_indices, skip_episode term_signal_action_indices: dict[str, int] = {} start_signal_action_indices: dict[str, int] = {} @@ -505,14 +499,13 @@ def annotate_episode_in_manual_mode( if not task_succeeded: print("\tThe final task was not completed.") - rejection_counts["task_not_completed"] += 1 + rejection_counts["task not completed"] += 1 return False print( f"\tOnly {len(marked_subtask_action_indices)} out of {expected_count} subtask signals" f' for eef "{eef_name}" were annotated.' ) - manual_remark_count += 1 print(f'\tThe episode will be replayed again for re-marking subtask signals for eef "{eef_name}".\n') num_steps = len(episode.data["actions"]) @@ -607,7 +600,7 @@ def sample_signals(action_index: int) -> None: task_succeeded = replay_episode(env, episode, success_term, on_step=sample_signals) if not task_succeeded: print("\tThe final task was not completed.") - rejection_counts["task_not_completed"] += 1 + rejection_counts["task not completed"] += 1 return False num_steps = len(episode.data["actions"]) @@ -619,22 +612,21 @@ def sample_signals(action_index: int) -> None: signal = torch.stack([step[signal_name][0] for step in sampled_signals]) edge = _first_rising_edge(signal) if edge is None: - active_at_start = bool(signal[0]) - reason = "already active at episode start" if active_at_start else "never fired" + reason = "already active at episode start" if bool(signal[0]) else "never fired" print(f'\tSubtask signal "{signal_name}" (eef "{eef_name}") {reason}.') - rejection_counts["signal_active_at_start" if active_at_start else "signal_never_fired"] += 1 + rejection_counts[f"signal {reason}"] += 1 return False eef_edges.append(edge) if any(later <= earlier for earlier, later in zip(eef_edges, eef_edges[1:])): print(f'\tSubtask signals for eef "{eef_name}" fired out of order: {dict(zip(eef_term_names, eef_edges))}.') - rejection_counts["signals_out_of_order"] += 1 + rejection_counts["signals out of order"] += 1 return False spacing_error = _validate_boundary_spacing(task_descriptor, eef_name, eef_edges, num_steps) if spacing_error is not None: print(f'\tRejected annotations for eef "{eef_name}": {spacing_error}.') - rejection_counts["boundaries_too_close"] += 1 + rejection_counts["boundaries too close"] += 1 return False print(f'\tDetected subtask signals for eef "{eef_name}": {dict(zip(eef_term_names, eef_edges))}') diff --git a/scripts/generate_dataset.py b/scripts/generate_dataset.py index 6f90b49..05b2fd7 100755 --- a/scripts/generate_dataset.py +++ b/scripts/generate_dataset.py @@ -33,7 +33,6 @@ """Launch Isaac Sim Simulator first.""" import argparse -import time from isaaclab.app import AppLauncher @@ -43,9 +42,6 @@ from autodata_interfaces.tasks.task_descriptor_utils import resolve_task_algorithm from autodata_utils.cli import callable_path -# Taken before the simulator launches, so the recorded setup time includes its startup. -PROCESS_START_TIME = 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"] @@ -144,8 +140,10 @@ 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 time # noqa: E402 import torch # noqa: E402 import traceback # noqa: E402 from pathlib import Path # noqa: E402 @@ -164,12 +162,7 @@ ) from autodata_interfaces.env.registration import load_callable, register_external_environment # noqa: E402 from autodata_utils import pose_math # noqa: E402 -from autodata_utils.generation_result import ( # noqa: E402 - AttemptLog, - attempts_file_for, - file_sha256, - write_generation_result, -) +from autodata_utils.generation_result import AttemptLog, attempts_file_for, write_generation_result # noqa: E402 from autodata_utils.tensor_utils import as_torch # noqa: E402 @@ -182,14 +175,11 @@ async def run_data_generator( success_term, pause_subtask: bool, stats: dict, - timing: dict, attempt_log: AttemptLog | None, ) -> None: """Repeatedly call ``data_generator.generate``, tally outcomes into ``stats`` and log each attempt.""" while True: attempt_start = time.monotonic() - if timing["generation_start"] is None: - timing["generation_start"] = attempt_start try: result = await data_generator.generate( env_id=env_id, @@ -213,15 +203,14 @@ async def run_data_generator( def _attempt_record(env_id: int, result: GenerationResult, start_time: float, end_time: float) -> dict: - """One attempt-log line; times [s] are measured from process start.""" + """One attempt-log line; times [s] are monotonic-clock readings.""" return { "env_id": env_id, "success": result.success, "failure_cause": result.failure_cause, "failed_subtask": result.failed_subtask, - "num_steps": result.num_steps, - "start_time_s": start_time - PROCESS_START_TIME, - "end_time_s": end_time - PROCESS_START_TIME, + "start_time_s": start_time, + "end_time_s": end_time, "initial_object_poses": _initial_object_poses(result.initial_state, env_id), "planner_calls": list(result.planner_calls), } @@ -238,6 +227,20 @@ def _initial_object_poses(initial_state: dict, env_id: int) -> dict[str, list[fl return poses +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, "cpu": cpu or platform.processor() or None} + + def _reset_ranges(env_cfg: Any) -> list[dict]: """Reset event terms that sample object poses from a ``pose_range``. @@ -245,7 +248,7 @@ def _reset_ranges(env_cfg: Any) -> list[dict]: """ ranges = [] events = getattr(env_cfg, "events", None) - for term_name, term in (vars(events) if events is not None else {}).items(): + for term in (vars(events) if events is not None else {}).values(): params = getattr(term, "params", None) if getattr(term, "mode", None) != "reset" or not params or not isinstance(params.get("pose_range"), dict): continue @@ -262,7 +265,6 @@ def _reset_ranges(env_cfg: Any) -> list[dict]: if init_state is not None: default_positions[name] = [float(value) for value in init_state.pos] ranges.append({ - "term": term_name, "func": func, "assets": assets, "pose_range": {axis: [float(bounds[0]), float(bounds[1])] for axis, bounds in params["pose_range"].items()}, @@ -308,7 +310,6 @@ def setup_async_generation( data_generator = DataGenerator(datastream=datastream, algorithm=algorithm) stats = {"num_success": 0, "num_failures": 0, "num_attempts": 0} - timing = {"generation_start": None} tasks = [] for env_id in range(num_envs): @@ -322,7 +323,6 @@ def setup_async_generation( success_term=success_term, pause_subtask=pause_subtask, stats=stats, - timing=timing, attempt_log=attempt_log, ) ) @@ -335,7 +335,6 @@ def setup_async_generation( "action_queue": env_action_queue, "info_pool": datastream.source_pool, "stats": stats, - "timing": timing, } @@ -463,7 +462,7 @@ def main() -> None: data_gen_tasks = asyncio.ensure_future(asyncio.gather(*async_components["tasks"])) generation_completed = False - generation_end = None + generation_start = time.monotonic() try: generation_completed = env_loop( env, @@ -488,7 +487,6 @@ def main() -> None: _close_motion_planners(motion_planners) if generation_completed and args_cli.result_file: - generation_start = async_components["timing"]["generation_start"] write_generation_result( result_file=args_cli.result_file, algorithm=args_cli.alg, @@ -500,18 +498,12 @@ def main() -> None: else None ), run_metadata={ - "seed": generation_policy_params.seed, - "stop_rule": "successes" if generation_policy_params.guarantee_success else "attempts", "device": args_cli.device, + **_hardware(args_cli.device), "num_envs": args_cli.num_envs, "step_dt": float(env.step_dt), - "setup_time_s": generation_start - PROCESS_START_TIME, "generation_time_s": generation_end - generation_start, - "source": { - "path": args_cli.input_file, - "sha256": file_sha256(args_cli.input_file), - "num_episodes": datastream.num_source_demos, - }, + "num_source_episodes": datastream.num_source_demos, "reset_ranges": _reset_ranges(env_cfg), "attempts_file": Path(attempts_file).name, }, diff --git a/scripts/validate_dataset.py b/scripts/validate_dataset.py index 253b016..39a9b03 100644 --- a/scripts/validate_dataset.py +++ b/scripts/validate_dataset.py @@ -1,17 +1,15 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Structurally validate AutoData HDF5 datasets and report the KPIs of a generated one. +"""Structurally validate AutoData HDF5 datasets. -Prints a summary table (episode count, env id, sim args) per file followed by any issues. With -``--report``, also prints the KPIs of a single generated dataset. +Prints a summary table (episode count, env id, sim args) per file followed by any issues. Usage:: python scripts/validate_dataset.py [ ...] python scripts/validate_dataset.py directory/*.hdf5 - python scripts/validate_dataset.py --report --result_file \\ - --source --embodiment [--report_file ] + python scripts/validate_dataset.py --report --result_file # KPIs The command exits nonzero if any supplied file is invalid. """ diff --git a/skills/autodata-generate-datasets/references/cli-reference.md b/skills/autodata-generate-datasets/references/cli-reference.md index a52ac13..47ec1a1 100644 --- a/skills/autodata-generate-datasets/references/cli-reference.md +++ b/skills/autodata-generate-datasets/references/cli-reference.md @@ -66,6 +66,5 @@ Add `--report` to print the KPIs of one generated dataset (see the KPIs docs pag ```bash python scripts/validate_dataset.py --report --result_file \ - --source --embodiment [--annotation_summary ] \ - [--report_file ] + --source --embodiment ``` From 1b55c3908976e85102ed0e865d8ca7cab69f514b Mon Sep 17 00:00:00 2001 From: Billy Okal Date: Mon, 19 Oct 2026 13:51:51 -0500 Subject: [PATCH 3/4] refactor: record KPI inputs with the shared datagen-info recorder Generation records obs/datagen_info with the recorder term annotation already uses, so tracking error compares achieved and commanded poses for generated and source data alike. Coverage takes its bounds from the attempted starts, failure share reports the cause only, planning time comes from a timer around the planner call, and annotation yield from episode counts. This drops the cuRobo, Datastream, adapter and annotation changes and the planner failure-type KPI. --- autodata_core/algorithms.py | 10 +- autodata_core/data_generator.py | 64 +---- autodata_interfaces/datastream/datastream.py | 37 +-- .../absolute_pose_embodiment_adapter.py | 3 - .../embodiments/embodiment_adapter.py | 8 - .../env/datagen_info_recorder.py | 49 ++++ .../motion_planners/curobo/curobo_planner.py | 15 +- autodata_tests/core/test_data_generator.py | 84 ++---- autodata_tests/e2e/test_annotation.py | 6 - .../e2e/test_mimicgen_data_generation.py | 5 +- .../e2e/test_skillgen_bin_data_generation.py | 11 +- .../interfaces/datastream/test_datastream.py | 14 - .../test_absolute_pose_embodiment_adapter.py | 4 - autodata_tests/utils/test_dataset_metrics.py | 55 ++-- autodata_tests/utils/test_validate_dataset.py | 7 +- autodata_tests/utils/utils.py | 9 - autodata_utils/dataset_metrics.py | 271 ++++++------------ docs/pages/concepts/data_generator.rst | 6 +- docs/pages/references/kpis.rst | 47 ++- scripts/annotate_demos.py | 65 +---- scripts/generate_dataset.py | 63 +--- scripts/validate_dataset.py | 6 +- .../references/cli-reference.md | 1 - 23 files changed, 250 insertions(+), 590 deletions(-) create mode 100644 autodata_interfaces/env/datagen_info_recorder.py diff --git a/autodata_core/algorithms.py b/autodata_core/algorithms.py index cb03a15..0bc4dca 100644 --- a/autodata_core/algorithms.py +++ b/autodata_core/algorithms.py @@ -15,6 +15,7 @@ from __future__ import annotations +import time from collections.abc import Iterator from typing import TYPE_CHECKING, Any @@ -185,7 +186,7 @@ def __init__(self, motion_planners: dict[int, Any]) -> None: motion_planners: One planner per env_id. Planners must expose ``update_world_and_plan_motion(target_pose, expected_attached_object, env_id, ...)`` and ``get_planned_poses() -> list[torch.Tensor[4, 4]]``. ``config.motion_noise_scale`` - and ``last_plan_summary`` (copied into the attempt's ``planner_calls``) are read if present. + is read if present. """ assert motion_planners, "SkillGen requires at least one motion planner" self.motion_planners = motion_planners @@ -238,6 +239,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,11 +247,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), ) - plan_summary = getattr(planner, "last_plan_summary", None) - if plan_summary is not None: - eef_state.planner_calls.append( - {"eef": eef_name, "subtask": eef_state.current_subtask_index, **plan_summary} - ) + 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 6456bd0..621a82b 100644 --- a/autodata_core/data_generator.py +++ b/autodata_core/data_generator.py @@ -52,19 +52,15 @@ 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 planning aborted the attempt, ``"execution"`` when every - subtask ran but the task never succeeded, None on success. - failed_subtask: Per-EEF index of the subtask where the attempt failed. For a planner failure, - the subtask being planned; for an execution failure, the first subtask whose termination - signal never fired (the last subtask if all fired). None when unknown. - planner_calls: One summary per motion-planner call made during the attempt. + failure_cause: ``"planner"`` when motion planning aborted the attempt, otherwise ``"execution"`` + for a failed attempt and None on success. + planning_time_s: Wall-clock time spent in motion planning [s]. """ initial_state: dict success: bool failure_cause: Literal["planner", "execution"] | None = None - failed_subtask: dict[str, int] | None = None - planner_calls: tuple[dict, ...] = () + planning_time_s: float = 0.0 @dataclass @@ -90,7 +86,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 - planner_calls: list[dict] = field(default_factory=list) + planning_time_s: float = 0.0 class DataGenerator: @@ -122,8 +118,6 @@ def __init__( self.algorithm = algorithm self.src_demo_datagen_info_pool = datastream.source_pool self.demo_keys = demo_keys - # Resolved on the first attempt, once the env has produced observations. - self._track_term_signals: bool | None = None self._validate_eef_count_for_algorithm() self._validate_coordination_for_algorithm() @@ -554,9 +548,6 @@ async def generate( eef_states = self._initialize_eef_states() selected_src_demo_inds: dict[str, int | None] = {name: None for name in self.datastream.get_eef_names()} buffers = _GenerationBuffers() - if self._track_term_signals is None: - self._track_term_signals = self.datastream.has_subtask_term_signals() - fired_term_signals: set[str] = set() randomized_subtask_boundaries: dict[str, np.ndarray] | None = None prev_pool_size = 0 @@ -593,8 +584,7 @@ async def generate( initial_state=initial_state, success=False, failure_cause="planner", - failed_subtask={eef_name: eef_state.current_subtask_index}, - planner_calls=self._collect_planner_calls(eef_states), + planning_time_s=sum(state.planning_time_s for state in eef_states.values()), ) next_trajectory, _is_motion_plan = result eef_state.current_trajectory = next_trajectory @@ -605,7 +595,6 @@ async def generate( eef_states=eef_states, ) multi_waypoint = MultiWaypoint(eef_waypoints) - self._record_step_poses(env_id, eef_waypoints) await asyncio.sleep(0) @@ -615,10 +604,6 @@ async def generate( env_id=env_id, env_action_queue=env_action_queue, ) - if self._track_term_signals: - signals = self.datastream.get_subtask_term_signals(env_ids=[env_id]) - fired_now = torch.stack([value[0] for value in signals.values()]).tolist() - fired_term_signals.update(name for name, fired in zip(signals, fired_now) if fired) buffers.success = buffers.success or exec_success if generation_policy.use_navigation_controller: hold, was_navigating = self._synchronize_navigation(eef_states, env_id, was_navigating) @@ -648,13 +633,11 @@ async def generate( if export_demo: env.recorder_manager.export_episodes(env_id_tensor) - failed = not buffers.success return GenerationResult( initial_state=initial_state, success=buffers.success, - failure_cause="execution" if failed else None, - failed_subtask=self._failed_subtasks(fired_term_signals) if failed and self._track_term_signals else None, - planner_calls=self._collect_planner_calls(eef_states), + failure_cause=None if buffers.success else "execution", + planning_time_s=sum(state.planning_time_s for state in eef_states.values()), ) # ------------------------------------------------------------------ @@ -710,37 +693,6 @@ def _build_runtime_subtask_constraints(self) -> dict: def _initialize_eef_states(self) -> dict[str, _EEFGenerationState]: return {name: _EEFGenerationState() for name in self.datastream.get_eef_names()} - def _record_step_poses(self, env_id: int, eef_waypoints: dict[str, Waypoint]) -> None: - """Append the noise-free waypoint pose commanded at this step, and the EEF pose before it, to the episode. - - Tracking-only channels command no pose, so they record no intended pose. - """ - commanded = self.datastream.get_eef_names_with_pose_action() - intended = {name: waypoint.pose.unsqueeze(0) for name, waypoint in eef_waypoints.items() if name in commanded} - achieved = {name: self.datastream.get_robot_eef_pose(env_ids=[env_id], eef_name=name) for name in eef_waypoints} - self.datastream.get_env().recorder_manager.add_to_episodes( - "obs/datagen_info", {"intended_eef_pose": intended, "eef_pose": achieved}, env_ids=[env_id] - ) - - def _failed_subtasks(self, fired_term_signals: set[str]) -> dict[str, int]: - """Per EEF, the first subtask whose termination signal never fired, else the last subtask.""" - failed: dict[str, int] = {} - for eef_name in self.datastream.get_eef_names(): - subtasks = self.datastream.get_subtasks(eef_name) - failed[eef_name] = next( - ( - index - for index, subtask in enumerate(subtasks) - if subtask.subtask_term_signal and subtask.subtask_term_signal not in fired_term_signals - ), - len(subtasks) - 1, - ) - return failed - - @staticmethod - def _collect_planner_calls(eef_states: dict[str, _EEFGenerationState]) -> tuple[dict, ...]: - return tuple(call for eef_state in eef_states.values() for call in eef_state.planner_calls) - def _maybe_refresh_randomized_subtask_boundaries( self, randomized_subtask_boundaries: dict[str, np.ndarray] | None, diff --git a/autodata_interfaces/datastream/datastream.py b/autodata_interfaces/datastream/datastream.py index e133354..9733fa7 100644 --- a/autodata_interfaces/datastream/datastream.py +++ b/autodata_interfaces/datastream/datastream.py @@ -191,11 +191,6 @@ def get_robot_eef_pose(self, env_ids: Sequence[int] | None, eef_name: str) -> to return self.embodiment_adapter.get_eef_poses(env_ids=env_ids)[eef_name] - def get_eef_names_with_pose_action(self) -> tuple[str, ...]: - """Return the EEFs whose target pose the action commands, excluding tracking-only channels.""" - - return self.embodiment_adapter.get_eef_names_with_pose_action() - def get_robot_joint_positions(self, env_ids: Sequence[int] | None = None) -> torch.Tensor: """Read the robot's current joint positions [rad] via the embodiment adapter. @@ -276,27 +271,12 @@ def get_object_poses(self, env_ids: Sequence[int] | None = None) -> dict[str, to object_pose_matrix[obj_name] = reference_inverse @ object_pose_matrix[obj_name] return object_pose_matrix - def _boundary_signal_names(self) -> list[str]: - """Termination-signal names of every non-final subtask; the final one ends with the trajectory.""" - - return [name for eef in self.get_eef_names() for name in self.get_term_signal_names(eef)[:-1] if name] - - def has_subtask_term_signals(self, obs_group: str = "subtask_terms") -> bool: - """Whether the env publishes every signal :meth:`get_subtask_term_signals` reads; not every env does.""" - - obs_buf = getattr(self.env, "obs_buf", None) - group = obs_buf.get(obs_group) if isinstance(obs_buf, dict) else None - if not isinstance(group, dict): - return False - signal_names = self._boundary_signal_names() - return bool(signal_names) and all(name in group for name in signal_names) - def get_subtask_term_signals( self, env_ids: Sequence[int] | None = None, obs_group: str = "subtask_terms" ) -> dict[str, torch.Tensor]: """Read the current per-subtask termination signals from the env observation buffer. - For every non-final subtask's termination-signal name declared by the task descriptor (across all + For every non-empty termination-signal name declared by the task descriptor (across all EEFs), reads the identically named boolean observation term from ``env.obs_buf[obs_group]``. The env must therefore publish one observation term per term-signal name in a non-concatenated group. Auto annotation samples this each replay @@ -322,12 +302,15 @@ def get_subtask_term_signals( "signal reads require the group to set 'concatenate_terms = False'." ) signals: dict[str, torch.Tensor] = {} - for signal_name in self._boundary_signal_names(): - assert signal_name in group, ( - f"Termination signal {signal_name!r} has no matching observation term in " - f"group {obs_group!r}. Available terms: {sorted(group.keys())}" - ) - signals[signal_name] = as_torch(group[signal_name])[index].reshape(-1).bool() + for eef_name in self.get_eef_names(): + for signal_name in self.get_term_signal_names(eef_name): + if not signal_name: + continue + assert signal_name in group, ( + f"Termination signal {signal_name!r} has no matching observation term in " + f"group {obs_group!r}. Available terms: {sorted(group.keys())}" + ) + signals[signal_name] = as_torch(group[signal_name])[index].reshape(-1).bool() return signals def get_robot_root_pose(self, env_ids: Sequence[int] | None = None) -> torch.Tensor: diff --git a/autodata_interfaces/embodiments/absolute_pose_embodiment_adapter.py b/autodata_interfaces/embodiments/absolute_pose_embodiment_adapter.py index 3e1409e..c1d6c09 100644 --- a/autodata_interfaces/embodiments/absolute_pose_embodiment_adapter.py +++ b/autodata_interfaces/embodiments/absolute_pose_embodiment_adapter.py @@ -80,9 +80,6 @@ def from_dict(cls, data: dict[str, Any]) -> AbsolutePoseEmbodimentAdapter: def get_eef_names(self) -> tuple[str, ...]: return tuple(self.eefs) - def get_eef_names_with_pose_action(self) -> tuple[str, ...]: - return tuple(name for name, spec in self.eefs.items() if spec.get("pose_action_slice") is not None) - def get_eef_poses(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]: assert self.env is not None, "Call bind_env(env) before reading state." index = slice(None) if env_ids is None else env_ids diff --git a/autodata_interfaces/embodiments/embodiment_adapter.py b/autodata_interfaces/embodiments/embodiment_adapter.py index 1694fd5..4f63034 100644 --- a/autodata_interfaces/embodiments/embodiment_adapter.py +++ b/autodata_interfaces/embodiments/embodiment_adapter.py @@ -93,14 +93,6 @@ def get_eef_names(self) -> tuple[str, ...]: Tuple of EEF names. Stable across calls; never empty. """ - def get_eef_names_with_pose_action(self) -> tuple[str, ...]: - """Return the end-effectors whose target pose the action commands. - - Tracking-only channels, whose actions carry only passthrough commands, are excluded. - """ - - return self.get_eef_names() - @abstractmethod def get_eef_poses(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]: """Read current end-effector poses [m, m, m] from the live env. diff --git a/autodata_interfaces/env/datagen_info_recorder.py b/autodata_interfaces/env/datagen_info_recorder.py new file mode 100644 index 0000000..9857fac --- /dev/null +++ b/autodata_interfaces/env/datagen_info_recorder.py @@ -0,0 +1,49 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Recorder that stores ``obs/datagen_info`` on every step of annotation and generation.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg +from isaaclab.managers import RecorderTerm, RecorderTermCfg +from isaaclab.utils import configclass + +if TYPE_CHECKING: + from autodata_interfaces.datastream import Datastream + + +class PreStepDatagenInfoRecorder(RecorderTerm): + """Records object, end-effector and commanded target poses from the Datastream 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 = { + "object_pose": self.datastream.get_object_poses(), + "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), + } + return "obs/datagen_info", datagen_info + + +@configclass +class PreStepDatagenInfoRecorderCfg(RecorderTermCfg): + """Configuration for :class:`PreStepDatagenInfoRecorder`.""" + + class_type: type[RecorderTerm] = PreStepDatagenInfoRecorder + + +@configclass +class DatagenInfoRecorderManagerCfg(ActionStateRecorderManagerCfg): + """Action/state recorder plus the per-step ``obs/datagen_info`` term.""" + + record_pre_step_datagen_info = PreStepDatagenInfoRecorderCfg() diff --git a/autodata_interfaces/motion_planners/curobo/curobo_planner.py b/autodata_interfaces/motion_planners/curobo/curobo_planner.py index f35595e..1a43077 100644 --- a/autodata_interfaces/motion_planners/curobo/curobo_planner.py +++ b/autodata_interfaces/motion_planners/curobo/curobo_planner.py @@ -5,7 +5,6 @@ import logging import numpy as np -import time import torch from dataclasses import dataclass from types import SimpleNamespace @@ -287,8 +286,6 @@ def __init__( # Initialize planning state self._current_plan: JointState | None = None self._plan_index: int = 0 - self._plan_failure_status: str | None = None - self.last_plan_summary: dict[str, Any] | None = None # Initialize visualization state self.frame_counter: int = 0 @@ -1332,11 +1329,9 @@ def _plan_to_contact_pose( self.logger.debug(f"Contact planning succeeded with {len(self._current_plan.position)} waypoints") else: self.logger.debug(f"Contact planning failed: {result.status}") - self._plan_failure_status = str(getattr(result.status, "value", result.status)) except Exception as e: self.logger.debug(f"Error during planning: {e}") - self._plan_failure_status = "exception" # Always restore sphere state after planning, regardless of success if contact: @@ -1848,8 +1843,7 @@ def update_world_and_plan_motion( Provides a high-level interface that handles the complete planning workflow: world synchronization, object attachment/detachment, gripper configuration, - and motion planning. Afterwards :attr:`last_plan_summary` holds the call's outcome, the - cuRobo status of the failed phase, and the wall-clock time [s]. + and motion planning. Args: target_pose: Target end-effector pose as 4x4 transformation matrix @@ -1861,8 +1855,6 @@ def update_world_and_plan_motion( Returns: True if complete planning pipeline succeeded, False if any step failed """ - start_time = time.perf_counter() - self._plan_failure_status = None # Always reset the plan before starting a new one to ensure a clean state self.reset_plan() @@ -1927,11 +1919,6 @@ def update_world_and_plan_motion( self._detach_objects() - self.last_plan_summary = { - "success": plan_success, - "status": "Success" if plan_success else self._plan_failure_status or "unknown", - "wall_time_s": time.perf_counter() - start_time, - } return plan_success # ===================================================================================== diff --git a/autodata_tests/core/test_data_generator.py b/autodata_tests/core/test_data_generator.py index 2fde05b..6667ca1 100644 --- a/autodata_tests/core/test_data_generator.py +++ b/autodata_tests/core/test_data_generator.py @@ -3,8 +3,8 @@ """Unit tests for :class:`autodata_core.data_generator.DataGenerator`. -Construction-time validation, navigation playback decisions and the planner-failure path of the -generation loop are covered here. Simulation integration is covered by the e2e suite. +Construction-time validation and navigation playback decisions are covered here. +Simulation integration is covered by the e2e suite. """ import asyncio @@ -15,7 +15,7 @@ import pytest -from autodata_core.algorithms import DexMimicGen, GenerationAlgorithm, MimicGen, SkillGen +from autodata_core.algorithms import DexMimicGen, MimicGen, SkillGen from autodata_core.data_generator import DataGenerator, GenerationResult, _EEFGenerationState from autodata_core.waypoint import MultiWaypoint, Waypoint from autodata_interfaces.tasks.generation_policy_spec import GenerationPolicy @@ -66,15 +66,6 @@ def get_generation_policy(self): def get_env(self): return self.env - def has_subtask_term_signals(self): - return False - - def get_eef_names_with_pose_action(self): - return tuple(self._subtasks) - - def get_robot_eef_pose(self, env_ids, eef_name): - return torch.eye(4).unsqueeze(0) - def _assert_constructed(generator: DataGenerator, datastream: _MockDatastream, algorithm) -> None: assert generator.algorithm is algorithm @@ -142,58 +133,6 @@ def test_skillgen_coordination_unsupported_error(): DataGenerator(datastream=datastream, algorithm=SkillGen(motion_planners={0: object()})) -class _PlanningFails(GenerationAlgorithm): - """Unregistered algorithm whose planner never finds a trajectory.""" - - def plan_subtask_trajectory(self, **kwargs): - return None - - -class _GenerationDatastream(_MockDatastream): - """Adds the reads ``DataGenerator.generate`` makes before the first plan.""" - - def __init__(self, subtasks_by_eef: dict) -> None: - super().__init__(subtasks_by_eef) - self.source_pool.subtask_boundaries = {eef: [[(0, 5), (5, 10)]] for eef in subtasks_by_eef} - self.env = SimpleNamespace(recorder_manager=SimpleNamespace(reset=lambda env_ids: None)) - - def get_scene_state(self, is_relative: bool = True): - return {"rigid_object": {}} - - -def test_planner_failure_records_cause_and_subtask(): - datastream = _GenerationDatastream({"franka": [_subtask("grasp_1"), _subtask("")]}) - generator = DataGenerator(datastream=datastream, algorithm=_PlanningFails()) - - async def run_attempt() -> GenerationResult: - reset_queue: asyncio.Queue = asyncio.Queue() - - async def serve_reset() -> None: - await reset_queue.get() - reset_queue.task_done() - - server = asyncio.create_task(serve_reset()) - result = await generator.generate( - env_id=0, success_term=None, env_reset_queue=reset_queue, env_action_queue=asyncio.Queue() - ) - await server - return result - - result = asyncio.run(run_attempt()) - - assert result.success is False - assert result.failure_cause == "planner" - assert result.failed_subtask == {"franka": 0} - - -def test_execution_failure_is_attributed_to_the_first_unfinished_subtask(): - datastream = _MockDatastream({"franka": [_subtask("grasp_1"), _subtask("stack_1"), _subtask("")]}) - generator = DataGenerator(datastream=datastream, algorithm=MimicGen()) - - assert generator._failed_subtasks({"grasp_1"}) == {"franka": 1} - assert generator._failed_subtasks({"grasp_1", "stack_1"}) == {"franka": 2} - - def test_generation_result_fields_and_frozen(): result = GenerationResult(initial_state={"foo": 1}, success=True) assert result.initial_state == {"foo": 1} @@ -257,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 @@ -280,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 aebe37c..f5383d3 100644 --- a/autodata_tests/e2e/test_annotation.py +++ b/autodata_tests/e2e/test_annotation.py @@ -20,7 +20,6 @@ def test_automatic_annotation_from_dataset(tmp_path): """Use the stored task name and export usable termination signals during replay.""" input_file = Path(TestPaths.test_data_dir) / "annotated_dataset_franka_stack_mimicgen.hdf5" output_file = tmp_path / "annotated.hdf5" - summary_file = tmp_path / "annotation_summary.json" run_subprocess([ TestPaths.python_path, str(Path(TestPaths.scripts_dir) / "annotate_demos.py"), @@ -37,15 +36,10 @@ def test_automatic_annotation_from_dataset(tmp_path): str(input_file), "--output_file", str(output_file), - "--summary_file", - str(summary_file), ]) assert_valid_dataset(str(output_file), min_num_demos=1) - summary = json.loads(summary_file.read_text(encoding="utf-8")) with h5py.File(input_file, "r") as source, h5py.File(output_file, "r") as annotated: assert len(annotated["data"]) == len(source["data"]) - assert summary["num_processed"] == summary["num_exported"] == len(source["data"]) - assert summary["rejections"] == {} assert json.loads(annotated["data"].attrs["env_args"])["env_name"] == "IsaacContrib-Stack-Cube-Franka-IK-Rel" for episode in annotated["data"].values(): signals = episode["obs/datagen_info/subtask_term_signals"] diff --git a/autodata_tests/e2e/test_mimicgen_data_generation.py b/autodata_tests/e2e/test_mimicgen_data_generation.py index a8cdceb..9d9babd 100644 --- a/autodata_tests/e2e/test_mimicgen_data_generation.py +++ b/autodata_tests/e2e/test_mimicgen_data_generation.py @@ -68,7 +68,7 @@ def _run_franka_cube_stack_mimicgen(num_envs: int, device: str) -> None: assert_valid_dataset(output_file, min_num_demos=GENERATION_NUM_TRIALS) with h5py.File(output_file, "r") as dataset: for episode in dataset["data"].values(): - assert len(episode["obs/datagen_info/intended_eef_pose/franka"]) == len(episode["actions"]) + assert len(episode["obs/datagen_info/target_eef_pose/franka"]) == len(episode["actions"]) # franka_ik_rel.yaml has no EEF offset, so the recorded pose must equal the pre-step observation. recorded = episode["obs/datagen_info/eef_pose/franka"][:, :3, 3] np.testing.assert_allclose(recorded, episode["obs/eef_pos"][:], atol=1e-6) @@ -90,11 +90,8 @@ def _run_franka_cube_stack_mimicgen(num_envs: int, device: str) -> None: report_file, ]) kpis = assert_kpi_report(report_file, APPLICABLE_KPIS) - assert any(term["func"].endswith(":randomize_object_pose") for term in result["run"]["reset_ranges"]) hardware = result["run"]["gpu"] if device == "cuda" else result["run"]["cpu"] assert hardware and hardware in kpis["demos_per_hour"]["note"] - for name in ("initial_state_coverage", "success_by_initial_configuration"): - assert "frame is wrong" not in json.dumps(kpis[name]), f"{name}: {kpis[name]}" @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 1088ce0..0235aba 100644 --- a/autodata_tests/e2e/test_skillgen_bin_data_generation.py +++ b/autodata_tests/e2e/test_skillgen_bin_data_generation.py @@ -17,12 +17,12 @@ from autodata_tests.utils.constants import TestPaths from autodata_tests.utils.subprocess import run_subprocess -from autodata_tests.utils.utils import assert_kpi_report, assert_valid_dataset, read_attempt_log +from autodata_tests.utils.utils import assert_kpi_report, assert_valid_dataset HEADLESS = True GENERATION_NUM_TRIALS = 1 # KPIs that a one-demo SkillGen run with a result file must produce. -APPLICABLE_KPIS = {"success_rate", "failure_share", "planner_failure_types", "planning_time"} +APPLICABLE_KPIS = {"success_rate", "failure_share", "planning_time"} def _run_franka_bin_stack_skillgen(num_envs: int, device: str) -> None: @@ -73,13 +73,6 @@ def _run_franka_bin_stack_skillgen(num_envs: int, device: str) -> None: "path": env_profile, "planner": "franka_stack_cube_bin", } - attempts = read_attempt_log(result_file) - # Every SkillGen attempt plans at least its first transit. - assert all(attempt["planner_calls"] for attempt in attempts) - for attempt in attempts: - last_call = attempt["planner_calls"][-1] - planner_failed = attempt["failure_cause"] == "planner" - assert last_call["success"] != planner_failed and (last_call["status"] == "Success") != planner_failed run_subprocess([ TestPaths.python_path, diff --git a/autodata_tests/interfaces/datastream/test_datastream.py b/autodata_tests/interfaces/datastream/test_datastream.py index 6ff07bd..56486fa 100644 --- a/autodata_tests/interfaces/datastream/test_datastream.py +++ b/autodata_tests/interfaces/datastream/test_datastream.py @@ -284,17 +284,3 @@ def test_get_robot_root_pose(): def test_get_scene_state_passthrough(): datastream, _, _, _ = _datastream() assert datastream.get_scene_state(is_relative=True) is _SCENE_STATE_SENTINEL - - -def test_subtask_signals_ignore_the_final_subtask(): - task_dict = _task_dict() - task_dict["subtasks"]["franka"][-1]["subtask_term_signal"] = "stack_2" - task = TaskDescriptor.from_dict(task_dict) - adapter = _embodiment_adapter() - env = _env() - env.obs_buf["subtask_terms"] = {"grasp_1": torch.tensor([True]), "stack_1": torch.tensor([False])} - pool = DataGenInfoPool(task_descriptor=task, embodiment_adapter=adapter, device="cpu", uses_start_signals=False) - datastream = Datastream(env=env, task_descriptor=task, embodiment_adapter=adapter, source_pool=pool) - - assert datastream.has_subtask_term_signals() - assert set(datastream.get_subtask_term_signals(env_ids=[0])) == {"grasp_1", "stack_1"} diff --git a/autodata_tests/interfaces/embodiments/test_absolute_pose_embodiment_adapter.py b/autodata_tests/interfaces/embodiments/test_absolute_pose_embodiment_adapter.py index 8e54941..b58fff6 100644 --- a/autodata_tests/interfaces/embodiments/test_absolute_pose_embodiment_adapter.py +++ b/autodata_tests/interfaces/embodiments/test_absolute_pose_embodiment_adapter.py @@ -40,10 +40,6 @@ def test_pose_conversion_preserves_scalar_grippers_and_body_commands(): assert all(torch.equal(noisy_channels[name], command[i]) for name, command in channels.items()) -def test_tracking_only_channel_commands_no_pose(): - assert _adapter().get_eef_names_with_pose_action() == ("left", "right") - - def test_controller_state_queries_select_requested_environment(): adapter = _adapter() state = torch.zeros(2, 2, 13) diff --git a/autodata_tests/utils/test_dataset_metrics.py b/autodata_tests/utils/test_dataset_metrics.py index d92ac63..917777d 100644 --- a/autodata_tests/utils/test_dataset_metrics.py +++ b/autodata_tests/utils/test_dataset_metrics.py @@ -32,16 +32,18 @@ def test_path_deviation_ignores_lag_along_the_path(): np.testing.assert_allclose(dataset_metrics.path_deviation(lagging, intended), 0.005, atol=1e-9) -def test_tracking_error_baseline_skips_channels_without_a_generated_intended_pose(): +def test_tracking_error_skips_tracking_only_channels(): path = np.linspace(0.0, 0.1, 30)[:, None] * np.array([1.0, 0.0, 0.0]) beside = path + np.array([0.0, 0.002, 0.0]) - generated = [{"eef_pose": {"left": beside}, "intended_eef_pose": {"left": path}}] - # A tracking-only channel's source target is a placeholder far from its achieved pose. - source = [{"eef_pose": {"left": beside, "body": path}, "target_eef_pose": {"left": path, "body": path + 1.0}}] + # A tracking-only channel commands no pose; its adapter writes an identity placeholder. + episode = { + "eef_pose": {"left": beside, "body": path}, + "target_eef_pose": {"left": path, "body": np.zeros_like(path)}, + } - kpi = dataset_metrics._tracking_error(generated, source) + kpi = dataset_metrics._tracking_error([episode], [episode]) - assert kpi["source_baseline"]["p95"] == pytest.approx(2.0) + assert kpi["value"]["p95"] == pytest.approx(2.0) and kpi["source_baseline"]["p95"] == pytest.approx(2.0) def _attempt(env_id: int, success: bool, start: float, duration: float) -> dict: @@ -71,41 +73,19 @@ def test_coverage_ratio_drops_when_successes_leave_a_region_empty(): def _attempts(x_positions: np.ndarray, successes: np.ndarray) -> list[dict]: + # The table only settles by about a millimetre, so it does not count as a varying object. return [ - {"success": bool(success), "initial_object_poses": {"cube": [float(x), 0.0, 0.0]}} - for x, success in zip(x_positions, successes) + { + "success": bool(success), + "initial_object_positions": {"cube": [float(x), 0.0], "table": [1e-3 * (i % 2), 0.0]}, + } + for i, (x, success) in enumerate(zip(x_positions, successes)) ] -def test_range_points_resolve_offset_and_absolute_ranges_and_reject_wrong_frames(): - attempts = _attempts(np.array([0.39, 0.41]), np.array([True, True])) - offset = { - "func": "isaaclab.envs.mdp.events:reset_root_state_uniform", - "pose_range": {"x": [-0.01, 0.01], "y": [0.0, 0.0]}, - "default_positions": {"cube": [0.4, 0.0, 0.0]}, - } - absolute = {**offset, "func": "isaaclab_tasks.contrib.stack.mdp.stack_events:randomize_object_pose"} - - points, bounds, axes = dataset_metrics._range_points(offset, "cube", attempts) - - np.testing.assert_allclose(points[:, 0], [-0.01, 0.01], atol=1e-9) - assert axes == ["x"] and bounds == [(-0.01, 0.01)] - assert "frame is wrong" in dataset_metrics._range_points(absolute, "cube", attempts) - - def test_success_by_configuration_flags_only_a_truly_weak_region(): rng = np.random.default_rng(2) x_positions = rng.uniform(0.0, 0.3, size=600) - result = { - "run": { - "reset_ranges": [{ - "func": "m:randomize_object_pose", - "assets": ["cube"], - "pose_range": {"x": [0.0, 0.3]}, - "default_positions": {}, - }] - } - } uniform = _attempts(x_positions, rng.uniform(size=600) < 0.7) weak_left = _attempts(x_positions, (rng.uniform(size=600) < 0.7) & (x_positions > 0.1)) # 9 of 11 at an overall rate of 65/67: unremarkable by an exact test, though a Wilson bound reads it as low. @@ -113,6 +93,7 @@ def test_success_by_configuration_flags_only_a_truly_weak_region(): np.array([0.05] * 11 + [0.15] * 28 + [0.25] * 28), np.array([True] * 9 + [False] * 2 + [True] * 56) ) - assert dataset_metrics._success_by_configuration(result, uniform, corrected=False)["value"] == 0 - assert dataset_metrics._success_by_configuration(result, weak_left, corrected=False)["value"] == 1 - assert dataset_metrics._success_by_configuration(result, near_perfect, corrected=False)["value"] == 0 + uniform_kpi = dataset_metrics._success_by_configuration(uniform, corrected=False) + assert uniform_kpi["value"] == 0 and set(uniform_kpi["per_asset"]) == {"cube"} + assert dataset_metrics._success_by_configuration(weak_left, corrected=False)["value"] == 1 + assert dataset_metrics._success_by_configuration(near_perfect, corrected=False)["value"] == 0 diff --git a/autodata_tests/utils/test_validate_dataset.py b/autodata_tests/utils/test_validate_dataset.py index cf351d9..cac35fc 100644 --- a/autodata_tests/utils/test_validate_dataset.py +++ b/autodata_tests/utils/test_validate_dataset.py @@ -62,9 +62,12 @@ def test_report_covers_a_run_whose_every_attempt_failed(tmp_path: Path) -> None: "run": {"num_envs": 1, "attempts_file": "result.attempts.jsonl"}, }) ) - attempt = {"env_id": 0, "success": False, "failure_cause": "execution", "failed_subtask": None, "planner_calls": []} + attempt = {"env_id": 0, "success": False, "failure_cause": "execution", "planning_time_s": 0.0} (tmp_path / "result.attempts.jsonl").write_text( - "".join(json.dumps({**attempt, "start_time_s": t, "end_time_s": t + 1.0}) + "\n" for t in range(3)) + "".join( + json.dumps({**attempt, "start_time_s": t, "end_time_s": t + 1.0, "initial_object_positions": {}}) + "\n" + for t in range(3) + ) ) report_path = tmp_path / "kpis.json" diff --git a/autodata_tests/utils/utils.py b/autodata_tests/utils/utils.py index 121580d..d293920 100644 --- a/autodata_tests/utils/utils.py +++ b/autodata_tests/utils/utils.py @@ -8,8 +8,6 @@ import math import os -from autodata_utils.generation_result import attempts_file_for - def assert_valid_dataset(output_file: str, min_num_demos: int) -> None: """Assert the run wrote a dataset containing at least min_num_demos entries.""" @@ -46,13 +44,6 @@ def read_generation_result(result_file: str) -> tuple[int, int]: return num_success, num_attempts -def read_attempt_log(result_file: str) -> list[dict]: - """Read the attempt log written next to a generation result file.""" - - with open(attempts_file_for(result_file), encoding="utf-8") as attempts_handle: - return [json.loads(line) for line in attempts_handle if line.strip()] - - def assert_kpi_report(report_file: str, applicable: set[str]) -> dict: """Assert every KPI is reported or explained, and the ``applicable`` ones carry finite values. diff --git a/autodata_utils/dataset_metrics.py b/autodata_utils/dataset_metrics.py index 422972b..0f479e9 100644 --- a/autodata_utils/dataset_metrics.py +++ b/autodata_utils/dataset_metrics.py @@ -24,13 +24,9 @@ """Below this many completed attempts in some env, the stop correction is itself biased.""" SAMPLES_PER_COVERAGE_CELL = 5 PATH_DEVIATION_WINDOW = 20 -"""Achieved steps are matched to intended path segments at most this many steps away.""" -RANGE_TOLERANCE = 1e-3 -"""Starts farther than this [m] outside their reset range mean the range frame was resolved wrongly.""" - -# Isaac Lab reset functions whose pose_range semantics are known. -_OFFSET_RANGE_FUNCS = ("reset_root_state_uniform",) -_ABSOLUTE_RANGE_FUNCS = ("randomize_object_pose",) +"""Achieved steps are matched to commanded path segments at most this many steps away.""" +MIN_START_EXTENT = 0.01 +"""Objects whose attempted starts span less than this [m] along an axis do not vary along it.""" def binomial_cdf(successes: int, trials: int, rate: float) -> float: @@ -77,18 +73,18 @@ def saturated_fraction(pose_actions: np.ndarray) -> float: return float(np.mean(np.isclose(np.abs(pose_actions), 1.0, rtol=0.0, atol=1e-6))) -def path_deviation(achieved: np.ndarray, intended: np.ndarray, window: int = PATH_DEVIATION_WINDOW) -> np.ndarray: - """Distance [m] from each achieved position to the intended path near the same step. +def path_deviation(achieved: np.ndarray, commanded: np.ndarray, window: int = PATH_DEVIATION_WINDOW) -> np.ndarray: + """Distance [m] from each achieved position to the commanded path near the same step. Args: achieved: ``[T, 3]`` achieved end-effector positions [m]. - intended: ``[T, 3]`` intended end-effector positions [m], step-aligned with ``achieved``. - window: Steps on either side whose intended segments are candidates. + commanded: ``[T, 3]`` commanded end-effector positions [m], step-aligned with ``achieved``. + window: Steps on either side whose commanded segments are candidates. """ - assert len(achieved) == len(intended), "achieved and intended positions must be step-aligned" - if len(intended) < 2: - return np.linalg.norm(achieved - intended, axis=1) - starts, ends = intended[:-1], intended[1:] + assert len(achieved) == len(commanded), "achieved and commanded positions must be step-aligned" + if len(commanded) < 2: + return np.linalg.norm(achieved - commanded, axis=1) + starts, ends = commanded[:-1], commanded[1:] deviations = np.empty(len(achieved)) for step, point in enumerate(achieved): lo, hi = max(0, step - window), min(len(starts), step + window + 1) @@ -123,10 +119,9 @@ def coverage_ratio( does not drive the ratio. Args: - success_points: ``[S, D]`` starts of successful attempts, in range coordinates, with - ``S >= min_coverage_successes(D)``. - attempt_points: ``[A, D]`` starts of all attempts, ``A >= S``. - bounds: Per-dimension ``(lo, hi)`` of the reset range. + success_points: ``[S, D]`` starts of successful attempts [m], with ``S >= min_coverage_successes(D)``. + attempt_points: ``[A, D]`` starts of all attempts [m], ``A >= S``. + bounds: Per-dimension ``(lo, hi)`` of the attempted starts [m]. seed: Seed of the random attempt subsets. """ num_success = len(success_points) @@ -158,7 +153,8 @@ def embodiment_clips_pose_actions(path: str) -> bool: def load_episodes(path: str) -> tuple[list[dict], dict]: """Per-episode actions and ``obs/datagen_info`` positions, plus the dataset's ``env_args``. - Annotation and generation both record ``eef_pose`` before each step, so source and generated data share a frame. + Annotation and generation record ``obs/datagen_info`` with the same recorder term before each step, so + source and generated data share a frame and a step. """ episodes = [] with h5py.File(path, "r") as h5_file: @@ -170,7 +166,7 @@ def load_episodes(path: str) -> tuple[list[dict], dict]: episode: dict[str, Any] = {"actions": demo["actions"][:]} datagen_info = demo["obs"].get("datagen_info") if datagen_info is not None: - for group in ("intended_eef_pose", "target_eef_pose", "eef_pose"): + for group in ("target_eef_pose", "eef_pose"): if group in datagen_info: episode[group] = {eef: datagen_info[group][eef][:][:, :3, 3] for eef in datagen_info[group]} if "eef_pose" in episode: @@ -246,14 +242,7 @@ def _success_rate(result: dict | None, stop_attempts: list[dict] | None) -> dict successes, attempts = sum(attempt["success"] for attempt in stop_attempts), len(stop_attempts) note = _CORRECTED_NOTE lo, hi = wilson_interval(successes, attempts) - return _kpi( - successes / attempts, - "fraction", - attempts, - interval_95=[lo, hi], - raw_rate=raw_rate, - note=note, - ) + return _kpi(successes / attempts, "fraction", attempts, interval_95=[lo, hi], raw_rate=raw_rate, note=note) def _demos_per_hour(result: dict | None) -> dict: @@ -283,27 +272,25 @@ def _demos_per_human_demo(result: dict | None) -> dict: return _kpi(result["num_success"] / num_source, "ratio", num_source, note=note) -def _annotation_yield(annotation_summary: dict | None) -> dict: - if annotation_summary is None: - return _not_applicable("needs --annotation_summary") - processed = annotation_summary["num_processed"] - if processed == 0: - return _not_applicable("no episodes were processed") - return _kpi( - annotation_summary["num_exported"] / processed, - "fraction", - processed, - rejections=annotation_summary["rejections"], - ) +def _annotation_yield(source_episodes: list[dict] | None, recorded_dataset: str | None) -> dict: + if source_episodes is None or recorded_dataset is None: + return _not_applicable("needs --source and --recorded_dataset") + with h5py.File(recorded_dataset, "r") as h5_file: + num_recorded = len(h5_file["data"]) + if num_recorded == 0: + return _not_applicable("the recorded dataset has no episodes") + note = "assumes --source is the annotation output of --recorded_dataset" + return _kpi(len(source_episodes) / num_recorded, "fraction", num_recorded, note=note) -def _pooled_path_deviation(episodes: list[dict], intended_key: str, eefs: set[str] | None = None) -> np.ndarray: +def _pooled_path_deviation(episodes: list[dict]) -> np.ndarray: deviations = [ - path_deviation(episode["eef_pose"][eef], episode[intended_key][eef]) + path_deviation(episode["eef_pose"][eef], target) for episode in episodes - if intended_key in episode and "eef_pose" in episode - for eef in episode[intended_key] - if eefs is None or eef in eefs + if "target_eef_pose" in episode and "eef_pose" in episode + for eef, target in episode["target_eef_pose"].items() + # Tracking-only channels command no pose; their adapter writes an identity placeholder. + if np.any(target) ] return np.concatenate(deviations) if deviations else np.empty(0) @@ -311,22 +298,16 @@ def _pooled_path_deviation(episodes: list[dict], intended_key: str, eefs: set[st def _tracking_error(episodes: list[dict], source_episodes: list[dict] | None) -> dict: if not episodes: return _not_applicable("the dataset has no episodes") - deviations = _pooled_path_deviation(episodes, "intended_eef_pose") + deviations = _pooled_path_deviation(episodes) if not len(deviations): - return _not_applicable("dataset has no obs/datagen_info/intended_eef_pose") - baseline = None - if source_episodes is not None: - # Generation records intended poses only for commanded EEFs; the source's target poses include - # placeholders for tracking-only channels. - commanded = {eef for episode in episodes for eef in episode.get("intended_eef_pose", {})} - source_deviations = _pooled_path_deviation(source_episodes, "target_eef_pose", commanded) - baseline = _percentiles(source_deviations, 1000.0) if len(source_deviations) else None + return _not_applicable("dataset has no obs/datagen_info/target_eef_pose") + source_deviations = _pooled_path_deviation(source_episodes) if source_episodes is not None else np.empty(0) + baseline = _percentiles(source_deviations, 1000.0) if len(source_deviations) else None return _kpi( _percentiles(deviations, 1000.0), - "mm from the noise-free intended path (position)", + "mm from the commanded path (position)", len(deviations), source_baseline=baseline, - note="generated values include injected action noise; the source baseline is against executed commands", ) @@ -361,7 +342,7 @@ def _motion_smoothness( source_step_dt: float | None, ) -> dict: if step_dt is None: - return _not_applicable("no step_dt in the result file or dataset sim_args") + return _not_applicable("the dataset has no sim_args") if not episodes: return _not_applicable("the dataset has no episodes") jerk = _pooled_jerk(episodes, step_dt) @@ -379,71 +360,30 @@ def _motion_smoothness( note=( None if source_episodes is None or source_step_dt is not None - else "source has no sim_args; its baseline assumes the generated step_dt" + else "source has no sim_args; its baseline assumes the generated control step" ), ) -def _range_points(reset_range: dict, asset: str, attempts: list[dict]) -> tuple | str: - """Starts of ``asset`` in reset-range coordinates over the range's non-zero-width x/y axes.""" - func_name = reset_range["func"].rsplit(":", 1)[-1] - if func_name in _OFFSET_RANGE_FUNCS: - origin = reset_range["default_positions"].get(asset) - if origin is None: - return f"no default position recorded for {asset}" - elif func_name in _ABSOLUTE_RANGE_FUNCS: - origin = [0.0, 0.0, 0.0] - else: - return f"unknown pose_range semantics of {reset_range['func']}" - axes = [ - (index, axis) - for index, axis in enumerate(("x", "y")) - if axis in reset_range["pose_range"] and reset_range["pose_range"][axis][1] > reset_range["pose_range"][axis][0] - ] - if not axes: - return f"{asset} has no x/y variation" - starts = np.array([attempt["initial_object_poses"][asset][:2] for attempt in attempts]) - np.array(origin[:2]) - points = starts[:, [index for index, _ in axes]] - bounds = [tuple(reset_range["pose_range"][axis]) for _, axis in axes] - in_range = np.all( - [ - (points[:, dim] >= lo - RANGE_TOLERANCE) & (points[:, dim] <= hi + RANGE_TOLERANCE) - for dim, (lo, hi) in enumerate(bounds) - ], - axis=0, - ) - if not in_range.all(): - return f"only {in_range.mean():.0%} of {asset} starts lie inside its reset range, so its frame is wrong" - return points, bounds, [axis for _, axis in axes] - - -def _configuration_inputs(result: dict | None, attempts: list[dict] | None) -> tuple[list, list[str]] | str: - if result is None or attempts is None: - return "needs --result_file with an attempt log" +def _varying_starts(attempts: list[dict] | None) -> list[tuple] | str: + """Per object, its attempted start positions along the x/y axes on which they vary, with their bounds.""" if not attempts: - return "the attempt log is empty" - reset_ranges = result.get("run", {}).get("reset_ranges", []) - entries, skipped = [], [] - for reset_range in reset_ranges: - for asset in reset_range["assets"]: - if not all(asset in attempt["initial_object_poses"] for attempt in attempts): - skipped.append(f"{asset} is not a rigid object") - continue - resolved = _range_points(reset_range, asset, attempts) - if isinstance(resolved, str): - skipped.append(resolved) - else: - entries.append((asset, *resolved)) - if not entries: - return "no reset range with x/y variation" + (f" ({'; '.join(skipped)})" if skipped else "") - return entries, skipped - - -def _success_by_configuration(result: dict | None, attempts: list[dict] | None, corrected: bool) -> dict: - resolved = _configuration_inputs(result, attempts) - if isinstance(resolved, str): - return _not_applicable(resolved) - entries, skipped = resolved + return "needs --result_file with a non-empty attempt log" + entries = [] + for asset in attempts[0]["initial_object_positions"]: + starts = np.array([attempt["initial_object_positions"][asset] for attempt in attempts]) + axes = [axis for axis in range(2) if np.ptp(starts[:, axis]) >= MIN_START_EXTENT] + if axes: + points = starts[:, axes] + bounds = [(float(points[:, dim].min()), float(points[:, dim].max())) for dim in range(len(axes))] + entries.append((asset, points, bounds, ["xy"[axis] for axis in axes])) + return entries or f"no object's start position varies by {MIN_START_EXTENT} m or more" + + +def _success_by_configuration(attempts: list[dict] | None, corrected: bool) -> dict: + entries = _varying_starts(attempts) + if isinstance(entries, str): + return _not_applicable(entries) success = np.array([attempt["success"] for attempt in attempts]) overall_rate = float(success.mean()) if overall_rate == 0.0: @@ -451,18 +391,16 @@ def _success_by_configuration(result: dict | None, attempts: list[dict] | None, per_asset = {} for asset, points, bounds, axes in entries: cells = grid_cells(points, bounds, 3) - table = [] - for cell in range(3 ** len(bounds)): - in_cell = cells == cell - n = int(in_cell.sum()) - table.append({"cell": cell, "attempts": n, "successes": int(success[in_cell].sum())}) - per_asset[asset] = {"axes": axes, "range": [list(bound) for bound in bounds], "cells": table} + table = [ + {"cell": cell, "attempts": int((cells == cell).sum()), "successes": int(success[cells == cell].sum())} + for cell in range(3 ** len(bounds)) + ] + per_asset[asset] = {"axes": axes, "bounds": bounds, "cells": table} dense = [ entry for grid in per_asset.values() for entry in grid["cells"] if entry["attempts"] >= SPARSE_CELL_ATTEMPTS ] if not dense: - reasons = [f"no grid cell has {SPARSE_CELL_ATTEMPTS} or more attempts", *skipped] - return _not_applicable("; ".join(reasons)) + return _not_applicable(f"no grid cell has {SPARSE_CELL_ATTEMPTS} or more attempts") for entry in dense: entry["rate"] = entry["successes"] / entry["attempts"] # One-sided exact test: how likely so few successes are if the cell had the overall rate. @@ -478,26 +416,20 @@ def _success_by_configuration(result: dict | None, attempts: list[dict] | None, ), overall_rate=overall_rate, per_asset=per_asset, - skipped=skipped, ) -def _initial_state_coverage(result: dict | None, attempts: list[dict] | None) -> dict: - resolved = _configuration_inputs(result, attempts) - if isinstance(resolved, str): - return _not_applicable(resolved) - entries, skipped = resolved +def _initial_state_coverage(attempts: list[dict] | None) -> dict: + entries = _varying_starts(attempts) + if isinstance(entries, str): + return _not_applicable(entries) success = np.array([attempt["success"] for attempt in attempts]) - per_asset = {} + per_asset, skipped = {}, [] for asset, points, bounds, axes in entries: if success.sum() < min_coverage_successes(len(bounds)): skipped.append(f"{asset} needs {min_coverage_successes(len(bounds))} successful attempts") continue - per_asset[asset] = { - "axes": axes, - "range": [list(bound) for bound in bounds], - **coverage_ratio(points[success], points, bounds), - } + per_asset[asset] = {"axes": axes, "bounds": bounds, **coverage_ratio(points[success], points, bounds)} if not per_asset: return _not_applicable("; ".join(skipped)) worst = min(entry["ratio"] for entry in per_asset.values()) @@ -512,52 +444,22 @@ def _failure_share(attempts: list[dict] | None, stop_attempts: list[dict] | None attempts = stop_attempts note = _CORRECTED_NOTE failures = [attempt for attempt in attempts if not attempt["success"]] - if not failures: - return _kpi({}, "fraction of failures", 0, note=note) - by_cause = Counter(attempt["failure_cause"] or "unknown" for attempt in failures) - by_subtask: dict[str, dict[str, dict[str, int]]] = {} - for attempt in failures: - for eef, subtask in (attempt["failed_subtask"] or {}).items(): - cause_table = by_subtask.setdefault(attempt["failure_cause"] or "unknown", {}).setdefault(eef, {}) - cause_table[str(subtask)] = cause_table.get(str(subtask), 0) + 1 + by_cause = Counter(attempt["failure_cause"] for attempt in failures) return _kpi( {cause: count / len(failures) for cause, count in by_cause.items()}, "fraction of failures", len(failures), - by_subtask=by_subtask, note=note, ) -def _planner_calls(attempts: list[dict] | None) -> list[dict]: - return [call for attempt in attempts or [] for call in attempt.get("planner_calls", [])] - - -def _planner_failure_types(attempts: list[dict] | None) -> dict: - calls = _planner_calls(attempts) - if not calls: - return _not_applicable("no motion-planner calls in the attempt log") - failed = [call for call in calls if not call["success"]] - return _kpi( - {status: count / len(failed) for status, count in Counter(call["status"] for call in failed).items()}, - "fraction of failed plans", - len(failed), - num_calls=len(calls), - ) - - def _planning_time(result: dict | None, attempts: list[dict] | None) -> dict: - calls = _planner_calls(attempts) - if not calls: - return _not_applicable("no motion-planner calls in the attempt log") - per_attempt = np.array([ - sum(call["wall_time_s"] for call in attempt["planner_calls"]) - for attempt in attempts - if attempt["planner_calls"] - ]) + planning = np.array([attempt["planning_time_s"] for attempt in attempts or [] if attempt["planning_time_s"] > 0]) + if not len(planning): + return _not_applicable("the run did no motion planning") generation_time = (result or {}).get("run", {}).get("generation_time_s") - share = float(per_attempt.sum() / generation_time) if generation_time else None - return _kpi(_percentiles(per_attempt), "s per attempt", len(per_attempt), share_of_generation_time=share) + share = float(planning.sum() / generation_time) if generation_time else None + return _kpi(_percentiles(planning), "s per attempt", len(planning), share_of_generation_time=share) def build_report( @@ -565,7 +467,7 @@ def build_report( result_file: str | None = None, source: str | None = None, embodiment: str | None = None, - annotation_summary: str | None = None, + recorded_dataset: str | None = None, ) -> dict[str, Any]: """Compute every KPI for a generated dataset. @@ -574,7 +476,7 @@ def build_report( result_file: Generation result JSON (``--result_file`` of ``generate_dataset.py``). source: Annotated source HDF5 dataset the generation started from; enables source baselines. embodiment: Embodiment YAML used for generation; needed for action saturation. - annotation_summary: Annotation summary JSON (``--summary_file`` of ``annotate_demos.py``). + recorded_dataset: Teleoperated HDF5 dataset that was annotated into ``source``; enables annotation yield. """ episodes, env_args = load_episodes(dataset) source_episodes, source_env_args = load_episodes(source) if source else (None, {}) @@ -596,29 +498,24 @@ def build_report( num_envs = result.get("run", {}).get("num_envs") if attempts and num_envs: stop_attempts = attempts_before_stop(attempts, num_envs) - annotation = None - if annotation_summary: - with open(annotation_summary, encoding="utf-8") as handle: - annotation = json.load(handle) - step_dt = (result or {}).get("run", {}).get("step_dt") or _step_dt_from_env_args(env_args) - source_step_dt = _step_dt_from_env_args(source_env_args) kpis = { "success_rate": _success_rate(result, stop_attempts), "demos_per_hour": _demos_per_hour(result), "demos_per_human_demo": _demos_per_human_demo(result), - "annotation_yield": _annotation_yield(annotation), + "annotation_yield": _annotation_yield(source_episodes, recorded_dataset), "tracking_error": _tracking_error(episodes, source_episodes), "action_saturation": _action_saturation( episodes, source_episodes, embodiment_clips_pose_actions(embodiment) if embodiment else None ), - "motion_smoothness": _motion_smoothness(episodes, step_dt, source_episodes, source_step_dt), - "initial_state_coverage": _initial_state_coverage(result, stop_attempts or attempts), + "motion_smoothness": _motion_smoothness( + episodes, _step_dt_from_env_args(env_args), source_episodes, _step_dt_from_env_args(source_env_args) + ), + "initial_state_coverage": _initial_state_coverage(stop_attempts or attempts), "success_by_initial_configuration": _success_by_configuration( - result, stop_attempts or attempts, corrected=stop_attempts is not None + stop_attempts or attempts, corrected=stop_attempts is not None ), "failure_share": _failure_share(attempts, stop_attempts), - "planner_failure_types": _planner_failure_types(attempts), "planning_time": _planning_time(result, attempts), } return {"dataset": dataset, "num_episodes": len(episodes), "kpis": kpis} diff --git a/docs/pages/concepts/data_generator.rst b/docs/pages/concepts/data_generator.rst index 9127a5f..2724fc1 100644 --- a/docs/pages/concepts/data_generator.rst +++ b/docs/pages/concepts/data_generator.rst @@ -123,9 +123,9 @@ Results and Statistics Three layers report what happened: -* **Per attempt** — ``GenerationResult``: ``success``, ``failure_cause``, ``failed_subtask``, - ``planner_calls`` and ``initial_state``. With ``--result_file``, each completed attempt is also - written as one line of ``.attempts.jsonl``. +* **Per attempt** — ``GenerationResult``: ``success``, ``failure_cause``, ``planning_time_s`` 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 diff --git a/docs/pages/references/kpis.rst b/docs/pages/references/kpis.rst index 7f776f9..deca735 100644 --- a/docs/pages/references/kpis.rst +++ b/docs/pages/references/kpis.rst @@ -2,23 +2,23 @@ KPIs ==== ``scripts/validate_dataset.py --report`` computes KPIs of a generated dataset from files that -generation and annotation write. It needs no simulator. +generation writes. It needs no simulator. Inputs ------ -- **Generated dataset.** Every episode records, per step under ``obs/datagen_info``, the - noise-free waypoint pose commanded (``intended_eef_pose``, only for end effectors whose action - commands a pose) and the control-link pose read before the step (``eef_pose``). +- **Generated dataset.** Generation records ``obs/datagen_info`` before every step with the same + recorder term as annotation: object poses, the control-link pose (``eef_pose``) and the target + pose the step's action commands (``target_eef_pose``). - **Result JSON and attempt log.** ``generate_dataset.py --result_file`` adds a ``run`` block to - the result JSON (device, GPU and CPU model, env count, control step, generation time, source - episode count, reset events that sample object poses from a ``pose_range``) and writes - ``.attempts.jsonl``, one line per completed attempt: env id, start and end - time, outcome, ``failure_cause`` (``planner`` or ``execution``), failed subtask per end - effector, initial object poses, and one summary per motion-planner call. + the result JSON (device, GPU and CPU model, env count, generation time, source episode count) + and writes ``.attempts.jsonl``, one line per completed attempt: env id, start + and end time, outcome, ``failure_cause`` (``planner`` or ``execution``), motion-planning time and + initial object positions. - **Source dataset** (``--source``): enables the source baselines. - **Embodiment YAML** (``--embodiment``): needed for ``action_saturation``. -- **Annotation summary** (``annotate_demos.py --summary_file``): needed for ``annotation_yield``. +- **Recorded dataset** (``--recorded_dataset``): the teleoperated demonstrations annotated into + ``--source``; needed for ``annotation_yield``. .. code-block:: bash @@ -49,10 +49,10 @@ Definitions * - ``demos_per_human_demo`` - Successful demonstrations over source demonstrations. * - ``annotation_yield`` - - Exported over processed episodes, with rejection reasons. + - Annotated demonstrations over recorded demonstrations. * - ``tracking_error`` - - p50 and p95 distance [mm] from each achieved end-effector position to the intended path. - Generated values include the injected action noise. + - p50 and p95 distance [mm] from each achieved end-effector position to the commanded path. + Tracking-only channels, which command no pose, are skipped. * - ``action_saturation`` - Share of pose-action entries at ±1, for adapters that clip pose actions. * - ``motion_smoothness`` @@ -61,26 +61,21 @@ Definitions - Per object, grid cells occupied by successful starts over cells occupied by as many attempted starts, on a grid of about five successes per cell. The worst object is reported. * - ``success_by_initial_configuration`` - - Number of reset-range cells (three per varying axis, at least 10 attempts each) whose - success count is improbably low at the overall rate: one-sided exact binomial test at - 0.025 divided by the number of cells. + - Number of cells (three per varying axis, at least 10 attempts each) whose success count is + improbably low at the overall rate: one-sided exact binomial test at 0.025 divided by the + number of cells. * - ``failure_share`` - - Share of failed attempts by cause, with counts by end effector and subtask. Execution - failures are attributed to a subtask only when the env publishes every non-final subtask's - termination signal in its ``subtask_terms`` observation group. - * - ``planner_failure_types`` - - Share of failed motion plans by cuRobo status. + - Share of failed attempts by cause. * - ``planning_time`` - - p50 and p95 planning wall-clock time per attempt [s], and its share of generation time. + - p50 and p95 motion-planning wall-clock time per attempt [s], and its share of generation time. Quality KPIs carry a ``source_baseline`` computed the same way on the source demonstrations. +The configuration KPIs grid each object's start positions along the x and y axes on which the +attempted starts span at least 1 cm, between the smallest and largest attempted start. + When generation stops, attempts still running are cancelled and never logged, and long attempts are the likelier to be running. ``success_rate``, ``failure_share`` and the configuration KPIs therefore count only attempts that started before the first env finished its last logged attempt. This needs at least three logged attempts from every env; otherwise they use all attempts and say the value is uncorrected. - -Ranges of ``randomize_object_pose`` are absolute in the env frame; ranges of -``reset_root_state_uniform`` are offsets from the asset's default position. The configuration KPIs -skip other reset functions and say why. diff --git a/scripts/annotate_demos.py b/scripts/annotate_demos.py index 06e2488..8ba775b 100644 --- a/scripts/annotate_demos.py +++ b/scripts/annotate_demos.py @@ -81,12 +81,6 @@ default="subtask_terms", help="Observation group holding the per-subtask boolean terms read in --auto mode.", ) -parser.add_argument( - "--summary_file", - type=str, - default=None, - help="Optional JSON path for the annotation summary (processed and exported counts, rejection reasons).", -) parser.add_argument( "--enable_cameras", @@ -107,26 +101,26 @@ import contextlib # noqa: E402 import gymnasium as gym # noqa: E402 -import json # noqa: E402 import math # noqa: E402 import os # noqa: E402 import sys # noqa: E402 import torch # noqa: E402 import traceback # noqa: E402 -from collections import Counter # 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.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 @@ -134,7 +128,6 @@ current_action_index = 0 marked_subtask_action_indices: list[int] = [] skip_episode = False -rejection_counts: Counter[str] = Counter() _datastream: Datastream | None = None @@ -160,34 +153,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, @@ -273,7 +238,7 @@ 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). @@ -296,6 +261,7 @@ def main() -> int: source_pool=empty_pool, uses_start_signals=annotate_start_signals, ) + PreStepDatagenInfoRecorder.datastream = _datastream env.reset() @@ -324,6 +290,7 @@ def main() -> int: for episode_index, episode_name in enumerate(episode_names): if not simulation_app.is_running() or simulation_app.is_exiting(): break + processed_episode_count += 1 print(f"\nAnnotating episode #{episode_index} ({episode_name})") episode = dataset_file_handler.load_episode(episode_name, env.device) @@ -335,7 +302,6 @@ def main() -> int: annotated = annotate_episode_in_manual_mode( env, episode, success_term, term_signal_names, start_signal_names, annotate_start_signals ) - processed_episode_count += 1 if annotated and not skip_episode: env.recorder_manager.set_success_to_episodes( @@ -351,15 +317,6 @@ 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 args_cli.summary_file: - summary = { - "num_processed": processed_episode_count, - "num_exported": exported_episode_count, - "rejections": dict(rejection_counts), - } - with open(args_cli.summary_file, "w", encoding="utf-8") as summary_handle: - json.dump(summary, summary_handle, indent=2, sort_keys=True) - summary_handle.write("\n") print("Exiting the app.") return exported_episode_count finally: @@ -473,7 +430,6 @@ def annotate_episode_in_manual_mode( task_succeeded = replay_episode(env, episode, success_term) if skip_episode: print("\tSkipping the episode.") - rejection_counts["skipped"] += 1 return False print(f"\tSubtasks marked at action indices: {marked_subtask_action_indices}") @@ -499,7 +455,6 @@ def annotate_episode_in_manual_mode( if not task_succeeded: print("\tThe final task was not completed.") - rejection_counts["task not completed"] += 1 return False print( @@ -600,7 +555,6 @@ def sample_signals(action_index: int) -> None: task_succeeded = replay_episode(env, episode, success_term, on_step=sample_signals) if not task_succeeded: print("\tThe final task was not completed.") - rejection_counts["task not completed"] += 1 return False num_steps = len(episode.data["actions"]) @@ -614,19 +568,16 @@ def sample_signals(action_index: int) -> None: if edge is None: reason = "already active at episode start" if bool(signal[0]) else "never fired" print(f'\tSubtask signal "{signal_name}" (eef "{eef_name}") {reason}.') - rejection_counts[f"signal {reason}"] += 1 return False eef_edges.append(edge) if any(later <= earlier for earlier, later in zip(eef_edges, eef_edges[1:])): print(f'\tSubtask signals for eef "{eef_name}" fired out of order: {dict(zip(eef_term_names, eef_edges))}.') - rejection_counts["signals out of order"] += 1 return False spacing_error = _validate_boundary_spacing(task_descriptor, eef_name, eef_edges, num_steps) if spacing_error is not None: print(f'\tRejected annotations for eef "{eef_name}": {spacing_error}.') - rejection_counts["boundaries too close"] += 1 return False print(f'\tDetected subtask signals for eef "{eef_name}": {dict(zip(eef_term_names, eef_edges))}') diff --git a/scripts/generate_dataset.py b/scripts/generate_dataset.py index 05b2fd7..9299140 100755 --- a/scripts/generate_dataset.py +++ b/scripts/generate_dataset.py @@ -160,8 +160,11 @@ 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 import pose_math # noqa: E402 from autodata_utils.generation_result import AttemptLog, attempts_file_for, write_generation_result # noqa: E402 from autodata_utils.tensor_utils import as_torch # noqa: E402 @@ -208,25 +211,16 @@ def _attempt_record(env_id: int, result: GenerationResult, start_time: float, en "env_id": env_id, "success": result.success, "failure_cause": result.failure_cause, - "failed_subtask": result.failed_subtask, + "planning_time_s": result.planning_time_s, "start_time_s": start_time, "end_time_s": end_time, - "initial_object_poses": _initial_object_poses(result.initial_state, env_id), - "planner_calls": list(result.planner_calls), + "initial_object_positions": { + name: as_torch(state["root_pose"])[env_id, :2].tolist() + for name, state in result.initial_state.get("rigid_object", {}).items() + }, } -def _initial_object_poses(initial_state: dict, env_id: int) -> dict[str, list[float]]: - """``[x, y, yaw]`` [m, m, rad] of each rigid object in the env frame at the start of the attempt.""" - poses = {} - for name, object_state in initial_state.get("rigid_object", {}).items(): - root_pose = as_torch(object_state["root_pose"])[env_id] - rotation = pose_math.matrix_from_quat(root_pose[3:7]) - yaw = torch.atan2(rotation[1, 0], rotation[0, 0]) - poses[name] = [float(root_pose[0]), float(root_pose[1]), float(yaw)] - return poses - - 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 @@ -241,38 +235,6 @@ def _hardware(device: str) -> dict[str, str | None]: return {"gpu": gpu, "cpu": cpu or platform.processor() or None} -def _reset_ranges(env_cfg: Any) -> list[dict]: - """Reset event terms that sample object poses from a ``pose_range``. - - ``default_positions`` lets a reader resolve ranges that are offsets from an asset's default pose. - """ - ranges = [] - events = getattr(env_cfg, "events", None) - for term in (vars(events) if events is not None else {}).values(): - params = getattr(term, "params", None) - if getattr(term, "mode", None) != "reset" or not params or not isinstance(params.get("pose_range"), dict): - continue - func = ( - term.func - if isinstance(term.func, str) - else f"{getattr(term.func, '__module__', '')}:{getattr(term.func, '__name__', '')}" - ) - asset_cfgs = list(params.get("asset_cfgs", [])) + ([params["asset_cfg"]] if "asset_cfg" in params else []) - assets = [asset_cfg.name for asset_cfg in asset_cfgs] - default_positions = {} - for name in assets: - init_state = getattr(getattr(env_cfg.scene, name, None), "init_state", None) - if init_state is not None: - default_positions[name] = [float(value) for value in init_state.pos] - ranges.append({ - "func": func, - "assets": assets, - "pose_range": {axis: [float(bounds[0]), float(bounds[1])] for axis, bounds in params["pose_range"].items()}, - "default_positions": default_positions, - }) - return ranges - - def build_datastream( env: Any, input_file: str, @@ -413,7 +375,9 @@ 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 DatagenInfoRecorderManagerCfg() + ), env_profile=env_profile, enable_cameras=args_cli.record_images, ) @@ -433,6 +397,7 @@ def main() -> None: task_descriptor=task_descriptor, embodiment_adapter=embodiment_adapter, ) + PreStepDatagenInfoRecorder.datastream = datastream # SkillGen needs one curobo planner per env. Mimic/DexMimic take no kwargs. motion_planners: dict | None = None @@ -501,10 +466,8 @@ def main() -> None: "device": args_cli.device, **_hardware(args_cli.device), "num_envs": args_cli.num_envs, - "step_dt": float(env.step_dt), "generation_time_s": generation_end - generation_start, "num_source_episodes": datastream.num_source_demos, - "reset_ranges": _reset_ranges(env_cfg), "attempts_file": Path(attempts_file).name, }, ) diff --git a/scripts/validate_dataset.py b/scripts/validate_dataset.py index 39a9b03..bd405ac 100644 --- a/scripts/validate_dataset.py +++ b/scripts/validate_dataset.py @@ -170,7 +170,9 @@ def main(argv: list[str] | None = None) -> int: parser.add_argument("--result_file", help="Generation result JSON written by generate_dataset.py --result_file.") parser.add_argument("--source", help="Annotated source HDF5 the generation started from; enables baselines.") parser.add_argument("--embodiment", help="Embodiment YAML used for generation; needed for action saturation.") - parser.add_argument("--annotation_summary", help="Summary JSON written by annotate_demos.py --summary_file.") + parser.add_argument( + "--recorded_dataset", help="Teleoperated HDF5 annotated into --source; enables annotation yield." + ) parser.add_argument("--report_file", help="Also write the KPI report to this JSON file.") args = parser.parse_args(argv) if args.report and len(args.files) != 1: @@ -194,7 +196,7 @@ def main(argv: list[str] | None = None) -> int: result_file=args.result_file, source=args.source, embodiment=args.embodiment, - annotation_summary=args.annotation_summary, + recorded_dataset=args.recorded_dataset, ) print() print(format_report(report)) diff --git a/skills/autodata-generate-datasets/references/cli-reference.md b/skills/autodata-generate-datasets/references/cli-reference.md index 47ec1a1..4c2da28 100644 --- a/skills/autodata-generate-datasets/references/cli-reference.md +++ b/skills/autodata-generate-datasets/references/cli-reference.md @@ -44,7 +44,6 @@ Annotate source demonstrations with subtask signals. | `--output_file ` | Annotated output. | | `--auto` | Annotate termination signals without keyboard input (`--viz none` supported). | | `--signal_obs_group ` | Observation group holding per-subtask boolean terms (default `subtask_terms`). | -| `--summary_file ` | Write processed/exported counts and rejection reasons as JSON. | Interactive controls: `N` play/resume, `B` pause, `S` mark subtask signal, `Q` skip episode. SkillGen start signals still require manual mode. From 7c0eab0e8d5e096ff700046749b12abe84528194 Mon Sep 17 00:00:00 2001 From: Billy Okal Date: Mon, 19 Oct 2026 15:22:52 -0500 Subject: [PATCH 4/4] refactor: rename annotation_yield to annotation_pass_rate --- autodata_utils/dataset_metrics.py | 6 +++--- docs/pages/references/kpis.rst | 6 +++--- scripts/validate_dataset.py | 2 +- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/autodata_utils/dataset_metrics.py b/autodata_utils/dataset_metrics.py index 0f479e9..31bb91d 100644 --- a/autodata_utils/dataset_metrics.py +++ b/autodata_utils/dataset_metrics.py @@ -272,7 +272,7 @@ def _demos_per_human_demo(result: dict | None) -> dict: return _kpi(result["num_success"] / num_source, "ratio", num_source, note=note) -def _annotation_yield(source_episodes: list[dict] | None, recorded_dataset: str | None) -> dict: +def _annotation_pass_rate(source_episodes: list[dict] | None, recorded_dataset: str | None) -> dict: if source_episodes is None or recorded_dataset is None: return _not_applicable("needs --source and --recorded_dataset") with h5py.File(recorded_dataset, "r") as h5_file: @@ -476,7 +476,7 @@ def build_report( result_file: Generation result JSON (``--result_file`` of ``generate_dataset.py``). source: Annotated source HDF5 dataset the generation started from; enables source baselines. embodiment: Embodiment YAML used for generation; needed for action saturation. - recorded_dataset: Teleoperated HDF5 dataset that was annotated into ``source``; enables annotation yield. + recorded_dataset: Teleoperated HDF5 dataset that was annotated into ``source``; enables the annotation pass rate. """ episodes, env_args = load_episodes(dataset) source_episodes, source_env_args = load_episodes(source) if source else (None, {}) @@ -503,7 +503,7 @@ def build_report( "success_rate": _success_rate(result, stop_attempts), "demos_per_hour": _demos_per_hour(result), "demos_per_human_demo": _demos_per_human_demo(result), - "annotation_yield": _annotation_yield(source_episodes, recorded_dataset), + "annotation_pass_rate": _annotation_pass_rate(source_episodes, recorded_dataset), "tracking_error": _tracking_error(episodes, source_episodes), "action_saturation": _action_saturation( episodes, source_episodes, embodiment_clips_pose_actions(embodiment) if embodiment else None diff --git a/docs/pages/references/kpis.rst b/docs/pages/references/kpis.rst index deca735..cc7ccd8 100644 --- a/docs/pages/references/kpis.rst +++ b/docs/pages/references/kpis.rst @@ -18,7 +18,7 @@ Inputs - **Source dataset** (``--source``): enables the source baselines. - **Embodiment YAML** (``--embodiment``): needed for ``action_saturation``. - **Recorded dataset** (``--recorded_dataset``): the teleoperated demonstrations annotated into - ``--source``; needed for ``annotation_yield``. + ``--source``; needed for ``annotation_pass_rate``. .. code-block:: bash @@ -48,8 +48,8 @@ Definitions or the CPU for ``--device cpu``) and env count. Compare only on the same hardware and env count. * - ``demos_per_human_demo`` - Successful demonstrations over source demonstrations. - * - ``annotation_yield`` - - Annotated demonstrations over recorded demonstrations. + * - ``annotation_pass_rate`` + - Share of recorded demonstrations that pass annotation into the source dataset. * - ``tracking_error`` - p50 and p95 distance [mm] from each achieved end-effector position to the commanded path. Tracking-only channels, which command no pose, are skipped. diff --git a/scripts/validate_dataset.py b/scripts/validate_dataset.py index bd405ac..afaff80 100644 --- a/scripts/validate_dataset.py +++ b/scripts/validate_dataset.py @@ -171,7 +171,7 @@ def main(argv: list[str] | None = None) -> int: parser.add_argument("--source", help="Annotated source HDF5 the generation started from; enables baselines.") parser.add_argument("--embodiment", help="Embodiment YAML used for generation; needed for action saturation.") parser.add_argument( - "--recorded_dataset", help="Teleoperated HDF5 annotated into --source; enables annotation yield." + "--recorded_dataset", help="Teleoperated HDF5 annotated into --source; enables the annotation pass rate." ) parser.add_argument("--report_file", help="Also write the KPI report to this JSON file.") args = parser.parse_args(argv)