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)