refactor(datasets): fold deterministic mode into EpisodeAwareSampler

Instead of a parallel DeterministicEpisodeAwareSampler class, extend the
existing EpisodeAwareSampler with a deterministic=True mode (seeded
Feistel permutation, epoch auto-advance, state_dict/load_state_dict).

The default mode is behavior-identical: same torch.randperm consumption
and the same generator contract accelerate synchronizes; the O(N) Python
index list is replaced by O(num_episodes) boundary arrays in both modes,
with `indices` kept as a back-compat property. Passing a generator
together with deterministic=True is rejected, and the state/seek methods
raise outside deterministic mode.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

tests/datasets/test_sampler.py

@@ -163,17 +163,19 @@ def test_partial_episode_drop_warns(caplog):
     assert "Episode 0" in caplog.text
 
 
-# --- DeterministicEpisodeAwareSampler ---
+# --- deterministic mode (seeded Feistel permutation) ---
+
+from functools import partial  # noqa: E402
+
+from lerobot.datasets.sampler import compute_sampler_state  # noqa: E402
+
+deterministic_sampler = partial(EpisodeAwareSampler, deterministic=True)
 
-from lerobot.datasets.sampler import (  # noqa: E402
-    DeterministicEpisodeAwareSampler,
-    compute_sampler_state,
-)
 
 EPISODE_BOUNDS = ([0, 2, 3], [2, 3, 6])  # episodes of 2, 1 and 3 frames
 
 
-def test_deterministic_sampler_unshuffled_matches_episode_aware():
+def test_deterministic_mode_unshuffled_matches_default_mode():
     for kwargs in (
         {},
         {"drop_n_first_frames": 1},
@@ -181,21 +183,34 @@ def test_deterministic_sampler_unshuffled_matches_episode_aware():
         {"episode_indices_to_use": [0, 2]},
     ):
         reference = EpisodeAwareSampler(*EPISODE_BOUNDS, shuffle=False, **kwargs)
-        sampler = DeterministicEpisodeAwareSampler(*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)
+    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 = DeterministicEpisodeAwareSampler([0], [num_frames], shuffle=True, seed=seed)
+        sampler = deterministic_sampler([0], [num_frames], shuffle=True, seed=seed)
         assert sorted(sampler) == list(range(num_frames))
 
 
 def test_deterministic_sampler_epochs_reproduce_and_differ():
-    sampler_a = DeterministicEpisodeAwareSampler([0], [100], shuffle=True, seed=42)
-    sampler_b = DeterministicEpisodeAwareSampler([0], [100], shuffle=True, seed=42)
+    sampler_a = deterministic_sampler([0], [100], shuffle=True, seed=42)
+    sampler_b = deterministic_sampler([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
@@ -203,15 +218,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(DeterministicEpisodeAwareSampler([0], [100], shuffle=True, seed=7)) != epoch_0
+    assert list(deterministic_sampler([0], [100], shuffle=True, seed=7)) != epoch_0
 
 
 def test_deterministic_sampler_resume_mid_epoch():
-    reference = DeterministicEpisodeAwareSampler(*EPISODE_BOUNDS, shuffle=True, seed=42)
+    reference = deterministic_sampler(*EPISODE_BOUNDS, shuffle=True, seed=42)
     epoch_0 = list(reference)
     epoch_1 = list(reference)
     for start in (0, 1, 4, len(epoch_0)):
-        resumed = DeterministicEpisodeAwareSampler(*EPISODE_BOUNDS, shuffle=True, seed=42)
+        resumed = deterministic_sampler(*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
@@ -222,7 +237,7 @@ def test_deterministic_sampler_constant_memory():
     # A trillion-frame dataset must instantiate instantly and seek anywhere in O(1):
     # only per-episode boundaries are stored, never per-frame indices.
     num_frames = 10**12
-    sampler = DeterministicEpisodeAwareSampler([0], [num_frames], shuffle=True, seed=0)
+    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})
     tail = list(sampler)
@@ -232,16 +247,16 @@ def test_deterministic_sampler_constant_memory():
 
 def test_deterministic_sampler_validation_matches_episode_aware():
     with pytest.raises(ValueError, match="drop_n_first_frames must be >= 0"):
-        DeterministicEpisodeAwareSampler([0], [10], drop_n_first_frames=-1)
+        deterministic_sampler([0], [10], drop_n_first_frames=-1)
     with pytest.raises(ValueError, match="drop_n_last_frames must be >= 0"):
-        DeterministicEpisodeAwareSampler([0], [10], drop_n_last_frames=-1)
+        deterministic_sampler([0], [10], drop_n_last_frames=-1)
     with pytest.raises(ValueError, match="No valid frames remain"):
-        DeterministicEpisodeAwareSampler([0, 1, 2], [1, 2, 3], drop_n_first_frames=1)
+        deterministic_sampler([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 = DeterministicEpisodeAwareSampler([0, 1], [1, 6], drop_n_first_frames=1, shuffle=False)
+        sampler = deterministic_sampler([0, 1], [1, 6], drop_n_first_frames=1, shuffle=False)
     assert list(sampler) == [2, 3, 4, 5]
     assert "Episode 0" in caplog.text