From 310ebb0f43726e4f3120b1f49bc06168b15449e7 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Tue, 20 Oct 2026 10:10:45 -0400 Subject: [PATCH] rldb/lance: action_clamp and history_floor at segment bounds Action chunks end at the anchor's segment end by default (repeat-last padding, flagged in action_pad_mask); action_clamp: recording runs them on. History reaches back into earlier segments by default; history_floor: task pads and masks it at the segment start, through a new EpisodeLeaf._history_floor hook. Claude-Session: https://claude.ai/code/session_01Kt1E8T2x8kdV9akAvbJyVd --- egomimic/rldb/lance/dataset.py | 56 +++++++++++++++++++++++++++++--- egomimic/rldb/leaf.py | 21 +++++++----- tests/unit/test_lance_dataset.py | 42 ++++++++++++++++++++++++ 3 files changed, 107 insertions(+), 12 deletions(-) diff --git a/egomimic/rldb/lance/dataset.py b/egomimic/rldb/lance/dataset.py index 7a0b943f8..db60cccff 100644 --- a/egomimic/rldb/lance/dataset.py +++ b/egomimic/rldb/lance/dataset.py @@ -330,6 +330,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 @@ -355,6 +358,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: @@ -373,6 +384,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() @@ -411,9 +424,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 @@ -532,9 +560,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 @@ -565,7 +602,13 @@ 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 @@ -573,6 +616,8 @@ def __init__( 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) @@ -809,6 +854,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: diff --git a/egomimic/rldb/leaf.py b/egomimic/rldb/leaf.py index 29bcf1a66..e1bd3d4be 100644 --- a/egomimic/rldb/leaf.py +++ b/egomimic/rldb/leaf.py @@ -94,11 +94,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 @@ -113,17 +117,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 diff --git a/tests/unit/test_lance_dataset.py b/tests/unit/test_lance_dataset.py index 23263fc8b..5ce07fcb6 100644 --- a/tests/unit/test_lance_dataset.py +++ b/tests/unit/test_lance_dataset.py @@ -396,3 +396,45 @@ 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] + 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] + with pytest.raises(ValueError, match="history_floor"): + LanceEpisodeResolver(root, key_map=KEY_MAP, history_floor="episode")