vllm-omni/examples/online_serving/pi0/README.md

π0 (Pi-Zero) VLA — OpenPI realtime serving

π0 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 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

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

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:

{
    "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.