Add π0 VLA model support (#4222)

Signed-off-by: Yichen Wang <yichen.wang@bytedance.com>
Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>

docs/models/supported_models.md

@@ -53,6 +53,7 @@ th {
 | `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — |
 | `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
 | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
+| `Pi0Pipeline` | π0 (Pi-Zero) VLA | `lerobot/pi0_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi0.md) |
 | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — |
 | `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | — |
 | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — |

examples/online_serving/pi0/README.md

@@ -0,0 +1,104 @@
+# π0 (Pi-Zero) VLA — OpenPI realtime serving
+
+[π0](https://www.physicalintelligence.company/blog/pi0) is a Vision-Language-Action
+model from Physical Intelligence: multi-camera images + a language instruction +
+robot proprioceptive state → a continuous action chunk via flow-matching denoising
+(it does **not** emit text tokens). This example serves π0 over the OpenPI realtime
+websocket protocol at `/v1/realtime/robot/openpi`.
+
+## Install extras
+
+The core `pip install -e .` does not include the OpenPI client used here:
+
+- `openpi-client` (from the [openpi](https://github.com/Physical-Intelligence/openpi)
+  repo: `pip install -e packages/openpi-client`), `websockets`, `msgpack`, `msgpack-numpy`
+
+## Weights
+
+π0 uses the LeRobot `lerobot/pi0_base` checkpoint (HF, ~13 GB). Either let the server
+download it (`MODEL=lerobot/pi0_base`) or point at a local copy
+(`MODEL=/path/to/pi0_base`).
+
+## Run the server
+
+```bash
+vllm serve lerobot/pi0_base --omni --port 8000 \
+    --served-model-name pi0 \
+    --deploy-config vllm_omni/deploy/pi0.yaml \
+    --enforce-eager --disable-log-stats
+```
+
+The deploy config (`vllm_omni/deploy/pi0.yaml`) declares a single diffusion stage
+(`Pi0Pipeline`), float32, `max_num_seqs: 1`, and the `policy_server_config` handshake
+metadata (3 cameras, `joint_position`, action horizon 50, action dim 32).
+
+### Commonly adjusted `model_config` keys
+
+These live under `stages[0].model_config` in `vllm_omni/deploy/pi0.yaml`:
+
+| Key | Default | Meaning |
+|---|---|---|
+| `chunk_size` | `50` | Action-chunk length (timesteps) the model predicts per inference. |
+| `num_inference_steps` | `10` | Flow-matching Euler denoising steps. |
+| `max_action_dim` | `32` | Action dimensionality (state/action are padded to this). |
+| `max_state_dim` | `32` | Proprioceptive-state dimensionality (zero-padded to this). |
+| `image_resolution` | `[224, 224]` | Per-camera input size (square; SigLIP). |
+| `tokenizer_max_length` | `48` | Max PaliGemma prompt tokens. |
+| `max_cameras` | `3` | Camera slots the model attends to (real + `-1`-padded). |
+| `image_feature_keys` | 3 `observation.images.*` keys | Camera order the model attends to. |
+| `image_key_map` | `{}` | Map raw obs camera keys → `image_feature_keys` (empty = verbatim). |
+
+The `policy_server_config` block below them is the OpenPI handshake metadata
+advertised to the client; keep its `action_horizon` / `action_dim` /
+`image_resolution` in sync with the `model_config` values above. The e2e test
+`tests/e2e/online_serving/test_pi0_expansion.py::test_pi0_openpi_online` connects to a
+live server and asserts the advertised metadata matches these `pi0.yaml` values.
+
+## Run the client
+
+```bash
+python examples/online_serving/pi0/openpi_client.py --host 127.0.0.1 --port 8000 \
+    --prompt "pick up the red block and place it in the bin"
+```
+
+It connects, prints the server metadata, sends robot observations, and prints the
+returned `[action_horizon, action_dim] = [50, 32]` action chunks. Replace the blank
+cameras / zero state in `_make_dummy_obs` with real frames (HWC uint8) and
+proprioceptive state to drive a robot.
+
+### Observation format
+
+The client sends a flat dict per inference:
+
+```python
+{
+    "observation.images.base_0_rgb":       np.uint8[H, W, 3],
+    "observation.images.left_wrist_0_rgb": np.uint8[H, W, 3],
+    "observation.images.right_wrist_0_rgb":np.uint8[H, W, 3],
+    "state":   np.float32[state_dim],   # zero-padded to max_state_dim=32 server-side
+    "prompt":  "pick up the red block",
+    "session_id": "<uuid>",             # accepted but ignored (π0 is stateless)
+}
+```
+
+Camera keys must match the server's `image_feature_keys` (the checkpoint's
+`input_features` order). If your robot uses different camera names, set
+`model_config.image_key_map` in `pi0.yaml` to map raw obs keys → feature keys.
+
+## Correctness
+
+π0's flow-matching kernel is bit-for-bit matched to the LeRobot `PI0Policy` reference
+(`max|Δ| = 7.15e-07`, CPU/float32, fixed noise; see
+`tests/diffusion/models/pi0/test_pi0_parity.py::test_pi0_vllm_omni_vs_lerobot`) and is
+version-stable across transformers releases (the version-stability checks in
+`tests/diffusion/models/pi0/test_pi0_units.py`). An OpenPI websocket online-serving e2e
+lives in `tests/e2e/online_serving/test_pi0_expansion.py::test_pi0_openpi_online`.
+
+## Limitations
+
+- **Normalization stats**: `lerobot/pi0_base` uses identity normalization, which is
+  fully supported (state passes through, actions are returned in the model's space).
+  Per-dataset `norm_stats` declared in a checkpoint's `config.json` are honored, but
+  stats stored only in LeRobot's `policy_preprocessor.json` companion safetensors are
+  not yet loaded — a fine-tuned checkpoint that relies on those would need that bridge
+  added before its actions are in real-world units.

examples/online_serving/pi0/openpi_client.py

@@ -0,0 +1,119 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Minimal OpenPI client for the π0 (Pi-Zero) VLA policy server.
+
+Connects to a running π0 server (see README.md), sends a single robot
+observation (3 cameras + state + language instruction), and prints the returned
+continuous action chunk ``[action_horizon, action_dim]``.
+
+The server speaks the OpenPI websocket protocol:
+    connect -> server sends msgpack(policy_server_config metadata)
+    infer   -> client sends msgpack(obs), server replies msgpack(actions ndarray)
+
+Usage:
+    python openpi_client.py --host 127.0.0.1 --port 8000 \
+        --prompt "pick up the red block"
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import sys
+import uuid
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+try:
+    # Avoid shadowing the installed ``openpi_client`` package with this dir.
+    example_dir = str(Path(__file__).resolve().parent)
+    removed = sys.path and sys.path[0] == example_dir
+    if removed:
+        sys.path.pop(0)
+    try:
+        from openpi_client.websocket_client_policy import WebsocketClientPolicy
+    finally:
+        if removed:
+            sys.path.insert(0, example_dir)
+except ImportError as exc:  # pragma: no cover - runtime dependency guard
+    raise ImportError("π0 OpenPI example requires `openpi-client`.") from exc
+
+logger = logging.getLogger("pi0_openpi_client")
+
+DEFAULT_HOST = "127.0.0.1"
+DEFAULT_PORT = 8000
+DEFAULT_PATH = "/v1/realtime/robot/openpi"
+DEFAULT_PROMPT = "pick up the red block and place it in the bin"
+
+# π0 / pi0_base camera identities (must match the server's image_feature_keys).
+CAMERA_KEYS = (
+    "observation.images.base_0_rgb",
+    "observation.images.left_wrist_0_rgb",
+    "observation.images.right_wrist_0_rgb",
+)
+STATE_DIM = 32
+
+
+def _make_dummy_obs(prompt: str, session_id: str, image_size: int = 224) -> dict[str, Any]:
+    """Build a single π0 observation with blank cameras + zero state.
+
+    Replace the zeros with real camera frames (HWC uint8) and proprioceptive
+    state to drive an actual robot.
+    """
+    obs: dict[str, Any] = {cam: np.zeros((image_size, image_size, 3), dtype=np.uint8) for cam in CAMERA_KEYS}
+    obs["state"] = np.zeros(STATE_DIM, dtype=np.float32)
+    obs["prompt"] = prompt
+    obs["session_id"] = session_id
+    return obs
+
+
+def _ws_uri(host: str, port: int, path: str) -> str:
+    """Build the full websocket URI. WebsocketClientPolicy takes the whole URI as
+    ``host`` (it has no separate path arg)."""
+    if host.startswith(("ws://", "wss://")):
+        return host
+    return f"ws://{host}:{port}{path}"
+
+
+def _as_action_array(response: Any) -> np.ndarray:
+    if isinstance(response, dict) and response.get("type") == "error":
+        raise RuntimeError(f"Inference failed: {response.get('message', response)}")
+    return np.asarray(response, dtype=np.float32)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="π0 OpenPI client")
+    parser.add_argument("--host", default=DEFAULT_HOST)
+    parser.add_argument("--port", type=int, default=DEFAULT_PORT)
+    parser.add_argument("--path", default=DEFAULT_PATH)
+    parser.add_argument("--prompt", default=DEFAULT_PROMPT)
+    parser.add_argument("--num-steps", type=int, default=2, help="Inference observations to send.")
+    args = parser.parse_args()
+    logging.basicConfig(level=logging.INFO, format="%(message)s")
+
+    policy = WebsocketClientPolicy(host=_ws_uri(args.host, args.port, args.path))
+    metadata = dict(policy.get_server_metadata())
+    logger.info("Server metadata: %s", metadata)
+
+    session_id = str(uuid.uuid4())
+    for step in range(args.num_steps):
+        obs = _make_dummy_obs(args.prompt, session_id)
+        actions = _as_action_array(policy.infer(obs))
+        logger.info(
+            "[step %d] actions shape=%s mean=%.4f std=%.4f",
+            step,
+            actions.shape,
+            float(actions.mean()),
+            float(actions.std()),
+        )
+        if not np.isfinite(actions).all():
+            raise RuntimeError("Server returned non-finite actions.")
+    logger.info("Done — π0 OpenPI round-trip OK.")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

recipes/lerobot/Pi0.md

@@ -0,0 +1,108 @@
+# Pi0 (π0, Pi-Zero) VLA
+
+> Vision-Language-Action policy serving over the OpenPI realtime websocket
+
+## Summary
+
+- Vendor: Physical Intelligence (checkpoint redistributed by LeRobot)
+- Model: `lerobot/pi0_base`
+- Task: Vision-Language-Action — multi-camera images + a language instruction +
+  robot proprioceptive state → a continuous action chunk
+  `[action_horizon, action_dim] = [50, 32]` via flow-matching denoising. π0 does
+  **not** emit text tokens.
+- Mode: Online serving via the OpenPI realtime robot API
+  (`/v1/realtime/robot/openpi` websocket)
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to deploy `lerobot/pi0_base` as a robot policy server: a robot
+client connects over the OpenPI websocket and, per control step, sends an
+observation (camera frames + state + instruction) and receives an action chunk.
+
+A single pipeline class (`Pi0Pipeline`) serves the policy. π0 is stateless across
+calls (no KV reuse), so the protocol's `session_id` / `reset` are accepted but
+ignored.
+
+## Weights
+
+π0 uses the LeRobot `lerobot/pi0_base` checkpoint (HF, ~13 GB). Either let the
+server download it (`MODEL=lerobot/pi0_base`) or point at a local copy.
+
+## Hardware Support
+
+### 1x H200 141GB (Online serving)
+
+#### Environment
+
+- OS: Ubuntu 22.04+
+- Python: 3.12+
+- Driver / runtime: NVIDIA CUDA environment
+- vLLM / vLLM-Omni: use the commit you are deploying from
+- OpenPI client extras (for the example client): `openpi-client`
+  (`pip install -e packages/openpi-client` from the
+  [openpi](https://github.com/Physical-Intelligence/openpi) repo), `websockets`,
+  `msgpack`, `msgpack-numpy`
+
+#### Command
+
+```bash
+vllm serve lerobot/pi0_base \
+  --omni \
+  --host 0.0.0.0 --port 8000 \
+  --served-model-name pi0 \
+  --deploy-config vllm_omni/deploy/pi0.yaml \
+  --enforce-eager --disable-log-stats
+```
+
+The deploy config (`vllm_omni/deploy/pi0.yaml`) declares a single diffusion stage
+(`Pi0Pipeline`), `float32`, `max_num_seqs: 1`, and the `policy_server_config`
+handshake metadata (3 cameras, `joint_position`, action horizon 50, action dim
+32) that enables the OpenPI realtime serving layer.
+
+#### Verification
+
+Run the example client (sends a synthetic observation and prints the returned
+action chunk):
+
+```bash
+python examples/online_serving/pi0/openpi_client.py \
+  --host 127.0.0.1 --port 8000 \
+  --prompt "pick up the red block and place it in the bin"
+```
+
+It connects, prints the server metadata, sends robot observations, and prints the
+returned `[action_horizon, action_dim] = [50, 32]` action chunks. The observation
+is a flat dict per inference:
+
+```python
+{
+    "observation.images.base_0_rgb":        np.uint8[H, W, 3],
+    "observation.images.left_wrist_0_rgb":  np.uint8[H, W, 3],
+    "observation.images.right_wrist_0_rgb": np.uint8[H, W, 3],
+    "state":   np.float32[state_dim],   # zero-padded to max_state_dim=32 server-side
+    "prompt":  "pick up the red block",
+    "session_id": "<uuid>",             # accepted but ignored (π0 is stateless)
+}
+```
+
+Camera keys must match the server's `image_feature_keys` (the checkpoint's
+`input_features` order). If your robot uses different camera names, set
+`model_config.image_key_map` in `pi0.yaml` to map raw obs keys → feature keys.
+
+## Correctness
+
+π0's flow-matching kernel is bit-for-bit matched to the LeRobot `PI0Policy`
+reference (`max|Δ| = 7.15e-07`, CPU/float32, fixed noise; see
+`tests/diffusion/models/pi0/test_pi0_parity.py::test_pi0_vllm_omni_vs_lerobot`).
+
+## References
+
+- Model blog: <https://www.physicalintelligence.company/blog/pi0>
+- Checkpoint: <https://huggingface.co/lerobot/pi0_base>
+- Pipeline: [`vllm_omni/diffusion/models/pi0/pipeline_pi0.py`](../../vllm_omni/diffusion/models/pi0/pipeline_pi0.py)
+- Deploy config: [`vllm_omni/deploy/pi0.yaml`](../../vllm_omni/deploy/pi0.yaml)
+- Example client: [`examples/online_serving/pi0/`](../../examples/online_serving/pi0/)
+- Tests: CPU units + LeRobot parity in [`tests/diffusion/models/pi0/`](../../tests/diffusion/models/pi0/),
+  OpenPI websocket e2e in
+  [`tests/e2e/online_serving/test_pi0_expansion.py`](../../tests/e2e/online_serving/test_pi0_expansion.py)

tests/diffusion/models/pi0/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

tests/diffusion/models/pi0/test_pi0_parity.py

@@ -0,0 +1,366 @@
+#!/usr/bin/env python
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""π0 LeRobot parity (in-process): the bit-for-bit correctness oracle.
+
+Verifies that vllm-omni's ``Pi0ForActionPrediction`` produces bit-for-bit
+matching action chunks with LeRobot's ``PI0Policy`` when fed:
+  * the same weights (``lerobot/pi0_base``)
+  * the same pre-processed inputs (images, masks, tokens, state)
+  * the same initial noise tensor
+
+π0 is a flow-matching model (Euler-integrated ODE from t=1 → t=0 with a fixed
+``num_steps``), so the output is deterministic once the noise is fixed and
+``torch.allclose`` on the final action chunk is a valid oracle. This is the
+authoritative correctness oracle for the vllm-omni π0 port (max|Δ| < 1e-4).
+
+The OpenPI websocket online-serving e2e lives separately in
+``tests/e2e/online_serving/test_pi0_expansion.py``.
+
+Run in a SEPARATE ``lerobot[pi]`` venv (avoids dep conflict with the vllm-omni
+env), with the vllm-omni ``pi0`` package importable::
+
+    python -m pytest tests/diffusion/models/pi0/test_pi0_parity.py -v -s
+
+Skipped automatically when LeRobot is not installed (e.g. the vllm-omni env).
+The run is CPU/float32 with fixed defaults; the only override is
+``PI0_PARITY_MODEL_PATH`` (a local pi0_base dir in LeRobot format, to skip the
+HF download; defaults to ``lerobot/pi0_base``).
+"""
+
+from __future__ import annotations
+
+import copy
+import importlib.util
+import os
+
+import pytest
+import torch
+
+# local_model: needs real weights + a lerobot venv (transformers 5.3.0), so it
+# runs locally rather than in the ready-CI (see docs/contributing/ci/CI_5levels.md).
+# Additionally gated on lerobot being importable.
+_HAS_LEROBOT = importlib.util.find_spec("lerobot") is not None
+
+pytestmark = [pytest.mark.local_model, pytest.mark.diffusion]
+
+
+# ─── Config (fixed; matches LeRobot defaults for ``lerobot/pi0_base``) ──
+DEVICE = "cpu"
+DTYPE_STR = "float32"
+ATOL = 1e-4
+NUM_STEPS = 10
+BATCH_SIZE = 2
+ACTION_DIM = 32
+STATE_DIM = 32
+ACTION_HORIZON = 50
+MAX_TOKEN_LEN = 48
+
+# The only knob: point at a local pi0_base dir (LeRobot format) to skip the HF
+# download; defaults to the HF repo id.
+MODEL_PATH = os.environ.get("PI0_PARITY_MODEL_PATH", "lerobot/pi0_base")
+
+
+def _resolve_checkpoint_dir() -> str:
+    """Return a local dir containing the pi0_base checkpoint (download if needed)."""
+    if os.path.isdir(MODEL_PATH):
+        return MODEL_PATH
+    from huggingface_hub import snapshot_download
+
+    return snapshot_download(repo_id=MODEL_PATH, repo_type="model")
+
+
+# ─── Dummy dataset stats (identity) ──────────────────────────────────
+def _dummy_dataset_stats() -> dict:
+    return {
+        "observation.state": {
+            "mean": torch.zeros(STATE_DIM),
+            "std": torch.ones(STATE_DIM),
+            "q01": torch.zeros(STATE_DIM),
+            "q99": torch.ones(STATE_DIM),
+        },
+        "action": {
+            "mean": torch.zeros(ACTION_DIM),
+            "std": torch.ones(ACTION_DIM),
+            "q01": torch.zeros(ACTION_DIM),
+            "q99": torch.ones(ACTION_DIM),
+        },
+        "images": {
+            cam: {
+                "mean": torch.zeros(3, 224, 224),
+                "std": torch.ones(3, 224, 224),
+                "q01": torch.zeros(3, 224, 224),
+                "q99": torch.ones(3, 224, 224),
+            }
+            for cam in ("base_0_rgb", "left_wrist_0_rgb", "right_wrist_0_rgb")
+        },
+    }
+
+
+def _create_dummy_batch(batch_size: int = BATCH_SIZE, device: str = DEVICE) -> dict:
+    """Reproducible dummy inputs — identical across both implementations."""
+    g = torch.Generator(device="cpu").manual_seed(0)
+    prompt = "Pick up the red block and place it in the bin"
+    return {
+        "observation.state": torch.randn(batch_size, STATE_DIM, generator=g, dtype=torch.float32).to(device),
+        "action": torch.randn(batch_size, ACTION_HORIZON, ACTION_DIM, generator=g, dtype=torch.float32).to(device),
+        "observation.images.base_0_rgb": torch.rand(batch_size, 3, 224, 224, generator=g, dtype=torch.float32).to(
+            device
+        ),
+        "observation.images.left_wrist_0_rgb": torch.rand(batch_size, 3, 224, 224, generator=g, dtype=torch.float32).to(
+            device
+        ),
+        "observation.images.right_wrist_0_rgb": torch.rand(
+            batch_size, 3, 224, 224, generator=g, dtype=torch.float32
+        ).to(device),
+        "task": [prompt for _ in range(batch_size)],
+    }
+
+
+# ─── LeRobot instantiation ────────────────────────────────────────────
+def _instantiate_lerobot():
+    from lerobot.policies.pi0 import PI0Policy
+    from lerobot.policies.pi0.processor_pi0 import make_pi0_pre_post_processors
+
+    policy = PI0Policy.from_pretrained(MODEL_PATH, strict=True)
+    policy.to(DEVICE)
+    policy.config.device = DEVICE
+    policy.eval()
+
+    pre, post = make_pi0_pre_post_processors(config=policy.config, dataset_stats=_dummy_dataset_stats())
+    return policy, pre, post
+
+
+# ─── vllm-omni instantiation ──────────────────────────────────────────
+def _instantiate_vllm_omni():
+    """Build the vllm-omni π0 model in isolation (no pipeline, no engine)."""
+    from vllm_omni.diffusion.models.pi0 import Pi0Config, Pi0ForActionPrediction
+
+    cfg = Pi0Config(
+        max_action_dim=ACTION_DIM,
+        max_state_dim=STATE_DIM,
+        chunk_size=ACTION_HORIZON,
+        num_inference_steps=NUM_STEPS,
+        dtype=DTYPE_STR,
+    )
+    model = Pi0ForActionPrediction(cfg)
+    model.to(DEVICE).eval()
+    _load_lerobot_weights(model)
+    return model
+
+
+def _load_lerobot_weights(model):
+    """Feed the ``lerobot/pi0_base`` safetensors into the vllm-omni model.
+
+    ``Pi0ForActionPrediction.load_weights`` handles the leading ``model.``
+    prefix and the flat→nested / lm_head→embed_tokens remaps, so we can pass
+    the raw checkpoint dict.
+    """
+    import safetensors.torch
+
+    cache_dir = _resolve_checkpoint_dir()
+    path = os.path.join(cache_dir, "model.safetensors")
+    state = safetensors.torch.load_file(path)
+    model.load_weights(list(state.items()))
+
+
+# ─── Helpers to extract LeRobot's pre-processed inputs ────────────────
+def _extract_lerobot_model_inputs(lerobot_policy, processed_batch):
+    """Mimic what ``PI0Policy.predict_action_chunk`` feeds into
+    ``self.model.sample_actions``. We use these *exact* tensors for vllm-omni
+    so any divergence must come from the core model, not preprocessing.
+    """
+    images, img_masks = lerobot_policy._preprocess_images(processed_batch)
+    from lerobot.utils.constants import (
+        OBS_LANGUAGE_ATTENTION_MASK,
+        OBS_LANGUAGE_TOKENS,
+    )
+
+    lang_tokens = processed_batch[OBS_LANGUAGE_TOKENS]
+    lang_masks = processed_batch[OBS_LANGUAGE_ATTENTION_MASK]
+    state = lerobot_policy.prepare_state(processed_batch)
+    return images, img_masks, lang_tokens, lang_masks, state
+
+
+# ─── Shared fixed-noise sampler ───────────────────────────────────────
+def _make_fixed_noise(batch_size: int, device: str) -> torch.Tensor:
+    g = torch.Generator(device="cpu").manual_seed(42)
+    return torch.randn(batch_size, ACTION_HORIZON, ACTION_DIM, generator=g, dtype=torch.float32).to(device)
+
+
+# ─── Main test ────────────────────────────────────────────────────────
+@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).")
+def test_pi0_vllm_omni_vs_lerobot():
+    print("\n[parity] Instantiating LeRobot…")
+    lerobot_policy, lerobot_pre, _ = _instantiate_lerobot()
+
+    print("[parity] Instantiating vllm-omni…")
+    omni_model = _instantiate_vllm_omni()
+
+    print("[parity] Preparing shared inputs…")
+    raw_batch = _create_dummy_batch()
+    processed_batch = lerobot_pre(copy.deepcopy(raw_batch))
+    images, img_masks, lang_tokens, lang_masks, state = _extract_lerobot_model_inputs(lerobot_policy, processed_batch)
+    noise = _make_fixed_noise(raw_batch["observation.state"].shape[0], DEVICE)
+
+    print(f"[parity] state.shape={state.shape}  lang_tokens.shape={lang_tokens.shape}")
+    print(f"[parity] images[0].shape={images[0].shape} (num_cams={len(images)})")
+    print(f"[parity] noise.shape={noise.shape}  dtype={noise.dtype}")
+
+    # ── LeRobot forward ──
+    print("[parity] Running LeRobot sample_actions…")
+    with torch.no_grad():
+        lerobot_actions = lerobot_policy.model.sample_actions(
+            images,
+            img_masks,
+            lang_tokens,
+            lang_masks,
+            state,
+            noise=noise,
+            num_steps=NUM_STEPS,
+        )
+    print(
+        f"[parity] LeRobot actions: shape={lerobot_actions.shape} "
+        f"mean={lerobot_actions.mean().item():.6f} std={lerobot_actions.std().item():.6f}"
+    )
+
+    # ── vllm-omni forward ──
+    print("[parity] Running vllm-omni sample_actions…")
+    with torch.no_grad():
+        omni_actions = omni_model.sample_actions(
+            images=images,
+            image_masks=img_masks,
+            lang_tokens=lang_tokens,
+            lang_masks=lang_masks,
+            state=state,
+            noise=noise,
+            num_steps=NUM_STEPS,
+        )
+    print(
+        f"[parity] vllm-omni actions: shape={omni_actions.shape} "
+        f"mean={omni_actions.mean().item():.6f} std={omni_actions.std().item():.6f}"
+    )
+
+    # ── Compare ──
+    diff = (lerobot_actions.float() - omni_actions.float()).abs()
+    max_diff = diff.max().item()
+    mean_diff = diff.mean().item()
+    print(f"[parity] |Δ| max={max_diff:.2e}  mean={mean_diff:.2e}  atol={ATOL:.1e}")
+    close = torch.allclose(lerobot_actions.float(), omni_actions.float(), atol=ATOL)
+    print(f"[parity] torch.allclose(atol={ATOL}): {close}")
+
+    if not close:
+        print("\n[parity] ⚠️  Outputs diverge — running per-stage diagnostics…")
+        _diagnose_divergence(
+            lerobot_policy.model,
+            omni_model,
+            images,
+            img_masks,
+            lang_tokens,
+            lang_masks,
+            state,
+            noise,
+        )
+
+    assert close, (
+        f"vllm-omni vs LeRobot actions differ beyond atol={ATOL}. max_diff={max_diff:.2e}  mean_diff={mean_diff:.2e}"
+    )
+
+
+# ─── Per-stage divergence diagnostics ─────────────────────────────────
+@torch.no_grad()
+def _diagnose_divergence(
+    lerobot_flow_model,
+    omni_model,
+    images,
+    img_masks,
+    lang_tokens,
+    lang_masks,
+    state,
+    noise,
+):
+    """Localize a numerical mismatch to a specific pipeline stage:
+    1. Prefix embeddings (SigLIP / embed_tokens / projector) — image vs lang.
+    2. Prefix KV cache layer 0 — PaliGemma LM attention.
+    3. A single denoise_step velocity at t=1.0 — action expert forward.
+    """
+    from vllm_omni.diffusion.models.pi0.modeling_pi0 import (
+        make_att_2d_masks,
+        prepare_attention_masks_4d,
+    )
+
+    # ── Stage 1: prefix embeddings ──
+    lr_prefix_embs, lr_prefix_pad, lr_prefix_att = lerobot_flow_model.embed_prefix(
+        images, img_masks, lang_tokens, lang_masks
+    )
+    sg_prefix_embs, sg_prefix_pad, sg_prefix_att = omni_model.embed_prefix(images, img_masks, lang_tokens, lang_masks)
+    total_diff = (lr_prefix_embs.float() - sg_prefix_embs.float()).abs().max().item()
+    print(f"[diag] prefix_embs max |Δ| = {total_diff:.2e}   (shape={tuple(sg_prefix_embs.shape)})")
+    print(f"[diag] prefix_pad_masks equal: {torch.equal(lr_prefix_pad, sg_prefix_pad)}")
+    print(f"[diag] prefix_att_masks equal: {torch.equal(lr_prefix_att.bool(), sg_prefix_att.bool())}")
+
+    num_cams = len(images)
+    img_len = 256 * num_cams
+    lang_len = lr_prefix_embs.shape[1] - img_len
+    img_diff = (lr_prefix_embs[:, :img_len].float() - sg_prefix_embs[:, :img_len].float()).abs().max().item()
+    lang_diff = (lr_prefix_embs[:, img_len:].float() - sg_prefix_embs[:, img_len:].float()).abs().max().item()
+    print(f"[diag]   image slice [:{img_len}] max|Δ| = {img_diff:.2e}   (num_cams={num_cams})")
+    print(f"[diag]   lang slice  [{img_len}:] max|Δ| = {lang_diff:.2e}   (len={lang_len})")
+
+    def _stats(name, t):
+        t = t.float()
+        print(
+            f"[diag] {name} prefix_embs: mean={t.mean().item():+.4f} "
+            f"std={t.std().item():.4f} min={t.min().item():+.2f} "
+            f"max={t.max().item():+.2f}"
+        )
+
+    _stats("LeRobot ", lr_prefix_embs)
+    _stats("vllm-omni", sg_prefix_embs)
+
+    # ── Stage 2: prefix KV cache ──
+    prefix_att_2d = make_att_2d_masks(sg_prefix_pad, sg_prefix_att)
+    prefix_pos = torch.cumsum(sg_prefix_pad, dim=1) - 1
+    prefix_att_4d = prepare_attention_masks_4d(prefix_att_2d)
+
+    lerobot_flow_model.paligemma_with_expert.paligemma.language_model.config._attn_implementation = "eager"
+    _, lr_kv = lerobot_flow_model.paligemma_with_expert.forward(
+        attention_mask=prefix_att_4d,
+        position_ids=prefix_pos,
+        past_key_values=None,
+        inputs_embeds=[lr_prefix_embs, None],
+        use_cache=True,
+    )
+    _, sg_kv = omni_model.paligemma_with_expert.forward(
+        attention_mask=prefix_att_4d,
+        position_ids=prefix_pos,
+        past_key_values=None,
+        inputs_embeds=[sg_prefix_embs, None],
+        use_cache=True,
+    )
+
+    def _layer0_kv(cache):
+        if isinstance(cache, list):
+            return cache[0]
+        layer0 = cache.layers[0]
+        return layer0.keys, layer0.values
+
+    try:
+        lr_k, lr_v = _layer0_kv(lr_kv)
+        sg_k, sg_v = _layer0_kv(sg_kv)
+        dk = (lr_k.float() - sg_k.float()).abs().max().item()
+        dv = (lr_v.float() - sg_v.float()).abs().max().item()
+        print(f"[diag] prefix KV layer0  K max|Δ|={dk:.2e}  V max|Δ|={dv:.2e}")
+    except Exception as e:  # noqa: BLE001
+        print(f"[diag] could not extract prefix KV for comparison: {e}")
+
+    # ── Stage 3: a single denoise_step at t=1.0 ──
+    bsize = state.shape[0]
+    t = torch.ones(bsize, dtype=torch.float32, device=state.device)
+    lr_vt = lerobot_flow_model.denoise_step(state, sg_prefix_pad, lr_kv, noise, t)
+    sg_vt = omni_model.denoise_step(state, sg_prefix_pad, sg_kv, noise, t)
+    print(f"[diag] denoise_step(t=1) v_t max|Δ| = {(lr_vt.float() - sg_vt.float()).abs().max().item():.2e}")
+
+
+if __name__ == "__main__":
+    test_pi0_vllm_omni_vs_lerobot()

tests/diffusion/models/pi0/test_pi0_units.py

@@ -0,0 +1,585 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU unit tests for the π0 (Pi-Zero) VLA model.
+
+Deliberately lightweight: pure utility code (config plumbing, mask building,
+sinusoidal embeddings, normalization, image preprocessing, weight-load remap)
+that runs on CPU with no weights and no lerobot. The full LeRobot parity check
+lives in ``tests/diffusion/models/pi0/test_pi0_parity.py`` and requires a model download.
+
+    pytest tests/diffusion/models/pi0/test_pi0_units.py -v
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.pi0.config import Pi0Config
+from vllm_omni.diffusion.models.pi0.modeling_pi0 import (
+    OPENPI_ATTENTION_MASK_VALUE,
+    Pi0ForActionPrediction,
+    _apply_norm,
+    _build_norm_buffers,
+    create_sinusoidal_pos_embedding,
+    get_gemma_config,
+    make_att_2d_masks,
+    prepare_attention_masks_4d,
+)
+from vllm_omni.diffusion.models.pi0.processor_pi0 import (
+    Pi0ImageProcessor,
+    build_model_inputs,
+    pil_image_to_tensor,
+    resize_with_pad,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+# ----------------------------------------------------------------------------
+# Pi0Config dataclass + checkpoint resolver
+# ----------------------------------------------------------------------------
+_LEROBOT_CFG = {
+    "type": "pi0",
+    "paligemma_variant": "gemma_2b",
+    "action_expert_variant": "gemma_300m",
+    "chunk_size": 50,
+    "max_action_dim": 32,
+    "max_state_dim": 32,
+    "num_inference_steps": 10,
+    "image_resolution": [224, 224],
+    "tokenizer_max_length": 48,
+    "dtype": "float32",
+    # Training-only keys that must be filtered out of the dataclass.
+    "optimizer_lr": 2.5e-5,
+    "scheduler_warmup_steps": 1000,
+    "input_features": {
+        "observation.images.base_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]},
+        "observation.images.left_wrist_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]},
+        "observation.images.right_wrist_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]},
+        "observation.state": {"type": "STATE", "shape": [32]},
+    },
+}
+
+_EXPECTED_CAMERA_ORDER = [
+    "observation.images.base_0_rgb",
+    "observation.images.left_wrist_0_rgb",
+    "observation.images.right_wrist_0_rgb",
+]
+
+
+def test_config_camera_order_from_input_features():
+    c = Pi0Config.from_model_config(_LEROBOT_CFG)
+    assert c.image_feature_keys == _EXPECTED_CAMERA_ORDER
+
+
+def test_config_filters_training_only_keys():
+    c = Pi0Config.from_model_config(_LEROBOT_CFG)
+    assert not hasattr(c, "optimizer_lr")
+    assert not hasattr(c, "scheduler_warmup_steps")
+
+
+def test_config_coerces_image_resolution_to_tuple():
+    c = Pi0Config.from_model_config(_LEROBOT_CFG)
+    assert c.image_resolution == (224, 224)
+    assert isinstance(c.image_resolution, tuple)
+
+
+def test_config_non_square_resolution_raises():
+    with pytest.raises(ValueError):
+        Pi0Config(image_resolution=(224, 256))
+
+
+@pytest.mark.parametrize("bad", [(224,), (224, 224, 3), 224])
+def test_config_malformed_resolution_raises(bad):
+    with pytest.raises(ValueError):
+        Pi0Config(image_resolution=bad)
+
+
+def test_config_defaults_and_empty():
+    c = Pi0Config.from_model_config(None)
+    assert c.chunk_size == 50
+    assert c.max_action_dim == 32
+    assert c.max_state_dim == 32
+    assert c.num_inference_steps == 10
+    assert c.image_resolution == (224, 224)
+    assert c.tokenizer_max_length == 48
+    assert c.max_cameras == 3
+    assert c.norm_stats is None
+    assert c.image_key_map == {}
+
+
+def test_config_deploy_yaml_override():
+    """Deploy-yaml model_config is authoritative (e.g. force dtype/variant)."""
+    c = Pi0Config.from_model_config({**_LEROBOT_CFG, "dtype": "bfloat16", "num_inference_steps": 4})
+    assert c.dtype == "bfloat16"
+    assert c.num_inference_steps == 4
+
+
+# ----------------------------------------------------------------------------
+# Kernel utilities: sinusoidal timestep embedding
+# ----------------------------------------------------------------------------
+def test_sinusoidal_embedding_shape():
+    emb = create_sinusoidal_pos_embedding(torch.tensor([0.0, 0.5, 1.0]), dimension=64)
+    assert emb.shape == (3, 64)
+
+
+def test_sinusoidal_embedding_values_at_zero():
+    emb = create_sinusoidal_pos_embedding(torch.tensor([0.0]), dimension=4)
+    assert torch.allclose(emb[0, :2].float(), torch.zeros(2), atol=1e-6)
+    assert torch.allclose(emb[0, 2:].float(), torch.ones(2), atol=1e-6)
+
+
+def test_sinusoidal_embedding_odd_dim_raises():
+    with pytest.raises(ValueError):
+        create_sinusoidal_pos_embedding(torch.tensor([1.0]), dimension=3)
+
+
+def test_sinusoidal_embedding_non_1d_raises():
+    with pytest.raises(ValueError):
+        create_sinusoidal_pos_embedding(torch.tensor([[0.0, 0.5]]), dimension=4)
+
+
+def test_sinusoidal_embedding_distinct_timesteps_differ():
+    emb = create_sinusoidal_pos_embedding(torch.tensor([0.1, 0.2, 0.5, 0.9]), dimension=16)
+    for i in range(emb.shape[0]):
+        for j in range(i + 1, emb.shape[0]):
+            assert not torch.allclose(emb[i], emb[j])
+
+
+# ----------------------------------------------------------------------------
+# Kernel utilities: attention mask construction
+# ----------------------------------------------------------------------------
+def test_make_att_2d_masks_bidirectional():
+    result = make_att_2d_masks(torch.ones(1, 4, dtype=torch.bool), torch.zeros(1, 4, dtype=torch.long))
+    assert result.all()
+
+
+def test_make_att_2d_masks_causal():
+    result = make_att_2d_masks(torch.ones(1, 4, dtype=torch.bool), torch.ones(1, 4, dtype=torch.long))
+    expected = torch.tril(torch.ones(4, 4, dtype=torch.bool))
+    assert torch.equal(result[0], expected)
+
+
+def test_make_att_2d_masks_prefix_lm():
+    result = make_att_2d_masks(torch.ones(1, 5, dtype=torch.bool), torch.tensor([[0, 0, 0, 1, 1]], dtype=torch.long))
+    assert result[0, 0, 0]
+    assert not result[0, 0, 3]
+    assert result[0, 3, 0]
+    assert not result[0, 3, 4]
+
+
+def test_make_att_2d_masks_respects_padding():
+    pad = torch.tensor([[1, 1, 1, 0]], dtype=torch.bool)  # last token is padding
+    att = torch.zeros(1, 4, dtype=torch.long)  # full bidirectional
+    result = make_att_2d_masks(pad, att)
+    assert not result[0, :, 3].any(), "Nothing should attend to padded key"
+    assert not result[0, 3, :].any(), "Padded query should attend to nothing"
+
+
+def test_prepare_attention_masks_4d():
+    mask_2d = torch.tensor([[[True, False], [True, True]]])
+    mask_4d = prepare_attention_masks_4d(mask_2d)
+    assert mask_4d.shape == (1, 1, 2, 2)
+    assert mask_4d[0, 0, 0, 0].item() == 0.0
+    assert mask_4d[0, 0, 1, 1].item() == 0.0
+    expected = torch.tensor(OPENPI_ATTENTION_MASK_VALUE, dtype=mask_4d.dtype)
+    assert torch.equal(mask_4d[0, 0, 0, 1], expected)
+
+
+def test_attention_mask_value_finite_and_negative():
+    assert math.isfinite(OPENPI_ATTENTION_MASK_VALUE)
+    assert OPENPI_ATTENTION_MASK_VALUE < -1e30
+
+
+# ----------------------------------------------------------------------------
+# Normalization round-trips
+# ----------------------------------------------------------------------------
+def test_build_norm_buffers_missing():
+    assert _build_norm_buffers(None, "state") is None
+    assert _build_norm_buffers({}, "state") is None
+    assert _build_norm_buffers({"state": None}, "state") is None
+
+
+def test_build_norm_buffers_mean_std():
+    stats = {"state": {"mode": "mean_std", "mean": [1.0, 2.0], "std": [0.5, 0.5]}}
+    out = _build_norm_buffers(stats, "state")
+    assert out["mode"] == "mean_std"
+    assert torch.equal(out["mean"], torch.tensor([1.0, 2.0]))
+    assert torch.equal(out["std"], torch.tensor([0.5, 0.5]))
+
+
+def test_build_norm_buffers_min_max():
+    stats = {"action": {"mode": "min_max", "min": [-1.0, 0.0], "max": [1.0, 2.0]}}
+    out = _build_norm_buffers(stats, "action")
+    assert out["mode"] == "min_max"
+    assert torch.equal(out["min"], torch.tensor([-1.0, 0.0]))
+    assert torch.equal(out["max"], torch.tensor([1.0, 2.0]))
+
+
+def test_apply_norm_mean_std_forward_inverse():
+    stats = _build_norm_buffers({"s": {"mode": "mean_std", "mean": [0.5, -0.2], "std": [2.0, 0.1]}}, "s")
+    x = torch.tensor([[1.0, 0.3]])
+    z = _apply_norm(x, stats, inverse=False)
+    assert torch.allclose(z, torch.tensor([[0.25, 5.0]]), atol=1e-5)
+    x_back = _apply_norm(z, stats, inverse=True)
+    assert torch.allclose(x_back, x, atol=1e-5)
+
+
+def test_apply_norm_min_max_forward_inverse():
+    stats = _build_norm_buffers({"s": {"mode": "min_max", "min": [0.0, -1.0], "max": [1.0, 1.0]}}, "s")
+    x = torch.tensor([[0.25, 0.5]])
+    z = _apply_norm(x, stats, inverse=False)
+    assert torch.allclose(z, torch.tensor([[-0.5, 0.5]]), atol=1e-5)
+    x_back = _apply_norm(z, stats, inverse=True)
+    assert torch.allclose(x_back, x, atol=1e-5)
+
+
+def test_apply_norm_preserves_padded_tail():
+    stats = _build_norm_buffers({"s": {"mode": "mean_std", "mean": [0.0, 0.0], "std": [1.0, 1.0]}}, "s")
+    x = torch.tensor([[1.0, 2.0, 99.0, -99.0]])
+    z = _apply_norm(x, stats, inverse=False)
+    assert torch.allclose(z[:, :2], x[:, :2])
+    assert torch.allclose(z[:, 2:], x[:, 2:])
+
+
+def test_apply_norm_noop_when_stats_none():
+    x = torch.randn(2, 5)
+    assert torch.equal(_apply_norm(x, None, inverse=False), x)
+    assert torch.equal(_apply_norm(x, None, inverse=True), x)
+
+
+# ----------------------------------------------------------------------------
+# Gemma variant configs
+# ----------------------------------------------------------------------------
+def test_gemma_2b_dims():
+    c = get_gemma_config("gemma_2b")
+    assert c.width == 2048
+    assert c.depth == 18
+
+
+def test_gemma_300m_dims():
+    c = get_gemma_config("gemma_300m")
+    assert c.width == 1024
+    assert c.depth == 18
+
+
+def test_gemma_unknown_variant_raises():
+    with pytest.raises(ValueError):
+        get_gemma_config("gemma_7b")
+
+
+# ----------------------------------------------------------------------------
+# load_weights remap rules (tiny model; needs torch + transformers)
+# ----------------------------------------------------------------------------
+def _tiny_pi0_model():
+    """Construct a tiny π0 model for remap tests (full Gemma depth but small
+    action/state dims — the PaliGemma backbone is always full-size)."""
+    cfg = Pi0Config(max_action_dim=8, max_state_dim=8, chunk_size=2)
+    return Pi0ForActionPrediction(cfg)
+
+
+@pytest.fixture(autouse=True)
+def _silence_pi0_loader(caplog):
+    # The remap tests feed a single synthetic weight; the loader's
+    # "missing params" warning is expected noise here.
+    logging.getLogger("vllm_omni.diffusion.models.pi0.modeling_pi0").setLevel(logging.ERROR)
+    yield
+
+
+@pytest.mark.slow
+def test_remap_paligemma_submodules_to_nested_layout():
+    model = _tiny_pi0_model()
+    params = dict(model.named_parameters())
+    target = next(k for k in params if k.startswith("paligemma_with_expert.paligemma.model.vision_tower."))
+    flat_key = "model." + target.replace(
+        "paligemma_with_expert.paligemma.model.vision_tower.",
+        "paligemma_with_expert.paligemma.vision_tower.",
+    )
+    payload = torch.full(params[target].shape, 0.12345)
+    before = params[target].detach().clone()
+    model.load_weights([(flat_key, payload)])
+    after = dict(model.named_parameters())[target]
+    assert not torch.equal(before, after), "flat→nested remap broken"
+    assert torch.allclose(after, payload)
+
+
+@pytest.mark.slow
+def test_lm_head_remap_to_embed_tokens():
+    model = _tiny_pi0_model()
+    params = dict(model.named_parameters())
+    target = "paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"
+    assert target in params
+    payload = torch.full(params[target].shape, 0.5)
+    before = params[target].detach().clone()
+    model.load_weights([("model.paligemma_with_expert.paligemma.lm_head.weight", payload)])
+    after = dict(model.named_parameters())[target]
+    assert not torch.equal(before, after), "lm_head→embed_tokens remap broken"
+    assert torch.allclose(after, payload)
+
+
+@pytest.mark.slow
+def test_vision_tower_version_robust_remap():
+    """transformers >=5.4 flattens SigLIP from vision_tower.vision_model.* to
+    vision_tower.*; the lerobot/pi0_base checkpoint uses the vision_model infix.
+    load_weights must adapt so vision weights aren't silently dropped (which
+    would leave the vision encoder at random init)."""
+    model = _tiny_pi0_model()
+    params = dict(model.named_parameters())
+    # A representative vision-tower param as it exists on the model.
+    target = next(
+        k
+        for k in params
+        if k.startswith("paligemma_with_expert.paligemma.model.vision_tower.") and "patch_embedding.weight" in k
+    )
+    vt = "paligemma_with_expert.paligemma.model.vision_tower."
+    rest = target[len(vt) :]
+    # Build the checkpoint key in whichever layout differs from the model's.
+    if rest.startswith("vision_model."):
+        ckpt_key = vt + rest[len("vision_model.") :]
+    else:
+        ckpt_key = vt + "vision_model." + rest
+    payload = torch.full(params[target].shape, 0.073)
+    before = params[target].detach().clone()
+    model.load_weights([("model." + ckpt_key, payload)])
+    after = dict(model.named_parameters())[target]
+    assert not torch.equal(before, after), "vision_tower version-robust remap broken"
+    assert torch.allclose(after, payload)
+
+
+# ----------------------------------------------------------------------------
+# Image processor
+# ----------------------------------------------------------------------------
+def test_resize_with_pad_square():
+    out = resize_with_pad(torch.ones(1, 3, 224, 224), 224, 224)
+    assert out.shape == (1, 3, 224, 224)
+
+
+def test_resize_with_pad_wide():
+    """Wide source (H < W): padding applied to top/bottom (rows)."""
+    out = resize_with_pad(torch.zeros(1, 3, 100, 200), 224, 224)
+    assert out.shape == (1, 3, 224, 224)
+    assert out[0, 0, 0, 112].item() == -1.0  # top row is padded
+
+
+def test_resize_with_pad_tall():
+    """Tall source (H > W): padding applied to left/right (cols)."""
+    out = resize_with_pad(torch.zeros(1, 3, 400, 100), 224, 224)
+    assert out.shape == (1, 3, 224, 224)
+    assert out[0, 0, 112, 0].item() == -1.0  # left-most column padded
+
+
+def test_resize_with_pad_rejects_non_4d():
+    with pytest.raises(ValueError):
+        resize_with_pad(torch.zeros(3, 100, 200), 224, 224)
+
+
+def test_pil_to_tensor():
+    import numpy as np
+    from PIL import Image
+
+    img = Image.fromarray(np.full((10, 10, 3), 128, dtype=np.uint8))
+    t = pil_image_to_tensor(img)
+    assert t.shape == (1, 3, 10, 10)
+    expected = 128.0 / 255.0 * 2.0 - 1.0
+    assert abs(t[0, 0, 5, 5].item() - expected) < 1e-4
+
+
+def test_pil_to_tensor_rgba_to_rgb():
+    import numpy as np
+    from PIL import Image
+
+    rgba = np.full((8, 8, 4), 200, dtype=np.uint8)
+    img = Image.fromarray(rgba, mode="RGBA")
+    t = pil_image_to_tensor(img)
+    assert t.shape == (1, 3, 8, 8)
+
+
+def test_make_empty_image():
+    proc = Pi0ImageProcessor(image_size=224)
+    empty = proc.make_empty_image()
+    assert empty.shape == (1, 3, 224, 224)
+    assert (empty == -1.0).all()
+
+
+def test_preprocess_single_resizes_ndarray():
+    import numpy as np
+
+    proc = Pi0ImageProcessor(image_size=224)
+    arr = np.full((120, 120, 3), 100, dtype=np.uint8)
+    t = proc.preprocess_single(arr)
+    assert t.shape == (1, 3, 224, 224)
+
+
+# ----------------------------------------------------------------------------
+# build_model_inputs: camera ordering, padding, masks, state
+# ----------------------------------------------------------------------------
+class _FakeTokenizer:
+    """Minimal stand-in for the PaliGemma tokenizer (no download)."""
+
+    def __call__(
+        self, text, padding=None, max_length=None, truncation=None, add_special_tokens=None, return_tensors=None
+    ):
+        ids = [2] + [ord(c) % 100 for c in text][: max_length - 1]
+        ids = ids[:max_length]
+        attn = [1] * len(ids)
+        while len(ids) < max_length:
+            ids.append(0)
+            attn.append(0)
+        return {"input_ids": ids, "attention_mask": attn}
+
+
+def test_build_model_inputs_camera_order_and_padding():
+    import numpy as np
+
+    config = Pi0Config.from_model_config(_LEROBOT_CFG)  # 3-camera config
+    device = torch.device("cpu")
+    # Provide only 2 of the 3 declared cameras.
+    robot_obs = {
+        "images": {
+            "observation.images.base_0_rgb": np.full((224, 224, 3), 50, dtype=np.uint8),
+            "observation.images.left_wrist_0_rgb": np.full((224, 224, 3), 80, dtype=np.uint8),
+        },
+        "state": [0.1] * 14,  # raw state shorter than max_state_dim=32
+        "prompt": "pick up the block",
+    }
+    images, image_masks, lang_tokens, lang_masks, state = build_model_inputs(
+        robot_obs, config, _FakeTokenizer(), device
+    )
+    assert len(images) == 3
+    assert all(img.shape == (1, 3, 224, 224) for img in images)
+    masks = [bool(m.item()) for m in image_masks]
+    assert masks == [True, True, False]
+    # The padded (3rd) camera must be all -1.
+    assert (images[2] == -1.0).all()
+    assert lang_tokens.shape == (1, 48)
+    assert lang_masks.shape == (1, 48)
+    assert state.shape == (1, 32)
+    # State padded with zeros beyond the 14 real entries.
+    assert torch.allclose(state[0, 14:], torch.zeros(18))
+
+
+def test_build_model_inputs_state_truncated():
+    config = Pi0Config.from_model_config(_LEROBOT_CFG)
+    robot_obs = {"images": {}, "state": [0.5] * 40, "prompt": "x"}
+    _, _, _, _, state = build_model_inputs(robot_obs, config, _FakeTokenizer(), torch.device("cpu"))
+    assert state.shape == (1, 32)
+
+
+def test_build_model_inputs_image_key_map():
+    import numpy as np
+
+    cfg = Pi0Config.from_model_config({**_LEROBOT_CFG, "image_key_map": {"cam_high": "observation.images.base_0_rgb"}})
+    robot_obs = {
+        "images": {"cam_high": np.full((224, 224, 3), 60, dtype=np.uint8)},
+        "prompt": "go",
+    }
+    _, image_masks, _, _, _ = build_model_inputs(robot_obs, cfg, _FakeTokenizer(), torch.device("cpu"))
+    masks = [bool(m.item()) for m in image_masks]
+    assert masks == [True, False, False]  # only base_0_rgb (mapped) is present
+
+
+# ----------------------------------------------------------------------------
+# Version-stability / self-consistency (gated on a real pi0_base checkpoint)
+#
+# The LeRobot parity gate (test_pi0_parity.py) only validates correctness at
+# transformers 5.3.0 (where LeRobot loads pi0_base cleanly). Production vllm-omni
+# ships transformers 5.6.2, and LeRobot is not a usable reference there (it
+# incorrectly loads the SigLIP vision tower at >=5.4). So we pin the kernel's deterministic
+# output to golden values established at 5.3.0 and re-check them here: running on
+# the production transformers proves the port is version-stable and guards the
+# vision_tower / embed_scale version-robust paths against regressions.
+#
+#   PI0_SELFCONSIST_CKPT=/path/to/pi0_base \
+#       pytest tests/diffusion/models/pi0/test_pi0_units.py -k version_stable -v -s
+# ----------------------------------------------------------------------------
+import os  # noqa: E402
+
+_SELFCONSIST_CKPT = os.environ.get("PI0_SELFCONSIST_CKPT")
+
+# Golden values established at transformers 5.3.0 (where parity = 7.15e-07).
+_GOLDEN_SUM = -67.808113
+_GOLDEN_STD = 0.252789
+_GOLDEN_MEAN = -0.021190
+_GOLDEN_EMBED_IMAGE_STD = 2.17603
+
+
+def _selfconsist_inputs(device):
+    # NB: draw order is load-bearing — the golden values were established with
+    # state drawn first, then the 3 images, then lang (shared RNG, seed 0).
+    g = torch.Generator(device="cpu").manual_seed(0)
+    b = 2
+    state = torch.randn(b, 32, generator=g)
+    images = [torch.rand(b, 3, 224, 224, generator=g) for _ in range(3)]
+    masks = [torch.ones(b, dtype=torch.bool) for _ in range(3)]
+    lang = torch.randint(0, 1000, (b, 48), generator=g)
+    lmask = torch.ones(b, 48, dtype=torch.bool)
+    ng = torch.Generator(device="cpu").manual_seed(42)
+    noise = torch.randn(b, 50, 32, generator=ng)
+    return (
+        [im.to(device) for im in images],
+        [m.to(device) for m in masks],
+        lang.to(device),
+        lmask.to(device),
+        state.to(device),
+        noise.to(device),
+    )
+
+
+def _selfconsist_model(device):
+    import safetensors.torch
+
+    model = Pi0ForActionPrediction(
+        Pi0Config(max_action_dim=32, max_state_dim=32, chunk_size=50, num_inference_steps=10, dtype="float32")
+    )
+    state = safetensors.torch.load_file(os.path.join(_SELFCONSIST_CKPT, "model.safetensors"))
+    filled = model.load_weights(list(state.items()))
+    expected = {n for n, _ in model.named_parameters() if "rotary_emb" not in n and not n.endswith(".inv_freq")}
+    missing = expected - (filled or set())
+    assert not missing, f"{len(missing)} params received no weight (e.g. {sorted(missing)[:3]})"
+    return model.to(device).eval()
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(not _SELFCONSIST_CKPT, reason="Set PI0_SELFCONSIST_CKPT to a pi0_base dir.")
+@torch.no_grad()
+def test_version_stable_actions():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = _selfconsist_model(device)
+    images, masks, lang, lmask, state, noise = _selfconsist_inputs(device)
+    actions = (
+        model.sample_actions(
+            images=images,
+            image_masks=masks,
+            lang_tokens=lang,
+            lang_masks=lmask,
+            state=state,
+            noise=noise,
+            num_steps=10,
+        )
+        .float()
+        .cpu()
+    )
+    s, m, std = actions.sum().item(), actions.mean().item(), actions.std().item()
+    assert abs(s - _GOLDEN_SUM) < 0.5, f"action sum drifted: {s} vs {_GOLDEN_SUM}"
+    assert abs(std - _GOLDEN_STD) < 0.01, f"action std drifted: {std} vs {_GOLDEN_STD}"
+    assert abs(m - _GOLDEN_MEAN) < 0.01, f"action mean drifted: {m} vs {_GOLDEN_MEAN}"
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(not _SELFCONSIST_CKPT, reason="Set PI0_SELFCONSIST_CKPT to a pi0_base dir.")
+@torch.no_grad()
+def test_version_stable_embed_image():
+    """Directly guards the SigLIP vision-tower load across the >=5.4 flattening."""
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = _selfconsist_model(device)
+    g = torch.Generator(device="cpu").manual_seed(7)
+    img = torch.rand(1, 3, 224, 224, generator=g, dtype=torch.float32).to(device)
+    std = model.paligemma_with_expert.embed_image(img).float().std().item()
+    assert abs(std - _GOLDEN_EMBED_IMAGE_STD) < 0.05, (
+        f"embed_image std drifted: {std} vs {_GOLDEN_EMBED_IMAGE_STD} — "
+        "likely a vision_tower weight-load regression under this transformers version."
+    )

tests/e2e/online_serving/test_pi0_expansion.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""E2E online serving test for π0 (Pi-Zero) OpenPI websocket serving.
+
+Boots ``vllm serve --omni --deploy-config pi0.yaml`` and drives the real OpenPI
+websocket (``/v1/realtime/robot/openpi``) — the same wire path a robot uses
+(handshake metadata → send observation → receive action chunk). Mirrors
+``tests/e2e/online_serving/test_dreamzero_expansion.py``. Needs a GPU + a
+pi0_base checkpoint; skipped in CI unless explicitly run.
+
+The in-process LeRobot parity oracle lives separately in
+``tests/diffusion/models/pi0/test_pi0_parity.py``.
+"""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import (
+    OmniServerParams,
+    get_open_port,
+    pi0_openpi_require_dependencies,
+    pi0_openpi_run_policy_session,
+    pi0_openpi_validate_session_result,
+)
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+MODEL = "lerobot/pi0_base"
+
+
+test_params = [
+    OmniServerParams(
+        model=MODEL,
+        port=8092,
+        server_args=[
+            "--deploy-config",
+            "vllm_omni/deploy/pi0.yaml",
+            "--served-model-name",
+            "pi0",
+            "--enforce-eager",
+            "--disable-log-stats",
+        ],
+        env_dict={
+            "ATTENTION_BACKEND": "torch",
+            "DIFFUSION_ATTENTION_BACKEND": "TORCH_SDPA",
+            "VLLM_DISABLE_COMPILE_CACHE": "1",
+            "MASTER_PORT": str(get_open_port()),
+        },
+    )
+]
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_pi0_openpi_online(omni_server):
+    try:
+        pi0_openpi_require_dependencies()
+    except ModuleNotFoundError as exc:
+        pytest.skip(str(exc))
+
+    result = pi0_openpi_run_policy_session(
+        host=omni_server.host,
+        port=omni_server.port,
+        prompt="pick up the red block and place it in the bin",
+        session_id="pi0-online-e2e",
+        num_steps=2,
+    )
+
+    # Asserts every returned chunk is [50, 32] + finite, and the handshake
+    # metadata matches pi0.yaml's policy_server_config.
+    pi0_openpi_validate_session_result(result)
+
+    metadata = result["metadata"]
+    assert tuple(metadata["image_resolution"]) == (224, 224)
+    assert metadata["needs_wrist_camera"] is True
+    assert metadata["needs_session_id"] is False
+    assert metadata["action_space"] == "joint_position"

tests/helpers/runtime.py

@@ -3265,6 +3265,146 @@ def iter_omni_runner(
         print("OmniRunner stopped")
 
 
+# ─────────────────────────────────────────────────────────────────────
+# π0 (Pi-Zero) OpenPI websocket policy client
+#
+# Minimal client for the OpenPI realtime robot protocol
+# (``/v1/realtime/robot/openpi``): connect → read handshake metadata → send
+# observation → receive an action chunk. Used by the π0 online-serving e2e
+# (tests/e2e/online_serving/test_pi0_expansion.py). π0's observation is 3 RGB
+# cameras + proprioceptive state + a language prompt, built as synthetic numpy
+# frames here (no video assets).
+# ─────────────────────────────────────────────────────────────────────
+PI0_OPENPI_PATH = "/v1/realtime/robot/openpi"
+PI0_OPENPI_DEFAULT_PROMPT = "pick up the red block and place it in the bin"
+# π0 / pi0_base camera identities (must match the server's image_feature_keys,
+# i.e. vllm_omni/deploy/pi0.yaml, whose image_key_map is empty).
+PI0_CAMERA_KEYS = (
+    "observation.images.base_0_rgb",
+    "observation.images.left_wrist_0_rgb",
+    "observation.images.right_wrist_0_rgb",
+)
+PI0_ACTION_HORIZON = 50
+PI0_ACTION_DIM = 32
+PI0_STATE_DIM = 32
+PI0_IMAGE_SIZE = 224
+
+
+def pi0_openpi_require_dependencies() -> None:
+    """Raise ModuleNotFoundError if the OpenPI websocket client deps are missing."""
+    missing = []
+    try:
+        import websockets.sync.client  # noqa: F401
+    except ImportError:
+        missing.append("websockets")
+    try:
+        from openpi_client import msgpack_numpy  # noqa: F401
+    except ImportError:
+        missing.append("openpi-client")
+    if missing:
+        raise ModuleNotFoundError(f"π0 OpenPI test dependencies are missing: {', '.join(missing)}")
+
+
+def _pi0_decode_action_response(response: bytes | str) -> np.ndarray:
+    from openpi_client import msgpack_numpy
+
+    if isinstance(response, str):
+        raise RuntimeError(f"Inference failed: {response}")
+    decoded = msgpack_numpy.unpackb(response)
+    if isinstance(decoded, dict) and decoded.get("type") == "error":
+        raise RuntimeError(f"Inference failed: {decoded.get('message', decoded)}")
+    return np.asarray(decoded, dtype=np.float32)
+
+
+def pi0_make_dummy_obs(*, prompt: str, session_id: str, image_size: int = PI0_IMAGE_SIZE) -> dict[str, Any]:
+    """A single π0 observation: 3 blank cameras (HWC uint8) + zero state + prompt."""
+    obs: dict[str, Any] = {cam: np.zeros((image_size, image_size, 3), dtype=np.uint8) for cam in PI0_CAMERA_KEYS}
+    obs["state"] = np.zeros(PI0_STATE_DIM, dtype=np.float32)
+    obs["prompt"] = prompt
+    obs["session_id"] = session_id
+    return obs
+
+
+def pi0_openpi_run_policy_session(
+    *,
+    host: str = "127.0.0.1",
+    port: int = 8000,
+    path: str = PI0_OPENPI_PATH,
+    prompt: str = PI0_OPENPI_DEFAULT_PROMPT,
+    session_id: str | None = None,
+    num_steps: int = 2,
+) -> dict[str, Any]:
+    """Connect, read handshake metadata, send ``num_steps`` observations."""
+    import uuid
+
+    import websockets.sync.client as websockets_client
+    from openpi_client import msgpack_numpy
+
+    session_id = session_id or str(uuid.uuid4())
+    uri = f"ws://{host}:{port}{path}"
+    packer = msgpack_numpy.Packer()
+    conn = websockets_client.connect(uri, compression=None, max_size=None, ping_interval=300, ping_timeout=3600)
+    try:
+        metadata = msgpack_numpy.unpackb(conn.recv())
+        if not isinstance(metadata, dict):
+            raise TypeError(f"Expected dict metadata from server, got {type(metadata)!r}")
+        actions = []
+        for _ in range(num_steps):
+            payload = pi0_make_dummy_obs(prompt=prompt, session_id=session_id)
+            payload["endpoint"] = "infer"
+            conn.send(packer.pack(payload))
+            actions.append(_pi0_decode_action_response(conn.recv()))
+        return {"metadata": dict(metadata), "actions": actions, "session_id": session_id}
+    finally:
+        conn.close()
+
+
+def pi0_openpi_validate_session_result(
+    result: dict[str, Any],
+    *,
+    expected_action_horizon: int = PI0_ACTION_HORIZON,
+    expected_action_dim: int = PI0_ACTION_DIM,
+) -> None:
+    """Assert the handshake metadata + every returned action chunk for π0."""
+    metadata = result["metadata"]
+    required_keys = (
+        "image_resolution",
+        "needs_wrist_camera",
+        "needs_session_id",
+        "action_space",
+        "action_horizon",
+        "action_dim",
+    )
+    missing = [key for key in required_keys if key not in metadata]
+    if missing:
+        raise AssertionError(f"Missing π0 metadata keys: {missing}")
+
+    if tuple(metadata["image_resolution"]) != (PI0_IMAGE_SIZE, PI0_IMAGE_SIZE):
+        raise AssertionError(f"Unexpected image_resolution: {metadata['image_resolution']!r}")
+    if not metadata["needs_wrist_camera"]:
+        raise AssertionError("π0 test expects needs_wrist_camera=True")
+    if metadata["needs_session_id"]:
+        raise AssertionError("π0 is stateless; needs_session_id must be False")
+    if metadata["action_space"] != "joint_position":
+        raise AssertionError(f"Unexpected action_space: {metadata['action_space']!r}")
+    if int(metadata["action_horizon"]) != expected_action_horizon:
+        raise AssertionError(f"Unexpected action_horizon: {metadata['action_horizon']!r}")
+    if int(metadata["action_dim"]) != expected_action_dim:
+        raise AssertionError(f"Unexpected action_dim: {metadata['action_dim']!r}")
+
+    actions = result["actions"]
+    if not actions:
+        raise AssertionError("No actions returned from the server")
+    for index, action in enumerate(actions):
+        if action.shape != (expected_action_horizon, expected_action_dim):
+            raise AssertionError(
+                f"Action {index} shape mismatch: expected "
+                f"{(expected_action_horizon, expected_action_dim)}, got {action.shape}"
+            )
+        if not np.isfinite(action).all():
+            raise AssertionError(f"Action {index} contains non-finite values")
+
+
 __all__ = [
     "DiffusionResponse",
     "HttpResponse",
@@ -3286,4 +3426,8 @@ __all__ = [
     "OpenAIClientHandler",
     "get_open_port",
     "dummy_messages_from_mix_data",
+    "pi0_openpi_require_dependencies",
+    "pi0_openpi_run_policy_session",
+    "pi0_openpi_validate_session_result",
+    "pi0_make_dummy_obs",
 ]

tests/model_tests/diffusion/test_alignment.py

@@ -77,6 +77,7 @@ EXCLUDED_MODELS = [
     "DreamZeroPipeline",
     "StableDiffusionXLPipeline",
     "Gr00tN1d7Pipeline",
+    "Pi0Pipeline",
     "SoulXSingerPipeline",
     "SoulXSingerSVCPipeline",
 ]

vllm_omni/config/pipeline_registry.py

@@ -33,6 +33,7 @@ from vllm.logger import init_logger
 from vllm_omni.config.stage_config import (
     PipelineConfig,
 )
+from vllm_omni.diffusion.models.pi0.pipeline_pi0 import PI0_PIPELINE
 from vllm_omni.model_executor.models.audex.pipeline import (
     AUDEX_S2S_PIPELINE,
     AUDEX_THINKER_ONLY_PIPELINE,
@@ -139,6 +140,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     "lance": LANCE_PIPELINE,
     "dreamzero": DREAMZERO_PIPELINE,
     "Gr00tN1d7": GR00T_N1D7_PIPELINE,
+    "pi0": PI0_PIPELINE,
     "gepard": GEPARD_PIPELINE,
     "glm_image": GLM_IMAGE_PIPELINE,
     "hunyuan_image_3_moe": HUNYUAN_IMAGE3_PIPELINE,

vllm_omni/deploy/pi0.yaml

@@ -0,0 +1,52 @@
+# π0 (Pi-Zero) VLA deploy: single diffusion stage.
+#
+# Topology (PI0_PIPELINE) is declared in vllm_omni/diffusion/models/pi0/pipeline_pi0.py.
+# π0 takes multi-camera images + a language instruction + robot state and emits
+# a continuous action chunk [chunk_size, max_action_dim] via flow matching.
+# Actions are exposed through the OpenPI realtime serving API
+# (/v1/realtime/robot/openpi); they ride in multimodal_output["actions"].
+#
+# Default: one GPU, TP=1, eager (flow matching needs no CUDA graph), bsz 1
+# (OpenPI policy serving is one observation -> one action reply). float32 to
+# match the LeRobot pi0_base reference numerics.
+
+pipeline: pi0
+async_chunk: false
+distributed_executor_backend: mp
+dtype: float32
+
+stages:
+  - stage_id: 0
+    devices: "0"
+    max_num_seqs: 1
+    enforce_eager: true
+    model_class_name: Pi0Pipeline
+    model_config:
+      paligemma_variant: gemma_2b
+      action_expert_variant: gemma_300m
+      chunk_size: 50
+      max_action_dim: 32
+      max_state_dim: 32
+      num_inference_steps: 10
+      image_resolution: [224, 224]
+      tokenizer_max_length: 48
+      max_cameras: 3
+      # Camera identities (the order the model attends to). Serving maps raw obs
+      # keys onto these via image_key_map; leave image_key_map empty to match
+      # incoming keys verbatim.
+      image_feature_keys:
+        - observation.images.base_0_rgb
+        - observation.images.left_wrist_0_rgb
+        - observation.images.right_wrist_0_rgb
+      image_key_map: {}
+      # OpenPI handshake metadata advertised to the robot client.
+      policy_server_config:
+        image_resolution: [224, 224]
+        n_external_cameras: 1
+        needs_wrist_camera: true
+        needs_stereo_camera: false
+        needs_session_id: false
+        action_space: joint_position
+        action_horizon: 50
+        action_dim: 32
+        max_cameras: 3

vllm_omni/diffusion/data.py

@@ -1318,6 +1318,12 @@ class OmniDiffusionConfig:
                         self.update_multimodal_support()
                     else:
                         raise
+                elif cfg.get("type") == "pi0":
+                    # π0 (Pi-Zero) VLA — the LeRobot config.json uses ``type: "pi0"``.
+                    if self.model_class_name is None:
+                        self.model_class_name = "Pi0Pipeline"
+                    self.set_tf_model_config(TransformerConfig())
+                    self.update_multimodal_support()
                 elif architectures and len(architectures) == 1:
                     architecture = architectures[0]
                     from vllm_omni.diffusion.registry import DiffusionModelRegistry

vllm_omni/diffusion/models/pi0/__init__.py

@@ -0,0 +1,48 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""π0 (Pi-Zero) VLA model for vllm-omni.
+
+PaliGemma (SigLIP vision + Gemma-2B LM) + Gemma-300M action expert + a
+flow-matching action head. Outputs a continuous action chunk
+``[horizon, action_dim]`` rather than tokens.
+"""
+
+from vllm_omni.diffusion.models.pi0.config import Pi0Config
+from vllm_omni.diffusion.models.pi0.modeling_pi0 import (
+    GemmaVariantConfig,
+    PaliGemmaWithActionExpert,
+    Pi0ForActionPrediction,
+    create_sinusoidal_pos_embedding,
+    get_gemma_config,
+    make_att_2d_masks,
+    prepare_attention_masks_4d,
+)
+from vllm_omni.diffusion.models.pi0.pipeline_pi0 import (
+    Pi0Pipeline,
+    get_pi0_post_process_func,
+)
+from vllm_omni.diffusion.models.pi0.processor_pi0 import (
+    Pi0ImageProcessor,
+    build_model_inputs,
+    pil_image_to_tensor,
+    resize_with_pad,
+    tokenize_prompt,
+)
+
+__all__ = [
+    "Pi0Config",
+    "Pi0ForActionPrediction",
+    "PaliGemmaWithActionExpert",
+    "GemmaVariantConfig",
+    "get_gemma_config",
+    "create_sinusoidal_pos_embedding",
+    "make_att_2d_masks",
+    "prepare_attention_masks_4d",
+    "Pi0ImageProcessor",
+    "build_model_inputs",
+    "pil_image_to_tensor",
+    "resize_with_pad",
+    "tokenize_prompt",
+    "Pi0Pipeline",
+    "get_pi0_post_process_func",
+]

vllm_omni/diffusion/models/pi0/config.py

@@ -0,0 +1,125 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Config surface for the π0 (Pi-Zero) VLA model in vllm-omni.
+
+Only the parameters that actually shape runtime behaviour live here; the
+transformer dimensions (hidden size, head count, ...) are derived from
+``paligemma_variant`` / ``action_expert_variant`` inside the model via
+``get_gemma_config`` (see ``modeling_pi0``).
+
+The resolver reads the **raw LeRobot ``config.json``** (the field surface of
+``lerobot.policies.pi0.PI0Config``): only the recognized dataclass fields are
+used; any other keys in the file are ignored.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from dataclasses import fields as dataclass_fields
+from pathlib import Path
+from typing import Any
+
+# LeRobot / OpenPI observation key conventions.
+ACTION = "action"
+OBS_STR = "observation"
+OBS_STATE = OBS_STR + ".state"
+OBS_IMAGES = OBS_STR + ".images"
+
+
+@dataclass
+class Pi0Config:
+    """π0 VLA config (dataclass, not an HF ``PretrainedConfig``).
+
+    Mirrors the runtime-relevant subset of LeRobot ``PI0Config`` plus a couple
+    of serving-side knobs (``max_cameras``, ``image_feature_keys``,
+    ``image_key_map``, ``norm_stats``).
+    """
+
+    # Backbone variants — mapped to Gemma dimensions by ``get_gemma_config``.
+    paligemma_variant: str = "gemma_2b"
+    action_expert_variant: str = "gemma_300m"
+
+    # Action chunk shape.
+    chunk_size: int = 50
+    max_action_dim: int = 32
+    max_state_dim: int = 32
+
+    # Flow-matching denoising schedule.
+    num_inference_steps: int = 10
+    # Sinusoidal timestep embedding periods (must match OpenPI/LeRobot).
+    min_period: float = 4e-3
+    max_period: float = 4.0
+
+    # Image preprocessing. π0/SigLIP only support square inputs.
+    image_resolution: tuple[int, int] = (224, 224)
+    tokenizer_max_length: int = 48
+    # Number of camera slots the model attends to (real + padded).
+    max_cameras: int = 3
+
+    # Weight dtype the checkpoint was saved in; parity passes this through to
+    # match LeRobot's ``PI0Config.dtype``.
+    dtype: str = "float32"
+
+    # Per-dataset normalization stats, optional (schema matches LeRobot's
+    # ``NormalizerProcessorStep``; consumed by ``_build_norm_buffers``). ``None``
+    # means identity / pass-through (the ``lerobot/pi0_base`` case).
+    norm_stats: dict | None = None
+
+    # Ordered list of image feature keys (e.g. ``observation.images.base_0_rgb``)
+    # taken from the checkpoint's ``input_features``. This is the **camera order**
+    # that must be reproduced for LeRobot parity.
+    image_feature_keys: list[str] | None = None
+
+    # Optional map from raw OpenPI obs keys → ``image_feature_keys`` entries, so
+    # serving can translate wire obs (e.g. ``observation/exterior_image_0_left``)
+    # into the checkpoint's camera identities. Empty → identity match on key.
+    image_key_map: dict[str, str] = field(default_factory=dict)
+
+    # Stored for reference / camera-order derivation; not used at inference.
+    input_features: dict[str, Any] = field(default_factory=dict)
+    output_features: dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        # Coerce list → tuple (JSON has no tuples) and validate squareness.
+        res = self.image_resolution
+        if not isinstance(res, (tuple, list)) or len(res) != 2 or res[0] != res[1]:
+            raise ValueError(f"π0 expects a square image_resolution (H == W); got {res!r}.")
+        self.image_resolution = (int(res[0]), int(res[1]))
+
+        # Derive the camera order from input_features if not given explicitly.
+        if self.image_feature_keys is None and self.input_features:
+            self.image_feature_keys = [key for key in self.input_features if key.startswith(OBS_IMAGES + ".")]
+
+    # ------------------------------------------------------------------
+    # Constructors
+    # ------------------------------------------------------------------
+    @classmethod
+    def from_pretrained(cls, checkpoint_dir: str | Path) -> Pi0Config:
+        """Build from a checkpoint directory's ``config.json``."""
+        checkpoint_dir = Path(checkpoint_dir)
+        config_path = checkpoint_dir / "config.json"
+        if not config_path.exists():
+            return cls()
+        with open(config_path, encoding="utf-8") as f:
+            raw = json.load(f)
+        return cls.from_model_config(raw)
+
+    @classmethod
+    def from_model_config(cls, model_config: dict[str, Any] | None) -> Pi0Config:
+        """Build from a config dict (LeRobot ``config.json`` or deploy yaml).
+
+        Keeps only the recognized dataclass fields (the LeRobot config also
+        carries many training-only keys) and coerces ``image_resolution`` to a
+        tuple.
+        """
+        if not model_config:
+            return cls()
+
+        raw = dict(model_config)
+        if "image_resolution" in raw:
+            raw["image_resolution"] = tuple(raw["image_resolution"])
+
+        allowed = {item.name for item in dataclass_fields(cls)}
+        filtered = {key: value for key, value in raw.items() if key in allowed}
+        return cls(**filtered)

vllm_omni/diffusion/models/pi0/modeling_pi0.py

@@ -0,0 +1,944 @@
+# 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

vllm_omni/diffusion/models/pi0/pipeline_pi0.py

@@ -0,0 +1,277 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""π0 (Pi-Zero) VLA pipeline for vllm-omni.
+
+Entry point for ``DiffusionEngine.step() → pipeline.forward(req)``. Mirrors the
+DreamZero contract: the pipeline owns ALL preprocessing. It reads the raw robot
+observation from ``req.sampling_params.extra_args["robot_obs"]`` (delivered by
+the OpenPI realtime serving layer), builds model inputs, runs flow-matching
+denoising, and returns ``DiffusionOutput(output={"actions": ndarray})`` — which
+``diffusion_engine`` promotes to ``multimodal_output["actions"]`` for the client.
+
+π0 is stateless across calls (no KV reuse), so ``session_id`` / ``reset`` from
+the OpenPI protocol are accepted but ignored.
+"""
+
+from __future__ import annotations
+
+import os
+
+import numpy as np
+import torch
+from torch import nn
+from vllm.logger import init_logger
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.models.pi0.config import Pi0Config
+from vllm_omni.diffusion.models.pi0.modeling_pi0 import Pi0ForActionPrediction
+from vllm_omni.diffusion.models.pi0.processor_pi0 import build_model_inputs
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+logger = init_logger(__name__)
+
+# Default tokenizer for the PaliGemma prefix (matches LeRobot Pi0).
+DEFAULT_PI0_TOKENIZER = "google/paligemma-3b-pt-224"
+
+
+def _pi0_post_process(x):
+    """Module-level identity post-process (picklable across the orchestrator's
+    multiprocess boundary — a local closure is not)."""
+    return x
+
+
+def get_pi0_post_process_func(od_config: OmniDiffusionConfig):
+    """π0 returns actions directly; post-processing is identity. The diffusion
+    engine resolves this via the registry (``_DIFFUSION_POST_PROCESS_FUNCS``) and
+    applies it engine-side, so the pipeline does NOT attach it to DiffusionOutput.
+    """
+    del od_config
+    return _pi0_post_process
+
+
+# Single-stage diffusion topology for online OpenPI serving.
+#
+# π0 is one diffusion stage (robot observation -> action chunk via flow
+# matching). It is registered in ``OMNI_PIPELINES`` (rather than relying on the
+# default single-stage fallback) because π0 is *online-served* via
+# ``vllm serve --deploy-config pi0.yaml``; that path
+# (``OmniStagePipelineFactory._create_from_registry``) requires a registered
+# pipeline so the deploy yaml's ``policy_server_config`` reaches the OpenPI
+# websocket layer. Offline-only models (e.g. ``internvla_a1``) need no topology.
+PI0_PIPELINE = PipelineConfig(
+    model_type="pi0",
+    model_arch="Pi0Pipeline",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="diffusion",
+            execution_type=StageExecutionType.DIFFUSION,
+            input_sources=(),
+            final_output=True,
+            final_output_type="action",
+            model_arch="Pi0Pipeline",
+        ),
+    ),
+)
+
+
+class Pi0Pipeline(nn.Module):
+    """π0 VLA pipeline: raw robot obs → continuous action chunk.
+
+    Registered as ``"Pi0Pipeline"`` in the diffusion registry. Weights are
+    self-loaded in ``__init__`` from the checkpoint's ``model.safetensors`` via
+    the kernel's ``load_weights`` (which handles the LeRobot key remaps).
+    """
+
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__()
+        self.od_config = od_config
+        self.prefix = prefix
+        # Resolve od_config.model to a LOCAL directory. A bare HF repo id (e.g.
+        # the documented ``lerobot/pi0_base``) must be snapshot-downloaded first;
+        # otherwise config/tokenizer/weights would silently fall back to random
+        # init (this pipeline self-loads from a local dir and has no
+        # ``weights_sources`` for the framework loader to download from).
+        self.model_dir = self._resolve_model_dir(od_config.model)
+        self.config = self._build_config(od_config)
+
+        custom_args = od_config.custom_pipeline_args or {}
+        self.tokenizer_source = str(custom_args.get("tokenizer", self._resolve_tokenizer_source()))
+
+        # Torch dtype/device from od_config.
+        self._torch_dtype = self._resolve_dtype(od_config)
+        self._device = self._resolve_device(od_config)
+
+        self.tokenizer = self._load_tokenizer()
+        self.model = self._initialize_model()
+
+    # ------------------------------------------------------------------
+    # Construction helpers
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _resolve_model_dir(model: str | None) -> str | None:
+        """Return a local directory for ``model``; download an HF repo id if needed.
+
+        Limits the snapshot to the files π0 actually reads (config, weights,
+        tokenizer) so a repo id doesn't pull unrelated artifacts.
+        """
+        if not model:
+            return None
+        if os.path.isdir(model):
+            return model
+        from huggingface_hub import snapshot_download
+
+        return snapshot_download(
+            repo_id=model,
+            allow_patterns=["*.json", "*.safetensors", "*.model", "tokenizer*"],
+        )
+
+    def _build_config(self, od_config: OmniDiffusionConfig) -> Pi0Config:
+        """Build Pi0Config from deploy-yaml model_config, falling back to the
+        checkpoint's config.json (raw LeRobot format)."""
+        if od_config.model_config:
+            config = Pi0Config.from_model_config(dict(od_config.model_config))
+            # The deploy yaml is authoritative, but must not silently drop
+            # checkpoint-derived fields it omits. Backfill the camera order and
+            # any normalization stats the yaml didn't specify.
+            if self.model_dir and (not config.image_feature_keys or config.norm_stats is None):
+                ckpt = Pi0Config.from_pretrained(self.model_dir)
+                if not config.image_feature_keys:
+                    config.image_feature_keys = ckpt.image_feature_keys
+                    if not config.input_features:
+                        config.input_features = ckpt.input_features
+                if config.norm_stats is None:
+                    config.norm_stats = ckpt.norm_stats
+            return config
+        if self.model_dir:
+            return Pi0Config.from_pretrained(self.model_dir)
+        return Pi0Config()
+
+    def _resolve_tokenizer_source(self) -> str:
+        """Prefer the checkpoint dir if it ships tokenizer files; else PaliGemma."""
+        if self.model_dir and os.path.isdir(self.model_dir):
+            if os.path.exists(os.path.join(self.model_dir, "tokenizer_config.json")):
+                return self.model_dir
+        return DEFAULT_PI0_TOKENIZER
+
+    @staticmethod
+    def _resolve_dtype(od_config: OmniDiffusionConfig) -> torch.dtype:
+        dt = od_config.dtype
+        if isinstance(dt, torch.dtype):
+            return dt
+        return getattr(torch, str(dt).split(".")[-1], torch.float32)
+
+    @staticmethod
+    def _resolve_device(od_config: OmniDiffusionConfig) -> torch.device:
+        from vllm_omni.diffusion.distributed.utils import get_local_device
+
+        try:
+            return get_local_device()
+        except Exception:  # noqa: BLE001
+            return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    def _load_tokenizer(self):
+        from transformers import AutoTokenizer
+
+        return AutoTokenizer.from_pretrained(self.tokenizer_source)
+
+    def has_real_checkpoint(self) -> bool:
+        return bool(self.model_dir) and os.path.exists(os.path.join(self.model_dir, "model.safetensors"))
+
+    def _initialize_model(self) -> Pi0ForActionPrediction:
+        model = Pi0ForActionPrediction(self.config)
+        if self.has_real_checkpoint():
+            self._load_checkpoint(model)
+        else:
+            logger.info("Pi0Pipeline: no model.safetensors under %s; using random init.", self.model_dir)
+        model.to(device=self._device, dtype=self._torch_dtype)
+        model.eval()
+        return model
+
+    def _load_checkpoint(self, model: Pi0ForActionPrediction) -> None:
+        import safetensors.torch
+
+        path = os.path.join(self.model_dir, "model.safetensors")
+        logger.info("Pi0Pipeline: loading π0 weights from %s", path)
+        state = safetensors.torch.load_file(path)
+        model.load_weights(list(state.items()))
+
+    # ------------------------------------------------------------------
+    # Framework weight-loading hook
+    # ------------------------------------------------------------------
+    def load_weights(self, weights=()):  # noqa: D401
+        """No-op for the diffusion loader: π0 self-loads its checkpoint in
+        ``__init__`` (the kernel's ``load_weights`` handles the LeRobot remaps).
+        We expose no ``weights_sources``, so the loader passes an empty iterator
+        here; returning ``None`` skips its strict unloaded-weights check.
+        """
+        for _ in weights:
+            pass
+        return None
+
+    # ------------------------------------------------------------------
+    # Inference
+    # ------------------------------------------------------------------
+    @torch.inference_mode()
+    def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
+        extra_args = getattr(req.sampling_params, "extra_args", None) or {}
+        robot_obs = extra_args.get("robot_obs")
+
+        if robot_obs is None:
+            # Dummy warmup path (no obs): return zeros so engine warmup/capture
+            # doesn't crash. Mirrors DreamZero's dummy-run handling.
+            first_prompt = req.prompts[0] if req.prompts else ""
+            prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
+            num_steps = getattr(req.sampling_params, "num_inference_steps", None)
+            if prompt == "dummy run" or num_steps == 1:
+                logger.info("Pi0Pipeline: dummy warmup request without robot_obs — returning zeros.")
+                return DiffusionOutput(
+                    output={
+                        "actions": np.zeros(
+                            (self.config.chunk_size, self.config.max_action_dim),
+                            dtype=np.float32,
+                        )
+                    },
+                )
+            return DiffusionOutput(
+                error="Pi0Pipeline.forward requires sampling_params.extra_args['robot_obs'].",
+            )
+
+        images, image_masks, lang_tokens, lang_masks, state = build_model_inputs(
+            robot_obs, self.config, self.tokenizer, self._device
+        )
+
+        # State normalization (identity for pi0_base) inside the model.
+        state = self.model._normalize_state(state)
+
+        noise = extra_args.get("noise")
+        if noise is not None and not isinstance(noise, torch.Tensor):
+            noise = torch.as_tensor(noise, dtype=torch.float32, device=self._device)
+        elif isinstance(noise, torch.Tensor):
+            noise = noise.to(device=self._device, dtype=torch.float32)
+
+        num_steps = extra_args.get("num_inference_steps")
+
+        actions = self.model.sample_actions(
+            images=images,
+            image_masks=image_masks,
+            lang_tokens=lang_tokens,
+            lang_masks=lang_masks,
+            state=state,
+            noise=noise,
+            num_steps=num_steps,
+        )
+        actions = self.model._unnormalize_actions(actions)
+
+        # (B=1, horizon, action_dim) → (horizon, action_dim) numpy for the wire.
+        actions_np = actions.squeeze(0).float().cpu().numpy()
+
+        # Note: post-processing is applied engine-side via the registry
+        # (_DIFFUSION_POST_PROCESS_FUNCS), so we don't attach it here (a local
+        # closure wouldn't survive the orchestrator's multiprocess pickling).
+        return DiffusionOutput(output={"actions": actions_np})

vllm_omni/diffusion/models/pi0/processor_pi0.py

@@ -0,0 +1,222 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Preprocessing for the π0 VLA model.
+
+Converts a raw robot observation (multi-camera images + language instruction +
+proprioceptive state) into the tensors that ``Pi0ForActionPrediction.sample_actions``
+consumes:
+
+  * images       : list of ``(1, 3, 224, 224)`` float tensors in ``[-1, 1]``,
+                   one per camera slot (real cameras first, then ``-1``-filled
+                   padding) in the checkpoint's ``image_feature_keys`` order
+  * image_masks  : list of ``(1,)`` bool tensors, ``True`` = real camera
+  * lang_tokens  : ``(1, L)`` long, PaliGemma tokenizer, ``\\n``-terminated
+  * lang_masks   : ``(1, L)`` bool
+  * state        : ``(1, max_state_dim)`` float32, zero-padded
+
+Image/prompt preprocessing matches OpenPI and LeRobot bit-for-bit — see the
+parity test (``test_pi0_parity.py``). A set of stateless helpers the diffusion
+pipeline calls directly (the DreamZero contract: the pipeline owns its
+preprocessing).
+
+Reference:
+  - OpenPI: openpi/src/openpi/shared/image_tools.py (resize_with_pad)
+  - OpenPI: openpi/src/openpi/models_pytorch/preprocessing_pytorch.py
+  - LeRobot: lerobot/src/lerobot/policies/pi0/modeling_pi0.py (Pi0NewLineProcessor)
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+logger = logging.getLogger(__name__)
+
+# Defaults straight from the π0 reference configs.
+PI0_IMAGE_SIZE = 224  # openpi/models/model.py IMAGE_RESOLUTION
+PI0_NUM_IMAGE_TOKENS = 256  # SigLIP So400m/14 on 224×224 → (224/14)**2
+PI0_IMAGE_TOKEN_INDEX = 257152  # openpi/models_pytorch/gemma_pytorch.py
+PI0_MAX_CAMERAS = 3  # openpi/models/model.py (3 camera slots)
+PI0_MAX_TOKEN_LEN = 48  # openpi/models/pi0_config.py
+
+
+def resize_with_pad(
+    images: torch.Tensor,
+    target_height: int,
+    target_width: int,
+    mode: str = "bilinear",
+) -> torch.Tensor:
+    """Resize ``(B, C, H, W)`` images to the target shape, preserving aspect
+    ratio with -1 padding on the short side.
+
+    Matches openpi ``image_tools.resize_with_pad_torch`` — the clamp to
+    [-1, 1] is what lets the padded region blend with SigLIP-normalized
+    pixels without adding signal at the boundary.
+    """
+    if images.ndim != 4:
+        raise ValueError(f"Expected 4-D (B,C,H,W), got {images.ndim}-D")
+    _, _, cur_h, cur_w = images.shape
+    ratio = max(cur_w / target_width, cur_h / target_height)
+    rh, rw = int(cur_h / ratio), int(cur_w / ratio)
+    align_corners = False if mode == "bilinear" else None
+    resized = F.interpolate(images, size=(rh, rw), mode=mode, align_corners=align_corners)
+    resized = resized.clamp(-1.0, 1.0)
+    ph, rem_h = divmod(target_height - rh, 2)
+    pw, rem_w = divmod(target_width - rw, 2)
+    return F.pad(resized, (pw, pw + rem_w, ph, ph + rem_h), value=-1.0)
+
+
+def pil_image_to_tensor(image: Image.Image) -> torch.Tensor:
+    """PIL → ``(1, C, H, W)`` float32 in ``[-1, 1]`` (SigLIP normalization)."""
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+    arr = np.array(image, dtype=np.float32) / 255.0 * 2.0 - 1.0
+    return torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
+
+
+class Pi0ImageProcessor:
+    """Minimal image preprocessor: image → normalized + padded ``[-1,1]`` tensor."""
+
+    def __init__(self, image_size: int = PI0_IMAGE_SIZE):
+        self.image_size = image_size
+
+    def preprocess_single(self, image: Any) -> torch.Tensor:
+        """Accept a PIL image, an HWC uint8/float ndarray, or a CHW tensor and
+        return a ``(1, 3, image_size, image_size)`` float tensor in ``[-1, 1]``."""
+        t = self._to_tensor(image)
+        if t.shape[2] != self.image_size or t.shape[3] != self.image_size:
+            t = resize_with_pad(t, self.image_size, self.image_size)
+        return t
+
+    def _to_tensor(self, image: Any) -> torch.Tensor:
+        if isinstance(image, Image.Image):
+            return pil_image_to_tensor(image)
+        if isinstance(image, np.ndarray):
+            arr = image
+            # HWC → normalize uint8/[0,255] to [-1,1]; assume already [-1,1] if float.
+            if arr.ndim == 3 and arr.shape[-1] in (1, 3):
+                if np.issubdtype(arr.dtype, np.integer) or arr.max() > 1.0:
+                    arr = arr.astype(np.float32) / 255.0 * 2.0 - 1.0
+                else:
+                    arr = arr.astype(np.float32)
+                return torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
+            return torch.as_tensor(arr, dtype=torch.float32)
+        if isinstance(image, torch.Tensor):
+            t = image
+            if t.ndim == 3:  # (C, H, W) → (1, C, H, W)
+                t = t.unsqueeze(0)
+            return t.to(dtype=torch.float32)
+        raise TypeError(f"Unsupported image type for π0 preprocessing: {type(image)}")
+
+    def make_empty_image(self) -> torch.Tensor:
+        """Fill tensor for an unused camera slot — pure -1, matches OpenPI/LeRobot."""
+        return torch.full((1, 3, self.image_size, self.image_size), -1.0)
+
+
+def tokenize_prompt(tokenizer, text: str, max_token_len: int = PI0_MAX_TOKEN_LEN):
+    """Return ``(input_ids, attention_mask)`` lists, length ``max_token_len``.
+
+    Matches LeRobot's ``Pi0NewLineProcessor``: append ``\\n`` if the caller
+    didn't, then run the PaliGemma tokenizer with right-padding.
+    """
+    prompt = text or ""
+    if not prompt.endswith("\n"):
+        prompt = prompt + "\n"
+    enc = tokenizer(
+        prompt,
+        padding="max_length",
+        max_length=max_token_len,
+        truncation=True,
+        add_special_tokens=True,
+        return_tensors=None,
+    )
+    return list(enc["input_ids"]), list(enc["attention_mask"])
+
+
+def _extract_images(robot_obs: dict, config) -> dict[str, Any]:
+    """Pull a ``{feature_key: image}`` map out of a raw robot obs.
+
+    Accepts two shapes:
+      * ``robot_obs["images"]`` is a dict keyed by camera name; or
+      * flat ``observation/<cam>`` style keys on ``robot_obs`` itself.
+    Keys are translated through ``config.image_key_map`` (raw obs key →
+    checkpoint feature key) so serving wire names map onto ``input_features``.
+    """
+    images = robot_obs.get("images")
+    if not isinstance(images, dict):
+        # Fall back to flat keys (anything that isn't a known scalar field).
+        images = {k: v for k, v in robot_obs.items() if k not in ("state", "prompt", "session_id", "reset", "images")}
+    key_map = getattr(config, "image_key_map", None) or {}
+    return {key_map.get(k, k): v for k, v in images.items()}
+
+
+def build_model_inputs(robot_obs: dict, config, tokenizer, device: torch.device):
+    """Convert a raw robot observation into ``sample_actions`` inputs.
+
+    Reproduces the LeRobot camera ordering exactly: iterate
+    ``config.image_feature_keys`` (the ordered camera identities from the
+    checkpoint's ``input_features``); for each, use the supplied image (mask
+    ``True``) or a ``-1``-filled empty image (mask ``False``). Then tokenize
+    the prompt and zero-pad/truncate the state to ``max_state_dim``.
+
+    Returns ``(images, image_masks, lang_tokens, lang_masks, state)`` with a
+    leading batch dim of 1 on the tensors.
+    """
+    image_size = int(config.image_resolution[0])
+    img_proc = Pi0ImageProcessor(image_size=image_size)
+
+    feature_keys = config.image_feature_keys or []
+    if not feature_keys:
+        # No declared camera order — fall back to whatever the obs provides,
+        # capped at max_cameras (preserves insertion order).
+        provided = _extract_images(robot_obs, config)
+        feature_keys = list(provided.keys())[: config.max_cameras]
+
+    obs_images = _extract_images(robot_obs, config)
+
+    images: list[torch.Tensor] = []
+    image_masks: list[torch.Tensor] = []
+    for key in feature_keys[: config.max_cameras]:
+        img = obs_images.get(key)
+        if img is not None:
+            tensor = img_proc.preprocess_single(img).to(device=device)
+            mask = True
+        else:
+            tensor = img_proc.make_empty_image().to(device=device)
+            mask = False
+        images.append(tensor)
+        image_masks.append(torch.tensor([mask], dtype=torch.bool, device=device))
+
+    # Pad up to max_cameras with empty slots if the checkpoint declares fewer
+    # cameras than the model attends to.
+    while len(images) < config.max_cameras:
+        images.append(img_proc.make_empty_image().to(device=device))
+        image_masks.append(torch.tensor([False], dtype=torch.bool, device=device))
+
+    prompt = robot_obs.get("prompt", "") or ""
+    ids, attn = tokenize_prompt(tokenizer, prompt, config.tokenizer_max_length)
+    lang_tokens = torch.tensor([ids], dtype=torch.long, device=device)
+    lang_masks = torch.tensor([attn], dtype=torch.bool, device=device)
+
+    state = _build_state(robot_obs.get("state"), config.max_state_dim, device)
+
+    return images, image_masks, lang_tokens, lang_masks, state
+
+
+def _build_state(raw_state, max_state_dim: int, device: torch.device) -> torch.Tensor:
+    """Zero-pad / truncate the raw state to ``(1, max_state_dim)`` float32."""
+    if raw_state is None:
+        return torch.zeros(1, max_state_dim, dtype=torch.float32, device=device)
+    state = torch.as_tensor(raw_state, dtype=torch.float32, device=device)
+    if state.dim() == 1:
+        state = state[None, :]
+    if state.shape[-1] < max_state_dim:
+        state = F.pad(state, (0, max_state_dim - state.shape[-1]))
+    elif state.shape[-1] > max_state_dim:
+        state = state[..., :max_state_dim]
+    return state

vllm_omni/diffusion/registry.py

@@ -181,6 +181,11 @@ _DIFFUSION_MODELS = {
         "pipeline_gr00t",
         "Gr00tN1d7Pipeline",
     ),
+    "Pi0Pipeline": (
+        "pi0",
+        "pipeline_pi0",
+        "Pi0Pipeline",
+    ),
     "LongCatImageEditPipeline": (
         "longcat_image",
         "pipeline_longcat_image_edit",
@@ -358,6 +363,11 @@ _NO_CACHE_ACCELERATION = {
     # Pipelines that do not support cache acceleration (cache_dit / tea_cache).
     "NextStep11Pipeline",
     "AudioXPipeline",
+    # π0 is a flow-matching VLA with a self-contained sample_actions loop and no
+    # DiT-style ``.transformer`` block list, so cache_dit / tea_cache cannot apply
+    # to it; list it here so a stray cache_backend override disables gracefully
+    # instead of erroring.
+    "Pi0Pipeline",
     "LingBotWorldCausalDMDPipeline",
 }
 
@@ -560,6 +570,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "LancePipeline": "get_lance_post_process_func",
     "MingImagePipeline": "get_ming_image_post_process_func",
     "InternVLAA1Pipeline": "get_internvla_a1_post_process_func",
+    "Pi0Pipeline": "get_pi0_post_process_func",
     "LongCatImageEditPipeline": "get_longcat_image_post_process_func",
     "StableDiffusion3Pipeline": "get_sd3_image_post_process_func",
     "FluxKontextPipeline": "get_flux_kontext_post_process_func",