Fix split-local LeRobot frame offsets - #19
Merged
Conversation
There was a problem hiding this comment.
Pull request overview
Fixes incorrect frame slicing for non-train LeRobot splits by consistently translating global episode frame row intervals into the split-local Hugging Face dataset coordinate system (using frame_dataset_offset) before performing grouped/batched reads.
Changes:
- Introduce a shared
_local_frame_interval(...)helper to convert and bounds-check global frame intervals into split-local dataset indices. - Apply the shared conversion to single-episode reads, grouped kinematics materialization, and exhaustive-window statistics batching.
- Add regression tests covering nonzero split offsets for validation materialization and exhaustive training stats.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
src/omg/data/lerobot_dataset.py |
Adds shared global→local frame interval conversion and applies it to grouped/batched readers to prevent empty reads on val/test splits. |
tests/data/test_lerobot_dataset.py |
Adds fixtures + regressions to ensure nonzero split offsets are handled correctly in grouped reads and stats batching. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
367
to
+371
| def _read_episode(self, sample: dict[str, Any]) -> dict[str, np.ndarray]: | ||
| episode_index = int(sample["episode_index"]) | ||
| if self._cached_episode_index == episode_index and self._cached_episode_data is not None: | ||
| return self._cached_episode_data | ||
| start = int(sample["data_start_row"]) - self.frame_dataset_offset | ||
| end = int(sample["data_end_row"]) - self.frame_dataset_offset | ||
| if start < 0 or end > len(self.frame_dataset) or end <= start: | ||
| raise IndexError( | ||
| "Episode frame interval is outside the loaded LeRobot split shards: " | ||
| f"episode={episode_index} local={start}:{end} loaded=0:{len(self.frame_dataset)}" | ||
| ) | ||
| start, end = self._local_frame_interval(sample["data_start_row"], sample["data_end_row"]) |
Comment on lines
+266
to
+267
| assert summary["episodes"] == 1 | ||
| assert summary["frames"] == 5 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Root cause
The release loader selects only parquet shards intersecting the requested split. Their Hugging Face dataset therefore uses local row coordinates and records the selected files' global starting row in
frame_dataset_offset.The ordinary
__getitem__path subtracted this offset, butiter_episode_kinematics_groupsanditer_stats_batchessliced the local dataset with global episode rows. This was invisible for the train split, which starts at row zero, and produced an emptyqpos_36tensor immediately for validation/test splits.Validation
git diff --checkpassed