rldb/lance: require_hands_in_frame and per-hand sample masks

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:

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

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")

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):