Add episode checks, validity metrics and a console summary of the metrics report

autodata_metrics/__init__.py

@@ -7,7 +7,17 @@
 Importing this package registers the built-in metrics.
 """
 
-from autodata_metrics import viability
+from autodata_metrics import validity, viability
 from autodata_metrics.evidence import Episode, Evidence
-from autodata_metrics.registry import AREAS, REGISTERED_METRICS, MetricSpec, NotApplicable, Value, metric
-from autodata_metrics.report import build_report, format_report, gate_violations
+from autodata_metrics.registry import (
+    AREAS,
+    REGISTERED_CHECKS,
+    REGISTERED_METRICS,
+    CheckSpec,
+    MetricSpec,
+    NotApplicable,
+    Value,
+    check,
+    metric,
+)
+from autodata_metrics.report import build_report, format_report, gate_results, load_gate

autodata_metrics/evidence.py

@@ -9,6 +9,7 @@
 import json
 import numpy as np
 from collections.abc import Mapping
+from functools import cached_property
 from typing import Any
 
 
@@ -22,10 +23,27 @@ def _has_path(record: Any, dotted_key: str) -> bool:
 
 
 def read_env_args(h5_file: h5py.File) -> dict[str, Any]:
-    """The dataset's ``env_args``, or empty when they are not a JSON object."""
+    """The dataset's ``env_args``, or empty when it records none.
+
+    Raises:
+        ValueError: If ``env_args`` is not a JSON object, or its ``autodata`` entry lacks the run id and exported
+            episode count; a run cannot be verified against corrupt metadata.
+    """
     raw_env_args = h5_file["data"].attrs.get("env_args", "{}")
-    env_args = json.loads(raw_env_args.decode() if isinstance(raw_env_args, bytes) else raw_env_args)
-    return env_args if isinstance(env_args, dict) else {}
+    try:
+        env_args = json.loads(raw_env_args.decode() if isinstance(raw_env_args, bytes) else raw_env_args)
+    except (TypeError, ValueError):
+        env_args = None
+    if not isinstance(env_args, dict):
+        raise ValueError(f"{h5_file.filename}: env_args are not a JSON object.")
+    entry = env_args.get("autodata")
+    if "autodata" in env_args and not (
+        isinstance(entry, dict)
+        and isinstance(entry.get("run_id"), str)
+        and isinstance(entry.get("exported_episodes"), int)
+    ):
+        raise ValueError(f"{h5_file.filename}: the autodata entry lacks its run_id or exported_episodes.")
+    return env_args
 
 
 class Episode:
@@ -41,12 +59,27 @@ def __contains__(self, key: str) -> bool:
     def __getitem__(self, key: str) -> np.ndarray:
         return self._group[key][()]
 
+    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])
+
+    def datasets(self) -> list[tuple[str, h5py.Dataset]]:
+        """Every dataset of the episode with its path, unread until accessed."""
+        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
+
 
 class Evidence:
     """A dataset, and optionally the generation run that produced it, as metrics read them.
 
     Attributes:
-        episodes: The dataset's episodes.
+        all_episodes: The dataset's episodes.
+        episodes: The episodes that passed every episode check; dataset metrics read these.
+        failures: Per episode check that applied to some episode, the failing episodes' names and reasons.
+        check_groups: Per episode check considered, its group; a check absent from ``failures`` did not apply.
         env_args: The dataset's ``env_args``; its ``autodata`` entry was written by the producer.
         result: Generation result, or None for a dataset without a run record or an interrupted run.
         run: The result's ``run`` block, or empty.
@@ -62,7 +95,10 @@ def __init__(
         source: Evidence | None = None,
     ) -> None:
         self.env_args = read_env_args(h5_file)
-        self.episodes = [Episode(group) for group in h5_file["data"].values()]
+        self.all_episodes = [Episode(group) for group in h5_file["data"].values()]
+        self.episodes = list(self.all_episodes)
+        self.failures: dict[str, dict[str, str]] = {}
+        self.check_groups: dict[str, str] = {}
         self.result = result
         self.run: Mapping[str, Any] = (result or {}).get("run", {})
         self.attempts = attempts
@@ -88,3 +124,13 @@ def has(self, key: str) -> bool:
         if head == "run":
             return _has_path(self.run, rest)
         return bool(self.episodes) and all(key in episode for episode in self.episodes)
+
+    @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
+        )

autodata_metrics/registry.py

@@ -15,7 +15,7 @@
 from typing import TYPE_CHECKING, Any, Literal
 
 if TYPE_CHECKING:
-    from autodata_metrics.evidence import Evidence
+    from autodata_metrics.evidence import Episode, Evidence
 
 AREAS = ("validity", "diversity", "viability", "diagnostics")
 
@@ -47,41 +47,60 @@ class MetricSpec:
     compute: Callable[[Evidence], Value]
     area: str
     signal: str
-    scope: Literal["dataset", "run"]
+    scope: Literal["episode", "dataset", "run"]
     unit: str
+    sample: str
     better: Literal["higher", "lower"] | None
     requires: tuple[str, ...]
     primary: str | None
     version: int
+    baseline: bool = False
+
+
+@dataclass(frozen=True)
+class CheckSpec:
+    """A registered episode check."""
+
+    name: str
+    check: Callable[[Episode, Evidence], str | None]
+    group: str
+    requires: tuple[str, ...]
 
 
 REGISTERED_METRICS: dict[str, MetricSpec] = {}
+REGISTERED_CHECKS: dict[str, CheckSpec] = {}
+CHECK_GROUPS = ("correctness", "image")
 
 
 def metric(
     name: str,
     *,
     area: str,
     signal: str,
-    scope: Literal["dataset", "run"],
+    scope: Literal["episode", "dataset", "run"],
     unit: str,
+    sample: str,
     better: Literal["higher", "lower"] | None = None,
     requires: tuple[str, ...] | list[str] = (),
     primary: str | None = None,
     version: int = 1,
+    baseline: bool = False,
 ) -> Callable[[Callable[[Evidence], Value]], Callable[[Evidence], Value]]:
     """Register the decorated function as the metric ``name``.
 
     Args:
         name: Unique metric name, used in reports and gates.
         area: One of :data:`AREAS`.
         signal: What the metric measures within its area, e.g. ``"task correctness"``.
-        scope: ``"dataset"`` metrics read episodes; ``"run"`` metrics read attempts and the run record.
+        scope: ``"episode"`` metrics summarize the episode checks over every episode; ``"dataset"`` metrics read
+            the episodes that passed them; ``"run"`` metrics read attempts and the run record.
         unit: Unit of the value.
+        sample: What the value's sample count counts, in the plural, e.g. ``"attempts"`` or ``"episodes"``.
         better: Which direction is better, if either.
         requires: Evidence keys; see :meth:`~autodata_metrics.evidence.Evidence.has`.
         primary: For a dict value, the entry a gate checks by default.
         version: Raise it whenever the definition changes, so gates pinned to the old one fail.
+        baseline: Also compute the metric on the source dataset.
     """
     assert area in AREAS, f"Metric {name!r} has area {area!r}; expected one of {AREAS}."
 
@@ -94,11 +113,36 @@ def register(compute: Callable[[Evidence], Value]) -> Callable[[Evidence], Value
             signal=signal,
             scope=scope,
             unit=unit,
+            sample=sample,
             better=better,
             requires=tuple(requires),
             primary=primary,
             version=version,
+            baseline=baseline,
         )
         return compute
 
     return register
+
+
+def check(
+    name: str, *, group: str = "correctness", requires: tuple[str, ...] | list[str] = ()
+) -> Callable[[Callable[[Episode, Evidence], str | None]], Callable[[Episode, Evidence], str | None]]:
+    """Register the decorated function as the episode check ``name``.
+
+    A check returns None for an episode that passes, a reason for one that fails, or raises :class:`NotApplicable`
+    for an episode it does not cover. Episodes that fail any check are left out of every dataset metric.
+
+    Args:
+        name: Unique check name.
+        group: One of :data:`CHECK_GROUPS`; the pass-rate metric the check counts toward.
+        requires: Evidence keys the check needs; without them it does not run.
+    """
+    assert group in CHECK_GROUPS, f"Check {name!r} has group {group!r}; expected one of {CHECK_GROUPS}."
+
+    def register(function: Callable[[Episode, Evidence], str | None]) -> Callable[[Episode, Evidence], str | None]:
+        assert name not in REGISTERED_CHECKS, f"Check {name!r} is already registered."
+        REGISTERED_CHECKS[name] = CheckSpec(name=name, check=function, group=group, requires=tuple(requires))
+        return function
+
+    return register