rldb/lance: gripper steps schema and robot reads

egomimic/rldb/lance/dataset.py

@@ -25,7 +25,7 @@
 import torch
 
 from egomimic.rldb.filters import DatasetFilter, PinError, _ensure_dataset_filter
-from egomimic.rldb.lance.schema import STEPS_KEY, steps_path, task_instance_id
+from egomimic.rldb.lance.schema import ARM, STEPS_KEY, steps_path, task_instance_id
 from egomimic.rldb.leaf import EpisodeLeaf
 
 logger = logging.getLogger(__name__)
@@ -45,11 +45,27 @@
         )
     },
 }
+# Robot steps columns already carry the zarr key names.
+COLUMNS_OF_END_EFFECTOR = {
+    "parallel_jaw_gripper": {
+        f"{s}.{f}": f"{s}.{f}" for s in ("left", "right") for f in ARM.names
+    }
+}
+
+
+def zarr_to_lance(end_effector: str) -> dict[str, str]:
+    return COLUMNS_OF_END_EFFECTOR.get(end_effector, ZARR_TO_LANCE)
+
+
+# A flag missing from a steps table (robots have no head) is True everywhere.
 VALIDITY = {"left": "left_valid", "right": "right_valid", "head": "head_pose_valid"}
 HANDS_OF = {
     "human_bimanual": ("left", "right"),
     "human_left_arm": ("left",),
     "human_right_arm": ("right",),
+    "eva_bimanual": ("left", "right"),
+    "eva_left_arm": ("left",),
+    "eva_right_arm": ("right",),
 }
 
 _TABLES: dict[tuple[str, int], "lance.LanceDataset"] = {}
@@ -261,14 +277,16 @@ def _fetch_numeric(items: list) -> list[dict]:
                 ]
             )
         )
-        cols = sorted({ZARR_TO_LANCE[zk] for j in js for zk in items[j][1]["needs"]})
+        cols = sorted(
+            {items[j][0].columns[zk] for j in js for zk in items[j][1]["needs"]}
+        )
         tbl = _table(uri, version).take(rows.tolist(), columns=cols)
         arrays = {c: _column(tbl, c) for c in cols}
         for j in js:
             leaf, plan = items[j]
             for zk, frames in plan["needs"].items():
                 at = np.searchsorted(rows, frames + leaf.row_start)
-                out[j][zk] = arrays[ZARR_TO_LANCE[zk]][at]
+                out[j][zk] = arrays[leaf.columns[zk]][at]
     return out
 
 
@@ -350,8 +368,11 @@ def __init__(
         action_clamp: str = "task",
         history_floor: str = "recording",
         hand_valid: np.ndarray | None = None,
+        columns: dict[str, str] | None = None,
     ):
         self.episode_path = name
+        # zarr key -> steps column; the steps table's lab decides it.
+        self.columns = columns or ZARR_TO_LANCE
         self.content_digest = content_digest
         self.steps_uri = steps_uri
         self.steps_version = steps_version
@@ -890,6 +911,7 @@ def resolve(
                 segments=sorted(passing),
                 action_clamp=self.action_clamp,
                 history_floor=self.history_floor,
+                columns=zarr_to_lance(row.end_effector),
                 hand_valid=np.stack(
                     [
                         valid[h] if h in hands else np.ones(n, bool)
@@ -930,7 +952,8 @@ def _steps_rows(self, recs: pd.DataFrame) -> dict:
         out = {}
         for key, group in recs.groupby(list(STEPS_KEY)):
             uri, version, steps = self._open(steps_path(*key))
-            cols = steps.to_table(columns=["recording_id", *VALIDITY.values()])
+            flags = {h: c for h, c in VALIDITY.items() if c in steps.schema.names}
+            cols = steps.to_table(columns=["recording_id", *flags.values()])
             n_rows = cols.num_rows
             first, change = _runs(cols.column("recording_id"))
             starts = dict(zip(first, change))
@@ -940,8 +963,12 @@ def _steps_rows(self, recs: pd.DataFrame) -> dict:
                 )
             counts = dict(zip(first, np.diff(np.r_[change, n_rows])))
             all_valid = {
-                h: cols.column(c).to_numpy(zero_copy_only=False)
-                for h, c in VALIDITY.items()
+                h: (
+                    cols.column(flags[h]).to_numpy(zero_copy_only=False)
+                    if h in flags
+                    else np.ones(n_rows, bool)
+                )
+                for h in VALIDITY
             }
             for row in group.itertuples():
                 rid, n = row.recording_id, int(row.num_frames)

egomimic/rldb/lance/schema.py

@@ -99,6 +99,31 @@ def _vec(n: int) -> pa.DataType:
     ]
 )
 
+# Robot arm, robot base frame: poses xyz (m) + quaternion wxyz, joints in rad, gripper in [0, 1].
+ARM = pa.struct(
+    [
+        ("obs_ee_pose", _vec(7)),
+        ("obs_joints", _vec(6)),
+        ("obs_gripper", pa.float32()),
+        ("cmd_ee_pose", _vec(7)),
+        ("cmd_joints", _vec(6)),
+        ("cmd_gripper", pa.float32()),
+    ]
+)
+
+# An arm the robot lacks is zeros with its flag False (Arrow's parquet reader rejects
+# null fixed-size lists); an all-zero pose frame is forward-filled and flagged False.
+EVA_STEPS = pa.schema(
+    [
+        ("recording_id", pa.string()),
+        ("timestamp_ns", pa.int64()),
+        ("left", ARM),
+        ("right", ARM),
+        ("left_valid", pa.bool_()),
+        ("right_valid", pa.bool_()),
+    ]
+)
+
 RECORDINGS = pa.schema(
     [
         ("recording_id", pa.string()),

tests/fixtures/lance_tables.py

@@ -7,6 +7,8 @@
 from egomimic.rldb.lance.schema import (
     ALIASES,
     ANNOTATIONS,
+    ARM,
+    EVA_STEPS,
     HUMAN_STEPS,
     RECORDINGS,
     TASKS,
@@ -76,6 +78,42 @@ def steps_table(rid, n, offset, valid=None):
     )
 
 
+def eva_steps_table(rid, n, offset, arms=("left", "right"), valid=None):
+    """Robot rows: obs_ee_pose x = frame + offset (+0.5 left, +0.25 right), cmd x one
+    frame ahead; an arm not in `arms` is zeros and flagged False."""
+    t = np.arange(n, dtype=np.float32)
+
+    def arm(side, s):
+        on = float(side in arms)
+        pose = lambda x: vec(  # noqa: E731
+            on * np.c_[x, np.zeros((n, 3)), np.ones(n), np.zeros((n, 2))], 7
+        )
+        return pa.StructArray.from_arrays(
+            [
+                pose(t + offset + s),
+                vec(np.zeros((n, 6)), 6),
+                pa.array(np.full(n, 0.5 * on, np.float32)),
+                pose(t + offset + s + 1),
+                vec(np.zeros((n, 6)), 6),
+                pa.array((on * t / n).astype(np.float32)),
+            ],
+            fields=list(ARM),
+        )
+
+    flags = {f"{a}_valid": np.full(n, a in arms) for a in ("left", "right")}
+    flags.update(valid or {})
+    return pa.table(
+        {
+            "recording_id": pa.array([rid] * n),
+            "timestamp_ns": pa.array(np.arange(n, dtype=np.int64) * PERIOD_NS),
+            "left": arm("left", 0.5),
+            "right": arm("right", 0.25),
+            **{k: pa.array(v) for k, v in flags.items()},
+        },
+        schema=EVA_STEPS,
+    )
+
+
 def task(rid, f0, f1, name="fold_clothes", description="folding clothes"):
     """A TASKS row over frames [f0, f1) of a `steps_table` recording."""
     t0 = f0 * PERIOD_NS
@@ -105,26 +143,41 @@ def write_catalog(
     tasks=lambda recs: [task(rid, 0, n) for rid, _, n in recs],
     aliases=lambda recs: [alias(rid, 0, n, h) for rid, h, n in recs],
     valid=lambda rid, n: {},
+    source="aria",
 ):
-    """Steps, recordings, tasks, aliases and annotations tables for two aria recordings
-    under ``tmp_path``; returns their sorted ``(recording_id, episode_hash, num_frames)``.
-    ``annotations(recs)``, ``tasks(recs)`` and ``aliases(recs)`` give those rows (by
-    default one task and one alias per whole recording); ``valid(rid, n)`` overrides
-    validity columns of a recording's steps."""
+    """Steps, recordings, tasks, aliases and annotations tables for two recordings of
+    lab rl2 from `source` under ``tmp_path``; returns their sorted ``(recording_id,
+    episode_hash, num_frames)``. ``annotations(recs)``, ``tasks(recs)`` and
+    ``aliases(recs)`` give those rows (by default one task and one alias per whole
+    recording); ``valid(rid, n)`` overrides validity columns of a recording's steps.
+    For "eva" (grippers) the first recording is bimanual and the second left-arm only,
+    each in its own steps table; aria ones (hands) are bimanual."""
     recs = [
-        (recording_id("aria", "1"), "2025-01-01-00-00-00-000000", 12),
-        (recording_id("aria", "2"), "2025-01-01-00-00-01-000000", 8),
+        (recording_id(source, "1"), "2025-01-01-00-00-00-000000", 12),
+        (recording_id(source, "2"), "2025-01-01-00-00-01-000000", 8),
     ]
+    embodiment = {recs[0][0]: "human_bimanual", recs[1][0]: "human_bimanual"}
+    end_effector = "human_hand"
+    if source == "eva":
+        end_effector = "parallel_jaw_gripper"
+        embodiment = {recs[0][0]: "eva_bimanual", recs[1][0]: "eva_left_arm"}
     recs.sort()
-    steps = pa.concat_tables(
-        [
-            steps_table(rid, n, 100 * i, valid(rid, n))
-            for i, (rid, _, n) in enumerate(recs)
-        ]
-    )
-    lance.write_dataset(
-        steps, tmp_path / steps_path("rl2", "human_bimanual", "human_hand")
-    )
+
+    def rows(i, rid, n):
+        if source != "eva":
+            return steps_table(rid, n, 100 * i, valid(rid, n))
+        arms = ("left", "right") if embodiment[rid] == "eva_bimanual" else ("left",)
+        return eva_steps_table(rid, n, 100 * i, arms, valid(rid, n))
+
+    for emb in sorted(set(embodiment.values())):
+        steps = pa.concat_tables(
+            [
+                rows(i, rid, n)
+                for i, (rid, _, n) in enumerate(recs)
+                if embodiment[rid] == emb
+            ]
+        )
+        lance.write_dataset(steps, tmp_path / steps_path("rl2", emb, end_effector))
     k = [200.0, 0, 240, 0, 200.0, 180, 0, 0, 1]
     lance.write_dataset(
         pa.Table.from_pylist(
@@ -133,8 +186,8 @@ def write_catalog(
                     recording_id=rid,
                     vendor_recording_id=h,
                     lab="rl2",
-                    embodiment="human_bimanual",
-                    end_effector="human_hand",
+                    embodiment=embodiment[rid],
+                    end_effector=end_effector,
                     rig_name="aria_gen1",
                     operator="op",
                     scene="1",

tests/unit/test_lance_commit.py

@@ -1,3 +1,4 @@
+import dataclasses
 import datetime
 
 import numpy as np
@@ -6,7 +7,7 @@
 
 lance = pytest.importorskip("lance")
 
-from fixtures.lance_tables import steps_table  # noqa: E402
+from fixtures.lance_tables import eva_steps_table, steps_table  # noqa: E402
 
 from egomimic.rldb.lance import commit  # noqa: E402
 from egomimic.rldb.lance.dataset import LanceEpisodeResolver  # noqa: E402
@@ -229,3 +230,16 @@ def test_drop_task_annotations_brings_old_tables_to_the_schema(tmp_path):
     assert lance.dataset(tasks).schema.names == TASKS.names
     assert commit.append(root, [second]) == [second.recording_id]
     commit.drop_task_annotations(root)
+
+
+def test_tag_covers_every_steps_table(tmp_path):
+    root = tmp_path / "lance"
+    eva = _delivery(
+        tmp_path, 1, 5, embodiment="eva_bimanual", end_effector="parallel_jaw_gripper"
+    )
+    eva = dataclasses.replace(eva, steps=eva_steps_table(eva.recording_id, 5, 0))
+    commit.append(root, [_delivery(tmp_path, 0, 10), eva])
+    eva_steps = steps_path("rl2", "eva_bimanual", "parallel_jaw_gripper")
+    tagged = commit.tag(root, "t1")
+    assert {STEPS, eva_steps} <= set(tagged)
+    assert "t1" in lance.dataset(root / eva_steps).tags.list()

tests/unit/test_lance_dataset.py

@@ -435,3 +435,69 @@ def test_norm_stats_leave_out_untracked_hand_chunks(catalog, tmp_path, episode_c
         np.mean(tracked, axis=0),
         rtol=1e-5,
     )
+
+
+EVA_KEY_MAP = {
+    "left.cmd_ee_pose": {
+        "key_type": "action_keys",
+        "zarr_key": "left.cmd_ee_pose",
+        "horizon": 3,
+    },
+    "left.cmd_gripper": {
+        "key_type": "action_keys",
+        "zarr_key": "left.cmd_gripper",
+        "horizon": 3,
+    },
+    "right.obs_ee_pose": {"key_type": "proprio_keys", "zarr_key": "right.obs_ee_pose"},
+    "left.obs_gripper": {"key_type": "proprio_keys", "zarr_key": "left.obs_gripper"},
+}
+
+
+def _eva(tmp_path, valid=lambda rid, n: {}):
+    recs = write_catalog(tmp_path, source="eva", valid=valid)
+    (a, _, _), (b, _, _) = _by_length(recs)
+    return tmp_path, a, b
+
+
+def _resolve_eva(root, key_map=EVA_KEY_MAP, **kw):
+    return LanceEpisodeResolver(root, key_map=key_map, cameras={}, **kw).resolve()
+
+
+def test_eva_reads_robot_columns_by_their_zarr_keys(tmp_path):
+    root, a, b = _eva(tmp_path)
+    leaf = _resolve_eva(root)[a]
+    assert _frames(leaf) == list(range(12))
+    s = _at(leaf, 2)
+    offset = 100 * sorted([a, b]).index(a)
+    np.testing.assert_allclose(
+        s["left.cmd_ee_pose"][:, 0].numpy(), offset + 0.5 + np.arange(3, 6)
+    )
+    assert s["left.cmd_gripper"].shape == (3, 1)
+    np.testing.assert_allclose(
+        s["left.cmd_gripper"][:, 0].numpy(), np.arange(2, 5) / 12
+    )
+    assert float(s["right.obs_ee_pose"][0]) == offset + 2.25
+    assert float(s["left.obs_gripper"][0]) == 0.5
+    assert s["action_hand_mask"].tolist() == [1, 1]
+
+
+def test_eva_arm_dropouts_gate_anchors_and_mask_the_arm(tmp_path):
+    def dropout(rid, n):
+        v = np.ones(n, bool)
+        v[4] = False
+        return {"right_valid": v} if n == 12 else {}
+
+    root, a, _ = _eva(tmp_path, valid=dropout)
+    assert _frames(_resolve_eva(root)[a]) == [0, 1, *range(5, 12)]
+    either = _resolve_eva(root, require_hands_in_frame="either")[a]
+    assert _at(either, 3)["action_hand_mask"].tolist() == [1, 0]
+
+
+def test_eva_single_arm_robot_never_masks_the_missing_arm(tmp_path):
+    root, _, b = _eva(tmp_path)
+    key_map = {k: v for k, v in EVA_KEY_MAP.items() if not k.startswith("right.")}
+    leaf = _resolve_eva(root, key_map=key_map)[b]
+    assert _frames(leaf) == list(range(8))
+    assert leaf[0]["action_hand_mask"].tolist() == [1, 1]
+    with pytest.raises(ValueError, match="require_hands_in_frame"):
+        _resolve_eva(root, key_map=key_map, require_hands_in_frame="right")