diff --git a/egomimic/rldb/lance/dataset.py b/egomimic/rldb/lance/dataset.py index db60cccff..fe7186a51 100644 --- a/egomimic/rldb/lance/dataset.py +++ b/egomimic/rldb/lance/dataset.py @@ -137,30 +137,62 @@ def _runs(ids: pa.ChunkedArray) -> tuple[list, np.ndarray]: return values, np.asarray(starts, dtype=np.int64) +REQUIRE_HANDS = ("none", "either", "both", "left", "right") + + +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, history=() + valid: dict[str, np.ndarray], + hands, + horizon: int, + require: str = "both", + history=(), ) -> np.ndarray: - """(K, 2) half-open frame intervals of anchors whose action chunk is valid for every hand. - - 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. Head and hands must also be valid - at t - o for each `history` offset o that lands inside the recording: invalid rows - hold filler poses, not NaN, so nothing downstream would reject them. + """(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"). The head pose must be valid at t whatever + `require` is, because actions are expressed in the anchor's head frame. Head and the + required hands must also be valid at t - o for each `history` offset o inside the + recording: invalid rows hold filler poses, not NaN, so nothing downstream would + reject them. """ n = len(valid["head"]) + if require == "none": + hands = () + elif 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)) + if oks: + ok &= ( + np.logical_or.reduce(oks) + if require == "either" + else np.logical_and.reduce(oks) + ) frame_ok = valid["head"].copy() - for h in hands: - frame_ok &= valid[h] + if hands: + frames = [valid[h] for h in hands] + frame_ok &= ( + np.logical_or.reduce(frames) + if require == "either" + else np.logical_and.reduce(frames) + ) for o in sorted({int(o) for o in history if 0 < o < n}): ok[o:] &= frame_ok[:-o] - edges = np.flatnonzero(np.diff(np.r_[0, ok.astype(np.int8), 0])) - return edges.reshape(-1, 2) + return _intervals(ok) def _mask(spans, n: int) -> np.ndarray: @@ -333,6 +365,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 @@ -366,6 +399,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: @@ -386,6 +426,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() @@ -435,6 +476,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": @@ -573,6 +625,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 always the head pose valid at the anchor. 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. @@ -604,7 +663,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"): @@ -618,6 +682,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) @@ -793,11 +858,13 @@ 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")), + hands, horizon, + self.require_hands_in_frame, history=self._history_offsets(float(row.fps)), ), n, @@ -857,6 +924,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 e1bd3d4be..8894095fa 100644 --- a/egomimic/rldb/leaf.py +++ b/egomimic/rldb/leaf.py @@ -86,6 +86,20 @@ 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(): + if spec.get("key_type") == "camera_keys": + continue + if int(spec.get("history") or 1) > 1: + return spec + return None + def _fps(self) -> float: """The episode's frame rate (zarr ``fps`` attr; 30 when absent).""" return float(self.metadata.get("fps") or 30) @@ -195,6 +209,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)) + if self.transform: for transform in self.transform or []: data = transform.transform(data) diff --git a/tests/unit/test_lance_dataset.py b/tests/unit/test_lance_dataset.py index 5ce07fcb6..6a00ad93f 100644 --- a/tests/unit/test_lance_dataset.py +++ b/tests/unit/test_lance_dataset.py @@ -438,3 +438,55 @@ def test_history_floor_at_the_task_start_when_asked(segmented): assert s["obs_head_pose"][:, 0].tolist() == [6, 6, 7] with pytest.raises(ValueError, match="history_floor"): LanceEpisodeResolver(root, key_map=KEY_MAP, 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_always_needs_the_head(): + valid = _hands() + valid["head"][7] = False + for require in ("either", "none"): + assert eligible_intervals(valid, ("left", "right"), 2, require)[-1, 1] == 7 + + +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 e4279f7c7..b19879301 100644 --- a/tests/unit/test_leaf_equivalence.py +++ b/tests/unit/test_leaf_equivalence.py @@ -71,6 +71,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):