Yichen WangandDidan Deng
2a9f6aaa8e
Signed-off-by: Yichen Wang <yichen.wang@bytedance.com> Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com> Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
105 lines
4.7 KiB
Markdown
105 lines
4.7 KiB
Markdown
# π0 (Pi-Zero) VLA — OpenPI realtime serving
|
|
|
|
[π0](https://www.physicalintelligence.company/blog/pi0) is a Vision-Language-Action
|
|
model from Physical Intelligence: multi-camera images + a language instruction +
|
|
robot proprioceptive state → a continuous action chunk via flow-matching denoising
|
|
(it does **not** emit text tokens). This example serves π0 over the OpenPI realtime
|
|
websocket protocol at `/v1/realtime/robot/openpi`.
|
|
|
|
## Install extras
|
|
|
|
The core `pip install -e .` does not include the OpenPI client used here:
|
|
|
|
- `openpi-client` (from the [openpi](https://github.com/Physical-Intelligence/openpi)
|
|
repo: `pip install -e packages/openpi-client`), `websockets`, `msgpack`, `msgpack-numpy`
|
|
|
|
## Weights
|
|
|
|
π0 uses the LeRobot `lerobot/pi0_base` checkpoint (HF, ~13 GB). Either let the server
|
|
download it (`MODEL=lerobot/pi0_base`) or point at a local copy
|
|
(`MODEL=/path/to/pi0_base`).
|
|
|
|
## Run the server
|
|
|
|
```bash
|
|
vllm serve lerobot/pi0_base --omni --port 8000 \
|
|
--served-model-name pi0 \
|
|
--deploy-config vllm_omni/deploy/pi0.yaml \
|
|
--enforce-eager --disable-log-stats
|
|
```
|
|
|
|
The deploy config (`vllm_omni/deploy/pi0.yaml`) declares a single diffusion stage
|
|
(`Pi0Pipeline`), float32, `max_num_seqs: 1`, and the `policy_server_config` handshake
|
|
metadata (3 cameras, `joint_position`, action horizon 50, action dim 32).
|
|
|
|
### Commonly adjusted `model_config` keys
|
|
|
|
These live under `stages[0].model_config` in `vllm_omni/deploy/pi0.yaml`:
|
|
|
|
| Key | Default | Meaning |
|
|
|---|---|---|
|
|
| `chunk_size` | `50` | Action-chunk length (timesteps) the model predicts per inference. |
|
|
| `num_inference_steps` | `10` | Flow-matching Euler denoising steps. |
|
|
| `max_action_dim` | `32` | Action dimensionality (state/action are padded to this). |
|
|
| `max_state_dim` | `32` | Proprioceptive-state dimensionality (zero-padded to this). |
|
|
| `image_resolution` | `[224, 224]` | Per-camera input size (square; SigLIP). |
|
|
| `tokenizer_max_length` | `48` | Max PaliGemma prompt tokens. |
|
|
| `max_cameras` | `3` | Camera slots the model attends to (real + `-1`-padded). |
|
|
| `image_feature_keys` | 3 `observation.images.*` keys | Camera order the model attends to. |
|
|
| `image_key_map` | `{}` | Map raw obs camera keys → `image_feature_keys` (empty = verbatim). |
|
|
|
|
The `policy_server_config` block below them is the OpenPI handshake metadata
|
|
advertised to the client; keep its `action_horizon` / `action_dim` /
|
|
`image_resolution` in sync with the `model_config` values above. The e2e test
|
|
`tests/e2e/online_serving/test_pi0_expansion.py::test_pi0_openpi_online` connects to a
|
|
live server and asserts the advertised metadata matches these `pi0.yaml` values.
|
|
|
|
## Run the client
|
|
|
|
```bash
|
|
python examples/online_serving/pi0/openpi_client.py --host 127.0.0.1 --port 8000 \
|
|
--prompt "pick up the red block and place it in the bin"
|
|
```
|
|
|
|
It connects, prints the server metadata, sends robot observations, and prints the
|
|
returned `[action_horizon, action_dim] = [50, 32]` action chunks. Replace the blank
|
|
cameras / zero state in `_make_dummy_obs` with real frames (HWC uint8) and
|
|
proprioceptive state to drive a robot.
|
|
|
|
### Observation format
|
|
|
|
The client sends a flat dict per inference:
|
|
|
|
```python
|
|
{
|
|
"observation.images.base_0_rgb": np.uint8[H, W, 3],
|
|
"observation.images.left_wrist_0_rgb": np.uint8[H, W, 3],
|
|
"observation.images.right_wrist_0_rgb":np.uint8[H, W, 3],
|
|
"state": np.float32[state_dim], # zero-padded to max_state_dim=32 server-side
|
|
"prompt": "pick up the red block",
|
|
"session_id": "<uuid>", # accepted but ignored (π0 is stateless)
|
|
}
|
|
```
|
|
|
|
Camera keys must match the server's `image_feature_keys` (the checkpoint's
|
|
`input_features` order). If your robot uses different camera names, set
|
|
`model_config.image_key_map` in `pi0.yaml` to map raw obs keys → feature keys.
|
|
|
|
## Correctness
|
|
|
|
π0's flow-matching kernel is bit-for-bit matched to the LeRobot `PI0Policy` reference
|
|
(`max|Δ| = 7.15e-07`, CPU/float32, fixed noise; see
|
|
`tests/diffusion/models/pi0/test_pi0_parity.py::test_pi0_vllm_omni_vs_lerobot`) and is
|
|
version-stable across transformers releases (the version-stability checks in
|
|
`tests/diffusion/models/pi0/test_pi0_units.py`). An OpenPI websocket online-serving e2e
|
|
lives in `tests/e2e/online_serving/test_pi0_expansion.py::test_pi0_openpi_online`.
|
|
|
|
## Limitations
|
|
|
|
- **Normalization stats**: `lerobot/pi0_base` uses identity normalization, which is
|
|
fully supported (state passes through, actions are returned in the model's space).
|
|
Per-dataset `norm_stats` declared in a checkpoint's `config.json` are honored, but
|
|
stats stored only in LeRobot's `policy_preprocessor.json` companion safetensors are
|
|
not yet loaded — a fine-tuned checkpoint that relies on those would need that bridge
|
|
added before its actions are in real-world units.
|