From dedf88d55516b2d5a99eda5326d310c3160f1586 Mon Sep 17 00:00:00 2001 From: Bartok9 Date: Fri, 10 Jul 2026 20:00:56 -0400 Subject: [PATCH] fix(training): apply DROID RLDS shuffle only when enabled MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Salvage of Physical-Intelligence/openpi#746 by @zavidnyi — rebased to main. The constructor exposes shuffle= but always shuffled the concatenated dataset. Gate dataset.shuffle on the flag so evaluation can be deterministic. --- src/openpi/training/droid_rlds_dataset.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/openpi/training/droid_rlds_dataset.py b/src/openpi/training/droid_rlds_dataset.py index 2cb08a16ac..ff179dbbee 100644 --- a/src/openpi/training/droid_rlds_dataset.py +++ b/src/openpi/training/droid_rlds_dataset.py @@ -230,7 +230,9 @@ def decode_images(traj): weights = [dataset.weight for dataset in datasets] final_dataset = dl.DLataset.sample_from_datasets(all_datasets, weights=weights) - final_dataset = final_dataset.shuffle(shuffle_buffer_size) + # Honor the constructor `shuffle` flag (default True). Skip for deterministic evaluation. + if shuffle: + final_dataset = final_dataset.shuffle(shuffle_buffer_size) final_dataset = final_dataset.batch(batch_size) # Note =>> Seems to reduce memory usage without affecting speed? final_dataset = final_dataset.with_ram_budget(1)