feat(streaming): seeded shard-order permutation per (seed, epoch, rank)

Shards were assigned to consumers in file-index order, so a sub-epoch
run over a corpus consolidated source-by-source trains on whatever the
first N% of files contains and drifts curriculum-style as sources change
under it. Permute the rank's shard list with a seeded RNG before worker
striding: a 30%-of-epoch run now sees a uniform 30% sample of files.

The permutation is seeded by (seed, epoch, rank) only - every DataLoader
worker of a rank must derive the identical list, since workers stride it
and disagreement would create overlapping shard assignments. It re-draws
each epoch, is the identity when shuffle=False, and stays deterministic
for fast-forward resume.

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

src/lerobot/datasets/streaming_dataset.py

@@ -400,6 +400,21 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
             state = _mix64(state ^ _mix64(salt))
         return np.random.default_rng(state)
 
+    def _shard_order(self, epoch: int, num_shards: int) -> list[int]:
+        """Seeded permutation of this rank's shard indices, re-drawn every epoch.
+
+        In a sub-epoch run over a corpus consolidated source-by-source, index-order shard
+        assignment means training on whatever the first N% of files contains; permuting the
+        shard order turns that into a uniform sample of files. Seeded by (seed, epoch, rank)
+        only — every DataLoader worker of the rank must agree on this list, because workers
+        stride it and disagreement would create overlapping shard assignments.
+        """
+        order = list(range(num_shards))
+        if self.shuffle:
+            state = _mix64(self.seed) ^ _mix64(0x5EED5EED) ^ _mix64(self.rank) ^ _mix64(epoch)
+            np.random.default_rng(_mix64(state)).shuffle(order)
+        return order
+
     def _make_video_decoder_cache(self) -> VideoDecoderCache:
         """Size the decoder cache to the pool's working set (pool episodes x cameras), capped at 128."""
         if self.video_decoder_cache_size is not None:
@@ -481,7 +496,9 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
             ds = split_dataset_by_node(ds, rank=self.rank, world_size=self.world_size)
 
         num_shards = ds.num_shards if self.max_num_shards is None else min(ds.num_shards, self.max_num_shards)
-        shard_indices = list(range(num_shards))
+        epoch = self._epoch
+        self._epoch += 1
+        shard_indices = self._shard_order(epoch, num_shards)
 
         # DataLoader workers within this rank further split the shards so they don't yield duplicates.
         worker_info = torch.utils.data.get_worker_info()
@@ -498,8 +515,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
         prefetcher = self._make_prefetcher()
         self._prefetcher = prefetcher
 
-        epoch = self._epoch
-        self._epoch += 1
         rng = self._consumer_rng(epoch, worker_id)
         # Workers beyond the shard count yield nothing and are stopped by the DataLoader, so the
         # batch round-robin effectively runs over min(num_workers, num_shards) active workers.

tests/datasets/test_streaming_native.py

@@ -378,3 +378,27 @@ def test_episode_grouping_native_and_fallback_agree(tmp_path, lerobot_dataset_fa
         return
     native = episode_signature(True)
     assert native == fallback
+
+
+def test_shard_order_permutation_properties(tmp_path, lerobot_dataset_factory):
+    """Shard order: a valid permutation, deterministic per (seed, epoch, rank), worker-independent
+    (workers stride the same list, so it must not depend on worker id), reshuffled across epochs,
+    and identity when shuffle is off."""
+    repo_id = f"{DUMMY_REPO_ID}-shardorder"
+    _make_local_dataset(lerobot_dataset_factory, tmp_path / "ds", repo_id, total_episodes=4, total_frames=80)
+
+    ds = StreamingLeRobotDataset(repo_id=repo_id, root=tmp_path / "ds", shuffle=True, seed=5)
+    num_shards = 32
+    order_epoch0 = ds._shard_order(0, num_shards)
+    assert sorted(order_epoch0) == list(range(num_shards))
+    assert ds._shard_order(0, num_shards) == order_epoch0  # deterministic
+    assert ds._shard_order(1, num_shards) != order_epoch0  # reshuffles per epoch
+    assert order_epoch0 != list(range(num_shards))  # actually permuted (P=1/32! of false alarm)
+
+    other_rank = StreamingLeRobotDataset(
+        repo_id=repo_id, root=tmp_path / "ds", shuffle=True, seed=5, rank=1, world_size=2
+    )
+    assert other_rank._shard_order(0, num_shards) != order_epoch0  # ranks decorrelated
+
+    unshuffled = StreamingLeRobotDataset(repo_id=repo_id, root=tmp_path / "ds", shuffle=False, seed=5)
+    assert unshuffled._shard_order(0, num_shards) == list(range(num_shards))