From abb330e444f91ad36e49a020111ed877fcf762d7 Mon Sep 17 00:00:00 2001 From: Jay Chooi Date: Thu, 6 Aug 2026 22:07:27 -0700 Subject: [PATCH] Skip short trailing windows instead of trimming to fewer chunks Windows near an episode tail can trim down to fewer chunks than the rest of the batch. With per_device_train_batch_size > 1 the collate np.stack then sees mixed shapes and crashes. Skip these windows entirely, like the existing <=8-frame case, and give the empty array an integer dtype since it is used as an index downstream. --- groot/vla/data/dataset/lerobot_sharded.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/groot/vla/data/dataset/lerobot_sharded.py b/groot/vla/data/dataset/lerobot_sharded.py index 409fa378..3564ef35 100755 --- a/groot/vla/data/dataset/lerobot_sharded.py +++ b/groot/vla/data/dataset/lerobot_sharded.py @@ -1220,11 +1220,11 @@ def add_step_set(anchor_index: int) -> None: if additional_idx < trajectory_length and unique_sorted.size < max_frames: unique_sorted = np.append(unique_sorted, additional_idx) else: - # Trim to 8n+1 format. Require at least 9 frames so (noisy_frames-1)//num_frame_per_block >= 1 - # for action/state model invariant (CausalWanModel); otherwise return empty so sample is skipped. - if unique_sorted.size <= 8: - return np.array([]) - unique_sorted = unique_sorted[:-7] + # Batch >1 requires uniform shapes across the batch: skip + # short windows entirely instead of trimming to fewer chunks + # (mixed chunk counts crash the collate np.stack). dtype must + # be integer: the empty array is used as an index downstream. + return np.array([], dtype=np.int64) # ensure that unique_sorted has 4n+1 frames assert unique_sorted.size % 8 == 1, f"unique_sorted size {unique_sorted.size} is not 4n+1"