refactor(datasets): use seeded torch.randperm instead of Feistel in EpisodeAwareSampler

Drop the Feistel permutation (and its SplitMix64 hash / cycle-walking) in favor of a
torch.randperm seeded from (seed, epoch). The deterministic mode keeps its key properties
- data order is a pure function of (seed, epoch), so it reproduces on every rank with no
  global-RNG synchronization, and
- state_dict / load_state_dict still resume sample-exactly, now by regenerating the epoch's
  permutation and slicing from the saved offset.

Construction stays O(num_episodes) (only episode boundaries are stored, never a per-frame
index list). The trade-off vs Feistel: the per-epoch shuffle is again O(num_frames) memory
(the randperm tensor) and no longer O(1)-seekable, in exchange for ~30 fewer LOC and a truly
uniform shuffle. Tests updated: the trillion-frame O(1) test is replaced with a
boundary-storage check and a scale resume-exactness test.

Co-authored-by: Cursor <cursoragent@cursor.com>

src/lerobot/datasets/sampler.py

@@ -22,38 +22,21 @@ import torch
 
 logger = logging.getLogger(__name__)
 
-_MASK_64 = (1 << 64) - 1
-_FEISTEL_ROUNDS = 4
-# Cycle-walking converges in <4 expected steps on the chosen domain; this bound is a generous
-# safety net that should never be hit in practice.
-_MAX_CYCLE_WALK_STEPS = 100
-
-
-def _mix64(x: int) -> int:
-    """SplitMix64 finalizer (64-bit integer hash)."""
-    x = (x + 0x9E3779B97F4A7C15) & _MASK_64
-    x ^= x >> 30
-    x = (x * 0xBF58476D1CE4E5B9) & _MASK_64
-    x ^= x >> 27
-    x = (x * 0x94D049BB133111EB) & _MASK_64
-    x ^= x >> 31
-    return x
-
 
 class EpisodeAwareSampler:
-    """Sampler over episode frames with O(num_episodes) memory.
+    """Sampler over episode frames that stores only per-episode boundaries.
 
-    Only episode boundaries are stored; logical positions map to frame indices on the fly, so
+    Logical positions map to frame indices on the fly (O(num_episodes) construction memory)
-    memory does not grow with the number of frames.
+    instead of materializing a Python list of every frame index.
 
-    By default (`deterministic=True`) shuffling uses a seeded Feistel permutation over
+    By default (`deterministic=True`) each epoch is shuffled with a `torch.randperm` seeded from
-    `[0, num_frames)`: the data order is a pure function of `(seed, epoch)`, needs no RNG
+    `(seed, epoch)`, so the data order is a pure function of `(seed, epoch)`: it reproduces on
-    synchronization across distributed ranks, and any position can be sought in O(1), enabling
+    every rank without synchronizing the global RNG, and `state_dict` / `load_state_dict` resume
-    sample-exact resume via `state_dict` / `load_state_dict`. Each completed `__iter__`
+    a run sample-exactly by regenerating the epoch's permutation and continuing from the saved
-    advances the epoch. The shuffle is pseudo-random rather than truly uniform — the standard
+    offset. Each call to `__iter__` advances the epoch. During a resumed epoch, `__len__` still
-    large-scale trade-off. During a resumed epoch, `__len__` still reports the full length.
+    reports the full length.
 
-    With `deterministic=False`, shuffling falls back to `torch.randperm` driven by `generator`
+    With `deterministic=False`, shuffling uses `torch.randperm` driven by `generator` instead
     (accelerate synchronizes the generator across ranks when preparing the dataloader).
     """
 
@@ -78,7 +61,8 @@ class EpisodeAwareSampler:
             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: Use the seeded Feistel permutation instead of `torch.randperm`.
+            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:
@@ -129,12 +113,6 @@ class EpisodeAwareSampler:
         self._epoch = 0
         self._start_index = 0
 
-        # Smallest even-bit-width power-of-two domain >= num_frames: equal Feistel halves,
-        # cycle-walking converges in <4 expected steps.
-        bits = max((self._num_frames - 1).bit_length(), 2)
-        self._half_bits = (bits + 1) // 2
-        self._half_mask = (1 << self._half_bits) - 1
-
     @property
     def indices(self) -> list[int]:
         """Materialized frame indices in unshuffled order; O(num_frames), introspection only."""
@@ -157,28 +135,11 @@ class EpisodeAwareSampler:
         if not self.deterministic:
             raise RuntimeError(f"{method} requires deterministic=True: an RNG order cannot be sought.")
 
-    def _round_keys(self, epoch: int) -> list[int]:
+    def _epoch_generator(self, epoch: int) -> torch.Generator:
-        state = _mix64(_mix64(self.seed) ^ _mix64(epoch))
+        # Derive a per-epoch seed from (seed, epoch) so the permutation is a pure function of both
-        keys = []
+        # and reproduces identically on every rank without touching the global RNG.
-        for _ in range(_FEISTEL_ROUNDS):
+        epoch_seed = int(np.random.SeedSequence([self.seed, epoch]).generate_state(1, dtype=np.uint64)[0])
-            state = _mix64(state)
+        return torch.Generator().manual_seed(epoch_seed)
-            keys.append(state)
-        return keys
-
-    def _permute(self, index: int, keys: list[int]) -> int:
-        # Feistel network with cycle-walking: a bijection on [0, num_frames).
-        half_bits, half_mask = self._half_bits, self._half_mask
-        for _ in range(_MAX_CYCLE_WALK_STEPS):
-            left, right = index >> half_bits, index & half_mask
-            for key in keys:
-                left, right = right, left ^ (_mix64(right ^ key) & half_mask)
-            index = (left << half_bits) | right
-            if index < self._num_frames:
-                return index
-        raise RuntimeError(
-            f"Feistel cycle-walking did not converge within {_MAX_CYCLE_WALK_STEPS} steps; "
-            "this should never happen for a valid domain."
-        )
 
     def _frame_index(self, position: int) -> int:
         episode = int(np.searchsorted(self._cum_lengths, position, side="right"))
@@ -203,9 +164,13 @@ class EpisodeAwareSampler:
                 yield self._frame_index(k)
 
     def _iter_deterministic_epoch(self, epoch: int, start: int) -> Iterator[int]:
-        keys = self._round_keys(epoch) if self.shuffle else None
+        if self.shuffle:
-        for k in range(start, self._num_frames):
+            order = torch.randperm(self._num_frames, generator=self._epoch_generator(epoch))
-            yield self._frame_index(self._permute(k, keys) if self.shuffle else k)
+            for k in range(start, self._num_frames):
+                yield self._frame_index(int(order[k]))
+        else:
+            for k in range(start, self._num_frames):
+                yield self._frame_index(k)
 
     def __len__(self) -> int:
         return self._num_frames

tests/datasets/test_sampler.py

@@ -169,7 +169,7 @@ def test_partial_episode_drop_warns(caplog):
     assert "Episode 0" in caplog.text
 
 
-# --- deterministic mode (seeded Feistel permutation) ---
+# --- deterministic mode (seeded torch.randperm) ---
 
 from functools import partial  # noqa: E402
 
@@ -239,19 +239,26 @@ def test_deterministic_sampler_resume_mid_epoch():
         assert list(resumed) == epoch_1
 
 
-def test_deterministic_sampler_constant_memory():
+def test_deterministic_sampler_construction_stores_only_boundaries():
-    # A trillion-frame dataset must instantiate instantly and seek anywhere in O(1):
+    # Construction is O(num_episodes), not O(num_frames): a million-frame single episode
-    # only per-episode boundaries are stored, never per-frame indices.
+    # instantiates from just its boundaries without materializing a per-frame index list.
-    num_frames = 10**12
+    num_frames = 1_000_000
     sampler = deterministic_sampler([0], [num_frames], shuffle=True, seed=0)
     assert len(sampler) == num_frames
-    sampler.load_state_dict({"epoch": 3, "start_index": num_frames - 3})
+    assert sampler._starts.shape == (1,) and sampler._cum_lengths.shape == (1,)
-    # Collect via the iterator: list(sampler) would call PyObject_LengthHint -> sampler.__len__
+
-    # (the full epoch length, here 10**12) and pre-allocate that many slots before iterating. The
+
-    # iterator itself exposes no length hint, so this stays O(1) like the resumed epoch it drains.
+def test_deterministic_sampler_resume_is_exact_at_scale():
-    tail = list(iter(sampler))
+    # Seeded randperm makes resume sample-exact at non-trivial sizes: regenerating the epoch's
-    assert len(tail) == 3
+    # permutation and slicing from the saved offset reproduces the remaining order exactly.
-    assert all(0 <= idx < num_frames for idx in tail)
+    num_frames = 100_000
+    reference = deterministic_sampler([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.load_state_dict({"epoch": 0, "start_index": start})
+    assert list(resumed) == epoch_0[start:]
 
 
 def test_deterministic_sampler_validation_matches_episode_aware():