From b5ff91b672bcec343ff6f0b29d3f56d21214fd74 Mon Sep 17 00:00:00 2001
From: Ilia Zavidnyi <45592678+zavidnyi@users.noreply.github.com>
Date: Sat, 25 Oct 2025 11:29:11 +0200
Subject: [PATCH] Shuffle should be applied conditionally

---
 src/openpi/training/droid_rlds_dataset.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/openpi/training/droid_rlds_dataset.py b/src/openpi/training/droid_rlds_dataset.py
index ebe863acf6..d0f7f63995 100644
--- a/src/openpi/training/droid_rlds_dataset.py
+++ b/src/openpi/training/droid_rlds_dataset.py
@@ -203,7 +203,8 @@ def decode_images(traj):
         dataset = dataset.frame_map(decode_images, num_parallel_calls)
 
         # Shuffle, batch
-        dataset = dataset.shuffle(shuffle_buffer_size)
+        if shuffle:
+            dataset = dataset.shuffle(shuffle_buffer_size)
         dataset = dataset.batch(batch_size)
         # Note =>> Seems to reduce memory usage without affecting speed?
         dataset = dataset.with_ram_budget(1)