diff --git a/autodata_core/algorithms.py b/autodata_core/algorithms.py index 65d6daf..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 @@ -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,6 +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), ) + eef_state.planning_time_s += time.perf_counter() - start_time if not planning_success: return None diff --git a/autodata_core/data_generator.py b/autodata_core/data_generator.py index 59481cb..621a82b 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,20 @@ async def _optional_lock(lock: asyncio.Lock | None): @dataclass(frozen=True) class GenerationResult: - """Output of one :meth:`DataGenerator.generate` invocation.""" + """Output of one :meth:`DataGenerator.generate` invocation. + + Attributes: + initial_state: Scene state of every env right after this env's reset. + success: Whether the task succeeded at any step of the attempt. + failure_cause: ``"planner"`` when motion planning aborted the attempt, otherwise ``"execution"`` + for a failed attempt and None on success. + planning_time_s: Wall-clock time spent in motion planning [s]. + """ initial_state: dict success: bool + failure_cause: Literal["planner", "execution"] | None = None + planning_time_s: float = 0.0 @dataclass @@ -76,6 +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 + planning_time_s: float = 0.0 class DataGenerator: @@ -569,7 +580,12 @@ async def generate( ) if result is None: # Planning failure (e.g. SkillGen motion planner). Abort with no success. - return GenerationResult(initial_state=initial_state, success=False) + return GenerationResult( + initial_state=initial_state, + success=False, + failure_cause="planner", + 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 eef_state.subtask_step_index = 0 @@ -620,6 +636,8 @@ async def generate( return GenerationResult( initial_state=initial_state, success=buffers.success, + failure_cause=None if buffers.success else "execution", + planning_time_s=sum(state.planning_time_s for state in eef_states.values()), ) # ------------------------------------------------------------------ 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_tests/core/test_data_generator.py b/autodata_tests/core/test_data_generator.py index 806c0e9..6667ca1 100644 --- a/autodata_tests/core/test_data_generator.py +++ b/autodata_tests/core/test_data_generator.py @@ -196,6 +196,21 @@ def test_navigation_does_not_skip_playback_after_cancellation_or_while_paused(pa assert all(state.subtask_step_index == 4 for state in states.values()) +def test_planning_failure_is_attributed_to_the_planner(monkeypatch): + generator = _navigation_generator(active=False, reached=True) + generator.datastream.env = SimpleNamespace(recorder_manager=Mock()) + monkeypatch.setattr(generator.algorithm, "plan_subtask_trajectory", lambda **kwargs: None) + monkeypatch.setattr(generator, "_reset_environment_for_generation", AsyncMock(return_value=(torch.tensor([0]), {}))) + + result = asyncio.run( + generator.generate( + env_id=0, success_term=None, env_reset_queue=asyncio.Queue(), env_action_queue=asyncio.Queue() + ) + ) + + assert result.success is False and result.failure_cause == "planner" + + def test_unresponsive_navigation_fails_and_finalizes_the_attempt(monkeypatch): generator = _navigation_generator(active=True, reached=False) generator.datastream.generation_policy.max_navigation_hold_steps = 3 @@ -219,7 +234,7 @@ def test_unresponsive_navigation_fails_and_finalizes_the_attempt(monkeypatch): ) ) - assert result.success is False + assert result.success is False and result.failure_cause == "execution" assert execute.await_count == 3 assert recorder.set_success_to_episodes.call_args.args[1].item() is False recorder.export_episodes.assert_called_once() diff --git a/autodata_tests/e2e/test_mimicgen_data_generation.py b/autodata_tests/e2e/test_mimicgen_data_generation.py index cfb02b9..9d9babd 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,31 @@ 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 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 +47,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 +66,32 @@ 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) + with h5py.File(output_file, "r") as dataset: + for episode in dataset["data"].values(): + 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) + with open(result_file, encoding="utf-8") as result_handle: + result = json.load(result_handle) + + 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) + hardware = result["run"]["gpu"] if device == "cuda" else result["run"]["cpu"] + assert hardware and hardware in kpis["demos_per_hour"]["note"] @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..0235aba 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 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", "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 = [ @@ -71,6 +74,18 @@ def _run_franka_bin_stack_skillgen(num_envs: int, device: str) -> None: "planner": "franka_stack_cube_bin", } + 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, + ]) + assert_kpi_report(report_file, APPLICABLE_KPIS) + @pytest.mark.with_subprocess def test_franka_bin_stack_skillgen_data_generation_single_env_cuda(): diff --git a/autodata_tests/utils/test_dataset_metrics.py b/autodata_tests/utils/test_dataset_metrics.py new file mode 100644 index 0000000..917777d --- /dev/null +++ b/autodata_tests/utils/test_dataset_metrics.py @@ -0,0 +1,99 @@ +# 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_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]) + # 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([episode], [episode]) + + 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: + 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]: + # The table only settles by about a millimetre, so it does not count as a varying object. + return [ + { + "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_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) + 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) + ) + + 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 3e1ea73..cac35fc 100644 --- a/autodata_tests/utils/test_validate_dataset.py +++ b/autodata_tests/utils/test_validate_dataset.py @@ -45,3 +45,37 @@ 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", "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, "initial_object_positions": {}}) + "\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..d293920 100644 --- a/autodata_tests/utils/utils.py +++ b/autodata_tests/utils/utils.py @@ -5,6 +5,7 @@ import h5py import json +import math import os @@ -41,3 +42,22 @@ def read_generation_result(result_file: str) -> tuple[int, int]: assert num_attempts > 0, f"Generation result has invalid num_attempts: {num_attempts}." assert num_success + num_failures == num_attempts, "Generation result counts do not add up." return num_success, num_attempts + + +def assert_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 diff --git a/autodata_utils/dataset_metrics.py b/autodata_utils/dataset_metrics.py new file mode 100644 index 0000000..31bb91d --- /dev/null +++ b/autodata_utils/dataset_metrics.py @@ -0,0 +1,553 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""KPIs of a generated AutoData dataset, computed offline from generation and annotation outputs. + +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 + +Z_95 = 1.959963984540054 +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 stop correction is itself biased.""" +SAMPLES_PER_COVERAGE_CELL = 5 +PATH_DEVIATION_WINDOW = 20 +"""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: + """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, 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]. + 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(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) + 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 [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) + 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} + + +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. + + 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) + 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 actions and ``obs/datagen_info`` positions, plus the dataset's ``env_args``. + + 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: + 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 ("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()] + + +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 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() + 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, 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") + 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"], + device=run["device"], + gpu=run["gpu"], + cpu=run["cpu"], + num_envs=run["num_envs"], + ) + + +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 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) + + +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: + 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]) -> np.ndarray: + deviations = [ + path_deviation(episode["eef_pose"][eef], target) + for episode in episodes + 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) + + +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) + if not len(deviations): + 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 commanded path (position)", + len(deviations), + source_baseline=baseline, + ) + + +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("the dataset has no 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 control step" + ), + ) + + +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 "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: + return _not_applicable("no attempt succeeded") + per_asset = {} + for asset, points, bounds, axes in entries: + cells = grid_cells(points, bounds, 3) + 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: + 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. + 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, + ) + + +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, 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, "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()) + 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"]] + 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), + note=note, + ) + + +def _planning_time(result: dict | None, attempts: list[dict] | None) -> dict: + 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(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( + dataset: str, + result_file: str | None = None, + source: str | None = None, + embodiment: str | None = None, + recorded_dataset: 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. + 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, {}) + 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) + + 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_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 + ), + "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( + stop_attempts or attempts, corrected=stop_attempts is not None + ), + "failure_share": _failure_share(attempts, stop_attempts), + "planning_time": _planning_time(result, attempts), + } + return {"dataset": dataset, "num_episodes": len(episodes), "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..65f545d 100644 --- a/autodata_utils/generation_result.py +++ b/autodata_utils/generation_result.py @@ -11,6 +11,7 @@ import tempfile from collections.abc import Mapping from pathlib import Path +from typing import Any def write_generation_result( @@ -19,6 +20,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 +33,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 +54,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 +70,21 @@ def write_generation_result( finally: with contextlib.suppress(FileNotFoundError): os.unlink(temporary_file_name) + + +def attempts_file_for(result_file: str) -> str: + """Path of the attempt log written next to ``result_file``.""" + return str(Path(result_file).with_suffix(".attempts.jsonl")) + + +class AttemptLog: + """Appends one JSON line per completed generation attempt.""" + + 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/concepts/data_generator.rst b/docs/pages/concepts/data_generator.rst index b182383..2724fc1 100644 --- a/docs/pages/concepts/data_generator.rst +++ b/docs/pages/concepts/data_generator.rst @@ -123,7 +123,9 @@ Results and Statistics Three layers report what happened: -* **Per attempt** — ``GenerationResult``: ``success`` and ``initial_state``. +* **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 @@ -146,7 +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. + 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 new file mode 100644 index 0000000..cc7ccd8 --- /dev/null +++ b/docs/pages/references/kpis.rst @@ -0,0 +1,81 @@ +KPIs +==== + +``scripts/validate_dataset.py --report`` computes KPIs of a generated dataset from files that +generation writes. It needs no simulator. + +Inputs +------ + +- **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, 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``. +- **Recorded dataset** (``--recorded_dataset``): the teleoperated demonstrations annotated into + ``--source``; needed for ``annotation_pass_rate``. + +.. 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 + +A dataset left empty because every attempt failed is still reported. A KPI that cannot be +computed reports ``not_applicable`` with a reason. + +Definitions +----------- + +.. list-table:: + :widths: 26 74 + :header-rows: 1 + + * - KPI + - Definition + * - ``success_rate`` + - Successful attempts over attempts, with a Wilson 95% interval. An attempt succeeds if the + task is successful at any step. + * - ``demos_per_hour`` + - 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. + * - ``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. + * - ``action_saturation`` + - Share of pose-action entries at ±1, for adapters that clip pose actions. + * - ``motion_smoothness`` + - p50 and p95 end-effector linear jerk [m/s³]. + * - ``initial_state_coverage`` + - 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 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. + * - ``planning_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. diff --git a/scripts/annotate_demos.py b/scripts/annotate_demos.py index b0d6b2b..8ba775b 100644 --- a/scripts/annotate_demos.py +++ b/scripts/annotate_demos.py @@ -110,15 +110,17 @@ 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 @@ -151,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, @@ -264,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). @@ -287,6 +261,7 @@ def main() -> int: source_pool=empty_pool, uses_start_signals=annotate_start_signals, ) + PreStepDatagenInfoRecorder.datastream = _datastream env.reset() diff --git a/scripts/generate_dataset.py b/scripts/generate_dataset.py index d6f5167..9299140 100755 --- a/scripts/generate_dataset.py +++ b/scripts/generate_dataset.py @@ -97,7 +97,10 @@ "--result_file", type=str, default=None, - help="Optional JSON path for the final structured generation result.", + help=( + "Optional JSON path for the final structured generation result. An attempt log with one JSON " + "line per attempt is written next to it (.attempts.jsonl)." + ), ) parser.add_argument( "--pause_subtask", @@ -137,14 +140,18 @@ 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 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, @@ -153,8 +160,13 @@ setup_env_config, setup_output_paths, ) +from autodata_interfaces.env.datagen_info_recorder import ( # noqa: E402 + DatagenInfoRecorderManagerCfg, + PreStepDatagenInfoRecorder, +) from autodata_interfaces.env.registration import load_callable, register_external_environment # noqa: E402 -from autodata_utils.generation_result import write_generation_result # noqa: E402 +from autodata_utils.generation_result import AttemptLog, attempts_file_for, write_generation_result # noqa: E402 +from autodata_utils.tensor_utils import as_torch # noqa: E402 async def run_data_generator( @@ -166,9 +178,11 @@ async def run_data_generator( success_term, pause_subtask: bool, stats: dict, + attempt_log: AttemptLog | None, ) -> None: - """Repeatedly call ``data_generator.generate`` and tally outcomes into ``stats``.""" + """Repeatedly call ``data_generator.generate``, tally outcomes into ``stats`` and log each attempt.""" while True: + attempt_start = time.monotonic() try: result = await data_generator.generate( env_id=env_id, @@ -187,6 +201,38 @@ async def run_data_generator( else: stats["num_failures"] += 1 stats["num_attempts"] += 1 + if attempt_log is not None: + attempt_log.write(_attempt_record(env_id, result, attempt_start, time.monotonic())) + + +def _attempt_record(env_id: int, result: GenerationResult, start_time: float, end_time: float) -> dict: + """One attempt-log line; times [s] are monotonic-clock readings.""" + return { + "env_id": env_id, + "success": result.success, + "failure_cause": result.failure_cause, + "planning_time_s": result.planning_time_s, + "start_time_s": start_time, + "end_time_s": end_time, + "initial_object_positions": { + name: as_torch(state["root_pose"])[env_id, :2].tolist() + for name, state in result.initial_state.get("rigid_object", {}).items() + }, + } + + +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 build_datastream( @@ -216,6 +262,7 @@ def setup_async_generation( success_term, algorithm, pause_subtask: bool, + attempt_log: AttemptLog | None = None, ) -> dict: """Build a single :class:`DataGenerator` over ``datastream`` and ``num_envs`` async tasks.""" asyncio_event_loop = asyncio.get_event_loop() @@ -238,6 +285,7 @@ def setup_async_generation( success_term=success_term, pause_subtask=pause_subtask, stats=stats, + attempt_log=attempt_log, ) ) tasks.append(task) @@ -327,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, ) @@ -347,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 @@ -362,6 +413,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 +422,12 @@ def main() -> None: success_term=success_term, algorithm=algorithm, pause_subtask=args_cli.pause_subtask, + attempt_log=attempt_log, ) data_gen_tasks = asyncio.ensure_future(asyncio.gather(*async_components["tasks"])) generation_completed = False + generation_start = time.monotonic() try: generation_completed = env_loop( env, @@ -383,6 +438,7 @@ def main() -> None: stats=async_components["stats"], data_gen_tasks=data_gen_tasks, ) + generation_end = time.monotonic() except asyncio.CancelledError: print("Async tasks cancelled.") finally: @@ -406,6 +462,14 @@ def main() -> None: if env_profile else None ), + run_metadata={ + "device": args_cli.device, + **_hardware(args_cli.device), + "num_envs": args_cli.num_envs, + "generation_time_s": generation_end - generation_start, + "num_source_episodes": datastream.num_source_demos, + "attempts_file": Path(attempts_file).name, + }, ) finally: env.close() diff --git a/scripts/validate_dataset.py b/scripts/validate_dataset.py index a53db2d..afaff80 100644 --- a/scripts/validate_dataset.py +++ b/scripts/validate_dataset.py @@ -9,6 +9,7 @@ python scripts/validate_dataset.py [ ...] python scripts/validate_dataset.py directory/*.hdf5 + python scripts/validate_dataset.py --report --result_file # KPIs The command exits nonzero if any supplied file is invalid. """ @@ -21,6 +22,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 +166,17 @@ 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( + "--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) + 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, + recorded_dataset=args.recorded_dataset, + ) + 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..4c2da28 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). | @@ -60,3 +60,10 @@ 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 +```