rldb/lance: action_clamp and history_floor at segment bounds

egomimic/rldb/lance/dataset.py

@@ -321,6 +321,9 @@ def __init__(
         content_digest: str | None = None,
         annotations: tuple[np.ndarray, np.ndarray, list[str]] | None = None,
         clamp_actions: bool = False,
+        segments: np.ndarray | None = None,
+        action_clamp: str = "task",
+        history_floor: str = "recording",
     ):
         self.episode_path = name
         self.content_digest = content_digest
@@ -346,6 +349,14 @@ def __init__(
         # annotations rows under the resolver's `annotations` config.
         self._ann = annotations
         self._clamp_actions = clamp_actions
+        # (S, 2) frame spans of the segments sampled, sorted by start.
+        self.segments = (
+            np.asarray(segments, np.int64).reshape(-1, 2)
+            if segments is not None
+            else np.array([[0, num_frames]], np.int64)
+        )
+        self.action_clamp = action_clamp
+        self.history_floor = history_floor
         super().__init__(key_map=key_map, transform_list=transform_list)
 
     def cache_fingerprint(self) -> str:
@@ -364,6 +375,8 @@ def cache_fingerprint(self) -> str:
             for k in sorted(self.videos):
                 h.update(f"{k}\0{self.videos[k]}\0".encode())
         h.update(self.anchors.tobytes())
+        h.update(f"{self.action_clamp}\0{self.history_floor}\0".encode())
+        h.update(self.segments.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()
@@ -398,9 +411,24 @@ def _annotation_text_for_frame(self, frame_idx: int):
             return []
         return [self._ann[2][j] for j in self._covering(frame_idx)]
 
+    def _segment(self, frame: int) -> tuple[int, int]:
+        """The innermost sampled segment holding `frame` (the recording when none does)."""
+        s, e = self.segments[:, 0], self.segments[:, 1]
+        on = (s <= frame) & (frame < e)
+        if not on.any():
+            return 0, self.total_frames
+        return int(s[on].max()), int(e[on].min())
+
+    def _history_floor(self, idx: int) -> int:
+        return self._segment(idx)[0] if self.history_floor == "task" else 0
+
     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 not self._clamp_actions or key_type != "action_keys" or self._ann is None:
+        if key_type != "action_keys":
+            return end
+        if self.action_clamp == "task":
+            end = min(end, self._segment(start_idx)[1])
+        if not self._clamp_actions or self._ann is None:
             return end
         cover = self._covering(start_idx)
         return min(end, int(self._ann[1][cover].min())) if len(cover) else end
@@ -519,9 +547,18 @@ class StepsRows(NamedTuple):
 class LanceEpisodeResolver:
     """Resolves recordings from local Lance tables into LanceRecordingDataset leaves.
 
-    Filter lambdas see each recording joined with its task, with SQL-style field names
-    (`episode_hash`, `lab`, `task`, ...), so existing filter configs apply. Recordings
-    are read from the steps table of their (lab, embodiment, end_effector).
+    Leaves are keyed by recording_id. Filter lambdas run once per segment of a
+    recording: each tasks row, plus each untasked gap (`task` None) when
+    `include_taskless_frames`. They see the recording joined with the segment, with
+    SQL-style field names (`episode_hash` is the recording_id, `lab`, `task`,
+    `task_instance_id`, ...), so existing filter configs apply. Pins and excludes name
+    a recording_id or an old clip episode_hash, meaning that clip's time range. Each
+    recording is read from the steps table of its (lab, embodiment, end_effector).
+
+    `action_clamp` "task" ends action chunks at the end of the anchor's segment
+    (repeat-last padding, flagged in `action_pad_mask`); "recording" runs them on.
+    `history_floor` "recording" lets history reach into earlier segments; "task" pads
+    and masks it at the segment start.
 
     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
@@ -552,14 +589,22 @@ def __init__(
         tag: str | None = None,
         annotations: dict | None = None,
         include_taskless_frames: bool = False,
+        action_clamp: str = "task",
+        history_floor: str = "recording",
     ):
+        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"):
+            raise ValueError(f"history_floor {history_floor!r}: 'task' or 'recording'")
         self.root = str(lance_root).rstrip("/")
         self.key_map = key_map
         self.transform_list = transform_list
         self.cameras = cameras or {"images.front_1": "front_1"}
         self.tag = tag
         self.versions = dict(versions or {})
         self.include_taskless_frames = include_taskless_frames
+        self.action_clamp = action_clamp
+        self.history_floor = history_floor
         self.annotations = None
         if annotations is not None:
             unknown = set(annotations) - set(self.ANNOTATION_DEFAULTS)
@@ -772,6 +817,9 @@ def resolve(
                 key_map=self.key_map,
                 transform_list=self.transform_list,
                 content_digest=getattr(row, "content_digest", None),
+                segments=sorted(passing),
+                action_clamp=self.action_clamp,
+                history_floor=self.history_floor,
                 **ann_kw,
             )
         if no_anchor:

egomimic/rldb/leaf.py

@@ -111,11 +111,15 @@ def _seconds_to_frames(self, seconds: float) -> int:
         """``seconds`` in this episode's fps, at least one frame."""
         return max(1, round(float(seconds) * self._fps()))
 
+    def _history_floor(self, idx: int) -> int:
+        """First frame history and lagged reads for ``idx`` may reach back to."""
+        return 0
+
     def _lagged_idx(self, idx: int, lag_s: float) -> int:
         """Frame ``lag_s`` seconds before ``idx`` (at least one frame back),
-        clamped at the episode start -- so the first samples read frame 0,
-        the same front padding a history window gets."""
-        return max(0, idx - self._seconds_to_frames(lag_s))
+        clamped at the history floor -- so the first samples read the floor
+        frame, the same front padding a history window gets."""
+        return max(self._history_floor(idx), idx - self._seconds_to_frames(lag_s))
 
     def _history_stride(self, spec: dict) -> int:
         """Frames between history steps: ``history_stride_s`` (seconds, so a
@@ -130,17 +134,18 @@ def _history_indices(self, idx: int, spec: dict) -> tuple[np.ndarray, np.ndarray
         ``(K,)`` float32 mask.
 
         The newest step is ``idx`` or, with ``lag_s``, that many seconds before
-        it. Steps before the episode start are masked 0 and repeat the oldest
-        real step (frame 0 when none is real), so the newest step is always
-        ``[-1]``.
+        it. Steps before the history floor (the episode start) are masked 0
+        and repeat the oldest real step (the floor frame when none is real), so
+        the newest step is always ``[-1]``.
         """
         history = int(spec["history"])
         newest = idx
         if spec.get("lag_s") is not None:
             newest = idx - self._seconds_to_frames(spec["lag_s"])
         steps = newest - self._history_stride(spec) * np.arange(history - 1, -1, -1)
-        real = steps >= 0
-        fill = steps[real][0] if real.any() else 0
+        floor = self._history_floor(idx)
+        real = steps >= floor
+        fill = steps[real][0] if real.any() else floor
         return np.where(real, steps, fill), real.astype(np.float32)
 
     @staticmethod

tests/unit/test_lance_dataset.py

@@ -308,3 +308,47 @@ def test_clip_pins_through_multidataset(segmented):
     assert _frames(video.datasets["clip_b2"]) == [4, 5, 6, 7]
     with pytest.raises(PinError):
         pinned_episode_subset(whole, ["nope"])
+
+
+HISTORY_KEY_MAP = {
+    **KEY_MAP,
+    "obs_head_pose": {
+        "key_type": "proprio_keys",
+        "zarr_key": "obs_head_pose",
+        "history": 3,
+    },
+}
+
+
+def _at(leaf, frame):
+    return leaf[_frames(leaf).index(frame)]
+
+
+def test_action_chunks_stop_at_the_task_end_unless_asked(segmented):
+    root, a, _ = segmented
+    leaf = _resolve(root)[a]
+    assert _at(leaf, 2)["action_pad_mask"].tolist() == [1, 1, 0]
+    head = _at(leaf, 2)["left.action_ee_pose"][:, 0].tolist()
+    assert head == [2.5, 3.5, 3.5]
+    through = _resolve(root, action_clamp="recording")[a]
+    assert _at(through, 2)["action_pad_mask"].tolist() == [1, 1, 1]
+    gaps = _resolve(root, include_taskless_frames=True)[a]
+    assert _at(gaps, 4)["action_pad_mask"].tolist() == [1, 1, 0]
+
+
+def test_history_floor_at_the_task_start_when_asked(segmented):
+    root, a, _ = segmented
+    reach_back = LanceEpisodeResolver(
+        root, key_map=HISTORY_KEY_MAP, cameras={}
+    ).resolve()[a]
+    s = _at(reach_back, 7)
+    assert s["obs_head_pose"][:, 0].tolist() == [5, 6, 7]
+    assert s["proprio_history_mask"].tolist() == [1, 1, 1]
+    floored = LanceEpisodeResolver(
+        root, key_map=HISTORY_KEY_MAP, cameras={}, history_floor="task"
+    ).resolve()[a]
+    s = _at(floored, 7)
+    assert s["obs_head_pose"][:, 0].tolist() == [6, 6, 7]
+    assert s["proprio_history_mask"].tolist() == [0, 1, 1]
+    with pytest.raises(ValueError, match="history_floor"):
+        LanceEpisodeResolver(root, history_floor="episode")