Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jeqcho wants to merge 1 commit into
base: main
Choose a base branch
from
Open

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

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions groot/vla/data/dataset/lerobot_sharded.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down