vllm-omni/recipes/lerobot/Pi05.md

# 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)