feat(streaming): exact-once epoch coverage for the byte-cache episode pool

The pool path sampled frames with replacement and never guaranteed a full
epoch (episodes rotated on a fixed cadence; frames drawn randomly, none
tracked). Add ExactCoveragePool: a deterministic planner that enumerates
every frame of every episode exactly once per epoch while keeping at most
pool_size episodes resident, so batch mixing stays high (uniform draw over
all remaining frames in the pool) but coverage is complete and reproducible.

Mechanics (the "evict only when all frames sampled" model): episodes are
admitted in a seeded global permutation; each resident episode carries a
seeded frame-index shuffle; each draw picks a resident episode with
probability proportional to its remaining frames and pops one; an episode
is evicted only when its last frame is emitted, then a new one is admitted;
the epoch ends when admission is exhausted and every resident episode drains.
Order is a pure function of (seed, epoch) -> resumable by deterministic
fast-forward. The planner does no I/O and exposes admission_order so callers
can prefetch episodes ahead of the sampling frontier.

Wired into the benchmark as --coverage {sampled,exact}: run_exact_coverage_stream
prefetches stream_prefetch_episodes beyond the frontier so a freshly admitted
episode's bytes are resident before it is drawn, then decodes each frame once,
paced to target.

Tests: 7 planner unit tests (exact-once coverage incl. a 45k-frame epoch,
pool-size bound, per-(seed,epoch) determinism with coverage preserved,
admission/eviction events, coupon-collector mixing, zero-length episodes) and
a mocked-cache structural test of run_exact_coverage_stream asserting
every-frame-once-per-camera plus the prefetch-before-decode invariant.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

scripts/bench_episode_byte_cache.py

@@ -32,6 +32,7 @@ import pyarrow.parquet as pq
 from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
 from lerobot.datasets.episode_video_streaming import (
     EpisodeByteCache,
+    ExactCoveragePool,
     EpisodeVideoManifest,
     NativeHTTPRangeFetcher,
     assert_hf_hub_range_cache_branch,
@@ -77,6 +78,13 @@ def parse_args() -> argparse.Namespace:
         help="Concurrent camera-fetch jobs. Total connections ~= workers x range-subranges; "
         "the HF bucket path saturates around 64 connections per host, so keep the product near 64.",
     )
+    parser.add_argument(
+        "--coverage",
+        choices=["sampled", "exact"],
+        default="sampled",
+        help="sampled: with-replacement random draws (no epoch guarantee). exact: every frame of "
+        "every shard episode decoded exactly once per epoch (deterministic, pool-bounded).",
+    )
     parser.add_argument(
         "--range-subranges",
         type=int,
@@ -350,6 +358,133 @@ def run_pool_random_decode(
     return result
 
 
+def run_exact_coverage_stream(
+    cache: EpisodeByteCache,
+    shard_episodes: Sequence[int],
+    *,
+    pool_size: int,
+    sample_count: int,
+    target_samples_s: float,
+    prefetch_ahead: int,
+    batch_size: int,
+    decode_workers: int,
+    seed: int,
+) -> dict[str, float]:
+    """Streaming keep-up with EXACT, exactly-once frame coverage (a real epoch).
+
+    Unlike run_pool_stream_simulation (with-replacement sampling, no coverage guarantee), this
+    drives the resident pool from ExactCoveragePool: every frame of every shard episode is decoded
+    exactly once, at most ``pool_size`` episodes resident, deterministic from ``seed``. Episodes are
+    prefetched ``prefetch_ahead`` beyond the sampling frontier so a freshly admitted episode's bytes
+    are already resident when it becomes eligible to be drawn.
+    """
+    manifest = cache.manifest
+    cams = manifest.video_keys
+    frame_counts = [
+        (int(ep), min(int(manifest.lookup(ep, c).frame_count) for c in cams)) for ep in shard_episodes
+    ]
+    total_frames = sum(n for _, n in frame_counts)
+    pool = ExactCoveragePool(frame_counts, pool_size, seed=seed)
+    counts = dict(frame_counts)
+    order = pool.admission_order
+    locks: dict[tuple[int, str], threading.Lock] = {}
+
+    prefetch_frontier = 0
+
+    def prefetch_upto(idx: int) -> None:
+        nonlocal prefetch_frontier
+        limit = min(idx, len(order))
+        while prefetch_frontier < limit:
+            cache.submit_prefetch(order[prefetch_frontier])
+            prefetch_frontier += 1
+
+    refill_wait_s = 0.0
+
+    def make_ready(ep: int) -> None:
+        nonlocal refill_wait_s
+        wait_start = time.perf_counter()
+        cache.ensure_ready(ep)
+        refill_wait_s += time.perf_counter() - wait_start
+        _open_resident_decoders(cache, [ep], decode_workers=decode_workers)
+        for c in cams:
+            locks[(ep, c)] = threading.Lock()
+
+    prefetch_upto(pool_size + prefetch_ahead)
+    for ep in pool.newly_admitted:
+        make_ready(ep)
+    pool.newly_admitted.clear()
+
+    decode_pool = ThreadPoolExecutor(max_workers=decode_workers) if decode_workers > 1 else None
+    deadline_miss_s = 0.0
+    samples_done = 0
+    decoded_samples: list[tuple[int, float]] = []
+    epoch_complete = False
+    start = time.perf_counter()
+    try:
+        while samples_done < sample_count and not epoch_complete:
+            batch_start = time.perf_counter()
+            current_batch_size = min(batch_size, sample_count - samples_done)
+            batch: list[tuple[int, float]] = []
+            for _ in range(current_batch_size):
+                try:
+                    ep, frame_index = next(pool)
+                except StopIteration:
+                    epoch_complete = True
+                    break
+                n = counts[ep]
+                batch.append((ep, frame_index / max(n - 1, 1)))
+                if pool.newly_admitted:
+                    for new_ep in pool.newly_admitted:
+                        make_ready(new_ep)
+                    pool.newly_admitted.clear()
+                    prefetch_upto(pool.admitted_count + prefetch_ahead)
+            if not batch:
+                break
+            if decode_pool is not None:
+                futures = [
+                    decode_pool.submit(_decode_training_sample, cache, ep, rel, locks) for ep, rel in batch
+                ]
+                for future in futures:
+                    future.result()
+            else:
+                for ep, rel in batch:
+                    _decode_training_sample(cache, ep, rel, locks)
+            decoded_samples.extend(batch)
+            samples_done += len(batch)
+            target_batch_s = len(batch) / target_samples_s if target_samples_s > 0 else 0.0
+            batch_elapsed = time.perf_counter() - batch_start
+            if target_batch_s > 0 and batch_elapsed < target_batch_s:
+                time.sleep(target_batch_s - batch_elapsed)
+            elif target_batch_s > 0:
+                deadline_miss_s += batch_elapsed - target_batch_s
+    finally:
+        if decode_pool is not None:
+            decode_pool.shutdown(wait=True)
+
+    elapsed = time.perf_counter() - start
+    result = {
+        "coverage_mode": "exact",
+        "target_samples_s": target_samples_s,
+        "actual_samples_s": samples_done / elapsed if elapsed > 0 else float("inf"),
+        "stream_wall_s": elapsed,
+        "refill_wait_s": refill_wait_s,
+        "deadline_miss_s": deadline_miss_s,
+        "samples_done": float(samples_done),
+        "shard_total_frames": float(total_frames),
+        "epoch_complete": 1.0 if epoch_complete else 0.0,
+        "prefetch_ahead": float(prefetch_ahead),
+        "batch_size": float(batch_size),
+        "decode_workers": float(decode_workers),
+        "kept_up": 1.0
+        if samples_done / elapsed >= target_samples_s * 0.98 and deadline_miss_s < elapsed * 0.02
+        else 0.0,
+    }
+    result.update(
+        {f"stream_{k}": v for k, v in _sampling_randomness(decoded_samples, batch_size=batch_size).items()}
+    )
+    return result
+
+
 def run_pool_stream_simulation(
     cache: EpisodeByteCache,
     resident_episodes: Sequence[int],
@@ -765,22 +900,43 @@ def run_fetch_pool(
             _log(
                 f"pool_stream: consuming {args.target_samples_s:.1f} samples/s while prefetching replacements"
             )
-            stream_sim = run_pool_stream_simulation(
-                cache,
-                episodes,
-                dataset_episode_count=dataset_episode_count,
-                num_episodes=benchmark_episode_count,
-                sample_count=args.stream_samples,
-                target_samples_s=args.target_samples_s,
-                samples_per_episode=args.pool_samples_per_episode,
-                prefetch_episodes=args.stream_prefetch_episodes,
-                shard_count=args.distributed_shard_count,
-                shard_index=args.distributed_shard_index,
-                shard_seed=args.seed,
-                batch_size=args.batch_size,
-                decode_workers=args.decode_workers,
-                seed=args.seed + 4,
-            )
+            if args.coverage == "exact":
+                _log("pool_stream: EXACT coverage (every frame once) while prefetching ahead")
+                shard_episodes = _episode_shard(
+                    dataset_episode_count,
+                    benchmark_episode_count,
+                    args.seed,
+                    shard_count=args.distributed_shard_count,
+                    shard_index=args.distributed_shard_index,
+                )
+                stream_sim = run_exact_coverage_stream(
+                    cache,
+                    shard_episodes,
+                    pool_size=len(episodes),
+                    sample_count=args.stream_samples,
+                    target_samples_s=args.target_samples_s,
+                    prefetch_ahead=args.stream_prefetch_episodes,
+                    batch_size=args.batch_size,
+                    decode_workers=args.decode_workers,
+                    seed=args.seed + 5,
+                )
+            else:
+                stream_sim = run_pool_stream_simulation(
+                    cache,
+                    episodes,
+                    dataset_episode_count=dataset_episode_count,
+                    num_episodes=benchmark_episode_count,
+                    sample_count=args.stream_samples,
+                    target_samples_s=args.target_samples_s,
+                    samples_per_episode=args.pool_samples_per_episode,
+                    prefetch_episodes=args.stream_prefetch_episodes,
+                    shard_count=args.distributed_shard_count,
+                    shard_index=args.distributed_shard_index,
+                    shard_seed=args.seed,
+                    batch_size=args.batch_size,
+                    decode_workers=args.decode_workers,
+                    seed=args.seed + 4,
+                )
     byte_count = _bytes_for(manifest, episodes)
     episode_mb = byte_count / len(episodes) / 1024**2
     job_count = max(timings["jobs"], 1.0)

src/lerobot/datasets/episode_video_streaming.py

@@ -981,6 +981,108 @@ class EpisodeVideoManifest:
         )
 
 
+class ExactCoveragePool:
+    """Deterministic, exactly-once frame coverage over a byte-cache episode pool.
+
+    The sampled/with-replacement pool (``run_pool_stream_simulation``) never guarantees a full
+    epoch: frames are drawn randomly and episodes rotate on a fixed cadence. This planner instead
+    enumerates *every frame of every episode exactly once per epoch* while keeping at most
+    ``pool_size`` episodes resident, so batch mixing stays high but coverage is complete and
+    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.
+      - 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
+        one frame.
+      - An episode is evicted only when its last frame is emitted; a new episode is then admitted.
+      - The epoch ends when the admission order is exhausted and every resident episode is drained.
+
+    Newly admitted episodes are surfaced via :attr:`newly_admitted` (drain it to drive prefetch)
+    and evictions via :attr:`evicted` (drain to release cache bytes). The planner does no I/O and
+    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).
+    """
+
+    def __init__(
+        self,
+        episode_frame_counts: Sequence[tuple[int, int]],
+        pool_size: int,
+        *,
+        seed: int,
+        epoch: int = 0,
+    ):
+        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.admission_order: list[int] = order.tolist()
+        self._admit_cursor = 0
+        self._remaining: dict[int, list[int]] = {}
+        self._remaining_total = 0
+        self.newly_admitted: list[int] = []
+        self.evicted: list[int] = []
+        for _ in range(max(1, pool_size)):
+            self._admit_one()
+
+    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)
+
+    @property
+    def remaining_total(self) -> int:
+        return self._remaining_total
+
+    @property
+    def admitted_count(self) -> int:
+        """Number of episodes pulled from the admission order so far (pool fills + rotations)."""
+        return self._admit_cursor
+
+    @property
+    def resident(self) -> list[int]:
+        return list(self._remaining)
+
+    def __iter__(self) -> "ExactCoveragePool":
+        return self
+
+    def __next__(self) -> tuple[int, int]:
+        if self._remaining_total == 0:
+            raise StopIteration
+        # Uniform draw over all remaining frames in the pool: walk the residents by cumulative
+        # 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):
+                chosen = ep
+                break
+            target -= len(frames)
+        frames = self._remaining[chosen]
+        frame_index = frames.pop()
+        self._remaining_total -= 1
+        if not frames:
+            del self._remaining[chosen]
+            self.evicted.append(chosen)
+            self._admit_one()
+        return chosen, frame_index
+
+
 class EpisodeByteCache:
     def __init__(
         self,

tests/datasets/test_exact_coverage_pool.py

@@ -0,0 +1,97 @@
+# 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.
+
+"""ExactCoveragePool: exactly-once frame coverage over a bounded episode pool."""
+
+from collections import Counter
+
+from lerobot.datasets.episode_video_streaming import ExactCoveragePool
+
+EPISODES = [(0, 5), (1, 3), (2, 8), (3, 1), (4, 6), (5, 4), (6, 7), (7, 2)]
+TOTAL = sum(n for _, n in EPISODES)
+EXPECTED = Counter((ep, i) for ep, n in EPISODES for i in range(n))
+
+
+def _drain(pool):
+    out, max_resident = [], 0
+    while True:
+        try:
+            out.append(next(pool))
+        except StopIteration:
+            break
+        max_resident = max(max_resident, len(pool.resident))
+    return out, max_resident
+
+
+def test_exact_once_coverage():
+    out, _ = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=42))
+    assert len(out) == TOTAL
+    assert Counter(out) == EXPECTED  # every (episode, frame) exactly once, no dups/misses
+
+
+def test_pool_never_exceeds_size():
+    _, max_resident = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=42))
+    assert max_resident <= 3
+
+
+def test_deterministic_per_seed_and_epoch():
+    a, _ = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=7))
+    b, _ = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=7))
+    c, _ = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=8))
+    d, _ = _drain(ExactCoveragePool(EPISODES, pool_size=3, seed=7, epoch=1))
+    assert a == b
+    assert a != c and a != d  # seed and epoch both change the order
+    assert Counter(c) == EXPECTED and Counter(d) == EXPECTED  # ... but coverage is preserved
+
+
+def test_admission_and_eviction_events():
+    pool = ExactCoveragePool(EPISODES, pool_size=3, seed=0)
+    admitted_ever, evicted_ever = set(), set()
+    # first three episodes admitted at construction
+    admitted_ever.update(pool.newly_admitted)
+    assert len(admitted_ever) == 3
+    while True:
+        pool.newly_admitted.clear()
+        pool.evicted.clear()
+        try:
+            next(pool)
+        except StopIteration:
+            break
+        admitted_ever.update(pool.newly_admitted)
+        evicted_ever.update(pool.evicted)
+    assert admitted_ever == {ep for ep, _ in EPISODES}  # every episode admitted exactly once
+    # every episode except the pool_size still resident at the end is evicted on exhaustion
+    assert len(evicted_ever) >= len(EPISODES) - 3
+
+
+def test_uniform_mixing_matches_coupon_collector():
+    # 64 equal episodes, pool 64, first 64 draws -> ~64*(1-(1-1/64)^64) ~= 41 distinct
+    big = [(e, 100) for e in range(64)]
+    pool = ExactCoveragePool(big, pool_size=64, seed=0)
+    head = [next(pool)[0] for _ in range(64)]
+    assert len(set(head)) >= 30  # far above sequential (=1); ~41 expected
+
+
+def test_large_epoch_bounded_and_complete():
+    big = [(e, 90) for e in range(500)]
+    out, max_resident = _drain(ExactCoveragePool(big, pool_size=64, seed=3))
+    assert len(out) == 500 * 90
+    assert len(set(out)) == 500 * 90  # exactly once
+    assert max_resident <= 64
+
+
+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})