bohrin 0.1.0__py3-none-any.whl

bohrin/__init__.py

@@ -0,0 +1,37 @@
+"""bohrin — the health check-up for robot-learning datasets (Layer 1).
+
+Point it at your demonstrations and get a plain-language list of the hidden defects before
+you train. No simulator, no training, no ground truth, no upload — the data never leaves
+the machine.
+
+    import bohrin
+    report = bohrin.scan("./my_teleop_data")
+    print(report.score)
+
+See ``docs/`` for the full architecture. Phase 1: reads local LeRobot datasets (v2.1 + v3)
+and runs a conformally-calibrated detector battery over the frozen Canonical IR.
+"""
+
+from __future__ import annotations
+
+from bohrin.api import scan
+from bohrin.detectors.base import AnalysisContext, Detector, Requirements
+from bohrin.detectors.registry import register
+from bohrin.ir.schema import Family, Severity
+from bohrin.report.model import Cluster, Finding, Report
+from bohrin.version import REPORT_SCHEMA_VERSION, __version__
+
+__all__ = [
+    "REPORT_SCHEMA_VERSION",
+    "AnalysisContext",
+    "Cluster",
+    "Detector",
+    "Family",
+    "Finding",
+    "Report",
+    "Requirements",
+    "Severity",
+    "__version__",
+    "register",
+    "scan",
+]

bohrin/_arrays.py

@@ -0,0 +1,26 @@
+"""Shared numpy typing aliases used across the IR and detectors.
+
+Keeping these in one place lets ``mypy --strict`` see precise array element types instead
+of a bare ``np.ndarray`` (which is generic and would leak ``Any``).
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+
+# A 1-D or 2-D array of float64 (the canonical dtype for actions/proprio after ingest).
+FloatArray = npt.NDArray[np.float64]
+
+# Integer arrays (episode indices, counts).
+IntArray = npt.NDArray[np.int64]
+
+# Boolean masks (anomaly flags).
+BoolArray = npt.NDArray[np.bool_]
+
+# An array of *unknown* dtype — only for the adapter boundary, where a source file may hold
+# uint8 pixels, int64 indices or float32 signals and the adapter's job is to coerce them.
+# Nothing past Stage ② should use this: the IR is float64 by contract.
+AnyArray = npt.NDArray[Any]

bohrin/_compat.py

@@ -0,0 +1,39 @@
+"""Small standard-library shims so the package runs on Python 3.10.
+
+A large share of the LeRobot / ROS / CUDA world is still pinned to 3.10-3.11, and a
+`requires-python` floor they cannot meet fails at ``pip install`` -- silently, from our
+point of view, because the user never files an issue. The floor is therefore 3.10 and
+this module carries the (very small) cost of that.
+
+``StrEnum`` (3.11+) and ``typing.Self`` (3.11+) are the only two constructs from later
+Python that the rest of the package uses.
+
+The ``StrEnum`` fallback reproduces the two behaviours we actually depend on: members
+compare equal to their ``str`` value, and ``str(member)`` yields that value rather than
+``"Class.MEMBER"`` (the plain ``str, Enum`` mixin does the latter on 3.10, which would
+corrupt every f-string in the report layer). Explicit methods, not assigned-from-``str``
+dunders, because mypy types ``str.__format__`` as unbound and rejects the assignment.
+"""
+
+from __future__ import annotations
+
+import enum
+import sys
+
+if sys.version_info >= (3, 11):
+    StrEnum = enum.StrEnum
+    from typing import Self as Self
+else:  # pragma: no cover - exercised by the 3.10 CI leg
+    from typing_extensions import Self as Self
+
+    class StrEnum(str, enum.Enum):
+        """Backport of :class:`enum.StrEnum` (Python 3.11+)."""
+
+        def __str__(self) -> str:
+            return str.__str__(self)
+
+        def __format__(self, format_spec: str) -> str:
+            return str.__format__(self, format_spec)
+
+
+__all__ = ["Self", "StrEnum"]

bohrin/_plugins.py

@@ -0,0 +1,39 @@
+"""Entry-point plugin discovery (docs/02 §10).
+
+The standard-library ``importlib.metadata`` route (non-provisional since 3.10). Both
+built-in adapters/detectors and third-party ones advertise themselves through the same
+entry-point groups — there is no privileged path. A plugin whose ``.load()`` raises is
+skipped with a warning rather than crashing the whole scan.
+"""
+
+from __future__ import annotations
+
+import warnings
+from importlib.metadata import entry_points
+
+
+def load_plugin_classes(group: str) -> dict[str, type]:
+    """Load every *class* advertised under ``group``, keyed by entry-point name.
+
+    Entries that fail to import, or that resolve to a non-class, are warned about and
+    skipped so one bad plugin can never take down the tool. Callers narrow the result to
+    the base they expect (adapter / detector) via ``issubclass``.
+    """
+    found: dict[str, type] = {}
+    for ep in entry_points(group=group):
+        try:
+            obj = ep.load()
+        except Exception as exc:  # a plugin must never crash discovery
+            warnings.warn(
+                f"bohrin: failed to load plugin {ep.name!r} from {group!r}: {exc}",
+                stacklevel=2,
+            )
+            continue
+        if not isinstance(obj, type):
+            warnings.warn(
+                f"bohrin: plugin {ep.name!r} in {group!r} is not a class; skipping",
+                stacklevel=2,
+            )
+            continue
+        found[ep.name] = obj
+    return found

bohrin/adapters/__init__.py

@@ -0,0 +1,21 @@
+"""Stage ① — the adapter layer (docs/02 §1, docs/01_DATA_LANDSCAPE.md)."""
+
+from __future__ import annotations
+
+from bohrin.adapters.base import Adapter, DatasetHandle, Sampler
+from bohrin.adapters.registry import (
+    UnknownFormatError,
+    discover,
+    register_adapter,
+    select_adapter,
+)
+
+__all__ = [
+    "Adapter",
+    "DatasetHandle",
+    "Sampler",
+    "UnknownFormatError",
+    "discover",
+    "register_adapter",
+    "select_adapter",
+]

bohrin/adapters/_arraysource.py

@@ -0,0 +1,198 @@
+"""Shared machinery for adapters whose source is "a bag of named arrays".
+
+HDF5, NumPy directories and Zarr replay buffers differ in *how* you list and slice arrays,
+but not in what happens next: run the schema mapper, slice per episode, wrap in a
+:class:`StepView`. That common half lives here so the three adapters stay thin and the IR
+construction is written — and tested — exactly once.
+
+The contract a concrete adapter implements is :class:`EpisodeArrays`: given an episode
+index, hand back its arrays. Everything else (mapping, dtype coercion, camera specs,
+provenance) is done for it.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator, Mapping, Sequence
+from dataclasses import dataclass
+from typing import Protocol
+
+import numpy as np
+
+from bohrin._arrays import AnyArray, FloatArray
+from bohrin.adapters._mapping import ArrayInfo, SchemaMapping, infer_mapping
+from bohrin.adapters.base import Sampler
+from bohrin.ir.episode import Episode, StepView, TaskLabel
+from bohrin.ir.schema import ActionSpace, CameraSpec, DatasetSchema, Provenance, SchemaHints
+
+
+class EpisodeArrays(Protocol):
+    """A source of per-episode named arrays."""
+
+    def episode_keys(self) -> Sequence[str]:
+        """Stable identifiers for each episode, in order."""
+        ...
+
+    def arrays(self, episode_key: str) -> Mapping[str, AnyArray]:
+        """The named arrays for one episode. Called lazily, one episode at a time."""
+        ...
+
+    def task_for(self, episode_key: str) -> str | None:
+        """The language instruction for this episode, if the format carries one."""
+        ...
+
+
+@dataclass(frozen=True, slots=True)
+class _ImageFrames:
+    """A lazy per-frame view over a ``(T, H, W[, C])`` array — satisfies ``LazyImage``."""
+
+    data: AnyArray
+    index: int
+
+    @property
+    def shape(self) -> tuple[int, int, int]:
+        frame = self.data[self.index]
+        if frame.ndim == 2:
+            return (int(frame.shape[0]), int(frame.shape[1]), 1)
+        return (int(frame.shape[0]), int(frame.shape[1]), int(frame.shape[2]))
+
+    def array(self) -> FloatArray:
+        frame = np.asarray(self.data[self.index], dtype=np.float64)
+        return frame if frame.ndim == 3 else frame[:, :, None]
+
+
+def _as_2d(arr: AnyArray) -> FloatArray:
+    """Coerce a per-step signal to ``(T, D)`` float64 — the IR's low-dimensional shape."""
+    out = np.asarray(arr, dtype=np.float64)
+    if out.ndim == 1:
+        return out.reshape(-1, 1)
+    if out.ndim > 2:  # flatten trailing dims, e.g. (T, 2, 3) poses → (T, 6)
+        return out.reshape(out.shape[0], -1)
+    return out
+
+
+def _as_1d(arr: AnyArray) -> FloatArray:
+    out = np.asarray(arr, dtype=np.float64).ravel()
+    return out
+
+
+def build_schema(
+    mapping: SchemaMapping,
+    sample: Mapping[str, AnyArray],
+    *,
+    control_hz: float | None,
+    embodiment: str | None,
+) -> DatasetSchema:
+    """Derive the dataset-wide schema from one representative episode."""
+    action = _as_2d(sample[mapping.action])
+    proprio = _as_2d(sample[mapping.proprio]) if mapping.proprio is not None and mapping.proprio in sample else None
+    cameras = tuple(
+        CameraSpec(key=key, height=int(sample[key].shape[1]), width=int(sample[key].shape[2]))
+        for key in mapping.images
+        if key in sample and sample[key].ndim >= 3
+    )
+    return DatasetSchema(
+        action_dim=int(action.shape[1]),
+        action_space=ActionSpace.UNKNOWN,  # custom containers rarely declare it; stay honest
+        proprio_dim=None if proprio is None else int(proprio.shape[1]),
+        cameras=cameras,
+        control_hz=control_hz,
+        embodiment=embodiment,
+    )
+
+
+def infer_control_hz(sample: Mapping[str, AnyArray], mapping: SchemaMapping) -> float | None:
+    """Estimate the control rate from timestamps, if the source has any."""
+    if mapping.timestamp is None or mapping.timestamp not in sample:
+        return None
+    ts = _as_1d(sample[mapping.timestamp])
+    if ts.size < 2:
+        return None
+    dt = float(np.median(np.diff(ts)))
+    return 1.0 / dt if dt > 0 else None
+
+
+class ArraySourceHandle:
+    """A :class:`DatasetHandle` over any :class:`EpisodeArrays` source."""
+
+    def __init__(
+        self,
+        source: EpisodeArrays,
+        *,
+        adapter_name: str,
+        uri: str,
+        declared: Mapping[str, object] | None = None,
+        embodiment: str | None = None,
+        no_vision: bool = False,
+        hints: SchemaHints | None = None,
+        splits: Mapping[str, Sequence[str]] | None = None,
+    ) -> None:
+        self._source = source
+        self._adapter_name = adapter_name
+        self._uri = uri
+        self._no_vision = no_vision
+        # episode key -> declared split name, inverted once so lookup is O(1) per episode.
+        self._split_of: dict[str, str] = {key: name for name, members in (splits or {}).items() for key in members}
+        self._hints = hints or SchemaHints.empty()
+        self._keys = list(source.episode_keys())
+        if not self._keys:
+            raise ValueError(f"{adapter_name}: no episodes found in {uri}")
+
+        first = source.arrays(self._keys[0])
+        infos = [ArrayInfo(key=k, shape=tuple(int(d) for d in v.shape)) for k, v in first.items()]
+        self._mapping = infer_mapping(infos, declared)
+        hz = infer_control_hz(first, self._mapping)
+        self._schema = build_schema(self._mapping, first, control_hz=hz, embodiment=embodiment)
+
+    @property
+    def mapping(self) -> SchemaMapping:
+        """The resolved role→key mapping (exposed for ``bohrin init`` and tests)."""
+        return self._mapping
+
+    def schema(self) -> DatasetSchema:
+        return self._schema
+
+    def profile_hints(self) -> SchemaHints:
+        return self._hints
+
+    def episode_count(self) -> int | None:
+        return len(self._keys)
+
+    def iter_episodes(self, *, sample: Sampler) -> Iterator[Episode]:
+        keep = set(sample.plan(len(self._keys)).tolist())
+        for index, key in enumerate(self._keys):
+            if index not in keep:
+                continue
+            yield self._episode(index, key)
+
+    def _episode(self, index: int, key: str) -> Episode:
+        arrays = self._source.arrays(key)
+        m = self._mapping
+        action = _as_2d(arrays[m.action])
+        n_steps = action.shape[0]
+
+        images: dict[str, Sequence[_ImageFrames]] = {}
+        depth: dict[str, Sequence[_ImageFrames]] = {}
+        if not self._no_vision:
+            for cam in m.images:
+                if cam in arrays:
+                    images[cam] = [_ImageFrames(arrays[cam], t) for t in range(min(n_steps, len(arrays[cam])))]
+            for cam in m.depth:
+                if cam in arrays:
+                    depth[cam] = [_ImageFrames(arrays[cam], t) for t in range(min(n_steps, len(arrays[cam])))]
+
+        steps = StepView(
+            action=action,
+            timestamp=_as_1d(arrays[m.timestamp]) if m.timestamp and m.timestamp in arrays else None,
+            proprio=_as_2d(arrays[m.proprio]) if m.proprio and m.proprio in arrays else None,
+            reward=_as_1d(arrays[m.reward]) if m.reward and m.reward in arrays else None,
+            images=images,
+            depth=depth,
+        )
+        task = self._source.task_for(key)
+        return Episode(
+            episode_id=key,
+            steps=steps,
+            source=Provenance(adapter=self._adapter_name, uri=self._uri, locator=key),
+            task=None if task is None else TaskLabel(text=task),
+            split=self._split_of.get(key),
+        )

bohrin/adapters/_mapping.py

@@ -0,0 +1,213 @@
+"""The schema mapper — how a *custom* container becomes Canonical IR (docs/02 §1.3).
+
+LeRobot and RLDS declare their own schema, so their adapters read it. Everything else —
+raw HDF5, a directory of ``.npz``, a Zarr replay buffer — is a bag of arrays with names
+chosen by whoever recorded it. This module answers one question for those formats:
+
+    given these array names and shapes, which one is the action, which the proprioception,
+    which the timestamps, and which are cameras?
+
+Two sources of truth, in priority order:
+
+1. **Declared** — a ``bohrin.yaml`` written by ``bohrin init``. Always wins. This is the
+   long-tail escape hatch: if inference is wrong, the user states the mapping once.
+2. **Inferred** — name matching against the conventions actually used across the public
+   corpora (robomimic, Diffusion Policy/UMI, DROID, and the ad-hoc layouts in between),
+   then a shape-based tie-break.
+
+Inference is deliberately **conservative**: when nothing matches confidently we return
+``None`` and let the caller raise, rather than silently profiling the wrong array. A
+detector battery pointed at the wrong column produces confident nonsense, which is worse
+than an error message telling the user to run ``bohrin init``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+
+#: Candidate names for each role, most specific first. Matching is case-insensitive and
+#: compares both the full key and its last path segment, so ``data/demo_0/actions`` and
+#: ``actions`` both resolve. Sources: robomimic (``actions``, ``obs/``), Diffusion Policy
+#: and UMI replay buffers (``action``, ``state``, ``robot0_eef_pos``), DROID/RLDS
+#: (``action``, ``observation/state``).
+_ACTION_NAMES = ("action", "actions", "act", "action_dict", "cmd", "command")
+_PROPRIO_NAMES = (
+    "observation.state",
+    "observation/state",
+    "proprio",
+    "proprioception",
+    "state",
+    "states",
+    "qpos",
+    "joint_positions",
+    "joint_position",
+    "robot_state",
+    "eef_pose",
+    "robot0_eef_pos",
+    "obs",
+)
+_TIMESTAMP_NAMES = ("timestamp", "timestamps", "time", "t", "stamp", "time_stamp")
+_REWARD_NAMES = ("reward", "rewards", "r")
+#: Substrings that mark an array as pixels rather than a low-dimensional signal.
+_IMAGE_HINTS = ("image", "rgb", "camera", "cam", "img", "pixels", "wrist", "front", "side", "top")
+_DEPTH_HINTS = ("depth", "disparity", "pointcloud", "point_cloud", "xyz")
+
+#: An array is treated as pixels if it has this many dims (T, H, W[, C]) and a plausible
+#: channel count. Shape is the tie-break when the *name* is uninformative.
+_MIN_IMAGE_NDIM = 3
+_IMAGE_CHANNELS = (1, 3, 4)
+
+
+@dataclass(frozen=True, slots=True)
+class ArrayInfo:
+    """What the mapper needs to know about one candidate array: its name and shape."""
+
+    key: str
+    shape: tuple[int, ...]
+
+    @property
+    def leaf(self) -> str:
+        """The last path segment, lowercased — ``data/demo_0/actions`` → ``actions``."""
+        return self.key.replace("\\", "/").rsplit("/", 1)[-1].lower()
+
+    @property
+    def ndim(self) -> int:
+        return len(self.shape)
+
+
+@dataclass(frozen=True, slots=True)
+class SchemaMapping:
+    """The resolved answer: which key plays which role."""
+
+    action: str
+    proprio: str | None = None
+    timestamp: str | None = None
+    reward: str | None = None
+    images: tuple[str, ...] = ()
+    depth: tuple[str, ...] = ()
+
+    @property
+    def used_keys(self) -> frozenset[str]:
+        """Every key this mapping claims — the caller may treat the rest as unused."""
+        named = {self.action, self.proprio, self.timestamp, self.reward}
+        return frozenset({k for k in named if k} | set(self.images) | set(self.depth))
+
+
+class UnmappableDatasetError(ValueError):
+    """Raised when no array can be identified as the action — the one required column.
+
+    Carries the candidate keys, because the fix is always "tell me which one it is".
+    """
+
+    def __init__(self, keys: Sequence[str]) -> None:
+        listed = ", ".join(sorted(keys)[:12]) or "(none)"
+        super().__init__(
+            "could not identify the action array in this dataset. "
+            f"Candidate arrays: {listed}. "
+            "Run `bohrin init <path>` to declare the mapping in a bohrin.yaml, "
+            "or pass --format to select a different adapter."
+        )
+
+
+def _match(info: ArrayInfo, names: Sequence[str]) -> int | None:
+    """Rank of the first matching name (lower is better), or ``None``."""
+    leaf, full = info.leaf, info.key.lower()
+    for rank, name in enumerate(names):
+        if leaf == name or full == name or full.endswith("/" + name):
+            return rank
+    return None
+
+
+def _is_image(info: ArrayInfo) -> bool:
+    if any(h in info.key.lower() for h in _DEPTH_HINTS):
+        return False
+    if any(h in info.key.lower() for h in _IMAGE_HINTS):
+        return info.ndim >= _MIN_IMAGE_NDIM
+    # Name says nothing — fall back to shape: (T, H, W) or (T, H, W, C).
+    if info.ndim == 4:
+        return info.shape[-1] in _IMAGE_CHANNELS
+    return False
+
+
+def _is_depth(info: ArrayInfo) -> bool:
+    return any(h in info.key.lower() for h in _DEPTH_HINTS) and info.ndim >= _MIN_IMAGE_NDIM
+
+
+def _pick(arrays: Sequence[ArrayInfo], names: Sequence[str], *, exclude: frozenset[str]) -> str | None:
+    """The best-matching key for a role, or ``None`` if nothing matches."""
+    ranked: list[tuple[int, str]] = []
+    for a in arrays:
+        if a.key in exclude:
+            continue
+        rank = _match(a, names)
+        if rank is not None:
+            ranked.append((rank, a.key))
+    return min(ranked)[1] if ranked else None
+
+
+def _widest_2d(arrays: Sequence[ArrayInfo], *, exclude: frozenset[str]) -> str | None:
+    """The widest ``(T, D)`` array — the shape-based fallback for the action column."""
+    candidates = [a for a in arrays if a.ndim == 2 and a.key not in exclude and not _is_image(a)]
+    if not candidates:
+        return None
+    return max(candidates, key=lambda a: (a.shape[1], a.key)).key
+
+
+def infer_mapping(
+    arrays: Sequence[ArrayInfo],
+    declared: Mapping[str, object] | None = None,
+    *,
+    allow_shape_fallback: bool = True,
+) -> SchemaMapping:
+    """Resolve array names to IR roles. Declared entries always beat inference.
+
+    ``declared`` is the ``schema_map`` section of a ``bohrin.yaml``; recognized keys are
+    ``action``, ``proprio``, ``timestamp``, ``reward``, ``images`` and ``depth``.
+
+    Raises :class:`UnmappableDatasetError` if the action column cannot be identified —
+    guessing it would mean profiling an arbitrary array and reporting the result as fact.
+    """
+    declared = declared or {}
+    present = {a.key for a in arrays}
+
+    def declared_str(role: str) -> str | None:
+        value = declared.get(role)
+        return str(value) if isinstance(value, str) else None
+
+    def declared_seq(role: str) -> tuple[str, ...]:
+        value = declared.get(role)
+        if isinstance(value, str):
+            return (value,)
+        if isinstance(value, (list, tuple)):
+            return tuple(str(v) for v in value)
+        return ()
+
+    action = declared_str("action") or _pick(arrays, _ACTION_NAMES, exclude=frozenset())
+    if action is None and allow_shape_fallback:
+        action = _widest_2d(arrays, exclude=frozenset())
+    if action is None or action not in present:
+        raise UnmappableDatasetError(sorted(present))
+
+    claimed = {action}
+    proprio = declared_str("proprio") or _pick(arrays, _PROPRIO_NAMES, exclude=frozenset(claimed))
+    if proprio:
+        claimed.add(proprio)
+    timestamp = declared_str("timestamp") or _pick(arrays, _TIMESTAMP_NAMES, exclude=frozenset(claimed))
+    if timestamp:
+        claimed.add(timestamp)
+    reward = declared_str("reward") or _pick(arrays, _REWARD_NAMES, exclude=frozenset(claimed))
+    if reward:
+        claimed.add(reward)
+
+    images = declared_seq("images") or tuple(sorted(a.key for a in arrays if a.key not in claimed and _is_image(a)))
+    depth = declared_seq("depth") or tuple(sorted(a.key for a in arrays if a.key not in claimed and _is_depth(a)))
+
+    return SchemaMapping(
+        action=action,
+        proprio=proprio,
+        timestamp=timestamp,
+        reward=reward,
+        images=images,
+        depth=depth,
+    )