Skip to content

Skip short trailing windows so batch > 1 collate sees uniform shapes - #87

Open
jeqcho wants to merge 1 commit into
dreamzero0:mainfrom
jeqcho:fix-collate-uniform-chunks
Open

Skip short trailing windows so batch > 1 collate sees uniform shapes#87
jeqcho wants to merge 1 commit into
dreamzero0:mainfrom
jeqcho:fix-collate-uniform-chunks

Conversation

@jeqcho

@jeqcho jeqcho commented Aug 7, 2026

Copy link
Copy Markdown

Fixes #80.

Root cause of the trailing-batch crash: near an episode tail, get_frame_indices trims the window (unique_sorted = unique_sorted[:-7]), producing a sample with fewer chunks than the rest of the batch. With per_device_train_batch_size > 1 the np.stack in dreamzero_cotrain.py then sees mixed shapes and raises ValueError: all input arrays must have the same shape. On our corpus this fired roughly once per 580 batches, so short smoke runs pass and long runs die.

Fix: skip short windows entirely (same mechanism as the existing <= 8 frames case) instead of trimming to a different chunk count. The empty array needs an integer dtype because it is used as an index downstream. The cost is discarding a small number of episode-tail anchors, which seems preferable to a crash; padding them to full length would change training semantics.

Note this is distinct from the variable-chunk design of the synthetic DROID dataset discussed in #69 — this path is the sharded LeRobot loader, where every sample is supposed to be full-length.

Validated on a 12k-step fine-tune of DreamZero-AgiBot at per-device batch 4 (global 32) over 124 bimanual YAM LeRobot repos with zero collate failures: https://github.com/robocurve/dreamzero-yam, released at https://huggingface.co/robocurve/dreamzero-yam-molmoact2.

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Dataloader crashes on trailing incomplete batch (np.stack shape mismatch)

1 participant