[RyanPCo]

torchvision draws one random parameter set per call, so applying the
configured train_image_augs to the stacked (B, 3, H, W) camera batch gave
every image in the batch the same jitter. Wrap train_image_augs in a new
PerSampleAugs module so each image draws its own parameters, and move the
inline aug block of _robomimic_to_hpt_data into _apply_image_augs with the
same gating. Deterministic eval augs stay batched.

Co-Authored-By: Claude Fable 5.1 [email protected]

🤖 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
• 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 👈 (View in Graphite)
• 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, 2 nit.

[RyanPCo]

low (test-quality): _hpt_stub wraps train_image_augs in PerSampleAugs itself, so no test checks that HPT.init does the wrapping, which is the actual fix. Move the wrapping into a small helper that both init and the stub call (or test it through the constructor path), and write test_missing_train_augs_pass_through as _hpt_stub(training=True) instead of re-inlining the stub.

[RyanPCo]

nit (clarity): In training with train_image_augs=None, this falls through to eval_image_augs. Before this PR that case raised, and nothing documents the fallback. Restructure so the choice is explicit: if short not in self.encoders: return images, then augs = self.train_image_augs if self.nets.training else self.eval_image_augs, then return augs(images) if augs is not None else images (or keep the eval fallback on purpose and say so). Also make the docstring a one-line summary.

[RyanPCo]

nit (comments): Drop "(ideas 0.10)" and the sentence about "the old inline" call, and remove the print() in test_per_sample_augs_jitters_each_image_independently.

[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

• _apply_image_augs applies eval augs during training when train augs are None: Restructured as suggested. In train mode, train_image_augs=None now means no augs (documented), with a one-line docstring. Test updated. perf: vectorized image augs and a torch.compile knob for HPT and Pi; Pi clips gradients like openpi #676's absent-camera masking and algo: RDT (Robotics Diffusion Transformer) on the HPT chassis #680's frames argument are kept on top.
• Test stub re-implements the init wrapping it should be testing: Wrapping extracted to image_augs.per_sample, used by both HPT.init and the stub. The duplicate stub test now uses _hpt_stub.
• Private plan reference and leftover print in tests: Removed the plan reference, the history sentence and the print.