Skip to content

Make dataset classes drop-in compatible with lerobot's dataset factory - #4

Merged
AyushExel merged 3 commits into
mainfrom
factory-drop-in
Jul 10, 2026
Merged

Make dataset classes drop-in compatible with lerobot's dataset factory#4
AyushExel merged 3 commits into
mainfrom
factory-drop-in

Conversation

@AyushExel

Copy link
Copy Markdown
Contributor

lerobot-train has no hook to swap its dataset class, so integrating the plugin previously required a user-side subclass shadowing absolute_to_relative_idx plus an adapter function for the constructor signature. This PR moves that interface responsibility into the plugin, where it belongs. After it, the entire training integration is:

import lerobot.datasets.factory as factory
from lerobot_lancedb import LeRobotLanceVideoDataset

factory.LeRobotDataset = LeRobotLanceVideoDataset  # done

if __name__ == "__main__":
    from lerobot.scripts.lerobot_train import main
    main()

Changes (both LeRobotLanceDataset and LeRobotLanceVideoDataset)

  1. Upstream-compatible constructor orderrepo_id first positional, root second, matching LeRobotDataset(repo_id, root=...) as the factory calls it. A positional local path is still detected (Path(arg).is_dir()) and treated as root, so pre-existing positional call sites keep working.
  2. absolute_to_relative_idx = None class attribute — lance serves every frame by absolute dataset index; the upstream property exists only for episode-filtered parquet readers and would try to build one. Semantically correct for lance, and lets EpisodeAwareSampler wiring read it safely.
  3. Factory kwargs toleratedvideo_backend, download_videos, force_cache_sync, depth_output_unit accepted and ignored (documented as no-ops for Lance reads).

Verification

  • 4-GPU accelerate launch run of lerobot-train (smolvla on the LIBERO lance dataset, --dataset.episodes filter active): trains normally, data_s=0.002.
  • All four call styles tested: factory-style (repo_id, root=...), keyword root=, positional path, and factory-style with parquet-only kwargs.
  • Test suite: 14 passed; test_subtask_index_round_trips errors identically on main under lerobot 0.6.0 (pre-existing meta.subtasks API drift, unrelated).

Context: this fell out of the end-to-end LIBERO example/blog (lancedb/training#6), where the user-side shim was the one wart in the migration story.

🤖 Generated with Claude Code

`lerobot-train` instantiates its dataset via `make_dataset`, which calls
`LeRobotDataset(repo_id, root=..., episodes=..., video_backend=..., ...)`
and reads `dataset.absolute_to_relative_idx` when wiring the
EpisodeAwareSampler. Three changes make both Lance classes usable there
directly — `factory.LeRobotDataset = LeRobotLanceVideoDataset` is now the
entire integration:

- constructor signature matches upstream arg order: `repo_id` is the first
  positional, `root` second (a positional local path is still detected and
  treated as root for pre-0.2 call-style compatibility)
- `absolute_to_relative_idx = None` class attribute: lance serves every
  frame by absolute index, so the upstream property (which would build a
  parquet reader) never applies
- parquet-only factory kwargs (`video_backend`, `download_videos`,
  `force_cache_sync`, `depth_output_unit`) are accepted and ignored

Verified with a 4-GPU accelerate run of lerobot-train (smolvla, episodes
filter active) and the existing test suite (the one failing test fails
identically on main under lerobot 0.6.0).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
AyushExel added a commit to lancedb/training that referenced this pull request Jul 10, 2026
…-lancedb#4)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
AyushExel and others added 2 commits July 10, 2026 17:16
…and bump to 0.2.0

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…args in test fixtures

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@AyushExel
AyushExel merged commit 10fc395 into main Jul 10, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant