Fix local LeRobot sample boundaries

cosmos-framework/cosmos_framework/data/vfm/action/_lerobot_local.py

@@ -196,22 +196,95 @@ def load_local_lerobot_collection(root: str | Path, subroots: list[str] | None =
     return roots, episodes, tasks, rows
 
 
+def _sample_stride_value(sample_stride: int) -> int:
+    stride = int(sample_stride)
+    if stride < 1:
+        raise ValueError(f"sample_stride must be >= 1, got {stride}")
+    return stride
+
+
+def _episode_key(row: dict[str, Any]) -> tuple[int, int]:
+    return int(row.get("_dataset_idx", 0)), int(row["episode_index"])
+
+
+def contiguous_sample_starts(rows: list[dict[str, Any]], chunk_length: int, sample_stride: int = 1) -> list[int]:
+    """Start row indices whose T+1 contiguous window stays inside one episode."""
+
+    stride = _sample_stride_value(sample_stride)
+    span = int(chunk_length) + 1
+    if span < 1:
+        raise ValueError(f"chunk_length must be >= 0, got {chunk_length}")
+
+    starts: list[int] = []
+    for row_idx in range(0, len(rows) - span + 1, stride):
+        selected = rows[row_idx : row_idx + span]
+        key = _episode_key(selected[0])
+        if all(_episode_key(row) == key for row in selected):
+            starts.append(row_idx)
+    return starts
+
+
+def fps_sample_starts(rows: list[dict[str, Any]], chunk_length: int, fps: float, sample_stride: int = 1) -> list[int]:
+    """Start row indices whose timestamp-sampled T+1 window fits in one episode."""
+
+    stride = _sample_stride_value(sample_stride)
+    fps_value = float(fps)
+    if fps_value <= 0:
+        raise ValueError(f"fps must be > 0, got {fps_value}")
+
+    episode_end_ts: dict[tuple[int, int], float] = {}
+    for row in rows:
+        episode_end_ts[_episode_key(row)] = float(row["timestamp"])
+
+    starts: list[int] = []
+    duration_s = float(chunk_length) / fps_value
+    for row_idx in range(0, len(rows), stride):
+        row = rows[row_idx]
+        if float(row["timestamp"]) + duration_s <= episode_end_ts[_episode_key(row)] + 1e-9:
+            starts.append(row_idx)
+    return starts
+
+
 def sample_count(rows: list[dict[str, Any]], chunk_length: int, sample_stride: int = 1) -> int:
     """Number of contiguous local samples that can provide T+1 observations."""
 
-    return max(0, (len(rows) - int(chunk_length) + int(sample_stride) - 1) // int(sample_stride))
+    return len(contiguous_sample_starts(rows, chunk_length, sample_stride))
 
 
-def rows_for_index(rows: list[dict[str, Any]], idx: int, chunk_length: int, sample_stride: int = 1) -> list[dict[str, Any]]:
+def _row_idx_for_sample(
+    rows: list[dict[str, Any]],
+    idx: int,
+    sample_stride: int,
+    sample_starts: list[int] | None,
+) -> int:
+    sample_idx = int(idx)
+    if sample_idx < 0:
+        raise IndexError(f"Index {idx} is negative")
+    if sample_starts is not None:
+        if sample_idx >= len(sample_starts):
+            raise IndexError(f"Index {idx} is out of range for {len(sample_starts)} samples")
+        return int(sample_starts[sample_idx])
+    row_idx = sample_idx * _sample_stride_value(sample_stride)
+    if row_idx >= len(rows):
+        raise IndexError(f"Index {idx} is out of range for {len(rows)} rows")
+    return row_idx
+
+
+def rows_for_index(
+    rows: list[dict[str, Any]],
+    idx: int,
+    chunk_length: int,
+    sample_stride: int = 1,
+    sample_starts: list[int] | None = None,
+) -> list[dict[str, Any]]:
     """Select T+1 contiguous rows for a flat local index."""
 
-    row_idx = int(idx) * int(sample_stride)
+    row_idx = _row_idx_for_sample(rows, idx, sample_stride, sample_starts)
     selected = rows[row_idx : row_idx + int(chunk_length) + 1]
     if len(selected) < int(chunk_length) + 1:
         raise IndexError(f"Index {idx} does not have {chunk_length + 1} rows")
-    episode_id = int(selected[0]["episode_index"])
-    dataset_idx = int(selected[0].get("_dataset_idx", 0))
-    if any(int(row["episode_index"]) != episode_id or int(row.get("_dataset_idx", 0)) != dataset_idx for row in selected):
+    episode_key = _episode_key(selected[0])
+    if any(_episode_key(row) != episode_key for row in selected):
         raise IndexError(f"Index {idx} crosses an episode or shard boundary")
     return selected
 
@@ -222,17 +295,13 @@ def rows_at_fps(
     chunk_length: int,
     fps: float,
     sample_stride: int = 1,
+    sample_starts: list[int] | None = None,
 ) -> list[dict[str, Any]]:
     """Select T+1 nearest timestamp rows at target FPS within one episode."""
 
-    start_row = rows[int(start_idx) * int(sample_stride)]
-    episode_id = int(start_row["episode_index"])
-    dataset_idx = int(start_row.get("_dataset_idx", 0))
-    episode_rows = [
-        row
-        for row in rows
-        if int(row["episode_index"]) == episode_id and int(row.get("_dataset_idx", 0)) == dataset_idx
-    ]
+    start_row = rows[_row_idx_for_sample(rows, start_idx, sample_stride, sample_starts)]
+    episode_key = _episode_key(start_row)
+    episode_rows = [row for row in rows if _episode_key(row) == episode_key]
     timestamps = np.asarray([float(row["timestamp"]) for row in episode_rows], dtype=np.float64)
     start_ts = float(start_row["timestamp"])
     target_ts = start_ts + np.arange(int(chunk_length) + 1, dtype=np.float64) / float(fps)

cosmos-framework/cosmos_framework/data/vfm/action/_lerobot_local_test.py

@@ -0,0 +1,54 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+import unittest
+
+from cosmos_framework.data.vfm.action._lerobot_local import (
+    contiguous_sample_starts,
+    fps_sample_starts,
+    rows_at_fps,
+    rows_for_index,
+    sample_count,
+)
+
+
+def _rows(lengths: list[int], *, fps: float = 10.0) -> list[dict]:
+    rows = []
+    index = 0
+    for episode_index, length in enumerate(lengths):
+        for frame in range(length):
+            rows.append(
+                {
+                    "index": index,
+                    "episode_index": episode_index,
+                    "timestamp": frame / fps,
+                }
+            )
+            index += 1
+    return rows
+
+
+class LocalLeRobotSampleIndexTest(unittest.TestCase):
+    def test_contiguous_samples_skip_episode_boundaries(self) -> None:
+        rows = _rows([20, 20])
+        starts = contiguous_sample_starts(rows, chunk_length=16)
+
+        self.assertEqual(starts, [0, 1, 2, 3, 20, 21, 22, 23])
+        self.assertEqual(sample_count(rows, chunk_length=16), len(starts))
+
+        selected = rows_for_index(rows, 4, chunk_length=16, sample_starts=starts)
+        self.assertEqual({row["episode_index"] for row in selected}, {1})
+
+        with self.assertRaisesRegex(IndexError, "crosses an episode"):
+            rows_for_index(rows, 12, chunk_length=16)
+
+    def test_fps_samples_skip_episode_tails(self) -> None:
+        rows = _rows([20], fps=10.0)
+        starts = fps_sample_starts(rows, chunk_length=16, fps=10.0)
+
+        self.assertEqual(starts, [0, 1, 2, 3])
+        selected = rows_at_fps(rows, 3, chunk_length=16, fps=10.0, sample_starts=starts)
+
+        self.assertEqual([row["index"] for row in selected], list(range(3, 20)))
+        with self.assertRaisesRegex(IndexError, "out of range"):
+            rows_at_fps(rows, 4, chunk_length=16, fps=10.0, sample_starts=starts)

cosmos-framework/cosmos_framework/data/vfm/action/bridge_orig_lerobot_dataset.py

@@ -14,11 +14,11 @@ from torch.utils.data import Dataset
 from cosmos_framework.data.vfm.action._lerobot_local import (
     build_result,
     choose_mode,
+    contiguous_sample_starts,
     load_local_lerobot,
     load_video_key,
     pick_caption,
     rows_for_index,
-    sample_count,
     task_text,
 )
 from cosmos_framework.data.vfm.action.pose_utils import (
@@ -102,6 +102,7 @@ class BridgeOrigLeRobotDataset(Dataset):
         self._tolerance_s = float(tolerance_s)
         self._viewpoint = viewpoint
         self._root, self._info, self._episodes, self._tasks, self._rows = load_local_lerobot(root)
+        self._sample_starts = contiguous_sample_starts(self._rows, self._chunk_length, self._sample_stride)
 
     @property
     def fps(self) -> float:
@@ -123,9 +124,12 @@ class BridgeOrigLeRobotDataset(Dataset):
     def action_dim(self) -> int:
         return 10
 
+    def __len__(self) -> int:
+        return len(self._sample_starts)
+
     def __getitem__(self, idx: int) -> dict[str, Any]:
         mode = choose_mode(self._mode)
-        observation_rows = rows_for_index(self._rows, idx, self._chunk_length, self._sample_stride)
+        observation_rows = rows_for_index(self._rows, idx, self._chunk_length, self._sample_stride, self._sample_starts)
         episode = self._episodes[int(observation_rows[0]["episode_index"])]
         action_rows = observation_rows[: self._chunk_length]
 
@@ -180,6 +184,3 @@ class BridgeOrigLeRobotDataset(Dataset):
             viewpoint=self._viewpoint,
             **extras,
         )
-
-    def __len__(self) -> int:
-        return sample_count(self._rows, self._chunk_length, self._sample_stride)

cosmos-framework/cosmos_framework/data/vfm/action/droid_lerobot_dataset.py

@@ -15,11 +15,11 @@ from torch.utils.data import Dataset
 from cosmos_framework.data.vfm.action._lerobot_local import (
     build_result,
     choose_mode,
+    contiguous_sample_starts,
     load_local_lerobot,
     load_video_key,
     pick_caption,
     rows_for_index,
-    sample_count,
     task_text,
 )
 from cosmos_framework.data.vfm.action.pose_utils import (
@@ -91,6 +91,7 @@ class DROIDLeRobotDataset(Dataset):
         if not (root_path / "meta" / "info.json").exists() and (root_path / "success" / "meta" / "info.json").exists():
             root_path = root_path / "success"
         self._root, self._info, self._episodes, self._tasks, self._rows = load_local_lerobot(root_path)
+        self._sample_starts = contiguous_sample_starts(self._rows, self._chunk_length, self._sample_stride)
 
     @property
     def fps(self) -> float:
@@ -114,7 +115,7 @@ class DROIDLeRobotDataset(Dataset):
 
     def __getitem__(self, idx: int) -> dict[str, Any]:
         mode = choose_mode(self._mode)
-        observation_rows = rows_for_index(self._rows, idx, self._chunk_length, self._sample_stride)
+        observation_rows = rows_for_index(self._rows, idx, self._chunk_length, self._sample_stride, self._sample_starts)
         episode = self._episodes[int(observation_rows[0]["episode_index"])]
         action_rows = observation_rows[: self._chunk_length]
 
@@ -180,4 +181,4 @@ class DROIDLeRobotDataset(Dataset):
         )
 
     def __len__(self) -> int:
-        return sample_count(self._rows, self._chunk_length, self._sample_stride)
+        return len(self._sample_starts)

cosmos-framework/cosmos_framework/data/vfm/action/fractal.py

@@ -14,11 +14,11 @@ from torch.utils.data import Dataset
 from cosmos_framework.data.vfm.action._lerobot_local import (
     build_result,
     choose_mode,
+    contiguous_sample_starts,
     load_local_lerobot,
     load_video_key,
     pick_caption,
     rows_for_index,
-    sample_count,
     task_text,
 )
 from cosmos_framework.data.vfm.action.pose_utils import PoseConvention, build_abs_pose_from_components, pose_abs_to_rel
@@ -64,6 +64,7 @@ class FractalLeRobotDataset(Dataset):
         self._viewpoint = viewpoint
         self._sample_stride = int(sample_stride)
         self._root, self._info, self._episodes, self._tasks, self._rows = load_local_lerobot(root)
+        self._sample_starts = contiguous_sample_starts(self._rows, self._chunk_length, self._sample_stride)
 
     @property
     def fps(self) -> float:
@@ -86,10 +87,10 @@ class FractalLeRobotDataset(Dataset):
         return 10
 
     def __len__(self) -> int:
-        return sample_count(self._rows, self._chunk_length, self._sample_stride)
+        return len(self._sample_starts)
 
     def __getitem__(self, idx: int) -> dict[str, Any]:
-        rows = rows_for_index(self._rows, idx, self._chunk_length, self._sample_stride)
+        rows = rows_for_index(self._rows, idx, self._chunk_length, self._sample_stride, self._sample_starts)
         episode = self._episodes[int(rows[0]["episode_index"])]
         video = load_video_key(self._root, self._info, episode, rows, _IMAGE_FEATURE, tolerance_s=1e-4)
 

cosmos-framework/cosmos_framework/data/vfm/action/robomind_franka_dataset.py

@@ -15,11 +15,11 @@ from torch.utils.data import Dataset
 from cosmos_framework.data.vfm.action._lerobot_local import (
     build_result,
     choose_mode,
+    fps_sample_starts,
     load_local_lerobot,
     load_video_key,
     pick_caption,
     rows_at_fps,
-    sample_count,
     task_text,
 )
 from cosmos_framework.data.vfm.action.pose_utils import PoseConvention, build_abs_pose_from_components, pose_abs_to_rel
@@ -68,6 +68,7 @@ class RoboMINDFrankaDataset(Dataset):
         self._viewpoint = viewpoint
         self._sample_stride = int(sample_stride)
         self._root, self._info, self._episodes, self._tasks, self._rows = load_local_lerobot(root)
+        self._sample_starts = fps_sample_starts(self._rows, self._chunk_length, self._fps, self._sample_stride)
         self._to_opencv = _ROBOMIND_FRANKA_TO_OPENCV[:3, :3]
 
     @property
@@ -91,7 +92,7 @@ class RoboMINDFrankaDataset(Dataset):
         return 10 if self._embodiment_type == "robomind-franka" else 20
 
     def __len__(self) -> int:
-        return sample_count(self._rows, self._chunk_length, self._sample_stride)
+        return len(self._sample_starts)
 
     def _camera_keys(self) -> tuple[str, str, str]:
         primary = "observation.images.camera_top" if self._embodiment_type == "robomind-franka" else "observation.images.camera_front"
@@ -143,7 +144,7 @@ class RoboMINDFrankaDataset(Dataset):
         return torch.from_numpy(action).float(), initial_pose_left, initial_pose_right
 
     def __getitem__(self, idx: int) -> dict[str, Any]:
-        rows = rows_at_fps(self._rows, idx, self._chunk_length, self._fps, self._sample_stride)
+        rows = rows_at_fps(self._rows, idx, self._chunk_length, self._fps, self._sample_stride, self._sample_starts)
         episode = self._episodes[int(rows[0]["episode_index"])]
         video = self._load_concat_video(episode, rows)
         built = self._build_action(rows)

cosmos-framework/cosmos_framework/data/vfm/action/umi_lerobot_dataset.py

@@ -14,11 +14,11 @@ from torch.utils.data import Dataset
 from cosmos_framework.data.vfm.action._lerobot_local import (
     build_result,
     choose_mode,
+    contiguous_sample_starts,
     load_local_lerobot,
     load_video_key,
     pick_caption,
     rows_for_index,
-    sample_count,
     task_text,
 )
 from cosmos_framework.data.vfm.action.pose_utils import PoseConvention, build_abs_pose_from_components, pose_abs_to_rel
@@ -55,6 +55,7 @@ class UMIFastLeRobotDataset(Dataset):
         self._viewpoint = viewpoint
         self._sample_stride = int(sample_stride)
         self._root, self._info, self._episodes, self._tasks, self._rows = load_local_lerobot(root)
+        self._sample_starts = contiguous_sample_starts(self._rows, self._chunk_length, self._sample_stride)
 
         # The packaged UMI trajectory is the right main camera trajectory.  For
         # viewer overlays we want the actual gripper/end-effector pose, so
@@ -93,10 +94,10 @@ class UMIFastLeRobotDataset(Dataset):
         return 10
 
     def __len__(self) -> int:
-        return sample_count(self._rows, self._chunk_length, self._sample_stride)
+        return len(self._sample_starts)
 
     def __getitem__(self, idx: int) -> dict[str, Any]:
-        rows = rows_for_index(self._rows, idx, self._chunk_length, self._sample_stride)
+        rows = rows_for_index(self._rows, idx, self._chunk_length, self._sample_stride, self._sample_starts)
         episode = self._episodes[int(rows[0]["episode_index"])]
         video = load_video_key(self._root, self._info, episode, rows, _IMAGE_FEATURE, tolerance_s=1e-4)
         pose = np.asarray([row[_POSE_FEATURE] for row in rows], dtype=np.float32)