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