# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Inference-only π0 (Pi-Zero) VLA math kernel for vllm-omni.

A self-contained inference kernel: only the math that turns a robot observation
into an action chunk, with no serving/request glue — validated bit-for-bit
against the LeRobot ``PI0Policy`` reference (``max|Δ| < 1e-4``).

π0 = PaliGemma (SigLIP vision + Gemma 2B LM) prefix + Gemma 300M action expert
suffix + flow-matching head. Inference:

1. Embed prefix (images + language) → bidirectional prefix tokens.
2. Forward prefix through PaliGemma → a per-layer ``list[(k, v)]`` KV cache.
3. For each Euler step ``t = 1.0, 1-dt, ..., 0``: embed the suffix
   (state + noisy actions + timestep), run the action expert with cross
   attention over the cached prefix K/V, predict velocity ``v_t``, and
   integrate ``x_t = x_t + dt * v_t``.
4. Return ``x_0`` as the action chunk ``(batch, action_horizon, action_dim)``.

We walk Gemma decoder layers ourselves in ``_compute_layer_*`` so we never go
through ``GemmaModel.forward`` (which has version-dependent mask/KV-cache
behaviour). Targets modern transformers (PaliGemmaForConditionalGeneration with
an inner PaliGemmaModel + DynamicCache).

Reference implementations:
   - OpenPI: openpi/src/openpi/models_pytorch/pi0_pytorch.py, gemma_pytorch.py
   - LeRobot: lerobot/src/lerobot/policies/pi0/modeling_pi0.py
Weight source: https://huggingface.co/lerobot/pi0_base
"""

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 = 48
DEFAULT_NUM_INFERENCE_STEPS = 10
DEFAULT_IMAGE_RESOLUTION = (224, 224)

# 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 _build_norm_buffers(norm_stats: dict | None, key: str) -> dict[str, torch.Tensor] | None:
    """Parse a ``norm_stats[key]`` entry into CPU tensors, or ``None``.

    Matches LeRobot's ``NormalizationMode``:
      - ``mean_std`` : ``forward = (x - mean) / std``, ``inverse = x * std + mean``
      - ``min_max``  : forward maps into ``[-1, 1]``, inverse maps back.
    """
    if not norm_stats or not isinstance(norm_stats, dict):
        return None
    entry = norm_stats.get(key)
    if not entry:
        return None
    mode = str(entry.get("mode", "mean_std")).lower()
    if mode == "mean_std":
        mean = entry.get("mean")
        std = entry.get("std")
        if mean is None or std is None:
            return None
        return {
            "mode": mode,
            "mean": torch.as_tensor(mean, dtype=torch.float32),
            "std": torch.as_tensor(std, dtype=torch.float32),
        }
    if mode == "min_max":
        lo = entry.get("min")
        hi = entry.get("max")
        if lo is None or hi is None:
            return None
        return {
            "mode": mode,
            "min": torch.as_tensor(lo, dtype=torch.float32),
            "max": torch.as_tensor(hi, dtype=torch.float32),
        }
    return None


def _apply_norm(
    x: torch.Tensor,
    stats: dict[str, torch.Tensor] | None,
    inverse: bool,
    eps: float = 1e-8,
) -> torch.Tensor:
    """Apply (un)normalization using the given stats. No-op if ``stats`` is None.

    Broadcasts over leading dims; stats vectors align with the last dim. When
    ``x`` has more last-dim entries than the stats (π0 pads state/action to
    ``max_dim``), only the first ``len(stats)`` entries are transformed and
    the padded tail is left untouched.
    """
    if stats is None:
        return x
    mode = stats["mode"]
    if mode == "mean_std":
        mean = stats["mean"].to(device=x.device, dtype=x.dtype)
        std = stats["std"].to(device=x.device, dtype=x.dtype)
        valid = mean.shape[0]
        head = x[..., :valid]
        if inverse:
            head = head * std + mean
        else:
            head = (head - mean) / (std + eps)
        if valid == x.shape[-1]:
            return head
        return torch.cat([head, x[..., valid:]], dim=-1)
    if mode == "min_max":
        lo = stats["min"].to(device=x.device, dtype=x.dtype)
        hi = stats["max"].to(device=x.device, dtype=x.dtype)
        valid = lo.shape[0]
        head = x[..., :valid]
        denom = (hi - lo).clamp_min(eps)
        if inverse:
            head = (head + 1.0) * 0.5 * denom + lo
        else:
            head = 2.0 * (head - lo) / denom - 1.0
        if valid == x.shape[-1]:
            return head
        return torch.cat([head, x[..., valid:]], dim=-1)
    return x


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``.
    Ref: openpi PI0Pytorch._prepare_attention_masks_4d
    """
    att_2d_masks_4d = att_2d_masks[:, None, :, :]
    return torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE)


# ──────────────────────────────────────────────────────────────────────
# Dual-backbone: PaliGemma + Action Expert
# Ref: openpi/models_pytorch/gemma_pytorch.py PaliGemmaWithExpertModel
# Ref: lerobot/policies/pi0/modeling_pi0.py PaliGemmaWithExpertModel
# ──────────────────────────────────────────────────────────────────────


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)``.
    Matches ``transformers.models.gemma.modeling_gemma.repeat_kv``.
    """
    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``.

    Important: we slice ``attention_mask[..., :key_states.shape[-2]]`` so that
    the same ``(B, 1, Q, prefix+suffix)`` mask produced by ``denoise_step``
    can be re-used both in the prefix pass (K length = prefix) and in the
    suffix pass (K length = prefix + suffix after we concat the cached prefix).
    Stock transformers' ``eager_attention_forward`` adds the mask without
    slicing and therefore breaks when the two dims disagree.
    """
    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 _compute_layer_prefix_only(layer_idx, hidden_states, attention_mask, position_ids, paligemma):
    """Run one PaliGemma LM layer on the prefix only, returning the layer
    output **and** the post-RoPE ``(k, v)`` so the caller can cache them for
    the suffix pass."""
    model = paligemma.model.language_model
    layer = model.layers[layer_idx]
    residual = hidden_states
    x = layer.input_layernorm(hidden_states)

    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(att) + residual
    after_resid = out
    out = layer.mlp(layer.post_attention_layernorm(out)) + after_resid
    return out, (k, v)


def _compute_layer_suffix_only(
    layer_idx,
    hidden_states,
    prefix_kv,
    attention_mask,
    position_ids,
    paligemma,
    gemma_expert,
):
    """Run one action-expert layer on the suffix, attending to the cached
    prefix K/V concatenated with freshly computed suffix K/V.

    ``prefix_kv`` is a ``(k_prefix, v_prefix)`` tuple, shape
    ``(B, num_kv_heads, prefix_len, head_dim)``, post-RoPE, produced by the
    corresponding ``_compute_layer_prefix_only`` call. Bypasses
    ``GemmaModel.forward`` entirely — no HF mask rebuild, no cache format
    gymnastics.
    """
    layer = gemma_expert.model.layers[layer_idx]
    residual = hidden_states
    x = layer.input_layernorm(hidden_states)

    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)

    out = layer.self_attn.o_proj(att) + residual
    after_resid = out
    out = layer.mlp(layer.post_attention_layernorm(out)) + after_resid
    return out


class PaliGemmaWithActionExpert(nn.Module):
    """Dual-backbone transformer: PaliGemma (Gemma 2B) + Action Expert (Gemma 300M).

    ``forward`` has two inference modes:
      - **prefix_only**: ``inputs_embeds=[prefix, None]`` + ``use_cache=True``
        → compute prefix hidden states + a layer-wise K/V cache.
      - **suffix_only**: ``inputs_embeds=[None, suffix]`` + the cache from the
        previous prefix pass → compute expert hidden states with cross-attention
        over the concatenated (prefix, suffix) K/V.

    Both modes walk Gemma decoder layers manually through ``_compute_layer_*``
    instead of ``GemmaModel.forward`` — that way we own the attention mask
    handling and the KV cache format (a plain ``list[(k, v)]`` one entry
    per layer).

    Ref: lerobot/policies/pi0/modeling_pi0.py PaliGemmaWithExpertModel
    """

    def __init__(self, vlm_config, action_expert_config):
        super().__init__()

        # Build HF PaliGemma config from the variant dims we expose.
        # NOTE: we do NOT set use_adarms / adarms_cond_dim — those are
        # OpenPI-only attributes used by π0.5, not π0.
        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"
        # transformers ≥ 5 uses ``dtype``; older versions still accept
        # ``torch_dtype``. We set the canonical name; HF's PretrainedConfig
        # normalizes one to the other internally.
        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 state/action embeddings we feed in.
        self.gemma_expert.model.embed_tokens = None

    def embed_image(self, pixel_values: torch.Tensor) -> torch.Tensor:
        """Encode images with SigLIP vision tower + PaliGemma projector.

        We run the two steps explicitly instead of calling
        ``PaliGemmaModel.get_image_features`` because that convenience
        method historically has returned either the projected tensor (old)
        or the raw vision-tower output (new). Being explicit makes us
        independent of which transformers release we're on.
        """
        # 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 with PaliGemma's embedding table, returning the
        ``* sqrt(hidden)``-scaled embedding (PaliGemma/Gemma convention).

        The scaling location moved across transformers releases:
          * ≤ 5.3: ``embed_tokens`` is a plain ``nn.Embedding`` and the
            ``* sqrt(hidden)`` normalizer lives inside ``GemmaModel.forward``
            (which we bypass) — so we apply it explicitly here.
          * ≥ 5.4: ``embed_tokens`` is a ``GemmaTextScaledWordEmbedding`` that
            self-applies ``embed_scale = hidden_size ** 0.5`` — applying it again
            would double-scale (≈45×). We detect this and skip the manual scale.
        This makes ``embed_language_tokens`` return the canonical scaled embedding
        on every transformers version.
        """
        embed_tokens = self.paligemma.model.language_model.embed_tokens
        lang_emb = embed_tokens(tokens)
        # If the embedding already self-scales (transformers ≥ 5.4), don't repeat it.
        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,
    ):
        """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:
            # Prefix-only: PaliGemma LM on (images + language) tokens; the
            # per-layer post-RoPE K/V is collected into a list that the
            # suffix pass will consume directly.
            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(
                "PaliGemmaWithActionExpert.forward only supports prefix-only "
                "or suffix-only dispatch; got both inputs_embeds populated."
            )
        # Suffix-only: action expert on (state + noisy_actions + time),
        # attending to the cached prefix K/V.
        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,
                paligemma=self.paligemma,
                gemma_expert=self.gemma_expert,
            )
        hidden_states = expert_lm.norm(hidden_states)
        return [None, hidden_states], None


# ──────────────────────────────────────────────────────────────────────
# Main π0 Model
# ──────────────────────────────────────────────────────────────────────


class Pi0ForActionPrediction(nn.Module):
    """π0 VLA model for robot action prediction via flow matching.

    Inference flow:
      1. Embed prefix (images + language) → prefix tokens.
      2. Forward prefix through PaliGemma → layer-wise KV cache.
      3. For each denoising step ``t = 1.0, 1-dt, ..., 0``:
         a. Embed suffix (state + x_t + timestep) → suffix tokens.
         b. Forward suffix through the action expert with the prefix cache.
         c. ``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 —
        # π0 runs in full precision for flow-matching parity.
        del quant_config
        self.config = config

        self.action_dim = getattr(config, "max_action_dim", DEFAULT_ACTION_DIM)
        # ``max_state_dim`` is independent from ``max_action_dim`` — e.g. Aloha
        # has state_dim=14, action_dim=32, both padded to their own max.  Fall
        # back to action_dim for older configs that only expose a single size.
        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 = PaliGemmaWithActionExpert(vlm_config, expert_config)

        # Action chunk projections (openpi pi0_pytorch.py).
        self.action_in_proj = nn.Linear(self.action_dim, self.expert_width)
        self.action_out_proj = nn.Linear(self.expert_width, self.action_dim)

        # State is projected to the expert dim once (π0 only, not π0.5).
        # Ref: lerobot modeling_pi0.py ``nn.Linear(config.max_state_dim, ...)``
        self.state_proj = nn.Linear(self.max_state_dim, self.expert_width)

        # Timestep + action fusion MLP (2W → W → W).
        self.action_time_mlp_in = nn.Linear(2 * self.expert_width, self.expert_width)
        self.action_time_mlp_out = nn.Linear(self.expert_width, self.expert_width)

        # Optional mean/std normalization stats (LeRobot NormalizerProcessorStep
        # semantics). Format on the config (all optional):
        #     norm_stats = {
        #         "state":  {"mode": "mean_std", "mean": [...], "std": [...]},
        #         "action": {"mode": "mean_std", "mean": [...], "std": [...]},
        #     }
        # ``mode`` may also be ``"min_max"``.
        #
        # Missing norm_stats is logged at INFO (not WARNING) because the most
        # common setup is "client pre-normalizes state and post-normalizes
        # actions using dataset stats" — that's a supported mode, not an error.
        self._state_norm = _build_norm_buffers(getattr(config, "norm_stats", None), "state")
        self._action_norm = _build_norm_buffers(getattr(config, "norm_stats", None), "action")
        if self._state_norm is None:
            logger.info(
                "π0: no state normalization stats on config.norm_stats — state values will pass through unchanged."
            )
        if self._action_norm is None:
            logger.info(
                "π0: no action normalization stats on config.norm_stats — "
                "returned actions are in the model's normalized space."
            )

    # ── State / action normalization ─────────────────────────────────

    def _normalize_state(self, state: torch.Tensor) -> torch.Tensor:
        return _apply_norm(state, self._state_norm, inverse=False)

    def _unnormalize_actions(self, actions: torch.Tensor) -> torch.Tensor:
        return _apply_norm(actions, self._action_norm, inverse=True)

    # ── 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 the prefix embeddings, per-token padding mask, and AR mask.

        Prefix tokens form a contiguous sequence
        ``[img_cam_0..., img_cam_1..., ..., lang_tokens...]`` with
        bidirectional attention; the returned ``att_masks`` are all zeros
        because each token is free to attend to every other prefix token (the
        suffix pass will put a causal boundary right before the state token).

        Ref: openpi PI0Pytorch.embed_prefix
        """
        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

        # ``embed_language_tokens`` already returns the ``* sqrt(hidden)``-scaled
        # embedding (version-robust: self-scaled at transformers ≥ 5.4, manually
        # scaled at ≤ 5.3). The image slice went through the projector and is
        # NOT scaled. See embed_language_tokens for the version details.
        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

    # ── Suffix embedding ─────────────────────────────────────────────

    def embed_suffix(
        self,
        state: torch.Tensor,
        noisy_actions: torch.Tensor,
        timestep: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Build the suffix embeddings + masks: ``[state_token, action_tokens×H]``.

        AR mask layout: ``[1, 1, 0, 0, ..., 0]`` — the state token is a causal
        boundary (no later token attends backwards through it onto prefix
        *by mistake*), the first action token starts a new causal block, and
        the rest of the action tokens attend to each other bidirectionally.

        Ref: openpi PI0Pytorch.embed_suffix
        """
        model_dtype = self.state_proj.weight.dtype
        state = state.to(dtype=model_dtype)
        noisy_actions = noisy_actions.to(dtype=model_dtype)
        device = state.device

        # State → (B, 1, W)
        state_emb = self.state_proj(state)[:, None, :]
        bsize = state_emb.shape[0]

        # Sinusoidal timestep → (B, W) → expand across the horizon.
        time_emb = create_sinusoidal_pos_embedding(
            timestep,
            self.action_in_proj.out_features,
            min_period=4e-3,
            max_period=4.0,
            device=device,
        ).to(dtype=model_dtype)

        # Fuse action and time via (2W → W → W) MLP with SiLU.
        action_emb = self.action_in_proj(noisy_actions)  # (B, H, W)
        action_time_emb = torch.cat([action_emb, time_emb[:, None, :].expand_as(action_emb)], dim=2)
        action_time_emb = self.action_time_mlp_in(action_time_emb)
        action_time_emb = F.silu(action_time_emb)
        action_time_emb = self.action_time_mlp_out(action_time_emb)

        embs = torch.cat([state_emb, action_time_emb], dim=1)
        pad_masks = torch.ones(bsize, embs.shape[1], dtype=torch.bool, device=device)
        att_masks = torch.tensor(
            [1] + [1] + [0] * (self.action_horizon - 1),
            dtype=embs.dtype,
            device=device,
        )[None, :].expand(bsize, -1)
        return embs, pad_masks, att_masks

    # ── Denoising step ───────────────────────────────────────────────

    def denoise_step(
        self,
        state: torch.Tensor,
        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``.

        Uses the prefix KV cache from ``sample_actions`` and only runs the
        action expert. Ref: openpi PI0Pytorch.denoise_step.
        """
        suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix(state, x_t, timestep)

        # Build the full (B, suffix_len, prefix_len + suffix_len) boolean mask:
        #   * suffix queries can see every *valid* prefix key (padded camera
        #     slots and padded language tokens are blocked).
        #   * within the suffix, the state token is a causal boundary; action
        #     tokens attend to each other bidirectionally.
        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,
        )

        # Drop the state token, keep only the action tokens, project down.
        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,
        state: torch.Tensor,
        noise: torch.Tensor | None = None,
        num_steps: int | 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 matches both OpenPI and LeRobot.
        Ref: openpi PI0Pytorch.sample_actions
        """
        if num_steps is None:
            num_steps = self.num_inference_steps

        bsize = state.shape[0]
        device = state.device
        if noise is None:
            noise = torch.randn(
                bsize,
                self.action_horizon,
                self.action_dim,
                dtype=torch.float32,
                device=device,
            )

        # 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(
                state=state,
                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]]):
        """Load weights from an ``lerobot/pi0_base`` safetensors checkpoint.

        Checkpoint keys (after stripping the leading ``model.`` prefix) look
        like ``paligemma_with_expert.paligemma.<sub>.*``. Modern transformers
        nests those sub-modules one level deeper (``paligemma.model.<sub>``),
        and PaliGemma ties ``lm_head.weight`` with ``embed_tokens.weight`` at
        ``post_init`` time. Two remap rules make the load lossless:

          1. ``paligemma.{vision_tower,multi_modal_projector,language_model}.*``
             → ``paligemma.model.{...}.*``
          2. ``paligemma.lm_head.weight``
             → ``paligemma.model.language_model.embed_tokens.weight``
             (the checkpoint does not store ``embed_tokens.weight`` on its
             own — only the tied ``lm_head.weight`` copy — and
             ``PaliGemmaForConditionalGeneration`` does not register
             ``lm_head.weight`` as a Parameter when tied, so without this
             rule the language embedding would silently remain at its
             random init.)

        On any remaining mismatch the loader logs a warning listing a sample
        of the unmatched checkpoint keys and model params; downstream parity
        tests surface these as a hard failure.
        """
        params_dict = dict(self.named_parameters())
        buffers_dict = dict(self.named_buffers())

        _PALIGEMMA_SUBMODULES = (
            "vision_tower",
            "multi_modal_projector",
            "language_model",
        )

        def _remap(name: str) -> str:
            # Strip the leading "model." that LeRobot's PI0Policy wrapper adds.
            if name.startswith("model."):
                name = name[len("model.") :]

            # Historical MLP alias used by some early π0 checkpoints.
            if name.startswith("time_mlp_in."):
                name = "action_time_mlp_in." + name[len("time_mlp_in.") :]
            elif name.startswith("time_mlp_out."):
                name = "action_time_mlp_out." + name[len("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) :]

            # Tied lm_head → embed_tokens redirect (see docstring).
            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 the SigLIP vision-tower nesting across transformers
            versions. transformers ≤5.3 wraps the encoder in an extra
            ``vision_tower.vision_model.*`` module; ≥5.4 flattens it to
            ``vision_tower.*``. The ``lerobot/pi0_base`` checkpoint uses the
            ``vision_model`` layout, so on newer transformers we must drop that
            infix (else 437 vision params silently load nothing). Adaptive:
            only rewrite when the rewritten key actually exists on the model.
            """
            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  # model already expects this exact layout
            if rest.startswith("vision_model."):
                # Checkpoint has the infix; model (≥5.4) doesn't → strip it.
                candidate = vt + rest[len("vision_model.") :]
            else:
                # Checkpoint lacks the infix; model (≤5.3) has it → insert it.
                candidate = vt + "vision_model." + rest
            return candidate if (candidate in params_dict or candidate in buffers_dict) else name

        loaded = 0
        skipped: list[str] = []
        filled_params: set = set()
        for name, loaded_weight in weights:
            mapped = _fix_vision_tower(_remap(name))
            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)

        # Reverse audit: any model param that got no checkpoint tensor at all
        # would be running with random init. ``rotary_emb.inv_freq`` and
        # friends are config-derived buffers, not stored in the checkpoint.
        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)

        if missing_params or skipped:
            parts: list[str] = []
            if skipped:
                parts.append(f"{len(skipped)} checkpoint key(s) had no home in the model (first 5: {skipped[:5]})")
            if missing_params:
                parts.append(f"{len(missing_params)} model param(s) received NO weight (first 5: {missing_params[:5]})")
            logger.warning(
                "π0 load_weights: %d tensors loaded — %s.",
                loaded,
                "; ".join(parts),
            )
        else:
            logger.info(
                "π0 load_weights: %d tensors loaded, 0 skipped, 0 missing.",
                loaded,
            )
        return filled_params