feat(dataset): integrate episode streaming into training

2026-07-26 11:16:00 +00:00 · 2026-07-23 16:10:39 +02:00
parent b85620657f
commit 3d70b21aac
34 changed files with 2771 additions and 942 deletions
@@ -37,6 +37,8 @@
 - sections:
  - local: lerobot-dataset-v3
    title: Using LeRobotDataset
+  - local: training_dataset_streaming
+    title: Training Dataset Streaming
  - local: porting_datasets_v3
    title: Porting Large Datasets
  - local: using_dataset_tools
@@ -131,15 +131,68 @@ for batch in data_loader:
    # model.forward(batch)
 ```

-## Stream a dataset (no downloads)
+## Stream a dataset during training

-Use `StreamingLeRobotDataset` to iterate directly from the Hub without local copies. This allows to stream large datasets without the need to downloading them onto disk or loading them onto memory, and is a key feature of the new dataset format.
+Enable training-time streaming with the public `--dataset.streaming=true` flag:
+
+```bash
+lerobot-train \
+  --dataset.repo_id=yaak-ai/L2D-v3 \
+  --dataset.streaming=true \
+  --policy.type=act \
+  --output_dir=outputs/train/act_streaming
+```
+
+This is separate from `--dataset.streaming_encoding=true`, which controls video encoding while
+recording. Training-time streaming leaves the map-style `LeRobotDataset` path unchanged.
+
+`StreamingLeRobotDataset` assigns complete episodes disjointly across distributed ranks and
+DataLoader workers, reads only the selected episode rows from Parquet, and keeps a bounded pool of
+episode video bytes. Every selected frame is visited exactly once per streaming epoch.
+
+On first use, LeRobot looks for a revision-matched MP4 index sidecar. If none is published with the
+dataset, it builds one in the revision-keyed local LeRobot cache under a process lock and installs it
+atomically. Training never uploads this sidecar. Dataset maintainers can build and publish one
+explicitly with `scripts/build_mp4_sidecar.py --push`.
+
+The default memory cap is 8 GiB per DataLoader worker. It can be adjusted along with episode mixing
+and prefetch:
+
+```bash
+lerobot-train \
+  --dataset.repo_id=yaak-ai/L2D-v3 \
+  --dataset.streaming=true \
+  --dataset.streaming_byte_budget_gb=4 \
+  --dataset.streaming_episode_pool_size=16 \
+  --dataset.streaming_prefetch_episodes=4 \
+  --policy.type=act \
+  --output_dir=outputs/train/act_streaming
+```
+
+Use `--dataset.streaming_data_root=hf://buckets/OWNER/BUCKET/PREFIX` when metadata lives in a
+dataset repository but Parquet and MP4 data are mirrored in an HF Bucket.
+
+To capture comparable data-pipeline results on a training host:
+
+```bash
+uv run python scripts/bench_streaming_dataset.py \
+  --repo-id=yaak-ai/L2D-v3 \
+  --batch-size=16 \
+  --num-workers=4 \
+  --summary-json=streaming-benchmark.json
+```
+
+The JSON records the code and dataset revisions, initialization/first-batch time, steady-state
+samples per second, batch-wait p50/p95, duplicate indices, and the exact cache/worker settings.
+Run end-to-end `lerobot-train` separately to measure GPU utilization and full training step time.
+
+The Python API uses the same implementation:

 ```python
 from lerobot.datasets import StreamingLeRobotDataset

 repo_id = "yaak-ai/L2D-v3"
-dataset = StreamingLeRobotDataset(repo_id)  # streams directly from the Hub
+dataset = StreamingLeRobotDataset(repo_id)
 ```

 <div style="display:flex; justify-content:center; gap:12px; flex-wrap:wrap;">
@@ -0,0 +1,127 @@
+# 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.