898 lines
41 KiB
Python
898 lines
41 KiB
Python
|
# SPDX-License-Identifier: Apache-2.0
|
|||
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|||
|
|
"""Inference-only π0.5 VLA math kernel for vllm-omni.
|
|||
|
|
|
|||
|
|
Only the math that turns a robot observation into an action chunk; no serving or
|
|||
|
|
request glue. Deliberately shaped like ``models/pi0/modeling_pi0.py`` so the two
|
|||
|
|
can later be factored into a shared Pi-family module (RFC step 2) on a
|
|||
|
|
"behaviour unchanged" review.
|
|||
|
|
|
|||
|
|
π0.5 = PaliGemma (SigLIP vision + Gemma 2B LM) prefix + Gemma 300M action expert
|
|||
|
|
suffix + flow-matching head. Where π0 projects the robot state through
|
|||
|
|
``state_proj``, π0.5 discretizes it into prompt tokens — so there is no
|
|||
|
|
``state_proj`` layer, ``sample_actions`` takes no ``state`` argument, and the
|
|||
|
|
suffix is action tokens only, which drops π0's leading state-token boundary from
|
|||
|
|
the suffix attention mask and leaves ``[1] + [0] * (horizon - 1)``.
|
|||
|
|
|
|||
|
|
π0.5 is nonetheless the *larger* model: the 37 AdaRMS ``dense`` projections add
|
|||
|
|
~116M parameters against the ~8K that ``state_proj`` saves. (LeRobot's README
|
|||
|
|
says otherwise; the checkpoint disagrees.)
|
|||
|
|
|
|||
|
|
Reference implementations:
|
|||
|
|
- OpenPI: openpi/src/openpi/models_pytorch/pi0_pytorch.py, gemma_pytorch.py
|
|||
|
|
- LeRobot: lerobot/src/lerobot/policies/pi05/modeling_pi05.py
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import logging
|
|||
|
|
import math
|
|||
|
|
from collections.abc import Iterable
|
|||
|
|
|
|||
|
|
import torch
|
|||
|
|
import torch.nn as nn
|
|||
|
|
import torch.nn.functional as F
|
|||
|
|
from transformers.models.auto import CONFIG_MAPPING
|
|||
|
|
from transformers.models.gemma.modeling_gemma import (
|
|||
|
|
GemmaForCausalLM,
|
|||
|
|
apply_rotary_pos_emb,
|
|||
|
|
)
|
|||
|
|
from transformers.models.paligemma.modeling_paligemma import (
|
|||
|
|
PaliGemmaForConditionalGeneration,
|
|||
|
|
)
|
|||
|
|
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
# ──────────────────────────────────────────────────────────────────────
|
|||
|
|
# Constants
|
|||
|
|
# ──────────────────────────────────────────────────────────────────────
|
|||
|
|
DEFAULT_ACTION_DIM = 32
|
|||
|
|
DEFAULT_ACTION_HORIZON = 50
|
|||
|
|
DEFAULT_MAX_TOKEN_LEN = 200 # π0 uses 48
|
|||
|
|
DEFAULT_NUM_INFERENCE_STEPS = 10
|
|||
|
|
DEFAULT_IMAGE_RESOLUTION = (224, 224) # openpi/models/model.py IMAGE_RESOLUTION
|
|||
|
|
DEFAULT_STATE_NUM_BINS = 256 # openpi PaliGemmaTokenizer.tokenize()
|
|||
|
|
|
|||
|
|
# Large negative value to fill masked-out positions in a float attention mask.
|
|||
|
|
# Matches OpenPI's constant exactly so that numerics line up during parity.
|
|||
|
|
# Ref: openpi/src/openpi/models/gemma.py
|
|||
|
|
OPENPI_ATTENTION_MASK_VALUE = -2.3819763e38
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ──────────────────────────────────────────────────────────────────────
|
|||
|
|
# Gemma variant configs (matches openpi/models/gemma.py get_config)
|
|||
|
|
# ──────────────────────────────────────────────────────────────────────
|
|||
|
|
class GemmaVariantConfig:
|
|||
|
|
def __init__(self, width, depth, mlp_dim, num_heads, num_kv_heads, head_dim):
|
|||
|
|
self.width = width
|
|||
|
|
self.depth = depth
|
|||
|
|
self.mlp_dim = mlp_dim
|
|||
|
|
self.num_heads = num_heads
|
|||
|
|
self.num_kv_heads = num_kv_heads
|
|||
|
|
self.head_dim = head_dim
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_gemma_config(variant: str) -> GemmaVariantConfig:
|
|||
|
|
if variant == "gemma_2b":
|
|||
|
|
return GemmaVariantConfig(2048, 18, 16384, 8, 1, 256)
|
|||
|
|
elif variant == "gemma_300m":
|
|||
|
|
return GemmaVariantConfig(1024, 18, 4096, 8, 1, 256)
|
|||
|
|
else:
|
|||
|
|
raise ValueError(f"Unknown variant: {variant}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ──────────────────────────────────────────────────────────────────────
|
|||
|
|
# Utility functions (match openpi/models_pytorch/pi0_pytorch.py)
|
|||
|
|
# ──────────────────────────────────────────────────────────────────────
|
|||
|
|
def create_sinusoidal_pos_embedding(
|
|||
|
|
time: torch.Tensor,
|
|||
|
|
dimension: int,
|
|||
|
|
min_period: float = 4e-3,
|
|||
|
|
max_period: float = 4.0,
|
|||
|
|
device: torch.device = None,
|
|||
|
|
) -> torch.Tensor:
|
|||
|
|
"""Compute a sine/cosine positional embedding for scalar timesteps.
|
|||
|
|
|
|||
|
|
Ref: openpi/models_pytorch/pi0_pytorch.py create_sinusoidal_pos_embedding
|
|||
|
|
"""
|
|||
|
|
if dimension % 2 != 0:
|
|||
|
|
raise ValueError(f"dimension ({dimension}) must be divisible by 2")
|
|||
|
|
if time.ndim != 1:
|
|||
|
|
raise ValueError("time tensor must be 1-D (batch_size,)")
|
|||
|
|
if device is None:
|
|||
|
|
device = time.device
|
|||
|
|
|
|||
|
|
# Use float64 for the log-linear sweep and for the inner products, to
|
|||
|
|
# match the reference implementation's numerical behaviour exactly.
|
|||
|
|
dtype = torch.float64
|
|||
|
|
fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)
|
|||
|
|
period = min_period * (max_period / min_period) ** fraction
|
|||
|
|
scaling_factor = 1.0 / period * 2 * math.pi
|
|||
|
|
sin_input = scaling_factor[None, :] * time[:, None].to(dtype)
|
|||
|
|
return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def make_att_2d_masks(pad_masks: torch.Tensor, att_masks: torch.Tensor) -> torch.Tensor:
|
|||
|
|
"""Build a 2D attention mask from a padding mask and an autoregressive mask.
|
|||
|
|
|
|||
|
|
Ref: openpi/models_pytorch/pi0_pytorch.py make_att_2d_masks
|
|||
|
|
"""
|
|||
|
|
if att_masks.ndim != 2:
|
|||
|
|
raise ValueError(f"att_masks must be 2-D, got {att_masks.ndim}-D")
|
|||
|
|
if pad_masks.ndim != 2:
|
|||
|
|
raise ValueError(f"pad_masks must be 2-D, got {pad_masks.ndim}-D")
|
|||
|
|
|
|||
|
|
cumsum = torch.cumsum(att_masks, dim=1)
|
|||
|
|
att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]
|
|||
|
|
pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None]
|
|||
|
|
return att_2d_masks & pad_2d_masks
|
|||
|
|
|
|||
|
|
|
|||
|
|
def prepare_attention_masks_4d(att_2d_masks: torch.Tensor) -> torch.Tensor:
|
|||
|
|
"""Convert ``(B, S, S)`` bool masks to ``(B, 1, S, S)`` float masks.
|
|||
|
|
|
|||
|
|
``True`` → 0.0 (attend), ``False`` → ``OPENPI_ATTENTION_MASK_VALUE``.
|
|||
|
|
"""
|
|||
|
|
att_2d_masks_4d = att_2d_masks[:, None, :, :]
|
|||
|
|
return torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class Pi05AdaRMSNorm(nn.Module):
|
|||
|
|
"""Adaptive RMSNorm conditioned on the flow-matching timestep.
|
|||
|
|
|
|||
|
|
π0 conditions on time by *concatenating* a time embedding onto each action
|
|||
|
|
embedding. π0.5 instead feeds the time embedding into every action-expert
|
|||
|
|
norm, which produces a per-layer ``(scale, shift, gate)`` triple::
|
|||
|
|
|
|||
|
|
y = norm(x) * (1 + scale) + shift
|
|||
|
|
out = residual + gate * sublayer(y)
|
|||
|
|
|
|||
|
|
``dense`` is zero-initialized, so an untrained model starts as the identity
|
|||
|
|
modulation with a closed gate — matching OpenPI's parameterization.
|
|||
|
|
|
|||
|
|
Note the shape of the unconditioned branch: ``normed * (1 + weight)`` with
|
|||
|
|
``weight`` zero-initialized, which is exactly ``transformers``'
|
|||
|
|
``GemmaRMSNorm``. That equivalence is why only the *expert* norms need
|
|||
|
|
replacing here and the PaliGemma prefix can keep stock Gemma layers.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, dim: int, eps: float = 1e-6, cond_dim: int | None = None):
|
|||
|
|
super().__init__()
|
|||
|
|
self.eps = eps
|
|||
|
|
self.dim = dim
|
|||
|
|
self.cond_dim = cond_dim
|
|||
|
|
if cond_dim is not None:
|
|||
|
|
self.dense = nn.Linear(cond_dim, dim * 3, bias=True)
|
|||
|
|
nn.init.zeros_(self.dense.weight)
|
|||
|
|
nn.init.zeros_(self.dense.bias)
|
|||
|
|
self.weight = None
|
|||
|
|
else:
|
|||
|
|
self.weight = nn.Parameter(torch.zeros(dim))
|
|||
|
|
self.dense = None
|
|||
|
|
|
|||
|
|
def _norm(self, x: torch.Tensor) -> torch.Tensor:
|
|||
|
|
var = torch.mean(torch.square(x.float()), dim=-1, keepdim=True)
|
|||
|
|
return x.float() * torch.rsqrt(var + self.eps)
|
|||
|
|
|
|||
|
|
def forward(
|
|||
|
|
self,
|
|||
|
|
x: torch.Tensor,
|
|||
|
|
cond: torch.Tensor | None = None,
|
|||
|
|
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
|||
|
|
"""Return ``(normed, gate)``; ``gate`` is ``None`` in the unconditioned case."""
|
|||
|
|
dtype = x.dtype
|
|||
|
|
normed = self._norm(x)
|
|||
|
|
if self.dense is None:
|
|||
|
|
normed = normed * (1.0 + self.weight.float())
|
|||
|
|
return normed.to(dtype), None
|
|||
|
|
|
|||
|
|
if cond is None:
|
|||
|
|
# A conditioned norm silently falling back to an unconditioned one
|
|||
|
|
# would drop the entire timestep signal and still return a
|
|||
|
|
# well-shaped, finite tensor.
|
|||
|
|
raise ValueError(
|
|||
|
|
"Pi05AdaRMSNorm was built with cond_dim="
|
|||
|
|
f"{self.cond_dim} but called without an AdaRMS conditioning vector."
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if cond.shape[-1] != self.cond_dim:
|
|||
|
|
raise ValueError(f"Expected AdaRMS cond dim {self.cond_dim}, got {cond.shape[-1]}")
|
|||
|
|
|
|||
|
|
modulation = self.dense(cond.to(self.dense.weight.dtype))
|
|||
|
|
if x.ndim == 3:
|
|||
|
|
# (B, 3*dim) → (B, 1, 3*dim), broadcast across the token axis: the
|
|||
|
|
# timestep is a per-sample scalar, identical for every action token.
|
|||
|
|
modulation = modulation.unsqueeze(1)
|
|||
|
|
scale, shift, gate = modulation.chunk(3, dim=-1)
|
|||
|
|
normed = normed * (1.0 + scale.float()) + shift.float()
|
|||
|
|
return normed.to(dtype), gate.to(dtype)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _gated_residual(residual: torch.Tensor, out: torch.Tensor, gate: torch.Tensor | None) -> torch.Tensor:
|
|||
|
|
if gate is None:
|
|||
|
|
return residual + out
|
|||
|
|
return residual + gate * out
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ──────────────────────────────────────────────────────────────────────
|
|||
|
|
# Dual-backbone: PaliGemma + AdaRMS action expert
|
|||
|
|
# ──────────────────────────────────────────────────────────────────────
|
|||
|
|
def _repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
|
|||
|
|
"""Repeat KV heads for grouped-query attention.
|
|||
|
|
|
|||
|
|
``(B, num_kv_heads, S, D) → (B, num_kv_heads * n_rep, S, D)``.
|
|||
|
|
"""
|
|||
|
|
if n_rep == 1:
|
|||
|
|
return hidden_states
|
|||
|
|
b, nh, s, d = hidden_states.shape
|
|||
|
|
hidden_states = hidden_states[:, :, None, :, :].expand(b, nh, n_rep, s, d)
|
|||
|
|
return hidden_states.reshape(b, nh * n_rep, s, d)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _attend(query_states, key_states, value_states, attention_mask, num_kv_groups, scaling):
|
|||
|
|
"""Manual eager attention: ``softmax(Q Kᵀ · scale + mask) · V``.
|
|||
|
|
|
|||
|
|
The mask is sliced to ``key_states.shape[-2]`` so the same
|
|||
|
|
``(B, 1, Q, prefix+suffix)`` mask works in both the prefix pass (K length =
|
|||
|
|
prefix) and the suffix pass (K length = prefix + suffix).
|
|||
|
|
"""
|
|||
|
|
k = _repeat_kv(key_states, num_kv_groups)
|
|||
|
|
v = _repeat_kv(value_states, num_kv_groups)
|
|||
|
|
attn_weights = torch.matmul(query_states, k.transpose(2, 3)) * scaling
|
|||
|
|
if attention_mask is not None:
|
|||
|
|
attn_weights = attn_weights + attention_mask[:, :, :, : k.shape[-2]]
|
|||
|
|
attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
|
|||
|
|
return torch.matmul(attn_weights, v)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _match(tensor: torch.Tensor, module: nn.Module) -> torch.Tensor:
|
|||
|
|
"""Cast ``tensor`` to the dtype ``module``'s weight expects."""
|
|||
|
|
return tensor.to(module.weight.dtype) if tensor.dtype != module.weight.dtype else tensor
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _compute_layer_prefix_only(layer_idx, hidden_states, attention_mask, position_ids, paligemma):
|
|||
|
|
"""Run one PaliGemma LM layer on the prefix, returning the layer output and
|
|||
|
|
the post-RoPE ``(k, v)`` for the suffix pass.
|
|||
|
|
|
|||
|
|
Identical to π0: the prefix backbone is unchanged in π0.5 (no AdaRMS —
|
|||
|
|
there is no timestep in the prefix).
|
|||
|
|
"""
|
|||
|
|
model = paligemma.model.language_model
|
|||
|
|
layer = model.layers[layer_idx]
|
|||
|
|
residual = hidden_states
|
|||
|
|
x = _match(layer.input_layernorm(hidden_states), layer.self_attn.q_proj)
|
|||
|
|
|
|||
|
|
hidden_shape = (*x.shape[:-1], -1, layer.self_attn.head_dim)
|
|||
|
|
q = layer.self_attn.q_proj(x).view(hidden_shape).transpose(1, 2)
|
|||
|
|
k = layer.self_attn.k_proj(x).view(hidden_shape).transpose(1, 2)
|
|||
|
|
v = layer.self_attn.v_proj(x).view(hidden_shape).transpose(1, 2)
|
|||
|
|
|
|||
|
|
cos, sin = model.rotary_emb(v, position_ids)
|
|||
|
|
q, k = apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1)
|
|||
|
|
|
|||
|
|
att = _attend(
|
|||
|
|
q,
|
|||
|
|
k,
|
|||
|
|
v,
|
|||
|
|
attention_mask,
|
|||
|
|
num_kv_groups=layer.self_attn.num_key_value_groups,
|
|||
|
|
scaling=1.0 / math.sqrt(layer.self_attn.head_dim),
|
|||
|
|
)
|
|||
|
|
att = att.transpose(1, 2).reshape(q.shape[0], -1, q.shape[1] * layer.self_attn.head_dim)
|
|||
|
|
|
|||
|
|
out = layer.self_attn.o_proj(_match(att, layer.self_attn.o_proj)) + residual
|
|||
|
|
after_resid = out
|
|||
|
|
normed = layer.post_attention_layernorm(out)
|
|||
|
|
out = layer.mlp(_match(normed, layer.mlp.up_proj)) + after_resid
|
|||
|
|
return out, (k, v)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _compute_layer_suffix_only(
|
|||
|
|
layer_idx,
|
|||
|
|
hidden_states,
|
|||
|
|
prefix_kv,
|
|||
|
|
attention_mask,
|
|||
|
|
position_ids,
|
|||
|
|
gemma_expert,
|
|||
|
|
adarms_cond,
|
|||
|
|
):
|
|||
|
|
"""Run one action-expert layer on the suffix with AdaRMS conditioning.
|
|||
|
|
|
|||
|
|
This is where π0.5 diverges from π0. Both norms are
|
|||
|
|
:class:`Pi05AdaRMSNorm` and each returns a gate that scales its sublayer's
|
|||
|
|
contribution to the residual stream.
|
|||
|
|
"""
|
|||
|
|
layer = gemma_expert.model.layers[layer_idx]
|
|||
|
|
|
|||
|
|
residual = hidden_states
|
|||
|
|
x, gate = layer.input_layernorm(hidden_states, adarms_cond)
|
|||
|
|
x = _match(x, layer.self_attn.q_proj)
|
|||
|
|
|
|||
|
|
hidden_shape = (*x.shape[:-1], -1, layer.self_attn.head_dim)
|
|||
|
|
q = layer.self_attn.q_proj(x).view(hidden_shape).transpose(1, 2)
|
|||
|
|
k_suf = layer.self_attn.k_proj(x).view(hidden_shape).transpose(1, 2)
|
|||
|
|
v_suf = layer.self_attn.v_proj(x).view(hidden_shape).transpose(1, 2)
|
|||
|
|
|
|||
|
|
# RoPE frequencies are shared between PaliGemma and the expert.
|
|||
|
|
cos, sin = gemma_expert.model.rotary_emb(v_suf, position_ids)
|
|||
|
|
q, k_suf = apply_rotary_pos_emb(q, k_suf, cos, sin, unsqueeze_dim=1)
|
|||
|
|
|
|||
|
|
# Concatenate cached prefix K/V (possibly different dtype) with suffix K/V.
|
|||
|
|
k_prefix, v_prefix = prefix_kv
|
|||
|
|
k = torch.cat([k_prefix.to(k_suf.dtype), k_suf], dim=2)
|
|||
|
|
v = torch.cat([v_prefix.to(v_suf.dtype), v_suf], dim=2)
|
|||
|
|
|
|||
|
|
att = _attend(
|
|||
|
|
q,
|
|||
|
|
k,
|
|||
|
|
v,
|
|||
|
|
attention_mask,
|
|||
|
|
num_kv_groups=layer.self_attn.num_key_value_groups,
|
|||
|
|
scaling=1.0 / math.sqrt(layer.self_attn.head_dim),
|
|||
|
|
)
|
|||
|
|
att = att.transpose(1, 2).reshape(q.shape[0], -1, q.shape[1] * layer.self_attn.head_dim)
|
|||
|
|
|
|||
|
|
hidden_states = _gated_residual(residual, layer.self_attn.o_proj(_match(att, layer.self_attn.o_proj)), gate)
|
|||
|
|
|
|||
|
|
residual = hidden_states
|
|||
|
|
x, gate = layer.post_attention_layernorm(hidden_states, adarms_cond)
|
|||
|
|
return _gated_residual(residual, layer.mlp(_match(x, layer.mlp.up_proj)), gate)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class PaliGemmaWithActionExpertPi05(nn.Module):
|
|||
|
|
"""Dual-backbone transformer: PaliGemma (Gemma 2B) + AdaRMS expert (300M).
|
|||
|
|
|
|||
|
|
Same two-mode dispatch as π0 (``prefix_only`` / ``suffix_only``), with one
|
|||
|
|
structural change: after building a stock ``GemmaForCausalLM`` expert, every
|
|||
|
|
norm in it is swapped for a :class:`Pi05AdaRMSNorm` carrying a ``dense``
|
|||
|
|
conditioning projection.
|
|||
|
|
|
|||
|
|
Swapping in place — rather than subclassing ``GemmaModel`` as #4419 does —
|
|||
|
|
keeps the module tree, and therefore the checkpoint key layout, identical to
|
|||
|
|
the expert's stock layout apart from the norms themselves.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, vlm_config, action_expert_config):
|
|||
|
|
super().__init__()
|
|||
|
|
|
|||
|
|
# PaliGemma prefix: identical to π0. It sees no timestep, so no AdaRMS.
|
|||
|
|
vlm_config_hf = CONFIG_MAPPING["paligemma"]()
|
|||
|
|
vlm_config_hf._vocab_size = 257152
|
|||
|
|
vlm_config_hf.image_token_index = 257152
|
|||
|
|
vlm_config_hf.text_config.hidden_size = vlm_config.width
|
|||
|
|
vlm_config_hf.text_config.intermediate_size = vlm_config.mlp_dim
|
|||
|
|
vlm_config_hf.text_config.num_attention_heads = vlm_config.num_heads
|
|||
|
|
vlm_config_hf.text_config.head_dim = vlm_config.head_dim
|
|||
|
|
vlm_config_hf.text_config.num_hidden_layers = vlm_config.depth
|
|||
|
|
vlm_config_hf.text_config.num_key_value_heads = vlm_config.num_kv_heads
|
|||
|
|
vlm_config_hf.text_config.hidden_activation = "gelu_pytorch_tanh"
|
|||
|
|
vlm_config_hf.text_config.dtype = "float32"
|
|||
|
|
vlm_config_hf.text_config.vocab_size = 257152
|
|||
|
|
vlm_config_hf.vision_config.intermediate_size = 4304
|
|||
|
|
vlm_config_hf.vision_config.projection_dim = 2048
|
|||
|
|
vlm_config_hf.vision_config.projector_hidden_act = "gelu_fast"
|
|||
|
|
vlm_config_hf.vision_config.dtype = "float32"
|
|||
|
|
|
|||
|
|
action_expert_config_hf = CONFIG_MAPPING["gemma"](
|
|||
|
|
head_dim=action_expert_config.head_dim,
|
|||
|
|
hidden_size=action_expert_config.width,
|
|||
|
|
intermediate_size=action_expert_config.mlp_dim,
|
|||
|
|
num_attention_heads=action_expert_config.num_heads,
|
|||
|
|
num_hidden_layers=action_expert_config.depth,
|
|||
|
|
num_key_value_heads=action_expert_config.num_kv_heads,
|
|||
|
|
vocab_size=257152,
|
|||
|
|
hidden_activation="gelu_pytorch_tanh",
|
|||
|
|
dtype="float32",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
self.paligemma = PaliGemmaForConditionalGeneration(config=vlm_config_hf)
|
|||
|
|
self.gemma_expert = GemmaForCausalLM(config=action_expert_config_hf)
|
|||
|
|
# The action expert doesn't embed tokens — it only consumes the
|
|||
|
|
# suffix action embeddings we feed in.
|
|||
|
|
self.gemma_expert.model.embed_tokens = None
|
|||
|
|
|
|||
|
|
self.adarms_cond_dim = action_expert_config.width
|
|||
|
|
self._install_adarms_norms(action_expert_config)
|
|||
|
|
|
|||
|
|
def _install_adarms_norms(self, action_expert_config) -> None:
|
|||
|
|
"""Replace every action-expert RMSNorm with a conditioned AdaRMS norm."""
|
|||
|
|
expert = self.gemma_expert.model
|
|||
|
|
eps = getattr(self.gemma_expert.config, "rms_norm_eps", 1e-6)
|
|||
|
|
width = action_expert_config.width
|
|||
|
|
for layer in expert.layers:
|
|||
|
|
layer.input_layernorm = Pi05AdaRMSNorm(width, eps=eps, cond_dim=self.adarms_cond_dim)
|
|||
|
|
layer.post_attention_layernorm = Pi05AdaRMSNorm(width, eps=eps, cond_dim=self.adarms_cond_dim)
|
|||
|
|
expert.norm = Pi05AdaRMSNorm(width, eps=eps, cond_dim=self.adarms_cond_dim)
|
|||
|
|
|
|||
|
|
def embed_image(self, pixel_values: torch.Tensor) -> torch.Tensor:
|
|||
|
|
"""Encode images with SigLIP vision tower + PaliGemma projector.
|
|||
|
|
|
|||
|
|
The two steps are run explicitly rather than via
|
|||
|
|
``PaliGemmaModel.get_image_features`` because that helper divides the
|
|||
|
|
projector output by ``sqrt(text_hidden_size)``. Being explicit keeps the
|
|||
|
|
scale unambiguous and matches π0 exactly (SigLIP is unchanged in π0.5).
|
|||
|
|
"""
|
|||
|
|
# Shapes: pixel_values (B, 3, 224, 224) → SigLIP (B, 256, 1152)
|
|||
|
|
# → projector (B, 256, 2048)
|
|||
|
|
vision_outputs = self.paligemma.model.vision_tower(pixel_values)
|
|||
|
|
image_features = vision_outputs.last_hidden_state
|
|||
|
|
return self.paligemma.model.multi_modal_projector(image_features)
|
|||
|
|
|
|||
|
|
def embed_language_tokens(self, tokens: torch.Tensor) -> torch.Tensor:
|
|||
|
|
"""Embed language tokens, returning the ``* sqrt(hidden)``-scaled embedding.
|
|||
|
|
|
|||
|
|
The scaling location moved across transformers releases: at ≤5.3 the
|
|||
|
|
normalizer lives inside ``GemmaModel.forward`` (which we bypass), and at
|
|||
|
|
≥5.4 ``GemmaTextScaledWordEmbedding`` self-applies it. Detect and avoid
|
|||
|
|
double-scaling.
|
|||
|
|
"""
|
|||
|
|
embed_tokens = self.paligemma.model.language_model.embed_tokens
|
|||
|
|
lang_emb = embed_tokens(tokens)
|
|||
|
|
if getattr(embed_tokens, "embed_scale", None) is None:
|
|||
|
|
lang_emb = lang_emb * math.sqrt(lang_emb.shape[-1])
|
|||
|
|
return lang_emb
|
|||
|
|
|
|||
|
|
def forward(
|
|||
|
|
self,
|
|||
|
|
attention_mask: torch.Tensor | None = None,
|
|||
|
|
position_ids: torch.LongTensor | None = None,
|
|||
|
|
past_key_values=None,
|
|||
|
|
inputs_embeds: list[torch.Tensor | None] | None = None,
|
|||
|
|
use_cache: bool = False,
|
|||
|
|
adarms_cond: torch.Tensor | None = None,
|
|||
|
|
):
|
|||
|
|
"""Dispatch to prefix_only / suffix_only and return
|
|||
|
|
``([prefix_out, suffix_out], past_key_values_or_None)``.
|
|||
|
|
"""
|
|||
|
|
num_layers = self.paligemma.config.text_config.num_hidden_layers
|
|||
|
|
pali_lm = self.paligemma.model.language_model
|
|||
|
|
expert_lm = self.gemma_expert.model
|
|||
|
|
|
|||
|
|
if inputs_embeds[1] is None:
|
|||
|
|
hidden_states = inputs_embeds[0]
|
|||
|
|
kv_list: list[tuple[torch.Tensor, torch.Tensor]] = []
|
|||
|
|
for layer_idx in range(num_layers):
|
|||
|
|
hidden_states, kv = _compute_layer_prefix_only(
|
|||
|
|
layer_idx,
|
|||
|
|
hidden_states,
|
|||
|
|
attention_mask,
|
|||
|
|
position_ids,
|
|||
|
|
paligemma=self.paligemma,
|
|||
|
|
)
|
|||
|
|
kv_list.append(kv)
|
|||
|
|
hidden_states = pali_lm.norm(hidden_states)
|
|||
|
|
return [hidden_states, None], (kv_list if use_cache else None)
|
|||
|
|
|
|||
|
|
if inputs_embeds[0] is not None:
|
|||
|
|
raise ValueError(
|
|||
|
|
"PaliGemmaWithActionExpertPi05.forward only supports prefix-only "
|
|||
|
|
"or suffix-only dispatch; got both inputs_embeds populated."
|
|||
|
|
)
|
|||
|
|
if not isinstance(past_key_values, list):
|
|||
|
|
raise TypeError(
|
|||
|
|
"suffix_only forward expects past_key_values to be the "
|
|||
|
|
"list[(k, v)] produced by a previous prefix_only forward; "
|
|||
|
|
f"got {type(past_key_values)}"
|
|||
|
|
)
|
|||
|
|
hidden_states = inputs_embeds[1]
|
|||
|
|
for layer_idx in range(num_layers):
|
|||
|
|
hidden_states = _compute_layer_suffix_only(
|
|||
|
|
layer_idx,
|
|||
|
|
hidden_states,
|
|||
|
|
past_key_values[layer_idx],
|
|||
|
|
attention_mask,
|
|||
|
|
position_ids,
|
|||
|
|
gemma_expert=self.gemma_expert,
|
|||
|
|
adarms_cond=adarms_cond,
|
|||
|
|
)
|
|||
|
|
hidden_states, _ = expert_lm.norm(hidden_states, adarms_cond)
|
|||
|
|
return [None, hidden_states], None
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ──────────────────────────────────────────────────────────────────────
|
|||
|
|
# Main π0.5 Model
|
|||
|
|
# ──────────────────────────────────────────────────────────────────────
|
|||
|
|
class Pi05ForActionPrediction(nn.Module):
|
|||
|
|
"""π0.5 VLA model for robot action prediction via flow matching.
|
|||
|
|
|
|||
|
|
Inference flow:
|
|||
|
|
1. Embed prefix (images + language, where the language already carries the
|
|||
|
|
discretized state) → prefix tokens.
|
|||
|
|
2. Forward prefix through PaliGemma → layer-wise KV cache.
|
|||
|
|
3. For each denoising step ``t = 1.0, 1-dt, ..., 0``:
|
|||
|
|
a. Embed the timestep → an AdaRMS conditioning vector.
|
|||
|
|
b. Embed the suffix (action tokens only).
|
|||
|
|
c. Forward the suffix through the AdaRMS action expert.
|
|||
|
|
d. ``x_t = x_t + dt * v_t`` (Euler integration).
|
|||
|
|
4. Return ``x_0`` as the predicted action chunk.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, config, quant_config=None, prefix: str = ""):
|
|||
|
|
super().__init__()
|
|||
|
|
# ``quant_config`` is accepted for interface compatibility but unused —
|
|||
|
|
# quant_config is not plumbed through; the weight dtype comes from the pipeline.
|
|||
|
|
del quant_config
|
|||
|
|
self.config = config
|
|||
|
|
|
|||
|
|
self.action_dim = getattr(config, "max_action_dim", DEFAULT_ACTION_DIM)
|
|||
|
|
self.max_state_dim = getattr(config, "max_state_dim", self.action_dim)
|
|||
|
|
self.action_horizon = getattr(config, "chunk_size", DEFAULT_ACTION_HORIZON)
|
|||
|
|
self.num_inference_steps = getattr(config, "num_inference_steps", DEFAULT_NUM_INFERENCE_STEPS)
|
|||
|
|
|
|||
|
|
paligemma_variant = getattr(config, "paligemma_variant", "gemma_2b")
|
|||
|
|
action_expert_variant = getattr(config, "action_expert_variant", "gemma_300m")
|
|||
|
|
vlm_config = get_gemma_config(paligemma_variant)
|
|||
|
|
expert_config = get_gemma_config(action_expert_variant)
|
|||
|
|
self.vlm_width = vlm_config.width
|
|||
|
|
self.expert_width = expert_config.width
|
|||
|
|
|
|||
|
|
# Dual backbone
|
|||
|
|
self.paligemma_with_expert = PaliGemmaWithActionExpertPi05(vlm_config, expert_config)
|
|||
|
|
|
|||
|
|
# Action chunk projections.
|
|||
|
|
self.action_in_proj = nn.Linear(self.action_dim, self.expert_width)
|
|||
|
|
self.action_out_proj = nn.Linear(self.expert_width, self.action_dim)
|
|||
|
|
|
|||
|
|
# π0.5 timestep MLP: (W → W → W) with SiLU, feeding AdaRMS.
|
|||
|
|
# π0 instead has action_time_mlp_{in,out} of shape (2W → W → W) because
|
|||
|
|
# it concatenates the time embedding onto the action embedding.
|
|||
|
|
# NOTE: there is deliberately **no** ``state_proj`` here — that is the
|
|||
|
|
# π0-only continuous-state path.
|
|||
|
|
self.time_mlp_in = nn.Linear(self.expert_width, self.expert_width)
|
|||
|
|
self.time_mlp_out = nn.Linear(self.expert_width, self.expert_width)
|
|||
|
|
|
|||
|
|
# ── Prefix embedding ─────────────────────────────────────────────
|
|||
|
|
def embed_prefix(
|
|||
|
|
self,
|
|||
|
|
images: list[torch.Tensor],
|
|||
|
|
image_masks: list[torch.Tensor],
|
|||
|
|
lang_tokens: torch.Tensor,
|
|||
|
|
lang_masks: torch.Tensor,
|
|||
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|||
|
|
"""Build prefix embeddings, per-token padding mask, and AR mask.
|
|||
|
|
|
|||
|
|
Prefix tokens form ``[img_cam_0..., ..., lang_tokens...]`` with fully
|
|||
|
|
bidirectional attention (all-zero ``att_masks``). Identical to π0 — the
|
|||
|
|
state is inside ``lang_tokens``, so nothing here changes shape-wise.
|
|||
|
|
|
|||
|
|
Cameras are embedded one at a time; each call is ``(B, 3, 224, 224)``.
|
|||
|
|
The number of slots is fixed by ``config.max_cameras`` for the deployed
|
|||
|
|
model; missing cameras occupy their slot with a false image mask.
|
|||
|
|
"""
|
|||
|
|
num_views = len(images)
|
|||
|
|
if len(image_masks) != num_views:
|
|||
|
|
raise ValueError(
|
|||
|
|
f"images and image_masks must contain the same number of views, got {num_views} and {len(image_masks)}."
|
|||
|
|
)
|
|||
|
|
max_cameras = int(self.config.max_cameras)
|
|||
|
|
if num_views != max_cameras:
|
|||
|
|
raise ValueError(f"Expected exactly max_cameras={max_cameras} image views, got {num_views}.")
|
|||
|
|
|
|||
|
|
embs: list[torch.Tensor] = []
|
|||
|
|
pad_masks: list[torch.Tensor] = []
|
|||
|
|
att_masks: list[int] = []
|
|||
|
|
|
|||
|
|
for img, img_mask in zip(images, image_masks):
|
|||
|
|
img_emb = self.paligemma_with_expert.embed_image(img)
|
|||
|
|
bsize, num_img_embs = img_emb.shape[:2]
|
|||
|
|
embs.append(img_emb)
|
|||
|
|
pad_masks.append(img_mask[:, None].expand(bsize, num_img_embs))
|
|||
|
|
att_masks += [0] * num_img_embs
|
|||
|
|
|
|||
|
|
lang_emb = self.paligemma_with_expert.embed_language_tokens(lang_tokens)
|
|||
|
|
embs.append(lang_emb)
|
|||
|
|
pad_masks.append(lang_masks)
|
|||
|
|
att_masks += [0] * lang_emb.shape[1]
|
|||
|
|
|
|||
|
|
embs = torch.cat(embs, dim=1)
|
|||
|
|
pad_masks = torch.cat(pad_masks, dim=1)
|
|||
|
|
att_masks = torch.tensor(att_masks, dtype=torch.bool, device=embs.device)
|
|||
|
|
att_masks = att_masks[None, :].expand(pad_masks.shape[0], -1)
|
|||
|
|
|
|||
|
|
return embs, pad_masks, att_masks
|
|||
|
|
|
|||
|
|
# ── Timestep + suffix embedding ──────────────────────────────────
|
|||
|
|
def embed_timestep(self, timestep: torch.Tensor) -> torch.Tensor:
|
|||
|
|
"""Timestep → AdaRMS conditioning vector ``(B, expert_width)``.
|
|||
|
|
|
|||
|
|
``silu(time_mlp_out(silu(time_mlp_in(sinusoid(t)))))``. The trailing
|
|||
|
|
SiLU is part of the reference implementation — dropping it is a silent
|
|||
|
|
numerical error, not a crash.
|
|||
|
|
"""
|
|||
|
|
model_dtype = self.action_in_proj.weight.dtype
|
|||
|
|
time_emb = create_sinusoidal_pos_embedding(
|
|||
|
|
timestep,
|
|||
|
|
self.action_in_proj.out_features,
|
|||
|
|
min_period=getattr(self.config, "min_period", 4e-3),
|
|||
|
|
max_period=getattr(self.config, "max_period", 4.0),
|
|||
|
|
device=timestep.device,
|
|||
|
|
).to(dtype=model_dtype)
|
|||
|
|
time_cond = self.time_mlp_in(time_emb)
|
|||
|
|
time_cond = F.silu(time_cond)
|
|||
|
|
time_cond = self.time_mlp_out(time_cond)
|
|||
|
|
return F.silu(time_cond)
|
|||
|
|
|
|||
|
|
def embed_suffix(
|
|||
|
|
self,
|
|||
|
|
noisy_actions: torch.Tensor,
|
|||
|
|
timestep: torch.Tensor,
|
|||
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
|||
|
|
"""Build the suffix: **action tokens only**, plus the AdaRMS condition.
|
|||
|
|
|
|||
|
|
π0's suffix is ``[state_token, action_tokens×H]`` with an AR mask of
|
|||
|
|
``[1, 1, 0...]``. π0.5 has no state token, so the suffix is
|
|||
|
|
``[action_tokens×H]`` and the mask is ``[1] + [0]*(H-1)``: the first
|
|||
|
|
action token opens a causal block and the rest attend bidirectionally
|
|||
|
|
within it.
|
|||
|
|
"""
|
|||
|
|
model_dtype = self.action_in_proj.weight.dtype
|
|||
|
|
noisy_actions = noisy_actions.to(dtype=model_dtype)
|
|||
|
|
|
|||
|
|
time_cond = self.embed_timestep(timestep)
|
|||
|
|
action_emb = self.action_in_proj(noisy_actions) # (B, H, W)
|
|||
|
|
|
|||
|
|
bsize, action_len = action_emb.shape[:2]
|
|||
|
|
pad_masks = torch.ones(bsize, action_len, dtype=torch.bool, device=action_emb.device)
|
|||
|
|
att_masks = torch.tensor(
|
|||
|
|
[1] + [0] * (self.action_horizon - 1),
|
|||
|
|
dtype=action_emb.dtype,
|
|||
|
|
device=action_emb.device,
|
|||
|
|
)[None, :].expand(bsize, -1)
|
|||
|
|
return action_emb, pad_masks, att_masks, time_cond
|
|||
|
|
|
|||
|
|
# ── Denoising step ───────────────────────────────────────────────
|
|||
|
|
def denoise_step(
|
|||
|
|
self,
|
|||
|
|
prefix_pad_masks: torch.Tensor,
|
|||
|
|
past_key_values,
|
|||
|
|
x_t: torch.Tensor,
|
|||
|
|
timestep: torch.Tensor,
|
|||
|
|
) -> torch.Tensor:
|
|||
|
|
"""Apply one flow-matching denoising step: predict ``v_t`` from ``x_t``.
|
|||
|
|
|
|||
|
|
Signature differs from π0's by exactly one argument: no ``state``.
|
|||
|
|
"""
|
|||
|
|
suffix_embs, suffix_pad_masks, suffix_att_masks, time_cond = self.embed_suffix(x_t, timestep)
|
|||
|
|
|
|||
|
|
batch_size = prefix_pad_masks.shape[0]
|
|||
|
|
suffix_len = suffix_pad_masks.shape[1]
|
|||
|
|
prefix_len = prefix_pad_masks.shape[1]
|
|||
|
|
|
|||
|
|
prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, suffix_len, prefix_len)
|
|||
|
|
suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks)
|
|||
|
|
full_att_2d_masks = torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2)
|
|||
|
|
|
|||
|
|
# Position IDs continue from where the prefix's last valid token left off.
|
|||
|
|
prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None]
|
|||
|
|
position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1
|
|||
|
|
|
|||
|
|
full_att_2d_masks_4d = prepare_attention_masks_4d(full_att_2d_masks)
|
|||
|
|
|
|||
|
|
outputs_embeds, _ = self.paligemma_with_expert.forward(
|
|||
|
|
attention_mask=full_att_2d_masks_4d,
|
|||
|
|
position_ids=position_ids,
|
|||
|
|
past_key_values=past_key_values,
|
|||
|
|
inputs_embeds=[None, suffix_embs],
|
|||
|
|
use_cache=False,
|
|||
|
|
adarms_cond=time_cond,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# Every suffix token is an action token in π0.5 (no state token to drop).
|
|||
|
|
suffix_out = outputs_embeds[1][:, -self.action_horizon :]
|
|||
|
|
suffix_out = suffix_out.to(dtype=self.action_out_proj.weight.dtype)
|
|||
|
|
return self.action_out_proj(suffix_out)
|
|||
|
|
|
|||
|
|
# ── Full action generation ───────────────────────────────────────
|
|||
|
|
@torch.no_grad()
|
|||
|
|
def sample_actions(
|
|||
|
|
self,
|
|||
|
|
images: list[torch.Tensor],
|
|||
|
|
image_masks: list[torch.Tensor],
|
|||
|
|
lang_tokens: torch.Tensor,
|
|||
|
|
lang_masks: torch.Tensor,
|
|||
|
|
noise: torch.Tensor | None = None,
|
|||
|
|
num_steps: int | None = None,
|
|||
|
|
generator: torch.Generator | list[torch.Generator] | None = None,
|
|||
|
|
) -> torch.Tensor:
|
|||
|
|
"""Generate an action chunk via iterative flow-matching denoising.
|
|||
|
|
|
|||
|
|
Convention: ``t=1`` is noise, ``t=0`` is the target — opposite of the
|
|||
|
|
published π0 paper but matching both OpenPI and LeRobot.
|
|||
|
|
|
|||
|
|
Takes no ``state``: π0.5's state rides inside ``lang_tokens``.
|
|||
|
|
"""
|
|||
|
|
if num_steps is None:
|
|||
|
|
num_steps = self.num_inference_steps
|
|||
|
|
|
|||
|
|
bsize = lang_tokens.shape[0]
|
|||
|
|
device = lang_tokens.device
|
|||
|
|
if noise is None:
|
|||
|
|
noise_shape = (self.action_horizon, self.action_dim)
|
|||
|
|
if isinstance(generator, list):
|
|||
|
|
if len(generator) != bsize:
|
|||
|
|
raise ValueError(f"Expected {bsize} generators, got {len(generator)}.")
|
|||
|
|
noise = torch.stack(
|
|||
|
|
[torch.randn(noise_shape, dtype=torch.float32, device=device, generator=item) for item in generator]
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
noise = torch.randn(
|
|||
|
|
bsize,
|
|||
|
|
*noise_shape,
|
|||
|
|
dtype=torch.float32,
|
|||
|
|
device=device,
|
|||
|
|
generator=generator,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 1. Prefix embeddings + mask building.
|
|||
|
|
prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(
|
|||
|
|
images, image_masks, lang_tokens, lang_masks
|
|||
|
|
)
|
|||
|
|
prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
|
|||
|
|
prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
|
|||
|
|
prefix_att_2d_masks_4d = prepare_attention_masks_4d(prefix_att_2d_masks)
|
|||
|
|
|
|||
|
|
# 2. Forward prefix through PaliGemma LM, producing a list[(k, v)] cache.
|
|||
|
|
_, past_key_values = self.paligemma_with_expert.forward(
|
|||
|
|
attention_mask=prefix_att_2d_masks_4d,
|
|||
|
|
position_ids=prefix_position_ids,
|
|||
|
|
past_key_values=None,
|
|||
|
|
inputs_embeds=[prefix_embs, None],
|
|||
|
|
use_cache=True,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 3. Euler-integrated denoising from t=1 down to t=0.
|
|||
|
|
dt = -1.0 / num_steps
|
|||
|
|
x_t = noise
|
|||
|
|
for step in range(num_steps):
|
|||
|
|
t = 1.0 + step * dt
|
|||
|
|
time_tensor = torch.full((bsize,), t, dtype=torch.float32, device=device)
|
|||
|
|
v_t = self.denoise_step(
|
|||
|
|
prefix_pad_masks=prefix_pad_masks,
|
|||
|
|
past_key_values=past_key_values,
|
|||
|
|
x_t=x_t,
|
|||
|
|
timestep=time_tensor,
|
|||
|
|
)
|
|||
|
|
x_t = x_t + dt * v_t
|
|||
|
|
return x_t
|
|||
|
|
|
|||
|
|
# ── Weight loading ───────────────────────────────────────────────
|
|||
|
|
def load_weights(
|
|||
|
|
self,
|
|||
|
|
weights: Iterable[tuple[str, torch.Tensor]],
|
|||
|
|
*,
|
|||
|
|
strict: bool = True,
|
|||
|
|
):
|
|||
|
|
"""Load and audit a LeRobot π0.5 safetensors checkpoint.
|
|||
|
|
|
|||
|
|
Same remap rules as π0 (strip the ``model.`` prefix, flatten→nested
|
|||
|
|
PaliGemma submodules, tied ``lm_head`` → ``embed_tokens``, version-robust
|
|||
|
|
SigLIP nesting), plus two π0.5-specific ones:
|
|||
|
|
|
|||
|
|
* ``action_time_mlp_{in,out}`` → ``time_mlp_{in,out}``: some
|
|||
|
|
checkpoints were exported under the π0 parameter names.
|
|||
|
|
* ``state_proj.*`` is reported, not silently dropped. A π0.5
|
|||
|
|
checkpoint should not contain it; its presence usually means a π0
|
|||
|
|
checkpoint was pointed at the π0.5 model class, which would
|
|||
|
|
otherwise run happily with a randomly-initialized action expert.
|
|||
|
|
|
|||
|
|
The action-expert norms are AdaRMS here, so they expose ``dense.weight``
|
|||
|
|
/ ``dense.bias`` and no plain ``weight``. A checkpoint that carries a
|
|||
|
|
plain expert-norm ``weight`` is a π0-shaped checkpoint; that too is
|
|||
|
|
rejected rather than skipped. ``strict=False`` exists only for focused
|
|||
|
|
remapping unit tests that intentionally provide a partial state dict;
|
|||
|
|
the serving path always uses the strict default.
|
|||
|
|
"""
|
|||
|
|
params_dict = dict(self.named_parameters())
|
|||
|
|
buffers_dict = dict(self.named_buffers())
|
|||
|
|
|
|||
|
|
_PALIGEMMA_SUBMODULES = ("vision_tower", "multi_modal_projector", "language_model")
|
|||
|
|
_EXPERT_PREFIX = "paligemma_with_expert.gemma_expert.model."
|
|||
|
|
|
|||
|
|
def _remap(name: str) -> str:
|
|||
|
|
# Strip the leading "model." that LeRobot's PI05Policy wrapper adds.
|
|||
|
|
if name.startswith("model."):
|
|||
|
|
name = name[len("model.") :]
|
|||
|
|
|
|||
|
|
# π0-style timestep MLP names → π0.5 names.
|
|||
|
|
if name.startswith("action_time_mlp_in."):
|
|||
|
|
name = "time_mlp_in." + name[len("action_time_mlp_in.") :]
|
|||
|
|
elif name.startswith("action_time_mlp_out."):
|
|||
|
|
name = "time_mlp_out." + name[len("action_time_mlp_out.") :]
|
|||
|
|
|
|||
|
|
# Nested PaliGemma layout.
|
|||
|
|
for sub in _PALIGEMMA_SUBMODULES:
|
|||
|
|
flat = f"paligemma_with_expert.paligemma.{sub}."
|
|||
|
|
nested = f"paligemma_with_expert.paligemma.model.{sub}."
|
|||
|
|
if name.startswith(flat) and not name.startswith(nested):
|
|||
|
|
return nested + name[len(flat) :]
|
|||
|
|
|
|||
|
|
if name == "paligemma_with_expert.paligemma.lm_head.weight":
|
|||
|
|
return "paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"
|
|||
|
|
|
|||
|
|
return name
|
|||
|
|
|
|||
|
|
def _fix_vision_tower(name: str) -> str:
|
|||
|
|
"""Reconcile SigLIP nesting across transformers versions (≤5.3 wraps
|
|||
|
|
the encoder in ``vision_tower.vision_model.*``; ≥5.4 flattens it)."""
|
|||
|
|
vt = "paligemma_with_expert.paligemma.model.vision_tower."
|
|||
|
|
if not name.startswith(vt):
|
|||
|
|
return name
|
|||
|
|
rest = name[len(vt) :]
|
|||
|
|
if name in params_dict or name in buffers_dict:
|
|||
|
|
return name
|
|||
|
|
if rest.startswith("vision_model."):
|
|||
|
|
candidate = vt + rest[len("vision_model.") :]
|
|||
|
|
else:
|
|||
|
|
candidate = vt + "vision_model." + rest
|
|||
|
|
return candidate if (candidate in params_dict or candidate in buffers_dict) else name
|
|||
|
|
|
|||
|
|
loaded = 0
|
|||
|
|
skipped: list[str] = []
|
|||
|
|
pi0_shaped: list[str] = []
|
|||
|
|
filled_params: set = set()
|
|||
|
|
|
|||
|
|
for name, loaded_weight in weights:
|
|||
|
|
mapped = _fix_vision_tower(_remap(name))
|
|||
|
|
|
|||
|
|
# Diagnose π0-shaped keys instead of dropping them quietly.
|
|||
|
|
is_expert_norm_weight = mapped.startswith(_EXPERT_PREFIX) and (
|
|||
|
|
mapped.endswith("input_layernorm.weight")
|
|||
|
|
or mapped.endswith("post_attention_layernorm.weight")
|
|||
|
|
or mapped == _EXPERT_PREFIX + "norm.weight"
|
|||
|
|
)
|
|||
|
|
if mapped.startswith("state_proj.") or is_expert_norm_weight:
|
|||
|
|
pi0_shaped.append(mapped)
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
if mapped in params_dict:
|
|||
|
|
param = params_dict[mapped]
|
|||
|
|
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
|||
|
|
weight_loader(param, loaded_weight)
|
|||
|
|
loaded += 1
|
|||
|
|
filled_params.add(mapped)
|
|||
|
|
elif mapped in buffers_dict:
|
|||
|
|
buffers_dict[mapped].copy_(loaded_weight)
|
|||
|
|
loaded += 1
|
|||
|
|
filled_params.add(mapped)
|
|||
|
|
else:
|
|||
|
|
skipped.append(mapped)
|
|||
|
|
|
|||
|
|
# LeRobot stores PaliGemma's tied text embedding as lm_head.weight; keep
|
|||
|
|
# lm_head filled so the tied-weight state stays consistent.
|
|||
|
|
embed_key = "paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"
|
|||
|
|
lm_head_key = "paligemma_with_expert.paligemma.lm_head.weight"
|
|||
|
|
if embed_key in filled_params and lm_head_key in params_dict and lm_head_key not in filled_params:
|
|||
|
|
params_dict[lm_head_key].data.copy_(params_dict[embed_key].data)
|
|||
|
|
filled_params.add(lm_head_key)
|
|||
|
|
|
|||
|
|
# Reverse audit: any model param that got no checkpoint tensor at all
|
|||
|
|
# would be running with random init.
|
|||
|
|
missing_params: list[str] = []
|
|||
|
|
for pname in params_dict:
|
|||
|
|
if pname in filled_params:
|
|||
|
|
continue
|
|||
|
|
if "rotary_emb" in pname or pname.endswith(".inv_freq"):
|
|||
|
|
continue
|
|||
|
|
missing_params.append(pname)
|
|||
|
|
|
|||
|
|
parts: list[str] = []
|
|||
|
|
if pi0_shaped:
|
|||
|
|
parts.append(
|
|||
|
|
f"{len(pi0_shaped)} checkpoint key(s) are π0-shaped, not π0.5-shaped (first 5: {pi0_shaped[:5]})"
|
|||
|
|
)
|
|||
|
|
if skipped:
|
|||
|
|
parts.append(f"{len(skipped)} checkpoint key(s) had no model target (first 5: {skipped[:5]})")
|
|||
|
|
if missing_params:
|
|||
|
|
parts.append(f"{len(missing_params)} model param(s) received no weight (first 5: {missing_params[:5]})")
|
|||
|
|
|
|||
|
|
if parts and strict:
|
|||
|
|
raise RuntimeError("Incomplete or incompatible π0.5 checkpoint: " + "; ".join(parts))
|
|||
|
|
if parts:
|
|||
|
|
logger.debug("π0.5 partial test load: %d tensors loaded — %s.", loaded, "; ".join(parts))
|
|||
|
|
else:
|
|||
|
|
logger.info("π0.5 load_weights: %d tensors loaded, 0 skipped, 0 missing.", loaded)
|
|||
|
|
return filled_params
|
|||
|
|
|
|||
|
|
|
|||
|
|
EntryClass = Pi05ForActionPrediction
|