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 <cursoragent@cursor.com>

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

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
+    Each epoch is shuffled with a `torch.randperm` seeded from `(seed, epoch)`, so the data order
-    `(seed, epoch)`, so the data order is a pure function of `(seed, epoch)`: it reproduces on
+    is a pure function of `(seed, epoch)`: it reproduces on every rank without synchronizing the
-    every rank without synchronizing the global RNG, and `state_dict` / `load_state_dict` resume
+    global RNG (no `generator` to sync across distributed ranks), and `state_dict` /
-    a run sample-exactly by regenerating the epoch's permutation and continuing from the saved
+    `load_state_dict` resume a run sample-exactly by regenerating the epoch's permutation and
-    offset. Each call to `__iter__` advances the epoch. During a resumed epoch, `__len__` still
+    continuing from the saved offset. Each call to `__iter__` advances the epoch. During a
-    reports the full length.
+    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).
     """
 
     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).
+            seed: Seed the permutation is derived from (together with the epoch).
-            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).
         """
         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]:
+    def _iter_epoch(self, epoch: int, start: int) -> 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]:
         if self.shuffle:
             order = torch.randperm(self._num_frames, generator=self._epoch_generator(epoch))
             for k in range(start, self._num_frames):

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:
+    if not cfg.dataset.streaming:
-        # Deterministic data order: no cross-rank RNG sync needed, sample-exact resume.
+        # 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

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():
+def test_shuffle_is_reproducible_across_instances():
-    # Two samplers shuffling with same-seed generators must yield identical permutations.
+    # The order is a pure function of (seed, epoch), so two fresh samplers (e.g. two ranks)
-    # This is what keeps batch shards disjoint across ranks in distributed training, where
+    # produce the same permutation without any generator synchronization.
-    # accelerate synchronizes the sampler's generator state instead of the global torch RNG.
+    sampler_a = EpisodeAwareSampler([0], [6], shuffle=True, seed=42)
-    sampler_a = EpisodeAwareSampler(
+    sampler_b = EpisodeAwareSampler([0], [6], shuffle=True, seed=42)
-        [0], [6], shuffle=True, deterministic=False, generator=torch.Generator().manual_seed(42)
+    epoch_0 = list(sampler_a)
-    )
+    assert list(sampler_b) == epoch_0
-    sampler_b = EpisodeAwareSampler(
-        [0], [6], shuffle=True, deterministic=False, generator=torch.Generator().manual_seed(42)
-    )
-    assert list(sampler_a) == list(sampler_b)
-
     # Desyncing the global RNG must not affect the permutation.
-    sampler_c = EpisodeAwareSampler(
+    sampler_c = EpisodeAwareSampler([0], [6], shuffle=True, seed=42)
-        [0], [6], shuffle=True, deterministic=False, generator=torch.Generator().manual_seed(42)
-    )
-    order_before = list(sampler_c)
-    sampler_c.generator.manual_seed(42)
     torch.randperm(1000)  # consume global RNG, as rank-asymmetric code (e.g. eval) would
-    assert list(sampler_c) == order_before
+    assert list(sampler_c) == epoch_0
-
-
-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}
 
 
 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) ---
+# --- seeded (seed, epoch) shuffling, resume, and state ---
-
-from functools import partial  # noqa: E402
 
 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_a = EpisodeAwareSampler([0], [100], shuffle=True, seed=42)
-    sampler_b = deterministic_sampler([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