Skip to content

Refactor pretraining loaders for self-contained LeRobot v3 datasets - #55

Merged
wayrise merged 1 commit into
mainfrom
refactor/final-pretraining-datasets
Oct 25, 2026
Merged

wayrise merged 1 commit into
mainfrom
refactor/final-pretraining-datasets

Conversation

@wayrise

@wayrise wayrise commented Oct 25, 2026

Copy link
Copy Markdown
Contributor

Pretraining now loads the self-contained, cleaned LeRobot v3 versions of AgiBotWorld-Beta, RoboCOIN, DROID, and InternData-A1. Readers consume the supplied statistics inside each dataset family without frame-filter CSVs, external statistics roots, or loader-side statistics generation. Retired pretraining readers and utilities are removed; benchmark statistics tools retain their shared accumulator.

Move the standalone YAMLs and mixture.yaml into configs/dataloader/. Selecting dataloader=mixture enables all four sources with full-data, proportional sampling. The training default remains RoboTwin. README updates document the pretraining command and announce the dataset releases and 1K-star milestone, including each dataset's Hugging Face link.

Review fixes also make video offsets honor each camera's canonical LeRobot timestamps, bound family statistics lookup, reject malformed or mismatched statistics, and support arbitrary leading batch dimensions in the Euler conversion helpers. Video readers retain the cumulative-offset fallback for older benchmark metadata. Mixture indexing now uses bounded temporary storage and an in-place shuffle while preserving the previous seeded ordering and sample counts.

Validation

  • Repository-wide Ruff and git diff --check pass. Compilation passes.
  • The core non-GPU suite completes with 1,897 passed, 18 skipped, 21 deselected, plus the existing Cosmos failure described below. All changed loader, configuration, statistics, and benchmark regression tests pass.
  • All 1,133 real buckets and 1,623,118 episode records pass the new timestamp validation. Their computed video offsets are unchanged from the previous loader.
  • Default four-source initialization covers 533,271,411 windows. Actual decoding passes for 24 direct windows and 17 trainer DataLoader batches (34 returned samples, two workers, all four sources); 173 camera windows decode successfully including prefetch.
  • Independent native IO tracing confirms no reads from external dataset/CSV/statistics/processing directories, no writes to the four datasets, and no external network connections. Source counts, weights, sample shapes, masks, and prompts match the pre-review smoke report.
  • A 10-million-window comparison produces an identical complete index SHA256; construction takes 0.269s versus 0.570s, with peak RSS 125,344 KiB versus 408,992 KiB in that isolated benchmark.

Existing baseline failure

tests/test_cosmos_predict25_registry.py::test_cosmos_predict25_build_with_bad_path_raises_clear_error fails with RuntimeError: No CUDA GPUs are available on this server: the installed Cosmos/Transformer Engine stack initializes CUDA before validating the deliberately nonexistent checkpoint path. The identical failure is reproduced in a detached, unmodified main worktree at 782e3133f674449e2986754cfc5f7c4e204835a2 using the same environment. The model and test files involved are unchanged by this PR. GPU/model training and exhaustive decoding of every dataset frame were not run.

@KraHsu KraHsu left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed all 49 changed files against 782e313, including the four dataset adapters, shared LeRobot v3 reader, statistics lookup and validation, action/state representations and masks, mixture indexing, Hydra configuration, and retained benchmark consumers. No actionable regressions found.

Local validation: 489 tests passed; Ruff, syntax parsing of 414 Python files, and git diff --check passed. The broader test attempt also had 11 failures and one collection error caused by missing torchvision/modelscope dependencies; additional model-dependent tests could not be completed. GitHub CI and Docker checks are green. Full-scale dataset decoding and GPU training were not independently reproduced.

@wayrise
wayrise merged commit 3dfc9a7 into main Oct 25, 2026
3 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.

3 participants