diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index 1fdb1ef4a..adc94f9e9 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -223,6 +223,17 @@ steps: -m "full_model and diffusion and H100 and B200 and cards_1" --run-level "full_model" + - label: ":robot_face: Robot Policy OpenPI · π0.5 · H100 · Single-GPU" + timeout_in_minutes: 120 + commands: + - pip install --no-deps "openpi-client @ git+https://github.com/Physical-Intelligence/openpi.git@215abfb217dbac7d5f1273282331b9b1866c0479#subdirectory=packages/openpi-client" + - >- + pytest -sv + tests/e2e/online_serving/test_pi05_expansion.py + -m "full_model and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · 2-GPU" source_file_dependencies: diffusion_qwen_image_function timeout_in_minutes: 120 diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 42971ed97..f481154fb 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -58,6 +58,7 @@ th { | `Magi2Pipeline` | MAGI-2 native Preview-stage T2VA and I2VA (272p/540p) | `sand-ai/MAGI-2-preview` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SandAI/MAGI-2-preview-L20X.md) | | `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) | +| `Pi05Pipeline` | π0.5 VLA | `lerobot/pi05_base` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/lerobot/Pi05.md) | | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | — | | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | — | | `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | diff --git a/recipes/README.md b/recipes/README.md index 078466e80..a7c10ca14 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -29,6 +29,8 @@ Example layout: ```text recipes/ + lerobot/ + Pi05.md Qwen/ Qwen3-Omni-H100.md Qwen3-TTS-A100.md @@ -53,6 +55,7 @@ recipes/ | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB | | [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) | | [`k2-fsa/OmniVoice.md`](./k2-fsa/OmniVoice.md) | Offline and online multilingual text-to-speech | 1x AMD MI300X (ROCm 7.2) | +| [`lerobot/Pi05.md`](./lerobot/Pi05.md) | Online OpenPI serving for the π0.5 vision-language-action policy | 1x H100 80GB | | [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU | | [`IndexTeam/IndexTTS-2_5.md`](./IndexTeam/IndexTTS-2_5.md) | Online + offline multilingual voice-cloned TTS with native speed and emotion control | 1x NVIDIA H20 96GB | | [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, FL2VA, image+audio Ref2VA, and multi-video Ref2VA serving | 1x GPU with CPU offload / 4x B300 | diff --git a/recipes/lerobot/Pi05.md b/recipes/lerobot/Pi05.md new file mode 100644 index 000000000..af648f5ad --- /dev/null +++ b/recipes/lerobot/Pi05.md @@ -0,0 +1,154 @@ +# Pi0.5 (π0.5) VLA + +> Vision-Language-Action policy serving over the OpenPI realtime websocket + +## Summary + +- Vendor: Physical Intelligence (checkpoint redistributed by LeRobot) +- Model: `lerobot/pi05_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.5 + 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/pi05_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 (`Pi05Pipeline`) serves the policy. π0.5 is stateless +across calls (first-order Markov, no KV reuse), so the protocol's `session_id` / +`reset` are accepted but ignored. + +This recipe covers the functional path only. Serving runs in eager mode, with no +CUDA Graph capture, no Triton kernels and no prefix or encoder caching, and both +`float32` and `bfloat16` are supported. Two π0.5 capabilities are not +implemented, and a checkpoint that declares either is rejected when it loads: + +- **MEM** short-horizon observation memory (`use_visual_memory`, + `use_proprioceptive_memory`) +- **RTC** real-time chunking (`rtc_config`) + +A checkpoint with `use_relative_actions` must also carry `action_feature_names` +in `config.json`. + +## Weights + +| Component | Params | Part of | +| --- | ---: | --- | +| Vision tower, SigLIP So400m/14 | 412 M | PaliGemma | +| Language model, Gemma-2B, with the multimodal projector | 2.51 B | PaliGemma | +| Action expert, Gemma-300M, with AdaRMS and the action/time heads | 693 M | π0.5 | +| Total | 3.62 B | 13.5 GiB in float32 | + +| Artifact | Repository | +| --- | --- | +| `model.safetensors`, 812 tensors, ~14.5 GB | `lerobot/pi05_base` | +| `policy_preprocessor.json` + `policy_preprocessor_step_*.safetensors` | `lerobot/pi05_base` | +| Tokenizer, gated | `google/paligemma-3b-pt-224` | + +## Hardware Support + +### 1x H100 80GB (Online serving) + +The nightly end-to-end job runs on one H100, and serving was also verified on an +RTX 4090 and an RTX PRO 5000. + +#### Environment + +- OS: Ubuntu 22.04+ +- Python: 3.12+ +- Driver / runtime: NVIDIA CUDA environment +- vLLM / vLLM-Omni: use the commit you are deploying from +- Hugging Face: `google/paligemma-3b-pt-224` license accepted, `hf auth login` +- OpenPI client extras (robot side, and for the e2e test): `openpi-client`, + `websockets`, `msgpack`, `msgpack-numpy`. CI pins `openpi-client` to OpenPI + revision `215abfb217dbac7d5f1273282331b9b1866c0479`. + +#### Command + +```bash +vllm serve lerobot/pi05_base \ + --omni \ + --host 0.0.0.0 --port 8000 \ + --deploy-config vllm_omni/deploy/pi05.yaml \ + --enforce-eager --disable-log-stats +``` + +The deploy config (`vllm_omni/deploy/pi05.yaml`) declares a single diffusion +stage (`Pi05Pipeline`), `float32`, `max_num_seqs: 1`, `tokenizer_max_length: 200` +and `state_num_bins: 256`, together with the `policy_server_config` handshake +metadata (3 cameras, `joint_position`, action horizon 50, action dim 32) that +enables the OpenPI realtime serving layer. Keep `action_horizon`, `action_dim` +and `image_resolution` in that block in sync with `model_config`. + +To serve in bfloat16, set `dtype: bfloat16` in the same file. + +#### Verification + +The observation is a flat dict per inference: + +```python +{ + "observation.images.base_0_rgb": np.uint8[H, W, 3], # or float32 in [0, 1] + "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], # must match the checkpoint's declared width + "prompt": "pick up the red block", + "session_id": "", # accepted but ignored (π0.5 is stateless) + "sampling_params": {"num_inference_steps": 4}, # optional engine knobs +} +``` + +Camera keys must match the server's `image_feature_keys`, which is the +checkpoint's `input_features` order. If your robot uses different camera names, +set `model_config.image_key_map` in `pi05.yaml` to map raw obs keys → feature +keys. Every request needs at least one configured camera and a non-null state. + +`state` must have the width `input_features["observation.state"]` declares, 32 +for `lerobot/pi05_base`. That checkpoint ships no normalization statistics, so +its values must already be normalized. + +PIL and uint8 images use the `[0, 255]` domain, while floating ndarray and +tensor images must be finite and in `[0, 1]`. Both are converted with +`x * 2 - 1`. + +Engine knobs live under `sampling_params`. It currently holds +`num_inference_steps`, which overrides the deploy-config denoising steps for a +single request. + +### Target-hardware validation + +Client E2E round trip over the OpenPI websocket: msgpack pack, send, engine, +denoise, unpack. Eager, batch 1, 10 denoising steps, p50 of 30 round +trips after 3 warmup. + +| GPU | dtype | Client E2E p50 (ms), 1 / 2 / 3 cameras | Device peak (GiB) | +| --- | --- | ---: | ---: | +| H100 80 GB HBM3 (sm_90) | float32 | 282 / 356 / 398 | 15.0 | +| H100 80 GB HBM3 (sm_90) | bfloat16 | 290 / 301 / 312 | 9.1 | +| RTX PRO 5000 72 GB (sm_120) | float32 | 185 / 231 / 281 | 14.8 | +| RTX PRO 5000 72 GB (sm_120) | bfloat16 | 156 / 177 / 203 | 8.9 | + +## Correctness + +π0.5's flow-matching kernel reproduces the LeRobot `PI05Policy` reference +exactly (`max|Δ| = 0`, CPU/float32, fixed noise, 1/2/3 cameras). The same file +also compares the prompt tokens carrying the discretized state, the +relative-action path, and a bfloat16 run against LeRobot in bfloat16: +`tests/diffusion/models/pi05/test_pi05_parity.py`. + +## References + +- Model blog: +- LeRobot π0.5 docs: +- Checkpoint: +- Pipeline: [`vllm_omni/diffusion/models/pi05/pipeline_pi05.py`](../../vllm_omni/diffusion/models/pi05/pipeline_pi05.py) +- Deploy config: [`vllm_omni/deploy/pi05.yaml`](../../vllm_omni/deploy/pi05.yaml) +- Tests: CPU units + LeRobot parity in [`tests/diffusion/models/pi05/`](../../tests/diffusion/models/pi05/), + OpenPI websocket e2e in + [`tests/e2e/online_serving/test_pi05_expansion.py`](../../tests/e2e/online_serving/test_pi05_expansion.py) diff --git a/tests/diffusion/models/pi05/__init__.py b/tests/diffusion/models/pi05/__init__.py new file mode 100644 index 000000000..4244ce4b4 --- /dev/null +++ b/tests/diffusion/models/pi05/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project diff --git a/tests/diffusion/models/pi05/test_pi05_parity.py b/tests/diffusion/models/pi05/test_pi05_parity.py new file mode 100644 index 000000000..fb30cb445 --- /dev/null +++ b/tests/diffusion/models/pi05/test_pi05_parity.py @@ -0,0 +1,500 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""π0.5 parity against LeRobot. + +Reference: **lerobot 0.6.1**, transformers 5.5.4, torch 2.11.0+cpu. The oracle +is LeRobot's own code, so a later version changing its π0.5 implementation +fails this without anything here regressing — diff ``policies/pi05/`` first. + +Compares final action chunks under fixed noise: float32 exactly, bfloat16 +against LeRobot's own bfloat16 layout within ``BF16_ATOL``. + +Needs a separate ``lerobot[pi]`` venv (its dependencies conflict with +vllm-omni's) with the vllm-omni ``pi05`` package importable:: + + python -m pytest tests/diffusion/models/pi05/test_pi05_parity.py -v -s + +Skips itself when LeRobot is absent. ``PI05_PARITY_MODEL_PATH`` points at a +local checkpoint to skip the HF download. +""" + +from __future__ import annotations + +import copy +import importlib.util +import os + +import pytest +import torch + +# local_model: needs real weights + a lerobot venv, so it runs locally rather +# than in the ready-CI. Additionally gated on lerobot being importable. +_HAS_LEROBOT = importlib.util.find_spec("lerobot") is not None + +# cpu: the comparison runs on CPU in float32 so both implementations see +# identical numerics; local_model: it needs a lerobot venv that CI does not have. +pytestmark = [pytest.mark.local_model, pytest.mark.diffusion, pytest.mark.cpu] + + +# ─── Config (fixed; matches LeRobot defaults for a π0.5 checkpoint) ──── +DEVICE = "cpu" +DTYPE_STR = "float32" +ATOL = 1e-4 +BF16_ATOL = 5e-2 # measured 2.95e-2, reproducible; margin is for version drift +NUM_STEPS = 10 +BATCH_SIZE = 2 +ACTION_DIM = 32 +STATE_DIM = 32 +ACTION_HORIZON = 50 +MAX_TOKEN_LEN = 200 # π0 uses 48 +NUM_STATE_BINS = 256 + +MODEL_PATH = os.environ.get("PI05_PARITY_MODEL_PATH", "lerobot/pi05_base") + +CAMERAS = ("base_0_rgb", "left_wrist_0_rgb", "right_wrist_0_rgb") + + +def _resolve_checkpoint_dir() -> str: + 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 ────────────────────────────────────────────── +def _dummy_dataset_stats() -> dict: + """Identity-ish stats. π0.5 normalizes with QUANTILES by default, so q01/q99 + carry the load here where π0's mean/std would. + + Using dummy stats means this test needs no checkpoint with real + ``norm_stats`` — it verifies numerical equivalence, not calibration. + """ + return { + "observation.state": { + "mean": torch.zeros(STATE_DIM), + "std": torch.ones(STATE_DIM), + "q01": -torch.ones(STATE_DIM), + "q99": torch.ones(STATE_DIM), + }, + "action": { + "mean": torch.zeros(ACTION_DIM), + "std": torch.ones(ACTION_DIM), + "q01": -torch.ones(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 CAMERAS + }, + } + + +def _create_dummy_batch(batch_size: int = BATCH_SIZE, num_views: int = 3, 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" + batch = { + "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), + "task": [prompt for _ in range(batch_size)], + } + for cam in CAMERAS[:num_views]: + batch[f"observation.images.{cam}"] = torch.rand(batch_size, 3, 224, 224, generator=g, dtype=torch.float32).to( + device + ) + return batch + + +# ─── Instantiation ──────────────────────────────────────────────────── +def _instantiate_lerobot(use_relative_actions: bool = False, action_feature_names=None, dtype: str = DTYPE_STR): + from lerobot.policies.pi05 import PI05Policy + from lerobot.policies.pi05.processor_pi05 import make_pi05_pre_post_processors + + policy = PI05Policy.from_pretrained(MODEL_PATH, strict=True, dtype=dtype) + policy.to(DEVICE) + policy.config.device = DEVICE + policy.config.use_relative_actions = use_relative_actions + if action_feature_names is not None: + # Must land on the config *before* the processors are built: + # make_pi05_pre_post_processors captures action_feature_names into + # RelativeActionsProcessorStep at construction time, and its _build_mask + # treats ``action_names is None`` as "shift every dim" — including the + # gripper that relative_exclude_joints is supposed to hold back. Setting + # it afterwards leaves the step with the all-True mask and silently + # compares two different transforms. + policy.config.action_feature_names = action_feature_names + policy.eval() + + pre, post = make_pi05_pre_post_processors(config=policy.config, dataset_stats=_dummy_dataset_stats()) + return policy, pre, post + + +def _instantiate_vllm_omni( + use_relative_actions: bool = False, + action_feature_names=None, + dtype: str = DTYPE_STR, +): + """Build the vllm-omni π0.5 model in isolation (no pipeline, no engine).""" + from vllm_omni.diffusion.models.pi05 import Pi05Config, Pi05ForActionPrediction + from vllm_omni.diffusion.models.pi05.pipeline_pi05 import _set_inference_dtype + + cfg = Pi05Config( + max_action_dim=ACTION_DIM, + max_state_dim=STATE_DIM, + chunk_size=ACTION_HORIZON, + num_inference_steps=NUM_STEPS, + tokenizer_max_length=MAX_TOKEN_LEN, + state_num_bins=NUM_STATE_BINS, + dtype=dtype, + use_relative_actions=use_relative_actions, + relative_exclude_joints=["gripper"] if action_feature_names else [], + action_feature_names=action_feature_names, + ) + model = Pi05ForActionPrediction(cfg) + _set_inference_dtype(model, getattr(torch, dtype)) + model.to(device=DEVICE) + _load_lerobot_weights(model) + model.eval() + return model, cfg + + +def _load_lerobot_weights(model): + import safetensors.torch + + path = os.path.join(_resolve_checkpoint_dir(), "model.safetensors") + state = safetensors.torch.load_file(path) + filled = model.load_weights(list(state.items())) + assert filled, "no weights were loaded — the remap rules are broken" + + +def _extract_lerobot_model_inputs(lerobot_policy, processed_batch): + """Mimic what ``PI05Policy.predict_action_chunk`` feeds into + ``self.model.sample_actions``. Using these *exact* tensors for vllm-omni + means any divergence must come from the core model, not preprocessing. + + Note π0.5 has no ``prepare_state`` step: the state is already inside + ``lang_tokens``. + """ + from lerobot.utils.constants import ( + OBS_LANGUAGE_ATTENTION_MASK, + OBS_LANGUAGE_TOKENS, + ) + + images, img_masks = lerobot_policy._preprocess_images(processed_batch) + return images, img_masks, processed_batch[OBS_LANGUAGE_TOKENS], processed_batch[OBS_LANGUAGE_ATTENTION_MASK] + + +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 parity test ───────────────────────────────────────────────── +@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).") +@pytest.mark.parametrize("num_views", [1, 2, 3]) +def test_pi05_vllm_omni_vs_lerobot(num_views): + lerobot_policy, lerobot_pre, _ = _instantiate_lerobot() + omni_model, _ = _instantiate_vllm_omni() + + raw_batch = _create_dummy_batch(num_views=num_views) + processed = lerobot_pre(copy.deepcopy(raw_batch)) + images, img_masks, lang_tokens, lang_masks = _extract_lerobot_model_inputs(lerobot_policy, processed) + noise = _make_fixed_noise(raw_batch["observation.state"].shape[0], DEVICE) + + # The input contract: 256 tokens per view + a constant 200 text tokens. + assert lang_tokens.shape[1] == MAX_TOKEN_LEN, "π0.5 must pad text to exactly 200 tokens" + + with torch.no_grad(): + lerobot_actions = lerobot_policy.model.sample_actions( + images, img_masks, lang_tokens, lang_masks, noise=noise, num_steps=NUM_STEPS + ) + omni_actions = omni_model.sample_actions( + images=images, + image_masks=img_masks, + lang_tokens=lang_tokens, + lang_masks=lang_masks, + noise=noise, + num_steps=NUM_STEPS, + ) + + diff = (lerobot_actions.float() - omni_actions.float()).abs() + print(f"[parity] views={num_views} |Δ| max={diff.max().item():.2e} mean={diff.mean().item():.2e}") + if not torch.allclose(lerobot_actions.float(), omni_actions.float(), atol=ATOL): + _diagnose_divergence(lerobot_policy.model, omni_model, images, img_masks, lang_tokens, lang_masks, noise) + assert torch.allclose(lerobot_actions.float(), omni_actions.float(), atol=ATOL), ( + f"actions differ beyond atol={ATOL}; max_diff={diff.max().item():.2e}" + ) + + +@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).") +def test_pi05_bfloat16_matches_lerobot_bfloat16(): + """Both sides in bfloat16, so this compares implementations rather than + measuring bfloat16's drift from float32.""" + # LeRobot applies the layout in PaliGemmaWithExpertModel.__init__ from + # config.dtype, so the policy must be built as bfloat16 rather than cast after. + lerobot_policy, lerobot_pre, _ = _instantiate_lerobot(dtype="bfloat16") + + omni_model, _ = _instantiate_vllm_omni(dtype="bfloat16") + + raw_batch = _create_dummy_batch(num_views=3) + processed = lerobot_pre(copy.deepcopy(raw_batch)) + images, img_masks, lang_tokens, lang_masks = _extract_lerobot_model_inputs(lerobot_policy, processed) + noise = _make_fixed_noise(raw_batch["observation.state"].shape[0], DEVICE) + + with torch.no_grad(): + lerobot_actions = lerobot_policy.model.sample_actions( + images, img_masks, lang_tokens, lang_masks, noise=noise, num_steps=NUM_STEPS + ) + omni_actions = omni_model.sample_actions( + images=images, + image_masks=img_masks, + lang_tokens=lang_tokens, + lang_masks=lang_masks, + noise=noise, + num_steps=NUM_STEPS, + ) + + diff = (lerobot_actions.float() - omni_actions.float()).abs() + print(f"[parity] bfloat16 |Δ| max={diff.max().item():.2e} mean={diff.mean().item():.2e}") + assert torch.allclose(lerobot_actions.float(), omni_actions.float(), atol=BF16_ATOL), ( + f"bfloat16 actions differ beyond atol={BF16_ATOL}; max_diff={diff.max().item():.2e}" + ) + + +@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).") +def test_pi05_prompt_parity(): + """π0.5's state is carried as prompt tokens, so the tokenized prompt is + itself part of the contract. A bin-index mismatch here (e.g. from running + the discretizer before the normalizer) would otherwise only surface as a + small, plausible-looking action difference. + """ + from transformers import AutoTokenizer + + from vllm_omni.diffusion.models.pi05.processor_pi05 import ( + apply_norm, + build_norm_stats, + build_pi05_prompt, + tokenize_prompt, + ) + + lerobot_policy, lerobot_pre, _ = _instantiate_lerobot() + raw_batch = _create_dummy_batch(batch_size=1) + processed = lerobot_pre(copy.deepcopy(raw_batch)) + + from lerobot.utils.constants import OBS_LANGUAGE_TOKENS + + lerobot_tokens = processed[OBS_LANGUAGE_TOKENS][0] + + stats = _dummy_dataset_stats()["observation.state"] + norm = build_norm_stats({"state": {"q01": stats["q01"], "q99": stats["q99"]}}, "state") + normalized = apply_norm(raw_batch["observation.state"][0].float(), norm) + prompt = build_pi05_prompt( + task=raw_batch["task"][0], + normalized_state=normalized.numpy(), + state_num_bins=NUM_STATE_BINS, + ) + tokenizer = AutoTokenizer.from_pretrained("google/paligemma-3b-pt-224", padding_side="right") + ids, _ = tokenize_prompt(tokenizer, prompt, MAX_TOKEN_LEN) + + assert len(ids) == MAX_TOKEN_LEN + assert torch.equal(torch.tensor(ids), lerobot_tokens.cpu()), "π0.5 prompt tokens diverge from LeRobot" + + +@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).") +@pytest.mark.parametrize("state_dim", [7, STATE_DIM]) +def test_pi05_prompt_parity_across_state_widths(state_dim): + """``pi05_base`` declares a 32-dim state, which is exactly ``max_state_dim``, + so the full-policy parity above cannot tell padding apart from the real + width. This drives LeRobot's tokenizer step directly, at both widths. + """ + from lerobot.lerobot_types import TransitionKey + from lerobot.policies.pi05.processor_pi05 import Pi05PrepareStateTokenizerProcessorStep + from lerobot.utils.constants import OBS_STATE + + from vllm_omni.diffusion.models.pi05.processor_pi05 import build_pi05_prompt + + g = torch.Generator().manual_seed(state_dim) + state = torch.randn(state_dim, generator=g, dtype=torch.float32) + task = "pick up the red block" + + step = Pi05PrepareStateTokenizerProcessorStep(max_state_dim=STATE_DIM) + transition = step( + { + TransitionKey.OBSERVATION: {OBS_STATE: state.unsqueeze(0)}, + TransitionKey.COMPLEMENTARY_DATA: {"task": [task]}, + } + ) + lerobot_prompt = transition[TransitionKey.COMPLEMENTARY_DATA]["task"][0] + + prompt = build_pi05_prompt(task=task, normalized_state=state.numpy()) + + assert prompt == lerobot_prompt + assert len(prompt.split("State: ")[1].split(";")[0].split()) == state_dim + + +@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).") +@pytest.mark.parametrize( + "lerobot_mode,stats,ours,state,state_dim", + [ + ("MEAN_STD", {"mean": [0.0], "std": [1e-7]}, {"mode": "mean_std", "mean": [0.0], "std": [1e-7]}, [1e-7], 1), + ("QUANTILES", {"q01": [0.0], "q99": [1e-7]}, {"mode": "quantile", "q01": [0.0], "q99": [1e-7]}, [1e-7], 1), + ("QUANTILES", {"q01": [2.0], "q99": [2.0]}, {"mode": "quantile", "q01": [2.0], "q99": [2.0]}, [2.0], 1), + ( + "QUANTILES", + {"q01": [0.0] * 4, "q99": [10.0] * 4}, + {"mode": "quantile", "q01": [0.0] * 4, "q99": [10.0] * 4}, + [-10.0, 0.0, 5.0, 20.0], + 4, + ), + ( + "MEAN_STD", + {"mean": [1.0, 2.0], "std": [0.5, 4.0]}, + {"mode": "mean_std", "mean": [1.0, 2.0], "std": [0.5, 4.0]}, + [3.0, -6.0], + 2, + ), + ("MIN_MAX", {"min": [-1.0], "max": [3.0]}, {"mode": "min_max", "min": [-1.0], "max": [3.0]}, [2.0], 1), + ], +) +def test_pi05_state_normalization_parity(lerobot_mode, stats, ours, state, state_dim): + """State normalization against LeRobot's ``NormalizerProcessorStep``. + + The full-policy parity test feeds both sides stats with a wide range, so it + cannot see the denominator rule. A wrongly normalized state still discretizes + into a well-formed prompt: for ``q01=0, q99=1e-7`` a threshold substitution + moves the state token from bin 255 to bin 0. + """ + from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature + from lerobot.processor.normalize_processor import NormalizerProcessorStep + from lerobot.utils.constants import OBS_STATE + + from vllm_omni.diffusion.models.pi05.processor_pi05 import apply_norm, build_norm_stats + + step = NormalizerProcessorStep( + features={OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(state_dim,))}, + norm_map={FeatureType.STATE: getattr(NormalizationMode, lerobot_mode)}, + stats={OBS_STATE: {k: torch.tensor(v, dtype=torch.float32) for k, v in stats.items()}}, + ) + expected = step._apply_transform( + torch.tensor(state, dtype=torch.float32), OBS_STATE, FeatureType.STATE, inverse=False + ) + + actual = apply_norm(torch.tensor(state, dtype=torch.float32), build_norm_stats({"state": ours}, "state")) + + assert torch.allclose(actual, expected, atol=1e-6) + + +@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).") +def test_pi05_relative_actions_parity(): + """With ``use_relative_actions=True`` the checkpoint's norm_stats live in + relative space, so the post-pipeline must add the raw state back. Compares + the full pre → model → post path against LeRobot's.""" + from lerobot.utils.constants import ( + OBS_LANGUAGE_ATTENTION_MASK, + OBS_LANGUAGE_TOKENS, + ) + + from vllm_omni.diffusion.models.pi05.processor_pi05 import ( + Pi05RelativeActions, + apply_norm, + build_norm_stats, + ) + + action_names = [f"joint_{i}" for i in range(ACTION_DIM - 1)] + ["gripper"] + lerobot_policy, lerobot_pre, lerobot_post = _instantiate_lerobot( + use_relative_actions=True, action_feature_names=action_names + ) + omni_model, cfg = _instantiate_vllm_omni(use_relative_actions=True, action_feature_names=action_names) + + raw_batch = _create_dummy_batch() + processed = lerobot_pre(copy.deepcopy(raw_batch)) + images, img_masks = lerobot_policy._preprocess_images(processed) + lang_tokens = processed[OBS_LANGUAGE_TOKENS] + lang_masks = processed[OBS_LANGUAGE_ATTENTION_MASK] + noise = _make_fixed_noise(raw_batch["observation.state"].shape[0], DEVICE) + + with torch.no_grad(): + lerobot_actions = lerobot_post( + lerobot_policy.model.sample_actions( + images, img_masks, lang_tokens, lang_masks, noise=noise, num_steps=NUM_STEPS + ) + ) + omni_raw = omni_model.sample_actions( + images=images, + image_masks=img_masks, + lang_tokens=lang_tokens, + lang_masks=lang_masks, + noise=noise, + num_steps=NUM_STEPS, + ) + + rel = Pi05RelativeActions( + enabled=True, + exclude_joints=cfg.relative_exclude_joints, + action_names=action_names, + max_action_dim=ACTION_DIM, + ) + omni_actions = apply_norm(omni_raw, build_norm_stats(cfg.norm_stats, "action"), inverse=True) + # Per-sample reference state, matching LeRobot: its + # RelativeActionsProcessorStep caches the whole (B, D) state on the input + # side and shifts each sample by its own row. Passing a single row here + # would shift every sample by sample 0's state, which only agrees for B=1. + omni_actions = rel.to_absolute(omni_actions, raw_batch["observation.state"]) + + lerobot_actions = torch.as_tensor(lerobot_actions).float() + assert torch.allclose(lerobot_actions, omni_actions.float(), atol=ATOL), ( + "relative-action path diverges from LeRobot; " + f"max_diff={(lerobot_actions - omni_actions.float()).abs().max().item():.2e}" + ) + + +# ─── Per-stage divergence diagnostics ───────────────────────────────── +@torch.no_grad() +def _diagnose_divergence(lerobot_flow_model, omni_model, images, img_masks, lang_tokens, lang_masks, noise): + """Localize a numerical mismatch to a specific stage: + 1. Prefix embeddings (SigLIP / embed_tokens / projector) — image vs lang. + 2. Prefix KV cache layer 0 — PaliGemma LM attention. + 3. The AdaRMS timestep conditioning vector. + 4. A single denoise_step velocity at t=1.0 — action expert forward. + """ + from vllm_omni.diffusion.models.pi05.modeling_pi05 import ( + make_att_2d_masks, + prepare_attention_masks_4d, + ) + + lr_embs, lr_pad, lr_att = lerobot_flow_model.embed_prefix(images, img_masks, lang_tokens, lang_masks) + sg_embs, sg_pad, sg_att = omni_model.embed_prefix(images, img_masks, lang_tokens, lang_masks) + print(f"[diag] prefix_embs max|Δ| = {(lr_embs.float() - sg_embs.float()).abs().max().item():.2e}") + print(f"[diag] prefix_pad_masks equal: {torch.equal(lr_pad, sg_pad)}") + + img_len = 256 * len(images) + print(f"[diag] image slice max|Δ| = {(lr_embs[:, :img_len] - sg_embs[:, :img_len]).abs().max().item():.2e}") + print(f"[diag] lang slice max|Δ| = {(lr_embs[:, img_len:] - sg_embs[:, img_len:]).abs().max().item():.2e}") + + prefix_att_4d = prepare_attention_masks_4d(make_att_2d_masks(sg_pad, sg_att)) + prefix_pos = torch.cumsum(sg_pad, dim=1) - 1 + _, sg_kv = omni_model.paligemma_with_expert.forward( + attention_mask=prefix_att_4d, + position_ids=prefix_pos, + past_key_values=None, + inputs_embeds=[sg_embs, None], + use_cache=True, + ) + + t = torch.ones(noise.shape[0], dtype=torch.float32) + sg_cond = omni_model.embed_timestep(t) + print(f"[diag] adarms cond: shape={tuple(sg_cond.shape)} mean={sg_cond.mean().item():+.4f}") + + sg_vt = omni_model.denoise_step(sg_pad, sg_kv, noise, t) + print(f"[diag] denoise_step(t=1) v_t mean={sg_vt.mean().item():+.6f}") + + +if __name__ == "__main__": + test_pi05_vllm_omni_vs_lerobot(3) diff --git a/tests/diffusion/models/pi05/test_pi05_units.py b/tests/diffusion/models/pi05/test_pi05_units.py new file mode 100644 index 000000000..26671a59c --- /dev/null +++ b/tests/diffusion/models/pi05/test_pi05_units.py @@ -0,0 +1,1107 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CPU unit tests for the π0.5 VLA model. + +Deliberately lightweight: config plumbing and the checkpoint-boundary rule, the +state discretization path, relative actions, and the weight-load remap. Runs on +CPU with no weights and no lerobot. The full LeRobot parity check lives in +``tests/diffusion/models/pi05/test_pi05_parity.py``. + +Most assertions here guard *silent* failure modes — cases where the wrong +behaviour still produces a well-shaped, finite action chunk: + + * normalization must run before state discretization; + * the normalization mode comes from the checkpoint, since π0.5 defaults to + quantile where π0 defaults to ``mean_std``; + * relative-action checkpoints must be honoured, and an unresolvable + ``relative_exclude_joints`` must raise rather than silently make every + dimension relative; + * MEM / RTC checkpoints must be rejected rather than served incorrectly. + + pytest tests/diffusion/models/pi05/test_pi05_units.py -v +""" + +from __future__ import annotations + +import logging +from typing import Any + +import numpy as np +import pytest +import torch + +from vllm_omni.diffusion.models.pi05 import modeling_pi05 +from vllm_omni.diffusion.models.pi05.config import ( + Pi05Config, + UnsupportedCheckpointCapabilityError, + load_lerobot_norm_stats, + resolve_excluded_action_indices, +) +from vllm_omni.diffusion.models.pi05.modeling_pi05 import ( + GemmaVariantConfig, + Pi05AdaRMSNorm, + Pi05ForActionPrediction, + create_sinusoidal_pos_embedding, + get_gemma_config, + make_att_2d_masks, +) +from vllm_omni.diffusion.models.pi05.processor_pi05 import ( + Pi05ImageProcessor, + Pi05Processor, + Pi05RelativeActions, + _assemble_model_inputs, + apply_norm, + as_state_vector, + build_norm_stats, + build_pi05_prompt, + discretize_state, + resize_with_pad, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +# ---------------------------------------------------------------------------- +# Pi05Config dataclass + checkpoint resolver +# ---------------------------------------------------------------------------- +_LEROBOT_CFG: dict[str, Any] = { + "type": "pi05", + "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": 200, + "dtype": "float32", + "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]}, + }, + "output_features": {"action": {"type": "ACTION", "shape": [32]}}, +} + +_EXPECTED_CAMERA_ORDER = [ + "observation.images.base_0_rgb", + "observation.images.left_wrist_0_rgb", + "observation.images.right_wrist_0_rgb", +] + +_ACTION_NAMES = ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "gripper"] + + +def test_config_parses_lerobot_pi05_json(): + """``_LEROBOT_CFG`` mirrors the runtime keys of a real checkpoint.""" + c = Pi05Config.from_model_config(_LEROBOT_CFG) + assert c.tokenizer_max_length == 200 + assert c.state_num_bins == 256 + assert c.image_feature_keys == _EXPECTED_CAMERA_ORDER + assert c.image_resolution == (224, 224) + + +def test_config_rejects_wrong_model_type(): + with pytest.raises(ValueError, match="type='pi05'"): + Pi05Config.from_model_config(dict(_LEROBOT_CFG, type="pi0")) + + +def test_config_rejects_unsupported_partial_action_chunk(): + with pytest.raises(UnsupportedCheckpointCapabilityError, match="n_action_steps"): + Pi05Config.from_model_config(dict(_LEROBOT_CFG, n_action_steps=10)) + + +@pytest.mark.parametrize( + "key,value", + [ + ("output_features", {"action": {"type": "ACTION", "shape": [33]}}), + ("input_features", {**_LEROBOT_CFG["input_features"], "observation.state": {"type": "STATE", "shape": [33]}}), + ], + ids=["action", "state"], +) +def test_config_rejects_a_schema_wider_than_the_model(key, value): + """``max_action_dim`` / ``max_state_dim`` are the widths the weights were + built for, so a wider checkpoint schema cannot be served.""" + with pytest.raises(ValueError, match="over max_state_dim"): + Pi05Config.from_model_config(dict(_LEROBOT_CFG, **{key: value})) + + +@pytest.mark.parametrize( + "key,value,attribute", + [ + ("output_features", {"action": {"type": "ACTION", "shape": [7]}}, "action_dim"), + ( + "input_features", + {**_LEROBOT_CFG["input_features"], "observation.state": {"type": "STATE", "shape": [7]}}, + "state_dim", + ), + ], + ids=["action", "state"], +) +def test_config_derives_the_real_width_from_the_schema(key, value, attribute): + """The declared width decides how many values reach the prompt and how many + columns come back, so a 7-joint checkpoint must not be padded out to 32.""" + config = Pi05Config.from_model_config(dict(_LEROBOT_CFG, **{key: value})) + assert getattr(config, attribute) == 7 + + +def test_config_rejects_inconsistent_openpi_metadata(): + with pytest.raises(ValueError, match="policy_server_config.action_dim"): + Pi05Config.from_model_config( + dict( + _LEROBOT_CFG, + output_features={"action": {"type": "ACTION", "shape": [7]}}, + policy_server_config={"action_dim": 32}, + ) + ) + + +def test_config_non_square_resolution_raises(): + with pytest.raises(ValueError): + Pi05Config(image_resolution=(224, 256)) + + +@pytest.mark.parametrize("key", ["max_cameras", "num_inference_steps"]) +def test_config_rejects_counts_below_one(key): + """``build_model_inputs`` emits one image slot per camera and the Euler loop + steps ``-1 / num_inference_steps``; neither is defined at zero.""" + with pytest.raises(ValueError, match=f"{key} must be at least 1"): + Pi05Config.from_model_config(dict(_LEROBOT_CFG, **{key: 0})) + + +@pytest.mark.parametrize("key", ["state", "action"]) +def test_config_rejects_norm_stats_that_do_not_match_the_declared_width(key): + """A checkpoint's statistics and its feature schema come from different + files, so they can disagree.""" + stats = {key: {"mode": "quantile", "q01": [0.0] * 2, "q99": [1.0] * 2}} + with pytest.raises(ValueError, match=rf"norm_stats\[{key}\]\[q01\] has 2 entries, expected 32"): + Pi05Config.from_model_config(dict(_LEROBOT_CFG, norm_stats=stats)) + + +# ---------------------------------------------------------------------------- +# LeRobot normalization-stats sidecar +# ---------------------------------------------------------------------------- +# LeRobot keeps stats out of config.json: policy_preprocessor.json holds the +# structure and one safetensors file per stateful step holds the numbers. +def _normalized(raw, stats): + """The serving path's state normalization, without the pipeline.""" + return apply_norm(torch.as_tensor(raw, dtype=torch.float32), build_norm_stats({"state": stats}, "state")).numpy() + + +def _write_lerobot_checkpoint( + tmp_path, + *, + norm_map: dict[str, str] | None = None, + stats: dict[str, dict[str, list[float]]] | None = None, + write_state_file: bool = True, +) -> str: + """Write a minimal LeRobot-shaped checkpoint dir; return its path.""" + import json as _json + + import safetensors.torch as _st + + tmp_path.mkdir(parents=True, exist_ok=True) + state_file = "policy_preprocessor_step_3_normalizer_processor.safetensors" + step: dict = { + "registry_name": "normalizer_processor", + "config": { + "eps": 1e-08, + "features": {}, + "norm_map": norm_map if norm_map is not None else {"STATE": "QUANTILES", "ACTION": "QUANTILES"}, + }, + } + if write_state_file: + step["state_file"] = state_file + flat = { + f"{feature}.{stat}": torch.tensor(values, dtype=torch.float32) + for feature, entry in (stats or {}).items() + for stat, values in entry.items() + } + _st.save_file(flat, str(tmp_path / state_file)) + + (tmp_path / "policy_preprocessor.json").write_text( + _json.dumps({"name": "policy_preprocessor", "steps": [step]}), encoding="utf-8" + ) + (tmp_path / "config.json").write_text(_json.dumps(_config_for(stats)), encoding="utf-8") + return str(tmp_path) + + +def _config_for(stats: dict | None) -> dict: + """``_LEROBOT_CFG`` with the state/action widths the stats actually cover.""" + if not stats: + return _LEROBOT_CFG + config = dict(_LEROBOT_CFG) + state = next(iter(stats.get("observation.state", {}).values()), None) + action = next(iter(stats.get("action", {}).values()), None) + if state is not None: + features = dict(config["input_features"]) + features["observation.state"] = {"type": "STATE", "shape": [len(state)]} + config["input_features"] = features + if action is not None: + config["output_features"] = {"action": {"type": "ACTION", "shape": [len(action)]}} + return config + + +# A LeRobot state_dict carries every stat compute_stats emits, not just the two +# the norm_map selects. +_FULL_DATASET_STATS = { + "observation.state": { + "mean": [10.0, 20.0], + "std": [1.0, 2.0], + "min": [-5.0, -5.0], + "max": [5.0, 5.0], + "q01": [-1.0, -2.0], + "q99": [1.0, 2.0], + }, + "action": { + "mean": [0.5, 0.5], + "std": [0.1, 0.1], + "min": [-1.0, -1.0], + "max": [1.0, 1.0], + "q01": [-0.25, -0.5], + "q99": [0.25, 0.5], + }, +} + + +@pytest.mark.parametrize( + "declared,expected", + [ + ("QUANTILES", {"mode": "quantile", "q01": [-1.0, -2.0], "q99": [1.0, 2.0]}), + ("MEAN_STD", {"mode": "mean_std", "mean": [10.0, 20.0], "std": [1.0, 2.0]}), + ("IDENTITY", None), + ], +) +def test_load_lerobot_norm_stats_takes_the_mode_from_norm_map(tmp_path, declared, expected): + """The regression guard. ``_FULL_DATASET_STATS`` carries all six statistics, + as a real state_dict does, so an implementation that sniffed the present keys + would answer mean_std for every case here. Only norm_map distinguishes them. + """ + path = _write_lerobot_checkpoint( + tmp_path, norm_map={"STATE": declared, "ACTION": "QUANTILES"}, stats=_FULL_DATASET_STATS + ) + stats = load_lerobot_norm_stats(path) + + if expected is None: + assert "state" not in stats + else: + assert stats["state"] == expected + assert stats["action"]["mode"] == "quantile" + + +@pytest.mark.parametrize("shape", ["no_sidecar", "no_state_file"]) +def test_load_lerobot_norm_stats_absent_stats_are_none(tmp_path, shape): + """``lerobot/pi05_base`` is the second shape: a normalizer step with no + stats attached. Neither is an error — the client is then expected to send an + already-normalized state.""" + if shape == "no_sidecar": + path = str(tmp_path) + else: + path = _write_lerobot_checkpoint(tmp_path, write_state_file=False) + + assert load_lerobot_norm_stats(path) is None + + +def test_load_lerobot_norm_stats_missing_declared_state_file_raises(tmp_path): + path = _write_lerobot_checkpoint(tmp_path, stats=_FULL_DATASET_STATS) + (tmp_path / "policy_preprocessor_step_3_normalizer_processor.safetensors").unlink() + + with pytest.raises(FileNotFoundError, match="normalizer state"): + load_lerobot_norm_stats(path) + + +def test_pipeline_rejects_missing_model_weights_before_initialization(tmp_path): + from vllm_omni.diffusion.models.pi05.pipeline_pi05 import Pi05Pipeline + + pipeline = object.__new__(Pi05Pipeline) + pipeline.model_dir = str(tmp_path) + + with pytest.raises(FileNotFoundError, match="model.safetensors"): + pipeline._initialize_model() + + +def test_build_config_merges_the_deploy_yaml_over_the_checkpoint(tmp_path): + """The deploy yaml carries the runtime knobs and the checkpoint carries the + feature schema and the stats. Both have to land on one config.""" + from types import SimpleNamespace + + from vllm_omni.diffusion.models.pi05.pipeline_pi05 import Pi05Pipeline + + pipeline = object.__new__(Pi05Pipeline) + pipeline.model_dir = _write_lerobot_checkpoint(tmp_path, stats=_FULL_DATASET_STATS) + + config = pipeline._build_config(SimpleNamespace(model_config={"type": "pi05", "num_inference_steps": 4})) + + assert config.num_inference_steps == 4 + assert (config.state_dim, config.action_dim) == (2, 2) + assert config.image_feature_keys == _EXPECTED_CAMERA_ORDER + assert config.norm_stats["state"]["mode"] == "quantile" + + +class _WideActionModel: + """Stands in for the 3.6B model, which is the one collaborator here that is + genuinely expensive to build. It returns the padded ``max_action_dim`` width + the real model emits, so the assertion is about the pipeline's cropping.""" + + def __init__(self, chunk_size: int, max_action_dim: int): + self._shape = (1, chunk_size, max_action_dim) + + def sample_actions(self, **kwargs): + del kwargs + return torch.ones(self._shape) + + +def test_pipeline_crops_actions_to_checkpoint_output_schema(monkeypatch): + from vllm_omni.diffusion.models.pi05 import pipeline_pi05 + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + config = Pi05Config.from_model_config( + dict(_LEROBOT_CFG, output_features={"action": {"type": "ACTION", "shape": [7]}}) + ) + pipeline = object.__new__(pipeline_pi05.Pi05Pipeline) + pipeline.config = config + pipeline.tokenizer = object() + pipeline._device = torch.device("cpu") + pipeline.processor = Pi05Processor(config, object(), torch.device("cpu")) + pipeline.model = _WideActionModel(config.chunk_size, config.max_action_dim) + monkeypatch.setattr( + pipeline.processor, + "build_model_inputs", + lambda robot_obs: ([torch.empty(0)], [torch.empty(0)], torch.empty(0), torch.empty(0)), + ) + request = OmniDiffusionRequest( + prompt="", + sampling_params=OmniDiffusionSamplingParams( + extra_args={"robot_obs": {"state": np.zeros(config.max_state_dim)}}, + ), + request_id="crop-test", + ) + + result = pipeline.forward(request) + + assert result.output["actions"].shape == (config.chunk_size, 7) + + +def test_load_lerobot_norm_stats_unknown_mode_raises(tmp_path): + """Fail loud: an unreproducible mode served anyway is silently wrong.""" + path = _write_lerobot_checkpoint( + tmp_path, norm_map={"STATE": "SOMETHING_NEW", "ACTION": "QUANTILES"}, stats=_FULL_DATASET_STATS + ) + with pytest.raises(ValueError, match="SOMETHING_NEW"): + load_lerobot_norm_stats(path) + + +def test_from_pretrained_backfills_sidecar_stats(tmp_path): + """End to end: the stats reach norm_stats without any deploy yaml.""" + path = _write_lerobot_checkpoint(tmp_path, stats=_FULL_DATASET_STATS) + config = Pi05Config.from_pretrained(path) + + assert config.norm_stats["state"]["mode"] == "quantile" + assert config.norm_stats["action"]["mode"] == "quantile" + + +def test_from_pretrained_does_not_override_explicit_norm_stats(tmp_path): + """config.json wins over the sidecar, matching LeRobot's override rule.""" + import json as _json + + path = _write_lerobot_checkpoint(tmp_path, stats=_FULL_DATASET_STATS) + explicit = {"state": {"mode": "min_max", "min": [0.0] * 32, "max": [1.0] * 32}} + (tmp_path / "config.json").write_text(_json.dumps({**_LEROBOT_CFG, "norm_stats": explicit}), encoding="utf-8") + + assert Pi05Config.from_pretrained(path).norm_stats == explicit + + +# ---------------------------------------------------------------------------- +# Checkpoint boundary: declared-but-unconsumed capabilities must raise +# ---------------------------------------------------------------------------- +@pytest.mark.parametrize( + "key,value", + [ + ("use_visual_memory", True), + ("use_proprioceptive_memory", True), + ("rtc_config", {"mode": "trained", "inference_delay": 2}), + ("n_obs_steps", 6), + ], +) +def test_config_rejects_unsupported_capability(key, value): + """MEM and RTC change what a correct action chunk looks like and are not + visible in the weights. Serving such a checkpoint must fail loudly.""" + with pytest.raises(UnsupportedCheckpointCapabilityError, match=key): + Pi05Config.from_model_config(dict(_LEROBOT_CFG, **{key: value})) + + +# ---------------------------------------------------------------------------- +# Relative actions +# ---------------------------------------------------------------------------- +def test_relative_actions_without_action_names_raises(): + """``action_feature_names`` is populated from dataset metadata at training + time. There is no dataset at serving time, so a relative-action checkpoint + that omits it cannot resolve ``relative_exclude_joints`` and must fail.""" + with pytest.raises(UnsupportedCheckpointCapabilityError, match="action_feature_names"): + Pi05Config.from_model_config(dict(_LEROBOT_CFG, use_relative_actions=True, action_feature_names=None)) + + +def test_relative_actions_empty_exclude_list_needs_no_action_names(): + c = Pi05Config.from_model_config(dict(_LEROBOT_CFG, use_relative_actions=True, relative_exclude_joints=[])) + assert c.use_relative_actions + + +def test_resolve_excluded_action_indices_unresolvable_raises(): + with pytest.raises(UnsupportedCheckpointCapabilityError): + resolve_excluded_action_indices(["no_such_joint"], _ACTION_NAMES) + + +def _relative_step(enabled=True): + return Pi05RelativeActions( + enabled=enabled, + exclude_joints=["gripper"], + action_names=_ACTION_NAMES, + max_action_dim=32, + ) + + +def test_relative_mask_excludes_gripper_and_padding(): + rel = _relative_step() + assert rel.relative_mask[0] + assert not rel.relative_mask[6], "gripper must stay absolute" + assert not rel.relative_mask[len(_ACTION_NAMES) :].any(), "padded dims carry no signal" + assert rel.num_relative_dims == 6 + + +def test_relative_absolute_round_trip_is_exact(): + """The two directions are one object precisely so they cannot disagree.""" + rel = _relative_step() + g = torch.Generator().manual_seed(0) + actions = torch.randn(1, 50, 32, generator=g) + state = torch.randn(32, generator=g) + assert torch.allclose(rel.to_absolute(rel.to_relative(actions, state), state), actions, atol=1e-6) + + +def test_relative_rejects_mismatched_action_dim(): + rel = _relative_step() + with pytest.raises(ValueError, match="max_action_dim"): + rel.to_absolute(torch.zeros(1, 4, 8), torch.zeros(32)) + + +# ---------------------------------------------------------------------------- +# State discretization — π0.5's defining input path +# ---------------------------------------------------------------------------- +def test_discretize_underflows_to_minus_one_and_saturates_at_the_top(): + """Out-of-range input is asymmetric, and that asymmetry is LeRobot's. + + ``np.digitize`` returns 0 for anything below the first bin edge, so the + ``- 1`` makes a below-range state land in bin ``-1``; above-range input + saturates at ``num_bins - 1`` on its own. The checkpoint was trained with + ``" -1"`` in the state prompt for under-range dims, so clipping that to 0 + would change the tokens the model sees — LeRobot parity catches it. + """ + out = discretize_state(np.array([-9.0, 9.0], dtype=np.float32), num_bins=256) + assert out.tolist() == [-1, 255] + + +def test_normalize_state_does_not_clip_out_of_range(): + """The parity-critical contract: normalization is affine and nothing else. + + Clamping to [-1, 1] here would hide an under-range dimension from + ``discretize_state``, turning LeRobot's ``-1`` bin into ``0`` and changing + the state tokens in the prompt. LeRobot applies + ``2.0 * (x - q01) / (q99 - q01) - 1.0`` with no clip. + """ + stats = {"q01": [0.0] * 32, "q99": [10.0] * 32} + out = _normalized([-10.0, 20.0] + [5.0] * 30, stats) + assert out[0] == pytest.approx(-3.0) # would be -1.0 if clipped + assert out[1] == pytest.approx(3.0) # would be 1.0 if clipped + # and the under-range dim must reach the -1 bin end to end + assert discretize_state(out, num_bins=256)[0] == -1 + + +def test_prompt_serializes_the_declared_state_width(): + """LeRobot's tokenizer step discretizes the state at its real width and + never pads to ``max_state_dim``.""" + prompt = build_pi05_prompt(task="x", normalized_state=np.zeros(7, dtype=np.float32)) + bins = prompt.split("State: ")[1].split(";")[0].split() + assert len(bins) == 7 + + +def test_prompt_template_matches_lerobot(): + prompt = build_pi05_prompt( + task=" pick_up the red\nblock ", + normalized_state=np.zeros(4, dtype=np.float32), + ) + assert prompt.startswith("Task: pick up the red block, State: ") + assert prompt.endswith(";\nAction: ") + + +# ---------------------------------------------------------------------------- +# Input contract: fixed camera slots x 256 tokens + constant 200 text tokens +# ---------------------------------------------------------------------------- +@pytest.mark.parametrize("views,expected", [(1, 456), (2, 712), (3, 968)]) +def test_prefix_length_is_fixed_by_max_cameras(views, expected): + """``max_cameras × 256`` image tokens plus a constant ``tokenizer_max_length``. + + The text segment is padded to its maximum and the camera slots are padded to + ``max_cameras``, so the prefix shape does not vary with the request — only + the attention mask does. + """ + keys = _EXPECTED_CAMERA_ORDER[:views] + config = Pi05Config(max_cameras=views, tokenizer_max_length=200, image_feature_keys=keys) + observation = { + "state": np.zeros(32, dtype=np.float32), + "images": {key: np.zeros((4, 4, 3), dtype=np.uint8) for key in keys}, + } + + images, masks, lang_tokens, _ = _assemble_model_inputs( + observation, config, _FakeTokenizer(), torch.device("cpu"), None + ) + + assert len(images) == views + assert all(bool(mask) for mask in masks) + assert lang_tokens.shape[-1] == 200 + assert 256 * len(images) + lang_tokens.shape[-1] == expected + + +# ---------------------------------------------------------------------------- +# Shared math helpers (same contracts as π0) +# ---------------------------------------------------------------------------- +def test_gemma_variant_dims(): + vlm = get_gemma_config("gemma_2b") + assert (vlm.width, vlm.depth, vlm.mlp_dim, vlm.num_kv_heads, vlm.head_dim) == (2048, 18, 16384, 1, 256) + expert = get_gemma_config("gemma_300m") + assert (expert.width, expert.depth, expert.mlp_dim) == (1024, 18, 4096) + + +def test_sinusoidal_embedding_shape_and_oddity(): + out = create_sinusoidal_pos_embedding(torch.tensor([0.0, 1.0]), 64) + assert out.shape == (2, 64) + with pytest.raises(ValueError): + create_sinusoidal_pos_embedding(torch.tensor([0.0]), 63) + + +def test_make_att_2d_masks_respects_padding(): + pad = torch.tensor([[True, True, False]]) + att = torch.tensor([[0, 0, 0]]) + out = make_att_2d_masks(pad, att) + assert not out[0, :, 2].any() + + +def test_build_norm_stats_rejects_unknown_mode(): + with pytest.raises(ValueError, match="Unsupported normalization mode"): + build_norm_stats({"action": {"mode": "qunatile"}}, "action") + + +@pytest.mark.parametrize( + "entry", + [ + {"mode": "mean_std", "mean": [0.0]}, + {"mode": "min_max", "min": [0.0]}, + {"mode": "quantile", "q01": [0.0]}, + ], +) +def test_build_norm_stats_rejects_incomplete_stats(entry): + with pytest.raises(ValueError, match="requires"): + build_norm_stats({"action": entry}, "action") + + +def test_normalize_quantile_round_trip_and_padded_tail(): + stats = build_norm_stats({"action": {"q01": [0.0, 0.0], "q99": [2.0, 4.0]}}, "action") + x = torch.tensor([[0.5, -0.5, 7.0]]) # third entry is padding beyond the stats + back = apply_norm(apply_norm(x, stats), stats, inverse=True) + assert torch.allclose(back, x, atol=1e-5) + assert back[0, 2] == 7.0, "padded tail must pass through untouched" + + +# ---------------------------------------------------------------------------- +# Image processor +# ---------------------------------------------------------------------------- +@pytest.mark.parametrize("value,expected", [(0.0, -1.0), (0.25, -0.5), (1.0, 1.0)]) +@pytest.mark.parametrize("container", ["numpy", "torch"]) +def test_float_image_uses_lerobot_zero_one_domain(value, expected, container): + if container == "numpy": + image = np.full((4, 4, 3), value, dtype=np.float32) + else: + image = torch.full((3, 4, 4), value, dtype=torch.float32) + + out = Pi05ImageProcessor(image_size=4).preprocess_single(image) + + assert out.shape == (1, 3, 4, 4) + assert torch.allclose(out, torch.full_like(out, expected)) + + +@pytest.mark.parametrize("bad", [-0.01, 1.01, np.nan, np.inf]) +@pytest.mark.parametrize("container", ["numpy", "torch"]) +def test_float_image_rejects_values_outside_lerobot_domain(bad, container): + if container == "numpy": + image = np.full((4, 4, 3), bad, dtype=np.float32) + else: + image = torch.full((3, 4, 4), bad, dtype=torch.float32) + + with pytest.raises(ValueError): + Pi05ImageProcessor(image_size=4).preprocess_single(image) + + +def test_resize_with_pad_pads_with_minus_one(): + out = resize_with_pad(torch.zeros(1, 3, 100, 200), 224, 224) + assert out.shape == (1, 3, 224, 224) + assert out[0, 0, 0, 0] == -1.0 + + +class _FakeTokenizer: + def __call__(self, text, **kwargs): + del text + length = kwargs["max_length"] + return {"input_ids": [0] * length, "attention_mask": [1] * length} + + +@pytest.mark.parametrize("explicit_none", [False, True]) +def test_build_model_inputs_preserves_missing_middle_camera_slot(explicit_none): + config = Pi05Config(image_feature_keys=_EXPECTED_CAMERA_ORDER, max_cameras=3) + images = { + _EXPECTED_CAMERA_ORDER[0]: np.zeros((4, 4, 3), dtype=np.uint8), + _EXPECTED_CAMERA_ORDER[2]: np.full((4, 4, 3), 255, dtype=np.uint8), + } + if explicit_none: + images[_EXPECTED_CAMERA_ORDER[1]] = None + observation = { + "images": images, + "state": np.zeros(32, dtype=np.float32), + } + + slots, masks, _, _ = _assemble_model_inputs(observation, config, _FakeTokenizer(), torch.device("cpu"), None) + + assert [bool(mask.item()) for mask in masks] == [True, False, True] + torch.testing.assert_close(slots[0], torch.full_like(slots[0], -1.0)) + torch.testing.assert_close(slots[1], torch.full_like(slots[1], -1.0)) + torch.testing.assert_close(slots[2], torch.full_like(slots[2], 1.0)) + + +def test_build_model_inputs_requires_configured_camera(): + config = Pi05Config(image_feature_keys=_EXPECTED_CAMERA_ORDER) + observation = { + "unknown_camera": np.zeros((4, 4, 3), dtype=np.uint8), + "state": np.zeros(32, dtype=np.float32), + } + with pytest.raises(ValueError, match="configured camera"): + _assemble_model_inputs(observation, config, _FakeTokenizer(), torch.device("cpu"), None) + + +# ---------------------------------------------------------------------------- +# Model behaviour + weight-load remap (tiny action dims; full Gemma backbone) +# ---------------------------------------------------------------------------- +# Module *structure* (no ``state_proj``, ``time_mlp_{in,out}`` widths, AdaRMS on +# the expert norms but not the prefix) is not asserted here: ``load_weights`` +# audits every model parameter against the checkpoint and raises on a missing, +# unmatched or π0-shaped key, so a structural drift fails loudly at load time +# rather than silently. ``test_pi0_shaped_keys_are_not_silently_loaded`` guards +# that audit, and ``test_pi05_parity.py`` is the numerical oracle. +# Both Gemma towers keep their real widths but drop from 18 layers to 2 for +# every test in this section. Nothing here depends on depth: the weight remap +# matches by name, and the AR mask, prefix length and ODE determinism are +# shape-driven. Full-scale numerical correctness is ``test_pi05_parity.py``'s +# job, against the real checkpoint. +# +# Two dimensions must NOT be shrunk, both because something outside these +# configs is sized against them: +# +# * ``width`` — the vision tower's ``multi_modal_projector`` projects image +# embeddings to the PaliGemma text width so the two can be concatenated. +# Narrowing the text tower alone makes that concat fail. +# * the vision tower itself — So400m/14 on 224x224 is what yields exactly 256 +# tokens per view, which ``test_prefix_length_and_valid_len_for_each_view_count`` +# exists to verify. +# +# Depths are equal across the two towers on purpose: the action expert runs +# layer-by-layer alongside the VLM, so a mismatch would not be a valid model. +_SHRUNK_GEMMA = { + "gemma_2b": GemmaVariantConfig(2048, 2, 1024, 8, 1, 256), + "gemma_300m": GemmaVariantConfig(1024, 2, 512, 8, 1, 256), +} + + +@pytest.fixture(scope="module", autouse=True) +def _shrink_gemma_backbones(): + """Patch the module-level lookup ``Pi05ForActionPrediction.__init__`` calls.""" + patch = pytest.MonkeyPatch() + patch.setattr(modeling_pi05, "get_gemma_config", lambda variant: _SHRUNK_GEMMA[variant]) + yield + patch.undo() + + +def _tiny_pi05_model(): + """A fresh model. Use for tests that load weights or otherwise mutate it.""" + return Pi05ForActionPrediction(Pi05Config(max_action_dim=8, max_state_dim=8, chunk_size=4, n_action_steps=4)) + + +@pytest.fixture(scope="module") +def tiny_model(): + """Shared read-only model, so the tests that only inspect it pay for one + construction between them rather than one each.""" + return _tiny_pi05_model() + + +@pytest.fixture(autouse=True) +def _silence_pi05_loader(): + logging.getLogger("vllm_omni.diffusion.models.pi05.modeling_pi05").setLevel(logging.CRITICAL) + yield + + +@pytest.mark.slow +def test_suffix_has_no_state_token(tiny_model): + """π0's suffix is ``[state, actions...]`` with AR mask ``[1, 1, 0...]``; + π0.5's is actions only with ``[1, 0...]``.""" + embs, pad, att, cond = tiny_model.embed_suffix(torch.randn(1, 4, 8), torch.tensor([1.0])) + assert embs.shape[1] == 4 + assert att[0].tolist() == [1, 0, 0, 0] + assert cond.shape == (1, tiny_model.expert_width) + + +# Not marked slow: constructs one small norm layer, no backbone. +def test_adarms_norm_returns_gate_only_when_conditioned(): + norm = Pi05AdaRMSNorm(8, cond_dim=4) + out, gate = norm(torch.randn(1, 3, 8), torch.randn(1, 4)) + assert gate is not None and gate.shape == (1, 1, 8) + + plain = Pi05AdaRMSNorm(8, cond_dim=None) + out, gate = plain(torch.randn(1, 3, 8)) + assert gate is None + + +@pytest.mark.slow +def test_lm_head_remap_to_embed_tokens(): + model = _tiny_pi05_model() + target = "paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight" + payload = torch.full(dict(model.named_parameters())[target].shape, 0.5) + model.load_weights([("model.paligemma_with_expert.paligemma.lm_head.weight", payload)], strict=False) + assert torch.allclose(dict(model.named_parameters())[target], payload) + + +@pytest.mark.slow +def test_pi0_shaped_keys_are_not_silently_loaded(): + """A π0 checkpoint pointed at the π0.5 class would leave the AdaRMS expert + randomly initialized. ``state_proj`` is the tell.""" + model = _tiny_pi05_model() + with pytest.raises(RuntimeError, match="π0-shaped"): + model.load_weights([("model.state_proj.weight", torch.zeros(1024, 8))]) + + +@pytest.mark.slow +@pytest.mark.parametrize("num_real_views", [1, 2, 3]) +def test_prefix_length_is_fixed_and_valid_len_tracks_masks(num_real_views, tiny_model): + """The deployed model has three slots; only their validity varies.""" + live_text = 120 + images = [torch.zeros(1, 3, 224, 224) for _ in range(3)] + masks = [torch.tensor([index < num_real_views]) for index in range(3)] + lang = torch.zeros(1, 200, dtype=torch.long) + lang_mask = torch.zeros(1, 200, dtype=torch.bool) + lang_mask[:, :live_text] = True + + embs, pad_masks, _ = tiny_model.embed_prefix(images, masks, lang, lang_mask) + assert embs.shape[1] == 256 * 3 + 200 + assert int(pad_masks.sum()) == 256 * num_real_views + live_text + + +@pytest.mark.slow +def test_sample_actions_shape_and_determinism(tiny_model): + """Flow matching is an ODE: fixed noise must give a bit-identical chunk.""" + model = tiny_model.eval() + images = [torch.zeros(1, 3, 224, 224) for _ in range(3)] + masks = [torch.tensor([True]), torch.tensor([False]), torch.tensor([False])] + lang = torch.zeros(1, 200, dtype=torch.long) + lang_mask = torch.ones(1, 200, dtype=torch.bool) + noise = torch.randn(1, 4, 8, generator=torch.Generator().manual_seed(42)) + + with torch.no_grad(): + a1 = model.sample_actions( + images=images, image_masks=masks, lang_tokens=lang, lang_masks=lang_mask, noise=noise, num_steps=2 + ) + a2 = model.sample_actions( + images=images, image_masks=masks, lang_tokens=lang, lang_masks=lang_mask, noise=noise, num_steps=2 + ) + assert a1.shape == (1, 4, 8) + assert torch.isfinite(a1).all() + assert torch.equal(a1, a2) + + +@pytest.mark.slow +def test_sample_actions_uses_request_generator(tiny_model): + model = tiny_model.eval() + images = [torch.zeros(1, 3, 224, 224) for _ in range(3)] + masks = [torch.tensor([True]), torch.tensor([False]), torch.tensor([False])] + lang = torch.zeros(1, 200, dtype=torch.long) + lang_mask = torch.ones(1, 200, dtype=torch.bool) + + def run(seed): + return model.sample_actions( + images=images, + image_masks=masks, + lang_tokens=lang, + lang_masks=lang_mask, + num_steps=2, + generator=torch.Generator().manual_seed(seed), + ) + + with torch.no_grad(): + first = run(42) + repeated = run(42) + different = run(43) + + assert torch.equal(first, repeated) + assert not torch.equal(first, different) + + +# ---------------------------------------------------------------------------- +# Serving dtype +# ---------------------------------------------------------------------------- +# The check that decides what runs lives on the pipeline, not on Pi05Config: +# Pi05Config.dtype records what the checkpoint declares and never reaches a +# cast, while _initialize_model casts using the top-level OmniDiffusionConfig. +@pytest.mark.parametrize( + "declared,expected", + [ + ("float32", torch.float32), + ("bfloat16", torch.bfloat16), + (torch.float32, torch.float32), + (torch.bfloat16, torch.bfloat16), + ], +) +def test_resolve_dtype_accepts_supported_serving_dtypes(declared, expected): + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.diffusion.models.pi05.pipeline_pi05 import Pi05Pipeline + + assert Pi05Pipeline._resolve_dtype(OmniDiffusionConfig(dtype=declared)) is expected + + +@pytest.mark.parametrize("declared", ["float16", "fp16", torch.float16]) +def test_resolve_dtype_rejects_unvalidated_dtypes(declared): + """float16 resolves to a real torch dtype, so a permissive lookup would run + the model in a precision nothing here has been checked in. + + Only dtypes that survive ``OmniDiffusionConfig`` normalization are worth + asserting on: it maps a string it does not recognize onto bfloat16 with a + warning, so e.g. ``"float64"`` never reaches this guard as float64. + """ + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.diffusion.models.pi05.pipeline_pi05 import Pi05Pipeline + + with pytest.raises(ValueError, match="Unsupported π0.5 dtype"): + Pi05Pipeline._resolve_dtype(OmniDiffusionConfig(dtype=declared)) + + +def test_bfloat16_layout_is_applied_before_checkpoint_load(monkeypatch, tmp_path): + """Checkpoint tensors must be copied directly into their final dtypes. + + LeRobot casts only the inner PaliGemma/expert module; the outer action and + time projections remain float32. The layout is established before loading + so float32-retained weights never pass through bfloat16. + """ + from vllm_omni.diffusion.models.pi05 import pipeline_pi05 + + class _TinyPi05(torch.nn.Module): + def __init__(self): + super().__init__() + self.paligemma_with_expert = torch.nn.Module() + self.paligemma_with_expert.backbone = torch.nn.Linear(2, 2) + self.paligemma_with_expert.vision_tower = torch.nn.Linear(2, 2) + self.paligemma_with_expert.input_layernorm = torch.nn.LayerNorm(2) + self.paligemma_with_expert.register_buffer("floating_buffer", torch.ones(1)) + self.action_in_proj = torch.nn.Linear(2, 2) + self.action_out_proj = torch.nn.Linear(2, 2) + self.time_mlp_in = torch.nn.Linear(2, 2) + self.time_mlp_out = torch.nn.Linear(2, 2) + + # DiffusersLoader constructs custom pipelines under the requested default + # dtype. Start fully in bfloat16 to verify that the helper restores every + # float32-retained tensor rather than relying on the construction default. + model = _TinyPi05().to(dtype=torch.bfloat16) + observed_dtypes = {} + + def fake_load_checkpoint(_self, model_to_load): + observed_dtypes.update({name: param.dtype for name, param in model_to_load.named_parameters()}) + + (tmp_path / "model.safetensors").touch() + pipeline = object.__new__(pipeline_pi05.Pi05Pipeline) + pipeline.model_dir = str(tmp_path) + pipeline.config = object() + pipeline._device = torch.device("cpu") + pipeline._torch_dtype = torch.bfloat16 + monkeypatch.setattr(pipeline_pi05, "Pi05ForActionPrediction", lambda _config: model) + monkeypatch.setattr(pipeline_pi05.Pi05Pipeline, "_load_checkpoint", fake_load_checkpoint) + + initialized = pipeline._initialize_model() + + assert observed_dtypes["paligemma_with_expert.backbone.weight"] is torch.bfloat16 + assert model.paligemma_with_expert.floating_buffer.dtype is torch.bfloat16 + for name in ( + "paligemma_with_expert.vision_tower.weight", + "paligemma_with_expert.input_layernorm.weight", + "action_in_proj.weight", + "action_out_proj.weight", + "time_mlp_in.weight", + "time_mlp_out.weight", + ): + assert observed_dtypes[name] is torch.float32 + assert initialized is model + assert not initialized.training + + +@pytest.mark.slow +def test_bfloat16_runs_and_tracks_float32(tiny_model): + """bfloat16 is a supported serving dtype, so it has to produce a usable + chunk — and one that still resembles float32. + + The bound is deliberately loose. This is not a parity check (bfloat16 + deviates from float32 by ~1-2% of the action range on the real checkpoint, + which is why recipes/lerobot/Pi05.md tells operators to validate it on their + own task). What it catches is gross breakage: a cast that silently did not + happen, activations overflowing to inf, or a dtype mismatch that would make + the two outputs unrelated. + """ + import copy + + images = [torch.zeros(1, 3, 224, 224) for _ in range(3)] + masks = [torch.tensor([True]), torch.tensor([False]), torch.tensor([False])] + lang = torch.zeros(1, 200, dtype=torch.long) + lang_mask = torch.ones(1, 200, dtype=torch.bool) + noise = torch.randn(1, 4, 8, generator=torch.Generator().manual_seed(7)) + + def run(model, dtype): + with torch.no_grad(): + return model.sample_actions( + images=[img.to(dtype) for img in images], + image_masks=masks, + lang_tokens=lang, + lang_masks=lang_mask, + noise=noise.to(dtype), + num_steps=2, + ).float() + + reference = run(tiny_model.eval(), torch.float32) + from vllm_omni.diffusion.models.pi05.pipeline_pi05 import _set_inference_dtype + + bf16_model = copy.deepcopy(tiny_model).eval() + _set_inference_dtype(bf16_model, torch.bfloat16) + + # The cast actually happened — otherwise everything below passes trivially. + assert any(param.dtype is torch.bfloat16 for param in bf16_model.parameters()) + assert bf16_model.action_in_proj.weight.dtype is torch.float32 + assert bf16_model.action_out_proj.weight.dtype is torch.float32 + assert bf16_model.time_mlp_in.weight.dtype is torch.float32 + assert bf16_model.time_mlp_out.weight.dtype is torch.float32 + + actual = run(bf16_model, torch.bfloat16) + assert actual.shape == reference.shape + assert torch.isfinite(actual).all() + + scale = float(reference.abs().max()) + assert float((actual - reference).abs().max()) < max(0.5 * scale, 1e-2) + + +# ---------------------------------------------------------------------------- +# State normalization runs π0's stats path +# ---------------------------------------------------------------------------- +@pytest.mark.parametrize( + "stats,state,expected", + [ + # A denominator that is small but not zero. ``normalize`` divides by + # ``std + eps``; replacing a small std with 1.0 would be a different map. + ({"mode": "mean_std", "mean": [0.0], "std": [1e-7]}, [1e-7], 1e-7 / (1e-7 + 1e-8)), + ({"mode": "quantile", "q01": [0.0], "q99": [1e-7]}, [1e-7], 1.0), + ({"mode": "quantile", "q01": [0.0], "q99": [1e-9]}, [1e-9], 1.0), + # An exactly zero range maps the bound to -1 instead of dividing by zero. + ({"mode": "quantile", "q01": [2.0], "q99": [2.0]}, [2.0], -1.0), + ], +) +def test_normalize_state_uses_the_shared_denominator_rule(stats, state, expected): + """State normalization and action unnormalization are the same affine map, + so a narrow range keeps producing LeRobot's answer.""" + assert _normalized(state, stats)[0] == pytest.approx(expected, rel=1e-5) + + +# ---------------------------------------------------------------------------- +# State wire contract +# ---------------------------------------------------------------------------- +@pytest.mark.parametrize("bad", [float("nan"), float("inf"), float("-inf")]) +def test_state_rejects_non_finite_values(bad): + """``np.digitize`` maps nan and inf to the top bin and -inf to -1, all legal + state tokens, so a dead sensor would reach the model as an ordinary prompt.""" + with pytest.raises(ValueError, match="NaN or Inf"): + as_state_vector([0.0, bad], 2) + + +def test_state_rejects_a_multi_row_matrix(): + """Flattening would read a (2, 16) request as a 32-value state.""" + with pytest.raises(ValueError, match=r"\(D,\) or \(1, D\)"): + as_state_vector(np.zeros((2, 16), dtype=np.float32), 32) + + +# ---------------------------------------------------------------------------- +# Per-request denoising steps reach the model +# ---------------------------------------------------------------------------- +class _SpyModel: + """Records the ``num_steps`` ``forward`` hands to ``sample_actions``.""" + + def __init__(self): + self.seen: list = [] + self.generators: list = [] + + def sample_actions(self, *, images, image_masks, lang_tokens, lang_masks, num_steps, generator): + del images, image_masks, lang_tokens, lang_masks + self.seen.append(num_steps) + self.generators.append(generator) + return torch.zeros(1, 50, 32) + + +def _spy_pipeline(spy: _SpyModel): + from vllm_omni.diffusion.models.pi05.pipeline_pi05 import Pi05Pipeline + + pipeline = object.__new__(Pi05Pipeline) + pipeline.config = Pi05Config.from_model_config(_LEROBOT_CFG) + pipeline.tokenizer = _FakeTokenizer() + pipeline._device = torch.device("cpu") + pipeline.model = spy + pipeline.processor = Pi05Processor(pipeline.config, _FakeTokenizer(), torch.device("cpu")) + return pipeline + + +def _spy_request(num_inference_steps, generator=None): + params = type("_Params", (), {})() + params.extra_args = { + "robot_obs": { + "state": np.zeros(32, dtype=np.float32), + "images": {key: np.zeros((4, 4, 3), dtype=np.uint8) for key in _EXPECTED_CAMERA_ORDER}, + } + } + params.num_inference_steps = num_inference_steps + params.generator = generator + request = type("_Request", (), {})() + request.sampling_params = params + return request + + +@pytest.mark.parametrize("requested", [None, 2, 25]) +def test_request_denoising_steps_reach_sample_actions(requested): + """The OpenPI E2E cannot see this: a pipeline that ignored the override and + ran the static default would still return a finite chunk of the right shape.""" + spy = _SpyModel() + output = _spy_pipeline(spy).forward(_spy_request(requested)) + + assert spy.seen == [requested] + assert output.output["actions"].shape == (50, 32) + + +def test_request_generator_reaches_sample_actions(): + generator = torch.Generator().manual_seed(42) + spy = _SpyModel() + + _spy_pipeline(spy).forward(_spy_request(2, generator)) + + assert spy.generators == [generator] + + +@pytest.mark.parametrize("bad", [0, -1, 2.5, True, "4"]) +def test_request_denoising_steps_must_be_a_positive_integer(bad): + spy = _SpyModel() + with pytest.raises(ValueError, match="num_inference_steps must be a positive integer"): + _spy_pipeline(spy).forward(_spy_request(bad)) + assert spy.seen == [] diff --git a/tests/e2e/online_serving/test_pi05_expansion.py b/tests/e2e/online_serving/test_pi05_expansion.py new file mode 100644 index 000000000..0c6fbfd34 --- /dev/null +++ b/tests/e2e/online_serving/test_pi05_expansion.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""E2E online serving test for π0.5 OpenPI websocket serving. + +Boots ``vllm serve --omni --deploy-config pi05.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_pi0_expansion.py``. Needs one H100 and the full +π0.5 checkpoint. + +The ``pi0_openpi_*`` helpers are protocol-level (handshake, send obs, receive +chunk) and not π0-specific, so π0.5 reuses them; only the deploy config and the +expected metadata differ. This file is selected explicitly by the nightly H100 +robot-policy job. + +The in-process LeRobot parity oracle lives separately in +``tests/diffusion/models/pi05/test_pi05_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/pi05_base" + + +test_params = [ + OmniServerParams( + model=MODEL, + port=8093, + server_args=[ + "--deploy-config", + "vllm_omni/deploy/pi05.yaml", + "--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_pi05_openpi_online(omni_server): + pi0_openpi_require_dependencies() + + 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="pi05-online-e2e", + num_steps=2, + num_inference_steps=2, + ) + + # Asserts every returned chunk is [50, 32] + finite, and the handshake + # metadata matches pi05.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/entrypoints/openai_api/test_openpi_serving.py b/tests/entrypoints/openai_api/test_openpi_serving.py index d274eadda..fa0148599 100644 --- a/tests/entrypoints/openai_api/test_openpi_serving.py +++ b/tests/entrypoints/openai_api/test_openpi_serving.py @@ -270,12 +270,14 @@ def test_build_request_uses_unique_engine_request_id_per_inference(): def test_build_request_clones_stage_defaults_before_protocol_fields(): default_params = OmniDiffusionSamplingParams( + num_inference_steps=12, + output_type="latent", extra_args={ "format_prompt_as_json": True, "session_id": "configured-session-must-not-win", "reset": False, "nested": {"values": []}, - } + }, ) od_config = SimpleNamespace(model_config={"policy_server_config": {}}) engine_client = SimpleNamespace( @@ -297,10 +299,18 @@ def test_build_request_clones_stage_defaults_before_protocol_fields(): session_id="wire-session", reset=False, ) + request_c = serving._build_request( + {"prompt": "pick up the object", "sampling_params": {"num_inference_steps": 4}}, + session_id="wire-session", + reset=False, + ) assert request_a.sampling_params.extra_args["format_prompt_as_json"] is True assert request_a.sampling_params.extra_args["session_id"] == "wire-session" assert request_a.sampling_params.extra_args["reset"] is True + assert request_a.sampling_params.num_inference_steps == 12 + assert request_a.sampling_params.output_type == "latent" + assert request_c.sampling_params.num_inference_steps == 4 request_a.sampling_params.extra_args["nested"]["values"].append("request-a") assert request_b.sampling_params.extra_args["nested"] == {"values": []} assert default_params.extra_args["nested"] == {"values": []} @@ -318,6 +328,46 @@ def test_build_request_forwards_seed_to_sampling_params(): assert isinstance(unseeded.sampling_params.seed, int) +# Per-request inference parameters live in their own namespace so robot feature +# names stay open-ended. Unknown keys inside it are ignored rather than rejected. +def test_build_request_forwards_num_inference_steps(): + serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=_engine_with_policy_config()) + + request = serving._build_request( + {"prompt": "pick up the object", "sampling_params": {"num_inference_steps": 4}}, + session_id="session-a", + reset=True, + ) + + assert request.sampling_params.num_inference_steps == 4 + # The namespace is an engine concern; the pipeline must not see it as a feature. + assert "sampling_params" not in request.sampling_params.extra_args["robot_obs"] + + +def test_build_request_omits_num_inference_steps_when_not_sent(): + """Absent must stay absent: the pipeline reads None as "use the configured + default", so materialising a null here would be a different request.""" + serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=_engine_with_policy_config()) + + request = serving._build_request({"prompt": "x"}, session_id="s", reset=True) + + assert request.sampling_params.num_inference_steps is None + + +def test_build_request_does_not_forward_arbitrary_observation_keys(): + """A robot observation may carry any feature name; none of them become + engine parameters.""" + serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=_engine_with_policy_config()) + + request = serving._build_request( + {"prompt": "x", "state": np.zeros(32), "guidance_scale": 7.5}, session_id="s", reset=True + ) + + extra_args = request.sampling_params.extra_args + assert set(extra_args) == {"reset", "session_id", "robot_obs"} + assert extra_args["robot_obs"]["guidance_scale"] == 7.5 + + def test_infer_keeps_session_state_but_uses_unique_engine_request_ids(): engine = RecordingEngine() serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=engine) diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 0632db55d..975a6ddd2 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -1166,6 +1166,7 @@ def pi0_openpi_run_policy_session( prompt: str = PI0_OPENPI_DEFAULT_PROMPT, session_id: str | None = None, num_steps: int = 2, + num_inference_steps: int | None = None, ) -> dict[str, Any]: """Connect, read handshake metadata, send ``num_steps`` observations.""" import uuid @@ -1184,6 +1185,8 @@ def pi0_openpi_run_policy_session( actions = [] for _ in range(num_steps): payload = pi0_make_dummy_obs(prompt=prompt, session_id=session_id) + if num_inference_steps is not None: + payload["sampling_params"] = {"num_inference_steps": num_inference_steps} payload["endpoint"] = "infer" conn.send(packer.pack(payload)) actions.append(_pi0_decode_action_response(conn.recv())) diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index f0a1285bc..ed266d734 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -82,6 +82,7 @@ EXCLUDED_MODELS = [ "StableDiffusionXLPipeline", "Gr00tN1d7Pipeline", "Pi0Pipeline", + "Pi05Pipeline", "SanaWmPipeline", # Audio pipeline conditioned by a separate Qwen2.5-Omni encoder stage; # covered by tests/diffusion/models/auk (CPU pipeline tests, parity) and diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 92aad9f4d..eada1019d 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -33,6 +33,7 @@ from vllm_omni.config.stage_config import ( PipelineConfig, ) from vllm_omni.diffusion.models.pi0_pipeline_config import PI0_PIPELINE +from vllm_omni.diffusion.models.pi05_pipeline_config import PI05_PIPELINE from vllm_omni.model_executor.models.audex.pipeline import ( AUDEX_S2S_PIPELINE, AUDEX_THINKER_ONLY_PIPELINE, @@ -153,6 +154,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "lingbot_world": LINGBOT_WORLD_PIPELINE, "Gr00tN1d7": GR00T_N1D7_PIPELINE, "pi0": PI0_PIPELINE, + "pi05": PI05_PIPELINE, # Cosmos3 policy checkpoints share HF metadata with the T2I/video Cosmos3 # checkpoints (which stay on the single-stage diffusion fallback), so this # entry is only reachable through a deploy yaml's ``pipeline:`` key diff --git a/vllm_omni/deploy/pi05.yaml b/vllm_omni/deploy/pi05.yaml new file mode 100644 index 000000000..aec8e88f0 --- /dev/null +++ b/vllm_omni/deploy/pi05.yaml @@ -0,0 +1,51 @@ +# π0.5 VLA: one diffusion stage, one GPU, eager, bsz 1 (OpenPI policy serving is +# one observation -> one action reply). Every key below is documented in +# recipes/lerobot/Pi05.md. + +pipeline: pi05 +async_chunk: false +distributed_executor_backend: mp +# float32 or bfloat16. bfloat16 is faster and halves memory at a numerical cost; +# read "Choosing a dtype" in the recipe before switching. +dtype: float32 + +stages: + - stage_id: 0 + devices: "0" + max_num_seqs: 1 + enforce_eager: true + model_class_name: Pi05Pipeline + model_config: + paligemma_variant: gemma_2b + action_expert_variant: gemma_300m + chunk_size: 50 + # Must equal chunk_size: this endpoint returns the whole chunk and keeps no + # action queue, so "execute the first N steps, then replan" is not available. + n_action_steps: 50 + max_action_dim: 32 + max_state_dim: 32 + num_inference_steps: 10 + image_resolution: [224, 224] + tokenizer_max_length: 200 + state_num_bins: 256 + max_cameras: 3 + image_feature_keys: + - observation.images.base_0_rgb + - observation.images.left_wrist_0_rgb + - observation.images.right_wrist_0_rgb + image_key_map: {} + use_relative_actions: false + relative_exclude_joints: + - gripper + # 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_action_dim: 32 + max_cameras: 3 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 0ffd0c54e..6aa1a27d9 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -1615,6 +1615,11 @@ class OmniDiffusionConfig: self.model_class_name = "Pi0Pipeline" self.set_tf_model_config(TransformerConfig()) self.update_multimodal_support() + elif cfg.get("type") == "pi05": + if self.model_class_name is None: + self.model_class_name = "Pi05Pipeline" + 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/pi05/__init__.py b/vllm_omni/diffusion/models/pi05/__init__.py new file mode 100644 index 000000000..ccd19acde --- /dev/null +++ b/vllm_omni/diffusion/models/pi05/__init__.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""π0.5 VLA model for vllm-omni. + +PaliGemma (SigLIP vision + Gemma-2B LM) + Gemma-300M action expert with AdaRMS +timestep conditioning + a flow-matching action head. Outputs a continuous action +chunk ``[horizon, action_dim]`` rather than tokens. + +Differs from π0 in five places: AdaRMS timestep conditioning (``time_mlp_*`` +instead of ``action_time_mlp_*``), AdaRMS norms throughout the action expert, a +200-token tokenizer budget, a discretized state carried in the prompt (no +``state_proj``), and relative-action support. +""" + +from vllm_omni.diffusion.models.pi05.config import ( + Pi05Config, + UnsupportedCheckpointCapabilityError, + resolve_excluded_action_indices, +) +from vllm_omni.diffusion.models.pi05.modeling_pi05 import ( + GemmaVariantConfig, + PaliGemmaWithActionExpertPi05, + Pi05AdaRMSNorm, + Pi05ForActionPrediction, + create_sinusoidal_pos_embedding, + get_gemma_config, + make_att_2d_masks, + prepare_attention_masks_4d, +) +from vllm_omni.diffusion.models.pi05.pipeline_pi05 import ( + Pi05Pipeline, + get_pi05_post_process_func, +) +from vllm_omni.diffusion.models.pi05.processor_pi05 import ( + Pi05ImageProcessor, + Pi05Processor, + Pi05RelativeActions, + apply_norm, + build_norm_stats, + build_pi05_prompt, + discretize_state, + pil_image_to_tensor, + resize_with_pad, + tokenize_prompt, +) + +__all__ = [ + "Pi05Config", + "UnsupportedCheckpointCapabilityError", + "resolve_excluded_action_indices", + "Pi05ForActionPrediction", + "PaliGemmaWithActionExpertPi05", + "Pi05AdaRMSNorm", + "GemmaVariantConfig", + "get_gemma_config", + "create_sinusoidal_pos_embedding", + "make_att_2d_masks", + "prepare_attention_masks_4d", + "Pi05ImageProcessor", + "Pi05Processor", + "Pi05RelativeActions", + "build_pi05_prompt", + "discretize_state", + "apply_norm", + "build_norm_stats", + "pil_image_to_tensor", + "resize_with_pad", + "tokenize_prompt", + "Pi05Pipeline", + "get_pi05_post_process_func", +] diff --git a/vllm_omni/diffusion/models/pi05/config.py b/vllm_omni/diffusion/models/pi05/config.py new file mode 100644 index 000000000..e86ddbae9 --- /dev/null +++ b/vllm_omni/diffusion/models/pi05/config.py @@ -0,0 +1,471 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Config surface for the π0.5 VLA model in vllm-omni. + +Deliberately shaped exactly like ``pi0/config.py``: a small dataclass that +consumes the raw LeRobot ``config.json`` (the field surface of +``lerobot.policies.pi05.PI05Config``) and keeps only the runtime-relevant +fields. Transformer dimensions are derived from ``paligemma_variant`` / +``action_expert_variant`` inside the model via ``get_gemma_config``. + +What π0.5 adds on top of the π0 config surface: + +* ``tokenizer_max_length = 200`` (π0 uses 48). +* ``state_num_bins`` — state is discretized into language tokens instead of + going through a ``state_proj`` layer. +* ``use_relative_actions`` / ``relative_exclude_joints`` / + ``action_feature_names`` — the relative-action contract. See + ``processor_pi05.Pi05RelativeActions``. +* Quantile normalization stats. LeRobot's π0.5 defaults ``STATE`` and + ``ACTION`` to ``NormalizationMode.QUANTILES`` where π0 uses ``MEAN_STD``. + +**Checkpoint boundary rule.** A capability that the checkpoint *declares* but +this implementation does not *consume* must raise, not be silently dropped. +π0.5 checkpoints can declare MEM (short-horizon observation memory) and RTC +(real-time chunking); neither is supported here, and serving such a checkpoint +anyway would produce plausible-looking wrong actions. See +``_reject_unsupported_capabilities``. +""" + +from __future__ import annotations + +import json +import logging +from collections.abc import Callable +from dataclasses import dataclass, field +from dataclasses import fields as dataclass_fields +from pathlib import Path +from typing import Any + +logger = logging.getLogger(__name__) + +# Serving dtypes, by name. Kept in sync with ``Pi05Pipeline.SUPPORTED_DTYPES``, +# which guards the path that performs the actual cast. Defined here rather than +# imported to keep config.py free of a pipeline import cycle. +SUPPORTED_DTYPE_NAMES = frozenset({"float32", "bfloat16"}) + +# LeRobot / OpenPI observation key conventions. +ACTION = "action" +OBS_STR = "observation" +OBS_STATE = OBS_STR + ".state" +OBS_IMAGES = OBS_STR + ".images" + + +class UnsupportedCheckpointCapabilityError(ValueError): + """A checkpoint declares a capability this implementation does not consume. + + Raised at load time rather than silently ignored: every one of these + capabilities changes what a *correct* action chunk looks like, and none of + them is visible in the weights alone. + """ + + +def resolve_excluded_action_indices( + exclude_joints: list[str] | None, + action_names: list[str] | None, +) -> list[int]: + """Map ``relative_exclude_joints`` names onto action-vector indices. + + Matching is exact name first, then substring (a checkpoint may name the + gripper dimension ``gripper_position`` while the config just says + ``gripper``). Single source of truth for both the config-time validation + and the runtime mask in ``processor_pi05.Pi05RelativeActions``. + + Returns an empty list when there is nothing to exclude. Raises when a name + cannot be resolved — an unresolvable exclusion would otherwise silently + become "make this dimension relative too". + """ + if not [name for name in (exclude_joints or []) if name]: + return [] + if not action_names: + raise UnsupportedCheckpointCapabilityError( + f"Cannot resolve relative_exclude_joints={exclude_joints!r} without action_feature_names." + ) + + # LeRobot walks the action names and excludes each one any token equals or is + # contained in, both lowercased. Walking the tokens instead and stopping at the + # first exact hit misses a second name the token is a substring of. + tokens = [str(name).lower() for name in exclude_joints if name] + lowered = [str(name).lower() for name in action_names] + indices = [i for i, name in enumerate(lowered) if any(t == name or t in name for t in tokens)] + unresolved = [t for t in tokens if not any(t == name or t in name for name in lowered)] + + if unresolved: + raise UnsupportedCheckpointCapabilityError( + f"relative_exclude_joints entries {unresolved!r} match no entry in action_feature_names={action_names!r}." + ) + return sorted(set(indices)) + + +@dataclass +class Pi05Config: + """π0.5 VLA config (dataclass, not an HF ``PretrainedConfig``).""" + + # 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 + # The stateless OpenPI endpoint returns one complete predicted chunk. + n_action_steps: 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.5/SigLIP only support square inputs. + image_resolution: tuple[int, int] = (224, 224) + # π0.5 pads text to 200 tokens (π0 uses 48). The prompt now also carries the + # serialized state, which is why it is so much longer. + tokenizer_max_length: int = 200 + # Number of camera slots the model attends to (real + padded). + max_cameras: int = 3 + + # Bins the normalized state is discretized into before it enters the prompt. + # Ref: openpi PaliGemmaTokenizer.tokenize(). + state_num_bins: int = 256 + + # Weight dtype the checkpoint was saved in. + dtype: str = "float32" + + # ── Relative actions ────────────────────────────────────────────── + # True when the checkpoint was trained on actions relative to the current + # state. Its ``norm_stats`` are then in relative space, so serving it without + # the transform silently yields wrong actions — the weights look identical + # either way. + use_relative_actions: bool = False + # Joint names kept absolute (gripper open/close is an absolute quantity). + relative_exclude_joints: list[str] = field(default_factory=lambda: ["gripper"]) + # Needed to resolve ``relative_exclude_joints``; there is no dataset to + # fall back on at serving time. + action_feature_names: list[str] | None = None + + # Per-dataset normalization stats (schema matches LeRobot's + # ``NormalizerProcessorStep``). ``None`` means identity / pass-through. + # π0.5 defaults to quantile mode; see ``processor_pi05.build_norm_stats``. + norm_stats: dict | None = None + + # Ordered camera-slot identities. Missing keys retain their slot and are + # represented by an empty image with a false mask during preprocessing. + image_feature_keys: list[str] | None = None + # Optional map from raw OpenPI obs keys → ``image_feature_keys`` entries. + image_key_map: dict[str, str] = field(default_factory=dict) + + # Checkpoint feature schemas. The input schema determines camera order and + # the state width serialized into the prompt; the output schema determines + # the unpadded action width returned on the wire. + input_features: dict[str, Any] = field(default_factory=dict) + output_features: dict[str, Any] = field(default_factory=dict) + # OpenPI handshake metadata from the deploy config. Construction validates + # it against the resolved checkpoint contract. + policy_server_config: dict[str, Any] = field(default_factory=dict) + + # Derived from output_features[ACTION].shape; never accepted as a second + # source of truth in config.json. + action_dim: int = field(init=False) + # Derived from input_features[OBS_STATE].shape — the number of values the + # prompt serializes, not a padding target. See processor_pi05.as_state_vector. + state_dim: int = field(init=False) + + def __post_init__(self) -> None: + # Coerce list → tuple (JSON has no tuples). Squareness is a real + # constraint: SigLIP only takes square inputs. + res = self.image_resolution + if not isinstance(res, (tuple, list)) or len(res) != 2 or res[0] != res[1]: + raise ValueError(f"π0.5 expects a square image_resolution (H == W); got {res!r}.") + self.image_resolution = (int(res[0]), int(res[1])) + + if self.n_action_steps != self.chunk_size: + raise UnsupportedCheckpointCapabilityError( + "The stateless OpenPI serving path returns one complete action chunk and " + f"does not support n_action_steps={self.n_action_steps} with " + f"chunk_size={self.chunk_size}." + ) + # Mirrors Pi05Pipeline.SUPPORTED_DTYPES. This field records what the + # *checkpoint* declares; the dtype the weights are actually cast to comes + # from the top-level OmniDiffusionConfig and is validated there. + if self.dtype not in SUPPORTED_DTYPE_NAMES: + raise ValueError(f"dtype must be one of {sorted(SUPPORTED_DTYPE_NAMES)}, got {self.dtype!r}.") + + # 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 + ".")] + + if self.max_cameras < 1: + raise ValueError(f"max_cameras must be at least 1, got {self.max_cameras!r}.") + if self.num_inference_steps < 1: + # ``dt = -1 / num_steps`` in the Euler loop. + raise ValueError(f"num_inference_steps must be at least 1, got {self.num_inference_steps!r}.") + + self._derive_widths() + self._validate_norm_stats() + self._validate_policy_server_config() + self._validate_relative_actions() + + def _validate_norm_stats(self) -> None: + """Statistics have to cover the widths the checkpoint declares.""" + for key, width in (("state", self.state_dim), ("action", self.action_dim)): + for name, value in (self.norm_stats or {}).get(key, {}).items(): + if name != "mode" and len(value) != width: + raise ValueError(f"norm_stats[{key}][{name}] has {len(value)} entries, expected {width}.") + + # ------------------------------------------------------------------ + # Validation + # ------------------------------------------------------------------ + def _derive_widths(self) -> None: + """Read the real state and action widths off the checkpoint schema. + + π0.5 serializes the state into the prompt at its real width, so a 7-joint + checkpoint must not be padded out to ``max_state_dim``. The action chunk + is cropped to ``action_dim`` on the way out. + """ + state = self.input_features.get(OBS_STATE) or {} + action = self.output_features.get(ACTION) or {} + if state.get("type", "STATE").upper() != "STATE": + raise ValueError(f"input_features[{OBS_STATE!r}].type must be STATE, got {state['type']!r}.") + if action.get("type", "ACTION").upper() != "ACTION": + raise ValueError(f"output_features[{ACTION!r}].type must be ACTION, got {action['type']!r}.") + self.state_dim = int(state["shape"][0]) if state else self.max_state_dim + self.action_dim = int(action["shape"][0]) if action else self.max_action_dim + if self.state_dim > self.max_state_dim or self.action_dim > self.max_action_dim: + raise ValueError( + f"Checkpoint declares state_dim={self.state_dim} / action_dim={self.action_dim}, " + f"over max_state_dim={self.max_state_dim} / max_action_dim={self.max_action_dim}." + ) + + def _validate_policy_server_config(self) -> None: + """The OpenPI handshake has to advertise the resolved contract.""" + expected = { + "action_horizon": self.chunk_size, + "action_dim": self.action_dim, + "max_action_dim": self.max_action_dim, + "max_cameras": self.max_cameras, + "image_resolution": list(self.image_resolution), + } + for key, value in expected.items(): + declared = (self.policy_server_config or {}).get(key) + if declared is None: + continue + if key == "image_resolution": + declared = list(declared) + if declared != value: + raise ValueError( + f"policy_server_config.{key}={declared!r} does not match the resolved value {value!r}." + ) + + def _validate_relative_actions(self) -> None: + """Resolve the excluded joints at load, so a bad name fails before serving.""" + if not self.use_relative_actions or not self.relative_exclude_joints: + return + if not self.action_feature_names: + raise UnsupportedCheckpointCapabilityError( + f"relative_exclude_joints={self.relative_exclude_joints!r} needs action_feature_names, " + "which the checkpoint does not carry." + ) + resolve_excluded_action_indices(self.relative_exclude_joints, self.action_feature_names) + + # ------------------------------------------------------------------ + # Constructors + # ------------------------------------------------------------------ + @classmethod + def from_pretrained(cls, checkpoint_dir: str | Path) -> Pi05Config: + """Build from a checkpoint directory's ``config.json``. + + Normalization stats are *not* in ``config.json``. LeRobot keeps them in + the processor sidecar, so they are loaded separately and backfilled + here — see :func:`load_lerobot_norm_stats`. + """ + checkpoint_dir = Path(checkpoint_dir) + config_path = checkpoint_dir / "config.json" + if not config_path.exists(): + raise FileNotFoundError(f"π0.5 checkpoint is missing required config: {config_path}.") + with open(config_path, encoding="utf-8") as f: + raw = json.load(f) + config = cls.from_model_config(raw) + + if config.norm_stats is None: + stats = load_lerobot_norm_stats(checkpoint_dir) + if stats: + config.norm_stats = stats + config._validate_norm_stats() + return config + + @classmethod + def from_model_config(cls, model_config: dict[str, Any] | None) -> Pi05Config: + """Build from a config dict (LeRobot ``config.json`` or deploy yaml).""" + if not model_config: + return cls() + + raw = dict(model_config) + _reject_unsupported_capabilities(raw) + + model_type = raw.pop("type", "pi05") + if model_type != "pi05": + raise ValueError(f"Expected a π0.5 checkpoint (type='pi05'), got type={model_type!r}.") + + if "image_resolution" in raw: + raw["image_resolution"] = tuple(raw["image_resolution"]) + + # A config.json carries the whole training recipe; none of it reaches + # inference. Capabilities we cannot serve are rejected above instead. + allowed = {item.name for item in dataclass_fields(cls) if item.init} + filtered = {key: value for key, value in raw.items() if key in allowed} + dropped = sorted(set(raw) - allowed) + if dropped: + logger.debug("π0.5 config: ignoring %d non-runtime key(s): %s", len(dropped), dropped) + return cls(**filtered) + + +# ---------------------------------------------------------------------- +# LeRobot normalization-stats sidecar +# ---------------------------------------------------------------------- +# Stats live beside config.json: ``policy_preprocessor.json`` names a +# safetensors file per stateful step, keyed ``"."``. +_PREPROCESSOR_JSON = "policy_preprocessor.json" + +# The mode must come from ``norm_map``: a state_dict carries all of +# mean/std/min/max/q01/q99 regardless, so guessing from the present keys would +# read a QUANTILES checkpoint as mean_std and apply a wrong affine map. +_NORM_MODE_FROM_LEROBOT: dict[str, tuple[str, str, str] | None] = { + "IDENTITY": None, + "MEAN_STD": ("mean_std", "mean", "std"), + "MIN_MAX": ("min_max", "min", "max"), + "QUANTILES": ("quantile", "q01", "q99"), +} +# Renaming happens in an earlier processor step, so these names are canonical +# here. VISUAL is absent by design: images are normalized in the processor. +_NORM_STATS_FEATURES = {OBS_STATE: ("STATE", "state"), ACTION: ("ACTION", "action")} + + +def load_lerobot_norm_stats(checkpoint_dir: str | Path) -> dict[str, dict[str, Any]] | None: + """Load normalization stats from a LeRobot checkpoint's processor sidecar. + + Returns a ``norm_stats``-shaped dict (``{"state": {...}, "action": {...}}``) + carrying an **explicit** ``mode``, or ``None`` when the checkpoint ships no + stats. ``lerobot/pi05_base`` is the latter case: its normalizer step has no + ``state_file``, i.e. normalization is identity and the client is expected to + send an already-normalized state. + + Raises on a mode we cannot reproduce rather than serving the checkpoint with + the wrong transform, which fails silently — a wrongly normalized state still + yields a plausible-looking action chunk. + """ + checkpoint_dir = Path(checkpoint_dir) + preprocessor_path = checkpoint_dir / _PREPROCESSOR_JSON + if not preprocessor_path.exists(): + return None + with open(preprocessor_path, encoding="utf-8") as f: + steps = json.load(f).get("steps", []) + + step = next((s for s in steps if s.get("registry_name") == "normalizer_processor"), {}) + state_file = step.get("state_file") + if not state_file: + logger.info( + "π0.5 config: %s declares no normalizer state file; the state passes through.", + _PREPROCESSOR_JSON, + ) + return None + state_path = checkpoint_dir / state_file + if not state_path.exists(): + raise FileNotFoundError(f"{preprocessor_path} declares normalizer state {state_file!r}, which is missing.") + + norm_map = {str(k).upper(): str(v).upper() for k, v in ((step.get("config") or {}).get("norm_map") or {}).items()} + + import safetensors.torch + + flat = safetensors.torch.load_file(str(state_path)) + + stats: dict[str, dict[str, Any]] = {} + for feature_name, (feature_type, stats_key) in _NORM_STATS_FEATURES.items(): + lerobot_mode = norm_map.get(feature_type, "IDENTITY") + if lerobot_mode not in _NORM_MODE_FROM_LEROBOT: + raise ValueError( + f"π0.5 cannot reproduce normalization mode {lerobot_mode!r} for {feature_type}. " + f"Expected one of {sorted(_NORM_MODE_FROM_LEROBOT)}." + ) + selected = _NORM_MODE_FROM_LEROBOT[lerobot_mode] + if selected is None: + continue + + mode, *stat_names = selected + missing = [name for name in stat_names if f"{feature_name}.{name}" not in flat] + if missing: + raise ValueError( + f"{feature_type} declares {lerobot_mode} normalization but the state file is missing {missing}." + ) + stats[stats_key] = {"mode": mode} | {name: flat[f"{feature_name}.{name}"].tolist() for name in stat_names} + + if not stats: + return None + logger.info( + "π0.5 config: loaded normalization stats from %s — %s.", + state_file, + ", ".join(f"{key}={value['mode']}" for key, value in sorted(stats.items())), + ) + return stats + + +# ---------------------------------------------------------------------- +# Checkpoint-boundary rule +# ---------------------------------------------------------------------- +# Each entry: config key → (is this value *enabled*?, why we cannot serve it). +# Every capability here changes what a *correct* action chunk looks like, and +# none is visible in the weights, so a checkpoint that enables one is refused +# rather than served plausibly wrong. See recipes/lerobot/Pi05.md for the +# longer explanations. +# +# The predicate is per key because these are not all flags: ``n_obs_steps`` is +# a count whose supported value is 1, and an empty ``rtc_config`` still selects +# RTC — LeRobot's ``RTCConfig`` defaults ``enabled=True``. +_UNSUPPORTED: dict[str, tuple[Callable[[Any], bool], str]] = { + "use_peft": (bool, "adapter-aware loading; this loader takes a merged checkpoint only"), + "empty_cameras": (bool, "mutates the input schema; declare camera features explicitly instead"), + "use_visual_memory": (bool, "MEM needs a per-session observation history, which this path does not keep"), + "use_proprioceptive_memory": ( + bool, + "MEM sends the state as a projected token, not the discretized prompt built here", + ), + "rtc_config": ( + lambda value: value is not None, + "RTC needs prefix guidance in the denoising loop and per-request chunk carry-over", + ), + "n_obs_steps": ( + lambda value: bool(value) and value > 1, + "an observation history; this path is first-order Markov", + ), +} + + +def _reject_unsupported_capabilities(raw: dict[str, Any]) -> None: + """Raise if the checkpoint enables something we do not consume. + + Scope note: this validates that we correctly consume *the checkpoint we were + handed*. Choosing which checkpoint to load belongs to the caller. + """ + problems = [ + f" - {key}={raw[key]!r}: {why}" + for key, (is_enabled, why) in _UNSUPPORTED.items() + if key in raw and is_enabled(raw[key]) + ] + if problems: + raise UnsupportedCheckpointCapabilityError( + "This checkpoint declares capabilities the vllm-omni π0.5 implementation " + "does not support:\n" + "\n".join(problems) + "\nServing it anyway would " + "produce plausible-looking but wrong actions." + ) + + # Not an error: it records that the checkpoint was trained with clean action + # prefixes sampled in, which is a property of training rather than a request + # to run RTC. Such a checkpoint is still correct to serve without RTC. + delay = raw.get("rtc_training_max_delay") + if delay: + logger.info( + "π0.5 config: checkpoint was trained with rtc_training_max_delay=%s. " + "Real-Time Chunking is not implemented here; serving proceeds without it.", + delay, + ) diff --git a/vllm_omni/diffusion/models/pi05/modeling_pi05.py b/vllm_omni/diffusion/models/pi05/modeling_pi05.py new file mode 100644 index 000000000..a792ae0ef --- /dev/null +++ b/vllm_omni/diffusion/models/pi05/modeling_pi05.py @@ -0,0 +1,898 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Inference-only π0.5 VLA math kernel for vllm-omni. + +Only the math that turns a robot observation into an action chunk; no serving or +request glue. Deliberately shaped like ``models/pi0/modeling_pi0.py`` so the two +can later be factored into a shared Pi-family module (RFC step 2) on a +"behaviour unchanged" review. + +π0.5 = PaliGemma (SigLIP vision + Gemma 2B LM) prefix + Gemma 300M action expert +suffix + flow-matching head. Where π0 projects the robot state through +``state_proj``, π0.5 discretizes it into prompt tokens — so there is no +``state_proj`` layer, ``sample_actions`` takes no ``state`` argument, and the +suffix is action tokens only, which drops π0's leading state-token boundary from +the suffix attention mask and leaves ``[1] + [0] * (horizon - 1)``. + +π0.5 is nonetheless the *larger* model: the 37 AdaRMS ``dense`` projections add +~116M parameters against the ~8K that ``state_proj`` saves. (LeRobot's README +says otherwise; the checkpoint disagrees.) + +Reference implementations: + - OpenPI: openpi/src/openpi/models_pytorch/pi0_pytorch.py, gemma_pytorch.py + - LeRobot: lerobot/src/lerobot/policies/pi05/modeling_pi05.py +""" + +from __future__ import annotations + +import logging +import math +from collections.abc import Iterable + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers.models.auto import CONFIG_MAPPING +from transformers.models.gemma.modeling_gemma import ( + GemmaForCausalLM, + apply_rotary_pos_emb, +) +from transformers.models.paligemma.modeling_paligemma import ( + PaliGemmaForConditionalGeneration, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +logger = logging.getLogger(__name__) + +# ────────────────────────────────────────────────────────────────────── +# Constants +# ────────────────────────────────────────────────────────────────────── +DEFAULT_ACTION_DIM = 32 +DEFAULT_ACTION_HORIZON = 50 +DEFAULT_MAX_TOKEN_LEN = 200 # π0 uses 48 +DEFAULT_NUM_INFERENCE_STEPS = 10 +DEFAULT_IMAGE_RESOLUTION = (224, 224) # openpi/models/model.py IMAGE_RESOLUTION +DEFAULT_STATE_NUM_BINS = 256 # openpi PaliGemmaTokenizer.tokenize() + +# Large negative value to fill masked-out positions in a float attention mask. +# Matches OpenPI's constant exactly so that numerics line up during parity. +# Ref: openpi/src/openpi/models/gemma.py +OPENPI_ATTENTION_MASK_VALUE = -2.3819763e38 + + +# ────────────────────────────────────────────────────────────────────── +# Gemma variant configs (matches openpi/models/gemma.py get_config) +# ────────────────────────────────────────────────────────────────────── +class GemmaVariantConfig: + def __init__(self, width, depth, mlp_dim, num_heads, num_kv_heads, head_dim): + self.width = width + self.depth = depth + self.mlp_dim = mlp_dim + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.head_dim = head_dim + + +def get_gemma_config(variant: str) -> GemmaVariantConfig: + if variant == "gemma_2b": + return GemmaVariantConfig(2048, 18, 16384, 8, 1, 256) + elif variant == "gemma_300m": + return GemmaVariantConfig(1024, 18, 4096, 8, 1, 256) + else: + raise ValueError(f"Unknown variant: {variant}") + + +# ────────────────────────────────────────────────────────────────────── +# Utility functions (match openpi/models_pytorch/pi0_pytorch.py) +# ────────────────────────────────────────────────────────────────────── +def create_sinusoidal_pos_embedding( + time: torch.Tensor, + dimension: int, + min_period: float = 4e-3, + max_period: float = 4.0, + device: torch.device = None, +) -> torch.Tensor: + """Compute a sine/cosine positional embedding for scalar timesteps. + + Ref: openpi/models_pytorch/pi0_pytorch.py create_sinusoidal_pos_embedding + """ + if dimension % 2 != 0: + raise ValueError(f"dimension ({dimension}) must be divisible by 2") + if time.ndim != 1: + raise ValueError("time tensor must be 1-D (batch_size,)") + if device is None: + device = time.device + + # Use float64 for the log-linear sweep and for the inner products, to + # match the reference implementation's numerical behaviour exactly. + dtype = torch.float64 + fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device) + period = min_period * (max_period / min_period) ** fraction + scaling_factor = 1.0 / period * 2 * math.pi + sin_input = scaling_factor[None, :] * time[:, None].to(dtype) + return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1) + + +def make_att_2d_masks(pad_masks: torch.Tensor, att_masks: torch.Tensor) -> torch.Tensor: + """Build a 2D attention mask from a padding mask and an autoregressive mask. + + Ref: openpi/models_pytorch/pi0_pytorch.py make_att_2d_masks + """ + if att_masks.ndim != 2: + raise ValueError(f"att_masks must be 2-D, got {att_masks.ndim}-D") + if pad_masks.ndim != 2: + raise ValueError(f"pad_masks must be 2-D, got {pad_masks.ndim}-D") + + cumsum = torch.cumsum(att_masks, dim=1) + att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None] + pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None] + return att_2d_masks & pad_2d_masks + + +def prepare_attention_masks_4d(att_2d_masks: torch.Tensor) -> torch.Tensor: + """Convert ``(B, S, S)`` bool masks to ``(B, 1, S, S)`` float masks. + + ``True`` → 0.0 (attend), ``False`` → ``OPENPI_ATTENTION_MASK_VALUE``. + """ + att_2d_masks_4d = att_2d_masks[:, None, :, :] + return torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE) + + +class Pi05AdaRMSNorm(nn.Module): + """Adaptive RMSNorm conditioned on the flow-matching timestep. + + π0 conditions on time by *concatenating* a time embedding onto each action + embedding. π0.5 instead feeds the time embedding into every action-expert + norm, which produces a per-layer ``(scale, shift, gate)`` triple:: + + y = norm(x) * (1 + scale) + shift + out = residual + gate * sublayer(y) + + ``dense`` is zero-initialized, so an untrained model starts as the identity + modulation with a closed gate — matching OpenPI's parameterization. + + Note the shape of the unconditioned branch: ``normed * (1 + weight)`` with + ``weight`` zero-initialized, which is exactly ``transformers``' + ``GemmaRMSNorm``. That equivalence is why only the *expert* norms need + replacing here and the PaliGemma prefix can keep stock Gemma layers. + """ + + def __init__(self, dim: int, eps: float = 1e-6, cond_dim: int | None = None): + super().__init__() + self.eps = eps + self.dim = dim + self.cond_dim = cond_dim + if cond_dim is not None: + self.dense = nn.Linear(cond_dim, dim * 3, bias=True) + nn.init.zeros_(self.dense.weight) + nn.init.zeros_(self.dense.bias) + self.weight = None + else: + self.weight = nn.Parameter(torch.zeros(dim)) + self.dense = None + + def _norm(self, x: torch.Tensor) -> torch.Tensor: + var = torch.mean(torch.square(x.float()), dim=-1, keepdim=True) + return x.float() * torch.rsqrt(var + self.eps) + + def forward( + self, + x: torch.Tensor, + cond: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Return ``(normed, gate)``; ``gate`` is ``None`` in the unconditioned case.""" + dtype = x.dtype + normed = self._norm(x) + if self.dense is None: + normed = normed * (1.0 + self.weight.float()) + return normed.to(dtype), None + + if cond is None: + # A conditioned norm silently falling back to an unconditioned one + # would drop the entire timestep signal and still return a + # well-shaped, finite tensor. + raise ValueError( + "Pi05AdaRMSNorm was built with cond_dim=" + f"{self.cond_dim} but called without an AdaRMS conditioning vector." + ) + + if cond.shape[-1] != self.cond_dim: + raise ValueError(f"Expected AdaRMS cond dim {self.cond_dim}, got {cond.shape[-1]}") + + modulation = self.dense(cond.to(self.dense.weight.dtype)) + if x.ndim == 3: + # (B, 3*dim) → (B, 1, 3*dim), broadcast across the token axis: the + # timestep is a per-sample scalar, identical for every action token. + modulation = modulation.unsqueeze(1) + scale, shift, gate = modulation.chunk(3, dim=-1) + normed = normed * (1.0 + scale.float()) + shift.float() + return normed.to(dtype), gate.to(dtype) + + +def _gated_residual(residual: torch.Tensor, out: torch.Tensor, gate: torch.Tensor | None) -> torch.Tensor: + if gate is None: + return residual + out + return residual + gate * out + + +# ────────────────────────────────────────────────────────────────────── +# Dual-backbone: PaliGemma + AdaRMS action expert +# ────────────────────────────────────────────────────────────────────── +def _repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """Repeat KV heads for grouped-query attention. + + ``(B, num_kv_heads, S, D) → (B, num_kv_heads * n_rep, S, D)``. + """ + if n_rep == 1: + return hidden_states + b, nh, s, d = hidden_states.shape + hidden_states = hidden_states[:, :, None, :, :].expand(b, nh, n_rep, s, d) + return hidden_states.reshape(b, nh * n_rep, s, d) + + +def _attend(query_states, key_states, value_states, attention_mask, num_kv_groups, scaling): + """Manual eager attention: ``softmax(Q Kᵀ · scale + mask) · V``. + + The mask is sliced to ``key_states.shape[-2]`` so the same + ``(B, 1, Q, prefix+suffix)`` mask works in both the prefix pass (K length = + prefix) and the suffix pass (K length = prefix + suffix). + """ + k = _repeat_kv(key_states, num_kv_groups) + v = _repeat_kv(value_states, num_kv_groups) + attn_weights = torch.matmul(query_states, k.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask[:, :, :, : k.shape[-2]] + attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + return torch.matmul(attn_weights, v) + + +def _match(tensor: torch.Tensor, module: nn.Module) -> torch.Tensor: + """Cast ``tensor`` to the dtype ``module``'s weight expects.""" + return tensor.to(module.weight.dtype) if tensor.dtype != module.weight.dtype else tensor + + +def _compute_layer_prefix_only(layer_idx, hidden_states, attention_mask, position_ids, paligemma): + """Run one PaliGemma LM layer on the prefix, returning the layer output and + the post-RoPE ``(k, v)`` for the suffix pass. + + Identical to π0: the prefix backbone is unchanged in π0.5 (no AdaRMS — + there is no timestep in the prefix). + """ + model = paligemma.model.language_model + layer = model.layers[layer_idx] + residual = hidden_states + x = _match(layer.input_layernorm(hidden_states), layer.self_attn.q_proj) + + hidden_shape = (*x.shape[:-1], -1, layer.self_attn.head_dim) + q = layer.self_attn.q_proj(x).view(hidden_shape).transpose(1, 2) + k = layer.self_attn.k_proj(x).view(hidden_shape).transpose(1, 2) + v = layer.self_attn.v_proj(x).view(hidden_shape).transpose(1, 2) + + cos, sin = model.rotary_emb(v, position_ids) + q, k = apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1) + + att = _attend( + q, + k, + v, + attention_mask, + num_kv_groups=layer.self_attn.num_key_value_groups, + scaling=1.0 / math.sqrt(layer.self_attn.head_dim), + ) + att = att.transpose(1, 2).reshape(q.shape[0], -1, q.shape[1] * layer.self_attn.head_dim) + + out = layer.self_attn.o_proj(_match(att, layer.self_attn.o_proj)) + residual + after_resid = out + normed = layer.post_attention_layernorm(out) + out = layer.mlp(_match(normed, layer.mlp.up_proj)) + after_resid + return out, (k, v) + + +def _compute_layer_suffix_only( + layer_idx, + hidden_states, + prefix_kv, + attention_mask, + position_ids, + gemma_expert, + adarms_cond, +): + """Run one action-expert layer on the suffix with AdaRMS conditioning. + + This is where π0.5 diverges from π0. Both norms are + :class:`Pi05AdaRMSNorm` and each returns a gate that scales its sublayer's + contribution to the residual stream. + """ + layer = gemma_expert.model.layers[layer_idx] + + residual = hidden_states + x, gate = layer.input_layernorm(hidden_states, adarms_cond) + x = _match(x, layer.self_attn.q_proj) + + hidden_shape = (*x.shape[:-1], -1, layer.self_attn.head_dim) + q = layer.self_attn.q_proj(x).view(hidden_shape).transpose(1, 2) + k_suf = layer.self_attn.k_proj(x).view(hidden_shape).transpose(1, 2) + v_suf = layer.self_attn.v_proj(x).view(hidden_shape).transpose(1, 2) + + # RoPE frequencies are shared between PaliGemma and the expert. + cos, sin = gemma_expert.model.rotary_emb(v_suf, position_ids) + q, k_suf = apply_rotary_pos_emb(q, k_suf, cos, sin, unsqueeze_dim=1) + + # Concatenate cached prefix K/V (possibly different dtype) with suffix K/V. + k_prefix, v_prefix = prefix_kv + k = torch.cat([k_prefix.to(k_suf.dtype), k_suf], dim=2) + v = torch.cat([v_prefix.to(v_suf.dtype), v_suf], dim=2) + + att = _attend( + q, + k, + v, + attention_mask, + num_kv_groups=layer.self_attn.num_key_value_groups, + scaling=1.0 / math.sqrt(layer.self_attn.head_dim), + ) + att = att.transpose(1, 2).reshape(q.shape[0], -1, q.shape[1] * layer.self_attn.head_dim) + + hidden_states = _gated_residual(residual, layer.self_attn.o_proj(_match(att, layer.self_attn.o_proj)), gate) + + residual = hidden_states + x, gate = layer.post_attention_layernorm(hidden_states, adarms_cond) + return _gated_residual(residual, layer.mlp(_match(x, layer.mlp.up_proj)), gate) + + +class PaliGemmaWithActionExpertPi05(nn.Module): + """Dual-backbone transformer: PaliGemma (Gemma 2B) + AdaRMS expert (300M). + + Same two-mode dispatch as π0 (``prefix_only`` / ``suffix_only``), with one + structural change: after building a stock ``GemmaForCausalLM`` expert, every + norm in it is swapped for a :class:`Pi05AdaRMSNorm` carrying a ``dense`` + conditioning projection. + + Swapping in place — rather than subclassing ``GemmaModel`` as #4419 does — + keeps the module tree, and therefore the checkpoint key layout, identical to + the expert's stock layout apart from the norms themselves. + """ + + def __init__(self, vlm_config, action_expert_config): + super().__init__() + + # PaliGemma prefix: identical to π0. It sees no timestep, so no AdaRMS. + vlm_config_hf = CONFIG_MAPPING["paligemma"]() + vlm_config_hf._vocab_size = 257152 + vlm_config_hf.image_token_index = 257152 + vlm_config_hf.text_config.hidden_size = vlm_config.width + vlm_config_hf.text_config.intermediate_size = vlm_config.mlp_dim + vlm_config_hf.text_config.num_attention_heads = vlm_config.num_heads + vlm_config_hf.text_config.head_dim = vlm_config.head_dim + vlm_config_hf.text_config.num_hidden_layers = vlm_config.depth + vlm_config_hf.text_config.num_key_value_heads = vlm_config.num_kv_heads + vlm_config_hf.text_config.hidden_activation = "gelu_pytorch_tanh" + vlm_config_hf.text_config.dtype = "float32" + vlm_config_hf.text_config.vocab_size = 257152 + vlm_config_hf.vision_config.intermediate_size = 4304 + vlm_config_hf.vision_config.projection_dim = 2048 + vlm_config_hf.vision_config.projector_hidden_act = "gelu_fast" + vlm_config_hf.vision_config.dtype = "float32" + + action_expert_config_hf = CONFIG_MAPPING["gemma"]( + head_dim=action_expert_config.head_dim, + hidden_size=action_expert_config.width, + intermediate_size=action_expert_config.mlp_dim, + num_attention_heads=action_expert_config.num_heads, + num_hidden_layers=action_expert_config.depth, + num_key_value_heads=action_expert_config.num_kv_heads, + vocab_size=257152, + hidden_activation="gelu_pytorch_tanh", + dtype="float32", + ) + + self.paligemma = PaliGemmaForConditionalGeneration(config=vlm_config_hf) + self.gemma_expert = GemmaForCausalLM(config=action_expert_config_hf) + # The action expert doesn't embed tokens — it only consumes the + # suffix action embeddings we feed in. + self.gemma_expert.model.embed_tokens = None + + self.adarms_cond_dim = action_expert_config.width + self._install_adarms_norms(action_expert_config) + + def _install_adarms_norms(self, action_expert_config) -> None: + """Replace every action-expert RMSNorm with a conditioned AdaRMS norm.""" + expert = self.gemma_expert.model + eps = getattr(self.gemma_expert.config, "rms_norm_eps", 1e-6) + width = action_expert_config.width + for layer in expert.layers: + layer.input_layernorm = Pi05AdaRMSNorm(width, eps=eps, cond_dim=self.adarms_cond_dim) + layer.post_attention_layernorm = Pi05AdaRMSNorm(width, eps=eps, cond_dim=self.adarms_cond_dim) + expert.norm = Pi05AdaRMSNorm(width, eps=eps, cond_dim=self.adarms_cond_dim) + + def embed_image(self, pixel_values: torch.Tensor) -> torch.Tensor: + """Encode images with SigLIP vision tower + PaliGemma projector. + + The two steps are run explicitly rather than via + ``PaliGemmaModel.get_image_features`` because that helper divides the + projector output by ``sqrt(text_hidden_size)``. Being explicit keeps the + scale unambiguous and matches π0 exactly (SigLIP is unchanged in π0.5). + """ + # Shapes: pixel_values (B, 3, 224, 224) → SigLIP (B, 256, 1152) + # → projector (B, 256, 2048) + vision_outputs = self.paligemma.model.vision_tower(pixel_values) + image_features = vision_outputs.last_hidden_state + return self.paligemma.model.multi_modal_projector(image_features) + + def embed_language_tokens(self, tokens: torch.Tensor) -> torch.Tensor: + """Embed language tokens, returning the ``* sqrt(hidden)``-scaled embedding. + + The scaling location moved across transformers releases: at ≤5.3 the + normalizer lives inside ``GemmaModel.forward`` (which we bypass), and at + ≥5.4 ``GemmaTextScaledWordEmbedding`` self-applies it. Detect and avoid + double-scaling. + """ + embed_tokens = self.paligemma.model.language_model.embed_tokens + lang_emb = embed_tokens(tokens) + if getattr(embed_tokens, "embed_scale", None) is None: + lang_emb = lang_emb * math.sqrt(lang_emb.shape[-1]) + return lang_emb + + def forward( + self, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values=None, + inputs_embeds: list[torch.Tensor | None] | None = None, + use_cache: bool = False, + adarms_cond: torch.Tensor | None = None, + ): + """Dispatch to prefix_only / suffix_only and return + ``([prefix_out, suffix_out], past_key_values_or_None)``. + """ + num_layers = self.paligemma.config.text_config.num_hidden_layers + pali_lm = self.paligemma.model.language_model + expert_lm = self.gemma_expert.model + + if inputs_embeds[1] is None: + hidden_states = inputs_embeds[0] + kv_list: list[tuple[torch.Tensor, torch.Tensor]] = [] + for layer_idx in range(num_layers): + hidden_states, kv = _compute_layer_prefix_only( + layer_idx, + hidden_states, + attention_mask, + position_ids, + paligemma=self.paligemma, + ) + kv_list.append(kv) + hidden_states = pali_lm.norm(hidden_states) + return [hidden_states, None], (kv_list if use_cache else None) + + if inputs_embeds[0] is not None: + raise ValueError( + "PaliGemmaWithActionExpertPi05.forward only supports prefix-only " + "or suffix-only dispatch; got both inputs_embeds populated." + ) + if not isinstance(past_key_values, list): + raise TypeError( + "suffix_only forward expects past_key_values to be the " + "list[(k, v)] produced by a previous prefix_only forward; " + f"got {type(past_key_values)}" + ) + hidden_states = inputs_embeds[1] + for layer_idx in range(num_layers): + hidden_states = _compute_layer_suffix_only( + layer_idx, + hidden_states, + past_key_values[layer_idx], + attention_mask, + position_ids, + gemma_expert=self.gemma_expert, + adarms_cond=adarms_cond, + ) + hidden_states, _ = expert_lm.norm(hidden_states, adarms_cond) + return [None, hidden_states], None + + +# ────────────────────────────────────────────────────────────────────── +# Main π0.5 Model +# ────────────────────────────────────────────────────────────────────── +class Pi05ForActionPrediction(nn.Module): + """π0.5 VLA model for robot action prediction via flow matching. + + Inference flow: + 1. Embed prefix (images + language, where the language already carries the + discretized state) → prefix tokens. + 2. Forward prefix through PaliGemma → layer-wise KV cache. + 3. For each denoising step ``t = 1.0, 1-dt, ..., 0``: + a. Embed the timestep → an AdaRMS conditioning vector. + b. Embed the suffix (action tokens only). + c. Forward the suffix through the AdaRMS action expert. + d. ``x_t = x_t + dt * v_t`` (Euler integration). + 4. Return ``x_0`` as the predicted action chunk. + """ + + def __init__(self, config, quant_config=None, prefix: str = ""): + super().__init__() + # ``quant_config`` is accepted for interface compatibility but unused — + # quant_config is not plumbed through; the weight dtype comes from the pipeline. + del quant_config + self.config = config + + self.action_dim = getattr(config, "max_action_dim", DEFAULT_ACTION_DIM) + self.max_state_dim = getattr(config, "max_state_dim", self.action_dim) + self.action_horizon = getattr(config, "chunk_size", DEFAULT_ACTION_HORIZON) + self.num_inference_steps = getattr(config, "num_inference_steps", DEFAULT_NUM_INFERENCE_STEPS) + + paligemma_variant = getattr(config, "paligemma_variant", "gemma_2b") + action_expert_variant = getattr(config, "action_expert_variant", "gemma_300m") + vlm_config = get_gemma_config(paligemma_variant) + expert_config = get_gemma_config(action_expert_variant) + self.vlm_width = vlm_config.width + self.expert_width = expert_config.width + + # Dual backbone + self.paligemma_with_expert = PaliGemmaWithActionExpertPi05(vlm_config, expert_config) + + # Action chunk projections. + self.action_in_proj = nn.Linear(self.action_dim, self.expert_width) + self.action_out_proj = nn.Linear(self.expert_width, self.action_dim) + + # π0.5 timestep MLP: (W → W → W) with SiLU, feeding AdaRMS. + # π0 instead has action_time_mlp_{in,out} of shape (2W → W → W) because + # it concatenates the time embedding onto the action embedding. + # NOTE: there is deliberately **no** ``state_proj`` here — that is the + # π0-only continuous-state path. + self.time_mlp_in = nn.Linear(self.expert_width, self.expert_width) + self.time_mlp_out = nn.Linear(self.expert_width, self.expert_width) + + # ── Prefix embedding ───────────────────────────────────────────── + def embed_prefix( + self, + images: list[torch.Tensor], + image_masks: list[torch.Tensor], + lang_tokens: torch.Tensor, + lang_masks: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build prefix embeddings, per-token padding mask, and AR mask. + + Prefix tokens form ``[img_cam_0..., ..., lang_tokens...]`` with fully + bidirectional attention (all-zero ``att_masks``). Identical to π0 — the + state is inside ``lang_tokens``, so nothing here changes shape-wise. + + Cameras are embedded one at a time; each call is ``(B, 3, 224, 224)``. + The number of slots is fixed by ``config.max_cameras`` for the deployed + model; missing cameras occupy their slot with a false image mask. + """ + num_views = len(images) + if len(image_masks) != num_views: + raise ValueError( + f"images and image_masks must contain the same number of views, got {num_views} and {len(image_masks)}." + ) + max_cameras = int(self.config.max_cameras) + if num_views != max_cameras: + raise ValueError(f"Expected exactly max_cameras={max_cameras} image views, got {num_views}.") + + embs: list[torch.Tensor] = [] + pad_masks: list[torch.Tensor] = [] + att_masks: list[int] = [] + + for img, img_mask in zip(images, image_masks): + img_emb = self.paligemma_with_expert.embed_image(img) + bsize, num_img_embs = img_emb.shape[:2] + embs.append(img_emb) + pad_masks.append(img_mask[:, None].expand(bsize, num_img_embs)) + att_masks += [0] * num_img_embs + + lang_emb = self.paligemma_with_expert.embed_language_tokens(lang_tokens) + embs.append(lang_emb) + pad_masks.append(lang_masks) + att_masks += [0] * lang_emb.shape[1] + + embs = torch.cat(embs, dim=1) + pad_masks = torch.cat(pad_masks, dim=1) + att_masks = torch.tensor(att_masks, dtype=torch.bool, device=embs.device) + att_masks = att_masks[None, :].expand(pad_masks.shape[0], -1) + + return embs, pad_masks, att_masks + + # ── Timestep + suffix embedding ────────────────────────────────── + def embed_timestep(self, timestep: torch.Tensor) -> torch.Tensor: + """Timestep → AdaRMS conditioning vector ``(B, expert_width)``. + + ``silu(time_mlp_out(silu(time_mlp_in(sinusoid(t)))))``. The trailing + SiLU is part of the reference implementation — dropping it is a silent + numerical error, not a crash. + """ + model_dtype = self.action_in_proj.weight.dtype + time_emb = create_sinusoidal_pos_embedding( + timestep, + self.action_in_proj.out_features, + min_period=getattr(self.config, "min_period", 4e-3), + max_period=getattr(self.config, "max_period", 4.0), + device=timestep.device, + ).to(dtype=model_dtype) + time_cond = self.time_mlp_in(time_emb) + time_cond = F.silu(time_cond) + time_cond = self.time_mlp_out(time_cond) + return F.silu(time_cond) + + def embed_suffix( + self, + noisy_actions: torch.Tensor, + timestep: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Build the suffix: **action tokens only**, plus the AdaRMS condition. + + π0's suffix is ``[state_token, action_tokens×H]`` with an AR mask of + ``[1, 1, 0...]``. π0.5 has no state token, so the suffix is + ``[action_tokens×H]`` and the mask is ``[1] + [0]*(H-1)``: the first + action token opens a causal block and the rest attend bidirectionally + within it. + """ + model_dtype = self.action_in_proj.weight.dtype + noisy_actions = noisy_actions.to(dtype=model_dtype) + + time_cond = self.embed_timestep(timestep) + action_emb = self.action_in_proj(noisy_actions) # (B, H, W) + + bsize, action_len = action_emb.shape[:2] + pad_masks = torch.ones(bsize, action_len, dtype=torch.bool, device=action_emb.device) + att_masks = torch.tensor( + [1] + [0] * (self.action_horizon - 1), + dtype=action_emb.dtype, + device=action_emb.device, + )[None, :].expand(bsize, -1) + return action_emb, pad_masks, att_masks, time_cond + + # ── Denoising step ─────────────────────────────────────────────── + def denoise_step( + self, + prefix_pad_masks: torch.Tensor, + past_key_values, + x_t: torch.Tensor, + timestep: torch.Tensor, + ) -> torch.Tensor: + """Apply one flow-matching denoising step: predict ``v_t`` from ``x_t``. + + Signature differs from π0's by exactly one argument: no ``state``. + """ + suffix_embs, suffix_pad_masks, suffix_att_masks, time_cond = self.embed_suffix(x_t, timestep) + + batch_size = prefix_pad_masks.shape[0] + suffix_len = suffix_pad_masks.shape[1] + prefix_len = prefix_pad_masks.shape[1] + + prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, suffix_len, prefix_len) + suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks) + full_att_2d_masks = torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2) + + # Position IDs continue from where the prefix's last valid token left off. + prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None] + position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1 + + full_att_2d_masks_4d = prepare_attention_masks_4d(full_att_2d_masks) + + outputs_embeds, _ = self.paligemma_with_expert.forward( + attention_mask=full_att_2d_masks_4d, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=[None, suffix_embs], + use_cache=False, + adarms_cond=time_cond, + ) + + # Every suffix token is an action token in π0.5 (no state token to drop). + suffix_out = outputs_embeds[1][:, -self.action_horizon :] + suffix_out = suffix_out.to(dtype=self.action_out_proj.weight.dtype) + return self.action_out_proj(suffix_out) + + # ── Full action generation ─────────────────────────────────────── + @torch.no_grad() + def sample_actions( + self, + images: list[torch.Tensor], + image_masks: list[torch.Tensor], + lang_tokens: torch.Tensor, + lang_masks: torch.Tensor, + noise: torch.Tensor | None = None, + num_steps: int | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + ) -> torch.Tensor: + """Generate an action chunk via iterative flow-matching denoising. + + Convention: ``t=1`` is noise, ``t=0`` is the target — opposite of the + published π0 paper but matching both OpenPI and LeRobot. + + Takes no ``state``: π0.5's state rides inside ``lang_tokens``. + """ + if num_steps is None: + num_steps = self.num_inference_steps + + bsize = lang_tokens.shape[0] + device = lang_tokens.device + if noise is None: + noise_shape = (self.action_horizon, self.action_dim) + if isinstance(generator, list): + if len(generator) != bsize: + raise ValueError(f"Expected {bsize} generators, got {len(generator)}.") + noise = torch.stack( + [torch.randn(noise_shape, dtype=torch.float32, device=device, generator=item) for item in generator] + ) + else: + noise = torch.randn( + bsize, + *noise_shape, + dtype=torch.float32, + device=device, + generator=generator, + ) + + # 1. Prefix embeddings + mask building. + prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix( + images, image_masks, lang_tokens, lang_masks + ) + prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks) + prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1 + prefix_att_2d_masks_4d = prepare_attention_masks_4d(prefix_att_2d_masks) + + # 2. Forward prefix through PaliGemma LM, producing a list[(k, v)] cache. + _, past_key_values = self.paligemma_with_expert.forward( + attention_mask=prefix_att_2d_masks_4d, + position_ids=prefix_position_ids, + past_key_values=None, + inputs_embeds=[prefix_embs, None], + use_cache=True, + ) + + # 3. Euler-integrated denoising from t=1 down to t=0. + dt = -1.0 / num_steps + x_t = noise + for step in range(num_steps): + t = 1.0 + step * dt + time_tensor = torch.full((bsize,), t, dtype=torch.float32, device=device) + v_t = self.denoise_step( + prefix_pad_masks=prefix_pad_masks, + past_key_values=past_key_values, + x_t=x_t, + timestep=time_tensor, + ) + x_t = x_t + dt * v_t + return x_t + + # ── Weight loading ─────────────────────────────────────────────── + def load_weights( + self, + weights: Iterable[tuple[str, torch.Tensor]], + *, + strict: bool = True, + ): + """Load and audit a LeRobot π0.5 safetensors checkpoint. + + Same remap rules as π0 (strip the ``model.`` prefix, flatten→nested + PaliGemma submodules, tied ``lm_head`` → ``embed_tokens``, version-robust + SigLIP nesting), plus two π0.5-specific ones: + + * ``action_time_mlp_{in,out}`` → ``time_mlp_{in,out}``: some + checkpoints were exported under the π0 parameter names. + * ``state_proj.*`` is reported, not silently dropped. A π0.5 + checkpoint should not contain it; its presence usually means a π0 + checkpoint was pointed at the π0.5 model class, which would + otherwise run happily with a randomly-initialized action expert. + + The action-expert norms are AdaRMS here, so they expose ``dense.weight`` + / ``dense.bias`` and no plain ``weight``. A checkpoint that carries a + plain expert-norm ``weight`` is a π0-shaped checkpoint; that too is + rejected rather than skipped. ``strict=False`` exists only for focused + remapping unit tests that intentionally provide a partial state dict; + the serving path always uses the strict default. + """ + params_dict = dict(self.named_parameters()) + buffers_dict = dict(self.named_buffers()) + + _PALIGEMMA_SUBMODULES = ("vision_tower", "multi_modal_projector", "language_model") + _EXPERT_PREFIX = "paligemma_with_expert.gemma_expert.model." + + def _remap(name: str) -> str: + # Strip the leading "model." that LeRobot's PI05Policy wrapper adds. + if name.startswith("model."): + name = name[len("model.") :] + + # π0-style timestep MLP names → π0.5 names. + if name.startswith("action_time_mlp_in."): + name = "time_mlp_in." + name[len("action_time_mlp_in.") :] + elif name.startswith("action_time_mlp_out."): + name = "time_mlp_out." + name[len("action_time_mlp_out.") :] + + # Nested PaliGemma layout. + for sub in _PALIGEMMA_SUBMODULES: + flat = f"paligemma_with_expert.paligemma.{sub}." + nested = f"paligemma_with_expert.paligemma.model.{sub}." + if name.startswith(flat) and not name.startswith(nested): + return nested + name[len(flat) :] + + if name == "paligemma_with_expert.paligemma.lm_head.weight": + return "paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight" + + return name + + def _fix_vision_tower(name: str) -> str: + """Reconcile SigLIP nesting across transformers versions (≤5.3 wraps + the encoder in ``vision_tower.vision_model.*``; ≥5.4 flattens it).""" + vt = "paligemma_with_expert.paligemma.model.vision_tower." + if not name.startswith(vt): + return name + rest = name[len(vt) :] + if name in params_dict or name in buffers_dict: + return name + if rest.startswith("vision_model."): + candidate = vt + rest[len("vision_model.") :] + else: + candidate = vt + "vision_model." + rest + return candidate if (candidate in params_dict or candidate in buffers_dict) else name + + loaded = 0 + skipped: list[str] = [] + pi0_shaped: list[str] = [] + filled_params: set = set() + + for name, loaded_weight in weights: + mapped = _fix_vision_tower(_remap(name)) + + # Diagnose π0-shaped keys instead of dropping them quietly. + is_expert_norm_weight = mapped.startswith(_EXPERT_PREFIX) and ( + mapped.endswith("input_layernorm.weight") + or mapped.endswith("post_attention_layernorm.weight") + or mapped == _EXPERT_PREFIX + "norm.weight" + ) + if mapped.startswith("state_proj.") or is_expert_norm_weight: + pi0_shaped.append(mapped) + continue + + if mapped in params_dict: + param = params_dict[mapped] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded += 1 + filled_params.add(mapped) + elif mapped in buffers_dict: + buffers_dict[mapped].copy_(loaded_weight) + loaded += 1 + filled_params.add(mapped) + else: + skipped.append(mapped) + + # LeRobot stores PaliGemma's tied text embedding as lm_head.weight; keep + # lm_head filled so the tied-weight state stays consistent. + embed_key = "paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight" + lm_head_key = "paligemma_with_expert.paligemma.lm_head.weight" + if embed_key in filled_params and lm_head_key in params_dict and lm_head_key not in filled_params: + params_dict[lm_head_key].data.copy_(params_dict[embed_key].data) + filled_params.add(lm_head_key) + + # Reverse audit: any model param that got no checkpoint tensor at all + # would be running with random init. + missing_params: list[str] = [] + for pname in params_dict: + if pname in filled_params: + continue + if "rotary_emb" in pname or pname.endswith(".inv_freq"): + continue + missing_params.append(pname) + + parts: list[str] = [] + if pi0_shaped: + parts.append( + f"{len(pi0_shaped)} checkpoint key(s) are π0-shaped, not π0.5-shaped (first 5: {pi0_shaped[:5]})" + ) + if skipped: + parts.append(f"{len(skipped)} checkpoint key(s) had no model target (first 5: {skipped[:5]})") + if missing_params: + parts.append(f"{len(missing_params)} model param(s) received no weight (first 5: {missing_params[:5]})") + + if parts and strict: + raise RuntimeError("Incomplete or incompatible π0.5 checkpoint: " + "; ".join(parts)) + if parts: + logger.debug("π0.5 partial test load: %d tensors loaded — %s.", loaded, "; ".join(parts)) + else: + logger.info("π0.5 load_weights: %d tensors loaded, 0 skipped, 0 missing.", loaded) + return filled_params + + +EntryClass = Pi05ForActionPrediction diff --git a/vllm_omni/diffusion/models/pi05/pipeline_pi05.py b/vllm_omni/diffusion/models/pi05/pipeline_pi05.py new file mode 100644 index 000000000..a4d32ed81 --- /dev/null +++ b/vllm_omni/diffusion/models/pi05/pipeline_pi05.py @@ -0,0 +1,311 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""π0.5 VLA pipeline for vllm-omni. + +Entry point for ``DiffusionEngine.step() → pipeline.forward(req)``. Mirrors the +DreamZero contract and the π0 pipeline: 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})``. + +π0.5 is stateless across calls (no KV reuse, first-order Markov), so +``session_id`` / ``reset`` from the OpenPI protocol are accepted but ignored. + +The post-processing order is load-bearing and matches LeRobot:: + + unnormalize → absolute actions → to_cpu + +``AbsoluteActionsProcessorStep`` must run *after* unnormalization, because a +relative-action checkpoint's ``norm_stats`` are computed in relative space. +""" + +from __future__ import annotations + +import json +import os +from dataclasses import fields as dataclass_fields + +import numpy as np +import torch +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.models.pi05.config import SUPPORTED_DTYPE_NAMES, Pi05Config +from vllm_omni.diffusion.models.pi05.modeling_pi05 import Pi05ForActionPrediction +from vllm_omni.diffusion.models.pi05.processor_pi05 import Pi05Processor +from vllm_omni.diffusion.models.pi05_pipeline_config import PI05_PIPELINE as PI05_PIPELINE +from vllm_omni.diffusion.request import OmniDiffusionRequest + +logger = init_logger(__name__) + +# π0.5 pins the PaliGemma tokenizer (LeRobot hardcodes it too). +DEFAULT_PI05_TOKENIZER = "google/paligemma-3b-pt-224" + +# float16 is absent deliberately, not by oversight: nothing here has been +# validated against it. The float32/bfloat16 trade-off is in the deploy config +# and recipes/lerobot/Pi05.md. +SUPPORTED_DTYPES = (torch.float32, torch.bfloat16) + +# The two lists guard different entry points — what the checkpoint declares +# versus the dtype actually cast to — so they must not drift apart. +assert {str(dtype).split(".")[-1] for dtype in SUPPORTED_DTYPES} == set(SUPPORTED_DTYPE_NAMES) + + +def _checkpoint_declared_keys(model_dir: str) -> set[str]: + """The keys the checkpoint's config.json actually carries.""" + path = os.path.join(model_dir, "config.json") + if not os.path.exists(path): + return set() + with open(path, encoding="utf-8") as f: + return set(json.load(f)) + + +def _comparable(value): + """Compare yaml and checkpoint values without tripping on list/tuple.""" + return list(value) if isinstance(value, (list, tuple)) else value + + +def _pi05_post_process(x): + """Module-level identity post-process (picklable across the orchestrator's + multiprocess boundary — a local closure is not).""" + return x + + +def get_pi05_post_process_func(od_config: OmniDiffusionConfig): + """π0.5 returns actions directly; post-processing is identity.""" + del od_config + return _pi05_post_process + + +_LEROBOT_FLOAT32_IN_BFLOAT16 = ( + "vision_tower", + "multi_modal_projector", + "input_layernorm", + "post_attention_layernorm", + "model.norm", +) + + +def _to_bfloat16_for_inference(model: Pi05ForActionPrediction) -> None: + """Apply LeRobot's mixed-precision inference layout explicitly.""" + inner_model = model.paligemma_with_expert + for name, param in inner_model.named_parameters(): + target_dtype = ( + torch.float32 if any(selector in name for selector in _LEROBOT_FLOAT32_IN_BFLOAT16) else torch.bfloat16 + ) + param.data = param.data.to(dtype=target_dtype) + for buffer in inner_model.buffers(): + if buffer.is_floating_point(): + buffer.data = buffer.data.to(dtype=torch.bfloat16) + + for name, param in model.named_parameters(): + if not name.startswith("paligemma_with_expert."): + param.data = param.data.to(dtype=torch.float32) + for name, buffer in model.named_buffers(): + if not name.startswith("paligemma_with_expert.") and buffer.is_floating_point(): + buffer.data = buffer.data.to(dtype=torch.float32) + + +def _set_inference_dtype(model: Pi05ForActionPrediction, dtype: torch.dtype) -> None: + if dtype == torch.float32: + model.to(dtype=torch.float32) + elif dtype == torch.bfloat16: + _to_bfloat16_for_inference(model) + else: + raise ValueError(f"Unsupported π0.5 inference dtype: {dtype!r}.") + + +class Pi05Pipeline(nn.Module): + """π0.5 VLA pipeline: raw robot obs → continuous action chunk. + + Registered as ``"Pi05Pipeline"`` in the diffusion registry. + """ + + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): + super().__init__() + self.od_config = od_config + self.prefix = prefix + 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())) + + 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() + + self.processor = Pi05Processor(self.config, self.tokenizer, self._device) + + # ------------------------------------------------------------------ + # 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.""" + if not model: + return None + if os.path.isdir(model): + return model + # Via repo_utils' shared HfApi rather than huggingface_hub directly, so + # the download carries vLLM-Omni's user agent like every other repo access. + from vllm_omni.transformers_utils.repo_utils import hf_api + + return hf_api().snapshot_download( + repo_id=model, + allow_patterns=["*.json", "*.safetensors", "*.model", "tokenizer*"], + ) + + def _build_config(self, od_config: OmniDiffusionConfig) -> Pi05Config: + """Read the config from the checkpoint, then let the deploy yaml override it.""" + checkpoint_config = Pi05Config.from_pretrained(self.model_dir) if self.model_dir else None + if checkpoint_config is None: + return Pi05Config.from_model_config(od_config.model_config) + if not od_config.model_config: + return checkpoint_config + + resolved = { + item.name: getattr(checkpoint_config, item.name) for item in dataclass_fields(Pi05Config) if item.init + } + # Only a key the checkpoint actually declares can disagree with the yaml. + # Comparing against the dataclass default instead would report every + # serving-only field, such as policy_server_config, on every start. + declared_keys = _checkpoint_declared_keys(self.model_dir) + for key, value in od_config.model_config.items(): + if key in resolved and key in declared_keys and _comparable(value) != _comparable(resolved[key]): + logger.warning( + "Pi05Pipeline: the deploy config sets %s=%r, overriding %r from the checkpoint.", + key, + value, + resolved[key], + ) + resolved.update(od_config.model_config) + return Pi05Config.from_model_config(resolved) + + 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_PI05_TOKENIZER + + @staticmethod + def _resolve_dtype(od_config: OmniDiffusionConfig) -> torch.dtype: + """Resolve the dtype the weights are actually cast to. + + This is the load-bearing check, not ``Pi05Config.dtype``: the cast in + :meth:`_initialize_model` reads the *top-level* ``OmniDiffusionConfig`` + field, so a guard on the model config alone would let an unsupported + dtype through. See :data:`SUPPORTED_DTYPES` for why float16 is excluded. + """ + dt = od_config.dtype + resolved = dt if isinstance(dt, torch.dtype) else getattr(torch, str(dt).split(".")[-1], None) + if resolved not in SUPPORTED_DTYPES: + raise ValueError( + f"Unsupported π0.5 dtype: {dt!r}. Supported: " + f"{', '.join(sorted(str(d).split('.')[-1] for d in SUPPORTED_DTYPES))}." + ) + return resolved + + @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 + + # padding_side="right" is part of the π0.5 spec: the prefix is a fixed + # 200-token block whose live tokens must start at index 0. + return AutoTokenizer.from_pretrained(self.tokenizer_source, padding_side="right") + + 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) -> Pi05ForActionPrediction: + if not self.has_real_checkpoint(): + expected = os.path.join(self.model_dir or "", "model.safetensors") + raise FileNotFoundError(f"π0.5 serving requires checkpoint weights at {expected}.") + model = Pi05ForActionPrediction(self.config) + _set_inference_dtype(model, self._torch_dtype) + model.to(device=self._device) + self._load_checkpoint(model) + model.eval() + return model + + def _load_checkpoint(self, model: Pi05ForActionPrediction) -> None: + import safetensors.torch + + path = os.path.join(self.model_dir, "model.safetensors") + logger.info("Pi05Pipeline: loading π0.5 weights from %s", path) + try: + state = safetensors.torch.load_file(path) + model.load_weights(state.items()) + except Exception as exc: + raise RuntimeError(f"Failed to load complete π0.5 checkpoint {path}: {exc}") from exc + + # ------------------------------------------------------------------ + # Framework weight-loading hook + # ------------------------------------------------------------------ + def load_weights(self, weights=()): # noqa: D401 + """No-op for the diffusion loader: π0.5 self-loads its checkpoint.""" + 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("Pi05Pipeline: dummy warmup request without robot_obs — returning zeros.") + return DiffusionOutput( + output={ + "actions": np.zeros( + (self.config.chunk_size, self.config.action_dim), + dtype=np.float32, + ) + }, + ) + return DiffusionOutput( + error="Pi05Pipeline.forward requires sampling_params.extra_args['robot_obs'].", + ) + + images, image_masks, lang_tokens, lang_masks = self.processor.build_model_inputs(robot_obs) + + num_steps = getattr(req.sampling_params, "num_inference_steps", None) + if num_steps is not None and ( + isinstance(num_steps, bool) or not isinstance(num_steps, (int, np.integer)) or int(num_steps) < 1 + ): + raise ValueError(f"num_inference_steps must be a positive integer, got {num_steps!r}.") + if num_steps is not None: + num_steps = int(num_steps) + + actions = self.model.sample_actions( + images=images, + image_masks=image_masks, + lang_tokens=lang_tokens, + lang_masks=lang_masks, + num_steps=num_steps, + generator=req.sampling_params.generator, + ) + + return DiffusionOutput(output={"actions": self.processor.build_model_outputs(actions, robot_obs)}) diff --git a/vllm_omni/diffusion/models/pi05/processor_pi05.py b/vllm_omni/diffusion/models/pi05/processor_pi05.py new file mode 100644 index 000000000..11ebad2db --- /dev/null +++ b/vllm_omni/diffusion/models/pi05/processor_pi05.py @@ -0,0 +1,568 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +r"""Preprocessing for the π0.5 VLA model. + +Converts a raw robot observation (multi-camera images + language instruction + +proprioceptive state) into the tensors ``Pi05ForActionPrediction.sample_actions`` +consumes. + +**The defining π0.5 difference:** the state is not projected by a ``state_proj`` +layer. It is normalized to ``[-1, 1]``, discretized into ``state_num_bins`` bins, +and serialized into the language prompt:: + + "Task: , State: ... ;\nAction: " + +so ``sample_actions`` receives no state tensor at all. + +**Normalization must precede discretization.** The discretizer bins over +``[-1, 1]`` and assumes the state is already in that range. Reversed, every bin +index is wrong and nothing raises. + +Reference: + - OpenPI: openpi/src/openpi/shared/image_tools.py (resize_with_pad) + - OpenPI: openpi/src/openpi/models/pi0_config.py (PaliGemmaTokenizer.tokenize) + - LeRobot: lerobot/src/lerobot/policies/pi05/processor_pi05.py +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image +from vllm.logger import init_logger + +from vllm_omni.diffusion.models.pi05.config import resolve_excluded_action_indices +from vllm_omni.diffusion.models.pi05.modeling_pi05 import ( + DEFAULT_IMAGE_RESOLUTION, + DEFAULT_MAX_TOKEN_LEN, + DEFAULT_STATE_NUM_BINS, +) + +logger = init_logger(__name__) + +# LeRobot NormalizerProcessorStep.eps. +NORM_EPS = 1e-8 + + +# ────────────────────────────────────────────────────────────────────── +# Image preprocessing (identical to π0 — SigLIP is unchanged in π0.5) +# ────────────────────────────────────────────────────────────────────── +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 Pi05ImageProcessor: + """Minimal image preprocessor: image → normalized + padded ``[-1,1]`` tensor.""" + + def __init__(self, image_size: int = DEFAULT_IMAGE_RESOLUTION[0]): + self.image_size = image_size + + def preprocess_single(self, image: Any) -> torch.Tensor: + """Convert one LeRobot-domain image to a normalized model tensor. + + PIL and uint8 inputs have domain ``[0, 255]``. Floating ndarray/tensor + inputs must be finite and in ``[0, 1]``. Values are converted + deterministically to ``[-1, 1]``; the pixel content is never used to + guess its input domain. + """ + 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): + if image.ndim != 3 or image.shape[-1] != 3: + raise ValueError(f"Expected an HWC ndarray with 3 channels, got shape {image.shape}.") + t = torch.from_numpy(np.ascontiguousarray(image)).permute(2, 0, 1).unsqueeze(0) + elif isinstance(image, torch.Tensor): + t = image.unsqueeze(0) if image.ndim == 3 else image + if t.ndim != 4 or t.shape[0] != 1 or t.shape[1] != 3: + raise ValueError(f"Expected a CHW tensor (optionally batched once), got shape {tuple(t.shape)}.") + else: + raise TypeError(f"Unsupported image type for π0.5 preprocessing: {type(image)}") + + if t.dtype == torch.uint8: + return t.to(dtype=torch.float32) / 255.0 * 2.0 - 1.0 + if not t.is_floating_point(): + raise TypeError(f"π0.5 images must use uint8 or floating dtype, got {t.dtype}.") + if not torch.isfinite(t).all(): + raise ValueError("Floating π0.5 image contains NaN or Inf.") + low, high = float(t.amin()), float(t.amax()) + if low < 0.0 or high > 1.0: + raise ValueError(f"Floating π0.5 images must be in [0, 1]; observed min={low}, max={high}.") + return t.to(dtype=torch.float32) * 2.0 - 1.0 + + 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) + + +# ────────────────────────────────────────────────────────────────────── +# Normalization (LeRobot NormalizerProcessorStep) +# ────────────────────────────────────────────────────────────────────── +@dataclass(frozen=True) +class NormStats: + """The affine map a checkpoint's statistics define. + + ``mean_std`` carries mean and std; the two range modes carry the lower and + upper bound (``min``/``max``, or the ``q01``/``q99`` quantiles π0.5 ships). + """ + + mode: str + lower: torch.Tensor + upper: torch.Tensor + + +def build_norm_stats(norm_stats: dict | None, key: str) -> NormStats | None: + """Read ``norm_stats[key]`` into tensors, or ``None`` when it carries none. + + The mode comes from the entry. A LeRobot state_dict ships mean, std, min, + max, q01 and q99 at once, so the statistic names present cannot select it. + Missing modes default to quantile, which is what LeRobot defaults π0.5's + STATE and ACTION to (π0 defaults to mean_std). + """ + entry = (norm_stats or {}).get(key) + if not entry: + return None + + mode = str(entry.get("mode", "quantile")).lower() + if mode == "mean_std": + names = ("mean", "std") + elif mode == "min_max": + names = ("min", "max") + elif mode == "quantile": + mode, names = "min_max", ("q01", "q99") + else: + raise ValueError(f"Unsupported normalization mode for {key!r}: {mode!r}.") + + bounds = [entry.get(name) for name in names] + if any(bound is None for bound in bounds): + raise ValueError(f"Normalization mode {mode!r} for {key!r} requires {names[0]!r} and {names[1]!r}.") + lower, upper = (torch.as_tensor(bound, dtype=torch.float32) for bound in bounds) + return NormStats(mode=mode, lower=lower, upper=upper) + + +def apply_norm(x: torch.Tensor, stats: NormStats | None, *, inverse: bool = False) -> torch.Tensor: + """Map raw units to the model's normalized space, or back with ``inverse``. + + Actions are padded to ``max_action_dim`` while the statistics cover only the + real width, so the tail passes through untouched. The eps rules are + LeRobot's: ``mean_std`` always divides by ``std + eps``, and the range modes + substitute ``eps`` only for an exactly zero range. + """ + if stats is None: + return x + lower = stats.lower.to(device=x.device, dtype=x.dtype) + upper = stats.upper.to(device=x.device, dtype=x.dtype) + valid = lower.shape[0] + head = x[..., :valid] + + if stats.mode == "mean_std": + head = head * upper + lower if inverse else (head - lower) / (upper + NORM_EPS) + else: + span = upper - lower + span = torch.where(span == 0, NORM_EPS, span) + head = (head + 1.0) * 0.5 * span + lower if inverse else 2.0 * (head - lower) / span - 1.0 + + if valid == x.shape[-1]: + return head + return torch.cat([head, x[..., valid:]], dim=-1) + + +# ────────────────────────────────────────────────────────────────────── +# State: normalize → discretize → prompt (π0.5's defining path) +# ────────────────────────────────────────────────────────────────────── +def _pad_or_truncate(raw_state: Any, width: int) -> np.ndarray: + """Zero-pad / truncate a vector to ``(width,)`` float32.""" + if raw_state is None: + return np.zeros((width,), dtype=np.float32) + if isinstance(raw_state, torch.Tensor): + raw_state = raw_state.detach().cpu().numpy() + state = np.asarray(raw_state, dtype=np.float32).reshape(-1) + if state.shape[0] < width: + state = np.pad(state, (0, width - state.shape[0])) + elif state.shape[0] > width: + state = state[:width] + return state.astype(np.float32) + + +def as_state_vector(raw_state: Any, state_dim: int) -> np.ndarray: + """Coerce a request's raw state to ``(state_dim,)`` float32, or raise. + + LeRobot's ``Pi05PrepareStateTokenizerProcessorStep`` discretizes whatever + width it is handed and never pads to ``max_state_dim``, so zero-filling a + 7-dim state up to 32 would append 25 state tokens the checkpoint never saw + in training. + """ + if raw_state is None: + raise ValueError("π0.5 requires a state; there is no default to fall back on.") + if isinstance(raw_state, torch.Tensor): + raw_state = raw_state.detach().cpu().numpy() + state = np.asarray(raw_state, dtype=np.float32) + if state.ndim == 2 and state.shape[0] == 1: + state = state[0] + if state.ndim != 1: + raise ValueError(f"π0.5 state must be shaped (D,) or (1, D); got {tuple(state.shape)}.") + if not np.isfinite(state).all(): + raise ValueError("π0.5 state contains NaN or Inf.") + if state.shape[0] != state_dim: + raise ValueError( + f"π0.5 state has {state.shape[0]} dimension(s), but the checkpoint declares " + f"{state_dim}. The state is serialized into the prompt at its real width, so " + "padding or truncating it would change the tokens the model sees." + ) + return state + + +def discretize_state(state: np.ndarray, *, num_bins: int = DEFAULT_STATE_NUM_BINS) -> np.ndarray: + """Discretize a ``[-1, 1]`` state into ``num_bins`` integer bins. + + Byte-for-byte LeRobot's ``processor_pi05.py``:: + + np.digitize(state_np, bins=np.linspace(-1, 1, 256 + 1)[:-1]) - 1 + + A state below ``-1`` lands in bin ``-1``, and that negative bin is part of + the contract: the checkpoint was trained with ``" -1"`` in the state prompt + for those dimensions. Clipping it to ``0`` changes the tokens the model + sees, so this must not clip. + + ``bins`` stays float64, matching LeRobot's default ``linspace`` dtype, so + boundary values fall on the same side. + + **Assumes the state is already normalized** — see the module docstring. + """ + bins = np.linspace(-1.0, 1.0, num_bins + 1)[:-1] + return (np.digitize(np.asarray(state, dtype=np.float32), bins=bins) - 1).astype(np.int64) + + +def build_pi05_prompt( + *, + task: str, + normalized_state: np.ndarray, + state_num_bins: int = DEFAULT_STATE_NUM_BINS, +) -> str: + """Build the π0.5 prompt: instruction + serialized discretized state. + + Matches LeRobot's ``Pi05PrepareStateTokenizerProcessorStep``, including the + task cleanup (``strip``, ``_`` → space, newline → space), the exact template + and the state values it serializes. The template already ends in a newline, + so — unlike π0 — there is no separate newline-appending step. + + ``normalized_state`` comes from ``Pi05ForActionPrediction._normalize_state``. + """ + cleaned_task = (task or "").strip().replace("_", " ").replace("\n", " ") + bins = discretize_state(normalized_state, num_bins=state_num_bins) + state_str = " ".join(str(int(x)) for x in bins.tolist()) + return f"Task: {cleaned_task}, State: {state_str};\nAction: " + + +def tokenize_prompt(tokenizer, text: str, max_token_len: int = DEFAULT_MAX_TOKEN_LEN): + """Return ``(input_ids, attention_mask)`` lists, length exactly ``max_token_len``. + + ``padding="max_length"`` is what makes the prefix a constant shape: the text + segment is always ``max_token_len`` tokens regardless of the instruction, so + only ``attention_mask.sum()`` varies per request. + """ + enc = tokenizer( + text, + 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"]) + + +# ────────────────────────────────────────────────────────────────────── +# Relative actions +# ────────────────────────────────────────────────────────────────────── +class Pi05RelativeActions: + """The relative/absolute action transform, as a single paired object. + + LeRobot builds one ``RelativeActionsProcessorStep`` and hands *the same + instance* to ``AbsoluteActionsProcessorStep``; the two directions must + agree on ``enabled`` and on which dimensions are excluded, so they are one + object here too. + + **Deviation from LeRobot, on purpose.** LeRobot's step keeps the reference + state on ``self`` between the pre- and post-pass. That is safe for a + single-threaded training loop and unsafe for a server: two in-flight + requests would share one reference state and silently corrupt each other's + actions. Here the state is passed explicitly to :meth:`to_absolute`, so the + object stays immutable after construction and is safe to share across + requests. + + Transform (LeRobot / OpenPI): ``relative = action - state`` on the way in, + ``absolute = relative + state`` on the way out, applied only to the + dimensions *not* named in ``exclude_joints``. Gripper open/close is an + absolute command, which is why it is excluded by default. + """ + + def __init__( + self, + *, + enabled: bool, + exclude_joints: list[str] | None = None, + action_names: list[str] | None = None, + max_action_dim: int = 32, + ): + self.enabled = bool(enabled) + self.exclude_joints = list(exclude_joints or []) + self.action_names = list(action_names) if action_names else None + self.max_action_dim = int(max_action_dim) + + # Boolean mask over action dims: True = this dim is relative to state. + mask = np.ones((self.max_action_dim,), dtype=bool) + if self.enabled: + for idx in resolve_excluded_action_indices(self.exclude_joints, self.action_names): + if 0 <= idx < self.max_action_dim: + mask[idx] = False + # Padding beyond the real action dimensions carries no signal; + # leaving it "relative" would add state noise into dead channels. + if self.action_names: + mask[len(self.action_names) :] = False + else: + mask[:] = False + self.relative_mask = mask + + @property + def num_relative_dims(self) -> int: + return int(self.relative_mask.sum()) + + def _state_row(self, state: Any, device, dtype) -> torch.Tensor: + """Raw state → ``(B, max_action_dim)`` aligned with the action dims. + + Accepts one state for the whole batch (``(D,)``, the serving path, + where the pipeline runs B=1) or one state per sample (``(B, D)``). + LeRobot caches the batched state and shifts each sample by its own row, + so the per-sample form has to be honoured: ``_pad_or_truncate`` + flattens, which would silently reduce ``(B, D)`` to sample 0's state and + apply it to every sample — wrong answers, no error. + """ + if isinstance(state, torch.Tensor): + state = state.detach().cpu().numpy() + arr = np.asarray(0.0 if state is None else state, dtype=np.float32) + if arr.ndim > 2: + raise ValueError(f"Expected state shaped (D,) or (B, D), got {tuple(arr.shape)}") + rows = arr if arr.ndim == 2 else arr[None, :] + padded = np.stack([_pad_or_truncate(row, self.max_action_dim) for row in rows]) + return torch.as_tensor(padded, device=device, dtype=dtype) + + def to_relative(self, actions: torch.Tensor, state: Any) -> torch.Tensor: + """``absolute → relative``. Input side (step 3). + + Not used on the inference path — there are no input actions to convert + at serving time — but it is what the transform *means*, and the parity + test exercises it as the inverse of :meth:`to_absolute`. + """ + if not self.enabled: + return actions + return self._shift(actions, state, sign=-1.0) + + def to_absolute(self, actions: torch.Tensor, state: Any) -> torch.Tensor: + """``relative → absolute``. Output side (step 2 of the post-pipeline). + + ``state`` must be the **raw** state — the same one the model was given + before normalization — because relative actions live in raw action space. + """ + if not self.enabled: + return actions + return self._shift(actions, state, sign=+1.0) + + def _shift(self, actions: torch.Tensor, state: Any, *, sign: float) -> torch.Tensor: + if actions.ndim != 3: + raise ValueError(f"Expected actions shaped (B, horizon, action_dim), got {tuple(actions.shape)}") + action_dim = actions.shape[-1] + if action_dim != self.max_action_dim: + raise ValueError( + f"Action dim {action_dim} does not match max_action_dim={self.max_action_dim}; " + "the relative mask would be misaligned." + ) + state_row = self._state_row(state, actions.device, actions.dtype) # (1, D) + mask = torch.as_tensor(self.relative_mask, device=actions.device) + delta = torch.where(mask, state_row, torch.zeros_like(state_row)) + # Broadcast over the action horizon: every step of the chunk is + # expressed relative to the same current state. + return actions + sign * delta[:, None, :] + + +# ────────────────────────────────────────────────────────────────────── +# Model input assembly +# ────────────────────────────────────────────────────────────────────── +def _extract_images(robot_obs: dict, config) -> dict[str, Any]: + """Pull a ``{feature_key: image}`` map out of a raw robot obs. + + This is functional step 1 (``rename_observations``): keys are translated + through ``config.image_key_map`` so serving wire names map onto the + checkpoint's ``input_features`` identities. + """ + images = robot_obs.get("images") + if not isinstance(images, dict): + images = {k: v for k, v in robot_obs.items() if _is_image_like(v)} + key_map = config.image_key_map or {} + return {key_map.get(k, k): v for k, v in images.items() if _is_image_like(v)} + + +def _is_image_like(value: Any) -> bool: + if isinstance(value, (Image.Image, torch.Tensor)): + return True + if isinstance(value, np.ndarray): + return value.ndim >= 3 + if isinstance(value, (list, tuple)): + try: + return np.asarray(value).ndim >= 3 + except Exception: # noqa: BLE001 + return False + return False + + +def _assemble_model_inputs(robot_obs: dict, config, tokenizer, device: torch.device, state_norm: NormStats | None): + """Convert a raw robot observation into ``sample_actions`` inputs. + + Returns ``(images, image_masks, lang_tokens, lang_masks)`` — note there is + **no state tensor**: π0.5 carries the state inside ``lang_tokens``. + + Camera slots follow ``config.image_feature_keys`` order. A missing key keeps + its semantic slot and receives an empty image with a false mask. The output + always contains exactly ``max_cameras`` slots for the deployed model. + """ + image_size = int(config.image_resolution[0]) + img_proc = Pi05ImageProcessor(image_size=image_size) + max_cameras = config.max_cameras + + feature_keys = config.image_feature_keys or [] + obs_images = _extract_images(robot_obs, config) + if not feature_keys: + # No declared camera order — fall back to whatever the obs provides, + # preserving insertion order. + feature_keys = list(obs_images.keys())[:max_cameras] + + camera_keys = feature_keys[:max_cameras] + if all(obs_images.get(key) is None for key in camera_keys): + raise ValueError("π0.5 observation must provide at least one configured camera image.") + if not config.image_feature_keys and len(obs_images) > max_cameras: + raise ValueError( + f"π0.5 observation provides {len(obs_images)} cameras, which exceeds max_cameras={max_cameras}." + ) + + images: list[torch.Tensor] = [] + image_masks: list[torch.Tensor] = [] + for key in camera_keys: + image = obs_images.get(key) + if image is None: + tensor = img_proc.make_empty_image().to(device=device) + mask = False + else: + tensor = img_proc.preprocess_single(image).to(device=device) + mask = True + images.append(tensor) + image_masks.append(torch.tensor([mask], dtype=torch.bool, device=device)) + + # A checkpoint may declare fewer named slots than the serving graph. + while len(images) < max_cameras: + images.append(img_proc.make_empty_image().to(device=device)) + image_masks.append(torch.tensor([False], dtype=torch.bool, device=device)) + + # Normalize and serialize the state into the tokenized prompt. + raw_state = as_state_vector(robot_obs.get("state"), config.state_dim) + normalized_state = apply_norm(torch.from_numpy(raw_state), state_norm).numpy() + prompt = build_pi05_prompt( + task=robot_obs.get("prompt", "") or "", + normalized_state=normalized_state, + state_num_bins=config.state_num_bins, + ) + 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) + + return images, image_masks, lang_tokens, lang_masks + + +# ────────────────────────────────────────────────────────────────────── +# The interface the pipeline uses +# ────────────────────────────────────────────────────────────────────── +class Pi05Processor: + """Everything between the OpenPI wire and the model. + + Owns the normalization statistics and the relative-action transform, so the + pipeline holds no preprocessing state of its own. + """ + + def __init__(self, config, tokenizer, device: torch.device): + self.config = config + self.tokenizer = tokenizer + self.device = device + self._state_norm = build_norm_stats(config.norm_stats, "state") + self._action_norm = build_norm_stats(config.norm_stats, "action") + self.relative_actions = Pi05RelativeActions( + enabled=config.use_relative_actions, + exclude_joints=config.relative_exclude_joints, + action_names=config.action_feature_names, + max_action_dim=config.max_action_dim, + ) + if self._action_norm is None: + logger.info("π0.5: no action normalization stats; returned actions stay in normalized space.") + if self.relative_actions.enabled: + logger.info( + "π0.5: relative actions enabled — %d of %d action dims are state-relative (excluded joints: %s).", + self.relative_actions.num_relative_dims, + config.max_action_dim, + config.relative_exclude_joints, + ) + + def build_model_inputs(self, robot_obs: dict): + """Raw observation → ``sample_actions`` inputs. + + No state tensor comes back: π0.5 carries the state inside ``lang_tokens``. + """ + return _assemble_model_inputs(robot_obs, self.config, self.tokenizer, self.device, self._state_norm) + + def build_model_outputs(self, actions: torch.Tensor, robot_obs: dict) -> np.ndarray: + """Model output → the action chunk the robot receives. + + ``AbsoluteActionsProcessorStep`` runs after unnormalization, because a + relative-action checkpoint's statistics are computed in relative space. + The raw state is the one the prompt encoded, before normalization. + """ + actions = apply_norm(actions, self._action_norm, inverse=True) + if self.relative_actions.enabled: + actions = self.relative_actions.to_absolute(actions, robot_obs.get("state")) + return actions.squeeze(0)[..., : self.config.action_dim].float().cpu().numpy() diff --git a/vllm_omni/diffusion/models/pi05_pipeline_config.py b/vllm_omni/diffusion/models/pi05_pipeline_config.py new file mode 100644 index 000000000..27035d509 --- /dev/null +++ b/vllm_omni/diffusion/models/pi05_pipeline_config.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Lightweight pipeline topology for π0.5.""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +# π0.5 is one diffusion stage (robot observation -> action chunk via flow +# matching). It is registered in ``OMNI_PIPELINES`` because π0.5 is online-served +# via ``vllm serve --deploy-config pi05.yaml``. Keep this topology outside the +# ``pi05`` package: importing a package submodule executes ``pi05.__init__``, +# which imports the runtime pipeline and ``diffusion.data``. +PI05_PIPELINE = PipelineConfig( + model_type="pi05", + model_arch="Pi05Pipeline", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="diffusion", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(), + final_output=True, + final_output_type="action", + model_arch="Pi05Pipeline", + ), + ), +) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 26b052a3a..f94952b1d 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -206,6 +206,11 @@ _DIFFUSION_MODELS = { "pipeline_pi0", "Pi0Pipeline", ), + "Pi05Pipeline": ( + "pi05", + "pipeline_pi05", + "Pi05Pipeline", + ), "LongCatImageEditPipeline": ( "longcat_image", "pipeline_longcat_image_edit", @@ -382,6 +387,7 @@ _NO_CACHE_ACCELERATION = { # to it; list it here so a stray cache_backend override disables gracefully # instead of erroring. "Pi0Pipeline", + "Pi05Pipeline", "LingBotWorldCausalDMDPipeline", } @@ -586,6 +592,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "MingImagePipeline": "get_ming_image_post_process_func", "InternVLAA1Pipeline": "get_internvla_a1_post_process_func", "Pi0Pipeline": "get_pi0_post_process_func", + "Pi05Pipeline": "get_pi05_post_process_func", "LongCatImageEditPipeline": "get_longcat_image_post_process_func", "StableDiffusion3Pipeline": "get_sd3_image_post_process_func", "FluxKontextPipeline": "get_flux_kontext_post_process_func", diff --git a/vllm_omni/entrypoints/openpi/serving.py b/vllm_omni/entrypoints/openpi/serving.py index 11fde2bf8..450cfbb79 100644 --- a/vllm_omni/entrypoints/openpi/serving.py +++ b/vllm_omni/entrypoints/openpi/serving.py @@ -194,6 +194,9 @@ class ServingRealtimeRobotOpenPI: # (e.g. a policy deploy yaml's ``extra_args``) and layer the OpenPI # protocol fields on top. seed = obs.pop("seed", None) + # Nested so engine knobs cannot collide with robot-defined obs keys. + sampling = obs.get("sampling_params") or {} + robot_obs = {key: value for key, value in obs.items() if key != "sampling_params"} sampling_params = OmniDiffusionSamplingParams() for default_params in get_default_sampling_params_list(self.engine_client): if isinstance(default_params, OmniDiffusionSamplingParams): @@ -205,15 +208,16 @@ class ServingRealtimeRobotOpenPI: { "reset": reset, "session_id": session_id, - "robot_obs": obs, + "robot_obs": robot_obs, } ) prompt = obs.get("prompt", "") - sampling_params = OmniDiffusionSamplingParams( - seed=int(seed) if seed is not None else None, - extra_args=extra_args, - ) + if seed is not None: + sampling_params.seed = int(seed) + if "num_inference_steps" in sampling: + sampling_params.num_inference_steps = sampling["num_inference_steps"] + sampling_params.extra_args = extra_args return OmniDiffusionRequest( prompt=prompt, sampling_params=sampling_params,