From d73bc7c0f6d2c93e24aa40931e8e7e776bf94784 Mon Sep 17 00:00:00 2001 From: jaynye Date: Wed, 2 Sep 2026 20:45:57 -0400 Subject: [PATCH 1/9] feat(calibration): per-episode calibration block with a legacy read shim Calibration is a measurement of the rig that recorded one episode, so it has to travel with the episode instead of living in a class constant. Add the `calibration` attribute block and the single reader every consumer goes through. The block names one reference frame and expresses every pose in it: `cameras[c].ref_T_cam` is a camera pose, `arm_bases[side]` is `ref_T_armbase`. `Calibration.base_T_cam(side)` composes the two into what the EVA transform pipeline consumes. No stored episode needs a rewrite. `read_calibration` prefers the block and falls back to lifting the legacy `intrinsics`/`extrinsics` pair, whose per-arm values are the front camera's pose in each arm base, so the lifted reference frame is `camera:front_1` and `arm_bases` is their inverse. `ZarrWriter` writes the block and derives the legacy pair from it, so readers that predate the block keep working. `ZarrEpisode`, `ZarrDataset` and the inspector now read through the shim rather than parsing the attributes themselves, which also drops two duplicate 3x3-to-3x4 normalizations. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_012rCh1QmnyqgPZ5TJpsfMQb --- docs/CONVENTIONS.md | 20 + egomimic/rldb/zarr/calibration.py | 477 ++++++++++++++++++ egomimic/rldb/zarr/test_calibration.py | 232 +++++++++ egomimic/rldb/zarr/zarr_dataset_multi.py | 58 ++- egomimic/rldb/zarr/zarr_writer.py | 50 +- .../inspector_lib/dataset_view.py | 60 ++- 6 files changed, 830 insertions(+), 67 deletions(-) create mode 100644 egomimic/rldb/zarr/calibration.py create mode 100644 egomimic/rldb/zarr/test_calibration.py diff --git a/docs/CONVENTIONS.md b/docs/CONVENTIONS.md index 7506ac3df..b7b0a0cf0 100644 --- a/docs/CONVENTIONS.md +++ b/docs/CONVENTIONS.md @@ -22,6 +22,26 @@ EVA episode extrinsics store one `base_T_cam` matrix for each arm. Human episodes store the head pose as `world_T_head`. The head frame and the camera frame are the same frame for egocentric human data. +### The calibration block + +An episode written after the `calibration` attribute existed names one +reference frame and expresses every pose in it. + +```text +calibration.reference_frame robot_base | slam_world | camera: +calibration.cameras[c].ref_T_cam camera c's pose in the reference frame +calibration.arm_bases[side] ref_T_armbase, the arm base pose in it +``` + +The camera that defines the reference frame needs no `ref_T_cam`. It is the +identity by definition. + +`Calibration.base_T_cam(side)` composes the two and returns what the EVA +transform pipeline consumes. An episode that predates the block reaches the +same value through the shim in `egomimic/rldb/zarr/calibration.py`: its +reference frame is `camera:front_1`, so `arm_bases[side]` is the inverse of +the stored `extrinsics[side]`. + ### Existing episodes need no migration The `ref_T_cam` naming records the direction the code already used. It changes diff --git a/egomimic/rldb/zarr/calibration.py b/egomimic/rldb/zarr/calibration.py new file mode 100644 index 000000000..abe9407e8 --- /dev/null +++ b/egomimic/rldb/zarr/calibration.py @@ -0,0 +1,477 @@ +"""Read and write the per-episode ``calibration`` metadata block. + +Calibration is a physical measurement of the rig that recorded one episode, so +it travels with the episode instead of living in a class constant. An episode +stores it under one ``calibration`` attribute:: + + "calibration": { + "reference_frame": "camera:front_1", + "cameras": { + "front_1": { + "K": [[fx, 0, cx, 0], [0, fy, cy, 0], [0, 0, 1, 0]], + "resolution": [W, H], + "rectified": true, + "ref_T_cam": [[...]], + }, + }, + "arm_bases": {"left": [[...]], "right": [[...]]}, + } + +Every matrix follows ``docs/CONVENTIONS.md``: ``A_T_B`` maps coordinates from +frame ``B`` to frame ``A``. ``ref_T_cam`` is the camera pose in the reference +frame, and ``arm_bases[side]`` is ``ref_T_armbase``, the arm base pose in the +reference frame. + +Episodes written before this block existed store ``intrinsics`` and +``extrinsics`` at the top level instead. :func:`read_calibration` reads either +form, so no stored episode needs a rewrite. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any + +import numpy as np + +#: Camera that legacy episodes calibrate and express `extrinsics` against. +LEGACY_REFERENCE_CAMERA = "front_1" + +#: Reference frames that do not name a camera. +STATIC_REFERENCE_FRAMES = frozenset({"robot_base", "slam_world"}) + +#: Prefix that makes a camera the reference frame, as in `camera:front_1`. +CAMERA_FRAME_PREFIX = "camera:" + +_CAMERA_FIELDS = frozenset({"K", "resolution", "rectified", "ref_T_cam"}) +_CALIBRATION_FIELDS = frozenset({"reference_frame", "cameras", "arm_bases"}) + + +class CalibrationError(ValueError): + """Report an invalid ``calibration`` block or legacy calibration attribute.""" + + +def _matrix(value: Any, shape: tuple[int, ...], where: str) -> np.ndarray: + """Return ``value`` as a finite float64 array of the given shape.""" + try: + arr = np.asarray(value, dtype=np.float64) + except (TypeError, ValueError) as exc: + raise CalibrationError(f"{where}: not a numeric array ({exc})") from exc + if arr.shape != shape: + raise CalibrationError( + f"{where}: expected shape {shape}, got {arr.shape}" + ) + if not np.isfinite(arr).all(): + raise CalibrationError(f"{where}: contains a non-finite value") + return arr + + +def _camera_matrix(value: Any, where: str) -> np.ndarray: + """Return a 3×4 camera matrix, padding a bare 3×3 with a zero column.""" + try: + arr = np.asarray(value, dtype=np.float64) + except (TypeError, ValueError) as exc: + raise CalibrationError(f"{where}: not a numeric array ({exc})") from exc + if arr.shape == (3, 3): + arr = np.hstack([arr, np.zeros((3, 1))]) + if arr.shape != (3, 4): + raise CalibrationError( + f"{where}: expected a 3x4 camera matrix (pad a bare 3x3 with " + f"np.hstack([K, np.zeros((3, 1))])), got shape {arr.shape}" + ) + if not np.isfinite(arr).all(): + raise CalibrationError(f"{where}: contains a non-finite value") + return arr + + +@dataclass(frozen=True) +class CameraCalibration: + """Store the calibration of one camera stream. + + Attributes: + name: The camera name. It matches the ``images.`` array key. + K: The 3×4 camera matrix, or ``None`` for a declared but uncalibrated + camera. + resolution: ``(width, height)`` in pixels, or ``None``. + rectified: Whether the stored frames are already rectified. + ref_T_cam: The 4×4 camera pose in the episode reference frame, or + ``None`` when the episode does not state it. + """ + + name: str + K: np.ndarray | None = None + resolution: tuple[int, int] | None = None + rectified: bool = True + ref_T_cam: np.ndarray | None = None + + def to_jsonable(self) -> dict[str, Any]: + """Return this camera as JSON-serializable episode metadata.""" + out: dict[str, Any] = {"rectified": bool(self.rectified)} + if self.K is not None: + out["K"] = self.K.tolist() + if self.resolution is not None: + out["resolution"] = [int(self.resolution[0]), int(self.resolution[1])] + if self.ref_T_cam is not None: + out["ref_T_cam"] = self.ref_T_cam.tolist() + return out + + +@dataclass(frozen=True) +class Calibration: + """Store the calibration of every camera and arm base in one episode. + + Attributes: + reference_frame: ``robot_base``, ``slam_world``, or + ``camera:``. Every pose in this block is expressed in + this frame. + cameras: A mapping from camera name to its calibration. + arm_bases: A mapping from ``"left"`` or ``"right"`` to a 4×4 + ``ref_T_armbase`` matrix. Human episodes leave this empty. + legacy: Whether :func:`read_calibration` built this from the + ``intrinsics`` and ``extrinsics`` attributes of an older episode. + """ + + reference_frame: str + cameras: Mapping[str, CameraCalibration] = field(default_factory=dict) + arm_bases: Mapping[str, np.ndarray] = field(default_factory=dict) + legacy: bool = False + + @property + def reference_camera(self) -> str | None: + """Return the camera that defines the reference frame, if any.""" + if self.reference_frame.startswith(CAMERA_FRAME_PREFIX): + return self.reference_frame[len(CAMERA_FRAME_PREFIX) :] + return None + + def default_camera(self) -> str | None: + """Return the camera to project against when a caller names none. + + The reference camera wins, then ``front_1``, then any camera whose name + contains ``front``, then the first declared camera. Projection and + visualization target the front image, so a front camera outranks a + wrist camera that happens to be declared first. + """ + reference = self.reference_camera + if reference is not None and reference in self.cameras: + return reference + if LEGACY_REFERENCE_CAMERA in self.cameras: + return LEGACY_REFERENCE_CAMERA + front = next((n for n in self.cameras if "front" in n.lower()), None) + if front is not None: + return front + return next(iter(self.cameras), None) + + def K(self, camera: str | None = None) -> np.ndarray | None: + """Return the 3×4 camera matrix of one camera. + + Args: + camera: A camera name. ``None`` selects :meth:`default_camera`. + + Returns: + The camera matrix, or ``None`` if the camera is absent or carries + no ``K``. + """ + name = camera if camera is not None else self.default_camera() + if name is None: + return None + entry = self.cameras.get(name) + return None if entry is None else entry.K + + def ref_T_cam(self, camera: str | None = None) -> np.ndarray | None: + """Return the pose of one camera in the reference frame. + + The reference camera is the identity by definition, so an episode that + references a camera need not state that camera's pose. + """ + name = camera if camera is not None else self.default_camera() + if name is None: + return None + entry = self.cameras.get(name) + if entry is not None and entry.ref_T_cam is not None: + return entry.ref_T_cam + if name == self.reference_camera: + return np.eye(4) + return None + + def base_T_cam(self, side: str, camera: str | None = None) -> np.ndarray | None: + """Return one camera's pose in the frame of one arm base. + + This is the quantity the EVA transform pipeline consumes, and it is + what the ``extrinsics`` attribute of a legacy episode stores directly. + + Args: + side: ``"left"`` or ``"right"``. + camera: A camera name. ``None`` selects :meth:`default_camera`. + + Returns: + A 4×4 ``armbase_T_cam`` matrix, or ``None`` when the episode + states no arm base or no camera pose. + """ + ref_T_armbase = self.arm_bases.get(side) + ref_T_cam = self.ref_T_cam(camera) + if ref_T_armbase is None or ref_T_cam is None: + return None + return np.linalg.inv(ref_T_armbase) @ ref_T_cam + + def intrinsics(self) -> dict[str, np.ndarray]: + """Return ``{camera: K}`` for every camera that carries a ``K``.""" + return { + name: cam.K for name, cam in self.cameras.items() if cam.K is not None + } + + def extrinsics(self, camera: str | None = None) -> dict[str, np.ndarray]: + """Return ``{side: base_T_cam}`` in the legacy attribute layout.""" + out = {} + for side in self.arm_bases: + base_T_cam = self.base_T_cam(side, camera) + if base_T_cam is not None: + out[side] = base_T_cam + return out + + def uncalibrated(self, cameras) -> list[str]: + """Return the given camera names that carry no ``K``, in order.""" + return [name for name in cameras if self.K(name) is None] + + def to_jsonable(self) -> dict[str, Any]: + """Return this calibration as JSON-serializable episode metadata.""" + out: dict[str, Any] = { + "reference_frame": self.reference_frame, + "cameras": { + name: cam.to_jsonable() for name, cam in self.cameras.items() + }, + } + if self.arm_bases: + out["arm_bases"] = { + side: np.asarray(T, dtype=np.float64).tolist() + for side, T in self.arm_bases.items() + } + return out + + +def _parse_reference_frame(value: Any, cameras, where: str) -> str: + if not isinstance(value, str) or not value: + raise CalibrationError(f"{where}: `reference_frame` must be a non-empty string") + if value in STATIC_REFERENCE_FRAMES: + return value + if value.startswith(CAMERA_FRAME_PREFIX): + camera = value[len(CAMERA_FRAME_PREFIX) :] + if camera and camera in cameras: + return value + raise CalibrationError( + f"{where}: reference_frame {value!r} names a camera that this episode " + f"does not declare; declared cameras are {sorted(cameras)}" + ) + raise CalibrationError( + f"{where}: unknown reference_frame {value!r}; expected one of " + f"{sorted(STATIC_REFERENCE_FRAMES)} or 'camera:'" + ) + + +def _parse_camera(name: str, block: Any, where: str) -> CameraCalibration: + if not isinstance(block, Mapping): + raise CalibrationError(f"{where}: must be a mapping, got {block!r}") + unknown = sorted(set(block) - _CAMERA_FIELDS) + if unknown: + raise CalibrationError( + f"{where}: unknown field(s) {unknown}; expected some of " + f"{sorted(_CAMERA_FIELDS)}" + ) + + K = block.get("K") + resolution = block.get("resolution") + if resolution is not None: + if ( + not isinstance(resolution, (list, tuple)) + or len(resolution) != 2 + or not all(isinstance(v, int) and v > 0 for v in resolution) + ): + raise CalibrationError( + f"{where}.resolution: expected [width, height] positive ints, " + f"got {resolution!r}" + ) + resolution = (int(resolution[0]), int(resolution[1])) + + rectified = block.get("rectified", True) + if not isinstance(rectified, bool): + raise CalibrationError( + f"{where}.rectified: expected a boolean, got {rectified!r}" + ) + + ref_T_cam = block.get("ref_T_cam") + return CameraCalibration( + name=name, + K=None if K is None else _camera_matrix(K, f"{where}.K"), + resolution=resolution, + rectified=rectified, + ref_T_cam=( + None + if ref_T_cam is None + else _matrix(ref_T_cam, (4, 4), f"{where}.ref_T_cam") + ), + ) + + +def parse_calibration(block: Any, where: str = "calibration") -> Calibration: + """Validate one ``calibration`` attribute block. + + Args: + block: The mapping stored under ``zarr.attrs["calibration"]``. + where: A label used in error messages. + + Returns: + The parsed calibration. + + Raises: + CalibrationError: If a field is unknown, missing, or malformed. + """ + if isinstance(block, Calibration): + return block + if not isinstance(block, Mapping): + raise CalibrationError(f"{where}: must be a mapping, got {block!r}") + unknown = sorted(set(block) - _CALIBRATION_FIELDS) + if unknown: + raise CalibrationError( + f"{where}: unknown field(s) {unknown}; expected some of " + f"{sorted(_CALIBRATION_FIELDS)}" + ) + + raw_cameras = block.get("cameras") + if not isinstance(raw_cameras, Mapping) or not raw_cameras: + raise CalibrationError( + f"{where}.cameras: expected a non-empty mapping from camera name to " + f"its calibration, got {raw_cameras!r}" + ) + cameras = { + name: _parse_camera(name, cam, f"{where}.cameras[{name!r}]") + for name, cam in raw_cameras.items() + } + + if "reference_frame" not in block: + raise CalibrationError(f"{where}: missing required field 'reference_frame'") + reference_frame = _parse_reference_frame( + block["reference_frame"], cameras, where + ) + + raw_arm_bases = block.get("arm_bases") or {} + if not isinstance(raw_arm_bases, Mapping): + raise CalibrationError( + f"{where}.arm_bases: expected a mapping from side to a 4x4 " + f"ref_T_armbase matrix, got {raw_arm_bases!r}" + ) + arm_bases = { + side: _matrix(T, (4, 4), f"{where}.arm_bases[{side!r}]") + for side, T in raw_arm_bases.items() + } + + return Calibration( + reference_frame=reference_frame, cameras=cameras, arm_bases=arm_bases + ) + + +def lift_legacy_calibration( + intrinsics: Any = None, extrinsics: Any = None +) -> Calibration | None: + """Build a :class:`Calibration` from the legacy attribute pair. + + ``intrinsics`` maps a camera name to its ``K``. A bare matrix is read as + the ``front_1`` camera, which is the only camera any legacy writer + calibrated. ``extrinsics`` maps an arm side to ``base_T_cam``, the front + camera's pose in that arm's base frame, so the reference frame is the front + camera and ``arm_bases[side]`` is the inverse of the stored matrix. + + Args: + intrinsics: The ``intrinsics`` attribute, or ``None``. + extrinsics: The ``extrinsics`` attribute, or ``None``. + + Returns: + The lifted calibration, or ``None`` if both attributes are absent. + + Raises: + CalibrationError: If either attribute is malformed. + """ + if not intrinsics and not extrinsics: + return None + + if isinstance(intrinsics, Mapping): + raw_cameras = list(intrinsics.items()) + elif intrinsics is not None: + raw_cameras = [(LEGACY_REFERENCE_CAMERA, intrinsics)] + else: + raw_cameras = [] + + cameras = { + str(name): CameraCalibration( + name=str(name), K=_camera_matrix(K, f"intrinsics[{name!r}]") + ) + for name, K in raw_cameras + } + + arm_bases: dict[str, np.ndarray] = {} + if extrinsics is not None: + if not isinstance(extrinsics, Mapping): + raise CalibrationError( + f"extrinsics: expected a mapping from arm side to a 4x4 " + f"base_T_cam matrix, got {extrinsics!r}" + ) + for side, base_T_cam in extrinsics.items(): + matrix = _matrix(base_T_cam, (4, 4), f"extrinsics[{side!r}]") + arm_bases[str(side)] = np.linalg.inv(matrix) + + # The legacy `extrinsics` values are per-arm poses of the front camera, so + # that camera is the reference frame. Declare it even when the episode + # calibrated no camera, otherwise `base_T_cam` has no pose to compose. + if arm_bases and LEGACY_REFERENCE_CAMERA not in cameras: + cameras[LEGACY_REFERENCE_CAMERA] = CameraCalibration( + name=LEGACY_REFERENCE_CAMERA + ) + + reference_frame = ( + f"{CAMERA_FRAME_PREFIX}{LEGACY_REFERENCE_CAMERA}" + if LEGACY_REFERENCE_CAMERA in cameras + else f"{CAMERA_FRAME_PREFIX}{next(iter(cameras))}" + ) + return Calibration( + reference_frame=reference_frame, + cameras=cameras, + arm_bases=arm_bases, + legacy=True, + ) + + +def read_calibration(attrs: Mapping[str, Any]) -> Calibration | None: + """Return the calibration of one episode from its Zarr attributes. + + This is the one reader every consumer goes through. It prefers the + ``calibration`` block and falls back to the legacy ``intrinsics`` and + ``extrinsics`` attributes, so episodes written before the block existed + stay readable without a rewrite. + + Args: + attrs: The episode's Zarr attributes. + + Returns: + The episode calibration, or ``None`` if the episode states none. + + Raises: + CalibrationError: If the stored calibration is malformed. + """ + block = attrs.get("calibration") + if block: + return parse_calibration(block) + return lift_legacy_calibration( + attrs.get("intrinsics"), attrs.get("extrinsics") + ) + + +__all__ = [ + "CAMERA_FRAME_PREFIX", + "LEGACY_REFERENCE_CAMERA", + "STATIC_REFERENCE_FRAMES", + "Calibration", + "CalibrationError", + "CameraCalibration", + "lift_legacy_calibration", + "parse_calibration", + "read_calibration", +] diff --git a/egomimic/rldb/zarr/test_calibration.py b/egomimic/rldb/zarr/test_calibration.py new file mode 100644 index 000000000..bf5b3bf0a --- /dev/null +++ b/egomimic/rldb/zarr/test_calibration.py @@ -0,0 +1,232 @@ +"""Test the ``calibration`` attribute block and the legacy read shim.""" + +import numpy as np +import pytest + +from egomimic.rldb.embodiment.eva import Eva +from egomimic.rldb.zarr.calibration import ( + Calibration, + CalibrationError, + CameraCalibration, + lift_legacy_calibration, + parse_calibration, + read_calibration, +) +from egomimic.rldb.zarr.zarr_dataset_multi import ZarrEpisode +from egomimic.rldb.zarr.zarr_writer import ZarrWriter + +K_FRONT = np.array( + [ + [200.0, 0.0, 160.0, 0.0], + [0.0, 200.0, 120.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + ] +) + + +def _calibration_block() -> dict: + return { + "reference_frame": "robot_base", + "cameras": { + "front_1": { + "K": K_FRONT.tolist(), + "resolution": [320, 240], + "rectified": True, + "ref_T_cam": np.eye(4).tolist(), + }, + "left_wrist": {"K": K_FRONT.tolist(), "rectified": False}, + }, + "arm_bases": {"left": np.eye(4).tolist()}, + } + + +def test_parse_round_trips_through_jsonable() -> None: + calibration = parse_calibration(_calibration_block()) + assert calibration.reference_frame == "robot_base" + assert calibration.reference_camera is None + assert calibration.cameras["front_1"].resolution == (320, 240) + assert calibration.cameras["left_wrist"].rectified is False + reparsed = parse_calibration(calibration.to_jsonable()) + assert reparsed.reference_frame == calibration.reference_frame + assert set(reparsed.cameras) == set(calibration.cameras) + for name, camera in calibration.cameras.items(): + other = reparsed.cameras[name] + assert (other.resolution, other.rectified) == ( + camera.resolution, + camera.rectified, + ) + np.testing.assert_array_equal(other.K, camera.K) + np.testing.assert_array_equal( + reparsed.arm_bases["left"], calibration.arm_bases["left"] + ) + + +def test_default_camera_prefers_the_front_camera() -> None: + calibration = parse_calibration( + { + "reference_frame": "robot_base", + "cameras": { + "left_wrist": {"K": K_FRONT.tolist()}, + "front_1": {"K": (2 * K_FRONT).tolist()}, + }, + } + ) + assert calibration.default_camera() == "front_1" + np.testing.assert_allclose(calibration.K(), 2 * K_FRONT) + + +def test_reference_camera_pose_is_the_identity() -> None: + calibration = parse_calibration( + { + "reference_frame": "camera:front_1", + "cameras": {"front_1": {"K": K_FRONT.tolist()}}, + } + ) + np.testing.assert_allclose(calibration.ref_T_cam("front_1"), np.eye(4)) + assert calibration.ref_T_cam("left_wrist") is None + + +@pytest.mark.parametrize( + "block, message", + [ + ({"cameras": {"front_1": {}}}, "reference_frame"), + ({"reference_frame": "robot_base", "cameras": {}}, "non-empty mapping"), + ( + {"reference_frame": "camera:missing", "cameras": {"front_1": {}}}, + "does not declare", + ), + ( + {"reference_frame": "elbow", "cameras": {"front_1": {}}}, + "unknown reference_frame", + ), + ( + {"reference_frame": "robot_base", "cameras": {"front_1": {"fx": 1}}}, + "unknown field", + ), + ( + { + "reference_frame": "robot_base", + "cameras": {"front_1": {"K": [[1.0, 2.0], [3.0, 4.0]]}}, + }, + "3x4 camera matrix", + ), + ( + { + "reference_frame": "robot_base", + "cameras": {"front_1": {}}, + "arm_bases": {"left": np.eye(3).tolist()}, + }, + r"expected shape \(4, 4\)", + ), + ], +) +def test_parse_rejects_malformed_blocks(block, message) -> None: + with pytest.raises(CalibrationError, match=message): + parse_calibration(block) + + +def test_a_bare_3x3_camera_matrix_is_padded() -> None: + calibration = parse_calibration( + { + "reference_frame": "robot_base", + "cameras": {"front_1": {"K": K_FRONT[:, :3].tolist()}}, + } + ) + np.testing.assert_allclose(calibration.K("front_1"), K_FRONT) + + +def test_legacy_attributes_lift_without_changing_base_T_cam() -> None: + calibration = lift_legacy_calibration( + intrinsics={"front_1": K_FRONT}, extrinsics=Eva.EXTRINSICS + ) + assert calibration.legacy + assert calibration.reference_frame == "camera:front_1" + np.testing.assert_allclose(calibration.K(), K_FRONT) + for side, base_T_cam in Eva.EXTRINSICS.items(): + # `arm_bases` holds ref_T_armbase, the inverse of the stored matrix. + np.testing.assert_allclose( + calibration.arm_bases[side], np.linalg.inv(base_T_cam), atol=1e-12 + ) + np.testing.assert_allclose( + calibration.base_T_cam(side), base_T_cam, atol=1e-12 + ) + + +def test_legacy_extrinsics_without_intrinsics_still_compose() -> None: + calibration = lift_legacy_calibration(extrinsics=Eva.EXTRINSICS) + assert calibration.K() is None + np.testing.assert_allclose( + calibration.base_T_cam("left"), Eva.EXTRINSICS["left"], atol=1e-12 + ) + + +def test_read_calibration_prefers_the_block_over_the_legacy_pair() -> None: + attrs = { + "calibration": _calibration_block(), + "intrinsics": {"front_1": (7 * K_FRONT).tolist()}, + } + calibration = read_calibration(attrs) + assert calibration is not None and not calibration.legacy + np.testing.assert_allclose(calibration.K(), K_FRONT) + assert read_calibration({}) is None + + +def test_uncalibrated_lists_cameras_without_a_matrix() -> None: + calibration = Calibration( + reference_frame="robot_base", + cameras={ + "front_1": CameraCalibration(name="front_1", K=K_FRONT), + "left_wrist": CameraCalibration(name="left_wrist"), + }, + ) + assert calibration.uncalibrated(["front_1", "left_wrist", "right_wrist"]) == [ + "left_wrist", + "right_wrist", + ] + + +def _write_episode(episode_path, **kwargs) -> None: + ZarrWriter.create_and_write( + episode_path=episode_path, + numeric_data={"left.obs_gripper": np.zeros((4, 1))}, + embodiment="eva_bimanual", + chunk_timesteps=4, + **kwargs, + ) + + +def test_the_writer_stores_the_block_and_the_legacy_pair(tmp_path) -> None: + episode_path = tmp_path / "calibrated.zarr" + _write_episode(episode_path, calibration=_calibration_block()) + + episode = ZarrEpisode(episode_path) + assert episode.metadata["calibration"]["reference_frame"] == "robot_base" + # Readers that predate the block still find the pair they expect. + np.testing.assert_allclose( + np.asarray(episode.metadata["intrinsics"]["front_1"]), K_FRONT + ) + np.testing.assert_allclose( + np.asarray(episode.metadata["extrinsics"]["left"]), np.eye(4) + ) + np.testing.assert_allclose(episode.calibration.K(), K_FRONT) + + +def test_a_legacy_episode_reads_through_the_shim(tmp_path) -> None: + episode_path = tmp_path / "legacy.zarr" + _write_episode( + episode_path, intrinsics={"front_1": K_FRONT}, extrinsics=Eva.EXTRINSICS + ) + + episode = ZarrEpisode(episode_path) + assert "calibration" not in episode.metadata + calibration = episode.calibration + assert calibration is not None and calibration.legacy + np.testing.assert_allclose(calibration.K(), K_FRONT) + np.testing.assert_allclose( + calibration.base_T_cam("right"), Eva.EXTRINSICS["right"], atol=1e-12 + ) + + +def test_the_writer_rejects_an_episode_with_no_camera_matrix(tmp_path) -> None: + with pytest.raises(ValueError, match="Camera intrinsics"): + _write_episode(tmp_path / "blank.zarr") diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index a4d5d0a1b..e0b8df907 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -43,6 +43,7 @@ # from action_chunk_transforms import Transform from egomimic.rldb.filters import DatasetFilter +from egomimic.rldb.zarr.calibration import Calibration, read_calibration from egomimic.utils.aws.aws_data_utils import load_env from egomimic.utils.aws.aws_sql import ( create_default_engine, @@ -1559,6 +1560,11 @@ def init_episode(self): self._image_keys = self._detect_image_keys() self._json_keys = self._detect_json_keys() + @property + def calibration(self) -> Calibration | None: + """Pass-through to ZarrEpisode.calibration (read from zarr metadata).""" + return self.episode_reader.calibration + @property def intrinsics(self) -> np.ndarray | dict[str, np.ndarray] | None: """Pass-through to ZarrEpisode.intrinsics (read from zarr metadata).""" @@ -1736,27 +1742,15 @@ def _next(reason: str, key: str = "") -> int: data["embodiment"] = get_embodiment_id(self.embodiment) # Per-episode camera intrinsics travel with the batch so a single # data-driven embodiment can project without a hardcoded class const. - # Single-camera -> (3,4) K matrix; no/multi-camera -> NaN sentinel, - # which _intrinsics_from_batch treats as "fall back to cls.INTRINSICS". - K = self.intrinsics - if isinstance(K, dict): - # Multi-camera rig: project against the front camera (viz/projection - # target the front image). Prefer a "front"-keyed entry, else the - # first available camera. - K = next( - (v for k, v in K.items() if "front" in str(k).lower()), - next(iter(K.values()), None) if K else None, - ) - if K is not None: - K = np.asarray(K, dtype=np.float32) - if K.shape == (3, 3): - # Normalize a 3x3 K to the canonical 3x4 (zeros last column); - # some contributors store 3x3 (e.g. microagi). - K = np.concatenate([K, np.zeros((3, 1), dtype=np.float32)], axis=1) - if K.shape != (3, 4): # unexpected -> sentinel (viz falls back to const) - K = np.full((3, 4), np.nan, dtype=np.float32) - else: + # The calibration shim picks the front camera and normalizes a bare + # 3x3 K; an episode without one gets the NaN sentinel, which + # _intrinsics_from_batch treats as "fall back to cls.INTRINSICS". + calibration = self.calibration + K = None if calibration is None else calibration.K() + if K is None: K = np.full((3, 4), np.nan, dtype=np.float32) + else: + K = np.asarray(K, dtype=np.float32) data["intrinsics"] = torch.from_numpy(np.ascontiguousarray(K)) ep_name = Path(self.episode_path).name data["episode_hash"] = ep_name[:-5] if ep_name.endswith(".zarr") else ep_name @@ -1829,6 +1823,7 @@ class ZarrEpisode: "_store", "metadata", "keys", + "_calibration", ) def __init__(self, path: str | Path): @@ -1841,6 +1836,17 @@ def __init__(self, path: str | Path): self._store = zarr.open_group(str(self._path), mode="r") self.metadata = dict(self._store.attrs) self.keys = self.metadata["features"] + self._calibration = read_calibration(self.metadata) + + @property + def calibration(self) -> Calibration | None: + """Per-episode calibration, read through the one shim. + + Returns the ``calibration`` attribute block, or the same information + lifted from the legacy ``intrinsics``/``extrinsics`` pair, or ``None`` + when the episode states neither. + """ + return self._calibration @property def intrinsics(self) -> np.ndarray | dict[str, np.ndarray] | None: @@ -1848,16 +1854,12 @@ def intrinsics(self) -> np.ndarray | dict[str, np.ndarray] | None: Camera intrinsics persisted in zarr metadata, deserialized to ndarray(s). Returns: - - np.ndarray for a single-camera episode, - - dict[str, np.ndarray] for multi-camera episodes, - - None if no intrinsics were written. + - dict[str, np.ndarray] keyed by camera name, + - None if the episode carries no calibrated camera. """ - raw = self.metadata.get("intrinsics") - if raw is None: + if self._calibration is None: return None - if isinstance(raw, dict): - return {k: np.asarray(v) for k, v in raw.items()} - return np.asarray(raw) + return self._calibration.intrinsics() or None def read( self, keys_with_ranges: dict[str, tuple[int, int | None]] diff --git a/egomimic/rldb/zarr/zarr_writer.py b/egomimic/rldb/zarr/zarr_writer.py index 39b375f60..0f95b2808 100644 --- a/egomimic/rldb/zarr/zarr_writer.py +++ b/egomimic/rldb/zarr/zarr_writer.py @@ -14,6 +14,11 @@ import zarr from zarr.core.dtype import VariableLengthBytes +from egomimic.rldb.zarr.calibration import ( + Calibration, + parse_calibration, +) + def _intrinsics_to_jsonable( intrinsics: np.ndarray | list | dict, @@ -309,6 +314,7 @@ def __init__( chunk_timesteps: int = 100, intrinsics: dict | None = None, extrinsics: dict | None = None, + calibration: Calibration | dict | None = None, verbose: bool = False, ): """Configure a writer for one Zarr v3 episode. @@ -326,6 +332,9 @@ def __init__( extrinsics: ``None`` or a mapping from keys to 4×4 ``ref_T_cam`` matrices. Robot episodes use arm names as keys. See ``docs/CONVENTIONS.md``. + calibration: ``None`` or the per-episode calibration to store under + the ``calibration`` attribute. A ``Calibration`` or the mapping + form of one. verbose: If true, print progress information during writes. """ self.episode_path = Path(episode_path) @@ -339,6 +348,9 @@ def __init__( self.chunk_timesteps = chunk_timesteps self.intrinsics = intrinsics self.extrinsics = extrinsics + self.calibration = ( + None if calibration is None else parse_calibration(calibration) + ) self.verbose = verbose # Track image shapes for metadata self._features: dict[str, dict[str, Any]] = {} @@ -712,6 +724,9 @@ def _build_metadata( "features": self._features, } + if self.calibration is not None: + metadata["calibration"] = self.calibration.to_jsonable() + if self.intrinsics is not None: metadata["intrinsics"] = _intrinsics_to_jsonable(self.intrinsics) @@ -719,15 +734,15 @@ def _build_metadata( metadata["extrinsics"] = _intrinsics_to_jsonable(self.extrinsics) # Apply overrides — but NEVER let them clobber the validated camera - # metadata. intrinsics/extrinsics are validated in create_and_write; a - # converter's metadata_override that happens to carry a stale or empty - # "intrinsics"/"extrinsics" key must not silently overwrite the validated + # metadata. calibration/intrinsics/extrinsics are validated in + # create_and_write; a converter's metadata_override that happens to + # carry a stale or empty key must not silently overwrite the validated # values (this was the Mecka clobber bug → empty intrinsics in zarr.json). if metadata_override: override = { k: v for k, v in metadata_override.items() - if k not in ("intrinsics", "extrinsics") + if k not in ("calibration", "intrinsics", "extrinsics") } metadata.update(override) @@ -746,8 +761,9 @@ def create_and_write( annotations: list[tuple[str, int, int]] | None = None, chunk_timesteps: int = 100, *, - intrinsics: dict, + intrinsics: dict | None = None, extrinsics: dict | None = None, + calibration: Calibration | dict | None = None, metadata_override: dict[str, Any] | None = None, ) -> Path: """Validate episode metadata and write one Zarr v3 episode. @@ -767,12 +783,17 @@ def create_and_write( annotations: ``(text, start_idx, end_idx)`` annotation tuples. chunk_timesteps: Number of frames in each numeric-array chunk. intrinsics: A non-empty mapping from camera keys to 3×4 camera - matrices. + matrices. Optional only when ``calibration`` supplies them. extrinsics: ``None`` or a non-empty mapping from keys to 4×4 ``ref_T_cam`` matrices. Robot episodes use arm names as keys. See ``docs/CONVENTIONS.md``. + calibration: ``None`` or the per-episode calibration described in + ``egomimic/rldb/zarr/calibration.py``. When it is present the + writer derives ``intrinsics`` and ``extrinsics`` from it for any + of the two that the caller omits, so readers that predate the + block keep working. metadata_override: Additional episode metadata. This mapping cannot - replace ``intrinsics`` or ``extrinsics``. + replace ``calibration``, ``intrinsics``, or ``extrinsics``. Returns: The path of the created ``.zarr`` directory. @@ -783,6 +804,7 @@ def create_and_write( ValueError: If ``intrinsics`` is empty or contains a non-3×4 value. ValueError: If ``extrinsics`` is not ``None`` or a non-empty mapping of 4×4 matrices. + CalibrationError: If ``calibration`` is malformed. """ # Validate the embodiment up front (it is guaranteed to be consumed by # the reader via get_embodiment_id) so a bad/empty/typo'd value fails @@ -795,11 +817,22 @@ def create_and_write( f"{[m.name.lower() for m in EMBODIMENT]}, got {embodiment!r}. " "See CONTRIBUTING_DATA.md §9." ) + # `calibration` is the current form; `intrinsics`/`extrinsics` are the + # legacy pair. Writing both keeps every reader working, so derive + # whichever the caller left out. + if calibration is not None: + calibration = parse_calibration(calibration) + if intrinsics is None: + intrinsics = calibration.intrinsics() + if extrinsics is None: + extrinsics = calibration.extrinsics() or None + if not isinstance(intrinsics, dict) or not intrinsics: raise ValueError( "Camera intrinsics must be a non-empty {camera_key: 3x4 K matrix} " 'dict for a contributed zarr episode, e.g. {"front_1": K}. ' - "See CONTRIBUTING_DATA.md §6.4." + "Pass `intrinsics=` or a `calibration=` block that carries a K " + "for at least one camera. See CONTRIBUTING_DATA.md §6.4." ) for cam_key, K in intrinsics.items(): try: @@ -843,6 +876,7 @@ def create_and_write( chunk_timesteps=chunk_timesteps, intrinsics=intrinsics, extrinsics=extrinsics, + calibration=calibration, ) writer.write( diff --git a/egomimic/scripts/data_visualization/inspector_lib/dataset_view.py b/egomimic/scripts/data_visualization/inspector_lib/dataset_view.py index 81551de95..b2f1987de 100644 --- a/egomimic/scripts/data_visualization/inspector_lib/dataset_view.py +++ b/egomimic/scripts/data_visualization/inspector_lib/dataset_view.py @@ -20,8 +20,9 @@ can't handle yet. Projection convention (human egocentric `human_bimanual`): the head IS the -camera. Per-episode intrinsics are read from `zarr.json` -(`grp.attrs["intrinsics"]["front_1"]`, a 3x4 K); poses in `*.obs_ee_pose` +camera. Per-episode intrinsics are read from `zarr.json` through +`read_calibration` (the `calibration` block, or the legacy `intrinsics` +attribute of an older episode); poses in `*.obs_ee_pose` (xyz + quat wxyz, SLAM-world) are transformed world → head frame using `obs_head_pose` (same xyz+wxyz layout) and projected with that K. There is no per-arm extrinsic. (This matches the validated `Human.viz` path; the old @@ -43,6 +44,7 @@ # `Embodiment.viz` method projects the poses and draws the overlay. # `projectaria_tools` is optional in the standalone visualization environment. # If an embodiment import fails, the browser shows an unavailable badge. +from egomimic.rldb.zarr.calibration import read_calibration from egomimic.utils.pose_utils import ee_pose_to_cam_frame from .images import ( @@ -283,22 +285,28 @@ def _badge(img_rgb, text: str): return img_rgb -def _intrinsics_from_zarr(grp): - """Per-episode 3x4 camera matrix K from zarr.json metadata - (`attrs["intrinsics"]["front_1"]`). A 3x3 K is normalized to 3x4 by - appending a zero column. Returns None if absent/unparseable.""" +def _calibration_from_zarr(grp): + """Read the episode calibration through the one shim. + + Returns: + A ``Calibration``, or ``None`` if the episode states none or the stored + calibration is malformed. The inspector renders episodes that a + validator would reject, so a bad block degrades the overlay rather than + raising. + """ try: - intr = dict(grp.attrs).get("intrinsics") - if not intr: - return None - K = np.asarray(intr["front_1"], dtype=float) + return read_calibration(dict(grp.attrs)) except Exception: return None - if K.shape == (3, 4): - return K - if K.shape == (3, 3): - return np.hstack([K, np.zeros((3, 1))]) - return None + + +def _intrinsics_from_zarr(grp): + """Per-episode 3x4 camera matrix K for the front camera. + + Returns None if the episode declares no calibrated camera. + """ + calibration = _calibration_from_zarr(grp) + return None if calibration is None else calibration.K() def _extrinsics_from_zarr(grp): @@ -306,26 +314,16 @@ def _extrinsics_from_zarr(grp): Returns: A dictionary with valid ``"left"`` and ``"right"`` 4×4 matrices. - The function returns ``None`` if no valid matrix is present. It ignores - missing, non-numeric, and incorrectly shaped values. + The function returns ``None`` if no valid matrix is present. """ - try: - extr = dict(grp.attrs).get("extrinsics") - if not extr: - return None - except Exception: + calibration = _calibration_from_zarr(grp) + if calibration is None: return None out = {} for arm in ("left", "right"): - m = extr.get(arm) - if m is None: - continue - try: - M = np.asarray(m, dtype=float) - except Exception: - continue - if M.shape == (4, 4): - out[arm] = M + base_T_cam = calibration.base_T_cam(arm) + if base_T_cam is not None: + out[arm] = base_T_cam return out or None From 0b62dd49e5134129ad9922d687d648553a370213 Mon Sep 17 00:00:00 2001 From: jaynye Date: Wed, 2 Sep 2026 20:47:59 -0400 Subject: [PATCH 2/9] feat(calibration): check that every image stream has a camera matrix `intrinsics` only had to be non-empty, so `eva_to_zarr.py` calibrates one camera while writing three image streams and passes. Coverage closes that: the declared `images.` keys are checked against the cameras that carry a K. The rule is opt-in. Every EVA episode in the corpus fails it today, so `strict=False` warns and `strict=True` raises. Flipping the default before the corpus is fixed would turn a good check into an outage. `uncalibrated_cameras` lives beside the calibration reader so the validator can reuse the same rule rather than restate it. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_012rCh1QmnyqgPZ5TJpsfMQb --- egomimic/rldb/zarr/calibration.py | 54 +++++++++++++++++++++++-- egomimic/rldb/zarr/test_calibration.py | 55 +++++++++++++++++++++++++- egomimic/rldb/zarr/zarr_writer.py | 50 +++++++++++++++++++++++ 3 files changed, 153 insertions(+), 6 deletions(-) diff --git a/egomimic/rldb/zarr/calibration.py b/egomimic/rldb/zarr/calibration.py index abe9407e8..71aee8f97 100644 --- a/egomimic/rldb/zarr/calibration.py +++ b/egomimic/rldb/zarr/calibration.py @@ -44,6 +44,10 @@ #: Prefix that makes a camera the reference frame, as in `camera:front_1`. CAMERA_FRAME_PREFIX = "camera:" +#: Prefix of the array key that stores one camera stream, as in +#: `images.front_1`. The text after it is the camera name. +IMAGE_KEY_PREFIX = "images." + _CAMERA_FIELDS = frozenset({"K", "resolution", "rectified", "ref_T_cam"}) _CALIBRATION_FIELDS = frozenset({"reference_frame", "cameras", "arm_bases"}) @@ -229,10 +233,6 @@ def extrinsics(self, camera: str | None = None) -> dict[str, np.ndarray]: out[side] = base_T_cam return out - def uncalibrated(self, cameras) -> list[str]: - """Return the given camera names that carry no ``K``, in order.""" - return [name for name in cameras if self.K(name) is None] - def to_jsonable(self) -> dict[str, Any]: """Return this calibration as JSON-serializable episode metadata.""" out: dict[str, Any] = { @@ -249,6 +249,49 @@ def to_jsonable(self) -> dict[str, Any]: return out +def camera_name(image_key: str) -> str | None: + """Return the camera that one stored image array belongs to. + + Args: + image_key: An array key such as ``"images.front_1"``. + + Returns: + The camera name, or ``None`` if the key names no image stream. + """ + marker = IMAGE_KEY_PREFIX + index = image_key.rfind(marker) + if index == -1: + return None + return image_key[index + len(marker) :] or None + + +def uncalibrated_cameras(image_keys, calibration: Calibration | None) -> list[str]: + """Return the image streams that no camera matrix covers. + + Coverage is the rule that makes per-episode calibration real: an episode + that stores three image streams and one ``K`` calibrates one camera in + three, and nothing downstream can tell. + + Args: + image_keys: The episode's image array keys. + calibration: The episode calibration, or ``None``. + + Returns: + The camera names, in the order the keys give them and without + duplicates, that carry no ``K``. + """ + missing: list[str] = [] + seen: set[str] = set() + for key in image_keys: + name = camera_name(key) + if name is None or name in seen: + continue + seen.add(name) + if calibration is None or calibration.K(name) is None: + missing.append(name) + return missing + + def _parse_reference_frame(value: Any, cameras, where: str) -> str: if not isinstance(value, str) or not value: raise CalibrationError(f"{where}: `reference_frame` must be a non-empty string") @@ -466,12 +509,15 @@ def read_calibration(attrs: Mapping[str, Any]) -> Calibration | None: __all__ = [ "CAMERA_FRAME_PREFIX", + "IMAGE_KEY_PREFIX", "LEGACY_REFERENCE_CAMERA", "STATIC_REFERENCE_FRAMES", "Calibration", "CalibrationError", "CameraCalibration", + "camera_name", "lift_legacy_calibration", "parse_calibration", "read_calibration", + "uncalibrated_cameras", ] diff --git a/egomimic/rldb/zarr/test_calibration.py b/egomimic/rldb/zarr/test_calibration.py index bf5b3bf0a..0ae8561c6 100644 --- a/egomimic/rldb/zarr/test_calibration.py +++ b/egomimic/rldb/zarr/test_calibration.py @@ -8,9 +8,11 @@ Calibration, CalibrationError, CameraCalibration, + camera_name, lift_legacy_calibration, parse_calibration, read_calibration, + uncalibrated_cameras, ) from egomimic.rldb.zarr.zarr_dataset_multi import ZarrEpisode from egomimic.rldb.zarr.zarr_writer import ZarrWriter @@ -171,7 +173,14 @@ def test_read_calibration_prefers_the_block_over_the_legacy_pair() -> None: assert read_calibration({}) is None -def test_uncalibrated_lists_cameras_without_a_matrix() -> None: +def test_camera_name_reads_the_stream_out_of_an_array_key() -> None: + assert camera_name("images.front_1") == "front_1" + assert camera_name("observations.images.left_wrist") == "left_wrist" + assert camera_name("left.obs_ee_pose") is None + assert camera_name("images.") is None + + +def test_uncalibrated_cameras_lists_streams_without_a_matrix() -> None: calibration = Calibration( reference_frame="robot_base", cameras={ @@ -179,10 +188,52 @@ def test_uncalibrated_lists_cameras_without_a_matrix() -> None: "left_wrist": CameraCalibration(name="left_wrist"), }, ) - assert calibration.uncalibrated(["front_1", "left_wrist", "right_wrist"]) == [ + keys = ["images.front_1", "images.left_wrist", "images.right_wrist"] + assert uncalibrated_cameras(keys, calibration) == ["left_wrist", "right_wrist"] + assert uncalibrated_cameras(keys, None) == [ + "front_1", "left_wrist", "right_wrist", ] + assert uncalibrated_cameras(["left.obs_gripper"], None) == [] + + +def test_coverage_is_a_warning_by_default_and_an_error_under_strict( + tmp_path, caplog +) -> None: + images = { + "images.front_1": np.zeros((4, 8, 8, 3), dtype=np.uint8), + "images.left_wrist": np.zeros((4, 8, 8, 3), dtype=np.uint8), + } + with caplog.at_level("WARNING"): + _write_episode( + tmp_path / "partial.zarr", + image_data=images, + intrinsics={"front_1": K_FRONT}, + ) + assert "left_wrist" in caplog.text + + with pytest.raises(ValueError, match=r"no camera matrix.*left_wrist"): + _write_episode( + tmp_path / "strict.zarr", + image_data=images, + intrinsics={"front_1": K_FRONT}, + strict=True, + ) + + +def test_full_coverage_passes_under_strict(tmp_path, caplog) -> None: + with caplog.at_level("WARNING"): + _write_episode( + tmp_path / "covered.zarr", + image_data={"images.front_1": np.zeros((4, 8, 8, 3), dtype=np.uint8)}, + calibration={ + "reference_frame": "camera:front_1", + "cameras": {"front_1": {"K": K_FRONT.tolist()}}, + }, + strict=True, + ) + assert "no camera matrix" not in caplog.text def _write_episode(episode_path, **kwargs) -> None: diff --git a/egomimic/rldb/zarr/zarr_writer.py b/egomimic/rldb/zarr/zarr_writer.py index 0f95b2808..8985429d4 100644 --- a/egomimic/rldb/zarr/zarr_writer.py +++ b/egomimic/rldb/zarr/zarr_writer.py @@ -6,6 +6,7 @@ """ import json +import logging from pathlib import Path from typing import Any, Literal @@ -16,9 +17,13 @@ from egomimic.rldb.zarr.calibration import ( Calibration, + lift_legacy_calibration, parse_calibration, + uncalibrated_cameras, ) +logger = logging.getLogger(__name__) + def _intrinsics_to_jsonable( intrinsics: np.ndarray | list | dict, @@ -287,6 +292,8 @@ def __exit__(self, exc_type, exc_val, exc_tb) -> bool: padding[:] = last_jpeg self._store[key][self._total_frames : padded] = padding + self._writer._check_camera_coverage(self._image_info) + # Write language annotations if self._writer.annotations is not None: self._writer._write_annotations(self._store, self._writer.annotations) @@ -315,6 +322,7 @@ def __init__( intrinsics: dict | None = None, extrinsics: dict | None = None, calibration: Calibration | dict | None = None, + strict: bool = False, verbose: bool = False, ): """Configure a writer for one Zarr v3 episode. @@ -335,6 +343,8 @@ def __init__( calibration: ``None`` or the per-episode calibration to store under the ``calibration`` attribute. A ``Calibration`` or the mapping form of one. + strict: If true, an image stream that no camera matrix covers is an + error instead of a warning. verbose: If true, print progress information during writes. """ self.episode_path = Path(episode_path) @@ -351,6 +361,7 @@ def __init__( self.calibration = ( None if calibration is None else parse_calibration(calibration) ) + self.strict = strict self.verbose = verbose # Track image shapes for metadata self._features: dict[str, dict[str, Any]] = {} @@ -431,6 +442,10 @@ def write( store, key, enc_arr, img_shape, padded_frames ) + self._check_camera_coverage( + list(image_data) + list(pre_encoded_image_data) + ) + # Write language annotations if provided if self.annotations is not None: self._write_annotations(store, self.annotations) @@ -703,6 +718,34 @@ def _write_annotations( "format": "annotation_v1", } + def _check_camera_coverage(self, image_keys) -> None: + """Check that every stored image stream has a camera matrix. + + `intrinsics` only has to be non-empty, so an episode that writes three + image streams and calibrates one passes. Coverage closes that: it is + checked against the image keys the episode actually stores. + + Args: + image_keys: The episode's image array keys. + + Raises: + ValueError: If ``strict`` is set and a stream carries no ``K``. + """ + calibration = self.calibration or lift_legacy_calibration( + self.intrinsics, self.extrinsics + ) + missing = uncalibrated_cameras(image_keys, calibration) + if not missing: + return + message = ( + f"{self.episode_path.name}: no camera matrix for image stream(s) " + f"{missing}. Every `images.` array needs a K in " + "`calibration.cameras`. See CONTRIBUTING_DATA.md §6.4." + ) + if self.strict: + raise ValueError(message) + logger.warning("%s Pass strict=True to make this an error.", message) + def _build_metadata( self, metadata_override: dict[str, Any] | None = None ) -> dict[str, Any]: @@ -764,6 +807,7 @@ def create_and_write( intrinsics: dict | None = None, extrinsics: dict | None = None, calibration: Calibration | dict | None = None, + strict: bool = False, metadata_override: dict[str, Any] | None = None, ) -> Path: """Validate episode metadata and write one Zarr v3 episode. @@ -792,6 +836,9 @@ def create_and_write( writer derives ``intrinsics`` and ``extrinsics`` from it for any of the two that the caller omits, so readers that predate the block keep working. + strict: If true, an image stream that no camera matrix covers is an + error instead of a warning. It stays opt-in because every EVA + episode in the corpus stores three streams and calibrates one. metadata_override: Additional episode metadata. This mapping cannot replace ``calibration``, ``intrinsics``, or ``extrinsics``. @@ -804,6 +851,8 @@ def create_and_write( ValueError: If ``intrinsics`` is empty or contains a non-3×4 value. ValueError: If ``extrinsics`` is not ``None`` or a non-empty mapping of 4×4 matrices. + ValueError: If ``strict`` is set and an image stream carries no + camera matrix. CalibrationError: If ``calibration`` is malformed. """ # Validate the embodiment up front (it is guaranteed to be consumed by @@ -877,6 +926,7 @@ def create_and_write( intrinsics=intrinsics, extrinsics=extrinsics, calibration=calibration, + strict=strict, ) writer.write( From 5dc10e9e06737a0100e40c0acb674fc6f6732ffe Mon Sep 17 00:00:00 2001 From: jaynye Date: Wed, 2 Sep 2026 20:51:56 -0400 Subject: [PATCH 3/9] feat(calibration): train against the rig the episode declares The transform pipeline expressed actions in the camera frame using `Eva.EXTRINSICS`, a class constant, and ignored what the episode stored. Two vendors on one platform have two rigs, and a class constant cannot tell them apart, so the factorization is unsound until the training path reads the episode. `ZarrDataset` now puts each arm's `base_T_cam` from the episode calibration into every sample, and a transform's `extra_batch_key` becomes a fallback that fills only keys the sample does not already carry. `Eva.EXTRINSICS` still covers an episode that declares no extrinsics. This changes numbers only for an episode whose stored rig differs from the constant. A test asserts the two agree: an episode holding `Eva.EXTRINSICS` and an episode holding none produce the same actions, which is every EVA episode in the corpus today. Two further tests show a second rig moves the actions and lands exactly where that rig predicts. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_012rCh1QmnyqgPZ5TJpsfMQb --- egomimic/rldb/embodiment/eva.py | 16 +++ egomimic/rldb/zarr/action_chunk_transforms.py | 32 ++++- egomimic/rldb/zarr/test_calibration.py | 115 +++++++++++++++++- egomimic/rldb/zarr/zarr_dataset_multi.py | 33 +++++ 4 files changed, 191 insertions(+), 5 deletions(-) diff --git a/egomimic/rldb/embodiment/eva.py b/egomimic/rldb/embodiment/eva.py index d63df477f..75a779831 100644 --- a/egomimic/rldb/embodiment/eva.py +++ b/egomimic/rldb/embodiment/eva.py @@ -26,7 +26,19 @@ class Eva(Embodiment): + """EVA X5 bimanual platform with a parallel jaw on each arm. + + ``INTRINSICS`` and ``EXTRINSICS`` are fallbacks for an episode that + declares no calibration of its own. Calibration measures the rig that + recorded one episode, so the episode's own values win: ``ZarrDataset`` + puts them in every sample and the transform pipeline reads them from + there. Two vendors on one platform have two rigs, and a class constant + cannot tell them apart. + """ + INTRINSICS = ARIA_INTRINSICS + #: Fallback rig: `base_T_cam` per arm, the camera pose in that arm's base + #: frame. See `docs/CONVENTIONS.md`. EXTRINSICS = { "left": np.array( [ @@ -248,6 +260,8 @@ def _build_eva_bimanual_eef_frame_transform_list( ) -> list[Transform]: """EVA bimanual transform pipeline with actions expressed relative to the current EEF pose (wrist frame), analogous to keypoints relative to wrist pose.""" + # A fallback for an episode that declares no extrinsics. `ZarrDataset` + # puts the episode's own rig in the sample, and that value wins. extrinsics = Eva.EXTRINSICS left_base_T_cam_pose = _matrix_to_xyzwxyz(extrinsics["left"][None, :])[0] right_base_T_cam_pose = _matrix_to_xyzwxyz(extrinsics["right"][None, :])[0] @@ -411,6 +425,8 @@ def _build_eva_bimanual_transform_list( is_quat: bool = True, ) -> list[Transform]: """Canonical EVA bimanual transform pipeline used by tests and notebooks.""" + # A fallback for an episode that declares no extrinsics. `ZarrDataset` + # puts the episode's own rig in the sample, and that value wins. extrinsics = Eva.EXTRINSICS left_base_T_cam_pose = _matrix_to_xyzwxyz(extrinsics["left"][None, :])[0] right_base_T_cam_pose = _matrix_to_xyzwxyz(extrinsics["right"][None, :])[0] diff --git a/egomimic/rldb/zarr/action_chunk_transforms.py b/egomimic/rldb/zarr/action_chunk_transforms.py index 9daf5e057..7d5955f93 100644 --- a/egomimic/rldb/zarr/action_chunk_transforms.py +++ b/egomimic/rldb/zarr/action_chunk_transforms.py @@ -35,6 +35,30 @@ xyzw_to_wxyz, ) +# --------------------------------------------------------------------------- +# Per-episode calibration +# --------------------------------------------------------------------------- + + +#: Batch key holding one arm's ``base_T_cam`` pose, the camera pose in that +#: arm's base frame. ``ZarrDataset`` fills it from the episode calibration; a +#: transform's ``extra_batch_key`` supplies an embodiment default when the +#: episode declares none. +def base_T_cam_pose_key(side: str) -> str: + """Return the batch key that carries one arm's ``base_T_cam`` pose.""" + return f"{side}_base_T_cam_pose" + + +def _apply_fallbacks(batch, extra_batch_key) -> None: + """Fill batch keys that the sample does not already carry. + + Calibration is per-episode, so a value the dataset read from the episode + must outrank the embodiment constant a transform was built with. + """ + for key, value in (extra_batch_key or {}).items(): + batch.setdefault(key, value) + + # --------------------------------------------------------------------------- # Base Transform # --------------------------------------------------------------------------- @@ -159,7 +183,9 @@ def __init__( frame. This pose represents ``reference_T_target``. chunk_world: Batch key for the input poses. transformed_key_name: Batch key for the output poses. - extra_batch_key: Values to add to the batch before the transform. + extra_batch_key: Fallback values for keys the batch does not + already carry. A per-episode value in the batch wins, so a + class constant here is a default and not an override. mode: Input and output layout. Use ``"xyz"``, ``"xyzypr"``, or ``"xyzwxyz"``. inverse: If true, compute ``target_T_chunk`` from @@ -190,7 +216,7 @@ def transform(self, batch): ValueError: If ``mode`` is not a supported pose layout. """ # Flatten the leading chunk dimensions into one pose dimension. - batch.update(self.extra_batch_key or {}) + _apply_fallbacks(batch, self.extra_batch_key) target_world_pose = np.asarray(batch[self.target_world]) chunk_world_poses = np.asarray(batch[self.chunk_world]) chunk_world_poses_shape = None @@ -437,7 +463,7 @@ def transform(self, batch): returns batch with new key containing transformed chunk world in target frame: (T, 14) """ - batch.update(self.extra_batch_key or {}) + _apply_fallbacks(batch, self.extra_batch_key) left_target_world = batch[self.left_target_world] right_target_world = batch[self.right_target_world] chunk_world = batch[self.chunk_world] diff --git a/egomimic/rldb/zarr/test_calibration.py b/egomimic/rldb/zarr/test_calibration.py index 0ae8561c6..f10e6c8f7 100644 --- a/egomimic/rldb/zarr/test_calibration.py +++ b/egomimic/rldb/zarr/test_calibration.py @@ -3,7 +3,10 @@ import numpy as np import pytest -from egomimic.rldb.embodiment.eva import Eva +from egomimic.rldb.embodiment.eva import Eva, _build_eva_bimanual_transform_list +from egomimic.rldb.zarr.action_chunk_transforms import ( + ActionChunkCoordinateFrameTransform, +) from egomimic.rldb.zarr.calibration import ( Calibration, CalibrationError, @@ -14,7 +17,7 @@ read_calibration, uncalibrated_cameras, ) -from egomimic.rldb.zarr.zarr_dataset_multi import ZarrEpisode +from egomimic.rldb.zarr.zarr_dataset_multi import ZarrDataset, ZarrEpisode from egomimic.rldb.zarr.zarr_writer import ZarrWriter K_FRONT = np.array( @@ -281,3 +284,111 @@ def test_a_legacy_episode_reads_through_the_shim(tmp_path) -> None: def test_the_writer_rejects_an_episode_with_no_camera_matrix(tmp_path) -> None: with pytest.raises(ValueError, match="Camera intrinsics"): _write_episode(tmp_path / "blank.zarr") + + +# -------------------------------------------------------------------------- +# The training path reads the rig the episode declares +# -------------------------------------------------------------------------- + +_EVA_KEY_MAP = { + "left.obs_ee_pose": {"zarr_key": "left.obs_ee_pose"}, + "right.obs_ee_pose": {"zarr_key": "right.obs_ee_pose"}, + "left.obs_gripper": {"zarr_key": "left.obs_gripper"}, + "right.obs_gripper": {"zarr_key": "right.obs_gripper"}, + "left.cmd_ee_pose": {"zarr_key": "left.cmd_ee_pose", "horizon": 4}, + "right.cmd_ee_pose": {"zarr_key": "right.cmd_ee_pose", "horizon": 4}, + "left.cmd_gripper": {"zarr_key": "left.cmd_gripper", "horizon": 4}, + "right.cmd_gripper": {"zarr_key": "right.cmd_gripper", "horizon": 4}, +} + + +def _pose_sequence(length: int, *, x_offset: float = 0.0) -> np.ndarray: + poses = np.zeros((length, 7)) + poses[:, 0] = x_offset + np.arange(length) * 0.01 + poses[:, 3] = 1.0 + return poses + + +def _eva_numeric_data(length: int) -> dict: + return { + "left.obs_ee_pose": _pose_sequence(length), + "right.obs_ee_pose": _pose_sequence(length, x_offset=0.2), + "left.obs_gripper": np.full((length, 1), 0.25), + "right.obs_gripper": np.full((length, 1), 0.75), + "left.cmd_ee_pose": _pose_sequence(length, x_offset=0.1), + "right.cmd_ee_pose": _pose_sequence(length, x_offset=0.3), + "left.cmd_gripper": np.linspace(0.0, 1.0, length)[:, None], + "right.cmd_gripper": np.linspace(1.0, 0.0, length)[:, None], + } + + +def _other_rig() -> dict: + """Return a rig displaced from `Eva.EXTRINSICS` by a measurable amount.""" + rig = {} + for side, base_T_cam in Eva.EXTRINSICS.items(): + moved = base_T_cam.copy() + moved[:3, 3] += np.array([0.05, -0.11, 0.07]) + rig[side] = moved + return rig + + +def _eva_actions(episode_path, extrinsics) -> np.ndarray: + """Write one EVA episode and return its transformed action chunk.""" + ZarrWriter.create_and_write( + episode_path=episode_path, + numeric_data=_eva_numeric_data(5), + embodiment="eva_bimanual", + intrinsics={"front_1": K_FRONT}, + extrinsics=extrinsics, + chunk_timesteps=4, + ) + dataset = ZarrDataset( + Episode_path=episode_path, + key_map=dict(_EVA_KEY_MAP), + transform_list=_build_eva_bimanual_transform_list( + chunk_length=6, stride=1, is_quat=True + ), + ) + return dataset[1]["actions_cartesian"].numpy() + + +def test_reading_the_episode_is_a_no_op_on_the_current_corpus(tmp_path) -> None: + stored = _eva_actions(tmp_path / "stored.zarr", Eva.EXTRINSICS) + absent = _eva_actions(tmp_path / "absent.zarr", None) + # Every EVA episode in the corpus stores exactly the class constant, so + # reading the episode instead of the constant changes no number today. + np.testing.assert_allclose(stored, absent, atol=1e-9) + + +def test_the_pipeline_follows_the_rig_the_episode_declares(tmp_path) -> None: + on_eva_rig = _eva_actions(tmp_path / "eva_rig.zarr", Eva.EXTRINSICS) + on_other_rig = _eva_actions(tmp_path / "other_rig.zarr", _other_rig()) + assert not np.allclose(on_eva_rig, on_other_rig, atol=1e-6) + + +def test_a_second_rig_matches_what_that_rig_predicts(tmp_path, monkeypatch) -> None: + other = _other_rig() + declared = _eva_actions(tmp_path / "declared.zarr", other) + + # An episode without extrinsics falls back to the class constant, so + # pointing the constant at the second rig must reproduce the same numbers. + monkeypatch.setattr(Eva, "EXTRINSICS", other) + fallback = _eva_actions(tmp_path / "fallback.zarr", None) + np.testing.assert_allclose(declared, fallback, atol=1e-9) + + +def test_a_transform_fallback_never_overrides_a_batch_value() -> None: + identity_pose = np.array([1.0, 2.0, 3.0, 1.0, 0.0, 0.0, 0.0]) + transform = ActionChunkCoordinateFrameTransform( + target_world="target", + chunk_world="chunk", + transformed_key_name="out", + extra_batch_key={"target": np.array([9.0, 9.0, 9.0, 1.0, 0.0, 0.0, 0.0])}, + mode="xyzwxyz", + ) + batch = {"target": identity_pose, "chunk": np.tile(identity_pose, (2, 1))} + + out = transform.transform(batch) + + np.testing.assert_array_equal(out["target"], identity_pose) + np.testing.assert_allclose(out["out"][:, :3], 0.0, atol=1e-12) diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index e0b8df907..45f3e5604 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -43,12 +43,14 @@ # from action_chunk_transforms import Transform from egomimic.rldb.filters import DatasetFilter +from egomimic.rldb.zarr.action_chunk_transforms import base_T_cam_pose_key from egomimic.rldb.zarr.calibration import Calibration, read_calibration from egomimic.utils.aws.aws_data_utils import load_env from egomimic.utils.aws.aws_sql import ( create_default_engine, episode_table_to_df, ) +from egomimic.utils.pose_utils import _matrix_to_xyzwxyz if TYPE_CHECKING: # Annotation-only import — avoids a runtime circular import with @@ -1559,12 +1561,38 @@ def init_episode(self): self.keys_dict = {k: (0, None) for k in self.episode_reader._collect_keys()} self._image_keys = self._detect_image_keys() self._json_keys = self._detect_json_keys() + self._extrinsic_poses = self._build_extrinsic_poses() @property def calibration(self) -> Calibration | None: """Pass-through to ZarrEpisode.calibration (read from zarr metadata).""" return self.episode_reader.calibration + def _build_extrinsic_poses(self) -> dict[str, np.ndarray]: + """Return this episode's per-arm ``base_T_cam`` poses for the batch. + + The transform pipeline expresses actions in the camera frame, so it + needs the rig that recorded this episode. Two vendors on one platform + have two rigs, and an embodiment class constant cannot tell them apart. + An episode that declares no extrinsics contributes nothing here and the + transform falls back to its embodiment default. + + Returns: + A mapping from ``_base_T_cam_pose`` to a 7-value + ``[xyz, qw, qx, qy, qz]`` pose. + """ + calibration = self.calibration + if calibration is None: + return {} + poses = {} + for side in ("left", "right"): + base_T_cam = calibration.base_T_cam(side) + if base_T_cam is not None: + poses[base_T_cam_pose_key(side)] = _matrix_to_xyzwxyz( + base_T_cam[None, :] + )[0] + return poses + @property def intrinsics(self) -> np.ndarray | dict[str, np.ndarray] | None: """Pass-through to ZarrEpisode.intrinsics (read from zarr metadata).""" @@ -1731,6 +1759,11 @@ def _next(reason: str, key: str = "") -> int: if retry: continue + # Per-episode extrinsics travel with the sample so the pipeline + # transforms into the frame of the rig that recorded this episode. + for key, pose in self._extrinsic_poses.items(): + data[key] = pose.copy() + if self.transform: for transform in self.transform or []: data = transform.transform(data) From 06335502ab8da7b3e372ebbe27176cc7def975fc Mon Sep 17 00:00:00 2001 From: jaynye Date: Wed, 2 Sep 2026 20:53:04 -0400 Subject: [PATCH 4/9] feat(calibration): declare each camera's projection model and distortion Collecting calibration is cheap and time-critical; consuming it is expensive and deferrable. A vendor who ships 10k episodes and later turns out to need distortion coefficients cannot supply them, because the rig has moved. Each camera in the block now declares `model` and `distortion`, defaulting to PINHOLE with no coefficients. The coefficient count is checked against the model, so a KANNALA_BRANDT camera with five coefficients or a pinhole camera with any is an error at ingest. No projection site reads either field. Nothing under `algo/` or `models/` reads intrinsics at all, so this changes no trained weight. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_012rCh1QmnyqgPZ5TJpsfMQb --- docs/CONVENTIONS.md | 13 +++++ egomimic/rldb/zarr/calibration.py | 68 +++++++++++++++++++++++++- egomimic/rldb/zarr/test_calibration.py | 48 ++++++++++++++++++ 3 files changed, 127 insertions(+), 2 deletions(-) diff --git a/docs/CONVENTIONS.md b/docs/CONVENTIONS.md index b7b0a0cf0..bbeea9edb 100644 --- a/docs/CONVENTIONS.md +++ b/docs/CONVENTIONS.md @@ -36,6 +36,19 @@ calibration.arm_bases[side] ref_T_armbase, the arm base pose in it The camera that defines the reference frame needs no `ref_T_cam`. It is the identity by definition. +Each camera also declares its projection model and distortion coefficients. + +```text +calibration.cameras[c].model PINHOLE | OPENCV | KANNALA_BRANDT +calibration.cameras[c].distortion coefficients in that model's order +calibration.cameras[c].rectified whether the stored frames are rectified +``` + +A camera that declares no model is `PINHOLE` with no coefficients. No +projection site honors a non-pinhole model yet. Collect the declaration anyway: +it measures a vendor's rig, and a rig that has moved cannot be recalibrated +after the fact. + `Calibration.base_T_cam(side)` composes the two and returns what the EVA transform pipeline consumes. An episode that predates the block reaches the same value through the shim in `egomimic/rldb/zarr/calibration.py`: its diff --git a/egomimic/rldb/zarr/calibration.py b/egomimic/rldb/zarr/calibration.py index 71aee8f97..cad94e0de 100644 --- a/egomimic/rldb/zarr/calibration.py +++ b/egomimic/rldb/zarr/calibration.py @@ -48,7 +48,22 @@ #: `images.front_1`. The text after it is the camera name. IMAGE_KEY_PREFIX = "images." -_CAMERA_FIELDS = frozenset({"K", "resolution", "rectified", "ref_T_cam"}) +#: Projection model of a camera, and the distortion coefficient counts it +#: accepts. Nothing projects a non-pinhole model yet. The declaration is +#: collected now because it measures a vendor's rig: if we discover later that +#: we need it, the rig has moved and the episodes cannot be recalibrated. +CAMERA_MODELS = { + "PINHOLE": frozenset({0}), + "OPENCV": frozenset({4, 5, 8, 12, 14}), + "KANNALA_BRANDT": frozenset({4}), +} + +#: Projection model assumed for a camera that declares none. +DEFAULT_CAMERA_MODEL = "PINHOLE" + +_CAMERA_FIELDS = frozenset( + {"K", "model", "distortion", "resolution", "rectified", "ref_T_cam"} +) _CALIBRATION_FIELDS = frozenset({"reference_frame", "cameras", "arm_bases"}) @@ -97,6 +112,10 @@ class CameraCalibration: name: The camera name. It matches the ``images.`` array key. K: The 3×4 camera matrix, or ``None`` for a declared but uncalibrated camera. + model: A key in ``CAMERA_MODELS``. No projection site honors a + non-pinhole model yet. + distortion: The distortion coefficients of ``model``, in that model's + order. Empty for a pinhole camera. resolution: ``(width, height)`` in pixels, or ``None``. rectified: Whether the stored frames are already rectified. ref_T_cam: The 4×4 camera pose in the episode reference frame, or @@ -105,15 +124,22 @@ class CameraCalibration: name: str K: np.ndarray | None = None + model: str = DEFAULT_CAMERA_MODEL + distortion: tuple[float, ...] = () resolution: tuple[int, int] | None = None rectified: bool = True ref_T_cam: np.ndarray | None = None def to_jsonable(self) -> dict[str, Any]: """Return this camera as JSON-serializable episode metadata.""" - out: dict[str, Any] = {"rectified": bool(self.rectified)} + out: dict[str, Any] = { + "model": self.model, + "rectified": bool(self.rectified), + } if self.K is not None: out["K"] = self.K.tolist() + if self.distortion: + out["distortion"] = [float(c) for c in self.distortion] if self.resolution is not None: out["resolution"] = [int(self.resolution[0]), int(self.resolution[1])] if self.ref_T_cam is not None: @@ -311,6 +337,31 @@ def _parse_reference_frame(value: Any, cameras, where: str) -> str: ) +def _parse_distortion(raw: Any, model: str, where: str) -> tuple[float, ...]: + """Validate the distortion coefficients declared for one camera model.""" + if raw is None: + raw = () + if isinstance(raw, (str, bytes)) or not isinstance(raw, (list, tuple, np.ndarray)): + raise CalibrationError( + f"{where}.distortion: expected a list of coefficients, got {raw!r}" + ) + try: + coefficients = tuple(float(c) for c in raw) + except (TypeError, ValueError) as exc: + raise CalibrationError( + f"{where}.distortion: contains a non-numeric coefficient ({exc})" + ) from exc + if not all(np.isfinite(coefficients)): + raise CalibrationError(f"{where}.distortion: contains a non-finite value") + allowed = CAMERA_MODELS[model] + if len(coefficients) not in allowed: + raise CalibrationError( + f"{where}.distortion: model {model} takes " + f"{sorted(allowed)} coefficients, got {len(coefficients)}" + ) + return coefficients + + def _parse_camera(name: str, block: Any, where: str) -> CameraCalibration: if not isinstance(block, Mapping): raise CalibrationError(f"{where}: must be a mapping, got {block!r}") @@ -322,6 +373,15 @@ def _parse_camera(name: str, block: Any, where: str) -> CameraCalibration: ) K = block.get("K") + + model = block.get("model", DEFAULT_CAMERA_MODEL) + if model not in CAMERA_MODELS: + raise CalibrationError( + f"{where}.model: unknown camera model {model!r}; expected one of " + f"{sorted(CAMERA_MODELS)}" + ) + distortion = _parse_distortion(block.get("distortion"), model, where) + resolution = block.get("resolution") if resolution is not None: if ( @@ -345,6 +405,8 @@ def _parse_camera(name: str, block: Any, where: str) -> CameraCalibration: return CameraCalibration( name=name, K=None if K is None else _camera_matrix(K, f"{where}.K"), + model=model, + distortion=distortion, resolution=resolution, rectified=rectified, ref_T_cam=( @@ -509,6 +571,8 @@ def read_calibration(attrs: Mapping[str, Any]) -> Calibration | None: __all__ = [ "CAMERA_FRAME_PREFIX", + "CAMERA_MODELS", + "DEFAULT_CAMERA_MODEL", "IMAGE_KEY_PREFIX", "LEGACY_REFERENCE_CAMERA", "STATIC_REFERENCE_FRAMES", diff --git a/egomimic/rldb/zarr/test_calibration.py b/egomimic/rldb/zarr/test_calibration.py index f10e6c8f7..8710514e2 100644 --- a/egomimic/rldb/zarr/test_calibration.py +++ b/egomimic/rldb/zarr/test_calibration.py @@ -130,6 +130,54 @@ def test_parse_rejects_malformed_blocks(block, message) -> None: parse_calibration(block) +def test_a_camera_declares_pinhole_and_no_distortion_by_default() -> None: + calibration = parse_calibration( + { + "reference_frame": "robot_base", + "cameras": {"front_1": {"K": K_FRONT.tolist()}}, + } + ) + camera = calibration.cameras["front_1"] + assert (camera.model, camera.distortion) == ("PINHOLE", ()) + assert calibration.to_jsonable()["cameras"]["front_1"]["model"] == "PINHOLE" + + +def test_a_declared_lens_model_round_trips() -> None: + block = { + "reference_frame": "robot_base", + "cameras": { + "front_1": { + "K": K_FRONT.tolist(), + "model": "KANNALA_BRANDT", + "distortion": [0.1, -0.02, 0.003, -0.0004], + "rectified": False, + } + }, + } + camera = parse_calibration(block).cameras["front_1"] + assert camera.model == "KANNALA_BRANDT" + assert camera.distortion == (0.1, -0.02, 0.003, -0.0004) + reparsed = parse_calibration(parse_calibration(block).to_jsonable()) + assert reparsed.cameras["front_1"].distortion == camera.distortion + + +@pytest.mark.parametrize( + "camera, message", + [ + ({"model": "FISHEYE"}, "unknown camera model"), + ({"distortion": [0.1, 0.2]}, "model PINHOLE takes"), + ({"model": "KANNALA_BRANDT", "distortion": [0.1]}, "takes"), + ({"distortion": "0.1"}, "expected a list"), + ({"model": "OPENCV", "distortion": [0.1, 0.2, 0.3, "x"]}, "non-numeric"), + ], +) +def test_a_camera_model_and_its_coefficients_must_agree(camera, message) -> None: + with pytest.raises(CalibrationError, match=message): + parse_calibration( + {"reference_frame": "robot_base", "cameras": {"front_1": camera}} + ) + + def test_a_bare_3x3_camera_matrix_is_padded() -> None: calibration = parse_calibration( { From 98204391e619e656a99594ddbd11904f1ac502b4 Mon Sep 17 00:00:00 2001 From: jaynye Date: Wed, 2 Sep 2026 20:58:19 -0400 Subject: [PATCH 5/9] feat(zarr): schema-driven episode validator `python -m egomimic.rldb.zarr.validate [--strict]` checks one episode against `schema/episode_v3.yaml`. The rules live in the schema file; the module reads them and holds none of its own, so a contributor can read and diff the contract instead of reverse-engineering it from Python. It replaces the 300-line `validate_episode()` pasted into CONTRIBUTING_DATA.md, which is the only validator a contributor can run today and which passed an all-identity episode with 28 checks OK and 0 errors. Array rules resolve their widths through the registry, so `arm_dof` and the keypoint topology come from `platforms.yaml` and `end_effectors.yaml` rather than being restated. `{side}` expands over the arms the arity declares, and a rule applies only when its `when` conditions hold, so a human episode is never asked for a gripper. Rules the corpus does not meet yet are declared `required: strict`: a warning by default and an error under `--strict`. Flipping them first would turn a good check into an outage, since the intrinsics coverage rule alone fails every `eva_fold` episode. `total_frames` stays the sole authoritative length: axis 0 is a lower bound so a padded tail passes, and every other axis is exact. Adds `ResolvedEmbodiment.arity` and `.sides`, which the array rules need to know whether an episode owes one arm or two. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_012rCh1QmnyqgPZ5TJpsfMQb --- egomimic/rldb/embodiment/embodiment.py | 29 ++ egomimic/rldb/zarr/schema/episode_v3.yaml | 166 +++++++ egomimic/rldb/zarr/test_validate.py | 275 +++++++++++ egomimic/rldb/zarr/validate.py | 565 ++++++++++++++++++++++ pyproject.toml | 3 +- 5 files changed, 1037 insertions(+), 1 deletion(-) create mode 100644 egomimic/rldb/zarr/schema/episode_v3.yaml create mode 100644 egomimic/rldb/zarr/test_validate.py create mode 100644 egomimic/rldb/zarr/validate.py diff --git a/egomimic/rldb/embodiment/embodiment.py b/egomimic/rldb/embodiment/embodiment.py index b8a802028..2f730a25a 100644 --- a/egomimic/rldb/embodiment/embodiment.py +++ b/egomimic/rldb/embodiment/embodiment.py @@ -135,6 +135,35 @@ def action_space(self) -> str: ) return spaces.pop() + @property + def arity(self) -> str | None: + """Return the arm configuration this embodiment name selects. + + Returns: + One of the platform's ``arity`` values, or ``None`` when + resolution started from a morphology block with no name. + """ + if self.embodiment_name is None: + return None + prefix = f"{self.platform.embodiment_prefix}_" + if self.embodiment_name.startswith(prefix): + return self.embodiment_name[len(prefix) :] + return None + + @property + def sides(self) -> tuple[str, ...]: + """Return the arms this embodiment carries. + + A single-arm arity carries one side. Anything else carries both, which + is what a morphology block with no name resolves to. + """ + arity = self.arity + if arity == "left_arm": + return ("left",) + if arity == "right_arm": + return ("right",) + return SIDES + @property def embodiment_class(self) -> type["Embodiment"]: """Import and return the platform's configured ``Embodiment`` class. diff --git a/egomimic/rldb/zarr/schema/episode_v3.yaml b/egomimic/rldb/zarr/schema/episode_v3.yaml new file mode 100644 index 000000000..799612c3e --- /dev/null +++ b/egomimic/rldb/zarr/schema/episode_v3.yaml @@ -0,0 +1,166 @@ +# The episode contract, as data. +# +# `egomimic/rldb/zarr/validate.py` reads this file and holds no rules of its +# own. A rule that lives here is one a contributor can read, diff and review. +# +# Run it with: +# python -m egomimic.rldb.zarr.validate [--strict] +# +# --------------------------------------------------------------------------- +# attributes +# required true | false | strict +# `strict` means the attribute is an error only under --strict and +# a warning otherwise. Use it for a rule the corpus does not meet +# yet: flipping such a rule to `true` before the corpus is fixed +# turns a good check into an outage. +# type str | int | bool | mapping +# min smallest accepted value for an int +# choices the accepted values +# check a named check that needs more than a type; see `validate.py` +# +# arrays +# key an array key. `{side}` expands over the episode's arms, and `*` +# matches any suffix among the keys the episode stores. +# required true | false | strict, as above. A wildcard key is never +# required. +# when conditions on the resolved embodiment; the rule is skipped when +# they do not hold +# shape expected dimensions, see below +# dtype int | float | object | any +# +# shape dimensions +# T `total_frames`. Axis 0 may be longer, because a stored array may +# carry a padded tail; every other axis is exact. +# arm_dof the platform's joints per arm +# aux_dof the platform's auxiliary chain +# ee_dof the side's end-effector joints +# kp3 three coordinates per slot in the side's keypoint topology +# exactly that many +# "*" any +# +# when conditions +# platform_kind robot | human +# has_arm_chain the platform declares `arm_dof` +# has_aux_chain the platform declares `aux` +# end_effector_class parallel_jaw | human_hand | dexterous_hand +# --------------------------------------------------------------------------- + +version: v3.1 +known_versions: [v3.0, v3.1] + +attributes: + embodiment: + required: true + type: str + check: embodiment_name + total_frames: + required: true + type: int + min: 1 + fps: + required: true + type: int + min: 1 + task_name: + required: true + type: str + task_description: + required: false + type: str + features: + required: true + type: mapping + schema_version: + required: strict + type: str + check: schema_version + morphology: + required: false + type: mapping + check: morphology + calibration: + required: strict + type: mapping + check: calibration + intrinsics: + required: false + type: mapping + extrinsics: + required: false + type: mapping + +checks: + # Every episode must state the rig that recorded it, through the current + # `calibration` block or the legacy pair. + - name: calibration_present + required: true + # One camera matrix per stored image stream. `eva_to_zarr.py` calibrates one + # camera of three, so every EVA episode in the corpus fails this today. + - name: camera_coverage + required: strict + +arrays: + - key: "{side}.obs_ee_pose" + required: true + shape: [T, 7] + dtype: float + # The hand root or palm pose. It is the only key human, parallel-jaw and + # dexterous episodes share, so every embodiment owes it. + + - key: "{side}.cmd_ee_pose" + required: false + shape: [T, 7] + dtype: float + + - key: "{side}.obs_joints" + when: {has_arm_chain: true} + required: false + shape: [T, arm_dof] + dtype: float + + - key: "{side}.cmd_joints" + when: {has_arm_chain: true} + required: false + shape: [T, arm_dof] + dtype: float + + - key: "{side}.obs_gripper" + when: {end_effector_class: parallel_jaw} + required: true + shape: [T, 1] + dtype: float + + - key: "{side}.cmd_gripper" + when: {end_effector_class: parallel_jaw} + required: true + shape: [T, 1] + dtype: float + + - key: "{side}.obs_keypoints" + when: {end_effector_class: human_hand} + required: true + shape: [T, kp3] + dtype: float + + - key: obs_head_pose + when: {platform_kind: human} + required: true + shape: [T, 7] + dtype: float + + # One clock per episode, in integer UTC nanoseconds. EVA episodes do not + # write it yet, so it is required only under --strict. + - key: obs_rgb_timestamps_ns + required: strict + shape: [T] + dtype: int + + - key: "images.*" + required: false + shape: [T] + dtype: object + + - key: annotations + required: false + shape: ["*"] + dtype: object diff --git a/egomimic/rldb/zarr/test_validate.py b/egomimic/rldb/zarr/test_validate.py new file mode 100644 index 000000000..ec84ead4d --- /dev/null +++ b/egomimic/rldb/zarr/test_validate.py @@ -0,0 +1,275 @@ +"""Test the schema-driven episode validator.""" + +import numpy as np +import pytest + +from egomimic.rldb.embodiment.eva import Eva +from egomimic.rldb.zarr.validate import ( + ERROR, + OK, + WARNING, + load_schema, + main, + validate_episode, +) +from egomimic.rldb.zarr.zarr_writer import ZarrWriter + +K = np.array([[200.0, 0.0, 160.0, 0.0], [0.0, 200.0, 120.0, 0.0], [0.0, 0.0, 1.0, 0.0]]) +LENGTH = 4 + + +def _levels(report) -> dict[str, str]: + return {f.check: f.level for f in report.findings} + + +def _poses(x_offset: float = 0.0) -> np.ndarray: + poses = np.zeros((LENGTH, 7)) + poses[:, 0] = x_offset + np.arange(LENGTH) * 0.01 + poses[:, 3] = 1.0 + return poses + + +def _eva_numeric() -> dict: + data = {} + for side, offset in (("left", 0.0), ("right", 0.2)): + data[f"{side}.obs_ee_pose"] = _poses(offset) + data[f"{side}.cmd_ee_pose"] = _poses(offset + 0.05) + data[f"{side}.obs_joints"] = np.zeros((LENGTH, 6)) + data[f"{side}.cmd_joints"] = np.zeros((LENGTH, 6)) + data[f"{side}.obs_gripper"] = np.zeros((LENGTH, 1)) + data[f"{side}.cmd_gripper"] = np.zeros((LENGTH, 1)) + data["obs_rgb_timestamps_ns"] = np.arange(LENGTH, dtype=np.int64) + return data + + +def _write_eva(path, *, numeric=None, images=None, **kwargs) -> None: + kwargs.setdefault("calibration", { + "reference_frame": "camera:front_1", + "cameras": {"front_1": {"K": K.tolist()}}, + "arm_bases": { + side: np.linalg.inv(T).tolist() for side, T in Eva.EXTRINSICS.items() + }, + }) + kwargs.setdefault("metadata_override", {"schema_version": "v3.1"}) + ZarrWriter.create_and_write( + episode_path=path, + numeric_data=_eva_numeric() if numeric is None else numeric, + image_data=images, + embodiment=kwargs.pop("embodiment", "eva_bimanual"), + chunk_timesteps=LENGTH, + **kwargs, + ) + + +def test_a_complete_eva_episode_passes_under_strict(tmp_path) -> None: + path = tmp_path / "eva.zarr" + _write_eva(path, images={"images.front_1": np.zeros((LENGTH, 8, 8, 3), np.uint8)}) + + report = validate_episode(path, strict=True) + + assert report.ok, report.text() + assert not report.warnings + assert _levels(report)["embodiment"] == OK + + +def test_a_missing_required_array_is_an_error(tmp_path) -> None: + numeric = _eva_numeric() + del numeric["right.obs_ee_pose"] + _write_eva(tmp_path / "eva.zarr", numeric=numeric) + + report = validate_episode(tmp_path / "eva.zarr") + + assert _levels(report)["right.obs_ee_pose"] == ERROR + assert not report.ok + + +def test_a_wrong_width_is_reported_with_the_dimension_that_set_it(tmp_path) -> None: + numeric = _eva_numeric() + numeric["left.obs_joints"] = np.zeros((LENGTH, 5)) + _write_eva(tmp_path / "eva.zarr", numeric=numeric) + + report = validate_episode(tmp_path / "eva.zarr") + + message = next( + f.message for f in report.findings if f.check == "left.obs_joints" + ) + assert "axis 1 is 5, expected 6 (arm_dof)" in message + + +def test_a_padded_tail_is_not_an_error(tmp_path) -> None: + path = tmp_path / "eva.zarr" + ZarrWriter.create_and_write( + episode_path=path, + numeric_data=_eva_numeric(), + embodiment="eva_bimanual", + chunk_timesteps=3, # 4 frames pad out to 6 + intrinsics={"front_1": K}, + extrinsics=Eva.EXTRINSICS, + ) + + report = validate_episode(path) + + assert _levels(report)["left.obs_ee_pose"] == OK + assert report.ok, report.text() + + +def test_an_array_shorter_than_total_frames_is_an_error(tmp_path) -> None: + path = tmp_path / "eva.zarr" + _write_eva(path) + store = __import__("zarr").open_group(str(path), mode="a") + store.attrs["total_frames"] = LENGTH + 10 + + report = validate_episode(path) + + assert "holds 4 frames for total_frames 14" in report.text() + + +def test_strict_promotes_the_rules_the_corpus_does_not_meet_yet(tmp_path) -> None: + path = tmp_path / "legacy.zarr" + ZarrWriter.create_and_write( + episode_path=path, + numeric_data=_eva_numeric(), + image_data={ + "images.front_1": np.zeros((LENGTH, 8, 8, 3), np.uint8), + "images.left_wrist": np.zeros((LENGTH, 8, 8, 3), np.uint8), + }, + embodiment="eva_bimanual", + chunk_timesteps=LENGTH, + intrinsics={"front_1": K}, + extrinsics=Eva.EXTRINSICS, + ) + + lenient = validate_episode(path) + strict = validate_episode(path, strict=True) + + assert lenient.ok + assert _levels(lenient)["camera_coverage"] == WARNING + assert _levels(lenient)["attrs.calibration"] == WARNING + assert not strict.ok + assert _levels(strict)["camera_coverage"] == ERROR + assert "left_wrist" in strict.text() + + +def test_a_single_arm_episode_owes_only_its_own_arm(tmp_path) -> None: + numeric = { + k: v for k, v in _eva_numeric().items() if not k.startswith("right.") + } + _write_eva(tmp_path / "left.zarr", numeric=numeric, embodiment="eva_left_arm") + + report = validate_episode(tmp_path / "left.zarr", strict=True) + + assert report.ok, report.text() + assert "right.obs_ee_pose" not in _levels(report) + + +def test_a_human_episode_owes_keypoints_and_a_head_pose(tmp_path) -> None: + path = tmp_path / "human.zarr" + numeric = { + "left.obs_ee_pose": _poses(), + "right.obs_ee_pose": _poses(0.2), + "left.obs_keypoints": np.zeros((LENGTH, 63)), + "obs_head_pose": _poses(), + "obs_rgb_timestamps_ns": np.arange(LENGTH, dtype=np.int64), + } + ZarrWriter.create_and_write( + episode_path=path, + numeric_data=numeric, + embodiment="human_bimanual", + chunk_timesteps=LENGTH, + intrinsics={"front_1": K}, + metadata_override={"schema_version": "v3.1"}, + ) + + report = validate_episode(path) + + levels = _levels(report) + assert levels["left.obs_keypoints"] == OK + assert levels["obs_head_pose"] == OK + # 21 MANO slots at three coordinates each. + assert levels["right.obs_keypoints"] == ERROR + # A human platform has no arm chain and no jaw, so neither rule runs. + assert "left.obs_gripper" not in levels + assert "left.obs_joints" not in levels + + +def test_an_unknown_embodiment_stops_before_the_array_rules(tmp_path) -> None: + path = tmp_path / "eva.zarr" + _write_eva(path) + store = __import__("zarr").open_group(str(path), mode="a") + store.attrs["embodiment"] = "sharpa_bimanual" + + report = validate_episode(path) + + levels = _levels(report) + assert levels["attrs.embodiment"] == ERROR + assert levels["embodiment"] == ERROR + assert "left.obs_ee_pose" not in levels + + +def test_a_morphology_block_must_agree_with_the_embodiment_name(tmp_path) -> None: + path = tmp_path / "eva.zarr" + _write_eva(path) + store = __import__("zarr").open_group(str(path), mode="a") + store.attrs["morphology"] = { + "platform": "human_body", + "end_effector": "mano_hand", + } + + report = validate_episode(path) + + assert _levels(report)["attrs.morphology"] == ERROR + assert "disagrees" in report.text() + + +def test_a_morphology_block_selects_the_end_effector(tmp_path) -> None: + path = tmp_path / "eva.zarr" + _write_eva(path) + store = __import__("zarr").open_group(str(path), mode="a") + store.attrs["morphology"] = { + "platform": "eva_x5", + "end_effector": {"left": "eva_parallel_jaw", "right": "eva_parallel_jaw"}, + "vendor": "rl2", + } + + report = validate_episode(path, strict=True) + + assert report.ok, report.text() + assert "eva_parallel_jaw" in report.text(verbose=True) + + +def test_a_directory_that_is_not_an_episode_reports_one_error(tmp_path) -> None: + report = validate_episode(tmp_path / "absent.zarr") + + assert not report.ok + assert _levels(report)["episode"] == ERROR + + +def test_the_cli_exit_code_follows_the_findings(tmp_path, capsys) -> None: + path = tmp_path / "eva.zarr" + _write_eva(path) + + assert main([str(path)]) == 0 + assert main([str(path), "--strict"]) == 0 + assert main([str(tmp_path / "absent.zarr")]) == 1 + assert "cannot open as a zarr group" in capsys.readouterr().out + + +def test_the_cli_can_report_json(tmp_path, capsys) -> None: + path = tmp_path / "eva.zarr" + _write_eva(path) + + main([str(path), "--json"]) + + payload = __import__("json").loads(capsys.readouterr().out) + assert payload[0]["ok"] is True + assert payload[0]["path"] == str(path) + + +@pytest.mark.parametrize("section", ["attributes", "checks", "arrays"]) +def test_every_schema_rule_declares_a_usable_requirement(section) -> None: + schema = load_schema() + rules = schema[section] + entries = rules.values() if isinstance(rules, dict) else rules + assert entries + for rule in entries: + assert rule.get("required", False) in (True, False, "strict") diff --git a/egomimic/rldb/zarr/validate.py b/egomimic/rldb/zarr/validate.py new file mode 100644 index 000000000..23b7e4a7d --- /dev/null +++ b/egomimic/rldb/zarr/validate.py @@ -0,0 +1,565 @@ +"""Validate one zarr episode against the schema in ``schema/episode_v3.yaml``. + +Run it with:: + + python -m egomimic.rldb.zarr.validate [--strict] + +The rules live in the schema file, not here. This module reads them, resolves +the episode's embodiment through the registry, and reports one finding per +rule. ``--strict`` promotes the rules the corpus does not meet yet from +warnings to errors, so a rule can land, be measured across the corpus, and +only then become the default. +""" + +from __future__ import annotations + +import argparse +import functools +import json +import sys +from collections.abc import Mapping +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import yaml +import zarr + +from egomimic.rldb.embodiment.embodiment import ( + Embodiment, + ResolvedEmbodiment, + canonical_embodiment_name, +) +from egomimic.rldb.embodiment.registry import load_embodiment_platforms +from egomimic.rldb.zarr.calibration import ( + CalibrationError, + read_calibration, + uncalibrated_cameras, +) + +SCHEMA_DIR = Path(__file__).parent / "schema" +SCHEMA_FILE = SCHEMA_DIR / "episode_v3.yaml" + +ERROR = "error" +WARNING = "warning" +OK = "ok" + +#: `required: strict` rules are warnings until `--strict` promotes them. They +#: exist so a rule can land before the corpus meets it. +_REQUIRED_VALUES = (True, False, "strict") + +_TYPE_NAMES = { + "str": str, + "int": int, + "bool": bool, + "mapping": Mapping, +} + + +class SchemaError(ValueError): + """Report an invalid rule in ``episode_v3.yaml``.""" + + +@dataclass(frozen=True) +class Finding: + """One rule's result. + + Attributes: + level: ``ok``, ``warning``, or ``error``. + check: The rule that produced this finding. + message: What the rule found, in one line. + """ + + level: str + check: str + message: str + + def __str__(self) -> str: + return f"{self.level.upper():<7} {self.check}: {self.message}" + + +@dataclass +class Report: + """Every finding for one episode. + + Attributes: + path: The episode directory. + findings: One finding per rule that ran, in schema order. + strict: Whether the strict rules were promoted to errors. + """ + + path: Path + findings: list[Finding] = field(default_factory=list) + strict: bool = False + + def add(self, level: str, check: str, message: str) -> None: + self.findings.append(Finding(level, check, message)) + + @property + def errors(self) -> list[Finding]: + return [f for f in self.findings if f.level == ERROR] + + @property + def warnings(self) -> list[Finding]: + return [f for f in self.findings if f.level == WARNING] + + @property + def ok(self) -> bool: + return not self.errors + + def summary(self) -> str: + return ( + f"{len(self.findings)} checks, {len(self.errors)} errors, " + f"{len(self.warnings)} warnings" + ) + + def text(self, verbose: bool = False) -> str: + """Render the report. + + Args: + verbose: If true, list the rules that passed as well. + """ + lines = [str(self.path)] + for finding in self.findings: + if finding.level != OK or verbose: + lines.append(f" {finding}") + lines.append(f" {self.summary()}") + return "\n".join(lines) + + def to_jsonable(self) -> dict[str, Any]: + return { + "path": str(self.path), + "strict": self.strict, + "ok": self.ok, + "findings": [ + {"level": f.level, "check": f.check, "message": f.message} + for f in self.findings + ], + } + + +@functools.lru_cache(maxsize=1) +def load_schema() -> dict: + """Load and check ``schema/episode_v3.yaml``. + + Returns: + The parsed schema. + + Raises: + SchemaError: If a rule declares an unusable ``required`` value. + """ + with SCHEMA_FILE.open("r") as f: + schema = yaml.safe_load(f) or {} + rules = list(schema.get("attributes", {}).items()) + rules += [(r.get("name"), r) for r in schema.get("checks", [])] + rules += [(r.get("key"), r) for r in schema.get("arrays", [])] + for name, rule in rules: + required = rule.get("required", False) + if required not in _REQUIRED_VALUES: + raise SchemaError( + f"{name!r}: `required` must be one of " + f"{list(_REQUIRED_VALUES)}, got {required!r}" + ) + return schema + + +def _level(required, strict: bool) -> str | None: + """Return the level a failed rule reports at, or ``None`` if it is optional.""" + if required is True: + return ERROR + if required == "strict": + return ERROR if strict else WARNING + return None + + +# --------------------------------------------------------------------------- +# Attribute rules +# --------------------------------------------------------------------------- + + +def _check_type(value, type_name: str | None) -> str | None: + if type_name is None: + return None + expected = _TYPE_NAMES.get(type_name) + if expected is None: + raise SchemaError(f"unknown type {type_name!r} in the schema") + # A bool is an int in Python; the schema means them separately. + if expected is int and isinstance(value, bool): + return "expected an int, got a bool" + if not isinstance(value, expected): + return f"expected {type_name}, got {type(value).__name__}" + return None + + +def _check_attribute(name: str, rule: dict, attrs: Mapping, report, context) -> None: + if name not in attrs: + level = _level(rule.get("required", False), report.strict) + if level is not None: + report.add(level, f"attrs.{name}", "missing") + return + + value = attrs[name] + problem = _check_type(value, rule.get("type")) + if problem is None and "min" in rule and value < rule["min"]: + problem = f"must be at least {rule['min']}, got {value}" + if problem is None and "choices" in rule and value not in rule["choices"]: + problem = f"expected one of {rule['choices']}, got {value!r}" + if problem is None and rule.get("check"): + problem = _ATTRIBUTE_CHECKS[rule["check"]](value, attrs, context) + + if problem is None: + report.add(OK, f"attrs.{name}", _describe(value)) + else: + report.add(ERROR, f"attrs.{name}", problem) + + +def _describe(value) -> str: + if isinstance(value, Mapping): + return f"{len(value)} entries" + text = str(value) + return text if len(text) <= 60 else text[:57] + "..." + + +def _check_embodiment_name(value, attrs, context) -> str | None: + if canonical_embodiment_name(value) not in load_embodiment_platforms(): + return ( + f"{value!r} is not owned by any platform in registry/platforms.yaml; " + f"known: {sorted(load_embodiment_platforms())}" + ) + return None + + +def _check_morphology(value, attrs, context) -> str | None: + try: + resolved = Embodiment.resolve(value) + except (TypeError, ValueError) as exc: + return str(exc) + named = context.get("named_platform") + if named is not None and resolved.platform.name != named.name: + return ( + f"platform {resolved.platform.name!r} disagrees with the platform " + f"{named.name!r} that embodiment {attrs.get('embodiment')!r} selects" + ) + return None + + +def _check_calibration(value, attrs, context) -> str | None: + try: + read_calibration(attrs) + except CalibrationError as exc: + return str(exc) + return None + + +def _check_schema_version(value, attrs, context) -> str | None: + known = load_schema().get("known_versions") or [] + if known and value not in known: + return f"unknown schema_version {value!r}; known: {known}" + return None + + +_ATTRIBUTE_CHECKS = { + "embodiment_name": _check_embodiment_name, + "morphology": _check_morphology, + "calibration": _check_calibration, + "schema_version": _check_schema_version, +} + + +# --------------------------------------------------------------------------- +# Named checks +# --------------------------------------------------------------------------- + + +def _check_calibration_present(rule, context, report) -> None: + calibration = context.get("calibration") + level = _level(rule.get("required", False), report.strict) + if calibration is None or not calibration.intrinsics(): + if level is not None: + report.add( + level, + "calibration_present", + "the episode states no camera matrix; calibration measures the " + "rig that recorded it and cannot be recovered later", + ) + return + report.add( + OK, + "calibration_present", + f"{len(calibration.intrinsics())} calibrated camera(s), " + f"reference_frame={calibration.reference_frame}", + ) + + +def _check_camera_coverage(rule, context, report) -> None: + missing = uncalibrated_cameras(context["array_keys"], context.get("calibration")) + level = _level(rule.get("required", False), report.strict) + if missing: + if level is not None: + report.add( + level, + "camera_coverage", + f"no camera matrix for image stream(s) {missing}", + ) + return + report.add(OK, "camera_coverage", "every image stream has a camera matrix") + + +_NAMED_CHECKS = { + "calibration_present": _check_calibration_present, + "camera_coverage": _check_camera_coverage, +} + + +# --------------------------------------------------------------------------- +# Array rules +# --------------------------------------------------------------------------- + + +def _condition_holds(when: Mapping, resolved: ResolvedEmbodiment, side) -> bool: + for name, expected in when.items(): + if name == "platform_kind": + actual = resolved.platform.kind + elif name == "has_arm_chain": + actual = resolved.platform.arm_dof is not None + elif name == "has_aux_chain": + actual = resolved.platform.aux is not None + elif name == "end_effector_class": + end_effector = resolved.end_effectors.get(side) if side else None + actual = None if end_effector is None else end_effector.ee_class + else: + raise SchemaError(f"unknown `when` condition {name!r} in the schema") + if actual != expected: + return False + return True + + +def _dimension(token, context, side) -> int | None: + """Resolve one schema shape token to a length, or ``None`` for unknown.""" + if isinstance(token, int): + return token + if token == "*": + return None + resolved: ResolvedEmbodiment = context["resolved"] + if token == "T": + return context["total_frames"] + if token == "arm_dof": + return resolved.platform.arm_dof + if token == "aux_dof": + return None if resolved.platform.aux is None else resolved.platform.aux.dof + if token == "ee_dof": + end_effector = resolved.end_effectors.get(side) if side else None + return None if end_effector is None else end_effector.dof + if token == "kp3": + end_effector = resolved.end_effectors.get(side) if side else None + if end_effector is None: + return None + return 3 * end_effector.keypoints.n_slots + raise SchemaError(f"unknown shape dimension {token!r} in the schema") + + +def _dtype_kind(dtype) -> str: + kind = getattr(dtype, "kind", "") + if kind in "iu": + return "int" + if kind == "f": + return "float" + if kind in "OSV": + return "object" + return kind or "any" + + +def _check_array(key: str, rule: dict, arrays: Mapping, report, context, side) -> None: + if key not in arrays: + level = _level(rule.get("required", False), report.strict) + if level is not None: + report.add(level, key, "missing") + return + + array = arrays[key] + shape = tuple(int(n) for n in array.shape) + expected = rule.get("shape") + problems = [] + if expected is not None: + if len(shape) != len(expected): + problems.append( + f"expected {len(expected)} dimension(s), got shape {shape}" + ) + else: + for axis, token in enumerate(expected): + want = _dimension(token, context, side) + if want is None: + continue + # `total_frames` is the sole authoritative length and a stored + # array may carry a padded tail, so axis 0 is a lower bound. + if axis == 0: + if shape[0] < want: + problems.append( + f"holds {shape[0]} frames for total_frames {want}" + ) + elif shape[axis] != want: + problems.append( + f"axis {axis} is {shape[axis]}, expected {want} ({token})" + ) + + want_dtype = rule.get("dtype") + if want_dtype and want_dtype != "any": + actual = _dtype_kind(array.dtype) + if actual != want_dtype: + problems.append(f"dtype is {array.dtype}, expected {want_dtype}") + + if problems: + report.add(ERROR, key, "; ".join(problems)) + else: + report.add(OK, key, f"shape {shape} {array.dtype}") + + +def _expand_key(template: str, arrays: Mapping, sides) -> list[tuple[str, str | None]]: + """Expand one schema key into the concrete keys it names. + + ``{side}`` expands over the episode's arms. ``*`` matches the keys the + episode stores, so a wildcard rule checks what is there and never demands + a key. + """ + if "{side}" in template: + candidates = [(template.format(side=side), side) for side in sides] + else: + candidates = [(template, None)] + out = [] + for key, side in candidates: + if "*" not in key: + out.append((key, side)) + continue + prefix, _, suffix = key.partition("*") + out.extend( + (name, side) + for name in arrays + if name.startswith(prefix) and name.endswith(suffix) + ) + return out + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + + +def validate_episode(path: str | Path, *, strict: bool = False) -> Report: + """Validate one zarr episode against ``schema/episode_v3.yaml``. + + Args: + path: The episode ``.zarr`` directory. + strict: If true, promote the rules the corpus does not meet yet from + warnings to errors. + + Returns: + A report holding one finding per rule that ran. + """ + path = Path(path) + report = Report(path=path, strict=strict) + try: + store = zarr.open_group(str(path), mode="r") + except Exception as exc: + report.add(ERROR, "episode", f"cannot open as a zarr group: {exc}") + return report + + schema = load_schema() + attrs = dict(store.attrs) + arrays = {name: store[name] for name in store.array_keys()} + + context: dict[str, Any] = {"array_keys": list(arrays)} + embodiment = attrs.get("embodiment") + if embodiment: + context["named_platform"] = load_embodiment_platforms().get( + canonical_embodiment_name(embodiment) + ) + try: + context["calibration"] = read_calibration(attrs) + except CalibrationError: + context["calibration"] = None + + for name, rule in schema.get("attributes", {}).items(): + _check_attribute(name, rule, attrs, report, context) + + for rule in schema.get("checks", []): + check = _NAMED_CHECKS.get(rule.get("name")) + if check is None: + raise SchemaError(f"unknown check {rule.get('name')!r} in the schema") + check(rule, context, report) + + resolved = _resolve(attrs, report) + if resolved is None: + return report + # Say which platform and end-effectors the array rules ran against, so a + # report explains itself without a second lookup. + report.add(OK, "embodiment", resolved.describe()) + context["resolved"] = resolved + context["total_frames"] = attrs.get("total_frames") + + for rule in schema.get("arrays", []): + for key, side in _expand_key(rule["key"], arrays, resolved.sides): + if side is not None and side not in resolved.end_effectors: + continue + when = rule.get("when") or {} + if when and not _condition_holds(when, resolved, side): + continue + _check_array(key, rule, arrays, report, context, side) + + return report + + +def _resolve(attrs: Mapping, report: Report) -> ResolvedEmbodiment | None: + """Resolve the episode's embodiment, preferring its morphology block.""" + for spec in (attrs.get("morphology"), attrs.get("embodiment")): + if not spec: + continue + try: + return Embodiment.resolve(spec) + except (TypeError, ValueError): + continue + report.add( + ERROR, + "embodiment", + "cannot resolve the episode's embodiment, so no array rule can run", + ) + return None + + +def main(argv: list[str] | None = None) -> int: + """Run the validator over one or more episodes. + + Returns: + ``0`` if every episode passed, ``1`` otherwise. + """ + parser = argparse.ArgumentParser( + prog="python -m egomimic.rldb.zarr.validate", + description=__doc__.splitlines()[0], + ) + parser.add_argument("paths", nargs="+", type=Path, help="episode .zarr paths") + parser.add_argument( + "--strict", + action="store_true", + help="promote the rules the corpus does not meet yet to errors", + ) + parser.add_argument( + "-v", + "--verbose", + action="store_true", + help="list the rules that passed as well", + ) + parser.add_argument( + "--json", action="store_true", help="print one JSON report per episode" + ) + args = parser.parse_args(argv) + + reports = [validate_episode(p, strict=args.strict) for p in args.paths] + if args.json: + print(json.dumps([r.to_jsonable() for r in reports], indent=2)) + else: + for report in reports: + print(report.text(verbose=args.verbose)) + return 0 if all(r.ok for r in reports) else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/pyproject.toml b/pyproject.toml index 1b6471417..d46a86b10 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,9 +69,10 @@ dependencies = [ [tool.setuptools.packages.find] where = ["."] -# Include the registry YAML files in built distributions. +# Include the registry and schema YAML files in built distributions. [tool.setuptools.package-data] "egomimic.rldb.embodiment.registry" = ["*.yaml"] +"egomimic.rldb.zarr" = ["schema/*.yaml"] [tool.pytest.ini_options] testpaths = ["egomimic/rldb/embodiment", "egomimic/rldb/zarr"] From ef1bcca5892f3226c1305212a590157921d11da9 Mon Sep 17 00:00:00 2001 From: jaynye Date: Wed, 2 Sep 2026 21:01:59 -0400 Subject: [PATCH 6/9] feat(zarr): degeneracy, timestamp and annotation rules Six rules that turn the failure modes we have actually been shipped into findings at ingest instead of a policy that trains badly six weeks later. - pose_degeneracy: a track that never moves, or that holds an exact identity rotation on more than 1% of frames, is a placeholder, not a measurement. - calibration_degeneracy: an exactly-identity extrinsic puts the camera at the arm base. It survives every other check, because the overlay it draws still looks plausible, bunched at the image centre. - intrinsics_signature: `fx == fy == W` with the principal point at the exact image centre is a synthesized camera. The `fx == W` conjunct is what keeps our own rectified Aria K, fx 266.5 at W 640, out of the net. - timestamps: strictly increasing, and a warning when the step is a power of two of at least 64 ns, which is float64 seconds converted to nanoseconds. A second stored time base is an error. - annotation_coverage: annotations cover 90% of the episode or the tail is trimmed. Opt-in, since EVA episodes pass no annotations at all. - annotation_text: no delimiter-encoded metadata. A skill taxonomy defined from one vendor's data is a taxonomy we will regret. Every rule and its threshold lives in the schema file. Run against the Sharpa sample the rules reproduce each measurement in the review by hand: 370 identity frames on all three pose tracks, identity extrinsics, the synthesized 480-pixel K on all three cameras, three duplicate timestamps under 256 ns quantization, 76% annotation coverage, and the ` | Skill: pick` suffix. The validator that shipped in CONTRIBUTING_DATA.md passes that episode with 28 checks OK and 0 errors. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_012rCh1QmnyqgPZ5TJpsfMQb --- egomimic/rldb/zarr/schema/episode_v3.yaml | 43 ++++ egomimic/rldb/zarr/test_validate.py | 183 +++++++++++++++- egomimic/rldb/zarr/validate.py | 248 +++++++++++++++++++++- 3 files changed, 471 insertions(+), 3 deletions(-) diff --git a/egomimic/rldb/zarr/schema/episode_v3.yaml b/egomimic/rldb/zarr/schema/episode_v3.yaml index 799612c3e..484284e78 100644 --- a/egomimic/rldb/zarr/schema/episode_v3.yaml +++ b/egomimic/rldb/zarr/schema/episode_v3.yaml @@ -94,11 +94,54 @@ checks: # `calibration` block or the legacy pair. - name: calibration_present required: true + # One camera matrix per stored image stream. `eva_to_zarr.py` calibrates one # camera of three, so every EVA episode in the corpus fails this today. - name: camera_coverage required: strict + # A pose track that never moves, or that holds an exact identity rotation on + # more than a slice of its frames, is a placeholder rather than a + # measurement. Nothing downstream can tell the two apart. + - name: pose_degeneracy + required: true + suffixes: [obs_ee_pose, cmd_ee_pose, obs_head_pose] + identity_rotation_fraction: 0.01 + + # An exactly-identity extrinsic places the camera at the arm base. It is the + # signature of a rig that was never measured, and it survives every other + # check because the projection it produces still looks plausible. + - name: calibration_degeneracy + required: true + + # `fx == fy == W` with the principal point at the exact image centre is a + # synthesized camera, not a calibrated one. The `fx == W` conjunct is what + # keeps our own rectified Aria K (fx 266.5, W 640) out of the net. + - name: intrinsics_signature + required: true + + # One clock per episode, in integer UTC nanoseconds, strictly increasing. A + # relative time base is derived on read, never stored. + - name: timestamps + required: true + key: obs_rgb_timestamps_ns + banned_key_substrings: [relative_timestamp] + + # Annotations cover the retained frames or the tail is trimmed. EVA episodes + # pass no annotations at all today, so this is opt-in. + - name: annotation_coverage + required: strict + key: annotations + minimum: 0.9 + + # A vendor who wants a skill taxonomy encodes one in the text, as in + # `pick up the cup | Skill: pick`. Ban the delimiter now and add a field when + # a second vendor asks for it. + - name: annotation_text + required: true + key: annotations + banned_delimiters: [" | "] + arrays: - key: "{side}.obs_ee_pose" required: true diff --git a/egomimic/rldb/zarr/test_validate.py b/egomimic/rldb/zarr/test_validate.py index ec84ead4d..829991985 100644 --- a/egomimic/rldb/zarr/test_validate.py +++ b/egomimic/rldb/zarr/test_validate.py @@ -23,9 +23,12 @@ def _levels(report) -> dict[str, str]: def _poses(x_offset: float = 0.0) -> np.ndarray: + """Return a moving pose track: xyz plus a rotating [qw, qx, qy, qz].""" + angles = np.linspace(0.1, 0.4, LENGTH) poses = np.zeros((LENGTH, 7)) poses[:, 0] = x_offset + np.arange(LENGTH) * 0.01 - poses[:, 3] = 1.0 + poses[:, 3] = np.cos(angles / 2) + poses[:, 6] = np.sin(angles / 2) return poses @@ -51,6 +54,7 @@ def _write_eva(path, *, numeric=None, images=None, **kwargs) -> None: }, }) kwargs.setdefault("metadata_override", {"schema_version": "v3.1"}) + kwargs.setdefault("annotations", [("fold the towel", 0, LENGTH)]) ZarrWriter.create_and_write( episode_path=path, numeric_data=_eva_numeric() if numeric is None else numeric, @@ -103,6 +107,7 @@ def test_a_padded_tail_is_not_an_error(tmp_path) -> None: numeric_data=_eva_numeric(), embodiment="eva_bimanual", chunk_timesteps=3, # 4 frames pad out to 6 + annotations=[("fold the towel", 0, LENGTH)], intrinsics={"front_1": K}, extrinsics=Eva.EXTRINSICS, ) @@ -135,6 +140,7 @@ def test_strict_promotes_the_rules_the_corpus_does_not_meet_yet(tmp_path) -> Non }, embodiment="eva_bimanual", chunk_timesteps=LENGTH, + annotations=[("fold the towel", 0, LENGTH)], intrinsics={"front_1": K}, extrinsics=Eva.EXTRINSICS, ) @@ -176,6 +182,7 @@ def test_a_human_episode_owes_keypoints_and_a_head_pose(tmp_path) -> None: numeric_data=numeric, embodiment="human_bimanual", chunk_timesteps=LENGTH, + annotations=[("wave", 0, LENGTH)], intrinsics={"front_1": K}, metadata_override={"schema_version": "v3.1"}, ) @@ -273,3 +280,177 @@ def test_every_schema_rule_declares_a_usable_requirement(section) -> None: assert entries for rule in entries: assert rule.get("required", False) in (True, False, "strict") + + +# -------------------------------------------------------------------------- +# Degeneracy, timestamps and annotations +# -------------------------------------------------------------------------- + + +def _finding(report, check: str) -> str: + return next(f.message for f in report.findings if f.check == check) + + +def test_a_pose_track_that_never_moves_is_an_error(tmp_path) -> None: + numeric = _eva_numeric() + numeric["left.obs_ee_pose"] = np.tile(_poses()[0], (LENGTH, 1)) + _write_eva(tmp_path / "still.zarr", numeric=numeric) + + report = validate_episode(tmp_path / "still.zarr") + + assert _levels(report)["pose_degeneracy"] == ERROR + assert "constant across all 4 frames" in _finding(report, "pose_degeneracy") + + +def test_identity_rotations_beyond_the_limit_are_an_error(tmp_path) -> None: + numeric = _eva_numeric() + identity = numeric["left.obs_ee_pose"].copy() + identity[:, 3:7] = [1.0, 0.0, 0.0, 0.0] + numeric["left.obs_ee_pose"] = identity + _write_eva(tmp_path / "identity.zarr", numeric=numeric) + + report = validate_episode(tmp_path / "identity.zarr") + + assert _levels(report)["pose_degeneracy"] == ERROR + assert "identity rotation on 100% of frames" in _finding(report, "pose_degeneracy") + + +def test_an_identity_extrinsic_is_an_error(tmp_path) -> None: + path = tmp_path / "identity_rig.zarr" + ZarrWriter.create_and_write( + episode_path=path, + numeric_data=_eva_numeric(), + embodiment="eva_bimanual", + chunk_timesteps=LENGTH, + annotations=[("fold the towel", 0, LENGTH)], + intrinsics={"front_1": K}, + extrinsics={"left": np.eye(4), "right": np.eye(4)}, + ) + + report = validate_episode(path) + + assert _levels(report)["calibration_degeneracy"] == ERROR + assert "left arm base" in _finding(report, "calibration_degeneracy") + + +def test_a_synthesized_camera_matrix_is_an_error(tmp_path) -> None: + width = height = 480 + synthetic = np.array( + [ + [width, 0.0, width / 2, 0.0], + [0.0, width, height / 2, 0.0], + [0.0, 0.0, 1.0, 0.0], + ] + ) + _write_eva( + tmp_path / "synthetic.zarr", + calibration={ + "reference_frame": "camera:front_1", + "cameras": { + "front_1": {"K": synthetic.tolist(), "resolution": [width, height]} + }, + }, + ) + + report = validate_episode(tmp_path / "synthetic.zarr") + + assert _levels(report)["intrinsics_signature"] == ERROR + assert "synthesized centred pinhole" in _finding(report, "intrinsics_signature") + + +def test_a_rectified_aria_camera_matrix_stays_out_of_the_net(tmp_path) -> None: + # fx is 266.5 at W 640, so the `fx == W` conjunct is what saves it. + aria = np.array( + [[266.5, 0.0, 320.0, 0.0], [0.0, 266.5, 240.0, 0.0], [0.0, 0.0, 1.0, 0.0]] + ) + _write_eva( + tmp_path / "aria.zarr", + calibration={ + "reference_frame": "camera:front_1", + "cameras": {"front_1": {"K": aria.tolist(), "resolution": [640, 480]}}, + }, + ) + + report = validate_episode(tmp_path / "aria.zarr") + + assert _levels(report)["intrinsics_signature"] == OK + + +def test_a_stalled_clock_is_an_error(tmp_path) -> None: + numeric = _eva_numeric() + numeric["obs_rgb_timestamps_ns"] = np.array( + [1_000, 2_000, 2_000, 1_500], dtype=np.int64 + ) + _write_eva(tmp_path / "clock.zarr", numeric=numeric) + + report = validate_episode(tmp_path / "clock.zarr") + + assert _levels(report)["timestamps"] == ERROR + assert "does not increase on 2 of 3 steps" in _finding(report, "timestamps") + + +def test_a_float64_quantized_clock_is_an_error(tmp_path) -> None: + numeric = _eva_numeric() + numeric["obs_rgb_timestamps_ns"] = ( + np.int64(1_700_000_000_000_000_000) + np.arange(LENGTH, dtype=np.int64) * 256 + ) + _write_eva(tmp_path / "quantized.zarr", numeric=numeric) + + report = validate_episode(tmp_path / "quantized.zarr") + + assert "quantized to 256 ns" in _finding(report, "timestamps") + + +def test_a_second_time_base_is_an_error(tmp_path) -> None: + numeric = _eva_numeric() + numeric["relative_timestamp_s"] = np.linspace(0.0, 0.1, LENGTH)[:, None] + _write_eva(tmp_path / "two_clocks.zarr", numeric=numeric) + + report = validate_episode(tmp_path / "two_clocks.zarr") + + assert _levels(report)["timestamps"] == ERROR + assert "one clock per episode" in _finding(report, "timestamps") + + +def test_annotation_coverage_below_the_minimum_is_reported(tmp_path) -> None: + _write_eva( + tmp_path / "thin.zarr", annotations=[("fold the towel", 0, LENGTH - 2)] + ) + + lenient = validate_episode(tmp_path / "thin.zarr") + strict = validate_episode(tmp_path / "thin.zarr", strict=True) + + assert _levels(lenient)["annotation_coverage"] == WARNING + assert _levels(strict)["annotation_coverage"] == ERROR + assert "cover 50% of the episode" in _finding(strict, "annotation_coverage") + + +def test_an_annotation_span_past_the_episode_is_reported(tmp_path) -> None: + _write_eva(tmp_path / "over.zarr", annotations=[("fold", 0, LENGTH + 5)]) + + report = validate_episode(tmp_path / "over.zarr", strict=True) + + assert "outside [0, 4)" in _finding(report, "annotation_coverage") + + +def test_delimiter_encoded_metadata_is_an_error(tmp_path) -> None: + _write_eva( + tmp_path / "skill.zarr", + annotations=[("pick up the cup | Skill: pick", 0, LENGTH)], + ) + + report = validate_episode(tmp_path / "skill.zarr") + + assert _levels(report)["annotation_text"] == ERROR + assert "Skill: pick" in _finding(report, "annotation_text") + + +def test_delimiter_encoded_metadata_in_the_task_description_is_an_error( + tmp_path, +) -> None: + _write_eva(tmp_path / "desc.zarr", task_description="fold | Skill: fold") + + report = validate_episode(tmp_path / "desc.zarr") + + assert _levels(report)["annotation_text"] == ERROR + assert "task_description" in _finding(report, "annotation_text") diff --git a/egomimic/rldb/zarr/validate.py b/egomimic/rldb/zarr/validate.py index 23b7e4a7d..cdae2c046 100644 --- a/egomimic/rldb/zarr/validate.py +++ b/egomimic/rldb/zarr/validate.py @@ -22,6 +22,7 @@ from pathlib import Path from typing import Any +import numpy as np import yaml import zarr @@ -32,6 +33,7 @@ ) from egomimic.rldb.embodiment.registry import load_embodiment_platforms from egomimic.rldb.zarr.calibration import ( + IMAGE_KEY_PREFIX, CalibrationError, read_calibration, uncalibrated_cameras, @@ -305,9 +307,246 @@ def _check_camera_coverage(rule, context, report) -> None: report.add(OK, "camera_coverage", "every image stream has a camera matrix") +_IDENTITY = np.eye(4) + +#: Quaternion layout is `[qw, qx, qy, qz]`; both signs are the same rotation. +_IDENTITY_QUATERNIONS = ( + np.array([1.0, 0.0, 0.0, 0.0]), + np.array([-1.0, 0.0, 0.0, 0.0]), +) + + +def _read(array, total_frames: int | None) -> np.ndarray: + """Read an array up to `total_frames`, which is the authoritative length.""" + end = array.shape[0] if total_frames is None else min(total_frames, array.shape[0]) + return np.asarray(array[:end]) + + +def _report_problems(rule, report, check: str, problems, passed: str) -> None: + """Record one finding for a named check.""" + if not problems: + report.add(OK, check, passed) + return + level = _level(rule.get("required", False), report.strict) + if level is not None: + report.add(level, check, "; ".join(problems)) + + +def _check_pose_degeneracy(rule, context, report) -> None: + suffixes = tuple(rule.get("suffixes") or ()) + threshold = float(rule.get("identity_rotation_fraction", 0.01)) + total_frames = context.get("total_frames") + problems = [] + checked = 0 + for key, array in context["arrays"].items(): + if not key.endswith(suffixes) or len(array.shape) != 2: + continue + poses = _read(array, total_frames) + if poses.shape[0] < 2 or poses.shape[1] != 7: + continue + checked += 1 + if len(np.unique(poses, axis=0)) == 1: + problems.append(f"{key} is constant across all {poses.shape[0]} frames") + rotations = poses[:, 3:7] + identity = np.zeros(len(rotations), dtype=bool) + for quaternion in _IDENTITY_QUATERNIONS: + identity |= np.all(rotations == quaternion, axis=1) + fraction = float(identity.mean()) + if fraction > threshold: + problems.append( + f"{key} holds an exact identity rotation on " + f"{fraction:.0%} of frames (limit {threshold:.0%})" + ) + _report_problems( + rule, report, "pose_degeneracy", problems, f"{checked} pose track(s) move" + ) + + +def _check_calibration_degeneracy(rule, context, report) -> None: + calibration = context.get("calibration") + if calibration is None: + return + problems = [] + for name, camera in calibration.cameras.items(): + # The reference camera is the identity by definition, so only a stored + # pose can be degenerate. + if ( + camera.ref_T_cam is not None + and name != calibration.reference_camera + and np.array_equal(camera.ref_T_cam, _IDENTITY) + ): + problems.append(f"cameras[{name!r}].ref_T_cam is exactly the identity") + for side in calibration.arm_bases: + base_T_cam = calibration.base_T_cam(side) + if base_T_cam is not None and np.allclose(base_T_cam, _IDENTITY, atol=0.0): + problems.append( + f"the camera sits exactly at the {side} arm base " + "(base_T_cam is the identity)" + ) + _report_problems( + rule, report, "calibration_degeneracy", problems, "no identity extrinsic" + ) + + +def _resolution(camera, context) -> tuple[int, int] | None: + """Return one camera's ``(width, height)``, from the block or the stream.""" + if camera.resolution is not None: + return camera.resolution + feature = context["features"].get(f"{IMAGE_KEY_PREFIX}{camera.name}") or {} + shape = feature.get("shape") or [] + if len(shape) >= 2: + return int(shape[1]), int(shape[0]) + return None + + +def _check_intrinsics_signature(rule, context, report) -> None: + calibration = context.get("calibration") + if calibration is None: + return + problems = [] + checked = 0 + for name, camera in calibration.cameras.items(): + resolution = _resolution(camera, context) + if camera.K is None or resolution is None: + continue + checked += 1 + width, height = resolution + fx, fy = camera.K[0, 0], camera.K[1, 1] + cx, cy = camera.K[0, 2], camera.K[1, 2] + if fx == fy == width and cx == width / 2 and cy == height / 2: + problems.append( + f"cameras[{name!r}].K is a synthesized centred pinhole " + f"(fx = fy = {fx:g} = W, principal point at the image centre)" + ) + _report_problems( + rule, + report, + "intrinsics_signature", + problems, + f"{checked} camera matri(ces) carry a measured focal length", + ) + + +def _check_timestamps(rule, context, report) -> None: + problems = [] + for substring in rule.get("banned_key_substrings") or (): + stored = [k for k in context["arrays"] if substring in k] + if stored: + problems.append( + f"{stored} store a second time base; keep one clock per episode " + "and derive a relative base on read" + ) + + key = rule.get("key", "obs_rgb_timestamps_ns") + array = context["arrays"].get(key) + if array is None: + _report_problems(rule, report, "timestamps", problems, "no clock stored") + return + + stamps = _read(array, context.get("total_frames")) + if stamps.ndim != 1 or stamps.shape[0] < 2: + _report_problems(rule, report, "timestamps", problems, "one stamp") + return + + steps = np.diff(stamps.astype(np.int64)) + stalled = int(np.count_nonzero(steps <= 0)) + if stalled: + problems.append(f"{key} does not increase on {stalled} of {len(steps)} steps") + + quantum = int(np.gcd.reduce(np.abs(steps))) if steps.size else 0 + if quantum >= 64 and quantum & (quantum - 1) == 0: + problems.append( + f"{key} is quantized to {quantum} ns, the signature of float64 " + "seconds converted to nanoseconds; compute the clock in integers" + ) + + _report_problems( + rule, + report, + "timestamps", + problems, + f"{len(stamps)} stamps increase strictly", + ) + + +def _annotations(context, key: str) -> list[dict]: + array = context["arrays"].get(key) + if array is None: + return [] + out = [] + for entry in np.asarray(array[:]): + if isinstance(entry, (bytes, bytearray, memoryview)): + entry = bytes(entry).decode("utf-8") + if isinstance(entry, str): + try: + entry = json.loads(entry) + except json.JSONDecodeError: + continue + if isinstance(entry, Mapping): + out.append(dict(entry)) + return out + + +def _check_annotation_coverage(rule, context, report) -> None: + key = rule.get("key", "annotations") + total_frames = context.get("total_frames") or 0 + annotations = _annotations(context, key) + problems = [] + covered = np.zeros(max(total_frames, 0), dtype=bool) + for annotation in annotations: + start = int(annotation.get("start_idx", -1)) + end = int(annotation.get("end_idx", -1)) + if start < 0 or end > total_frames or end <= start: + problems.append( + f"span [{start}, {end}) is outside [0, {total_frames}) or empty" + ) + continue + covered[start:end] = True + fraction = float(covered.mean()) if covered.size else 0.0 + minimum = float(rule.get("minimum", 0.9)) + if fraction < minimum: + problems.append( + f"{len(annotations)} annotation(s) cover {fraction:.0%} of the " + f"episode (minimum {minimum:.0%}); trim the tail or annotate it" + ) + _report_problems( + rule, + report, + "annotation_coverage", + problems, + f"{len(annotations)} annotation(s) cover {fraction:.0%}", + ) + + +def _check_annotation_text(rule, context, report) -> None: + delimiters = tuple(rule.get("banned_delimiters") or ()) + texts = [ + (f"annotation {i}", a.get("text", "")) + for i, a in enumerate(_annotations(context, rule.get("key", "annotations"))) + ] + texts.append(("task_description", context["attrs"].get("task_description") or "")) + problems = [] + for where, text in texts: + for delimiter in delimiters: + if delimiter in str(text): + problems.append( + f"{where} encodes metadata after {delimiter!r}: {text!r}" + ) + break + _report_problems( + rule, report, "annotation_text", problems, "no delimiter-encoded metadata" + ) + + _NAMED_CHECKS = { "calibration_present": _check_calibration_present, "camera_coverage": _check_camera_coverage, + "pose_degeneracy": _check_pose_degeneracy, + "calibration_degeneracy": _check_calibration_degeneracy, + "intrinsics_signature": _check_intrinsics_signature, + "timestamps": _check_timestamps, + "annotation_coverage": _check_annotation_coverage, + "annotation_text": _check_annotation_text, } @@ -467,7 +706,13 @@ def validate_episode(path: str | Path, *, strict: bool = False) -> Report: attrs = dict(store.attrs) arrays = {name: store[name] for name in store.array_keys()} - context: dict[str, Any] = {"array_keys": list(arrays)} + context: dict[str, Any] = { + "array_keys": list(arrays), + "arrays": arrays, + "attrs": attrs, + "features": attrs.get("features") or {}, + "total_frames": attrs.get("total_frames"), + } embodiment = attrs.get("embodiment") if embodiment: context["named_platform"] = load_embodiment_platforms().get( @@ -494,7 +739,6 @@ def validate_episode(path: str | Path, *, strict: bool = False) -> Report: # report explains itself without a second lookup. report.add(OK, "embodiment", resolved.describe()) context["resolved"] = resolved - context["total_frames"] = attrs.get("total_frames") for rule in schema.get("arrays", []): for key, side in _expand_key(rule["key"], arrays, resolved.sides): From 5fc82f368f7f0e0cc0fd6d90c677b40441ee89dd Mon Sep 17 00:00:00 2001 From: jaynye Date: Wed, 2 Sep 2026 21:04:25 -0400 Subject: [PATCH 7/9] feat(zarr): data_status, and refuse a structural sample everywhere A sample sent to show the shape of a delivery is real data in the schema sense and not real data in the training sense, and until now nothing downstream could tell the two apart. Sharpa reached for this with six free-text `*_status` attributes that nothing reads. `data_status` is `complete` or `structural_sample`. The writer records it and rejects anything else. A non-complete episode gets no staging row and is skipped by the resolver, so it cannot reach a training run by any path. An episode written before the attribute existed reads as `complete`: the corpus predates the distinction and every episode in it was delivered as finished data. The validator therefore asks for the attribute only under `--strict`, while an unknown value is always an error. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_012rCh1QmnyqgPZ5TJpsfMQb --- egomimic/rldb/zarr/episode_attrs.py | 48 ++++++++ egomimic/rldb/zarr/schema/episode_v3.yaml | 6 + egomimic/rldb/zarr/test_data_status.py | 112 ++++++++++++++++++ egomimic/rldb/zarr/validate.py | 3 +- egomimic/rldb/zarr/zarr_dataset_multi.py | 22 ++++ egomimic/rldb/zarr/zarr_writer.py | 22 ++++ .../stage_processed_folder.py | 9 +- 7 files changed, 220 insertions(+), 2 deletions(-) create mode 100644 egomimic/rldb/zarr/episode_attrs.py create mode 100644 egomimic/rldb/zarr/test_data_status.py diff --git a/egomimic/rldb/zarr/episode_attrs.py b/egomimic/rldb/zarr/episode_attrs.py new file mode 100644 index 000000000..ee0857733 --- /dev/null +++ b/egomimic/rldb/zarr/episode_attrs.py @@ -0,0 +1,48 @@ +"""Episode attribute vocabulary shared by the writer, the reader and the validator. + +Keeping the vocabulary here rather than in any one of them means the writer +does not import the reader, and nobody restates a literal. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +#: A finished recording. Only these episodes get a database row and reach a +#: dataset. +DATA_STATUS_COMPLETE = "complete" + +#: A sample sent to show the shape of a delivery. It is real data in the schema +#: sense and not real data in the training sense, and nothing downstream could +#: previously tell the two apart. +DATA_STATUS_STRUCTURAL_SAMPLE = "structural_sample" + +DATA_STATUS_VALUES = (DATA_STATUS_COMPLETE, DATA_STATUS_STRUCTURAL_SAMPLE) + + +def data_status(attrs: Mapping) -> str: + """Return one episode's data status. + + An episode written before the attribute existed is ``complete``: the + corpus predates the distinction and every episode in it was delivered as + finished data. + + Args: + attrs: The episode's Zarr attributes. + """ + value = attrs.get("data_status") + return DATA_STATUS_COMPLETE if not value else str(value) + + +def is_complete(attrs: Mapping) -> bool: + """Return whether one episode is a finished recording.""" + return data_status(attrs) == DATA_STATUS_COMPLETE + + +__all__ = [ + "DATA_STATUS_COMPLETE", + "DATA_STATUS_STRUCTURAL_SAMPLE", + "DATA_STATUS_VALUES", + "data_status", + "is_complete", +] diff --git a/egomimic/rldb/zarr/schema/episode_v3.yaml b/egomimic/rldb/zarr/schema/episode_v3.yaml index 484284e78..8152718a1 100644 --- a/egomimic/rldb/zarr/schema/episode_v3.yaml +++ b/egomimic/rldb/zarr/schema/episode_v3.yaml @@ -61,6 +61,12 @@ attributes: required: true type: int min: 1 + # `complete` or `structural_sample`. An episode written before this + # attribute existed reads as complete, so it is required only under --strict. + data_status: + required: strict + type: str + choices: [complete, structural_sample] task_name: required: true type: str diff --git a/egomimic/rldb/zarr/test_data_status.py b/egomimic/rldb/zarr/test_data_status.py new file mode 100644 index 000000000..885932fc5 --- /dev/null +++ b/egomimic/rldb/zarr/test_data_status.py @@ -0,0 +1,112 @@ +"""Test that a structural sample is refused everywhere a complete episode lands.""" + +import numpy as np +import pytest +import zarr + +from egomimic.rldb.zarr.episode_attrs import ( + DATA_STATUS_COMPLETE, + DATA_STATUS_STRUCTURAL_SAMPLE, + data_status, + is_complete, +) +from egomimic.rldb.zarr.validate import ERROR, WARNING, validate_episode +from egomimic.rldb.zarr.zarr_dataset_multi import ( + EpisodeResolver, + ZarrDataset, + ZarrEpisode, +) +from egomimic.rldb.zarr.zarr_writer import ZarrWriter + +K = np.array([[200.0, 0.0, 160.0, 0.0], [0.0, 200.0, 120.0, 0.0], [0.0, 0.0, 1.0, 0.0]]) +LENGTH = 4 + + +def _write(path, **kwargs) -> None: + ZarrWriter.create_and_write( + episode_path=path, + numeric_data={"left.obs_gripper": np.zeros((LENGTH, 1))}, + embodiment="eva_bimanual", + chunk_timesteps=LENGTH, + intrinsics={"front_1": K}, + **kwargs, + ) + + +def test_an_episode_without_the_attribute_reads_as_complete() -> None: + assert data_status({}) == DATA_STATUS_COMPLETE + assert is_complete({}) + assert not is_complete({"data_status": DATA_STATUS_STRUCTURAL_SAMPLE}) + + +def test_the_writer_records_the_status(tmp_path) -> None: + _write(tmp_path / "complete.zarr") + _write(tmp_path / "sample.zarr", data_status=DATA_STATUS_STRUCTURAL_SAMPLE) + + assert ZarrEpisode(tmp_path / "complete.zarr").data_status == DATA_STATUS_COMPLETE + assert ( + ZarrEpisode(tmp_path / "sample.zarr").data_status + == DATA_STATUS_STRUCTURAL_SAMPLE + ) + + +def test_the_writer_rejects_an_unknown_status(tmp_path) -> None: + with pytest.raises(ValueError, match="data_status must be one of"): + _write(tmp_path / "bad.zarr", data_status="mostly_done") + + +def test_the_resolver_refuses_a_structural_sample(tmp_path, caplog) -> None: + _write(tmp_path / "complete.zarr") + _write(tmp_path / "sample.zarr", data_status=DATA_STATUS_STRUCTURAL_SAMPLE) + resolver = EpisodeResolver(tmp_path, key_map={}) + + with caplog.at_level("WARNING"): + datasets = resolver._load_zarr_datasets( + search_path=tmp_path, valid_folder_names={"complete", "sample"} + ) + + assert set(datasets) == {"complete"} + assert "not 'complete'" in caplog.text + + +def test_the_dataset_exposes_the_status(tmp_path) -> None: + _write(tmp_path / "sample.zarr", data_status=DATA_STATUS_STRUCTURAL_SAMPLE) + + dataset = ZarrDataset(Episode_path=tmp_path / "sample.zarr", key_map={}) + + assert dataset.data_status == DATA_STATUS_STRUCTURAL_SAMPLE + + +def test_the_validator_reads_and_checks_the_status(tmp_path) -> None: + _write(tmp_path / "sample.zarr", data_status=DATA_STATUS_STRUCTURAL_SAMPLE) + report = validate_episode(tmp_path / "sample.zarr") + assert next( + f.level for f in report.findings if f.check == "attrs.data_status" + ) == "ok" + + _write(tmp_path / "legacy.zarr") + store = zarr.open_group(str(tmp_path / "legacy.zarr"), mode="a") + del store.attrs["data_status"] + lenient = validate_episode(tmp_path / "legacy.zarr") + strict = validate_episode(tmp_path / "legacy.zarr", strict=True) + statuses = { + "lenient": next( + f.level for f in lenient.findings if f.check == "attrs.data_status" + ), + "strict": next( + f.level for f in strict.findings if f.check == "attrs.data_status" + ), + } + assert statuses == {"lenient": WARNING, "strict": ERROR} + + +def test_the_validator_rejects_an_unknown_status(tmp_path) -> None: + _write(tmp_path / "odd.zarr") + store = zarr.open_group(str(tmp_path / "odd.zarr"), mode="a") + store.attrs["data_status"] = "in_review" + + report = validate_episode(tmp_path / "odd.zarr") + + finding = next(f for f in report.findings if f.check == "attrs.data_status") + assert finding.level == ERROR + assert "expected one of" in finding.message diff --git a/egomimic/rldb/zarr/validate.py b/egomimic/rldb/zarr/validate.py index cdae2c046..eab84b078 100644 --- a/egomimic/rldb/zarr/validate.py +++ b/egomimic/rldb/zarr/validate.py @@ -1,4 +1,4 @@ -"""Validate one zarr episode against the schema in ``schema/episode_v3.yaml``. +"""Validate one zarr episode against the rules in schema/episode_v3.yaml. Run it with:: @@ -686,6 +686,7 @@ def _expand_key(template: str, arrays: Mapping, sides) -> list[tuple[str, str | def validate_episode(path: str | Path, *, strict: bool = False) -> Report: """Validate one zarr episode against ``schema/episode_v3.yaml``. + Args: path: The episode ``.zarr`` directory. strict: If true, promote the rules the corpus does not meet yet from diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index 45f3e5604..5d56b3eff 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -45,6 +45,7 @@ from egomimic.rldb.filters import DatasetFilter from egomimic.rldb.zarr.action_chunk_transforms import base_T_cam_pose_key from egomimic.rldb.zarr.calibration import Calibration, read_calibration +from egomimic.rldb.zarr.episode_attrs import data_status, is_complete from egomimic.utils.aws.aws_data_utils import load_env from egomimic.utils.aws.aws_sql import ( create_default_engine, @@ -234,6 +235,17 @@ def _load_zarr_datasets(self, search_path: Path, valid_folder_names: set[str]): key_map=self.key_map, transform_list=self.transform_list, ) + # A structural sample is real data in the schema sense and not + # real data in the training sense. Refuse it here so it cannot + # reach a training run through any resolver. + if not is_complete(ds_obj.metadata): + logger.warning( + "Skipping %s: data_status is %r, not 'complete'", + p, + ds_obj.data_status, + ) + skipped.append(p.name) + continue datasets[name] = ds_obj except Exception as e: logger.error(f"Failed to load dataset at {p}: {e}") @@ -1563,6 +1575,11 @@ def init_episode(self): self._json_keys = self._detect_json_keys() self._extrinsic_poses = self._build_extrinsic_poses() + @property + def data_status(self) -> str: + """Pass-through to ZarrEpisode.data_status (read from zarr metadata).""" + return self.episode_reader.data_status + @property def calibration(self) -> Calibration | None: """Pass-through to ZarrEpisode.calibration (read from zarr metadata).""" @@ -1871,6 +1888,11 @@ def __init__(self, path: str | Path): self.keys = self.metadata["features"] self._calibration = read_calibration(self.metadata) + @property + def data_status(self) -> str: + """Return ``complete`` or ``structural_sample`` for this episode.""" + return data_status(self.metadata) + @property def calibration(self) -> Calibration | None: """Per-episode calibration, read through the one shim. diff --git a/egomimic/rldb/zarr/zarr_writer.py b/egomimic/rldb/zarr/zarr_writer.py index 8985429d4..6537bcdfa 100644 --- a/egomimic/rldb/zarr/zarr_writer.py +++ b/egomimic/rldb/zarr/zarr_writer.py @@ -21,6 +21,10 @@ parse_calibration, uncalibrated_cameras, ) +from egomimic.rldb.zarr.episode_attrs import ( + DATA_STATUS_COMPLETE, + DATA_STATUS_VALUES, +) logger = logging.getLogger(__name__) @@ -322,6 +326,7 @@ def __init__( intrinsics: dict | None = None, extrinsics: dict | None = None, calibration: Calibration | dict | None = None, + data_status: str = DATA_STATUS_COMPLETE, strict: bool = False, verbose: bool = False, ): @@ -343,6 +348,10 @@ def __init__( calibration: ``None`` or the per-episode calibration to store under the ``calibration`` attribute. A ``Calibration`` or the mapping form of one. + data_status: ``complete`` for a finished recording, or + ``structural_sample`` for one sent to show the shape of a + delivery. Only a complete episode gets a database row and + reaches a dataset. strict: If true, an image stream that no camera matrix covers is an error instead of a warning. verbose: If true, print progress information during writes. @@ -361,6 +370,12 @@ def __init__( self.calibration = ( None if calibration is None else parse_calibration(calibration) ) + if data_status not in DATA_STATUS_VALUES: + raise ValueError( + f"data_status must be one of {list(DATA_STATUS_VALUES)}, got " + f"{data_status!r}" + ) + self.data_status = data_status self.strict = strict self.verbose = verbose # Track image shapes for metadata @@ -760,6 +775,7 @@ def _build_metadata( """ metadata = { "embodiment": self.embodiment, + "data_status": self.data_status, "total_frames": self.total_frames, "fps": self.fps, "task_name": self.task_name, @@ -807,6 +823,7 @@ def create_and_write( intrinsics: dict | None = None, extrinsics: dict | None = None, calibration: Calibration | dict | None = None, + data_status: str = DATA_STATUS_COMPLETE, strict: bool = False, metadata_override: dict[str, Any] | None = None, ) -> Path: @@ -836,6 +853,10 @@ def create_and_write( writer derives ``intrinsics`` and ``extrinsics`` from it for any of the two that the caller omits, so readers that predate the block keep working. + data_status: ``complete`` for a finished recording, or + ``structural_sample`` for one sent to show the shape of a + delivery. Only a complete episode gets a database row and + reaches a dataset. strict: If true, an image stream that no camera matrix covers is an error instead of a warning. It stays opt-in because every EVA episode in the corpus stores three streams and calibrates one. @@ -926,6 +947,7 @@ def create_and_write( intrinsics=intrinsics, extrinsics=extrinsics, calibration=calibration, + data_status=data_status, strict=strict, ) diff --git a/egomimic/scripts/backfill_scripts/stage_processed_folder.py b/egomimic/scripts/backfill_scripts/stage_processed_folder.py index 67359f8e3..04eec02fe 100644 --- a/egomimic/scripts/backfill_scripts/stage_processed_folder.py +++ b/egomimic/scripts/backfill_scripts/stage_processed_folder.py @@ -87,7 +87,9 @@ def find_zarr_prefixes(s3, base): def read_batch(prefixes, folder, cfg): """Read the zarr.json for a batch of prefixes -> list[row dict]. FULLY self-contained (only stdlib + boto3, imported inside) so it runs on any Ray - worker regardless of whether egomimic / ~/.egoverse_env are present there.""" + worker regardless of whether egomimic / ~/.egoverse_env are present there. + + An episode whose ``data_status`` is not ``complete`` yields no row.""" import json as _json import re as _re from datetime import datetime as _dt @@ -121,6 +123,11 @@ def read_batch(prefixes, folder, cfg): ).replace(tzinfo=_tz.utc).isoformat() except ValueError: created_at = None + # A structural sample never gets a row. Mirrors + # `egomimic.rldb.zarr.episode_attrs.DATA_STATUS_COMPLETE`, restated as + # a literal because this function must run on a bare Ray worker. + if (attrs.get("data_status") or "complete") != "complete": + continue feats = attrs.get("features") out.append({ "episode_hash": episode_hash, From 1c49f2652917e12c8faec2dc87f8828a5001b8da Mon Sep 17 00:00:00 2001 From: jaynye Date: Thu, 3 Sep 2026 17:19:28 -0400 Subject: [PATCH 8/9] tightened documentation --- egomimic/rldb/embodiment/embodiment.py | 13 +- egomimic/rldb/embodiment/eva.py | 20 +-- egomimic/rldb/zarr/action_chunk_transforms.py | 16 +-- egomimic/rldb/zarr/calibration.py | 121 ++++++++---------- egomimic/rldb/zarr/episode_attrs.py | 22 ++-- egomimic/rldb/zarr/schema/episode_v3.yaml | 75 +++++------ egomimic/rldb/zarr/test_calibration.py | 6 +- egomimic/rldb/zarr/test_data_status.py | 2 +- egomimic/rldb/zarr/test_validate.py | 9 +- egomimic/rldb/zarr/validate.py | 81 ++++++------ egomimic/rldb/zarr/zarr_dataset_multi.py | 42 +++--- egomimic/rldb/zarr/zarr_writer.py | 66 ++++------ .../stage_processed_folder.py | 15 ++- .../inspector_lib/dataset_view.py | 20 +-- 14 files changed, 234 insertions(+), 274 deletions(-) diff --git a/egomimic/rldb/embodiment/embodiment.py b/egomimic/rldb/embodiment/embodiment.py index 2f730a25a..b4c0ec3bd 100644 --- a/egomimic/rldb/embodiment/embodiment.py +++ b/egomimic/rldb/embodiment/embodiment.py @@ -137,11 +137,11 @@ def action_space(self) -> str: @property def arity(self) -> str | None: - """Return the arm configuration this embodiment name selects. + """Return the arity suffix encoded in ``embodiment_name``. Returns: - One of the platform's ``arity`` values, or ``None`` when - resolution started from a morphology block with no name. + The text after ``_``, or ``None`` when + the resolved value has no matching embodiment name. """ if self.embodiment_name is None: return None @@ -152,10 +152,11 @@ def arity(self) -> str | None: @property def sides(self) -> tuple[str, ...]: - """Return the arms this embodiment carries. + """Return side candidates implied by the encoded arity. - A single-arm arity carries one side. Anything else carries both, which - is what a morphology block with no name resolves to. + ``left_arm`` and ``right_arm`` select one side. All other values return + both candidates; callers resolving a morphology mapping must still + filter candidates not present in ``end_effectors``. """ arity = self.arity if arity == "left_arm": diff --git a/egomimic/rldb/embodiment/eva.py b/egomimic/rldb/embodiment/eva.py index 75a779831..0f9b07498 100644 --- a/egomimic/rldb/embodiment/eva.py +++ b/egomimic/rldb/embodiment/eva.py @@ -26,19 +26,15 @@ class Eva(Embodiment): - """EVA X5 bimanual platform with a parallel jaw on each arm. + """Dataset transforms and visualization for the EVA X5 platform. - ``INTRINSICS`` and ``EXTRINSICS`` are fallbacks for an episode that - declares no calibration of its own. Calibration measures the rig that - recorded one episode, so the episode's own values win: ``ZarrDataset`` - puts them in every sample and the transform pipeline reads them from - there. Two vendors on one platform have two rigs, and a class constant - cannot tell them apart. + ``EXTRINSICS`` is used only when a sample has no episode-specific + ``base_T_cam`` poses. ``INTRINSICS`` is the visualization fallback when a + sample has no camera matrix. Both constants describe the original EVA rig. """ INTRINSICS = ARIA_INTRINSICS - #: Fallback rig: `base_T_cam` per arm, the camera pose in that arm's base - #: frame. See `docs/CONVENTIONS.md`. + #: Compatibility fallback: one ``base_T_cam`` pose per arm. EXTRINSICS = { "left": np.array( [ @@ -260,8 +256,7 @@ def _build_eva_bimanual_eef_frame_transform_list( ) -> list[Transform]: """EVA bimanual transform pipeline with actions expressed relative to the current EEF pose (wrist frame), analogous to keypoints relative to wrist pose.""" - # A fallback for an episode that declares no extrinsics. `ZarrDataset` - # puts the episode's own rig in the sample, and that value wins. + # Supply class-default poses only for sample keys absent at transform time. extrinsics = Eva.EXTRINSICS left_base_T_cam_pose = _matrix_to_xyzwxyz(extrinsics["left"][None, :])[0] right_base_T_cam_pose = _matrix_to_xyzwxyz(extrinsics["right"][None, :])[0] @@ -425,8 +420,7 @@ def _build_eva_bimanual_transform_list( is_quat: bool = True, ) -> list[Transform]: """Canonical EVA bimanual transform pipeline used by tests and notebooks.""" - # A fallback for an episode that declares no extrinsics. `ZarrDataset` - # puts the episode's own rig in the sample, and that value wins. + # Supply class-default poses only for sample keys absent at transform time. extrinsics = Eva.EXTRINSICS left_base_T_cam_pose = _matrix_to_xyzwxyz(extrinsics["left"][None, :])[0] right_base_T_cam_pose = _matrix_to_xyzwxyz(extrinsics["right"][None, :])[0] diff --git a/egomimic/rldb/zarr/action_chunk_transforms.py b/egomimic/rldb/zarr/action_chunk_transforms.py index 7d5955f93..6c5292ea1 100644 --- a/egomimic/rldb/zarr/action_chunk_transforms.py +++ b/egomimic/rldb/zarr/action_chunk_transforms.py @@ -35,26 +35,14 @@ xyzw_to_wxyz, ) -# --------------------------------------------------------------------------- -# Per-episode calibration -# --------------------------------------------------------------------------- - -#: Batch key holding one arm's ``base_T_cam`` pose, the camera pose in that -#: arm's base frame. ``ZarrDataset`` fills it from the episode calibration; a -#: transform's ``extra_batch_key`` supplies an embodiment default when the -#: episode declares none. def base_T_cam_pose_key(side: str) -> str: - """Return the batch key that carries one arm's ``base_T_cam`` pose.""" + """Return the sample key for ``base_T_cam`` as ``[xyz, qw, qx, qy, qz]``.""" return f"{side}_base_T_cam_pose" def _apply_fallbacks(batch, extra_batch_key) -> None: - """Fill batch keys that the sample does not already carry. - - Calibration is per-episode, so a value the dataset read from the episode - must outrank the embodiment constant a transform was built with. - """ + """Add fallback values for absent keys without replacing sample values.""" for key, value in (extra_batch_key or {}).items(): batch.setdefault(key, value) diff --git a/egomimic/rldb/zarr/calibration.py b/egomimic/rldb/zarr/calibration.py index cad94e0de..cae947c14 100644 --- a/egomimic/rldb/zarr/calibration.py +++ b/egomimic/rldb/zarr/calibration.py @@ -1,4 +1,4 @@ -"""Read and write the per-episode ``calibration`` metadata block. +"""Parse and serialize per-episode camera and arm-base calibration. Calibration is a physical measurement of the rig that recorded one episode, so it travels with the episode instead of living in a class constant. An episode @@ -9,6 +9,8 @@ "cameras": { "front_1": { "K": [[fx, 0, cx, 0], [0, fy, cy, 0], [0, 0, 1, 0]], + "model": "PINHOLE", + "distortion": [], "resolution": [W, H], "rectified": true, "ref_T_cam": [[...]], @@ -17,10 +19,10 @@ "arm_bases": {"left": [[...]], "right": [[...]]}, } -Every matrix follows ``docs/CONVENTIONS.md``: ``A_T_B`` maps coordinates from -frame ``B`` to frame ``A``. ``ref_T_cam`` is the camera pose in the reference -frame, and ``arm_bases[side]`` is ``ref_T_armbase``, the arm base pose in the -reference frame. +Every rigid transform follows ``docs/CONVENTIONS.md``: ``A_T_B`` maps +coordinates from frame ``B`` to frame ``A``. ``ref_T_cam`` is the camera pose +in the reference frame, and ``arm_bases[side]`` is ``ref_T_armbase``, the arm +base pose in the reference frame. Episodes written before this block existed store ``intrinsics`` and ``extrinsics`` at the top level instead. :func:`read_calibration` reads either @@ -35,30 +37,26 @@ import numpy as np -#: Camera that legacy episodes calibrate and express `extrinsics` against. +#: Conventional camera name assigned to bare legacy intrinsics and extrinsics. LEGACY_REFERENCE_CAMERA = "front_1" -#: Reference frames that do not name a camera. +#: Supported reference frames that do not name a camera. STATIC_REFERENCE_FRAMES = frozenset({"robot_base", "slam_world"}) -#: Prefix that makes a camera the reference frame, as in `camera:front_1`. +#: Prefix for a camera reference frame, as in ``camera:front_1``. CAMERA_FRAME_PREFIX = "camera:" -#: Prefix of the array key that stores one camera stream, as in -#: `images.front_1`. The text after it is the camera name. +#: Prefix of an image-array key; the suffix is its camera name. IMAGE_KEY_PREFIX = "images." -#: Projection model of a camera, and the distortion coefficient counts it -#: accepts. Nothing projects a non-pinhole model yet. The declaration is -#: collected now because it measures a vendor's rig: if we discover later that -#: we need it, the rig has moved and the episodes cannot be recalibrated. +#: Supported projection models mapped to accepted distortion-vector lengths. CAMERA_MODELS = { "PINHOLE": frozenset({0}), "OPENCV": frozenset({4, 5, 8, 12, 14}), "KANNALA_BRANDT": frozenset({4}), } -#: Projection model assumed for a camera that declares none. +#: Projection model used when a camera omits ``model``. DEFAULT_CAMERA_MODEL = "PINHOLE" _CAMERA_FIELDS = frozenset( @@ -68,11 +66,11 @@ class CalibrationError(ValueError): - """Report an invalid ``calibration`` block or legacy calibration attribute.""" + """Raised when current or legacy calibration metadata is malformed.""" def _matrix(value: Any, shape: tuple[int, ...], where: str) -> np.ndarray: - """Return ``value`` as a finite float64 array of the given shape.""" + """Convert ``value`` to a finite float64 array with exactly ``shape``.""" try: arr = np.asarray(value, dtype=np.float64) except (TypeError, ValueError) as exc: @@ -87,7 +85,7 @@ def _matrix(value: Any, shape: tuple[int, ...], where: str) -> np.ndarray: def _camera_matrix(value: Any, where: str) -> np.ndarray: - """Return a 3×4 camera matrix, padding a bare 3×3 with a zero column.""" + """Return finite ``[K_3x3 | 0]``; append a zero column to a 3×3 input.""" try: arr = np.asarray(value, dtype=np.float64) except (TypeError, ValueError) as exc: @@ -106,16 +104,16 @@ def _camera_matrix(value: Any, where: str) -> np.ndarray: @dataclass(frozen=True) class CameraCalibration: - """Store the calibration of one camera stream. + """Normalized calibration metadata for one camera stream. Attributes: name: The camera name. It matches the ``images.`` array key. - K: The 3×4 camera matrix, or ``None`` for a declared but uncalibrated - camera. - model: A key in ``CAMERA_MODELS``. No projection site honors a - non-pinhole model yet. - distortion: The distortion coefficients of ``model``, in that model's - order. Empty for a pinhole camera. + K: The 3×4 ``[K_3x3 | 0]`` matrix, or ``None`` when undeclared. + model: A key in ``CAMERA_MODELS``. Current projection code does not + branch on this value. + distortion: The model-specific distortion vector. Its length is + validated against ``CAMERA_MODELS``; current projection code does + not consume it. resolution: ``(width, height)`` in pixels, or ``None``. rectified: Whether the stored frames are already rectified. ref_T_cam: The 4×4 camera pose in the episode reference frame, or @@ -131,7 +129,7 @@ class CameraCalibration: ref_T_cam: np.ndarray | None = None def to_jsonable(self) -> dict[str, Any]: - """Return this camera as JSON-serializable episode metadata.""" + """Serialize this camera to plain values accepted by Zarr attributes.""" out: dict[str, Any] = { "model": self.model, "rectified": bool(self.rectified), @@ -149,12 +147,12 @@ def to_jsonable(self) -> dict[str, Any]: @dataclass(frozen=True) class Calibration: - """Store the calibration of every camera and arm base in one episode. + """Normalized calibration metadata for one episode. Attributes: reference_frame: ``robot_base``, ``slam_world``, or - ``camera:``. Every pose in this block is expressed in - this frame. + ``camera:``. Every rigid pose in this block is + expressed in this frame. cameras: A mapping from camera name to its calibration. arm_bases: A mapping from ``"left"`` or ``"right"`` to a 4×4 ``ref_T_armbase`` matrix. Human episodes leave this empty. @@ -175,12 +173,11 @@ def reference_camera(self) -> str | None: return None def default_camera(self) -> str | None: - """Return the camera to project against when a caller names none. + """Select a camera when the caller does not name one. - The reference camera wins, then ``front_1``, then any camera whose name - contains ``front``, then the first declared camera. Projection and - visualization target the front image, so a front camera outranks a - wrist camera that happens to be declared first. + Selection order is the declared reference camera, ``front_1``, the + first name containing ``front`` (case-insensitive), and finally the + first declared camera. """ reference = self.reference_camera if reference is not None and reference in self.cameras: @@ -193,7 +190,7 @@ def default_camera(self) -> str | None: return next(iter(self.cameras), None) def K(self, camera: str | None = None) -> np.ndarray | None: - """Return the 3×4 camera matrix of one camera. + """Return one camera's normalized 3×4 ``[K_3x3 | 0]`` matrix. Args: camera: A camera name. ``None`` selects :meth:`default_camera`. @@ -209,10 +206,10 @@ def K(self, camera: str | None = None) -> np.ndarray | None: return None if entry is None else entry.K def ref_T_cam(self, camera: str | None = None) -> np.ndarray | None: - """Return the pose of one camera in the reference frame. + """Return one camera's pose in the episode reference frame. - The reference camera is the identity by definition, so an episode that - references a camera need not state that camera's pose. + A missing pose for the reference camera is synthesized as identity; + another missing pose returns ``None``. """ name = camera if camera is not None else self.default_camera() if name is None: @@ -225,7 +222,7 @@ def ref_T_cam(self, camera: str | None = None) -> np.ndarray | None: return None def base_T_cam(self, side: str, camera: str | None = None) -> np.ndarray | None: - """Return one camera's pose in the frame of one arm base. + """Compose one camera's pose in the selected arm-base frame. This is the quantity the EVA transform pipeline consumes, and it is what the ``extrinsics`` attribute of a legacy episode stores directly. @@ -235,8 +232,8 @@ def base_T_cam(self, side: str, camera: str | None = None) -> np.ndarray | None: camera: A camera name. ``None`` selects :meth:`default_camera`. Returns: - A 4×4 ``armbase_T_cam`` matrix, or ``None`` when the episode - states no arm base or no camera pose. + A 4×4 ``base_T_cam`` matrix, or ``None`` when the selected arm base + or camera pose is unavailable. """ ref_T_armbase = self.arm_bases.get(side) ref_T_cam = self.ref_T_cam(camera) @@ -245,13 +242,13 @@ def base_T_cam(self, side: str, camera: str | None = None) -> np.ndarray | None: return np.linalg.inv(ref_T_armbase) @ ref_T_cam def intrinsics(self) -> dict[str, np.ndarray]: - """Return ``{camera: K}`` for every camera that carries a ``K``.""" + """Return ``{camera_name: K}`` for cameras with intrinsic matrices.""" return { name: cam.K for name, cam in self.cameras.items() if cam.K is not None } def extrinsics(self, camera: str | None = None) -> dict[str, np.ndarray]: - """Return ``{side: base_T_cam}`` in the legacy attribute layout.""" + """Return composable ``{side: base_T_cam}`` values in legacy layout.""" out = {} for side in self.arm_bases: base_T_cam = self.base_T_cam(side, camera) @@ -260,7 +257,7 @@ def extrinsics(self, camera: str | None = None) -> dict[str, np.ndarray]: return out def to_jsonable(self) -> dict[str, Any]: - """Return this calibration as JSON-serializable episode metadata.""" + """Serialize this calibration to plain values accepted by Zarr attributes.""" out: dict[str, Any] = { "reference_frame": self.reference_frame, "cameras": { @@ -276,7 +273,7 @@ def to_jsonable(self) -> dict[str, Any]: def camera_name(image_key: str) -> str | None: - """Return the camera that one stored image array belongs to. + """Extract the camera name after the final ``images.`` in an array key. Args: image_key: An array key such as ``"images.front_1"``. @@ -292,11 +289,7 @@ def camera_name(image_key: str) -> str | None: def uncalibrated_cameras(image_keys, calibration: Calibration | None) -> list[str]: - """Return the image streams that no camera matrix covers. - - Coverage is the rule that makes per-episode calibration real: an episode - that stores three image streams and one ``K`` calibrates one camera in - three, and nothing downstream can tell. + """Return unique image-stream camera names without a declared ``K``. Args: image_keys: The episode's image array keys. @@ -338,7 +331,7 @@ def _parse_reference_frame(value: Any, cameras, where: str) -> str: def _parse_distortion(raw: Any, model: str, where: str) -> tuple[float, ...]: - """Validate the distortion coefficients declared for one camera model.""" + """Return a finite distortion vector whose length is valid for ``model``.""" if raw is None: raw = () if isinstance(raw, (str, bytes)) or not isinstance(raw, (list, tuple, np.ndarray)): @@ -418,7 +411,7 @@ def _parse_camera(name: str, block: Any, where: str) -> CameraCalibration: def parse_calibration(block: Any, where: str = "calibration") -> Calibration: - """Validate one ``calibration`` attribute block. + """Validate and normalize one ``calibration`` attribute block. Args: block: The mapping stored under ``zarr.attrs["calibration"]``. @@ -479,11 +472,11 @@ def lift_legacy_calibration( ) -> Calibration | None: """Build a :class:`Calibration` from the legacy attribute pair. - ``intrinsics`` maps a camera name to its ``K``. A bare matrix is read as - the ``front_1`` camera, which is the only camera any legacy writer - calibrated. ``extrinsics`` maps an arm side to ``base_T_cam``, the front - camera's pose in that arm's base frame, so the reference frame is the front - camera and ``arm_bases[side]`` is the inverse of the stored matrix. + ``intrinsics`` maps camera names to their ``K`` matrices. A bare matrix is + assigned to ``front_1``. ``extrinsics`` maps each arm side to + ``base_T_cam`` for ``front_1``; lifting therefore uses + ``camera:front_1`` as the reference frame and stores the inverse matrix as + ``arm_bases[side]``. Args: intrinsics: The ``intrinsics`` attribute, or ``None``. @@ -523,9 +516,8 @@ def lift_legacy_calibration( matrix = _matrix(base_T_cam, (4, 4), f"extrinsics[{side!r}]") arm_bases[str(side)] = np.linalg.inv(matrix) - # The legacy `extrinsics` values are per-arm poses of the front camera, so - # that camera is the reference frame. Declare it even when the episode - # calibrated no camera, otherwise `base_T_cam` has no pose to compose. + # Legacy extrinsics are poses of ``front_1`` in each arm base. Declare that + # camera even without intrinsics so ``base_T_cam`` remains composable. if arm_bases and LEGACY_REFERENCE_CAMERA not in cameras: cameras[LEGACY_REFERENCE_CAMERA] = CameraCalibration( name=LEGACY_REFERENCE_CAMERA @@ -545,12 +537,11 @@ def lift_legacy_calibration( def read_calibration(attrs: Mapping[str, Any]) -> Calibration | None: - """Return the calibration of one episode from its Zarr attributes. + """Parse current or legacy calibration from one episode's attributes. - This is the one reader every consumer goes through. It prefers the - ``calibration`` block and falls back to the legacy ``intrinsics`` and - ``extrinsics`` attributes, so episodes written before the block existed - stay readable without a rewrite. + A non-empty ``calibration`` block takes precedence. Otherwise the legacy + ``intrinsics`` and ``extrinsics`` attributes are lifted into the same + representation. Args: attrs: The episode's Zarr attributes. diff --git a/egomimic/rldb/zarr/episode_attrs.py b/egomimic/rldb/zarr/episode_attrs.py index ee0857733..5bcbd916f 100644 --- a/egomimic/rldb/zarr/episode_attrs.py +++ b/egomimic/rldb/zarr/episode_attrs.py @@ -1,31 +1,27 @@ -"""Episode attribute vocabulary shared by the writer, the reader and the validator. +"""Shared names and normalization helpers for episode attributes. -Keeping the vocabulary here rather than in any one of them means the writer -does not import the reader, and nobody restates a literal. +The writer and readers depend on this module without depending on each other. """ from __future__ import annotations from collections.abc import Mapping -#: A finished recording. Only these episodes get a database row and reach a -#: dataset. +#: A finished recording eligible for staging and resolver-based loading. DATA_STATUS_COMPLETE = "complete" -#: A sample sent to show the shape of a delivery. It is real data in the schema -#: sense and not real data in the training sense, and nothing downstream could -#: previously tell the two apart. +#: A delivery-format example that staging and dataset resolvers must exclude. DATA_STATUS_STRUCTURAL_SAMPLE = "structural_sample" DATA_STATUS_VALUES = (DATA_STATUS_COMPLETE, DATA_STATUS_STRUCTURAL_SAMPLE) def data_status(attrs: Mapping) -> str: - """Return one episode's data status. + """Return the stored status, defaulting a missing or falsey value to complete. - An episode written before the attribute existed is ``complete``: the - corpus predates the distinction and every episode in it was delivered as - finished data. + The default preserves the behavior of episodes written before + ``data_status`` was introduced. This helper does not validate a non-empty + stored value; the writer and schema validator do that separately. Args: attrs: The episode's Zarr attributes. @@ -35,7 +31,7 @@ def data_status(attrs: Mapping) -> str: def is_complete(attrs: Mapping) -> bool: - """Return whether one episode is a finished recording.""" + """Return whether the normalized status is exactly ``complete``.""" return data_status(attrs) == DATA_STATUS_COMPLETE diff --git a/egomimic/rldb/zarr/schema/episode_v3.yaml b/egomimic/rldb/zarr/schema/episode_v3.yaml index 8152718a1..e233f6623 100644 --- a/egomimic/rldb/zarr/schema/episode_v3.yaml +++ b/egomimic/rldb/zarr/schema/episode_v3.yaml @@ -1,7 +1,8 @@ -# The episode contract, as data. +# The declarative portion of the episode contract. # -# `egomimic/rldb/zarr/validate.py` reads this file and holds no rules of its -# own. A rule that lives here is one a contributor can read, diff and review. +# This file selects fields and named checks, supplies their parameters, and +# assigns failure severity. `validate.py` implements the schema interpreter and +# the named check predicates referenced here. # # Run it with: # python -m egomimic.rldb.zarr.validate [--strict] @@ -9,22 +10,27 @@ # --------------------------------------------------------------------------- # attributes # required true | false | strict -# `strict` means the attribute is an error only under --strict and -# a warning otherwise. Use it for a rule the corpus does not meet -# yet: flipping such a rule to `true` before the corpus is fixed -# turns a good check into an outage. +# For a missing attribute or array, `true` emits an error, `false` +# emits no finding, and `strict` emits a warning unless --strict +# promotes it to an error. Present attributes and arrays must meet +# their declared type and shape regardless of this setting. # type str | int | bool | mapping # min smallest accepted value for an int # choices the accepted values -# check a named check that needs more than a type; see `validate.py` +# check an attribute predicate implemented in `validate.py` +# +# checks +# name a predicate implemented in `validate.py` +# required controls the severity of a predicate failure as described above +# predicate-specific parameters # # arrays -# key an array key. `{side}` expands over the episode's arms, and `*` -# matches any suffix among the keys the episode stores. +# key an array key. `{side}` expands over arity-derived candidates, +# filtered by resolved end-effectors; `*` selects stored matches. # required true | false | strict, as above. A wildcard key is never # required. -# when conditions on the resolved embodiment; the rule is skipped when -# they do not hold +# when conditions on the resolved embodiment; skip the rule when any +# condition does not hold # shape expected dimensions, see below # dtype int | float | object | any # @@ -96,53 +102,50 @@ attributes: type: mapping checks: - # Every episode must state the rig that recorded it, through the current - # `calibration` block or the legacy pair. + # Require at least one camera matrix, read from the current `calibration` + # block or lifted from the legacy attributes. - name: calibration_present required: true - # One camera matrix per stored image stream. `eva_to_zarr.py` calibrates one - # camera of three, so every EVA episode in the corpus fails this today. + # Match each stored `images.` stream to a camera entry carrying K. + # This remains a warning without --strict during the compatibility rollout. - name: camera_coverage required: strict - # A pose track that never moves, or that holds an exact identity rotation on - # more than a slice of its frames, is a placeholder rather than a - # measurement. Nothing downstream can tell the two apart. + # For a listed 7D pose track with at least two retained rows, reject identical + # rows or more than 1% exact +identity/-identity quaternions. - name: pose_degeneracy required: true suffixes: [obs_ee_pose, cmd_ee_pose, obs_head_pose] identity_rotation_fraction: 0.01 - # An exactly-identity extrinsic places the camera at the arm base. It is the - # signature of a rig that was never measured, and it survives every other - # check because the projection it produces still looks plausible. + # Reject an exact identity pose stored for a non-reference camera, or a + # numerically identity `base_T_cam` composed for an arm and the default camera. - name: calibration_degeneracy required: true - # `fx == fy == W` with the principal point at the exact image centre is a - # synthesized camera, not a calibrated one. The `fx == W` conjunct is what - # keeps our own rectified Aria K (fx 266.5, W 640) out of the net. + # Reject the exact synthetic-intrinsics signature `fx == fy == width`, + # `cx == width / 2`, and `cy == height / 2`. - name: intrinsics_signature required: true - # One clock per episode, in integer UTC nanoseconds, strictly increasing. A - # relative time base is derived on read, never stored. + # Reject non-increasing retained timestamps, stored keys containing + # `relative_timestamp`, and a power-of-two step GCD of at least 64 ns (the + # heuristic signature of float seconds converted to integer nanoseconds). - name: timestamps required: true key: obs_rgb_timestamps_ns banned_key_substrings: [relative_timestamp] - # Annotations cover the retained frames or the tail is trimmed. EVA episodes - # pass no annotations at all today, so this is opt-in. + # Valid annotation intervals must jointly cover at least 90% of + # `[0, total_frames)`. Keep this transitional rule strict-only. - name: annotation_coverage required: strict key: annotations minimum: 0.9 - # A vendor who wants a skill taxonomy encodes one in the text, as in - # `pick up the cup | Skill: pick`. Ban the delimiter now and add a field when - # a second vendor asks for it. + # Forbid the exact substring ` | ` in annotation text and task_description; + # structured metadata needs a separate schema field. - name: annotation_text required: true key: annotations @@ -153,8 +156,8 @@ arrays: required: true shape: [T, 7] dtype: float - # The hand root or palm pose. It is the only key human, parallel-jaw and - # dexterous episodes share, so every embodiment owes it. + # Observed end-effector-root pose. Require one for every resolved side, + # independent of the end-effector class. - key: "{side}.cmd_ee_pose" required: false @@ -197,8 +200,8 @@ arrays: shape: [T, 7] dtype: float - # One clock per episode, in integer UTC nanoseconds. EVA episodes do not - # write it yet, so it is required only under --strict. + # Canonical RGB timestamp array. When present it must be one-dimensional, + # integer-valued, and at least `total_frames` long. Presence is strict-only. - key: obs_rgb_timestamps_ns required: strict shape: [T] diff --git a/egomimic/rldb/zarr/test_calibration.py b/egomimic/rldb/zarr/test_calibration.py index 8710514e2..051ac7220 100644 --- a/egomimic/rldb/zarr/test_calibration.py +++ b/egomimic/rldb/zarr/test_calibration.py @@ -371,7 +371,7 @@ def _eva_numeric_data(length: int) -> dict: def _other_rig() -> dict: - """Return a rig displaced from `Eva.EXTRINSICS` by a measurable amount.""" + """Translate each fallback camera pose by ``[0.05, -0.11, 0.07]`` metres.""" rig = {} for side, base_T_cam in Eva.EXTRINSICS.items(): moved = base_T_cam.copy() @@ -403,8 +403,8 @@ def _eva_actions(episode_path, extrinsics) -> np.ndarray: def test_reading_the_episode_is_a_no_op_on_the_current_corpus(tmp_path) -> None: stored = _eva_actions(tmp_path / "stored.zarr", Eva.EXTRINSICS) absent = _eva_actions(tmp_path / "absent.zarr", None) - # Every EVA episode in the corpus stores exactly the class constant, so - # reading the episode instead of the constant changes no number today. + # Storing the class-default rig and omitting it must exercise the same + # transform because omission selects that fallback. np.testing.assert_allclose(stored, absent, atol=1e-9) diff --git a/egomimic/rldb/zarr/test_data_status.py b/egomimic/rldb/zarr/test_data_status.py index 885932fc5..65607c667 100644 --- a/egomimic/rldb/zarr/test_data_status.py +++ b/egomimic/rldb/zarr/test_data_status.py @@ -1,4 +1,4 @@ -"""Test that a structural sample is refused everywhere a complete episode lands.""" +"""Test status writing, legacy defaults, validation, and resolver filtering.""" import numpy as np import pytest diff --git a/egomimic/rldb/zarr/test_validate.py b/egomimic/rldb/zarr/test_validate.py index 829991985..2b43e38b8 100644 --- a/egomimic/rldb/zarr/test_validate.py +++ b/egomimic/rldb/zarr/test_validate.py @@ -23,7 +23,7 @@ def _levels(report) -> dict[str, str]: def _poses(x_offset: float = 0.0) -> np.ndarray: - """Return a moving pose track: xyz plus a rotating [qw, qx, qy, qz].""" + """Return a moving ``[x, y, z, qw, qx, qy, qz]`` pose track.""" angles = np.linspace(0.1, 0.4, LENGTH) poses = np.zeros((LENGTH, 7)) poses[:, 0] = x_offset + np.arange(LENGTH) * 0.01 @@ -192,9 +192,10 @@ def test_a_human_episode_owes_keypoints_and_a_head_pose(tmp_path) -> None: levels = _levels(report) assert levels["left.obs_keypoints"] == OK assert levels["obs_head_pose"] == OK - # 21 MANO slots at three coordinates each. + # The registry declares 21 MANO slots with three coordinates per slot. assert levels["right.obs_keypoints"] == ERROR - # A human platform has no arm chain and no jaw, so neither rule runs. + # The human platform has neither an arm chain nor a parallel jaw, so those + # conditional rules do not run. assert "left.obs_gripper" not in levels assert "left.obs_joints" not in levels @@ -359,7 +360,7 @@ def test_a_synthesized_camera_matrix_is_an_error(tmp_path) -> None: def test_a_rectified_aria_camera_matrix_stays_out_of_the_net(tmp_path) -> None: - # fx is 266.5 at W 640, so the `fx == W` conjunct is what saves it. + # ``fx`` is 266.5 at width 640, so the full synthetic signature is false. aria = np.array( [[266.5, 0.0, 320.0, 0.0], [0.0, 266.5, 240.0, 0.0], [0.0, 0.0, 1.0, 0.0]] ) diff --git a/egomimic/rldb/zarr/validate.py b/egomimic/rldb/zarr/validate.py index eab84b078..92fce03fe 100644 --- a/egomimic/rldb/zarr/validate.py +++ b/egomimic/rldb/zarr/validate.py @@ -1,14 +1,14 @@ -"""Validate one zarr episode against the rules in schema/episode_v3.yaml. +"""Validate a Zarr episode against ``schema/episode_v3.yaml``. Run it with:: python -m egomimic.rldb.zarr.validate [--strict] -The rules live in the schema file, not here. This module reads them, resolves -the episode's embodiment through the registry, and reports one finding per -rule. ``--strict`` promotes the rules the corpus does not meet yet from -warnings to errors, so a rule can land, be measured across the corpus, and -only then become the default. +The schema declares attributes, arrays, conditions, thresholds, and severity. +This module interprets those declarations and implements the named predicates. +``--strict`` promotes failures marked ``required: strict`` from warnings to +errors; validation of a present attribute or array is always an error when its +declared type or shape is wrong. """ from __future__ import annotations @@ -46,8 +46,7 @@ WARNING = "warning" OK = "ok" -#: `required: strict` rules are warnings until `--strict` promotes them. They -#: exist so a rule can land before the corpus meets it. +#: Accepted values for a schema rule's ``required`` field. _REQUIRED_VALUES = (True, False, "strict") _TYPE_NAMES = { @@ -59,16 +58,16 @@ class SchemaError(ValueError): - """Report an invalid rule in ``episode_v3.yaml``.""" + """Raised when ``episode_v3.yaml`` contains an unsupported declaration.""" @dataclass(frozen=True) class Finding: - """One rule's result. + """One validation result emitted for an episode. Attributes: level: ``ok``, ``warning``, or ``error``. - check: The rule that produced this finding. + check: The attribute, array key, or named predicate being reported. message: What the rule found, in one line. """ @@ -82,12 +81,14 @@ def __str__(self) -> str: @dataclass class Report: - """Every finding for one episode. + """Validation findings for one episode. Attributes: path: The episode directory. - findings: One finding per rule that ran, in schema order. - strict: Whether the strict rules were promoted to errors. + findings: Findings in validation order. A schema rule may emit zero, + one, or multiple findings after key expansion. + strict: Whether ``required: strict`` failures are errors instead of + warnings. """ path: Path @@ -116,7 +117,7 @@ def summary(self) -> str: ) def text(self, verbose: bool = False) -> str: - """Render the report. + """Format this report for terminal output. Args: verbose: If true, list the rules that passed as well. @@ -142,13 +143,13 @@ def to_jsonable(self) -> dict[str, Any]: @functools.lru_cache(maxsize=1) def load_schema() -> dict: - """Load and check ``schema/episode_v3.yaml``. + """Load the schema and validate every declared ``required`` value. Returns: The parsed schema. Raises: - SchemaError: If a rule declares an unusable ``required`` value. + SchemaError: If a rule declares an unsupported ``required`` value. """ with SCHEMA_FILE.open("r") as f: schema = yaml.safe_load(f) or {} @@ -166,7 +167,7 @@ def load_schema() -> dict: def _level(required, strict: bool) -> str | None: - """Return the level a failed rule reports at, or ``None`` if it is optional.""" + """Return a requirement failure's severity, or ``None`` when optional.""" if required is True: return ERROR if required == "strict": @@ -317,13 +318,13 @@ def _check_camera_coverage(rule, context, report) -> None: def _read(array, total_frames: int | None) -> np.ndarray: - """Read an array up to `total_frames`, which is the authoritative length.""" + """Read the available prefix, capped at the authoritative frame count.""" end = array.shape[0] if total_frames is None else min(total_frames, array.shape[0]) return np.asarray(array[:end]) def _report_problems(rule, report, check: str, problems, passed: str) -> None: - """Record one finding for a named check.""" + """Record a named check's pass or its requirement-level failure.""" if not problems: report.add(OK, check, passed) return @@ -368,8 +369,8 @@ def _check_calibration_degeneracy(rule, context, report) -> None: return problems = [] for name, camera in calibration.cameras.items(): - # The reference camera is the identity by definition, so only a stored - # pose can be degenerate. + # Identity is implicit for the reference camera; only an explicitly + # stored pose for another camera is subject to this branch. if ( camera.ref_T_cam is not None and name != calibration.reference_camera @@ -389,7 +390,7 @@ def _check_calibration_degeneracy(rule, context, report) -> None: def _resolution(camera, context) -> tuple[int, int] | None: - """Return one camera's ``(width, height)``, from the block or the stream.""" + """Return ``(width, height)`` from calibration or image feature metadata.""" if camera.resolution is not None: return camera.resolution feature = context["features"].get(f"{IMAGE_KEY_PREFIX}{camera.name}") or {} @@ -574,7 +575,7 @@ def _condition_holds(when: Mapping, resolved: ResolvedEmbodiment, side) -> bool: def _dimension(token, context, side) -> int | None: - """Resolve one schema shape token to a length, or ``None`` for unknown.""" + """Resolve a shape token; return ``None`` for a wildcard or absent spec.""" if isinstance(token, int): return token if token == "*": @@ -629,7 +630,7 @@ def _check_array(key: str, rule: dict, arrays: Mapping, report, context, side) - want = _dimension(token, context, side) if want is None: continue - # `total_frames` is the sole authoritative length and a stored + # ``total_frames`` is the authoritative length and a stored # array may carry a padded tail, so axis 0 is a lower bound. if axis == 0: if shape[0] < want: @@ -654,11 +655,10 @@ def _check_array(key: str, rule: dict, arrays: Mapping, report, context, side) - def _expand_key(template: str, arrays: Mapping, sides) -> list[tuple[str, str | None]]: - """Expand one schema key into the concrete keys it names. + """Expand a schema key into the concrete array keys it selects. - ``{side}`` expands over the episode's arms. ``*`` matches the keys the - episode stores, so a wildcard rule checks what is there and never demands - a key. + ``{side}`` expands over the resolved side candidates. ``*`` selects stored + keys only, so a wildcard rule never creates a missing-key finding. """ if "{side}" in template: candidates = [(template.format(side=side), side) for side in sides] @@ -684,16 +684,17 @@ def _expand_key(template: str, arrays: Mapping, sides) -> list[tuple[str, str | def validate_episode(path: str | Path, *, strict: bool = False) -> Report: - """Validate one zarr episode against ``schema/episode_v3.yaml``. - - + """Validate one Zarr episode against ``schema/episode_v3.yaml``. Args: path: The episode ``.zarr`` directory. - strict: If true, promote the rules the corpus does not meet yet from - warnings to errors. + strict: If true, treat ``required: strict`` failures as errors rather + than warnings. Returns: - A report holding one finding per rule that ran. + The findings emitted while validating the episode. + + Raises: + SchemaError: If the schema contains an unsupported declaration. """ path = Path(path) report = Report(path=path, strict=strict) @@ -736,8 +737,8 @@ def validate_episode(path: str | Path, *, strict: bool = False) -> Report: resolved = _resolve(attrs, report) if resolved is None: return report - # Say which platform and end-effectors the array rules ran against, so a - # report explains itself without a second lookup. + # Record the resolved platform and end-effectors used by subsequent array + # conditions and dimensions. report.add(OK, "embodiment", resolved.describe()) context["resolved"] = resolved @@ -754,7 +755,7 @@ def validate_episode(path: str | Path, *, strict: bool = False) -> Report: def _resolve(attrs: Mapping, report: Report) -> ResolvedEmbodiment | None: - """Resolve the episode's embodiment, preferring its morphology block.""" + """Resolve morphology first, then fall back to the embodiment name.""" for spec in (attrs.get("morphology"), attrs.get("embodiment")): if not spec: continue @@ -771,10 +772,10 @@ def _resolve(attrs: Mapping, report: Report) -> ResolvedEmbodiment | None: def main(argv: list[str] | None = None) -> int: - """Run the validator over one or more episodes. + """Validate CLI paths and return an exit status. Returns: - ``0`` if every episode passed, ``1`` otherwise. + ``0`` if every report has no errors, otherwise ``1``. """ parser = argparse.ArgumentParser( prog="python -m egomimic.rldb.zarr.validate", diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index 5d56b3eff..c9c4642b1 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/egomimic/rldb/zarr/zarr_dataset_multi.py @@ -235,9 +235,8 @@ def _load_zarr_datasets(self, search_path: Path, valid_folder_names: set[str]): key_map=self.key_map, transform_list=self.transform_list, ) - # A structural sample is real data in the schema sense and not - # real data in the training sense. Refuse it here so it cannot - # reach a training run through any resolver. + # Resolver-based loading accepts only the exact ``complete`` + # status; structural samples and unknown statuses are skipped. if not is_complete(ds_obj.metadata): logger.warning( "Skipping %s: data_status is %r, not 'complete'", @@ -1577,22 +1576,21 @@ def init_episode(self): @property def data_status(self) -> str: - """Pass-through to ZarrEpisode.data_status (read from zarr metadata).""" + """Return the episode status, including the legacy default.""" return self.episode_reader.data_status @property def calibration(self) -> Calibration | None: - """Pass-through to ZarrEpisode.calibration (read from zarr metadata).""" + """Return parsed current or legacy episode calibration.""" return self.episode_reader.calibration def _build_extrinsic_poses(self) -> dict[str, np.ndarray]: - """Return this episode's per-arm ``base_T_cam`` poses for the batch. + """Build transform-input poses for the calibration's default camera. - The transform pipeline expresses actions in the camera frame, so it - needs the rig that recorded this episode. Two vendors on one platform - have two rigs, and an embodiment class constant cannot tell them apart. - An episode that declares no extrinsics contributes nothing here and the - transform falls back to its embodiment default. + For each arm with enough calibration to compose ``base_T_cam``, convert + the matrix to the pose layout consumed by the transform pipeline. + Missing calibration leaves the sample key absent so the configured + transform may supply its compatibility fallback. Returns: A mapping from ``_base_T_cam_pose`` to a 7-value @@ -1776,8 +1774,8 @@ def _next(reason: str, key: str = "") -> int: if retry: continue - # Per-episode extrinsics travel with the sample so the pipeline - # transforms into the frame of the rig that recorded this episode. + # Inject episode-derived poses before transforms run. Transform + # fallbacks use ``setdefault`` and therefore preserve these values. for key, pose in self._extrinsic_poses.items(): data[key] = pose.copy() @@ -1790,11 +1788,10 @@ def _next(reason: str, key: str = "") -> int: data[k] = torch.from_numpy(v).to(torch.float32) data["embodiment"] = get_embodiment_id(self.embodiment) - # Per-episode camera intrinsics travel with the batch so a single - # data-driven embodiment can project without a hardcoded class const. - # The calibration shim picks the front camera and normalizes a bare - # 3x3 K; an episode without one gets the NaN sentinel, which - # _intrinsics_from_batch treats as "fall back to cls.INTRINSICS". + # ``Calibration.K()`` applies the default-camera selection and the + # parser has already normalized a 3x3 K to 3x4. A missing matrix is + # represented by NaNs, which ``_intrinsics_from_batch`` interprets + # as a request for the embodiment-class fallback. calibration = self.calibration K = None if calibration is None else calibration.K() if K is None: @@ -1890,16 +1887,15 @@ def __init__(self, path: str | Path): @property def data_status(self) -> str: - """Return ``complete`` or ``structural_sample`` for this episode.""" + """Return the status, defaulting a missing or falsey value to ``complete``.""" return data_status(self.metadata) @property def calibration(self) -> Calibration | None: - """Per-episode calibration, read through the one shim. + """Return calibration normalized by :func:`read_calibration`. - Returns the ``calibration`` attribute block, or the same information - lifted from the legacy ``intrinsics``/``extrinsics`` pair, or ``None`` - when the episode states neither. + The result comes from the current block or a lifted legacy attribute + pair. It is ``None`` when neither representation is present. """ return self._calibration diff --git a/egomimic/rldb/zarr/zarr_writer.py b/egomimic/rldb/zarr/zarr_writer.py index 6537bcdfa..a59eb2b75 100644 --- a/egomimic/rldb/zarr/zarr_writer.py +++ b/egomimic/rldb/zarr/zarr_writer.py @@ -342,18 +342,15 @@ def __init__( chunk_timesteps: Number of frames in each numeric-array chunk. intrinsics: A mapping from camera keys to 3×4 camera matrices. This low-level constructor permits ``None``. - extrinsics: ``None`` or a mapping from keys to 4×4 ``ref_T_cam`` - matrices. Robot episodes use arm names as keys. See - ``docs/CONVENTIONS.md``. - calibration: ``None`` or the per-episode calibration to store under - the ``calibration`` attribute. A ``Calibration`` or the mapping - form of one. - data_status: ``complete`` for a finished recording, or - ``structural_sample`` for one sent to show the shape of a - delivery. Only a complete episode gets a database row and - reaches a dataset. - strict: If true, an image stream that no camera matrix covers is an - error instead of a warning. + extrinsics: The legacy ``{side: base_T_cam}`` mapping, or ``None``. + See ``docs/CONVENTIONS.md``. + calibration: A ``Calibration`` or its attribute-block mapping. The + writer stores it under ``calibration``. + data_status: Initial episode status. The staging path and dataset + resolvers accept only ``complete`` and exclude + ``structural_sample``. + strict: If true, reject an image stream with no matching camera + matrix; otherwise log a warning. verbose: If true, print progress information during writes. """ self.episode_path = Path(episode_path) @@ -734,11 +731,11 @@ def _write_annotations( } def _check_camera_coverage(self, image_keys) -> None: - """Check that every stored image stream has a camera matrix. + """Compare stored image streams with the effective camera calibration. - `intrinsics` only has to be non-empty, so an episode that writes three - image streams and calibrates one passes. Coverage closes that: it is - checked against the image keys the episode actually stores. + The current ``calibration`` block is authoritative when present; + otherwise the legacy ``intrinsics`` and ``extrinsics`` attributes are + lifted. Each ``images.`` key must match a camera carrying ``K``. Args: image_keys: The episode's image array keys. @@ -792,11 +789,8 @@ def _build_metadata( if self.extrinsics is not None: metadata["extrinsics"] = _intrinsics_to_jsonable(self.extrinsics) - # Apply overrides — but NEVER let them clobber the validated camera - # metadata. calibration/intrinsics/extrinsics are validated in - # create_and_write; a converter's metadata_override that happens to - # carry a stale or empty key must not silently overwrite the validated - # values (this was the Mecka clobber bug → empty intrinsics in zarr.json). + # Preserve the current and legacy calibration representations supplied + # to the writer; metadata overrides may add fields but not replace them. if metadata_override: override = { k: v @@ -845,21 +839,15 @@ def create_and_write( chunk_timesteps: Number of frames in each numeric-array chunk. intrinsics: A non-empty mapping from camera keys to 3×4 camera matrices. Optional only when ``calibration`` supplies them. - extrinsics: ``None`` or a non-empty mapping from keys to 4×4 - ``ref_T_cam`` matrices. Robot episodes use arm names as keys. - See ``docs/CONVENTIONS.md``. - calibration: ``None`` or the per-episode calibration described in - ``egomimic/rldb/zarr/calibration.py``. When it is present the - writer derives ``intrinsics`` and ``extrinsics`` from it for any - of the two that the caller omits, so readers that predate the - block keep working. - data_status: ``complete`` for a finished recording, or - ``structural_sample`` for one sent to show the shape of a - delivery. Only a complete episode gets a database row and - reaches a dataset. - strict: If true, an image stream that no camera matrix covers is an - error instead of a warning. It stays opt-in because every EVA - episode in the corpus stores three streams and calibrates one. + extrinsics: The legacy non-empty ``{side: base_T_cam}`` mapping, or + ``None``. See ``docs/CONVENTIONS.md``. + calibration: A ``Calibration`` or its attribute-block mapping. For + each omitted legacy argument, the writer derives an equivalent + ``intrinsics`` or ``extrinsics`` value for older readers. + data_status: ``complete`` or ``structural_sample``. The staging path + and dataset resolvers accept only ``complete`` episodes. + strict: If true, reject an image stream with no matching camera + matrix; otherwise log a warning. metadata_override: Additional episode metadata. This mapping cannot replace ``calibration``, ``intrinsics``, or ``extrinsics``. @@ -872,6 +860,7 @@ def create_and_write( ValueError: If ``intrinsics`` is empty or contains a non-3×4 value. ValueError: If ``extrinsics`` is not ``None`` or a non-empty mapping of 4×4 matrices. + ValueError: If ``data_status`` is not a supported value. ValueError: If ``strict`` is set and an image stream carries no camera matrix. CalibrationError: If ``calibration`` is malformed. @@ -887,9 +876,8 @@ def create_and_write( f"{[m.name.lower() for m in EMBODIMENT]}, got {embodiment!r}. " "See CONTRIBUTING_DATA.md §9." ) - # `calibration` is the current form; `intrinsics`/`extrinsics` are the - # legacy pair. Writing both keeps every reader working, so derive - # whichever the caller left out. + # Derive only omitted legacy fields. The current block remains the + # source used by readers that understand it. if calibration is not None: calibration = parse_calibration(calibration) if intrinsics is None: diff --git a/egomimic/scripts/backfill_scripts/stage_processed_folder.py b/egomimic/scripts/backfill_scripts/stage_processed_folder.py index 04eec02fe..075e80d77 100644 --- a/egomimic/scripts/backfill_scripts/stage_processed_folder.py +++ b/egomimic/scripts/backfill_scripts/stage_processed_folder.py @@ -85,11 +85,13 @@ def find_zarr_prefixes(s3, base): def read_batch(prefixes, folder, cfg): - """Read the zarr.json for a batch of prefixes -> list[row dict]. FULLY - self-contained (only stdlib + boto3, imported inside) so it runs on any Ray - worker regardless of whether egomimic / ~/.egoverse_env are present there. + """Build staging rows from the episode metadata under each S3 prefix. - An episode whose ``data_status`` is not ``complete`` yields no row.""" + The function imports its dependencies locally so a Ray worker needs neither + the ``egomimic`` package nor ``~/.egoverse_env``. It omits every episode + whose status is not exactly ``complete``; a missing or falsey status + defaults to ``complete`` for compatibility with older episodes. + """ import json as _json import re as _re from datetime import datetime as _dt @@ -123,9 +125,8 @@ def read_batch(prefixes, folder, cfg): ).replace(tzinfo=_tz.utc).isoformat() except ValueError: created_at = None - # A structural sample never gets a row. Mirrors - # `egomimic.rldb.zarr.episode_attrs.DATA_STATUS_COMPLETE`, restated as - # a literal because this function must run on a bare Ray worker. + # Keep the literal aligned with ``DATA_STATUS_COMPLETE`` without + # importing the package on a bare Ray worker. if (attrs.get("data_status") or "complete") != "complete": continue feats = attrs.get("features") diff --git a/egomimic/scripts/data_visualization/inspector_lib/dataset_view.py b/egomimic/scripts/data_visualization/inspector_lib/dataset_view.py index b2f1987de..be5032b6a 100644 --- a/egomimic/scripts/data_visualization/inspector_lib/dataset_view.py +++ b/egomimic/scripts/data_visualization/inspector_lib/dataset_view.py @@ -286,13 +286,12 @@ def _badge(img_rgb, text: str): def _calibration_from_zarr(grp): - """Read the episode calibration through the one shim. + """Read current or legacy calibration without breaking the inspector. Returns: - A ``Calibration``, or ``None`` if the episode states none or the stored - calibration is malformed. The inspector renders episodes that a - validator would reject, so a bad block degrades the overlay rather than - raising. + The normalized calibration, or ``None`` if metadata access fails, no + calibration is present, or its stored representation is malformed. + Invalid calibration disables the overlay instead of the episode view. """ try: return read_calibration(dict(grp.attrs)) @@ -301,20 +300,21 @@ def _calibration_from_zarr(grp): def _intrinsics_from_zarr(grp): - """Per-episode 3x4 camera matrix K for the front camera. + """Return the default camera's normalized 3×4 ``[K_3x3 | 0]`` matrix. - Returns None if the episode declares no calibrated camera. + Return ``None`` if the episode has no valid calibration or if the selected + camera has no matrix. """ calibration = _calibration_from_zarr(grp) return None if calibration is None else calibration.K() def _extrinsics_from_zarr(grp): - """Read per-arm ``base_T_cam`` matrices from episode metadata. + """Compose per-arm poses of the calibration's default camera. Returns: - A dictionary with valid ``"left"`` and ``"right"`` 4×4 matrices. - The function returns ``None`` if no valid matrix is present. + A dictionary mapping available ``"left"`` and ``"right"`` arms to + 4×4 ``base_T_cam`` matrices, or ``None`` if none can be composed. """ calibration = _calibration_from_zarr(grp) if calibration is None: From 375af2390d34aef1bde8ecdb83477cb2d0a76ab0 Mon Sep 17 00:00:00 2001 From: jaynye Date: Thu, 3 Sep 2026 18:24:59 -0400 Subject: [PATCH 9/9] factorized camera, schema, and length validation --- egomimic/rldb/zarr/schema/episode_v3.yaml | 104 +++++++++--- egomimic/rldb/zarr/test_calibration.py | 14 +- egomimic/rldb/zarr/test_data_status.py | 20 ++- egomimic/rldb/zarr/test_validate.py | 142 ++++++++++++---- egomimic/rldb/zarr/validate.py | 190 ++++++++++++++++------ egomimic/rldb/zarr/zarr_writer.py | 31 ++-- 6 files changed, 369 insertions(+), 132 deletions(-) diff --git a/egomimic/rldb/zarr/schema/episode_v3.yaml b/egomimic/rldb/zarr/schema/episode_v3.yaml index e233f6623..786a57e6d 100644 --- a/egomimic/rldb/zarr/schema/episode_v3.yaml +++ b/egomimic/rldb/zarr/schema/episode_v3.yaml @@ -1,19 +1,26 @@ # The declarative portion of the episode contract. # # This file selects fields and named checks, supplies their parameters, and -# assigns failure severity. `validate.py` implements the schema interpreter and -# the named check predicates referenced here. +# assigns each rule a severity class. `validate.py` implements the schema +# interpreter and the named check predicates referenced here. # # Run it with: -# python -m egomimic.rldb.zarr.validate [--strict] +# python -m egomimic.rldb.zarr.validate [RULE FLAGS] # # --------------------------------------------------------------------------- # attributes -# required true | false | strict -# For a missing attribute or array, `true` emits an error, `false` -# emits no finding, and `strict` emits a warning unless --strict -# promotes it to an error. Present attributes and arrays must meet -# their declared type and shape regardless of this setting. +# required true | false. A missing required attribute or array emits a +# finding. An optional field emits no finding when absent. +# severity integrity | adoption | coverage +# Integrity findings are always errors. Adoption and coverage +# findings are errors unless their named rule is waived for this +# invocation. Present fields must meet their declared type and +# shape even when their presence requirement is waived. +# name Stable rule name when the attribute or key is not the desired +# CLI flag stem. Dashes replace underscores in CLI flags. +# why User-facing reason for a waivable rule. This is required because +# the schema generates CLI help and must explain what accepting a +# gap gives up without relying on a separate rollout document. # type str | int | bool | mapping # min smallest accepted value for an int # choices the accepted values @@ -21,14 +28,16 @@ # # checks # name a predicate implemented in `validate.py` -# required controls the severity of a predicate failure as described above +# required whether a predicate failure emits a finding +# severity controls whether that finding can be waived, as described above # predicate-specific parameters # # arrays # key an array key. `{side}` expands over arity-derived candidates, # filtered by resolved end-effectors; `*` selects stored matches. -# required true | false | strict, as above. A wildcard key is never -# required. +# required true | false, as above. A wildcard key is never required. +# severity controls missing-key severity. Shape and dtype failures are +# integrity errors for every array, including a waived one. # when conditions on the resolved embodiment; skip the rule when any # condition does not hold # shape expected dimensions, see below @@ -57,48 +66,69 @@ known_versions: [v3.0, v3.1] attributes: embodiment: required: true + severity: integrity type: str check: embodiment_name total_frames: required: true + severity: integrity type: int min: 1 fps: required: true + severity: integrity type: int min: 1 - # `complete` or `structural_sample`. An episode written before this - # attribute existed reads as complete, so it is required only under --strict. data_status: - required: strict + required: true + severity: adoption + why: >- + Require an explicit data status so an incomplete structural sample + cannot inherit the legacy "complete" default and enter training. type: str choices: [complete, structural_sample] task_name: required: true + severity: integrity type: str task_description: required: false + severity: integrity type: str features: required: true + severity: integrity type: mapping schema_version: - required: strict + required: true + severity: adoption + why: >- + Require a schema version so readers can select the episode contract + explicitly instead of inferring it from whichever fields are present. type: str check: schema_version morphology: required: false + severity: integrity type: mapping check: morphology calibration: - required: strict + name: calibration_block + required: true + severity: adoption + why: >- + Require the calibration block so camera models, distortion, resolution, + and frame relationships are explicit; waive only when legacy calibration + attributes contain enough information for the compatibility reader. type: mapping check: calibration intrinsics: required: false + severity: integrity type: mapping extrinsics: required: false + severity: integrity type: mapping checks: @@ -106,16 +136,22 @@ checks: # block or lifted from the legacy attributes. - name: calibration_present required: true + severity: integrity # Match each stored `images.` stream to a camera entry carrying K. - # This remains a warning without --strict during the compatibility rollout. - name: camera_coverage - required: strict + required: true + severity: coverage + why: >- + Require a camera matrix for every image stream because projection and + reprojection checks cannot interpret an uncalibrated view; waive only + after explicitly accepting that physical calibration gap. # For a listed 7D pose track with at least two retained rows, reject identical # rows or more than 1% exact +identity/-identity quaternions. - name: pose_degeneracy required: true + severity: integrity suffixes: [obs_ee_pose, cmd_ee_pose, obs_head_pose] identity_rotation_fraction: 0.01 @@ -123,24 +159,32 @@ checks: # numerically identity `base_T_cam` composed for an arm and the default camera. - name: calibration_degeneracy required: true + severity: integrity # Reject the exact synthetic-intrinsics signature `fx == fy == width`, # `cx == width / 2`, and `cy == height / 2`. - name: intrinsics_signature required: true + severity: integrity # Reject non-increasing retained timestamps, stored keys containing # `relative_timestamp`, and a power-of-two step GCD of at least 64 ns (the # heuristic signature of float seconds converted to integer nanoseconds). - name: timestamps required: true + severity: integrity key: obs_rgb_timestamps_ns banned_key_substrings: [relative_timestamp] # Valid annotation intervals must jointly cover at least 90% of - # `[0, total_frames)`. Keep this transitional rule strict-only. + # `[0, total_frames)`. - name: annotation_coverage - required: strict + required: true + severity: coverage + why: >- + Require annotations to cover at least 90% of retained frames because an + unlabelled tail gives those frames no reliable language target; waive + only after accepting the gap instead of annotating or trimming the tail. key: annotations minimum: 0.9 @@ -148,12 +192,14 @@ checks: # structured metadata needs a separate schema field. - name: annotation_text required: true + severity: integrity key: annotations banned_delimiters: [" | "] arrays: - key: "{side}.obs_ee_pose" required: true + severity: integrity shape: [T, 7] dtype: float # Observed end-effector-root pose. Require one for every resolved side, @@ -161,58 +207,72 @@ arrays: - key: "{side}.cmd_ee_pose" required: false + severity: integrity shape: [T, 7] dtype: float - key: "{side}.obs_joints" when: {has_arm_chain: true} required: false + severity: integrity shape: [T, arm_dof] dtype: float - key: "{side}.cmd_joints" when: {has_arm_chain: true} required: false + severity: integrity shape: [T, arm_dof] dtype: float - key: "{side}.obs_gripper" when: {end_effector_class: parallel_jaw} required: true + severity: integrity shape: [T, 1] dtype: float - key: "{side}.cmd_gripper" when: {end_effector_class: parallel_jaw} required: true + severity: integrity shape: [T, 1] dtype: float - key: "{side}.obs_keypoints" when: {end_effector_class: human_hand} required: true + severity: integrity shape: [T, kp3] dtype: float - key: obs_head_pose when: {platform_kind: human} required: true + severity: integrity shape: [T, 7] dtype: float # Canonical RGB timestamp array. When present it must be one-dimensional, - # integer-valued, and at least `total_frames` long. Presence is strict-only. + # integer-valued, and at least `total_frames` long. - key: obs_rgb_timestamps_ns - required: strict + name: rgb_timestamps + required: true + severity: coverage + why: >- + Require integer UTC-nanosecond RGB timestamps so streams can be aligned + without reconstructing the capture clock after the recording is over. shape: [T] dtype: int - key: "images.*" required: false + severity: integrity shape: [T] dtype: object - key: annotations required: false + severity: integrity shape: ["*"] dtype: object diff --git a/egomimic/rldb/zarr/test_calibration.py b/egomimic/rldb/zarr/test_calibration.py index 051ac7220..f16304c21 100644 --- a/egomimic/rldb/zarr/test_calibration.py +++ b/egomimic/rldb/zarr/test_calibration.py @@ -249,9 +249,7 @@ def test_uncalibrated_cameras_lists_streams_without_a_matrix() -> None: assert uncalibrated_cameras(["left.obs_gripper"], None) == [] -def test_coverage_is_a_warning_by_default_and_an_error_under_strict( - tmp_path, caplog -) -> None: +def test_writer_camera_coverage_can_be_required_explicitly(tmp_path, caplog) -> None: images = { "images.front_1": np.zeros((4, 8, 8, 3), dtype=np.uint8), "images.left_wrist": np.zeros((4, 8, 8, 3), dtype=np.uint8), @@ -266,14 +264,16 @@ def test_coverage_is_a_warning_by_default_and_an_error_under_strict( with pytest.raises(ValueError, match=r"no camera matrix.*left_wrist"): _write_episode( - tmp_path / "strict.zarr", + tmp_path / "required.zarr", image_data=images, intrinsics={"front_1": K_FRONT}, - strict=True, + require_camera_coverage=True, ) -def test_full_coverage_passes_under_strict(tmp_path, caplog) -> None: +def test_full_coverage_passes_when_camera_coverage_is_required( + tmp_path, caplog +) -> None: with caplog.at_level("WARNING"): _write_episode( tmp_path / "covered.zarr", @@ -282,7 +282,7 @@ def test_full_coverage_passes_under_strict(tmp_path, caplog) -> None: "reference_frame": "camera:front_1", "cameras": {"front_1": {"K": K_FRONT.tolist()}}, }, - strict=True, + require_camera_coverage=True, ) assert "no camera matrix" not in caplog.text diff --git a/egomimic/rldb/zarr/test_data_status.py b/egomimic/rldb/zarr/test_data_status.py index 65607c667..285eb3426 100644 --- a/egomimic/rldb/zarr/test_data_status.py +++ b/egomimic/rldb/zarr/test_data_status.py @@ -87,17 +87,19 @@ def test_the_validator_reads_and_checks_the_status(tmp_path) -> None: _write(tmp_path / "legacy.zarr") store = zarr.open_group(str(tmp_path / "legacy.zarr"), mode="a") del store.attrs["data_status"] - lenient = validate_episode(tmp_path / "legacy.zarr") - strict = validate_episode(tmp_path / "legacy.zarr", strict=True) + required = validate_episode(tmp_path / "legacy.zarr") + waived = validate_episode( + tmp_path / "legacy.zarr", requirements={"data_status": False} + ) statuses = { - "lenient": next( - f.level for f in lenient.findings if f.check == "attrs.data_status" + "required": next( + f.level for f in required.findings if f.check == "attrs.data_status" ), - "strict": next( - f.level for f in strict.findings if f.check == "attrs.data_status" + "waived": next( + f.level for f in waived.findings if f.check == "attrs.data_status" ), } - assert statuses == {"lenient": WARNING, "strict": ERROR} + assert statuses == {"required": ERROR, "waived": WARNING} def test_the_validator_rejects_an_unknown_status(tmp_path) -> None: @@ -105,7 +107,9 @@ def test_the_validator_rejects_an_unknown_status(tmp_path) -> None: store = zarr.open_group(str(tmp_path / "odd.zarr"), mode="a") store.attrs["data_status"] = "in_review" - report = validate_episode(tmp_path / "odd.zarr") + report = validate_episode( + tmp_path / "odd.zarr", requirements={"data_status": False} + ) finding = next(f for f in report.findings if f.check == "attrs.data_status") assert finding.level == ERROR diff --git a/egomimic/rldb/zarr/test_validate.py b/egomimic/rldb/zarr/test_validate.py index 2b43e38b8..5f5761eb1 100644 --- a/egomimic/rldb/zarr/test_validate.py +++ b/egomimic/rldb/zarr/test_validate.py @@ -8,9 +8,14 @@ ERROR, OK, WARNING, + Report, + SchemaError, + _build_parser, + _level, load_schema, main, validate_episode, + waivable_rules, ) from egomimic.rldb.zarr.zarr_writer import ZarrWriter @@ -60,16 +65,16 @@ def _write_eva(path, *, numeric=None, images=None, **kwargs) -> None: numeric_data=_eva_numeric() if numeric is None else numeric, image_data=images, embodiment=kwargs.pop("embodiment", "eva_bimanual"), - chunk_timesteps=LENGTH, + chunk_timesteps=kwargs.pop("chunk_timesteps", LENGTH), **kwargs, ) -def test_a_complete_eva_episode_passes_under_strict(tmp_path) -> None: +def test_a_complete_eva_episode_passes_the_default_contract(tmp_path) -> None: path = tmp_path / "eva.zarr" _write_eva(path, images={"images.front_1": np.zeros((LENGTH, 8, 8, 3), np.uint8)}) - report = validate_episode(path, strict=True) + report = validate_episode(path) assert report.ok, report.text() assert not report.warnings @@ -102,15 +107,7 @@ def test_a_wrong_width_is_reported_with_the_dimension_that_set_it(tmp_path) -> N def test_a_padded_tail_is_not_an_error(tmp_path) -> None: path = tmp_path / "eva.zarr" - ZarrWriter.create_and_write( - episode_path=path, - numeric_data=_eva_numeric(), - embodiment="eva_bimanual", - chunk_timesteps=3, # 4 frames pad out to 6 - annotations=[("fold the towel", 0, LENGTH)], - intrinsics={"front_1": K}, - extrinsics=Eva.EXTRINSICS, - ) + _write_eva(path, chunk_timesteps=3) # 4 frames pad out to 6 report = validate_episode(path) @@ -129,7 +126,7 @@ def test_an_array_shorter_than_total_frames_is_an_error(tmp_path) -> None: assert "holds 4 frames for total_frames 14" in report.text() -def test_strict_promotes_the_rules_the_corpus_does_not_meet_yet(tmp_path) -> None: +def test_each_waiver_lowers_only_its_named_rule_to_a_warning(tmp_path) -> None: path = tmp_path / "legacy.zarr" ZarrWriter.create_and_write( episode_path=path, @@ -145,15 +142,23 @@ def test_strict_promotes_the_rules_the_corpus_does_not_meet_yet(tmp_path) -> Non extrinsics=Eva.EXTRINSICS, ) - lenient = validate_episode(path) - strict = validate_episode(path, strict=True) + default = validate_episode(path) + waived = validate_episode( + path, + requirements={ + "camera_coverage": False, + "calibration_block": False, + "schema_version": False, + }, + ) - assert lenient.ok - assert _levels(lenient)["camera_coverage"] == WARNING - assert _levels(lenient)["attrs.calibration"] == WARNING - assert not strict.ok - assert _levels(strict)["camera_coverage"] == ERROR - assert "left_wrist" in strict.text() + assert not default.ok + assert _levels(default)["camera_coverage"] == ERROR + assert _levels(default)["attrs.calibration"] == ERROR + assert waived.ok + assert _levels(waived)["camera_coverage"] == WARNING + assert _levels(waived)["attrs.calibration"] == WARNING + assert "left_wrist" in waived.text() def test_a_single_arm_episode_owes_only_its_own_arm(tmp_path) -> None: @@ -162,7 +167,7 @@ def test_a_single_arm_episode_owes_only_its_own_arm(tmp_path) -> None: } _write_eva(tmp_path / "left.zarr", numeric=numeric, embodiment="eva_left_arm") - report = validate_episode(tmp_path / "left.zarr", strict=True) + report = validate_episode(tmp_path / "left.zarr") assert report.ok, report.text() assert "right.obs_ee_pose" not in _levels(report) @@ -239,7 +244,7 @@ def test_a_morphology_block_selects_the_end_effector(tmp_path) -> None: "vendor": "rl2", } - report = validate_episode(path, strict=True) + report = validate_episode(path) assert report.ok, report.text() assert "eva_parallel_jaw" in report.text(verbose=True) @@ -257,7 +262,7 @@ def test_the_cli_exit_code_follows_the_findings(tmp_path, capsys) -> None: _write_eva(path) assert main([str(path)]) == 0 - assert main([str(path), "--strict"]) == 0 + assert main([str(path), "--camera-coverage"]) == 0 assert main([str(tmp_path / "absent.zarr")]) == 1 assert "cannot open as a zarr group" in capsys.readouterr().out @@ -271,6 +276,7 @@ def test_the_cli_can_report_json(tmp_path, capsys) -> None: payload = __import__("json").loads(capsys.readouterr().out) assert payload[0]["ok"] is True assert payload[0]["path"] == str(path) + assert payload[0]["requirements"] == dict.fromkeys(waivable_rules(), True) @pytest.mark.parametrize("section", ["attributes", "checks", "arrays"]) @@ -280,7 +286,64 @@ def test_every_schema_rule_declares_a_usable_requirement(section) -> None: entries = rules.values() if isinstance(rules, dict) else rules assert entries for rule in entries: - assert rule.get("required", False) in (True, False, "strict") + assert isinstance(rule.get("required", False), bool) + assert rule["severity"] in ("integrity", "adoption", "coverage") + + +def test_waivable_rules_own_both_cli_flags_and_explain_why() -> None: + rules = waivable_rules() + assert set(rules) == { + "schema_version", + "data_status", + "calibration_block", + "camera_coverage", + "annotation_coverage", + "rgb_timestamps", + } + + help_text = " ".join(_build_parser().format_help().split()) + for name, rule in rules.items(): + flag = name.replace("_", "-") + assert f"--{flag}" in help_text + assert f"--no-{flag}" in help_text + assert " ".join(rule["why"].split()) in help_text + + args = _build_parser().parse_args( + ["episode.zarr", "--no-camera-coverage", "--annotation-coverage"] + ) + assert args.require_camera_coverage is False + assert args.require_annotation_coverage is True + assert args.require_schema_version is True + + +def test_integrity_rules_cannot_be_waived(tmp_path) -> None: + with pytest.raises(SchemaError, match="unknown waivable rule 'pose_degeneracy'"): + validate_episode( + tmp_path / "episode.zarr", + requirements={"pose_degeneracy": False}, + ) + + +def test_rule_levels_follow_the_invocation_without_weakening_integrity( + tmp_path, +) -> None: + schema = load_schema() + requirements = dict.fromkeys(waivable_rules(schema), True) + report = Report(path=tmp_path / "episode.zarr", requirements=requirements) + + calibration = schema["attributes"]["calibration"] + assert _level(calibration, report, "calibration") == ERROR + report.requirements["calibration_block"] = False + assert _level(calibration, report, "calibration") == WARNING + + integrity = schema["checks"][0] + assert integrity["name"] == "calibration_present" + assert _level(integrity, report, integrity["name"]) == ERROR + + +def test_the_retired_strict_flag_is_rejected() -> None: + with pytest.raises(SystemExit): + _build_parser().parse_args(["episode.zarr", "--strict"]) # -------------------------------------------------------------------------- @@ -402,6 +465,20 @@ def test_a_float64_quantized_clock_is_an_error(tmp_path) -> None: assert "quantized to 256 ns" in _finding(report, "timestamps") +def test_a_timestamp_waiver_does_not_waive_a_wrong_dtype(tmp_path) -> None: + numeric = _eva_numeric() + numeric["obs_rgb_timestamps_ns"] = np.arange(LENGTH, dtype=np.float64) + _write_eva(tmp_path / "float_clock.zarr", numeric=numeric) + + report = validate_episode( + tmp_path / "float_clock.zarr", + requirements={"rgb_timestamps": False}, + ) + + assert _levels(report)["obs_rgb_timestamps_ns"] == ERROR + assert "expected int" in _finding(report, "obs_rgb_timestamps_ns") + + def test_a_second_time_base_is_an_error(tmp_path) -> None: numeric = _eva_numeric() numeric["relative_timestamp_s"] = np.linspace(0.0, 0.1, LENGTH)[:, None] @@ -418,18 +495,21 @@ def test_annotation_coverage_below_the_minimum_is_reported(tmp_path) -> None: tmp_path / "thin.zarr", annotations=[("fold the towel", 0, LENGTH - 2)] ) - lenient = validate_episode(tmp_path / "thin.zarr") - strict = validate_episode(tmp_path / "thin.zarr", strict=True) + required = validate_episode(tmp_path / "thin.zarr") + waived = validate_episode( + tmp_path / "thin.zarr", + requirements={"annotation_coverage": False}, + ) - assert _levels(lenient)["annotation_coverage"] == WARNING - assert _levels(strict)["annotation_coverage"] == ERROR - assert "cover 50% of the episode" in _finding(strict, "annotation_coverage") + assert _levels(required)["annotation_coverage"] == ERROR + assert _levels(waived)["annotation_coverage"] == WARNING + assert "cover 50% of the episode" in _finding(waived, "annotation_coverage") def test_an_annotation_span_past_the_episode_is_reported(tmp_path) -> None: _write_eva(tmp_path / "over.zarr", annotations=[("fold", 0, LENGTH + 5)]) - report = validate_episode(tmp_path / "over.zarr", strict=True) + report = validate_episode(tmp_path / "over.zarr") assert "outside [0, 4)" in _finding(report, "annotation_coverage") diff --git a/egomimic/rldb/zarr/validate.py b/egomimic/rldb/zarr/validate.py index 92fce03fe..82c8ee302 100644 --- a/egomimic/rldb/zarr/validate.py +++ b/egomimic/rldb/zarr/validate.py @@ -2,13 +2,13 @@ Run it with:: - python -m egomimic.rldb.zarr.validate [--strict] + python -m egomimic.rldb.zarr.validate [RULE FLAGS] -The schema declares attributes, arrays, conditions, thresholds, and severity. -This module interprets those declarations and implements the named predicates. -``--strict`` promotes failures marked ``required: strict`` from warnings to -errors; validation of a present attribute or array is always an error when its -declared type or shape is wrong. +The schema declares attributes, arrays, conditions, thresholds, and severity +classes. Integrity failures are always errors. Adoption and coverage failures +are errors by default, and each such rule supplies a matched ``--`` / +``--no-`` CLI pair. Validation of a present attribute or array is always +an error when its declared type or shape is wrong. """ from __future__ import annotations @@ -16,6 +16,7 @@ import argparse import functools import json +import re import sys from collections.abc import Mapping from dataclasses import dataclass, field @@ -46,8 +47,10 @@ WARNING = "warning" OK = "ok" -#: Accepted values for a schema rule's ``required`` field. -_REQUIRED_VALUES = (True, False, "strict") +#: Accepted values for a schema rule's ``severity`` class. +_SEVERITY_VALUES = ("integrity", "adoption", "coverage") +_INTEGRITY = "integrity" +_RULE_NAME = re.compile(r"^[a-z][a-z0-9_]*$") _TYPE_NAMES = { "str": str, @@ -87,13 +90,13 @@ class Report: path: The episode directory. findings: Findings in validation order. A schema rule may emit zero, one, or multiple findings after key expansion. - strict: Whether ``required: strict`` failures are errors instead of - warnings. + requirements: Whether each waivable rule is required. A false value + means that the invocation reports that rule's failures as warnings. """ path: Path findings: list[Finding] = field(default_factory=list) - strict: bool = False + requirements: dict[str, bool] = field(default_factory=dict) def add(self, level: str, check: str, message: str) -> None: self.findings.append(Finding(level, check, message)) @@ -132,7 +135,7 @@ def text(self, verbose: bool = False) -> str: def to_jsonable(self) -> dict[str, Any]: return { "path": str(self.path), - "strict": self.strict, + "requirements": self.requirements, "ok": self.ok, "findings": [ {"level": f.level, "check": f.check, "message": f.message} @@ -143,36 +146,103 @@ def to_jsonable(self) -> dict[str, Any]: @functools.lru_cache(maxsize=1) def load_schema() -> dict: - """Load the schema and validate every declared ``required`` value. + """Load the schema and validate its requirement declarations. Returns: The parsed schema. Raises: - SchemaError: If a rule declares an unsupported ``required`` value. + SchemaError: If a rule has an invalid requirement or severity class. """ with SCHEMA_FILE.open("r") as f: schema = yaml.safe_load(f) or {} - rules = list(schema.get("attributes", {}).items()) - rules += [(r.get("name"), r) for r in schema.get("checks", [])] - rules += [(r.get("key"), r) for r in schema.get("arrays", [])] - for name, rule in rules: + waivable_names = set() + for default_name, rule in _schema_rules(schema): required = rule.get("required", False) - if required not in _REQUIRED_VALUES: + if not isinstance(required, bool): raise SchemaError( - f"{name!r}: `required` must be one of " - f"{list(_REQUIRED_VALUES)}, got {required!r}" + f"{default_name!r}: `required` must be true or false, got {required!r}" + ) + severity = rule.get("severity") + if severity not in _SEVERITY_VALUES: + raise SchemaError( + f"{default_name!r}: `severity` must be one of " + f"{list(_SEVERITY_VALUES)}, got {severity!r}" + ) + if severity == _INTEGRITY: + continue + if not required: + raise SchemaError( + f"{default_name!r}: a waivable rule must be required by default" + ) + name = _rule_name(default_name, rule) + if not _RULE_NAME.fullmatch(name): + raise SchemaError( + f"{default_name!r}: rule name {name!r} must use snake_case" + ) + if name in waivable_names: + raise SchemaError(f"duplicate waivable rule name {name!r}") + waivable_names.add(name) + if not isinstance(rule.get("why"), str) or not rule["why"].strip(): + raise SchemaError( + f"{default_name!r}: a waivable rule needs a non-empty `why`" ) return schema -def _level(required, strict: bool) -> str | None: - """Return a requirement failure's severity, or ``None`` when optional.""" - if required is True: +def _schema_rules(schema: Mapping): + """Yield ``(default_name, rule)`` pairs in schema order.""" + yield from schema.get("attributes", {}).items() + yield from ((rule.get("name"), rule) for rule in schema.get("checks", [])) + yield from ((rule.get("key"), rule) for rule in schema.get("arrays", [])) + + +def _rule_name(default_name: str, rule: Mapping) -> str: + """Return the stable name used by requirement overrides and CLI flags.""" + return rule.get("name", default_name) + + +def waivable_rules(schema: Mapping | None = None) -> dict[str, dict]: + """Return the schema rules whose failure severity callers may lower. + + The schema owns the rule names and explanations so the validator cannot add + a waivable rule without also exposing a documented CLI flag for it. + """ + schema = load_schema() if schema is None else schema + return { + _rule_name(default_name, rule): rule + for default_name, rule in _schema_rules(schema) + if rule.get("severity") != _INTEGRITY + } + + +def _requirements( + schema: Mapping, overrides: Mapping[str, bool] | None +) -> dict[str, bool]: + """Build one explicit error-or-warning decision per waivable rule.""" + requirements = dict.fromkeys(waivable_rules(schema), True) + for name, required in (overrides or {}).items(): + if name not in requirements: + raise SchemaError( + f"unknown waivable rule {name!r}; expected one of " + f"{sorted(requirements)}" + ) + if not isinstance(required, bool): + raise SchemaError( + f"requirement for {name!r} must be a bool, got {required!r}" + ) + requirements[name] = required + return requirements + + +def _level(rule: Mapping, report: Report, default_name: str) -> str | None: + """Return a missing or failed requirement's invocation-specific level.""" + if not rule.get("required", False): + return None + if rule["severity"] == _INTEGRITY: return ERROR - if required == "strict": - return ERROR if strict else WARNING - return None + name = _rule_name(default_name, rule) + return ERROR if report.requirements[name] else WARNING # --------------------------------------------------------------------------- @@ -196,7 +266,7 @@ def _check_type(value, type_name: str | None) -> str | None: def _check_attribute(name: str, rule: dict, attrs: Mapping, report, context) -> None: if name not in attrs: - level = _level(rule.get("required", False), report.strict) + level = _level(rule, report, name) if level is not None: report.add(level, f"attrs.{name}", "missing") return @@ -276,7 +346,7 @@ def _check_schema_version(value, attrs, context) -> str | None: def _check_calibration_present(rule, context, report) -> None: calibration = context.get("calibration") - level = _level(rule.get("required", False), report.strict) + level = _level(rule, report, "calibration_present") if calibration is None or not calibration.intrinsics(): if level is not None: report.add( @@ -296,7 +366,7 @@ def _check_calibration_present(rule, context, report) -> None: def _check_camera_coverage(rule, context, report) -> None: missing = uncalibrated_cameras(context["array_keys"], context.get("calibration")) - level = _level(rule.get("required", False), report.strict) + level = _level(rule, report, "camera_coverage") if missing: if level is not None: report.add( @@ -328,7 +398,7 @@ def _report_problems(rule, report, check: str, problems, passed: str) -> None: if not problems: report.add(OK, check, passed) return - level = _level(rule.get("required", False), report.strict) + level = _level(rule, report, check) if level is not None: report.add(level, check, "; ".join(problems)) @@ -611,7 +681,7 @@ def _dtype_kind(dtype) -> str: def _check_array(key: str, rule: dict, arrays: Mapping, report, context, side) -> None: if key not in arrays: - level = _level(rule.get("required", False), report.strict) + level = _level(rule, report, key) if level is not None: report.add(level, key, "missing") return @@ -683,28 +753,32 @@ def _expand_key(template: str, arrays: Mapping, sides) -> list[tuple[str, str | # --------------------------------------------------------------------------- -def validate_episode(path: str | Path, *, strict: bool = False) -> Report: +def validate_episode( + path: str | Path, *, requirements: Mapping[str, bool] | None = None +) -> Report: """Validate one Zarr episode against ``schema/episode_v3.yaml``. + Args: path: The episode ``.zarr`` directory. - strict: If true, treat ``required: strict`` failures as errors rather - than warnings. + requirements: Per-rule severity decisions. ``True`` reports a failure + as an error; ``False`` waives it to a warning. Omitted rules remain + required. Integrity rules cannot be overridden. Returns: The findings emitted while validating the episode. Raises: - SchemaError: If the schema contains an unsupported declaration. + SchemaError: If the schema or a requirement override is invalid. """ path = Path(path) - report = Report(path=path, strict=strict) + schema = load_schema() + report = Report(path=path, requirements=_requirements(schema, requirements)) try: store = zarr.open_group(str(path), mode="r") except Exception as exc: report.add(ERROR, "episode", f"cannot open as a zarr group: {exc}") return report - schema = load_schema() attrs = dict(store.attrs) arrays = {name: store[name] for name in store.array_keys()} @@ -771,22 +845,26 @@ def _resolve(attrs: Mapping, report: Report) -> ResolvedEmbodiment | None: return None -def main(argv: list[str] | None = None) -> int: - """Validate CLI paths and return an exit status. - - Returns: - ``0`` if every report has no errors, otherwise ``1``. - """ +def _build_parser() -> argparse.ArgumentParser: + """Build the CLI parser from the schema's waivable rules.""" parser = argparse.ArgumentParser( prog="python -m egomimic.rldb.zarr.validate", description=__doc__.splitlines()[0], ) parser.add_argument("paths", nargs="+", type=Path, help="episode .zarr paths") - parser.add_argument( - "--strict", - action="store_true", - help="promote the rules the corpus does not meet yet to errors", - ) + for name, rule in waivable_rules().items(): + flag = name.replace("_", "-") + # argparse performs %-interpolation on help text. Schema rationales are + # ordinary prose and may contain percentages such as the 90% coverage + # threshold, so escape them before handing the text to argparse. + help_text = rule["why"].replace("%", "%%") + parser.add_argument( + f"--{flag}", + dest=f"require_{name}", + action=argparse.BooleanOptionalAction, + default=True, + help=f"[{rule['severity']}] {help_text} (default: required)", + ) parser.add_argument( "-v", "--verbose", @@ -796,9 +874,19 @@ def main(argv: list[str] | None = None) -> int: parser.add_argument( "--json", action="store_true", help="print one JSON report per episode" ) - args = parser.parse_args(argv) + return parser + + +def main(argv: list[str] | None = None) -> int: + """Validate CLI paths and return an exit status. - reports = [validate_episode(p, strict=args.strict) for p in args.paths] + Returns: + ``0`` if every report has no errors, otherwise ``1``. + """ + parser = _build_parser() + args = parser.parse_args(argv) + requirements = {name: getattr(args, f"require_{name}") for name in waivable_rules()} + reports = [validate_episode(p, requirements=requirements) for p in args.paths] if args.json: print(json.dumps([r.to_jsonable() for r in reports], indent=2)) else: diff --git a/egomimic/rldb/zarr/zarr_writer.py b/egomimic/rldb/zarr/zarr_writer.py index a59eb2b75..58b80f079 100644 --- a/egomimic/rldb/zarr/zarr_writer.py +++ b/egomimic/rldb/zarr/zarr_writer.py @@ -327,7 +327,7 @@ def __init__( extrinsics: dict | None = None, calibration: Calibration | dict | None = None, data_status: str = DATA_STATUS_COMPLETE, - strict: bool = False, + require_camera_coverage: bool = False, verbose: bool = False, ): """Configure a writer for one Zarr v3 episode. @@ -349,8 +349,9 @@ def __init__( data_status: Initial episode status. The staging path and dataset resolvers accept only ``complete`` and exclude ``structural_sample``. - strict: If true, reject an image stream with no matching camera - matrix; otherwise log a warning. + require_camera_coverage: If true, reject an image stream with no + matching camera matrix; otherwise log a warning. Coverage is + named explicitly because this option changes no other check. verbose: If true, print progress information during writes. """ self.episode_path = Path(episode_path) @@ -373,7 +374,7 @@ def __init__( f"{data_status!r}" ) self.data_status = data_status - self.strict = strict + self.require_camera_coverage = require_camera_coverage self.verbose = verbose # Track image shapes for metadata self._features: dict[str, dict[str, Any]] = {} @@ -741,7 +742,8 @@ def _check_camera_coverage(self, image_keys) -> None: image_keys: The episode's image array keys. Raises: - ValueError: If ``strict`` is set and a stream carries no ``K``. + ValueError: If camera coverage is required and a stream carries no + ``K``. """ calibration = self.calibration or lift_legacy_calibration( self.intrinsics, self.extrinsics @@ -754,9 +756,11 @@ def _check_camera_coverage(self, image_keys) -> None: f"{missing}. Every `images.` array needs a K in " "`calibration.cameras`. See CONTRIBUTING_DATA.md §6.4." ) - if self.strict: + if self.require_camera_coverage: raise ValueError(message) - logger.warning("%s Pass strict=True to make this an error.", message) + logger.warning( + "%s Pass require_camera_coverage=True to make this an error.", message + ) def _build_metadata( self, metadata_override: dict[str, Any] | None = None @@ -818,7 +822,7 @@ def create_and_write( extrinsics: dict | None = None, calibration: Calibration | dict | None = None, data_status: str = DATA_STATUS_COMPLETE, - strict: bool = False, + require_camera_coverage: bool = False, metadata_override: dict[str, Any] | None = None, ) -> Path: """Validate episode metadata and write one Zarr v3 episode. @@ -846,8 +850,9 @@ def create_and_write( ``intrinsics`` or ``extrinsics`` value for older readers. data_status: ``complete`` or ``structural_sample``. The staging path and dataset resolvers accept only ``complete`` episodes. - strict: If true, reject an image stream with no matching camera - matrix; otherwise log a warning. + require_camera_coverage: If true, reject an image stream with no + matching camera matrix; otherwise log a warning. Coverage is + named explicitly because this option changes no other check. metadata_override: Additional episode metadata. This mapping cannot replace ``calibration``, ``intrinsics``, or ``extrinsics``. @@ -861,8 +866,8 @@ def create_and_write( ValueError: If ``extrinsics`` is not ``None`` or a non-empty mapping of 4×4 matrices. ValueError: If ``data_status`` is not a supported value. - ValueError: If ``strict`` is set and an image stream carries no - camera matrix. + ValueError: If camera coverage is required and an image stream + carries no camera matrix. CalibrationError: If ``calibration`` is malformed. """ # Validate the embodiment up front (it is guaranteed to be consumed by @@ -936,7 +941,7 @@ def create_and_write( extrinsics=extrinsics, calibration=calibration, data_status=data_status, - strict=strict, + require_camera_coverage=require_camera_coverage, ) writer.write(