mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-06 09:37:06 +00:00
Pepijn
06aa6a0425
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>
98 lines
3.7 KiB
Python
98 lines
3.7 KiB
Python
# 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})
|