Report dataset KPIs from generation records

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,13 +239,15 @@ 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,
             env_id=env_id,
             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
 

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()),
         )
 
     # ------------------------------------------------------------------

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()

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()

autodata_tests/e2e/test_mimicgen_data_generation.py

@@ -1,24 +1,41 @@
 # 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
 
 import pytest
 
 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

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():