Paper Rob Openpi

paper_rob__openpi

OpenPI provides open-source Vision-Language-Action (VLA) foundation models from Physical Intelligence. Three model variants -- pi0, pi0-FAST, and pi0.5 -- share a PaliGemma backbone and differ in their action decoding strategy. Pre-trained on 10k+ hours of multi-platform robot data, they can be fine-tuned on custom tasks with as little as a single GPU.

Paper Info

Method Overview

All three models share a two-stage training recipe analogous to LLM pre-training and post-training:

1. Pre-training on 10k+ hours of robot data across 7 platforms and 68 tasks.
2. Post-training (fine-tuning) to adapt to a specific downstream task or robot.

The backbone is PaliGemma (SigLip So400m/14 vision encoder + Gemma 2B language model). An action expert (a smaller Gemma 300M model) decodes actions conditioned on the VLM's representation. The three variants differ in how the action expert produces actions:

Key insight: Using a separate action expert with flow matching decouples action generation from language understanding, enabling the VLM backbone to retain its semantic knowledge while learning fine-grained motor control.

Paper-Code Mapping

Setup

Dependencies

• Python 3.11+
• JAX (for default training) or PyTorch (>= transformers 4.53.2, for DDP training)
• CUDA 12.x+
• Key packages: flax, optax, orbax-checkpoint, tyro, einops, augmax, wandb, safetensors
• LeRobot (included as git submodule for data loading)

Installation

# Clone with LeRobot submodule
git clone --recurse-submodules https://github.com/Physical-Intelligence/openpi.git
cd openpi
git submodule update --init --recursive

# Install with uv (recommended)
GIT_LFS_SKIP_SMUDGE=1 uv sync
GIT_LFS_SKIP_SMUDGE=1 uv pip install -e .

Docker: see docs/docker.md.

Hardware Requirements

Available Checkpoints

Base models (starting points for fine-tuning):

Fine-tuned experts (ready for inference):

Checkpoints auto-download from GCS on first use via download.maybe_download().

Usage Scenarios

Serve a Trained Policy (Inference)

uv run scripts/serve_policy.py policy:checkpoint \
  --policy.config=pi05_droid \
  --policy.dir=gs://openpi-assets/checkpoints/pi05_droid

WebSocket server starts on port 8000. Health check: GET /healthz.

Fine-Tune with JAX

# 1. Compute normalization stats (required before first train)
uv run scripts/compute_norm_stats.py --config-name pi05_libero

# 2. Train
XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 uv run scripts/train.py pi05_libero \
  --exp-name=my_experiment

Fine-Tune with PyTorch (Multi-GPU)

# Multi-GPU
torchrun --nproc_per_node=8 scripts/train_pytorch.py pi05_libero \
  --exp-name=my_experiment

# Single GPU
python scripts/train_pytorch.py pi05_libero --exp-name=my_experiment

Convert JAX Checkpoint to PyTorch

python examples/convert_jax_model_to_pytorch.py \
  --checkpoint_dir gs://openpi-assets/checkpoints/pi0_base \
  --output_path ./pi0_base_pytorch \
  --config_name pi0_base

Key TrainConfig Fields

Code Integration Guide

Client-Side Inference (Recommended)

The simplest integration path is to run the policy server and connect via openpi-client:

# pip install openpi-client
from openpi_client import websocket_client_policy as wcp

policy = wcp.WebsocketClientPolicy(host="SERVER_IP", port=8000)

action = policy.infer({
    "observation/image": camera_image,       # (H, W, 3) uint8
    "observation/state": robot_state,        # (state_dim,) float32
    "prompt": "pick up the red cup",
})["actions"]  # shape: (horizon, action_dim) float32

Direct Python API (In-Process)

from openpi.training import config as _config
from openpi.policies import policy_config
from openpi.shared import download

# Load config and checkpoint
config = _config.get_config("pi05_droid")
ckpt_dir = download.maybe_download("gs://openpi-assets/checkpoints/pi05_droid")

# create_trained_policy auto-detects JAX vs PyTorch checkpoint
policy = policy_config.create_trained_policy(config, ckpt_dir)

# Run inference
import numpy as np
result = policy.infer({
    "observation/exterior_image_1_left": np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8),
    "observation/wrist_image_left": np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8),
    "observation/joint_position": np.zeros(7, dtype=np.float32),
    "observation/gripper_position": np.zeros(1, dtype=np.float32),
    "prompt": "pick up the block",
})
actions = result["actions"]  # (horizon, action_dim)

Data Format

Input observation dict (after repack transforms):

Standard 16-dim action space (used by ALOHA-type configs):

Data Pipeline

Raw dataset -> Repack transforms -> Data transforms -> Normalize -> Model transforms -> Model
                (key remapping)     (robot-specific)   (z-score     (tokenize, resize,
                                                        or quantile) pad)

Defined in DataConfig with three transforms.Group fields: repack_transforms, data_transforms, model_transforms. Each group has inputs (applied before model) and outputs (applied after model, inference only).

Integration Notes

• The repo uses uv for dependency management; running scripts outside uv requires manual path/dependency setup.
• Image inputs can be uint8 [0, 255] -- Observation.from_dict() auto-converts to float32 [-1, 1].
• The openpi-client package (packages/openpi-client/) is independently installable via pip install openpi-client.
• For PyTorch models, create_trained_policy() auto-detects by checking for model.safetensors in the checkpoint dir.
• Different robot platforms require different observation key names; check the corresponding *_policy.py files and repack transforms.

Config System

Dataclass-based with Tyro CLI parsing (not Hydra/YAML).

from openpi.training import config as _config

# Load a named config
config = _config.get_config("pi05_libero")

# CLI: tyro.extras.overridable_config_cli(...)
config = _config.cli()

All configs are defined in _CONFIGS list in src/openpi/training/config.py. Key data config factories:

Core Architecture

Images ──> SigLip So400m/14 ──┐
                               ├──> PaliGemma (Gemma 2B, prefix-LM attn) ──┐
Language prompt ───────────────┘                                            │
                                                                            v
                                                               ┌─── Action Expert (Gemma 300M) ───┐
                                                               │  pi0:   flow matching             │
                                                               │         state_proj + action_time_  │
                                                               │         mlp merges timestep+action │
                                                               │  pi0.5: adaRMSNorm injects        │
                                                               │         timestep; no state_proj    │
                                                               │  FAST:  autoregressive decoding    │
                                                               │         with FAST tokenizer        │
                                                               └───────────────────────────────────┘
                                                                            │
                                                                  Action chunk output
                                                               (action_horizon x action_dim)

Flow matching (pi0/pi0.5): Time sampled from Beta(1.5, 1), noisy actions x_t = t*noise + (1-t)*actions, velocity target u_t = noise - actions. Inference: Euler ODE integration from t=1 (noise) to t=0 (clean) with KV-cached prefix.

Autoregressive (pi0-FAST): Actions tokenized by FAST tokenizer into discrete tokens, predicted via standard next-token prediction with cross-entropy loss. EOS token (id=1) for early stopping.

Repo Structure

Fine-Tuning Configs Reference

Tips & Gotchas

• Always run compute_norm_stats.py before training on new data. Training will use incorrect normalization without it.
• Set XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 for JAX training to avoid OOM.
• Use GIT_LFS_SKIP_SMUDGE=1 during install -- weights download on demand from GCS.
• LoRA reduces VRAM from ~70 GB to ~22.5 GB with minimal quality loss. Set ema_decay=None when using LoRA.
• Checkpoints auto-download from GCS on first use and cache to ~/.cache/openpi/.
• Pi0.5 generally outperforms pi0 due to knowledge insulation (adaRMSNorm decouples timestep injection).
• For custom robots: provide normalization stats matching your action/state dims. Use AssetsConfig(assets_dir=..., asset_id=...) to reuse existing stats.
• RLDS data loader requires num_workers=0 -- it handles multi-processing internally.
• The max_token_len for pi0-FAST should be ~180 for single-arm robots, ~250 for dual-arm. Too small clips tokens; too large wastes memory.
• Pi0/pi0.5 use action_dim=32 by default (padded); pi0-FAST uses the actual action dim (e.g., 7 for single-arm).
• Flow matching uses t=1 for noise and t=0 for clean data (opposite of the pi0 paper convention).
• Gemma variants: "gemma_2b" (PaliGemma backbone), "gemma_300m" (action expert), "gemma_2b_lora" / "gemma_300m_lora" (LoRA versions), "dummy" (for debugging).
• PyTorch training does not support EMA; use JAX training for EMA-based fine-tuning.
• For DDP PyTorch training, batch size is automatically split across GPUs (effective_batch_size = batch_size // world_size).