diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 62645aae6..d9fdb09a4 100644 --- a/docs/models/supported_models.md +++ b/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` | ✅︎ | ✅︎ | | ✅︎ | — | diff --git a/examples/online_serving/pi0/README.md b/examples/online_serving/pi0/README.md new file mode 100644 index 000000000..1f424a88b --- /dev/null +++ b/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": "", # 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. diff --git a/examples/online_serving/pi0/openpi_client.py b/examples/online_serving/pi0/openpi_client.py new file mode 100644 index 000000000..f2154b33d --- /dev/null +++ b/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()) diff --git a/recipes/lerobot/Pi0.md b/recipes/lerobot/Pi0.md new file mode 100644 index 000000000..6341e57a0 --- /dev/null +++ b/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": "", # 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: +- Checkpoint: +- 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) diff --git a/tests/diffusion/models/pi0/__init__.py b/tests/diffusion/models/pi0/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/tests/diffusion/models/pi0/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/diffusion/models/pi0/test_pi0_parity.py b/tests/diffusion/models/pi0/test_pi0_parity.py new file mode 100644 index 000000000..3b6358b50 --- /dev/null +++ b/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() diff --git a/tests/diffusion/models/pi0/test_pi0_units.py b/tests/diffusion/models/pi0/test_pi0_units.py new file mode 100644 index 000000000..a1fa18f1a --- /dev/null +++ b/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." + ) diff --git a/tests/e2e/online_serving/test_pi0_expansion.py b/tests/e2e/online_serving/test_pi0_expansion.py new file mode 100644 index 000000000..3a9c3b177 --- /dev/null +++ b/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" diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index c48411c9b..4d06eb86a 100644 --- a/tests/helpers/runtime.py +++ b/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", ] diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index 423ccc7a8..7faf175f8 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -77,6 +77,7 @@ EXCLUDED_MODELS = [ "DreamZeroPipeline", "StableDiffusionXLPipeline", "Gr00tN1d7Pipeline", + "Pi0Pipeline", "SoulXSingerPipeline", "SoulXSingerSVCPipeline", ] diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 69e3ebe11..e8e3448b2 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/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, diff --git a/vllm_omni/deploy/pi0.yaml b/vllm_omni/deploy/pi0.yaml new file mode 100644 index 000000000..8d93855d8 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index a5b6d06b5..0ea85405b 100644 --- a/vllm_omni/diffusion/data.py +++ b/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 diff --git a/vllm_omni/diffusion/models/pi0/__init__.py b/vllm_omni/diffusion/models/pi0/__init__.py new file mode 100644 index 000000000..3862d1836 --- /dev/null +++ b/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", +] diff --git a/vllm_omni/diffusion/models/pi0/config.py b/vllm_omni/diffusion/models/pi0/config.py new file mode 100644 index 000000000..ef0484f8e --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/models/pi0/modeling_pi0.py b/vllm_omni/diffusion/models/pi0/modeling_pi0.py new file mode 100644 index 000000000..a45d931ef --- /dev/null +++ b/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..*``. Modern transformers + nests those sub-modules one level deeper (``paligemma.model.``), + 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 diff --git a/vllm_omni/diffusion/models/pi0/pipeline_pi0.py b/vllm_omni/diffusion/models/pi0/pipeline_pi0.py new file mode 100644 index 000000000..0877df4d6 --- /dev/null +++ b/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}) diff --git a/vllm_omni/diffusion/models/pi0/processor_pi0.py b/vllm_omni/diffusion/models/pi0/processor_pi0.py new file mode 100644 index 000000000..0133e5729 --- /dev/null +++ b/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/`` 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 diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 2c4bffb4c..47063bc4a 100644 --- a/vllm_omni/diffusion/registry.py +++ b/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",