feat(datasets): add DeterministicEpisodeAwareSampler with O(1) memory and sample-exact resume

Add a sampler that never materializes frame indices: it stores only
per-episode boundaries (numpy, a few bytes per episode) and maps logical
positions to frame indices on the fly with searchsorted. Shuffling uses a
seeded Feistel permutation over [0, num_frames) (cycle-walking to the
exact domain), so the data order is a pure function of (seed, epoch):

- no RNG state to synchronize across distributed ranks,
- constant memory and zero epoch-boundary cost at any dataset size,
- O(1) seek to any position, enabling sample-exact resume.

Opt in with --deterministic_sampler=true. On resume, lerobot-train maps
the checkpointed step back to (epoch, start_index) via
compute_sampler_state and continues at the exact sample where the run
left off (up to accelerate's even_batches padding at epoch boundaries).
The shuffle is pseudo-random rather than a true uniform permutation, the
standard trade-off in large-scale training loaders.

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

tests/datasets/test_sampler.py

@@ -161,3 +161,109 @@ def test_partial_episode_drop_warns(caplog):
     # Episode 0 is skipped (1 frame, drop 1), Episode 1 keeps frames 2-5
     assert sampler.indices == [2, 3, 4, 5]
     assert "Episode 0" in caplog.text
+
+
+# --- DeterministicEpisodeAwareSampler ---
+
+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():
+    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 = DeterministicEpisodeAwareSampler(*EPISODE_BOUNDS, shuffle=False, **kwargs)
+        assert list(sampler) == list(reference), kwargs
+        assert len(sampler) == len(reference), kwargs
+
+
+@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)
+        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)
+    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
+    assert epoch_1 != epoch_0
+    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
+
+
+def test_deterministic_sampler_resume_mid_epoch():
+    reference = DeterministicEpisodeAwareSampler(*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.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
+        assert list(resumed) == epoch_1
+
+
+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)
+    assert len(sampler) == num_frames
+    sampler.load_state_dict({"epoch": 3, "start_index": num_frames - 3})
+    tail = list(sampler)
+    assert len(tail) == 3
+    assert all(0 <= idx < num_frames for idx in tail)
+
+
+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)
+    with pytest.raises(ValueError, match="drop_n_last_frames must be >= 0"):
+        DeterministicEpisodeAwareSampler([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)
+
+
+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)
+    assert list(sampler) == [2, 3, 4, 5]
+    assert "Episode 0" in caplog.text
+
+
+def test_compute_sampler_state():
+    # 100 frames, batch 10, 2 ranks -> 10 underlying batches, 5 per rank per epoch.
+    assert compute_sampler_state(step=0, num_frames=100, batch_size=10, num_processes=2) == {
+        "epoch": 0,
+        "start_index": 0,
+    }
+    # step 7 -> epoch 1, 2 per-rank batches in = 2 * 10 * 2 = 40 samples in
+    assert compute_sampler_state(step=7, num_frames=100, batch_size=10, num_processes=2) == {
+        "epoch": 1,
+        "start_index": 40,
+    }
+    # uneven epoch: 95 frames -> 10 underlying batches (last short), still 5 per rank
+    assert compute_sampler_state(step=12, num_frames=95, batch_size=10, num_processes=2) == {
+        "epoch": 2,
+        "start_index": 40,
+    }
+    # uneven sharding: 105 frames -> 11 underlying batches, 6 per rank (even_batches pads)
+    assert compute_sampler_state(step=11, num_frames=105, batch_size=10, num_processes=2) == {
+        "epoch": 1,
+        "start_index": 100,
+    }