move filtering idle actions after flatten

Author: verityw

src/openpi/training/droid_rlds_dataset.py

@@ -116,6 +116,10 @@ def chunk_actions(traj):
 
         dataset = dataset.traj_map(chunk_actions, num_parallel_calls)
 
+        # Flatten: map from trajectory dataset to dataset of individual action chunks
+        dataset = dataset.flatten(num_parallel_calls=num_parallel_calls)
+
+        # Filter out frames where actions are idle. Must be done after flattening, as filter should apply per-frame.
         def filter_idle(traj):
             """Filter out chunks with idle actions.
             --> we filter if at least first half of chunk does not move.
@@ -127,9 +131,6 @@ def filter_idle(traj):
 
         dataset = dataset.filter(filter_idle)
 
-        # Flatten: map from trajectory dataset to dataset of individual action chunks
-        dataset = dataset.flatten(num_parallel_calls=num_parallel_calls)
-
         # Decode images: RLDS saves encoded images, only decode now for efficiency
         def decode_images(traj):
             traj["observation"]["image"] = tf.io.decode_image(