Success only filter for lerobot dataset conversion.#22
Open
k1000dai wants to merge 11 commits into
Open
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
Adds options and logic to support filtering to successful episodes during LeRobot v2.1 conversion, while maintaining LeRobot’s requirement for contiguous episode_index/task_index and enabling safe re-runs into an existing output directory.
Changes:
- Add
success_onlyfiltering during conversion and remap episode/task indices to contiguous ranges in the output. - Add
overwritesupport to clear prior LeRobot outputs safely when re-running conversion. - Adjust train split computation to avoid empty train splits on very small (filtered) datasets.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
tests/test_lerobot_v21.py |
Updates existing conversion fixture to use overwrite, and adds a new test covering success_only conversion output. |
src/openarm_dataset/lerobot_v21.py |
Implements success_only filtering, contiguous episode/task remapping, overwrite clearing logic, and updates split computation. |
src/openarm_dataset/convert.py |
Exposes --success-only and --overwrite CLI flags and passes them through to Dataset.write(...) for lerobot_v2.1. |
Contributor
|
|
Contributor
Author
|
OK, will separate the PR. |
k1000dai
changed the title
Contributor
Author
|
Remove |
kou
reviewed
May 12, 2026
| num_episodes = len(records) | ||
| total_chunks = max((num_episodes - 1) // CHUNK_SIZE + 1, 0) if num_episodes else 0 | ||
| train_end = int(num_episodes * train_split) | ||
| train_end = round(num_episodes * train_split) |
Comment on lines
+115
to
+135
| def _build_remaps(dataset: Dataset, records): | ||
| """Build remapping dicts from original episode/task indices to contiguous indices. | ||
|
|
||
| When records is a filtered subset of episodes (e.g., success_only=True), | ||
| original indices may be sparse. LeRobot v2.1 expects episode/task indices | ||
| to be contiguous starting from 0. When records contains all episodes the | ||
| returned maps are the identity. | ||
| """ | ||
| remap_episode_index = {original: new for new, (original, *_) in enumerate(records)} | ||
| seen = set() | ||
| used_task_indices = [] | ||
| for original_episode_index, *_ in records: | ||
| original_task_index = int( | ||
| dataset.meta.episodes[original_episode_index]["task_index"] | ||
| ) | ||
| if original_task_index not in seen: | ||
| seen.add(original_task_index) | ||
| used_task_indices.append(original_task_index) | ||
| used_task_indices.sort() | ||
| remap_task_index = {original: new for new, original in enumerate(used_task_indices)} | ||
| return remap_episode_index, remap_task_index |
Comment on lines
+59
to
+64
| parser.add_argument( | ||
| "--success-only", | ||
| help="Include only successful episodes in the output dataset (default: False) if the output format is lerobot_v2.1", | ||
| action="store_true", | ||
| default=False, | ||
| ) |
kou
reviewed
May 12, 2026
Contributor
kou
left a comment
kou
left a comment
There was a problem hiding this comment.
It may be better that we create a class instead of using module private functions to reduce arguments for module private functions.
Can we work on it as a separated task?
| ): | ||
| records = [] | ||
| for episode_index in range(dataset.meta.num_episodes): | ||
| success = bool(dataset.meta.episodes[episode_index]["success"]) |
| seen = set() | ||
| used_task_indices = [] | ||
| for original_episode_index, *_ in records: | ||
| original_task_index = int( |
abetomo
reviewed
May 12, 2026
|
|
||
| # collect downsampled data for each episode | ||
| records = _collect_downsampled_data(dataset, fps, joint_keys) | ||
| records = _collect_downsampled_data(dataset, fps, joint_keys, success_only) |
Contributor
There was a problem hiding this comment.
Should we return an error when there are 0 records?
Contributor
There was a problem hiding this comment.
Is empty LeRobotDataset v2.1 invalid?
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
values starting from 0, the conversion now remaps original sparse indices to a contiguous range whenever episodes are filtered out (and falls back to the identity map when nothing is filtered).