From afc17c518714147ec30fc2999f2dd72f34bc23b3 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Tue, 20 Oct 2026 10:18:27 -0400 Subject: [PATCH] rldb/lance: require_hands_in_frame and per-hand sample masks require_hands_in_frame (none | either | both | left | right, default both) picks the anchors: the named hands of the embodiment valid over the whole action chunk, plus the head pose at the anchor unless none. Samples carry action_hand_mask ([left, right] valid over the real chunk) and proprio_hand_valid ([left, right] per proprio frame) for the losses, stats and bounds check. Claude-Session: https://claude.ai/code/session_01Kt1E8T2x8kdV9akAvbJyVd --- egomimic/rldb/lance/dataset.py | 92 ++++++++++++++++++++++++----- egomimic/rldb/leaf.py | 7 +++ tests/unit/test_lance_dataset.py | 52 ++++++++++++++++ tests/unit/test_leaf_equivalence.py | 4 ++ 4 files changed, 141 insertions(+), 14 deletions(-) diff --git a/egomimic/rldb/lance/dataset.py b/egomimic/rldb/lance/dataset.py index d9d9fb01e..9ea873bff 100644 --- a/egomimic/rldb/lance/dataset.py +++ b/egomimic/rldb/lance/dataset.py @@ -137,21 +137,46 @@ def _runs(ids: pa.ChunkedArray) -> tuple[list, np.ndarray]: return values, np.asarray(starts, dtype=np.int64) -def eligible_intervals(valid: dict[str, np.ndarray], hands, horizon: int) -> np.ndarray: - """(K, 2) half-open frame intervals of anchors whose action chunk is valid for every hand. +REQUIRE_HANDS = ("none", "either", "both", "left", "right") - The chunk is [t, min(t + horizon, n)); the head pose must be valid at t because - actions are expressed in the anchor's head frame. + +def _chunk_valid(v: np.ndarray, horizon: int) -> np.ndarray: + """Per anchor t: `v` holds over the whole chunk [t, min(t + horizon, n)).""" + n = len(v) + c = np.r_[0, np.cumsum(v.astype(np.int64))] + end = np.minimum(np.arange(n) + horizon, n) + return (c[end] - c[:n]) == (end - np.arange(n)) + + +def eligible_intervals( + valid: dict[str, np.ndarray], hands, horizon: int, require: str = "both" +) -> np.ndarray: + """(K, 2) half-open frame intervals of the anchors `require` admits. + + `hands` are the embodiment's hands. `require` names which of them must be valid + over the action chunk [t, min(t + horizon, n)): all ("both"), any ("either"), one + ("left" / "right"), or none ("none", which also skips the head check). Otherwise + the head pose must be valid at t, because actions are expressed in the anchor's + head frame. """ n = len(valid["head"]) + if require == "none": + return np.array([[0, n]] if n else [], np.int64).reshape(-1, 2) + if require in ("left", "right"): + if require not in hands: + raise ValueError( + f"require_hands_in_frame {require!r}: the embodiment has hands {tuple(hands)}" + ) + hands = (require,) + oks = [_chunk_valid(valid[h], horizon) for h in hands] ok = valid["head"].copy() - for h in hands: - v = valid[h].astype(np.int64) - c = np.r_[0, np.cumsum(v)] - end = np.minimum(np.arange(n) + horizon, n) - ok &= (c[end] - c[:n]) == (end - np.arange(n)) - edges = np.flatnonzero(np.diff(np.r_[0, ok.astype(np.int8), 0])) - return edges.reshape(-1, 2) + if oks: + ok &= ( + np.logical_or.reduce(oks) + if require == "either" + else np.logical_and.reduce(oks) + ) + return _intervals(ok) def _mask(spans, n: int) -> np.ndarray: @@ -324,6 +349,7 @@ def __init__( segments: np.ndarray | None = None, action_clamp: str = "task", history_floor: str = "recording", + hand_valid: np.ndarray | None = None, ): self.episode_path = name self.content_digest = content_digest @@ -357,6 +383,13 @@ def __init__( ) self.action_clamp = action_clamp self.history_floor = history_floor + # (num_frames, 2) [left, right] tracked per frame; a hand the embodiment + # lacks is always True so it never masks anything. + self.hand_valid = ( + np.asarray(hand_valid, bool) + if hand_valid is not None + else np.ones((num_frames, 2), bool) + ) super().__init__(key_map=key_map, transform_list=transform_list) def cache_fingerprint(self) -> str: @@ -377,6 +410,7 @@ def cache_fingerprint(self) -> str: h.update(self.anchors.tobytes()) h.update(f"{self.action_clamp}\0{self.history_floor}\0".encode()) h.update(self.segments.tobytes()) + h.update(b"hands\0" + np.packbits(self.hand_valid).tobytes()) if self._clamp_actions and self._ann is not None: h.update(b"clamp\0" + self._ann[0].tobytes() + self._ann[1].tobytes()) return h.hexdigest() @@ -422,6 +456,17 @@ def _segment(self, frame: int) -> tuple[int, int]: def _history_floor(self, idx: int) -> int: return self._segment(idx)[0] if self.history_floor == "task" else 0 + def _hand_masks(self, idx: int) -> dict: + out = {} + horizon = self._action_horizon() + if horizon is not None: + end = self._chunk_end_idx(idx, horizon, "action_keys") + out["action_hand_mask"] = self.hand_valid[idx:end].all(0).astype(np.float32) + spec = self._proprio_history_spec() + frames = self._history_indices(idx, spec)[0] if spec is not None else idx + out["proprio_hand_valid"] = self.hand_valid[frames].astype(np.float32) + return out + def _chunk_end_idx(self, start_idx: int, horizon: int, key_type: str | None) -> int: end = super()._chunk_end_idx(start_idx, horizon, key_type) if key_type != "action_keys": @@ -560,6 +605,13 @@ class LanceEpisodeResolver: `history_floor` "recording" lets history reach into earlier segments; "task" pads and masks it at the segment start. + `require_hands_in_frame` (none | either | both | left | right, default both) picks + the anchors: the named hands of the embodiment valid over the whole action chunk, + and the head pose valid at the anchor unless "none". It only gates sampling. Each + sample carries `action_hand_mask` ([left, right] valid over the chunk) and + `proprio_hand_valid` ([left, right] per proprio frame), which the losses, norm + stats and bounds check use to leave out an untracked hand. + Each table is read at `versions[name]` (a number or tag), else at `tag`, else at its latest version. Tags are resolved to numbers once, here, so a tag moved later does not change what an existing resolver reads. @@ -591,7 +643,12 @@ def __init__( include_taskless_frames: bool = False, action_clamp: str = "task", history_floor: str = "recording", + require_hands_in_frame: str = "both", ): + if require_hands_in_frame not in REQUIRE_HANDS: + raise ValueError( + f"require_hands_in_frame {require_hands_in_frame!r}: one of {REQUIRE_HANDS}" + ) if action_clamp not in ("task", "recording"): raise ValueError(f"action_clamp {action_clamp!r}: 'task' or 'recording'") if history_floor not in ("task", "recording"): @@ -605,6 +662,7 @@ def __init__( self.include_taskless_frames = include_taskless_frames self.action_clamp = action_clamp self.history_floor = history_floor + self.require_hands_in_frame = require_hands_in_frame self.annotations = None if annotations is not None: unknown = set(annotations) - set(self.ANNOTATION_DEFAULTS) @@ -760,10 +818,9 @@ def resolve( if filters.episode_hashes: mask &= _mask([clock.frames(r) for r in pins[rid]], n) mask &= ~_mask([clock.frames(r) for r in excludes.get(rid, [])], n) + hands = HANDS_OF.get(row.embodiment, ("left", "right")) eligible = _mask( - eligible_intervals( - valid, HANDS_OF.get(row.embodiment, ("left", "right")), horizon - ), + eligible_intervals(valid, hands, horizon, self.require_hands_in_frame), n, ) anchors = _intervals(mask & eligible) @@ -820,6 +877,13 @@ def resolve( segments=sorted(passing), action_clamp=self.action_clamp, history_floor=self.history_floor, + hand_valid=np.stack( + [ + valid[h] if h in hands else np.ones(n, bool) + for h in ("left", "right") + ], + axis=1, + ), **ann_kw, ) if no_anchor: diff --git a/egomimic/rldb/leaf.py b/egomimic/rldb/leaf.py index da47cd955..e43b84aba 100644 --- a/egomimic/rldb/leaf.py +++ b/egomimic/rldb/leaf.py @@ -86,6 +86,11 @@ def _action_pad_mask(self, idx: int) -> np.ndarray | None: mask[: max(n_real, 0)] = 1.0 return mask + def _hand_masks(self, idx: int) -> dict: + """Per-hand validity keys for sample ``idx`` (``action_hand_mask``, + ``proprio_hand_valid``); none when the store tracks no per-hand validity.""" + return {} + def _proprio_history_spec(self) -> dict | None: """The first non-camera keymap entry with ``history`` > 1, else None.""" for spec in self.key_map.values(): @@ -220,6 +225,8 @@ def _finish(self, idx: int, data: dict) -> dict: # attention, so each image window carries its own mask. data[f"{k}{HISTORY_MASK_SUFFIX}"] = self._history_indices(idx, spec)[1] + data.update(self._hand_masks(idx)) + history_mask = self._proprio_history_mask(idx) if history_mask is not None: # Informational: the front padding already looks exactly like a diff --git a/tests/unit/test_lance_dataset.py b/tests/unit/test_lance_dataset.py index 4d71f3948..a1b452282 100644 --- a/tests/unit/test_lance_dataset.py +++ b/tests/unit/test_lance_dataset.py @@ -352,3 +352,55 @@ def test_history_floor_at_the_task_start_when_asked(segmented): assert s["proprio_history_mask"].tolist() == [0, 1, 1] with pytest.raises(ValueError, match="history_floor"): LanceEpisodeResolver(root, history_floor="episode") + + +def _hands(n=8): + valid = {h: np.ones(n, bool) for h in ("head", "left", "right")} + valid["left"][2] = False + valid["right"][5] = False + return valid + + +@pytest.mark.parametrize( + "require, want", + [ + ("both", [[0, 1], [3, 4], [6, 8]]), + ("either", [[0, 8]]), + ("left", [[0, 1], [3, 8]]), + ("right", [[0, 4], [6, 8]]), + ("none", [[0, 8]]), + ], +) +def test_require_hands_in_frame_rules(require, want): + iv = eligible_intervals(_hands(), ("left", "right"), 2, require) + np.testing.assert_array_equal(iv, want) + + +def test_require_hands_needs_the_head_unless_none(): + valid = _hands() + valid["head"][7] = False + assert eligible_intervals(valid, ("left", "right"), 2, "either")[-1, 1] == 7 + assert eligible_intervals(valid, ("left", "right"), 2, "none")[-1, 1] == 8 + + +def test_require_a_hand_the_embodiment_lacks_is_an_error(): + with pytest.raises(ValueError, match="left"): + eligible_intervals(_hands(), ("right",), 2, "left") + + +def test_samples_carry_per_hand_masks(catalog): + root, recs = catalog + rid = recs[0][0] + both = _resolve(root)[rid] + s = both[0] + assert s["action_hand_mask"].tolist() == [1, 1] + assert s["proprio_hand_valid"].tolist() == [1, 1] + either = _resolve(root, require_hands_in_frame="either")[rid] + assert _frames(either) == list(range(recs[0][2])) + # left is untracked at frame 4: in the chunk of anchor 3, at the anchor 4. + assert _at(either, 3)["action_hand_mask"].tolist() == [0, 1] + assert _at(either, 3)["proprio_hand_valid"].tolist() == [1, 1] + assert _at(either, 4)["proprio_hand_valid"].tolist() == [0, 1] + assert _at(either, 5)["action_hand_mask"].tolist() == [1, 1] + with pytest.raises(ValueError, match="require_hands_in_frame"): + _resolve(root, require_hands_in_frame="at_least_one") diff --git a/tests/unit/test_leaf_equivalence.py b/tests/unit/test_leaf_equivalence.py index 21e225bff..845b53909 100644 --- a/tests/unit/test_leaf_equivalence.py +++ b/tests/unit/test_leaf_equivalence.py @@ -65,6 +65,10 @@ def test_zarr_and_lance_leaves_agree(tmp_path): assert isinstance(leaf, EpisodeLeaf) and isinstance(twin, EpisodeLeaf) for i in range(len(leaf)): a, b = leaf[i], twin[leaf.anchor(i)] + # Only Lance tracks per-hand validity; every action chunk sampled here + # has both hands (history can reach back to an untracked frame). + assert bool((a.pop("action_hand_mask") == 1).all()) + a.pop("proprio_hand_valid") assert a.keys() == b.keys() for k in a: if isinstance(a[k], torch.Tensor):