Files
lerobot/docs/source/training_dataset_streaming.mdx
T

128 lines
6.1 KiB
Plaintext

# Training Dataset Streaming
Training-time dataset streaming lets `lerobot-train` consume a LeRobotDataset v3 without first
downloading its complete Parquet and video payload. Enable it through the existing public switch:
```bash
lerobot-train \
--dataset.repo_id=OWNER/DATASET \
--dataset.streaming=true \
--policy.type=act \
--output_dir=outputs/train/act_streaming
```
This feature is independent of `--dataset.streaming_encoding=true`. `streaming_encoding` controls
how videos are written while recording; `dataset.streaming` controls how an existing dataset is read
during training. Recording and rollout encoders are not used by this training path.
## How an epoch is read
Each distributed rank and DataLoader worker owns a disjoint set of complete episodes. Within that
set, an exact-coverage sampler mixes a bounded pool of episodes while visiting every selected frame
once per epoch. Parquet columns and compressed MP4 byte ranges are prefetched from the same episode
admission frontier. Temporal history and future windows are resolved inside the complete episode,
including the same boundary padding masks as map-style loading.
The default map-style `LeRobotDataset` behavior is unchanged when `--dataset.streaming=false`.
## MP4 sidecars and the first run
Video streaming uses a small MP4 index sidecar. Dataset initialization first checks the
revision-keyed local cache, then looks for a valid published sidecar. If neither is available,
LeRobot builds the sidecar locally while holding a process lock and installs it atomically. A
failed or interrupted build does not replace the previous valid file.
Training is read-only: it never uploads a sidecar or modifies the dataset repository. On a cluster
with node-local caches, the first job may build once per node. A shared LeRobot cache avoids that
duplication.
Dataset maintainers can build a sidecar ahead of time:
```bash
uv run python scripts/build_mp4_sidecar.py \
--repo-id=OWNER/DATASET \
--revision=COMMIT_SHA \
--data-root=hf://datasets/OWNER/DATASET@COMMIT_SHA \
--output=/tmp/dataset-mp4-sidecar.npz
```
Publication is always explicit. Add `--push` only after validating the complete-dataset sidecar.
Subset sidecars cannot be published.
## Configuration
The production defaults are:
| Option | Default | Meaning |
| ----------------------------- | ------: | --------------------------------------------------- |
| `streaming_episode_pool_size` | 32 | Complete episodes mixed by each worker |
| `streaming_prefetch_episodes` | 8 | Episodes fetched ahead of the active pool |
| `streaming_byte_budget_gb` | 8 | Maximum synthesized MP4 bytes per worker |
| `streaming_data_root` | unset | Optional local, Hub, bucket, or fsspec payload root |
Memory limits are per DataLoader worker. For example, four workers with the default byte budget can
use up to 32 GiB for compressed episode video bytes, plus Parquet batches, decoders, and framework
overhead. Start with fewer workers or a smaller budget on memory-constrained hosts:
```bash
lerobot-train \
--dataset.repo_id=OWNER/DATASET \
--dataset.streaming=true \
--dataset.streaming_episode_pool_size=16 \
--dataset.streaming_prefetch_episodes=4 \
--dataset.streaming_byte_budget_gb=4 \
--num_workers=2 \
--policy.type=act \
--output_dir=outputs/train/act_streaming
```
If metadata remains in a dataset repository while payload files are mirrored elsewhere, set
`--dataset.streaming_data_root`. Supported values include local paths, revision-qualified
`hf://datasets/...` roots, `hf://buckets/...` roots, and fsspec URLs.
## Resume and shuffle migration
The earlier streaming reader used a bounded row shuffle buffer. The episode reader instead has
deterministic exact-coverage ordering derived from the seed and epoch. Checkpoint resume restores the
per-rank sample offset using the checkpoint batch size. Changing distributed world size or
DataLoader worker count changes episode ownership; changing batch size changes the batch boundaries.
For sample-exact comparisons, resume with the same world size, worker count, and batch size.
## Benchmarking
Measure the integrated path on the same hosts used for training:
```bash
uv run python scripts/bench_streaming_dataset.py \
--repo-id=OWNER/DATASET \
--revision=COMMIT_SHA \
--batch-size=16 \
--num-workers=4 \
--summary-json=streaming-benchmark.json
```
The output records source revisions, startup and first-batch latency, steady-state throughput,
batch-wait percentiles, duplicate indices, memory high-water marks, and the exact settings. Run a
separate end-to-end training A/B to include policy compute, device transfer, and optimizer time.
Do not compare results unless the code revision, dataset revision, hardware, and settings match.
For lower-level byte-fetch, retry, decoder-open, and refill diagnostics, build the sidecar explicitly
and run `scripts/bench_episode_byte_cache.py` with the same `--repo-id`, `--revision`,
`--data-root`, and `--sidecar-path`. Treat this as a stage profiler; throughput claims should come
from the integrated dataset benchmark and end-to-end training A/B.
## Troubleshooting
- **The first batch takes a long time:** check logs for a sidecar build. Reuse a shared cache or
publish a validated complete sidecar explicitly.
- **A sidecar lock times out:** another process may still be indexing the same revision. Confirm it
is healthy before removing a stale lock.
- **A rank owns no data:** reduce the number of ranks or workers so every rank owns at least one
selected episode.
- **The byte budget is exceeded:** lower the episode pool, increase the per-worker byte budget, or
use fewer/larger episodes per worker.
- **Remote reads cannot authenticate:** verify the normal Hugging Face token or fsspec credentials
are available in every worker environment. Credentials are never embedded in the sidecar.
- **Refill stalls are high:** compare p95/p99 batch wait, reduce network contention, raise prefetch
gradually, and verify that the dataset sidecar matches the exact revision.