feat: paligemma base weights for pi0.5 training

[AnikethCheluva]

Add a path to pretrain pi0.5 from Google's raw PaliGemma base weights
instead of starting from PI's pi05_base checkpoint.

• external/openpi bumped to include the new convert_paligemma_to_pytorch.py
  script (downloads pt_224.npz from GCS and produces a PI0Pytorch
  safetensors with the backbone loaded and action expert / projection
  heads at random init).
• pi0.5_paligemma_init.yaml: new hydra model config tuned for
  from-scratch pretraining (higher LR, longer warmup, grad clipping on).
• paligemma_base_training.md: plan doc explaining the two checkpoints,
  the conversion flow, and training considerations.

[AnikethCheluva]

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

• feat: paligemma base weights for pi0.5 training #609 👈 (View in Graphite)
• make config changes #176
• refactoring aciton transforms #604 : 1 other dependent PR (#605 )
• main

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

[github-actions]

Claude Code Review

PR Review: paligemma base weights for pi0.5 training

Summary

Adds a from-scratch pretraining path for pi0.5 that starts from Google's PaliGemma weights (backbone only) with a randomly initialized action expert, via a new openpi conversion script and a Hydra model config.

Key concerns

1. Hardcoded absolute user path in shared config. pytorch_weight_path: /storage/home/hcoda1/5/acheluva3/r-dxu345-0/... bakes @AnikethCheluva's PACE homedir into a committed config. Anyone else on the team using model=pi0.5_paligemma_init will fail to load. The plan doc even acknowledges "point it wherever you saved" — that comment should be reflected in the YAML (env var, ??? required field, or ${oc.env:PI_CKPT_ROOT}/paligemma_base_pytorch). Note pi0.5_base.yaml has the same anti-pattern, but new configs shouldn't perpetuate it.

2. Conversion script isn't in the diff. The description and plan doc both reference external/openpi/examples/convert_paligemma_to_pytorch.py as the core deliverable, but only the submodule SHA bump is shown. I can't verify:

   • Key-prefix mapping (paligemma_with_expert.paligemma. — is that actually how PI0Pytorch names them? convert_jax_model_to_pytorch.py is the ground truth to cross-check.)
   • Whether load_state_dict(strict=False) swallows unexpected keys silently (backbone key rename drift → silently random-init backbone → looks like it's training but VLM is cold. This is the exact failure mode that's hardest to catch from loss curves alone.)
   • That save_model produces a file safetensors.torch.load_model (strict=True in pi.py) can actually reload without missing keys.

   Please either include the script in this PR or link the openpi PR for review alongside.

3. Verification is aspirational, not automated. The "Verification checklist" is a manual list. Given the silent-failure risk in (2), at minimum add an assertion in the conversion script that:

   • Every key in the freshly-instantiated PI0Pytorch state dict is either (a) covered by the prefixed PaliGemma keys, or (b) in an explicit allowlist of expected-random keys (action_in_proj, action_out_proj, time_mlp_in, time_mlp_out, action expert layers).
   • No unexpected keys were dropped from the source PaliGemma state.
     Print counts of loaded vs random params so it's visible in logs.

4. num_training_steps: 200000 with a data config unspecified. Cosine schedule collapses to ~0 LR at step 200k. If someone runs this on a small dataset for 60k steps (the old horizon), they get a lower effective LR than pi0.5_base after warmup, defeating the purpose. Worth a comment in the YAML that this horizon assumes a specific data scale.

Suggestions

• Move the hardcoded path to an env var or make it a required Hydra field (pytorch_weight_path: ???) with the resolution documented in the plan doc.
• Add a README.md inside egomimic/algo/pi_checkpoints/ (or update pi05.md) rather than a top-level paligemma_base_training.md — top-level markdown files tend to rot. At minimum, cross-link from pi05.md.
• In the conversion script, write the source HF commit SHA (not just "paligemma_base") into config.json for reproducibility.
• Add a smoke test or a scripts/verify_pi_checkpoint.py that loads the safetensors into PI0Pytorch and asserts key coverage — cheap insurance against submodule drift.
• Nit: enable_grad_norm: true appears both as a top-level key and is discussed as "consider setting" in the plan doc — plan doc should be updated to match (it's on, not optional).

Verdict: Request Changes

The design is sound and the plan doc is unusually thorough (nice), but the hardcoded path is a blocker for a shared config, and I can't sign off on correctness without seeing the conversion script's key-mapping and load assertions. Once the openpi script is reviewable and the path is parameterized, this should be quick to re-approve.

---

Reviewed by Claude · Review workflow