[Model] Add π0.5 (Pi0.5) VLA model support (#6950)

Signed-off-by: chenchaox <chenchaox@nvidia.com>
2026-09-17 05:18:51 -07:00
1 parent 873e9bff7c
commit d0755b5716
@@ -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)