[RyanPCo]

Adds a pi0.5 arm that starts from the released PaliGemma VLM instead of the pi0.5 base checkpoint. This separates what VLM pretraining buys from what pi0.5's robot pretraining buys.

• Loaded from google/paligemma-3b-pt-224: the SigLIP tower, Gemma-2B and the projector.
• Left at init: everything pi0.5 adds, meaning the Gemma-300M action expert, action_in/out_proj and the adaRMS time MLPs.

Usage

model/pi0.5_paligemma_init.yaml is an overlay. Compose it after any pi0.5_* model config:

defaults:
  - pi0.5_bc_mecka_6d
  - pi0.5_paligemma_init
  - _self_

• It changes only the init source. The LR schedule stays pi0.5_base's.
• $EGOVERSE_PALIGEMMA_WEIGHTS points it at a local snapshot.
• Prefetch the checkpoint once before a multi-rank run (command in pi05.md).

Loader

• pytorch_weight_path and paligemma_weight_path are mutually exclusive. Setting both is an error; setting neither warns.
• The load is strict. It streams shard by shard and fails if any parameter is left untouched or any checkpoint tensor goes unused.
• The release's SigLIP names are mapped onto whichever layout the installed transformers builds. On 4.53.2 (the pin) the tree keeps a nested vision_model. level; on 5.x it is flattened. An earlier version always mapped to the 5.x names and failed 21 tensors on the pin. The key-string unit tests couldn't see that, so a test now loads a tiny real PaliGemmaForConditionalGeneration; it passes on 4.53.2 and 5.17.0.
• The released vocabulary is padded to 257216 rows. embed_tokens and lm_head are truncated to openpi's 257152. The dropped rows are <image> plus 63 padding tokens, and openpi never looks them up.

Gated repo and mirrors

• google/paligemma-3b-pt-224 is license-gated. Downloading it needs an HF token that has accepted the license.
• The mirror fallback is off by default. With paligemma_allow_mirror: true, a 401/403 on the canonical repo falls back to third-party mirrors. Rate limits and outages don't trigger it. Every mirror shard is sha256-checked against the canonical repo's metadata before use.
• Only *.safetensors are downloaded.

transformers 4.53 vs 5.x

• Supported setup: the pin, 4.53.2 plus openpi's transformers_replace patch.
• On 5.x, only this arm loads. The pi0.5 base checkpoint stores the nested SigLIP names, and its strict load stops at startup.
• openpi_compat.py shims for a 5.x venv (all to delete once the fork's patch is ported forward):
  • get_image_features returns an output object, not a tensor.
  • The language_model / vision_tower / multi_modal_projector forwarders are gone.
  • openpi's fp32 keep-list no longer matches the flattened SigLIP embeddings, which then trained in bf16.
• The first commit makes the test harness's HF-offline switch work on either version.
• pi05.md documents all of this.

Verification

• Training run: a full trainHydra run to exit 0 on an H100 (4 epochs, all four validation heads, checkpoint and videos written), on transformers 5.17.
• Real checkpoint: the real 12 GB checkpoint loaded into the real model (tests/integration/test_paligemma_init_real.py), on 5.17. It also asserts the SigLIP embeddings stay fp32.
• Unit tests: pass on 4.53.2 and 5.17.0.

🤖 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

• pi: initialize from PaliGemma with a scratch action expert #677 : 2 dependent PRs (#678 , #680 ) 👈 (View in Graphite)
• perf: vectorized image augs and a torch.compile knob for HPT and Pi; Pi clips gradients like openpi #676
• loader: one zarr read per key, not one per keymap entry #674
• eval: build the pinned-episode video batches only where they are used #673
• checkpoints: last.ckpt as a symlink, not a second full copy #672
• trainHydra: metric_frames_per_episode: auto, and warn on empty prompts #671
• transforms: per-step wrist framing for proprio history #670
• viz: one keypoint renderer, the horizon trace #669
• configs: history and backbone-LR knobs on the keypoint wrist HPT model #668
• trainer: ddp_bf16_compress, DDP with a bf16 gradient all-reduce #667
• eval: overlay videos play at the source rate under DDP #666
• hpt: build the policy on CPU and let Lightning place it #664 : 1 other dependent PR (#665 )
• hpt: never re-initialise pretrained stems; hash-check them; mask text padding #645 : 1 other dependent PR (#646 )
• 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
• 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 high, 2 nit.

[RyanPCo]

high (correctness): Any rebuild of a PaliGemma-init PI from its saved config downloads PaliGemma again. paligemma_weight_path is saved in the checkpoint's config_tree, and two paths miss it. trainHydra._build_model_config_tree nulls only pytorch_weight_path on resume/requeue: that key is already null for this arm, so has_weights is False and nothing is nulled. rollout._patch_checkpoint_paths returns early when pytorch_weight_path is None. As a result, every resumed rank and the robot rollout call snapshot_download for the gated google/paligemma-3b-pt-224 and load 3B weights that the checkpoint state_dict then overwrites. Without network access or a Gemma-licensed HF token, this fails. Please null paligemma_weight_path alongside pytorch_weight_path in both places when weights come from a checkpoint, and add a resolve test that sets ckpt_path on the paligemma-init config.

[RyanPCo]

nit (dead-code): self.model was just built from PI0Pytorch(model_cfg) above, so it can never be DistributedDataParallel here. Since this block is being moved anyway, pass self.model straight to both loaders and drop target.

[RyanPCo]

nit (test-quality): This duplicates test_the_mirror_is_opt_in, because allow_mirror already defaults to False. test_only_the_shards_are_downloaded (line 305) only re-asserts a module constant; either delete it or assert the allow_patterns the stubbed snapshot_download receives. In test_a_gated_canonical_repo_falls_back_to_a_matching_mirror, good and its digest entry are overwritten by the update on line 349; build digests in one comprehension.

[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

• Loading a PaliGemma-init checkpoint downloads PaliGemma from the Hub again: trainHydra nulls paligemma_weight_path too when weights come from ckpt; rollout._patch_checkpoint_paths nulls it (rollout.py is outside the PR, but the fix needs it); tests added for both
• DDP unwrap on a just-constructed model: loads go straight into self.model; the pre-existing copy in _resize_action_projections is left alone (outside the PR)
• Duplicate and tautological tests: duplicate mirror-off test removed; dead 'good' digests removed; shards test now checks the allow_patterns kwarg snapshot_download receives