From 3d0e8681f3c78915f0807476938e27190d370e14 Mon Sep 17 00:00:00 2001 From: pepijn Date: Thu, 11 Jun 2026 15:31:15 +0000 Subject: [PATCH] refactor(datasets): make EpisodeAwareSampler always deterministic With Feistel gone, deterministic and legacy modes were both just torch.randperm and the deterministic path strictly dominated (reproducible across ranks via the (seed, epoch) seed, no accelerate generator sync, resumable). Collapse to a single path and drop the redundant flag: - remove the `deterministic` and `generator` constructor args, `_iter_default`, and `_require_deterministic`; `set_epoch` / `state_dict` / `load_state_dict` are now unconditional - remove the `deterministic_sampler` train config field and the legacy generator branch in lerobot_train.py (non-streaming map datasets always use the sampler) - drop the now-obsolete generator/legacy tests Note: removes the `generator` kwarg from EpisodeAwareSampler (back-compat break vs main); the order is now a pure function of (seed, epoch), so no cross-rank RNG sync is needed. Co-authored-by: Cursor --- src/lerobot/configs/train.py | 4 -- src/lerobot/datasets/sampler.py | 47 +++----------- src/lerobot/scripts/lerobot_train.py | 20 +----- tests/datasets/test_sampler.py | 94 +++++++--------------------- 4 files changed, 35 insertions(+), 130 deletions(-) diff --git a/src/lerobot/configs/train.py b/src/lerobot/configs/train.py index d074bc8a9..bac1a946b 100644 --- a/src/lerobot/configs/train.py +++ b/src/lerobot/configs/train.py @@ -99,10 +99,6 @@ class TrainPipelineConfig(HubMixin): batch_size: int = 8 prefetch_factor: int = 4 persistent_workers: bool = True - # Deterministic data order (pure function of seed and epoch): immune to cross-rank RNG - # desync and enables sample-exact resume. Set to false for the legacy RNG-based shuffle. - # Ignored when dataset.streaming is enabled. - deterministic_sampler: bool = True steps: int = 100_000 eval_freq: int = 20_000 log_freq: int = 200 diff --git a/src/lerobot/datasets/sampler.py b/src/lerobot/datasets/sampler.py index 170bbfb9f..9d102e58f 100644 --- a/src/lerobot/datasets/sampler.py +++ b/src/lerobot/datasets/sampler.py @@ -29,15 +29,12 @@ class EpisodeAwareSampler: Logical positions map to frame indices on the fly (O(num_episodes) construction memory) instead of materializing a Python list of every frame index. - By default (`deterministic=True`) each epoch is shuffled with a `torch.randperm` seeded from - `(seed, epoch)`, so the data order is a pure function of `(seed, epoch)`: it reproduces on - every rank without synchronizing the global RNG, and `state_dict` / `load_state_dict` resume - a run sample-exactly by regenerating the epoch's permutation and continuing from the saved - offset. Each call to `__iter__` advances the epoch. During a resumed epoch, `__len__` still - reports the full length. - - With `deterministic=False`, shuffling uses `torch.randperm` driven by `generator` instead - (accelerate synchronizes the generator across ranks when preparing the dataloader). + Each epoch is shuffled with a `torch.randperm` seeded from `(seed, epoch)`, so the data order + is a pure function of `(seed, epoch)`: it reproduces on every rank without synchronizing the + global RNG (no `generator` to sync across distributed ranks), and `state_dict` / + `load_state_dict` resume a run sample-exactly by regenerating the epoch's permutation and + continuing from the saved offset. Each call to `__iter__` advances the epoch. During a + resumed epoch, `__len__` still reports the full length. """ def __init__( @@ -48,8 +45,6 @@ class EpisodeAwareSampler: drop_n_first_frames: int = 0, drop_n_last_frames: int = 0, shuffle: bool = False, - generator: torch.Generator | None = None, - deterministic: bool = True, seed: int = 0, ): """ @@ -60,17 +55,12 @@ class EpisodeAwareSampler: drop_n_first_frames: Frames to drop from the start of each episode. drop_n_last_frames: Frames to drop from the end of each episode. shuffle: Whether to shuffle the indices. - generator: Generator for non-deterministic shuffling (global torch RNG when None). - deterministic: Seed the shuffle from `(seed, epoch)` for reproducible, resumable - order instead of a `generator`-driven `torch.randperm`. - seed: Seed the deterministic permutation is derived from (together with the epoch). + seed: Seed the permutation is derived from (together with the epoch). """ if drop_n_first_frames < 0: raise ValueError(f"drop_n_first_frames must be >= 0, got {drop_n_first_frames}") if drop_n_last_frames < 0: raise ValueError(f"drop_n_last_frames must be >= 0, got {drop_n_last_frames}") - if deterministic and generator is not None: - raise ValueError("generator is unused in deterministic mode; pass seed instead.") from_indices = np.asarray(dataset_from_indices, dtype=np.int64) to_indices = np.asarray(dataset_to_indices, dtype=np.int64) @@ -107,8 +97,6 @@ class EpisodeAwareSampler: self._cum_lengths = np.cumsum(lengths[used]) self._num_frames = int(self._cum_lengths[-1]) self.shuffle = shuffle - self.generator = generator - self.deterministic = deterministic self.seed = seed self._epoch = 0 self._start_index = 0 @@ -119,22 +107,15 @@ class EpisodeAwareSampler: return [self._frame_index(k) for k in range(self._num_frames)] def set_epoch(self, epoch: int) -> None: - self._require_deterministic("set_epoch") self._epoch = epoch def state_dict(self) -> dict: - self._require_deterministic("state_dict") return {"epoch": self._epoch, "start_index": self._start_index} def load_state_dict(self, state: dict) -> None: - self._require_deterministic("load_state_dict") self._epoch = state["epoch"] self._start_index = state["start_index"] - def _require_deterministic(self, method: str) -> None: - if not self.deterministic: - raise RuntimeError(f"{method} requires deterministic=True: an RNG order cannot be sought.") - def _epoch_generator(self, epoch: int) -> torch.Generator: # Derive a per-epoch seed from (seed, epoch) so the permutation is a pure function of both # and reproduces identically on every rank without touching the global RNG. @@ -147,23 +128,13 @@ class EpisodeAwareSampler: return int(self._starts[episode]) + position_in_episode def __iter__(self) -> Iterator[int]: - if not self.deterministic: - return self._iter_default() # Advance epoch state eagerly, not on first consumption of the generator. epoch, start = self._epoch, self._start_index self._epoch += 1 self._start_index = 0 - return self._iter_deterministic_epoch(epoch, start) + return self._iter_epoch(epoch, start) - def _iter_default(self) -> Iterator[int]: - if self.shuffle: - for i in torch.randperm(self._num_frames, generator=self.generator): - yield self._frame_index(int(i)) - else: - for k in range(self._num_frames): - yield self._frame_index(k) - - def _iter_deterministic_epoch(self, epoch: int, start: int) -> Iterator[int]: + def _iter_epoch(self, epoch: int, start: int) -> Iterator[int]: if self.shuffle: order = torch.randperm(self._num_frames, generator=self._epoch_generator(epoch)) for k in range(start, self._num_frames): diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index bf056b05b..6787d1c4a 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -388,8 +388,9 @@ 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 cfg.deterministic_sampler and not cfg.dataset.streaming: - # Deterministic data order: no cross-rank RNG sync needed, sample-exact resume. + if not cfg.dataset.streaming: + # Deterministic data order (pure function of seed and epoch): no cross-rank RNG sync + # needed and sample-exact resume. shuffle = False sampler = EpisodeAwareSampler( dataset.meta.episodes["dataset_from_index"], @@ -417,21 +418,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): f"Resuming data order at epoch {sampler_state['epoch']}, " f"sample {sampler_state['start_index']}" ) - elif hasattr(active_cfg, "drop_n_last_frames"): - shuffle = False - # Legacy RNG shuffle: a dedicated generator lets accelerate synchronize it across ranks. - sampler_generator = torch.Generator() - if cfg.seed is not None: - sampler_generator.manual_seed(cfg.seed) - sampler = EpisodeAwareSampler( - dataset.meta.episodes["dataset_from_index"], - dataset.meta.episodes["dataset_to_index"], - episode_indices_to_use=dataset.episodes, - drop_n_last_frames=active_cfg.drop_n_last_frames, - shuffle=True, - deterministic=False, - generator=sampler_generator, - ) else: shuffle = True sampler = None diff --git a/tests/datasets/test_sampler.py b/tests/datasets/test_sampler.py index 066ef7733..3bce6e7cb 100644 --- a/tests/datasets/test_sampler.py +++ b/tests/datasets/test_sampler.py @@ -114,34 +114,17 @@ def test_shuffle(): assert set(sampler) == {0, 1, 2, 3, 4, 5} -def test_shuffle_with_generator_is_deterministic(): - # Two samplers shuffling with same-seed generators must yield identical permutations. - # This is what keeps batch shards disjoint across ranks in distributed training, where - # accelerate synchronizes the sampler's generator state instead of the global torch RNG. - sampler_a = EpisodeAwareSampler( - [0], [6], shuffle=True, deterministic=False, generator=torch.Generator().manual_seed(42) - ) - sampler_b = EpisodeAwareSampler( - [0], [6], shuffle=True, deterministic=False, generator=torch.Generator().manual_seed(42) - ) - assert list(sampler_a) == list(sampler_b) - +def test_shuffle_is_reproducible_across_instances(): + # The order is a pure function of (seed, epoch), so two fresh samplers (e.g. two ranks) + # produce the same permutation without any generator synchronization. + sampler_a = EpisodeAwareSampler([0], [6], shuffle=True, seed=42) + sampler_b = EpisodeAwareSampler([0], [6], shuffle=True, seed=42) + epoch_0 = list(sampler_a) + assert list(sampler_b) == epoch_0 # Desyncing the global RNG must not affect the permutation. - sampler_c = EpisodeAwareSampler( - [0], [6], shuffle=True, deterministic=False, generator=torch.Generator().manual_seed(42) - ) - order_before = list(sampler_c) - sampler_c.generator.manual_seed(42) + sampler_c = EpisodeAwareSampler([0], [6], shuffle=True, seed=42) torch.randperm(1000) # consume global RNG, as rank-asymmetric code (e.g. eval) would - assert list(sampler_c) == order_before - - -def test_generator_attribute_defaults_to_none(): - # accelerate detects synchronizable samplers via `hasattr(sampler, "generator")`, - # so the attribute must exist even when no generator is passed. - sampler = EpisodeAwareSampler([0], [6], shuffle=True, deterministic=False) - assert sampler.generator is None - assert set(sampler) == {0, 1, 2, 3, 4, 5} + assert list(sampler_c) == epoch_0 def test_negative_drop_first_frames_raises(): @@ -169,54 +152,23 @@ def test_partial_episode_drop_warns(caplog): assert "Episode 0" in caplog.text -# --- deterministic mode (seeded torch.randperm) --- - -from functools import partial # noqa: E402 +# --- seeded (seed, epoch) shuffling, resume, and state --- from lerobot.datasets.sampler import compute_sampler_state # noqa: E402 -deterministic_sampler = partial(EpisodeAwareSampler, deterministic=True) - - EPISODE_BOUNDS = ([0, 2, 3], [2, 3, 6]) # episodes of 2, 1 and 3 frames -def test_deterministic_mode_unshuffled_matches_default_mode(): - for kwargs in ( - {}, - {"drop_n_first_frames": 1}, - {"drop_n_last_frames": 1}, - {"episode_indices_to_use": [0, 2]}, - ): - reference = EpisodeAwareSampler(*EPISODE_BOUNDS, shuffle=False, **kwargs) - sampler = deterministic_sampler(*EPISODE_BOUNDS, shuffle=False, **kwargs) - assert list(sampler) == list(reference), kwargs - assert len(sampler) == len(reference), kwargs - - -def test_deterministic_mode_rejects_generator(): - with pytest.raises(ValueError, match="generator is unused in deterministic mode"): - deterministic_sampler(*EPISODE_BOUNDS, shuffle=True, generator=torch.Generator()) - - -def test_state_methods_require_deterministic_mode(): - sampler = EpisodeAwareSampler(*EPISODE_BOUNDS, shuffle=True, deterministic=False) - with pytest.raises(RuntimeError, match="deterministic=True"): - sampler.set_epoch(1) - with pytest.raises(RuntimeError, match="deterministic=True"): - sampler.state_dict() - - @pytest.mark.parametrize("num_frames", [1, 2, 3, 37, 64, 100]) def test_deterministic_sampler_shuffle_is_permutation(num_frames): for seed in (0, 1, 1234): - sampler = deterministic_sampler([0], [num_frames], shuffle=True, seed=seed) + sampler = EpisodeAwareSampler([0], [num_frames], shuffle=True, seed=seed) assert sorted(sampler) == list(range(num_frames)) def test_deterministic_sampler_epochs_reproduce_and_differ(): - sampler_a = deterministic_sampler([0], [100], shuffle=True, seed=42) - sampler_b = deterministic_sampler([0], [100], shuffle=True, seed=42) + sampler_a = EpisodeAwareSampler([0], [100], shuffle=True, seed=42) + sampler_b = EpisodeAwareSampler([0], [100], shuffle=True, seed=42) epoch_0 = list(sampler_a) assert list(sampler_b) == epoch_0 # same (seed, epoch) -> same order on any process epoch_1 = list(sampler_a) # __iter__ auto-advances the epoch @@ -224,15 +176,15 @@ def test_deterministic_sampler_epochs_reproduce_and_differ(): assert sorted(epoch_1) == sorted(epoch_0) sampler_a.set_epoch(0) assert list(sampler_a) == epoch_0 - assert list(deterministic_sampler([0], [100], shuffle=True, seed=7)) != epoch_0 + assert list(EpisodeAwareSampler([0], [100], shuffle=True, seed=7)) != epoch_0 def test_deterministic_sampler_resume_mid_epoch(): - reference = deterministic_sampler(*EPISODE_BOUNDS, shuffle=True, seed=42) + reference = EpisodeAwareSampler(*EPISODE_BOUNDS, shuffle=True, seed=42) epoch_0 = list(reference) epoch_1 = list(reference) for start in (0, 1, 4, len(epoch_0)): - resumed = deterministic_sampler(*EPISODE_BOUNDS, shuffle=True, seed=42) + resumed = EpisodeAwareSampler(*EPISODE_BOUNDS, shuffle=True, seed=42) resumed.load_state_dict({"epoch": 0, "start_index": start}) assert list(resumed) == epoch_0[start:] # the resumed sampler continues into the same epoch 1 as the uninterrupted one @@ -243,7 +195,7 @@ def test_deterministic_sampler_construction_stores_only_boundaries(): # Construction is O(num_episodes), not O(num_frames): a million-frame single episode # instantiates from just its boundaries without materializing a per-frame index list. num_frames = 1_000_000 - sampler = deterministic_sampler([0], [num_frames], shuffle=True, seed=0) + sampler = EpisodeAwareSampler([0], [num_frames], shuffle=True, seed=0) assert len(sampler) == num_frames assert sampler._starts.shape == (1,) and sampler._cum_lengths.shape == (1,) @@ -252,27 +204,27 @@ def test_deterministic_sampler_resume_is_exact_at_scale(): # Seeded randperm makes resume sample-exact at non-trivial sizes: regenerating the epoch's # permutation and slicing from the saved offset reproduces the remaining order exactly. num_frames = 100_000 - reference = deterministic_sampler([0], [num_frames], shuffle=True, seed=0) + reference = EpisodeAwareSampler([0], [num_frames], shuffle=True, seed=0) epoch_0 = list(reference) assert sorted(epoch_0) == list(range(num_frames)) start = num_frames - 5 - resumed = deterministic_sampler([0], [num_frames], shuffle=True, seed=0) + resumed = EpisodeAwareSampler([0], [num_frames], shuffle=True, seed=0) resumed.load_state_dict({"epoch": 0, "start_index": start}) assert list(resumed) == epoch_0[start:] def test_deterministic_sampler_validation_matches_episode_aware(): with pytest.raises(ValueError, match="drop_n_first_frames must be >= 0"): - deterministic_sampler([0], [10], drop_n_first_frames=-1) + EpisodeAwareSampler([0], [10], drop_n_first_frames=-1) with pytest.raises(ValueError, match="drop_n_last_frames must be >= 0"): - deterministic_sampler([0], [10], drop_n_last_frames=-1) + EpisodeAwareSampler([0], [10], drop_n_last_frames=-1) with pytest.raises(ValueError, match="No valid frames remain"): - deterministic_sampler([0, 1, 2], [1, 2, 3], drop_n_first_frames=1) + EpisodeAwareSampler([0, 1, 2], [1, 2, 3], drop_n_first_frames=1) def test_deterministic_sampler_partial_episode_drop_warns(caplog): with caplog.at_level(logging.WARNING, logger="lerobot.datasets.sampler"): - sampler = deterministic_sampler([0, 1], [1, 6], drop_n_first_frames=1, shuffle=False) + sampler = EpisodeAwareSampler([0, 1], [1, 6], drop_n_first_frames=1, shuffle=False) assert list(sampler) == [2, 3, 4, 5] assert "Episode 0" in caplog.text