pi: initialize from PaliGemma with a scratch action expert

README.md

@@ -101,6 +101,7 @@ Every cluster-specific default in the configs can be overridden without editing
 | `EGOVERSE_CACHE_DIR` | `paths.cache_dir`; norm stats cached by content under `<cache_dir>/norm_stats/` | `$EGOVERSE_LOG_DIR/cache` |
 | `EGOVERSE_ROOT` | `paths.root_dir` | the directory you launch from |
 | `EGOVERSE_PI05_WEIGHTS` | pi0.5 base weights directory (`model.robomimic_model.config.pytorch_weight_path`) | lab PACE path |
+| `EGOVERSE_PALIGEMMA_WEIGHTS` | PaliGemma checkpoint for the from-scratch pi arm (`...config.paligemma_weight_path`), a local dir or an HF repo id | `google/paligemma-3b-pt-224` |
 | `WANDB_ENTITY`, `WANDB_PROJECT` | `logger.wandb.entity` / `.project` | `rl2-group` / `zarr_test` |
 
 ### Other Settings

egomimic/algo/pi.py

@@ -16,6 +16,8 @@
 from transformers import AutoTokenizer
 
 from egomimic.algo.algo import Algo
+from egomimic.models import openpi_compat
+from egomimic.models.paligemma_init import load_paligemma_weights, select_init_source
 from egomimic.models.preprocess_pi_obs import (
     PI_CAMERA_SLOTS,
     _concat_proprio,
@@ -38,6 +40,10 @@
 logger.setLevel(logging.INFO)
 logger.propagate = True  # Explicitly enable propagation (default, but ensures it works)
 
+# openpi targets transformers 4.53 (this repo's pin); these shims let it run on a
+# 5.x venv too. See egomimic/models/openpi_compat.py.
+openpi_compat.apply()
+
 
 class PI(Algo):
     """ """
@@ -176,29 +182,39 @@ def __init__(
 
         self.model = openpi.models_pytorch.pi0_pytorch.PI0Pytorch(model_cfg)
 
-        if self.config.pytorch_weight_path is not None:
+        # Three mutually exclusive starting points, in decreasing order of what
+        # is pretrained: the full pi0.5 base checkpoint, PaliGemma only (the
+        # action expert and the action/time projections stay at their init), or
+        # nothing at all.
+        paligemma_weight_path = getattr(self.config, "paligemma_weight_path", None)
+        init_source = select_init_source(
+            self.config.pytorch_weight_path, paligemma_weight_path
+        )
+        if init_source == "pi05_base":
             model_path = os.path.join(
                 self.config.pytorch_weight_path, "model.safetensors"
             )
             if not os.path.isfile(model_path):
                 raise FileNotFoundError(
                     f"Pretrained weight file not found: {model_path}"
                 )
-            target = (
-                self.model.module
-                if isinstance(self.model, torch.nn.parallel.DistributedDataParallel)
-                else self.model
-            )
-            safetensors.torch.load_model(target, model_path)
+            safetensors.torch.load_model(self.model, model_path)
             logger.info(
                 "Loaded pretrained weights from %s (%d parameters)",
                 model_path,
-                sum(p.numel() for p in target.parameters()),
+                sum(p.numel() for p in self.model.parameters()),
+            )
+        elif init_source == "paligemma":
+            load_paligemma_weights(
+                self.model,
+                paligemma_weight_path,
+                allow_mirror=getattr(self.config, "paligemma_allow_mirror", False),
             )
         else:
             logger.warning(
-                "No pytorch_weight_path: no base weights loaded. The weights must "
-                "come from a checkpoint; otherwise this trains from scratch."
+                "Neither pytorch_weight_path nor paligemma_weight_path is set: no "
+                "base weights loaded. The weights must come from a checkpoint; "
+                "otherwise this trains from scratch, VLM included."
             )
         # openpi's PyTorch port hard-codes 32-wide action projections
         # regardless of ``Pi0Config.action_dim``. Wider action spaces (the

egomimic/hydra_configs/model/pi0.5_base.yaml

@@ -5,6 +5,17 @@ robomimic_model:
     pytorch_training_precision: bfloat16
     # directory holding model.safetensors of the pi0.5 base ($EGOVERSE_PI05_WEIGHTS overrides)
     pytorch_weight_path: ${oc.env:EGOVERSE_PI05_WEIGHTS,/storage/home/hcoda1/4/paphiwetsa3/r-dxu345-0/projects/EgoVerse/egomimic/algo/pi_checkpoints/pi05_base_pytorch}
+    # The alternative to pytorch_weight_path: a local directory or an HF repo id
+    # holding the paligemma-3b-*-224 safetensors. Only the VLM prefix (SigLIP
+    # tower, multimodal projector, Gemma-2B) is loaded from it; everything pi0.5
+    # adds -- the action expert, the action in/out projections, the adaRMS time
+    # MLPs -- trains from scratch. Setting both paths is an error. Flip it on
+    # with the model/pi0.5_paligemma_init.yaml overlay.
+    paligemma_weight_path: null
+    # google/paligemma-3b-pt-224 is license-gated. true = without an HF token
+    # that has accepted the Gemma terms, fall back to a third-party mirror whose
+    # shards sha256-match the canonical repo; false (default) = raise.
+    paligemma_allow_mirror: false
     model:
       pi05: true
       # Width of the model's action vector. 32 is openpi's pretrained width

egomimic/hydra_configs/model/pi0.5_bc_mecka_6d_paligemma_init.yaml

@@ -0,0 +1,15 @@
+defaults:
+  - pi0.5_bc_mecka_6d
+  - pi0.5_paligemma_init
+  - _self_
+
+# pi0.5_bc_mecka_6d with the pi0.5 base checkpoint swapped for a PaliGemma-only
+# init -- the ablation that asks how much of the mecka 6D result comes from
+# pi0.5's robot pretraining rather than from the VLM. Everything else (the
+# cartesian 6D action layout, the action converter, the prompt blocks, the
+# optimizer) is the parent's; see pi0.5_paligemma_init.yaml for the two keys
+# this arm changes and for the lr knobs worth sweeping.
+#
+# Pairs with the same data configs as the parent (data=mecka_all_6d,
+# mecka_fold_freeform*_6d) and evaluator=eval_pi_wristframe_6d. Norm stats are
+# data-config-specific, not model-specific, so this arm reuses the parent's.

egomimic/hydra_configs/model/pi0.5_paligemma_init.yaml

@@ -0,0 +1,30 @@
+# OVERLAY, not a runnable model config. Compose it AFTER any pi0.5_* config to
+# swap the pi0.5 base checkpoint for a PaliGemma-only init: the VLM prefix comes
+# from the public paligemma-3b-pt-224 release and the action expert, the action
+# in/out projections and the adaRMS time MLPs start from their PyTorch init.
+#
+#   defaults:
+#     - pi0.5_bc_mecka_6d
+#     - pi0.5_paligemma_init
+#     - _self_
+#
+# It sets only the two init keys, so it never clobbers the config it is layered
+# on (action_dim, ac_keys, domains, the prompt flags all survive). Per-run,
+# without a new file:
+#   model.robomimic_model.config.pytorch_weight_path=null \
+#   model.robomimic_model.config.paligemma_weight_path=google/paligemma-3b-pt-224
+#
+# The optimizer is left at pi0.5_base's fine-tuning recipe (constant 3e-5 after
+# 2000 warmup steps). A randomly initialized action expert may well want a
+# higher peak and a longer warmup; that is a per-run decision, e.g.
+#   model.optimizer.lr=1e-4 model.scheduler.num_warmup_steps=5000
+#
+# Multi-rank runs should pre-fetch the checkpoint once, or every rank races on
+# the same Hugging Face download lock:
+#   python -c "from egomimic.models.paligemma_init import resolve_paligemma_dir as r; print(r('google/paligemma-3b-pt-224'))"
+robomimic_model:
+  config:
+    pytorch_weight_path: null
+    # $EGOVERSE_PALIGEMMA_WEIGHTS overrides: point it at a local snapshot dir to
+    # skip the Hub entirely.
+    paligemma_weight_path: ${oc.env:EGOVERSE_PALIGEMMA_WEIGHTS,google/paligemma-3b-pt-224}

egomimic/models/openpi_compat.py

@@ -0,0 +1,129 @@
+"""Adapters that let openpi's PyTorch pi0 run against transformers 5.
+
+`external/openpi` pins `transformers==4.53.2` and ships a `transformers_replace`
+patch written against it (its own `check.py` asserts that version), which is also
+this repo's pin. A venv on transformers 5 works too, but there the stock modules
+have moved on, so a few of openpi's call sites no longer line up. Nothing here
+runs below transformers 5, and each shim is a no-op when it finds openpi's own
+spelling already in place -- so porting the fork's patch forward makes this
+module inert. Delete it then.
+
+Applied from `egomimic.algo.pi`, which is the only importer of openpi.
+"""
+
+import logging
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+_APPLIED = False
+
+
+def _patch_embed_image() -> bool:
+    """`get_image_features` returns a `BaseModelOutputWithPooling`, not a tensor.
+
+    transformers 5's `PaliGemmaModel.get_image_features` returns the vision
+    output object with the projected features stashed on `pooler_output`, while
+    openpi's `embed_prefix` reads `.shape` off the result. Unwrap it.
+    """
+    from openpi.models_pytorch.gemma_pytorch import PaliGemmaWithExpertModel
+
+    original = PaliGemmaWithExpertModel.embed_image
+
+    def embed_image(self, image: torch.Tensor) -> torch.Tensor:
+        features = original(self, image)
+        if isinstance(features, torch.Tensor):
+            return features
+        return features.pooler_output
+
+    PaliGemmaWithExpertModel.embed_image = embed_image
+    return True
+
+
+def _patch_paligemma_bc_properties() -> bool:
+    """`paligemma.language_model` / `.vision_tower` / `.multi_modal_projector`.
+
+    transformers 5 moved those submodules under `PaliGemmaModel` and dropped the
+    back-compat properties that used to forward to them from
+    `PaliGemmaForConditionalGeneration`. openpi still reads all three off the
+    outer model. Put the forwarders back.
+    """
+    from transformers import PaliGemmaForConditionalGeneration
+
+    added = []
+    for name in ("language_model", "vision_tower", "multi_modal_projector"):
+        if hasattr(PaliGemmaForConditionalGeneration, name):
+            continue
+        setattr(
+            PaliGemmaForConditionalGeneration,
+            name,
+            property(lambda self, _name=name: getattr(self.model, _name)),
+        )
+        added.append(name)
+    return bool(added)
+
+
+# openpi keeps SigLIP's patch and position embeddings in float32 by substring on
+# `vision_tower.vision_model.embeddings.*`; the flattened 5.x spelling of the same
+# three tensors.
+_FP32_KEEP_FLAT = (
+    "vision_tower.embeddings.patch_embedding.",
+    "vision_tower.embeddings.position_embedding.",
+)
+
+
+def _patch_fp32_keep_list() -> bool:
+    """openpi's float32 keep-list misses the SigLIP embeddings.
+
+    `to_bfloat16_for_selected_params` casts everything to bf16 and then restores
+    float32 on names containing `vision_tower.vision_model.embeddings.*`.
+    transformers 5 flattened `SiglipVisionModel`, so nothing matches and the
+    patch embedding and position embedding stay bf16 -- the tensors openpi
+    protects because small updates vanish in bf16. Re-apply the keep under the
+    flattened names; with the nested names present they match nothing.
+    """
+    from openpi.models_pytorch.gemma_pytorch import PaliGemmaWithExpertModel
+
+    original = PaliGemmaWithExpertModel.to_bfloat16_for_selected_params
+
+    def to_bfloat16_for_selected_params(self, precision="bfloat16"):
+        original(self, precision)
+        if precision != "bfloat16":
+            return
+        for name, param in self.named_parameters():
+            if any(selector in name for selector in _FP32_KEEP_FLAT):
+                param.data = param.data.to(dtype=torch.float32)
+
+    PaliGemmaWithExpertModel.to_bfloat16_for_selected_params = (
+        to_bfloat16_for_selected_params
+    )
+    return True
+
+
+def apply() -> None:
+    """Install every shim once. Safe to call repeatedly."""
+    global _APPLIED
+    if _APPLIED:
+        return
+    _APPLIED = True
+
+    import transformers
+
+    major = int(transformers.__version__.split(".")[0])
+    if major < 5:
+        return
+
+    shims = (
+        ("embed_image", _patch_embed_image),
+        ("paligemma_bc_properties", _patch_paligemma_bc_properties),
+        ("fp32_keep_list", _patch_fp32_keep_list),
+    )
+    applied = [name for name, fn in shims if fn()]
+    if applied:
+        logger.info(
+            "openpi targets transformers 4.53 and this environment has %s; "
+            "applied compatibility shims: %s",
+            transformers.__version__,
+            ", ".join(applied),
+        )