From 8deb7587c243aac0027f4ba2df1284a38ce1735f Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Tue, 20 Oct 2026 20:12:04 -0400 Subject: [PATCH] rldb/lance: steps/eva schema and robot reads steps/eva stores each arm's measured and commanded EE pose, joints and gripper under the zarr key names, plus left_valid/right_valid for forward-filled dropout frames. The loader maps zarr keys to columns per lab, treats a validity flag the table lacks (robots have no head) as always true, and knows the eva_* arm sets. tag, compact and cleanup now cover every lab's steps table. Claude-Session: https://claude.ai/code/session_01Kt1E8T2x8kdV9akAvbJyVd --- egomimic/rldb/lance/dataset.py | 40 +++++++++++--- egomimic/rldb/lance/schema.py | 25 +++++++++ tests/fixtures/lance_tables.py | 89 +++++++++++++++++++++++++------- tests/unit/test_lance_commit.py | 16 +++++- tests/unit/test_lance_dataset.py | 66 +++++++++++++++++++++++ 5 files changed, 211 insertions(+), 25 deletions(-) diff --git a/egomimic/rldb/lance/dataset.py b/egomimic/rldb/lance/dataset.py index fb812426e..16c3b26e5 100644 --- a/egomimic/rldb/lance/dataset.py +++ b/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,28 @@ ) }, } +# 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 steps_columns(end_effector: str) -> dict[str, str]: + """{zarr key: steps column} for a recording with this end effector.""" + 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"] = {} @@ -277,14 +294,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 @@ -366,8 +385,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, from `steps_columns`; human columns by default. + self.columns = columns or ZARR_TO_LANCE self.content_digest = content_digest self.steps_uri = steps_uri self.steps_version = steps_version @@ -945,6 +967,7 @@ def resolve( segments=sorted(passing), action_clamp=self.action_clamp, history_floor=self.history_floor, + columns=steps_columns(row.end_effector), hand_valid=np.stack( [ valid[h] if h in hands else np.ones(n, bool) @@ -985,7 +1008,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)) @@ -995,8 +1019,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) diff --git a/egomimic/rldb/lance/schema.py b/egomimic/rldb/lance/schema.py index e7575692c..f704765ce 100644 --- a/egomimic/rldb/lance/schema.py +++ b/egomimic/rldb/lance/schema.py @@ -102,6 +102,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()), diff --git a/tests/fixtures/lance_tables.py b/tests/fixtures/lance_tables.py index dbdc9a99c..aeeff0b41 100644 --- a/tests/fixtures/lance_tables.py +++ b/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", diff --git a/tests/unit/test_lance_commit.py b/tests/unit/test_lance_commit.py index 83297ce3a..265e7adee 100644 --- a/tests/unit/test_lance_commit.py +++ b/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 PERIOD_NS, steps_table # noqa: E402 +from fixtures.lance_tables import PERIOD_NS, eva_steps_table, steps_table # noqa: E402 from egomimic.rldb.lance import commit # noqa: E402 from egomimic.rldb.lance.dataset import LanceEpisodeResolver # noqa: E402 @@ -291,3 +292,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() diff --git a/tests/unit/test_lance_dataset.py b/tests/unit/test_lance_dataset.py index f15253af8..c4802a5ad 100644 --- a/tests/unit/test_lance_dataset.py +++ b/tests/unit/test_lance_dataset.py @@ -521,3 +521,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")