refactor(dataset): make streaming pool rank-owned

docs/source/training_dataset_streaming.mdx

@@ -17,11 +17,16 @@ during training. Recording and rollout encoders are not used by this training pa
 
 ## 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.
+Each distributed rank owns a deterministic, frame-balanced set of complete episodes. One logical
+exact-coverage pool per rank mixes those episodes while visiting every selected frame once per
+rank-local coverage epoch. Parquet columns and compressed MP4 byte ranges are prefetched from the
+same byte-aware admission frontier. Temporal history and future windows are resolved inside the
+complete episode, including the same boundary padding masks as map-style loading.
+
+When `--num_workers` is nonzero, training uses one dedicated DataLoader process per rank. Its
+bounded result queue holds decoded batches while the policy trains. The configured worker count is
+instead used as internal Parquet and byte-range fetch concurrency, so increasing it does not create
+independent samplers or multiply the cache budget.
 
 The default map-style `LeRobotDataset` behavior is unchanged when `--dataset.streaming=false`.
 
@@ -55,14 +60,15 @@ The production defaults are:
 
 | Option                        | Default | Meaning                                             |
 | ----------------------------- | ------: | --------------------------------------------------- |
-| `streaming_episode_pool_size` |      32 | Complete episodes mixed by each worker              |
+| `streaming_episode_pool_size` |      32 | Maximum complete episodes mixed by each rank        |
 | `streaming_prefetch_episodes` |       8 | Episodes fetched ahead of the active pool           |
-| `streaming_byte_budget_gb`    |       8 | Maximum synthesized MP4 bytes per worker            |
+| `streaming_byte_budget_gb`    |       8 | Maximum synthesized MP4 bytes per rank              |
 | `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:
+The active episode set is capped by both episode count and the exact synthesized mini-MP4 sizes
+computed from the sidecar. An episode larger than the complete rank budget fails before training
+fetches its payload. The cache, decoder LRU, and decoded-batch queue remain independently bounded.
+Start with a smaller pool or budget on memory-constrained hosts:
 
 ```bash
 lerobot-train \
@@ -71,7 +77,7 @@ lerobot-train \
   --dataset.streaming_episode_pool_size=16 \
   --dataset.streaming_prefetch_episodes=4 \
   --dataset.streaming_byte_budget_gb=4 \
-  --num_workers=2 \
+  --num_workers=4 \
   --policy.type=act \
   --output_dir=outputs/train/act_streaming
 ```
@@ -85,8 +91,12 @@ If metadata remains in a dataset repository while payload files are mirrored els
 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.
+batch size changes ownership or batch boundaries. For sample-exact comparisons, resume with the
+same world size and batch size. Keep the same internal fetch concurrency when comparing performance.
+
+Streaming checkpoints created by the earlier multi-worker sampler record a different ownership
+topology and are rejected for sample-exact resume. Start a new run from the saved policy weights
+rather than claiming that its dataset stream resumes exactly.
 
 ## Benchmarking
 
@@ -97,7 +107,8 @@ uv run python scripts/bench_streaming_dataset.py \
   --repo-id=OWNER/DATASET \
   --revision=COMMIT_SHA \
   --batch-size=16 \
-  --num-workers=4 \
+  --num-workers=1 \
+  --fetch-workers=4 \
   --summary-json=streaming-benchmark.json
 ```
 
@@ -117,10 +128,10 @@ from the integrated dataset benchmark and end-to-end training A/B.
   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.
+- **A rank owns no data:** reduce the number of ranks so every rank owns at least one selected
+  episode.
+- **The byte budget is exceeded:** lower the episode pool, increase the per-rank byte budget, or
+  use an explicitly prepared payload layout with smaller episode ranges.
 - **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

examples/training/train_with_streaming.py

@@ -83,10 +83,12 @@ def main():
     optimizer = torch.optim.Adam(policy.parameters(), lr=1e-4)
     dataloader = torch.utils.data.DataLoader(
         dataset,
-        num_workers=4,
+        # One worker owns the rank-level pool. Internal fetch concurrency is configured through
+        # StreamingLeRobotDataset.max_num_shards (the lerobot-train CLI derives it from num_workers).
+        num_workers=1,
         batch_size=16,
         pin_memory=device.type != "cpu",
-        prefetch_factor=2,  # loads batches with multiprocessing while policy trains
+        prefetch_factor=2,  # bounded decoded-batch queue while the policy trains
         persistent_workers=True,
     )
 

scripts/bench_streaming_dataset.py

@@ -22,6 +22,7 @@ import statistics
 import subprocess
 import sys
 import time
+from collections.abc import Sequence
 from pathlib import Path
 
 import torch
@@ -37,7 +38,14 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--data-root", default=None)
     parser.add_argument("--episodes", type=int, default=None, help="Use the first N episodes.")
     parser.add_argument("--batch-size", type=int, default=16)
-    parser.add_argument("--num-workers", type=int, default=4)
+    parser.add_argument(
+        "--num-workers",
+        type=int,
+        choices=(0, 1),
+        default=1,
+        help="Rank-level DataLoader process count. Use 1 for the production pipeline.",
+    )
+    parser.add_argument("--fetch-workers", type=int, default=4)
     parser.add_argument("--prefetch-factor", type=int, default=2)
     parser.add_argument("--episode-pool-size", type=int, default=32)
     parser.add_argument("--prefetch-episodes", type=int, default=8)
@@ -48,7 +56,7 @@ def parse_args() -> argparse.Namespace:
     return parser.parse_args()
 
 
-def percentile(values: list[float], quantile: float) -> float:
+def percentile(values: Sequence[float], quantile: float) -> float:
     if not values:
         return 0.0
     ordered = sorted(values)
@@ -83,6 +91,8 @@ def child_process_max_rss_mb() -> float:
 
 def main() -> None:
     args = parse_args()
+    if args.fetch_workers <= 0:
+        raise ValueError("--fetch-workers must be positive")
     episodes = list(range(args.episodes)) if args.episodes is not None else None
 
     init_start = time.perf_counter()
@@ -95,7 +105,7 @@ def main() -> None:
         episode_pool_size=args.episode_pool_size,
         prefetch_episodes=args.prefetch_episodes,
         byte_budget_gb=args.byte_budget_gb,
-        max_num_shards=max(1, args.num_workers),
+        max_num_shards=max(1, args.fetch_workers),
         return_uint8=True,
     )
     dataset_init_s = time.perf_counter() - init_start
@@ -112,6 +122,7 @@ def main() -> None:
     waits: list[float] = []
     measured_samples = 0
     measured_indices: list[int] = []
+    unique_episodes_per_batch: list[int] = []
     first_batch_s = 0.0
     exhausted = False
 
@@ -129,7 +140,9 @@ def main() -> None:
             if batch_index >= args.warmup_batches:
                 waits.append(wait_s)
                 indices = batch["index"].reshape(-1).tolist()
+                episode_indices = batch["episode_index"].reshape(-1).tolist()
                 measured_indices.extend(int(index) for index in indices)
+                unique_episodes_per_batch.append(len({int(index) for index in episode_indices}))
                 measured_samples += len(indices)
     finally:
         shutdown = getattr(iterator, "_shutdown_workers", None)
@@ -155,12 +168,18 @@ def main() -> None:
         "batch_wait_p95_ms": percentile(waits, 0.95) * 1000,
         "batch_wait_p99_ms": percentile(waits, 0.99) * 1000,
         "duplicate_indices": measured_samples - len(set(measured_indices)),
+        "unique_episodes_per_batch_mean": (
+            statistics.fmean(unique_episodes_per_batch) if unique_episodes_per_batch else 0.0
+        ),
+        "unique_episodes_per_batch_p50": percentile(unique_episodes_per_batch, 0.50),
+        "unique_episodes_per_batch_p95": percentile(unique_episodes_per_batch, 0.95),
         "epoch_exhausted": exhausted,
         "main_process_max_rss_mb": main_process_max_rss_mb(),
         "worker_process_max_rss_mb": child_process_max_rss_mb(),
         "config": {
             "batch_size": args.batch_size,
             "num_workers": args.num_workers,
+            "fetch_workers": args.fetch_workers,
             "prefetch_factor": args.prefetch_factor,
             "episode_pool_size": args.episode_pool_size,
             "prefetch_episodes": args.prefetch_episodes,

src/lerobot/configs/default.py

@@ -46,11 +46,11 @@ class DatasetConfig:
     streaming: bool = False
     # Optional data-plane root used by training-time streaming. Metadata still resolves from repo_id/root.
     streaming_data_root: str | None = None
-    # Number of complete episodes mixed by the exact-coverage sampler in each DataLoader worker.
+    # Number of complete episodes mixed by the rank-level exact-coverage sampler.
     streaming_episode_pool_size: int = 32
     # Complete episodes fetched ahead of the current admission frontier.
     streaming_prefetch_episodes: int = 8
-    # Hard per-worker cap for synthesized episode-video bytes.
+    # Hard per-rank cap for synthesized episode-video bytes.
     streaming_byte_budget_gb: float = 8.0
     # Fraction of episodes held out per task for offline evaluation (0.0 = disabled).
     eval_split: float = 0.0

src/lerobot/datasets/streaming_dataset.py

@@ -15,7 +15,7 @@
 # limitations under the License.
 import io
 import os
-from collections.abc import Callable, Iterator
+from collections.abc import Callable, Iterator, Mapping
 from concurrent.futures import Future, ThreadPoolExecutor
 from pathlib import Path
 
@@ -41,12 +41,31 @@ from .utils import check_version_compatibility
 from .video_utils import decode_video_frames_pyav
 
 
+def _balanced_episode_shards(
+    episode_indices: list[int],
+    episode_frame_counts: Mapping[int, int],
+    *,
+    world_size: int,
+) -> list[list[int]]:
+    """Assign whole episodes deterministically with greedy frame-count balancing."""
+    if world_size <= 0:
+        raise ValueError("world_size must be positive")
+    shards: list[list[int]] = [[] for _ in range(world_size)]
+    shard_frames = [0] * world_size
+    for episode in sorted(episode_indices, key=lambda item: (-episode_frame_counts[item], item)):
+        rank = min(range(world_size), key=lambda item: (shard_frames[item], item))
+        shards[rank].append(episode)
+        shard_frames[rank] += episode_frame_counts[episode]
+    return shards
+
+
 class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
     """Episode-scoped streaming reader for LeRobot datasets.
 
-    Metadata is cached locally, while each worker reads only the Parquet rows and MP4 byte ranges
-    needed for the complete episodes it owns. Episode ownership is disjoint across distributed
-    ranks and dataloader workers, and every selected frame is yielded exactly once per iteration.
+    Metadata is cached locally, while each rank reads only the Parquet rows and MP4 byte ranges
+    needed for the complete episodes it owns. One DataLoader worker per rank owns the logical pool;
+    its bounded result queue provides decoded-batch prefetch without creating independent samplers.
+    Episode ownership is disjoint and every selected frame is yielded exactly once per iteration.
     MP4 sidecars are resolved automatically and built in a revision-keyed local cache when absent.
 
     Example:
@@ -123,8 +142,8 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
                 and fsspec URLs.
             episode_pool_size (int | None, optional): Number of complete episodes in the sampling pool.
             prefetch_episodes (int, optional): Episodes prefetched beyond the active pool.
-            byte_budget_gb (float, optional): Per-worker upper bound for synthesized episode video bytes.
-            repeat (bool, optional): Repeat worker-local exact-coverage epochs without yielding a
+            byte_budget_gb (float, optional): Per-rank upper bound for synthesized episode video bytes.
+            repeat (bool, optional): Repeat rank-local exact-coverage epochs without yielding a
                 short final training batch. The training factory enables this; direct iteration is
                 finite by default.
         """
@@ -266,39 +285,40 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
         self._active_epoch = epoch
 
         worker = torch.utils.data.get_worker_info()
-        worker_index = worker.id if worker is not None else 0
-        worker_count = worker.num_workers if worker is not None else 1
+        if worker is not None and worker.num_workers > 1:
+            raise RuntimeError(
+                "StreamingLeRobotDataset uses one rank-level sampling pool and supports at most "
+                "one DataLoader worker per rank"
+            )
         resume_offset = self._resume_offset
         self._resume_offset = 0
-        worker_resume_offset, logical_worker_index = self._worker_resume_state(
-            resume_offset,
-            batch_size=self._resume_batch_size,
-            worker_index=worker_index,
-            worker_count=worker_count,
-        )
-        consumer_episodes, _consumer_index, _consumer_count = self._consumer_episodes(
-            worker_id=logical_worker_index,
-            num_workers=worker_count,
-        )
+        consumer_episodes, _rank, _world_size = self._rank_episodes()
         consumer_frame_count = sum(self._episode_frame_count(episode) for episode in consumer_episodes)
         if consumer_frame_count:
-            worker_epoch_delta, worker_resume_offset = divmod(
-                worker_resume_offset,
-                consumer_frame_count,
-            )
+            worker_epoch_delta, resume_offset = divmod(resume_offset, consumer_frame_count)
         else:
             worker_epoch_delta = 0
         epoch += worker_epoch_delta
         self._active_epoch = epoch
         if self.shuffle:
             self._next_epoch = epoch + 1
+
+        max_workers = min(self.max_num_shards, max(1, self.episode_pool_size + self.prefetch_episodes))
+        video_cache = self._make_video_cache(consumer_episodes, max_workers)
+        episode_byte_sizes = (
+            {episode: video_cache.manifest.episode_byte_size(episode) for episode in consumer_episodes}
+            if video_cache is not None
+            else None
+        )
         planner = ExactCoveragePool(
             [(episode, self._episode_frame_count(episode)) for episode in consumer_episodes],
             pool_size=self.episode_pool_size,
             seed=self.seed,
             epoch=epoch,
+            episode_byte_sizes=episode_byte_sizes,
+            byte_budget=self.byte_budget if episode_byte_sizes is not None else None,
         )
-        for _ in range(worker_resume_offset):
+        for _ in range(resume_offset):
             try:
                 next(planner)
             except StopIteration:
@@ -307,11 +327,9 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
         planner.evicted.clear()
 
         parquet_reader = EpisodeParquetReader(self._data_root, columns=self._projected_columns)
-        max_workers = min(self.max_num_shards, max(1, self.episode_pool_size + self.prefetch_episodes))
         executor = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="lerobot-parquet")
         episode_futures: dict[int, Future[datasets.Dataset]] = {}
-        scheduled_cursor = 0
-        video_cache = self._make_video_cache(consumer_episodes, max_workers)
+        scheduled_episodes: set[int] = set()
         retained_video_episodes: set[int] = set()
         if video_cache is not None:
             for episode_index in planner.resident:
@@ -326,17 +344,14 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
             return future
 
         def schedule_frontier() -> None:
-            nonlocal scheduled_cursor
-            target = min(
-                len(planner.admission_order),
-                planner.admitted_count + self.prefetch_episodes,
-            )
-            while scheduled_cursor < target:
-                episode_index = planner.admission_order[scheduled_cursor]
+            frontier = [*planner.resident, *planner.prefetch_candidates(self.prefetch_episodes)]
+            for episode_index in frontier:
+                if episode_index in scheduled_episodes:
+                    continue
                 submit(episode_index)
                 if video_cache is not None:
                     video_cache.submit_prefetch(episode_index)
-                scheduled_cursor += 1
+                scheduled_episodes.add(episode_index)
 
         schedule_frontier()
         try:
@@ -379,12 +394,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
             if video_cache is not None:
                 video_cache.close()
 
-    def _consumer_episodes(
-        self,
-        *,
-        worker_id: int | None = None,
-        num_workers: int | None = None,
-    ) -> tuple[list[int], int, int]:
+    def _rank_episodes(self) -> tuple[list[int], int, int]:
         if torch.distributed.is_available() and torch.distributed.is_initialized():
             rank = torch.distributed.get_rank()
             world_size = torch.distributed.get_world_size()
@@ -393,49 +403,23 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
             world_size = int(os.environ.get("WORLD_SIZE", "1"))
         if world_size <= 0 or rank < 0 or rank >= world_size:
             raise ValueError(f"Invalid distributed rank/world size: rank={rank}, world_size={world_size}")
-
-        worker = torch.utils.data.get_worker_info()
-        if worker_id is None:
-            worker_id = worker.id if worker is not None else 0
-        if num_workers is None:
-            num_workers = worker.num_workers if worker is not None else 1
-        consumer_index = rank * num_workers + worker_id
-        consumer_count = world_size * num_workers
-        return self._selected_episodes[consumer_index::consumer_count], consumer_index, consumer_count
-
-    @staticmethod
-    def _worker_resume_state(
-        offset: int,
-        *,
-        batch_size: int,
-        worker_index: int,
-        worker_count: int,
-    ) -> tuple[int, int]:
-        completed_batches, partial_batch = divmod(offset, batch_size)
-        completed_cycles, next_logical_worker = divmod(completed_batches, worker_count)
-        logical_worker = (worker_index + next_logical_worker) % worker_count
-        worker_batches = completed_cycles + int(logical_worker < next_logical_worker)
-        worker_offset = worker_batches * batch_size
-        if partial_batch and logical_worker == next_logical_worker:
-            worker_offset += partial_batch
-        return worker_offset, logical_worker
+        counts = {episode: self._episode_frame_count(episode) for episode in self._selected_episodes}
+        shards = _balanced_episode_shards(self._selected_episodes, counts, world_size=world_size)
+        return shards[rank], rank, world_size
 
     def _episode_frame_count(self, episode_index: int) -> int:
         episode = self.meta.episodes[episode_index]
         return int(episode["dataset_to_index"] - episode["dataset_from_index"])
 
     def num_frames_for_rank(self, rank: int, world_size: int, num_workers: int) -> int:
-        """Return frames owned by one training rank under episode-scoped worker sharding."""
+        """Return frames owned by one training rank under balanced whole-episode sharding."""
         if world_size <= 0 or rank < 0 or rank >= world_size:
             raise ValueError(f"Invalid distributed rank/world size: rank={rank}, world_size={world_size}")
-        workers = max(1, num_workers)
-        consumer_count = world_size * workers
-        consumer_indices = set(range(rank * workers, (rank + 1) * workers))
-        return sum(
-            self._episode_frame_count(episode)
-            for position, episode in enumerate(self._selected_episodes)
-            if position % consumer_count in consumer_indices
-        )
+        if num_workers > 1:
+            raise ValueError("Rank-level streaming supports at most one DataLoader worker per rank")
+        counts = {episode: self._episode_frame_count(episode) for episode in self._selected_episodes}
+        shards = _balanced_episode_shards(self._selected_episodes, counts, world_size=world_size)
+        return sum(counts[episode] for episode in shards[rank])
 
     def _load_episode_dataset(
         self,

src/lerobot/scripts/lerobot_train.py

@@ -460,22 +460,16 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     else:
         shuffle = True
         sampler = None
-        train_num_workers = cfg.num_workers
-        if train_num_workers > 0:
-            max_nonempty_workers = dataset.num_episodes // accelerator.num_processes
-            if max_nonempty_workers == 0:
-                raise ValueError(
-                    "At least one distributed rank owns no streaming episode. "
-                    "Reduce the distributed world size."
-                )
-            train_num_workers = min(train_num_workers, max_nonempty_workers)
-            if train_num_workers != cfg.num_workers and is_main_process:
-                logging.warning(
-                    "Reducing streaming DataLoader workers from %d to %d so every rank/worker "
-                    "owns at least one complete episode.",
-                    cfg.num_workers,
-                    train_num_workers,
-                )
+        # One dedicated DataLoader process owns the rank-level planner/cache. Its bounded result
+        # queue provides decoded-batch prefetch; max_num_shards controls the internal Parquet/range
+        # fetch executors independently.
+        train_num_workers = min(cfg.num_workers, 1)
+        if cfg.num_workers > 1 and is_main_process:
+            logging.info(
+                "Using one streaming DataLoader worker per rank; %d configured workers remain "
+                "available as the dataset's internal fetch concurrency.",
+                cfg.num_workers,
+            )
         rank_frames = dataset.num_frames_for_rank(
             accelerator.process_index,
             accelerator.num_processes,

src/lerobot/streaming/episode_video.py

@@ -17,7 +17,7 @@ import posixpath
 import threading
 import time
 from collections import OrderedDict
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from concurrent.futures import Future, ThreadPoolExecutor, as_completed
 from dataclasses import dataclass
 from datetime import UTC, datetime
@@ -33,7 +33,13 @@ import numpy as np
 from huggingface_hub import HfApi, HfFileSystem, constants
 from huggingface_hub.utils import get_session, hf_raise_for_status
 
-from lerobot.streaming.mp4 import Mp4Index, Mp4SampleSlice, fetch_mp4_index, synthesize_mp4
+from lerobot.streaming.mp4 import (
+    Mp4Index,
+    Mp4SampleSlice,
+    fetch_mp4_index,
+    synthesize_mp4,
+    synthesized_mp4_size,
+)
 
 if TYPE_CHECKING:
     from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
@@ -1041,6 +1047,16 @@ class EpisodeVideoManifest:
             source_start_pts=span.source_start_pts,
         )
 
+    def episode_byte_size(self, episode_index: int) -> int:
+        """Exact synthesized video bytes retained while an episode is active."""
+        return sum(
+            synthesized_mp4_size(
+                self.mp4_index(episode_index, camera_key),
+                self.sample_slice(episode_index, camera_key),
+            )
+            for camera_key in self.video_keys
+        )
+
 
 class ExactCoveragePool:
     """Deterministic, exactly-once frame coverage over a byte-cache episode pool.
@@ -1052,7 +1068,8 @@ class ExactCoveragePool:
     reproducible.
 
     Mechanics (this is the "evict only when all frames sampled" model):
-      - Episodes are admitted in a seeded global permutation; the first ``pool_size`` fill the pool.
+      - Episodes are admitted in a seeded global permutation until either ``pool_size`` or the
+        optional indexed-byte budget is reached.
       - Each resident episode carries a seeded shuffle of its own frame indices.
       - Each draw picks a resident episode with probability proportional to its *remaining* frames
         (i.e. a uniform draw over all remaining frames in the pool, the map-style ideal) and pops
@@ -1065,9 +1082,9 @@ class ExactCoveragePool:
     is fully unit-testable. It yields ``(episode_index, frame_index)``; map to a decode timestamp
     with ``frame_index / max(frame_count - 1, 1)``.
 
-    Determinism: the order is a pure function of ``(seed, epoch)`` and the episode->frame_count
-    map. Resume is a deterministic fast-forward: re-instantiate with the same seed/epoch and skip
-    ``n`` samples (tabular only, no decode).
+    Determinism: the order is a pure function of ``(seed, epoch)``, the episode frame counts, and
+    optional byte sizes/budget. Resume is a deterministic fast-forward: re-instantiate with the
+    same inputs and skip ``n`` samples (tabular only, no decode).
     """
 
     def __init__(
@@ -1077,34 +1094,71 @@ class ExactCoveragePool:
         *,
         seed: int,
         epoch: int = 0,
+        episode_byte_sizes: Mapping[int, int] | None = None,
+        byte_budget: int | None = None,
     ):
         self._counts = {int(ep): int(n) for ep, n in episode_frame_counts if int(n) > 0}
         self._rng = np.random.default_rng([seed, epoch])
         order = np.array(sorted(self._counts), dtype=np.int64)
         self._rng.shuffle(order)
-        # Full admission order is exposed so a caller can prefetch episodes ahead of when they
-        # enter the sampling pool (a freshly admitted episode is immediately eligible to be drawn,
-        # so its bytes must already be resident).
+        self.pool_size = max(1, pool_size)
+        self._byte_budget = byte_budget
+        if byte_budget is not None and byte_budget <= 0:
+            raise ValueError("byte_budget must be positive")
+        if byte_budget is not None and episode_byte_sizes is None:
+            raise ValueError("episode_byte_sizes are required when byte_budget is set")
+        self._byte_sizes = {
+            episode: int(episode_byte_sizes[episode]) if episode_byte_sizes is not None else 0
+            for episode in self._counts
+        }
+        if any(size < 0 for size in self._byte_sizes.values()):
+            raise ValueError("episode byte sizes must be non-negative")
+        if byte_budget is not None:
+            oversized = next(
+                ((episode, size) for episode, size in self._byte_sizes.items() if size > byte_budget),
+                None,
+            )
+            if oversized is not None:
+                episode, size = oversized
+                raise ValueError(
+                    f"Episode {episode} requires {size} bytes, exceeding the byte budget {byte_budget}"
+                )
+
+        # Preserve the full seeded order for benchmark/tooling compatibility. Byte-aware admission
+        # may temporarily skip an entry, but every episode remains in this deterministic frontier.
         self.admission_order: list[int] = order.tolist()
-        self._admit_cursor = 0
-        self._remaining: dict[int, list[int]] = {}
+        self._pending: list[int] = list(self.admission_order)
+        self._admitted_count = 0
+        self._remaining: dict[int, tuple[np.ndarray, int]] = {}
         self._remaining_total = 0
+        self._resident_bytes = 0
         self.newly_admitted: list[int] = []
         self.evicted: list[int] = []
-        for _ in range(max(1, pool_size)):
-            self._admit_one()
+        self._admit_available()
 
-    def _admit_one(self) -> None:
-        if self._admit_cursor >= len(self.admission_order):
-            return
-        ep = self.admission_order[self._admit_cursor]
-        self._admit_cursor += 1
-        n = self._counts[ep]
-        frames = np.arange(n, dtype=np.int64)
-        self._rng.shuffle(frames)
-        self._remaining[ep] = frames.tolist()
-        self._remaining_total += n
-        self.newly_admitted.append(ep)
+    def _admit_available(self) -> None:
+        while len(self._remaining) < self.pool_size and self._pending:
+            available_bytes = None if self._byte_budget is None else self._byte_budget - self._resident_bytes
+            pending_index = next(
+                (
+                    index
+                    for index, episode in enumerate(self._pending)
+                    if available_bytes is None or self._byte_sizes[episode] <= available_bytes
+                ),
+                None,
+            )
+            if pending_index is None:
+                return
+
+            episode = self._pending.pop(pending_index)
+            frame_count = self._counts[episode]
+            frames = np.arange(frame_count, dtype=np.int64)
+            self._rng.shuffle(frames)
+            self._remaining[episode] = (frames, frame_count)
+            self._remaining_total += frame_count
+            self._resident_bytes += self._byte_sizes[episode]
+            self._admitted_count += 1
+            self.newly_admitted.append(episode)
 
     @property
     def remaining_total(self) -> int:
@@ -1113,12 +1167,22 @@ class ExactCoveragePool:
     @property
     def admitted_count(self) -> int:
         """Number of episodes pulled from the admission order so far (pool fills + rotations)."""
-        return self._admit_cursor
+        return self._admitted_count
 
     @property
     def resident(self) -> list[int]:
         return list(self._remaining)
 
+    @property
+    def resident_bytes(self) -> int:
+        return self._resident_bytes
+
+    def prefetch_candidates(self, count: int) -> list[int]:
+        """Return the next deterministic pending frontier without admitting it."""
+        if count <= 0:
+            return []
+        return self._pending[:count]
+
     def __iter__(self) -> ExactCoveragePool:
         return self
 
@@ -1129,18 +1193,22 @@ class ExactCoveragePool:
         # remaining count. O(pool_size) per draw (~1024) -> negligible next to decode.
         target = int(self._rng.integers(self._remaining_total))
         chosen = None
-        for ep, frames in self._remaining.items():
-            if target < len(frames):
+        for ep, (_frames, remaining) in self._remaining.items():
+            if target < remaining:
                 chosen = ep
                 break
-            target -= len(frames)
-        frames = self._remaining[chosen]
-        frame_index = frames.pop()
+            target -= remaining
+        frames, remaining = self._remaining[chosen]
+        remaining -= 1
+        frame_index = int(frames[remaining])
         self._remaining_total -= 1
-        if not frames:
+        if remaining == 0:
             del self._remaining[chosen]
             self.evicted.append(chosen)
-            self._admit_one()
+            self._resident_bytes -= self._byte_sizes[chosen]
+            self._admit_available()
+        else:
+            self._remaining[chosen] = (frames, remaining)
         return chosen, frame_index
 
 

src/lerobot/streaming/mp4.py

@@ -329,10 +329,49 @@ def synthesize_mp4(index: Mp4Index, sample_slice: Mp4SampleSlice, mdat_payload:
     sync = index.sync_samples[(index.sync_samples >= lo) & (index.sync_samples < hi)] - lo + 1
     moov = _make_moov(index, durations, sizes, rel_offsets, sync, mdat_data_offset=0)
     header_size = len(index.ftyp) + len(moov)
-    moov = _make_moov(index, durations, sizes, rel_offsets, sync, mdat_data_offset=header_size + 8)
+    mdat_header_size = 8 if len(mdat_payload) + 8 <= 0xFFFFFFFF else 16
+    moov = _make_moov(
+        index,
+        durations,
+        sizes,
+        rel_offsets,
+        sync,
+        mdat_data_offset=header_size + mdat_header_size,
+    )
     return index.ftyp + moov + _box(b"mdat", mdat_payload)
 
 
+def synthesized_mp4_size(index: Mp4Index, sample_slice: Mp4SampleSlice) -> int:
+    """Return the exact synthesized mini-MP4 size without fetching its media payload."""
+    lo = sample_slice.sample_lo
+    hi = sample_slice.sample_hi + 1
+    if lo < 0 or hi > len(index.sample_sizes) or lo >= hi:
+        raise ValueError(f"Invalid sample range [{lo}, {hi}) for {index.file_path}")
+
+    offsets = index.sample_offsets[lo:hi]
+    sizes = index.sample_sizes[lo:hi]
+    rel_offsets = offsets - sample_slice.byte_offset
+    if int(rel_offsets.min()) != 0:
+        raise ValueError("Sample slice must start at the minimum referenced sample offset")
+    if int((rel_offsets + sizes).max()) > sample_slice.byte_length:
+        raise ValueError("Sample slice does not cover all referenced samples")
+
+    durations = index.sample_durations[lo:hi]
+    sync = index.sync_samples[(index.sync_samples >= lo) & (index.sync_samples < hi)] - lo + 1
+    moov = _make_moov(index, durations, sizes, rel_offsets, sync, mdat_data_offset=0)
+    header_size = len(index.ftyp) + len(moov)
+    mdat_header_size = 8 if sample_slice.byte_length + 8 <= 0xFFFFFFFF else 16
+    moov = _make_moov(
+        index,
+        durations,
+        sizes,
+        rel_offsets,
+        sync,
+        mdat_data_offset=header_size + mdat_header_size,
+    )
+    return len(index.ftyp) + len(moov) + mdat_header_size + sample_slice.byte_length
+
+
 def iter_boxes(
     data: bytes,
     start: int,

tests/datasets/test_episode_video_streaming.py

@@ -38,6 +38,7 @@ from lerobot.streaming.mp4 import (
     _vmhd,
     parse_mp4_index,
     synthesize_mp4,
+    synthesized_mp4_size,
 )
 
 
@@ -90,6 +91,15 @@ def test_synthesized_mp4_rebases_one_chunk_per_sample_offsets():
     np.testing.assert_array_equal(mini_index.sample_sizes, np.array([10, 10, 10]))
 
 
+def test_synthesized_mp4_size_matches_materialized_bytes():
+    mp4 = parse_mp4_index("test.mp4", _minimal_mp4([10_000, 10_050, 10_025]))
+    sample_slice = mp4.sample_slice(0.0, 2.0, keyframe_pad_s=0, keyframe_pad_fraction=0)
+
+    mini = synthesize_mp4(mp4, sample_slice, b"x" * sample_slice.byte_length)
+
+    assert synthesized_mp4_size(mp4, sample_slice) == len(mini)
+
+
 def test_parser_accepts_co64_chunk_offsets():
     mp4 = parse_mp4_index("test.mp4", _minimal_mp4([10_000, 10_050, 10_025], use_co64=True))
 

tests/datasets/test_exact_coverage_pool.py

@@ -16,6 +16,8 @@
 
 from collections import Counter
 
+import pytest
+
 from lerobot.streaming.episode_video import ExactCoveragePool
 
 EPISODES = [(0, 5), (1, 3), (2, 8), (3, 1), (4, 6), (5, 4), (6, 7), (7, 2)]
@@ -95,3 +97,50 @@ def test_zero_length_episodes_skipped():
     pool = ExactCoveragePool([(0, 3), (1, 0), (2, 2)], pool_size=8, seed=0)
     out, _ = _drain(pool)
     assert Counter(out) == Counter({(0, 0): 1, (0, 1): 1, (0, 2): 1, (2, 0): 1, (2, 1): 1})
+
+
+def test_byte_aware_admission_never_exceeds_budget():
+    sizes = {0: 7, 1: 6, 2: 4, 3: 3, 4: 2}
+    pool = ExactCoveragePool(
+        EPISODES[:5],
+        pool_size=4,
+        seed=9,
+        episode_byte_sizes=sizes,
+        byte_budget=10,
+    )
+    out = []
+    max_resident_bytes = pool.resident_bytes
+    while pool.remaining_total:
+        out.append(next(pool))
+        max_resident_bytes = max(max_resident_bytes, pool.resident_bytes)
+
+    assert Counter(out) == Counter((ep, frame) for ep, count in EPISODES[:5] for frame in range(count))
+    assert max_resident_bytes <= 10
+    assert len(pool.admission_order) == len(EPISODES[:5])
+
+
+def test_byte_aware_admission_rejects_one_oversized_episode():
+    with pytest.raises(ValueError, match="Episode 1.*byte budget"):
+        ExactCoveragePool(
+            [(0, 2), (1, 3)],
+            pool_size=2,
+            seed=0,
+            episode_byte_sizes={0: 4, 1: 11},
+            byte_budget=10,
+        )
+
+
+def test_prefetch_candidates_follow_deterministic_pending_frontier():
+    pool = ExactCoveragePool(
+        EPISODES,
+        pool_size=2,
+        seed=17,
+        episode_byte_sizes={episode: 1 for episode, _ in EPISODES},
+        byte_budget=2,
+    )
+
+    candidates = pool.prefetch_candidates(3)
+
+    assert len(candidates) == 3
+    assert not set(candidates) & set(pool.resident)
+    assert candidates == pool.prefetch_candidates(3)

tests/datasets/test_streaming_production.py

@@ -19,7 +19,7 @@ import torch
 
 pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
 
-from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
+from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset, _balanced_episode_shards
 from lerobot.utils.utils import cycle
 from tests.fixtures.constants import DUMMY_REPO_ID
 
@@ -209,6 +209,28 @@ def test_streaming_reads_video_bytes_from_configured_fsspec_root(
         assert torch.equal(sample[camera_key], reference[camera_key])
 
 
+def test_streaming_rejects_episode_larger_than_rank_byte_budget(
+    tmp_path: Path, lerobot_dataset_factory
+) -> None:
+    root = tmp_path / "dataset"
+    lerobot_dataset_factory(
+        root=root,
+        repo_id=DUMMY_REPO_ID,
+        total_episodes=2,
+        total_frames=10,
+    )
+    streaming = StreamingLeRobotDataset(
+        DUMMY_REPO_ID,
+        root=root,
+        shuffle=False,
+        buffer_size=2,
+        byte_budget_gb=1 / 1024**3,
+    )
+
+    with pytest.raises(ValueError, match="Episode .*byte budget"):
+        next(iter(streaming))
+
+
 def test_streaming_rank_shards_are_disjoint(tmp_path: Path, lerobot_dataset_factory, monkeypatch) -> None:
     root = tmp_path / "dataset"
     map_dataset = lerobot_dataset_factory(
@@ -239,9 +261,22 @@ def test_streaming_rank_shards_are_disjoint(tmp_path: Path, lerobot_dataset_fact
     assert per_rank[0] | per_rank[1] == set(range(len(map_dataset)))
 
 
-def test_streaming_workers_do_not_duplicate_frames(tmp_path: Path, lerobot_dataset_factory) -> None:
+def test_rank_shards_are_greedily_balanced_by_frame_count() -> None:
+    shards = _balanced_episode_shards(
+        [0, 1, 2, 3, 4],
+        {0: 100, 1: 90, 2: 20, 3: 10, 4: 5},
+        world_size=2,
+    )
+
+    assert {episode for shard in shards for episode in shard} == {0, 1, 2, 3, 4}
+    assert set(shards[0]).isdisjoint(shards[1])
+    totals = [sum({0: 100, 1: 90, 2: 20, 3: 10, 4: 5}[episode] for episode in shard) for shard in shards]
+    assert max(totals) - min(totals) <= 15
+
+
+def test_streaming_rejects_multiple_sampling_workers(tmp_path: Path, lerobot_dataset_factory) -> None:
     root = tmp_path / "dataset"
-    map_dataset = lerobot_dataset_factory(
+    lerobot_dataset_factory(
         root=root,
         repo_id=DUMMY_REPO_ID,
         total_episodes=8,
@@ -256,10 +291,8 @@ def test_streaming_workers_do_not_duplicate_frames(tmp_path: Path, lerobot_datas
     )
     loader = torch.utils.data.DataLoader(streaming, batch_size=None, num_workers=2)
 
-    indices = [int(item["index"]) for item in loader]
-
-    assert len(indices) == len(map_dataset)
-    assert set(indices) == set(range(len(map_dataset)))
+    with pytest.raises(RuntimeError, match="one DataLoader worker per rank"):
+        list(loader)
 
 
 def test_streaming_persistent_workers_advance_epochs(tmp_path: Path, lerobot_dataset_factory) -> None:
@@ -281,7 +314,7 @@ def test_streaming_persistent_workers_advance_epochs(tmp_path: Path, lerobot_dat
     loader = torch.utils.data.DataLoader(
         streaming,
         batch_size=None,
-        num_workers=2,
+        num_workers=1,
         persistent_workers=True,
     )
     try:
@@ -317,7 +350,7 @@ def test_streaming_worker_exception_propagates_and_workers_stop(
     loader = torch.utils.data.DataLoader(
         streaming,
         batch_size=None,
-        num_workers=2,
+        num_workers=1,
         persistent_workers=True,
     )
     try:
@@ -376,7 +409,7 @@ def test_streaming_worker_resume_reproduces_remaining_stream(
     )
 
     def load(dataset: StreamingLeRobotDataset) -> list[int]:
-        loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, num_workers=2)
+        loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, num_workers=1)
         if batch_size is None:
             return [int(item["index"]) for item in loader]
         return [int(index) for batch in loader for index in batch["index"]]
@@ -450,7 +483,7 @@ def test_streaming_worker_resume_after_epoch_boundary(tmp_path: Path, lerobot_da
         loader = torch.utils.data.DataLoader(
             dataset,
             batch_size=4,
-            num_workers=2,
+            num_workers=1,
             persistent_workers=True,
         )
         try:
@@ -502,7 +535,7 @@ def test_streaming_local_training_step_smoke(tmp_path: Path, lerobot_dataset_fac
         buffer_size=2,
         repeat=True,
     )
-    loader = torch.utils.data.DataLoader(dataset, batch_size=4, num_workers=2)
+    loader = torch.utils.data.DataLoader(dataset, batch_size=4, num_workers=1)
     iterator = iter(loader)
     try:
         batch = next(iterator)