[Model] Add π0.5 (Pi0.5) VLA model support (#6950)
Signed-off-by: chenchaox <chenchaox@nvidia.com>
This commit is contained in:
Chenchao Xu authored and GitHub committed
1 parent
873e9bff7c
commit
d0755b5716
22 files changed
+4342
-6
No files matched your search
@@ -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
|
||||
|
||||
@@ -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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
|
||||
@@ -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 |
|
||||
|
||||
@@ -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": "<uuid>", # 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: <https://www.physicalintelligence.company/blog/pi05>
|
||||
- LeRobot π0.5 docs: <https://github.com/huggingface/lerobot/blob/main/docs/source/policy_pi05_README.md>
|
||||
- Checkpoint: <https://huggingface.co/lerobot/pi05_base>
|
||||
- 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)
|
||||
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
@@ -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)
|
||||
File diff suppressed because it is too large.
Load diff
@@ -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"
|
||||
@@ -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)
|
||||
|
||||
@@ -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()))
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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 ``"<feature_name>.<stat_name>"``.
|
||||
_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,
|
||||
)
|
||||
@@ -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
|
||||
@@ -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 "<missing-model-dir>", "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)})
|
||||
@@ -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: <instruction>, State: <b0> <b1> ... <bN>;\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()
|
||||
@@ -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",
|
||||
),
|
||||
),
|
||||
)
|
||||
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in new issue
Block a user