feat(streaming): default episode pool 1024 and wire streaming into lerobot-train

Raise the default episode_pool_size to 1024 (DatasetConfig + StreamingLeRobotDataset)
for better default shuffle quality at scale.

Streaming is now a first-class option of the main train script: when cfg.dataset.streaming
is set, the dataloader is not handed to accelerate (the dataset is already rank-disjoint via
split_dataset_by_node, so IterableDatasetShard would drop (N-1)/N of each rank's stream),
batches are moved to device manually, and the episode-aware sampler is skipped. Remove the
standalone examples/scaling/train_streaming_multinode.py example in favor of this wiring.

Co-authored-by: Cursor <cursoragent@cursor.com>

examples/scaling/train_streaming_multinode.py

@@ -1,192 +0,0 @@
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Distributed, resumable streaming training on a large HF-hosted dataset.
-
-This example shows how to train (or just stress the data pipeline) over a multi-TB dataset that never
-touches local disk, scaling across GPUs and nodes with Accelerate. It demonstrates the large-scale
-streaming features of :class:`StreamingLeRobotDataset`:
-
-- per-rank sharding via ``split_dataset_by_node`` (each GPU streams disjoint data; ``rank``/``world_size``
-  are auto-resolved from the Accelerate state, so nothing needs to be passed explicitly);
-- DataLoader-worker shard splitting (no duplicate frames within a rank);
-- native `datasets` resume: the loader checkpoints stream state via ``state_dict()`` (``torchdata`` StatefulDataLoader when available, so ``num_workers > 0`` resumes too);
-- an explicit video-decoder cache size so the working set of open decoders does not thrash.
-
-Launch with Accelerate (single node, N GPUs):
-
-    accelerate launch --num_processes=8 examples/scaling/train_streaming_multinode.py \
-        --repo_id=lerobot/droid_1.0.1 --batch_size=64
-
-Multinode runs launch the same script with your cluster's accelerate/SLURM setup.
-
-Pass ``--dummy`` to skip the model entirely and measure pure dataloading throughput.
-"""
-
-import argparse
-import time
-from pathlib import Path
-
-import torch
-from accelerate import Accelerator
-from torch.utils.data import DataLoader
-
-from lerobot.datasets import LeRobotDatasetMetadata, StreamingLeRobotDataset
-from lerobot.utils.constants import ACTION
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--repo_id", type=str, default="lerobot/droid_1.0.1")
-    parser.add_argument(
-        "--root", type=str, default=None, help="Local/prewarmed dataset root (else stream from Hub)."
-    )
-    parser.add_argument("--output_dir", type=str, default="outputs/train/streaming_multinode")
-    parser.add_argument("--steps", type=int, default=1000)
-    parser.add_argument("--batch_size", type=int, default=64, help="Per-process batch size.")
-    parser.add_argument("--num_workers", type=int, default=8)
-    parser.add_argument(
-        "--episode_pool_size",
-        type=int,
-        default=64,
-        help="Whole episodes open per consumer (randomness knob).",
-    )
-    parser.add_argument("--video_decoder_cache_size", type=int, default=None)
-    parser.add_argument("--n_action_steps", type=int, default=16, help="Action-chunk length (delta horizon).")
-    parser.add_argument("--save_freq", type=int, default=200)
-    parser.add_argument("--log_freq", type=int, default=20)
-    parser.add_argument("--resume_from", type=str, default=None, help="Checkpoint dir to resume from.")
-    parser.add_argument("--dummy", action="store_true", help="Skip the model; measure dataloading only.")
-    return parser.parse_args()
-
-
-def make_dataloader(
-    args: argparse.Namespace, meta: LeRobotDatasetMetadata
-) -> tuple[DataLoader, StreamingLeRobotDataset]:
-    # Supervise an action chunk; delta_timestamps drive the SARM-style temporal window.
-    delta_timestamps = {ACTION: [t / meta.fps for t in range(args.n_action_steps)]}
-    # rank / world_size are resolved automatically from the Accelerate state inside the dataset.
-    dataset = StreamingLeRobotDataset(
-        args.repo_id,
-        root=args.root,
-        delta_timestamps=delta_timestamps,
-        episode_pool_size=args.episode_pool_size,
-        video_decoder_cache_size=args.video_decoder_cache_size,
-        tolerance_s=1e-3,
-    )
-    # torchdata's StatefulDataLoader checkpoints each worker's dataset state through the
-    # dataset's native state_dict protocol, making resume work with num_workers > 0. Fall back
-    # to the plain DataLoader (resume then requires num_workers=0).
-    try:
-        from torchdata.stateful_dataloader import StatefulDataLoader
-
-        loader_cls = StatefulDataLoader
-    except ImportError:
-        loader_cls = DataLoader
-    loader = loader_cls(
-        dataset,
-        batch_size=args.batch_size,
-        num_workers=args.num_workers,
-        pin_memory=True,
-        drop_last=True,
-        prefetch_factor=2 if args.num_workers > 0 else None,
-    )
-    return loader, dataset
-
-
-def main() -> None:
-    args = parse_args()
-    accelerator = Accelerator()
-    output_dir = Path(args.output_dir)
-    if accelerator.is_main_process:
-        output_dir.mkdir(parents=True, exist_ok=True)
-
-    meta = LeRobotDatasetMetadata(args.repo_id, root=args.root)
-    loader, dataset = make_dataloader(args, meta)
-
-    if args.dummy:
-        model = optimizer = None
-    else:
-        from lerobot.policies.act import ACTConfig, ACTPolicy
-        from lerobot.utils.feature_utils import dataset_to_policy_features
-
-        features = dataset_to_policy_features(meta.features)
-        output_features = {k: ft for k, ft in features.items() if k == ACTION}
-        input_features = {k: ft for k, ft in features.items() if k not in output_features}
-        cfg = ACTConfig(input_features=input_features, output_features=output_features)
-        model = ACTPolicy(cfg)
-        optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
-        # Do NOT prepare the dataloader: the dataset is already rank-disjoint via
-        # split_dataset_by_node, and accelerate's IterableDatasetShard would keep only every
-        # world_size-th batch of it (silently training on 1/N of the data while decoding all
-        # of it). Batches are moved to the device manually in the loop.
-        model, optimizer = accelerator.prepare(model, optimizer)
-
-    # Resume: native datasets stream state, saved per rank. With torchdata's StatefulDataLoader
-    # the state covers every worker; with the plain DataLoader it is exact for num_workers=0.
-    can_checkpoint_loader = hasattr(loader, "state_dict")
-    if args.resume_from is not None:
-        state_path = Path(args.resume_from) / f"dataset_state_rank{accelerator.process_index}.pt"
-        state = torch.load(state_path, weights_only=False)  # plain dict of stream offsets  # nosec B614
-        if can_checkpoint_loader:
-            loader.load_state_dict(state)
-        else:
-            dataset.load_state_dict(state)
-        accelerator.print(f"Resumed dataset stream from {state_path}")
-
-    step = 0
-    frames_seen = 0
-    window_start = time.perf_counter()
-    done = False
-    while not done:
-        for batch in loader:
-            if model is not None:
-                batch = {k: (v.to(accelerator.device) if torch.is_tensor(v) else v) for k, v in batch.items()}
-                loss, _ = model.forward(batch)
-                accelerator.backward(loss)
-                optimizer.step()
-                optimizer.zero_grad()
-
-            step += 1
-            frames_seen += args.batch_size
-            if step % args.log_freq == 0:
-                elapsed = time.perf_counter() - window_start
-                fps_per_proc = (args.log_freq * args.batch_size) / max(elapsed, 1e-9)
-                total_fps = fps_per_proc * accelerator.num_processes
-                accelerator.print(
-                    f"step {step} | {fps_per_proc:.1f} frames/s/proc | {total_fps:.1f} frames/s total"
-                    + ("" if model is None else f" | loss {loss.item():.3f}")
-                )
-                window_start = time.perf_counter()
-
-            if step % args.save_freq == 0:
-                ckpt = output_dir / f"checkpoint-{step}"
-                if accelerator.is_main_process:
-                    ckpt.mkdir(parents=True, exist_ok=True)
-                accelerator.wait_for_everyone()
-                # Every rank saves its own stream state: shard positions differ per rank.
-                state = loader.state_dict() if can_checkpoint_loader else dataset.state_dict()
-                torch.save(state, ckpt / f"dataset_state_rank{accelerator.process_index}.pt")
-                if model is not None and accelerator.is_main_process:
-                    accelerator.unwrap_model(model).save_pretrained(ckpt)
-
-            if step >= args.steps:
-                done = True
-                break
-
-    accelerator.print(f"End of training: {step} steps, ~{frames_seen} frames/proc")
-
-
-if __name__ == "__main__":
-    main()

src/lerobot/configs/default.py

@@ -42,7 +42,7 @@ class DatasetConfig:
     # Whole episodes each streaming consumer keeps open to shuffle across (the randomness knob).
     # Larger mixes more episodes per batch at the cost of cold-start latency; RAM stays small because
     # the pool holds tabular rows only. Ignored when streaming is False.
-    streaming_episode_pool_size: int = 64
+    streaming_episode_pool_size: int = 1024
 
     def __post_init__(self) -> None:
         if self.episodes is not None:

src/lerobot/datasets/streaming_dataset.py

@@ -79,7 +79,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
         dataset = StreamingLeRobotDataset(
             repo_id="your-dataset-repo-id",
             delta_timestamps={"action": [0.0, 0.1, 0.2]},
-            episode_pool_size=64,
+            episode_pool_size=1024,
         )
         for sample in dataset:
             ...
@@ -97,7 +97,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
         revision: str | None = None,
         force_cache_sync: bool = False,
         streaming: bool = True,
-        episode_pool_size: int | None = 64,
+        episode_pool_size: int | None = 1024,
         frame_shuffle_buffer_size: int | None = None,
         buffer_size: int | None = None,
         max_num_shards: int | None = None,
@@ -128,7 +128,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
             episode_pool_size (int, optional): Whole episodes each consumer keeps open to shuffle
                 across — the randomness knob. Larger mixes more episodes per batch (closer to
                 map-style uniform) at the cost of cold-start latency and frame-buffer RAM.
-                Defaults to 64.
+                Defaults to 1024.
             frame_shuffle_buffer_size (int | None, optional): Frame-level shuffle buffer after the
                 episode pool. Defaults to ``episode_pool_size x average episode length`` (capped),
                 which matches the pool's mixing radius.
@@ -178,7 +178,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
         self.shuffle = shuffle
 
         self.streaming = streaming
-        self.episode_pool_size = max(1, episode_pool_size) if episode_pool_size else 64
+        self.episode_pool_size = max(1, episode_pool_size) if episode_pool_size else 1024
         self._return_uint8 = return_uint8
 
         self.rank, self.world_size = self._resolve_distributed(rank, world_size)

src/lerobot/scripts/lerobot_train.py

@@ -387,7 +387,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})")
 
     # create dataloader for offline training
-    if hasattr(active_cfg, "drop_n_last_frames"):
+    if hasattr(active_cfg, "drop_n_last_frames") and not cfg.dataset.streaming:
         shuffle = False
         # A dedicated generator (rather than the global torch RNG) lets accelerator.prepare
         # synchronize the shuffle permutation across ranks, keeping batch shards disjoint even
@@ -426,9 +426,16 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
 
     # Prepare everything with accelerator
     accelerator.wait_for_everyone()
-    policy, optimizer, dataloader, lr_scheduler = accelerator.prepare(
+    if cfg.dataset.streaming:
-        policy, optimizer, dataloader, lr_scheduler
+        # The streaming IterableDataset is already rank-disjoint via split_dataset_by_node, so we must
-    )
+        # NOT hand the dataloader to accelerate: its IterableDatasetShard would keep only every
+        # world_size-th batch of each rank's already-disjoint stream (silently training on 1/N of the
+        # data while decoding all of it). Batches are moved to the device manually in the loop below.
+        policy, optimizer, lr_scheduler = accelerator.prepare(policy, optimizer, lr_scheduler)
+    else:
+        policy, optimizer, dataloader, lr_scheduler = accelerator.prepare(
+            policy, optimizer, dataloader, lr_scheduler
+        )
     dl_iter = cycle(dataloader)
 
     policy.train()
@@ -468,6 +475,9 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     for _ in range(step, cfg.steps):
         start_time = time.perf_counter()
         batch = next(dl_iter)
+        if cfg.dataset.streaming:
+            # The streaming dataloader is not accelerate-prepared (see above), so move to device here.
+            batch = {k: (v.to(device, non_blocking=True) if torch.is_tensor(v) else v) for k, v in batch.items()}
         for cam_key in dataset.meta.camera_keys:
             if cam_key in batch and batch[cam_key].dtype == torch.uint8:
                 batch[cam_key] = batch[cam_key].to(dtype=torch.float32) / 255.0