[RyanPCo]

Adds an opt-in mask so the HPT loss can skip the repeat-last action padding at episode and annotation tails. The dataset now emits the mask; the loss uses it only when use_pad_mask: true.

Behaviour changes

None by default. use_pad_mask is false in hpt_bc_mecka_6d_300M.yaml and hpt_bc_keypoints_wrist_300M.yaml, and every other HPT config takes the __init__ default of false. With it off, the mask is dropped from the batch and pad_mask stays all ones, which is exactly F.mse_loss.

Turning it on changes the loss: the repeat-last padding stops being supervised, so runs with and without it are not comparable. On the three pinned flagship episodes, 6.4–7.6 % of supervised action steps are repeat-last padding and 12.8–15.2 % of samples contain some. The layout-aware eval metrics would have to be masked in the same change, or train and val would measure different things.

What changed

• Dataset: ZarrDataset emits action_pad_mask (1.0 per real frame, 0.0 per repeat-last pad) next to every horizoned action chunk. It honours the annotation-cutoff EOS clamp.
• Interpolation: a new InterpolatePadMask transform carries the mask through the 30 → 100 upsampling on the interpolators' own grid. An output step is real only if both source frames it interpolates between are real. It is a no-op without the key and idempotent, so a mode that chains two builders (include_ee_pose) upsamples once. It is added to the eva and human cartesian and keypoint transform lists.
• HPT: with use_pad_mask on, process_batch_for_training builds pad_mask (B, S, 1) from the dataset mask instead of ones, and raises if the mask length doesn't match the action horizon.
• Loss: DenoisingPolicy.loss_fn takes the mask and averages the squared error over the valid steps only. An all-padding chunk contributes 0, not NaN.
• Padding is still repeat-last and no tail samples are dropped.

Notes

• It ships off to match openpi, which requests action_is_pad from LeRobotDataset but never uses it (Pi0.compute_loss averages over the whole horizon).
• LeRobot's own policies do mask, but they take .mean() over all elements including the zeroed ones, so their loss shrinks as padding grows. This divides by the valid count instead, which keeps the gradient magnitude constant.

Tests

• New tests/unit/test_loss_masking.py: the dataset mask (tail chunks, annotation cutoff, survival through bounds check and normalization), InterpolatePadMask (30 → 100, stride, no-op, idempotent), loss_fn against F.mse_loss, and the use_pad_mask flag in process_batch_for_training and the model yamls.
• tests/unit/test_action_chunk_transforms.py: the eva and human transform-list snapshots now include InterpolatePadMask.

🤖 Generated with Claude Code

https://claude.ai/code/session_01HjCQBqQCr2igFEMR7yUovw

[RyanPCo]

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

• hpt: never re-initialise pretrained stems; hash-check them; mask text padding #645 : 2 dependent PRs (#646 , #664 )
• obs: proprio history (K proprio steps, one token each, learned time embedding) #650
• loss: mask the repeat-last action padding out of the HPT loss #649 👈 (View in Graphite)
• resnet: keep a frozen backbone in eval mode; give backbones their own LR group #652
• augs: per-sample image augmentation for HPT train-mode augs #651
• eval: deterministic validation (seeded prompts + flow sampling); log-only grad-norm MAD #644
• data: mecka_fold_flagship_top3_hpt_6d, the 3-operator flagship train split #643
• val: subsampled metric loaders + pinned contiguous video loaders #642
• infra: mecka 6D / keypoint recipes, 300M models, defaults, and the rollout revert #640
• eval: layout-aware frame-invariant metrics, and a train-split viz head #659
• data: stop letting one bad episode or a rotation channel take a run down #658
• norm: guard degenerate channels, horizon pooling, and no rot6d normalization #657
• keypoints: the 144-D wrist-first hand-keypoint action becomes the human default #656
• mecka: fix the double-mirrored LEFT wrist frame, and correct existing zarrs #655
• 6d: PI goes 6D-only; Eva and Human gain the continuous-rotation modes #654
• 6d: continuous rotation representation, bimanual layouts, shared norm helpers #653
• train: give each Slurm slurm job gets its own cache dir to avoid conflict #638 : 1 other dependent PR (#639 )
• write git version into the run dir for reproducibility #637
• train: don't load PI base weights on resume and verify that checkpoint path exists #636
• configs: declare proprio and actions dims once per HPT model config #634 : 1 other dependent PR (#635 )
• cache norm stats for data configs to minimize recomputation #633
• import cloud packages lazily to support local training #631 : 1 other dependent PR (#632 )
• cache dataset resolver for faster training initialization #630
• validate episode hashes in data config before training #628
• modify pi camera map to match data #629
• add tests to ensure training works between every model and data vendor #626
• run CI tests on every PR #625
• add env variables for user-specific paths #624
• build: set pyproject and uv as the dependency managers #623
• fix(configs): fix paths for data hydra yaml configs #622
• refactor(algo): delete ACT #627 : 1 other dependent PR (#621 )
• tests: create dedicated test folder with unit and integration tests #620
• add pyright type checker #619
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[RyanPCo]

Automated adversarial review (read-only; findings were checked by a second skeptical pass before posting): 1 low, 3 nit.

[RyanPCo]

low (correctness): use_pad_mask only takes effect through DenoisingPolicy.compute_loss. PolicyHead.compute_loss (MLPPolicyHead / MultiBlockTransformerDecoder) still computes LOSS(pred_action, target_action) and ignores data["pad_mask"]. On a non-diffusion HPT config the flag is accepted and the mask is built and validated, but the padded tail is still fully supervised. Either apply the mask in PolicyHead.compute_loss or raise when use_pad_mask is set and the head is not a DenoisingPolicy.

[RyanPCo]

nit (duplication): InterpolatePadMask(new_chunk_length=chunk_length, stride=stride) is pasted into five builders (eva.py x2, human.py x3). A future builder that forgets it only shows up as a horizon-mismatch ValueError once use_pad_mask is on. Since the transform is idempotent, consider appending it once in a shared helper, or after the builder in get_transform_list.

[RyanPCo]

nit (comments): "downstream sees the same keys as before" describes the diff. Say why instead, e.g. # action_pad_mask is consumed here; only pad_mask goes downstream.

[RyanPCo]

nit (test-quality): InterpolatePadMask exists now, so _pad_mask_transform_cls ("so the module still collects before the transform exists") is leftover TDD scaffolding. Import it at module top and drop the helper.

[RyanPCo]

Review findings: status after the fixup pass. The fixes are squashed into this PR's commits, and the whole stack was force-pushed.

Fixed

• use_pad_mask is silently a no-op for non-denoising HPT heads: HPT.init now raises ValueError when use_pad_mask is set without diffusion. No unit test: building a real HPT is heavy.
• Comment references the previous behavior: Reworded as suggested.
• TDD scaffolding left in the test module: InterpolatePadMask is imported at module top; helper deleted.

Not fixed

• InterpolatePadMask hand-inserted into each builder: Nit, not cheap: a central append would have to thread each builder's chunk_length and stride through get_transform_list. The transform is idempotent and a no-op without a mask.