Add the diversity metrics and the start-position bias diagnostic

autodata_interfaces/env/datagen_info_recorder.py

@@ -16,7 +16,8 @@
 
 
 class PreStepDatagenInfoRecorder(RecorderTerm):
-    """Records end-effector and target poses, and optionally rigid and deformable object state, before each step.
+    """Records end-effector, commanded target and object poses, and optionally deformable nodal positions, before
+    each step, in one frame.
 
     Attributes:
         datastream: Datastream to read. Set it on the class once the Datastream exists, since Isaac Lab
@@ -30,9 +31,8 @@ def record_pre_step(self):
         datagen_info = {
             "eef_pose": self.datastream.embodiment_adapter.get_eef_poses(env_ids=None),
             "target_eef_pose": self.datastream.action_to_target_eef_pose(self._env.action_manager.action),
+            "object_pose": self.datastream.get_object_poses(),
         }
-        if self.cfg.record_object_poses:
-            datagen_info["object_pose"] = self.datastream.get_object_poses()
         if self.cfg.record_object_nodal_positions:
             datagen_info["object_nodal_position"] = self.datastream.get_object_nodal_positions()
         return "obs/datagen_info", datagen_info
@@ -43,8 +43,6 @@ class PreStepDatagenInfoRecorderCfg(RecorderTermCfg):
     """Configuration for :class:`PreStepDatagenInfoRecorder`."""
 
     class_type: type[RecorderTerm] = PreStepDatagenInfoRecorder
-    record_object_poses: bool = True
-    """Annotation needs object poses for the source pool; generation leaves them out, as no metric reads them."""
 
     record_object_nodal_positions: bool = False
     """Record deformable nodal positions for source annotations; pose-only metric recording leaves them out."""
@@ -88,13 +86,13 @@ def add_step_success_recorder(recorder_cfg: object, success_term: TerminationTer
 
 
 def add_pose_recorder(recorder_cfg: object) -> None:
-    """Add the per-step end-effector and commanded target pose term to any recorder manager configuration.
+    """Add the per-step end-effector, commanded target and object pose term to any recorder manager configuration.
 
-    Object poses are left out, as no metric reads them. A configuration that already has the
-    ``record_pre_step_datagen_info`` term, such as :class:`DatagenInfoRecorderManagerCfg`, keeps its own.
+    A configuration that already has the ``record_pre_step_datagen_info`` term, such as
+    :class:`DatagenInfoRecorderManagerCfg`, keeps its own.
 
     Args:
         recorder_cfg: Recorder manager configuration the env will be built with.
     """
     if not hasattr(recorder_cfg, "record_pre_step_datagen_info"):
-        recorder_cfg.record_pre_step_datagen_info = PreStepDatagenInfoRecorderCfg(record_object_poses=False)
+        recorder_cfg.record_pre_step_datagen_info = PreStepDatagenInfoRecorderCfg()

autodata_metrics/__init__.py

@@ -7,7 +7,7 @@
 Importing this package registers the built-in metrics.
 """
 
-from autodata_metrics import validity, viability
+from autodata_metrics import diversity, validity, viability
 from autodata_metrics.evidence import Episode, Evidence
 from autodata_metrics.registry import (
     AREAS,

autodata_metrics/evidence.py

@@ -9,7 +9,8 @@
 import h5py
 import json
 import numpy as np
-from collections.abc import Mapping
+from collections.abc import Iterator, Mapping
+from contextlib import contextmanager
 from functools import cached_property
 from typing import Any
 
@@ -46,6 +47,7 @@ class Episode:
     def __init__(self, group: h5py.Group) -> None:
         self.name = group.name.rsplit("/", 1)[-1]
         self._group = group
+        self._listing: list[tuple[str, h5py.Dataset]] | None = None
 
     def __contains__(self, key: str) -> bool:
         return key in self._group
@@ -57,15 +59,41 @@ def children(self, key: str) -> list[str]:
         """Names under the group ``key``, such as the end effectors under ``obs/datagen_info/eef_pose``."""
         return list(self._group[key])
 
+    @contextmanager
+    def listed(self) -> Iterator[None]:
+        """Within the block, :meth:`datasets` walks the episode once, so several checks share one listing."""
+        try:
+            self._listing = self.datasets()
+        except Exception:  # An episode that cannot be walked is walked by each check, which fails it with the error.
+            self._listing = None
+        try:
+            yield
+        finally:
+            self._listing = None
+
     def datasets(self) -> list[tuple[str, h5py.Dataset]]:
         """Every dataset of the episode with its path, unread until accessed."""
+        if self._listing is not None:
+            return self._listing
         found: list[tuple[str, h5py.Dataset]] = []
         self._group.visititems(
             lambda path, item: found.append((path, item)) if isinstance(item, h5py.Dataset) else None
         )
         return found
 
 
+def _episode(data: h5py.Group, name: str) -> Episode:
+    """The episode ``name`` of the ``data`` group.
+
+    Raises:
+        InvalidRecord: If the episode's group cannot be opened.
+    """
+    try:
+        return Episode(data[name])
+    except (KeyError, OSError, RuntimeError) as exc:
+        raise InvalidRecord(f"{data.file.filename} episode {name} cannot be read: {exc}") from exc
+
+
 class Evidence:
     """A dataset, and optionally the generation run that produced it, as metrics read them.
 
@@ -89,7 +117,7 @@ def __init__(
         source: Evidence | None = None,
     ) -> None:
         self.env_args = read_env_args(h5_file)
-        self.all_episodes = [Episode(group) for group in h5_file["data"].values()]
+        self.all_episodes = [_episode(h5_file["data"], name) for name in h5_file["data"]]
         self.episodes = list(self.all_episodes)
         self.failures: dict[str, dict[str, str]] = {}
         self.check_groups: dict[str, str] = {}
@@ -139,12 +167,20 @@ def as_baseline_for(self, run: Evidence) -> Evidence:
         baseline.env_args = {**self.env_args, "autodata": entry}
         return baseline
 
+    @cached_property
+    def baseline(self) -> Evidence | None:
+        """The source dataset as a baseline for this run, built once on first use, after the episode checks."""
+        return self.source.as_baseline_for(self) if self.source is not None else None
+
     @cached_property
     def camera_paths(self) -> frozenset[str]:
-        """Paths of the camera streams, 4-D ``uint8`` arrays, that any episode records."""
-        return frozenset(
-            path
-            for episode in self.all_episodes
-            for path, dataset in episode.datasets()
-            if dataset.dtype == np.uint8 and dataset.ndim == 4
-        )
+        """Paths of the camera streams, 4-D ``uint8`` arrays, that any readable episode records."""
+        paths = set()
+        for episode in self.all_episodes:
+            try:
+                paths.update(
+                    path for path, dataset in episode.datasets() if dataset.dtype == np.uint8 and dataset.ndim == 4
+                )
+            except Exception:  # An episode that cannot be read fails its own checks.
+                continue
+        return frozenset(paths)

autodata_metrics/records.py

@@ -72,6 +72,7 @@ class MapOf:
     "start_time_s": NUMBER,
     "end_time_s": NUMBER,
     "extras": Maybe({"planning_time_s": Maybe(NON_NEGATIVE), "planning_failed": Maybe(bool)}),
+    "initial_object_positions": Maybe(MapOf(ListOf(NUMBER, 2))),
 }
 ANNOTATION = {
     "mode": Maybe(("auto", "manual")),
@@ -118,7 +119,8 @@ def parse(value: Any, shape: Any, where: str) -> Any:
                 parsed[name] = field
         return parsed
     if isinstance(shape, ListOf):
-        _require(isinstance(value, list) and shape.length in (None, len(value)), value, "a list", where)
+        expected = "a list" if shape.length is None else f"a list of {shape.length}"
+        _require(isinstance(value, list) and shape.length in (None, len(value)), value, expected, where)
         return [parse(item, shape.shape, f"{where}[{index}]") for index, item in enumerate(value)]
     if isinstance(shape, MapOf):
         _require(isinstance(value, dict), value, "an object", where)

autodata_metrics/report.py

@@ -124,25 +124,28 @@ def check_same_run(evidence: Evidence) -> list[str]:
 
 
 def run_checks(evidence: Evidence, checks: Mapping[str, CheckSpec]) -> None:
-    """Run the episode checks and keep only the episodes that pass all of them in ``evidence.episodes``."""
+    """Run the episode checks and keep only the episodes that pass all of them in ``evidence.episodes``.
+
+    Episodes are checked one at a time, each listed once for all its checks and released after them.
+    """
     for spec in checks.values():
         evidence.check_groups[spec.name] = spec.group
-        if any(not evidence.has(key) for key in spec.requires):
-            continue
-        failed: dict[str, str] = {}
-        applied = False
-        for episode in evidence.all_episodes:
-            try:
-                reason = spec.check(episode, evidence)
-            except NotApplicable:
-                continue
-            except Exception as exc:  # A broken check fails its episodes visibly instead of sinking the report.
-                reason = f"{type(exc).__name__}: {exc}"
-            applied = True
-            if reason is not None:
-                failed[episode.name] = reason
-        if applied:
-            evidence.failures[spec.name] = failed
+    applicable = {spec.name: spec for spec in checks.values() if all(evidence.has(key) for key in spec.requires)}
+    failed: dict[str, dict[str, str]] = {name: {} for name in applicable}
+    applied: set[str] = set()
+    for episode in evidence.all_episodes:
+        with episode.listed():
+            for name, spec in applicable.items():
+                try:
+                    reason = spec.check(episode, evidence)
+                except NotApplicable:
+                    continue
+                except Exception as exc:  # A broken check fails its episodes visibly instead of sinking the report.
+                    reason = f"{type(exc).__name__}: {exc}"
+                applied.add(name)
+                if reason is not None:
+                    failed[name][episode.name] = reason
+    evidence.failures.update({name: failed[name] for name in applicable if name in applied})
     failing = {name for failed in evidence.failures.values() for name in failed}
     evidence.episodes = [episode for episode in evidence.all_episodes if episode.name not in failing]
 
@@ -181,8 +184,8 @@ def evaluate(spec: MetricSpec, evidence: Evidence) -> dict[str, Any]:
         "n": int(value.n),
         "details": _finite_or_none(value.details),
     }
-    if spec.baseline and evidence.source is not None:
-        row["baseline"] = evaluate(spec, evidence.source.as_baseline_for(evidence)).get("value")
+    if spec.baseline and evidence.baseline is not None:
+        row["baseline"] = evaluate(spec, evidence.baseline).get("value")
     return row
 
 

autodata_metrics/viability.py

@@ -10,12 +10,17 @@
 from collections import Counter
 from typing import Any
 
+from autodata_metrics.diversity import MIN_START_EXTENT_M
 from autodata_metrics.evidence import Evidence
 from autodata_metrics.registry import NotApplicable, Value, metric
 
 Z_95 = 1.959963984540054
 MIN_ATTEMPTS_PER_ENV_FOR_STOP_CORRECTION = 3
 """Below this many logged attempts in some env, the stop correction is itself biased."""
+BIAS_BINS = 3
+"""Cells per moving axis of the start-position grid that ``generation_bias`` tests."""
+SPARSE_CELL_ATTEMPTS = 10
+"""Fewest attempts a start-position cell needs for ``generation_bias`` to test its success rate."""
 
 
 def wilson_interval(successes: int, trials: int, z: float = Z_95) -> tuple[float, float]:
@@ -28,6 +33,23 @@ def wilson_interval(successes: int, trials: int, z: float = Z_95) -> tuple[float
     return max(0.0, center - half_width), min(1.0, center + half_width)
 
 
+def binomial_cdf(successes: int, trials: int, rate: float) -> float:
+    """Probability of at most ``successes`` successes in ``trials`` trials that each succeed at ``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 attempts_before_stop(attempts: list[dict[str, Any]], num_envs: int) -> list[dict[str, Any]] | None:
     """Attempts that started before the first env completed its last logged attempt.
 
@@ -173,3 +195,61 @@ def planning_time(ev: Evidence) -> Value:
     generation_time = ev.run.get("generation_time_s")
     share = float(times.sum() / generation_time) if generation_time else None
     return Value({"p50": float(p50), "p95": float(p95)}, n=len(times), details={"share_of_generation_time": share})
+
+
+@metric(
+    "generation_bias",
+    area="diagnostics",
+    signal="failure causes",
+    scope="run",
+    unit="low start cells",
+    sample="attempts",
+    better="lower",
+    requires=["attempts.initial_object_positions"],
+)
+def generation_bias(ev: Evidence) -> Value:
+    """Start-position cells whose success rate is significantly below the overall rate.
+
+    Each moving object's attempted starts are split into :data:`BIAS_BINS` cells per moving axis. Every cell with
+    :data:`SPARSE_CELL_ATTEMPTS` or more attempts gets a one-sided exact binomial test against the overall rate,
+    stop-corrected when the attempt log allows, Bonferroni-corrected at 2.5%. With nine cells tested, a cell that never
+    succeeds is flagged from 10 attempts at an overall rate of 50%, and from 56 at 10%. The details name the worst
+    tested cell, flagged or not.
+    """
+    attempts, corrected = _counted_attempts(ev)
+    success = np.array([attempt["success"] for attempt in attempts])
+    if not success.any():
+        raise NotApplicable("no attempt succeeded")
+    rate = float(success.mean())
+    names = set.intersection(*(set(attempt["initial_object_positions"]) for attempt in attempts))
+    if not names:
+        raise NotApplicable("no object's start position is in every attempt")
+    tested = []
+    for name in sorted(names):
+        starts = np.array([attempt["initial_object_positions"][name] for attempt in attempts])
+        moving = np.ptp(starts, axis=0) >= MIN_START_EXTENT_M
+        if not moving.any():
+            continue
+        axes = [axis for axis, moves in zip("xy", moving) if moves]
+        low, extent = starts[:, moving].min(axis=0), np.ptp(starts[:, moving], axis=0)
+        bins = np.clip(((starts[:, moving] - low) / extent * BIAS_BINS).astype(int), 0, BIAS_BINS - 1)
+        for cell in np.unique(bins, axis=0):
+            members = (bins == cell).all(axis=1)
+            count, successes = int(members.sum()), int(success[members].sum())
+            if count >= SPARSE_CELL_ATTEMPTS:
+                p_value = binomial_cdf(successes, count, rate)
+                label = dict(zip(axes, cell.tolist()))
+                tested.append(
+                    {"object": name, "cell": label, "attempts": count, "successes": successes, "p_value": p_value}
+                )
+    if not tested:
+        raise NotApplicable(f"no moving object's start cell has {SPARSE_CELL_ATTEMPTS} attempts")
+    low_cells = [entry for entry in tested if entry["p_value"] < 0.025 / len(tested)]
+    details = {
+        "success_rate": rate,
+        "stop_corrected": corrected,
+        "tested_cells": len(tested),
+        "worst_cell": min(tested, key=lambda entry: entry["p_value"]),
+        "low_cells": low_cells,
+    }
+    return Value(len(low_cells), n=len(attempts), details=details)

autodata_tests/e2e/test_mimicgen_data_generation.py

@@ -8,7 +8,7 @@
 
 from autodata_tests.utils.constants import TestPaths
 from autodata_tests.utils.subprocess import run_subprocess
-from autodata_tests.utils.utils import assert_ok_metrics, assert_recorded_run, assert_valid_dataset
+from autodata_tests.utils.utils import PAIR_METRICS, assert_ok_metrics, assert_recorded_run, assert_valid_dataset
 
 HEADLESS = True
 GENERATION_NUM_TRIALS = 1
@@ -23,6 +23,7 @@
     "action_saturation",
     "motion_smoothness",
     "joint_limit_ratio",
+    "exact_duplicates",
 }
 
 
@@ -64,7 +65,7 @@ def _run_franka_cube_stack_mimicgen(num_envs: int, device: str) -> None:
 
         assert_valid_dataset(output_file, min_num_demos=GENERATION_NUM_TRIALS)
         assert_recorded_run(output_file, result_file)
-        assert_ok_metrics(output_file, result_file, METRICS)
+        assert_ok_metrics(output_file, result_file, METRICS, PAIR_METRICS)
 
 
 @pytest.mark.with_subprocess

autodata_tests/e2e/test_skillgen_bin_data_generation.py

@@ -17,7 +17,7 @@
 
 from autodata_tests.utils.constants import TestPaths
 from autodata_tests.utils.subprocess import run_subprocess
-from autodata_tests.utils.utils import assert_ok_metrics, assert_recorded_run, assert_valid_dataset
+from autodata_tests.utils.utils import PAIR_METRICS, assert_ok_metrics, assert_recorded_run, assert_valid_dataset
 
 HEADLESS = True
 GENERATION_NUM_TRIALS = 1
@@ -34,6 +34,7 @@
     "action_saturation",
     "motion_smoothness",
     "joint_limit_ratio",
+    "exact_duplicates",
 }
 
 
@@ -80,7 +81,7 @@ def _run_franka_bin_stack_skillgen(num_envs: int, device: str) -> None:
         attempts = assert_recorded_run(output_file, result_file)
         # Every SkillGen attempt plans at least one transit, and the algorithm reports how long it took.
         assert all(attempt["extras"]["planning_time_s"] > 0 for attempt in attempts)
-        assert_ok_metrics(output_file, result_file, METRICS)
+        assert_ok_metrics(output_file, result_file, METRICS, PAIR_METRICS)
 
         with open(result_file, encoding="utf-8") as result_handle:
             result = json.load(result_handle)