Fine Tuning Serving Openpi

OpenPI Fine-Tuning and Serving

End-to-end workflows for fine-tuning and serving Physical Intelligence's OpenPI models (pi0, pi0-fast, pi0.5) on robot manipulation tasks from the public openpi repository. Covers blank-machine setup, JAX training, PyTorch training, checkpoint conversion, and policy inference serving.

Quick start

Clone the public repo, install the workspace, then serve a pretrained policy:

git clone --recurse-submodules https://github.com/Physical-Intelligence/openpi.git
cd openpi
GIT_LFS_SKIP_SMUDGE=1 uv sync
GIT_LFS_SKIP_SMUDGE=1 uv pip install -e .
uv run scripts/serve_policy.py --env DROID

from openpi_client import websocket_client_policy

client = websocket_client_policy.WebsocketClientPolicy(host="localhost", port=8000)
result = client.infer(observation)
actions = result["actions"]  # numpy array of shape (chunk_size, action_dim)

Core concepts

Model family: OpenPI implements three model variants from Physical Intelligence:

Key design choices:

• Dual backend: JAX (primary, official training) and PyTorch (community, deployment-friendly)
• Config-driven: All training/serving parameters defined in src/openpi/training/config.py
• Norm stats: Every config requires precomputed normalization statistics before training
• WebSocket serving: Policy servers expose a WebSocket API for low-latency inference

Training loop invariant: After every config or dataset change, always re-run this cycle: 1. Compute norm stats → 2. Train → 3. Serve checkpoint → 4. Validate inference

Compute requirements

Workflow 0: Blank-machine setup

Copy this checklist and track progress:

Setup Progress:
- [ ] Step 1: Clone the public openpi repo with submodules
- [ ] Step 2: Install uv and sync the workspace
- [ ] Step 3: Install the editable package
- [ ] Step 4: Verify core imports and serving entrypoint

Step 1: Clone repo

git clone --recurse-submodules https://github.com/Physical-Intelligence/openpi.git
cd openpi

If you already cloned without submodules:

git submodule update --init --recursive

Step 2: Sync dependencies

GIT_LFS_SKIP_SMUDGE=1 uv sync

Step 3: Install editable package

GIT_LFS_SKIP_SMUDGE=1 uv pip install -e .

Step 4: Verify installation

uv run python -c "from openpi.training import config as _config; print(_config.get_config('pi05_droid').name)"
uv run scripts/serve_policy.py --help

When to use vs alternatives

Use this skill when:

• Fine-tuning pi0, pi0-fast, or pi0.5 on LeRobot or RLDS datasets
• Serving OpenPI policies for ALOHA, DROID, or LIBERO evaluation
• Converting JAX checkpoints to PyTorch format
• Debugging OpenPI training issues (norm stats, memory, config)

Use `fine-tuning-openvla-oft` instead when:

• Fine-tuning OpenVLA with continuous action heads and LoRA
• Reproducing OpenVLA-OFT paper results on LIBERO or ALOHA

Use `evaluating-cosmos-policy` instead when:

• Evaluating NVIDIA Cosmos Policy on simulation benchmarks

---

Workflow 1: JAX fine-tuning on LeRobot data

Copy this checklist and track progress:

JAX Fine-Tuning Progress:
- [ ] Step 1: Select and copy closest training config
- [ ] Step 2: Update dataset mapping and base checkpoint
- [ ] Step 3: Compute normalization statistics
- [ ] Step 4: Launch JAX training
- [ ] Step 5: Serve checkpoint and run inference sanity check

Step 1: Select config

Copy the closest config from src/openpi/training/config.py:

Step 2: Update dataset and transforms

# In src/openpi/training/config.py, modify your config:
TrainConfig(
    name="my_custom_config",
    model_type="pi05",
    data=LeRobotDataConfig(
        repo_id="your-org/your-dataset",
        # Adjust transforms to match your data format
    ),
    weight_loader=Pi05WeightLoader(),  # Match model type
)

Set repo_id for your dataset and ensure weight_loader matches the model type (pi0 vs pi0.5).

Step 3: Compute normalization statistics

uv run scripts/compute_norm_stats.py --config-name <config_name>

This must run before every training launch when config, dataset, or transforms change.

Step 4: Launch JAX training

XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 uv run scripts/train.py <config_name> \
  --exp-name=<run_name> \
  --overwrite

For full DROID RLDS training, add the rlds dependency group:

uv run --group rlds scripts/compute_norm_stats.py \
  --config-name pi05_full_droid_finetune \
  --max-frames 10000000

XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 uv run --group rlds scripts/train.py \
  pi05_full_droid_finetune \
  --exp-name=<run_name> --overwrite

Step 5: Serve and validate

uv run scripts/serve_policy.py policy:checkpoint \
  --policy.config=<config_name> \
  --policy.dir=checkpoints/<config_name>/<run_name>/<step>

Verify with a test client:

from openpi_client import websocket_client_policy

client = websocket_client_policy.WebsocketClientPolicy(host="localhost", port=8000)
# Build observation matching your config's expected keys
obs = {"image": img_array, "state": state_array, "prompt": "pick up the cup"}
result = client.infer(obs)
print(f"Action shape: {result['actions'].shape}")  # (chunk_size, action_dim)

---

Workflow 2: PyTorch training and checkpoint conversion

Copy this checklist and track progress:

PyTorch Setup Progress:
- [ ] Step 1: Sync dependencies and verify transformer version
- [ ] Step 2: Apply OpenPI transformer patches
- [ ] Step 3: Convert JAX checkpoint to PyTorch format
- [ ] Step 4: Launch PyTorch training or serve converted checkpoint

Step 1: Sync dependencies

uv sync
uv pip show transformers

Step 2: Apply required patches

OpenPI PyTorch requires custom modifications to the installed transformers package:

cp -r ./src/openpi/models_pytorch/transformers_replace/* \
  .venv/lib/python3.11/site-packages/transformers/

Step 3: Convert JAX checkpoint

uv run examples/convert_jax_model_to_pytorch.py \
  --checkpoint_dir <jax_checkpoint_dir> \
  --config_name <config_name> \
  --output_path <pytorch_checkpoint_dir>

Step 4: Train or serve

Single GPU training:

uv run scripts/train_pytorch.py <config_name> --exp_name <run_name>

Multi-GPU distributed training:

uv run torchrun --standalone --nnodes=1 --nproc_per_node=<num_gpus> \
  scripts/train_pytorch.py <config_name> --exp_name <run_name>

Programmatic inference with converted checkpoint:

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

config = _config.get_config("pi05_droid")
policy = policy_config.create_trained_policy(config, "<pytorch_checkpoint_dir>")
result = policy.infer(example)
actions = result["actions"]  # numpy array

Checkpoints follow the convention: checkpoints/<config_name>/<exp_name>/<step>/.

---

Workflow 3: Policy inference serving

Copy this checklist and track progress:

Inference Server Progress:
- [ ] Step 1: Choose target environment and checkpoint
- [ ] Step 2: Start policy server
- [ ] Step 3: Confirm server is reachable
- [ ] Step 4: Integrate client into robot or simulation code

Step 1: Choose environment

Default environment presets:

Step 2: Start server

Default mode (uses preset checkpoint):

uv run scripts/serve_policy.py --env ALOHA

Explicit checkpoint mode (custom or local model):

uv run scripts/serve_policy.py policy:checkpoint \
  --policy.config=pi05_libero \
  --policy.dir=checkpoints/pi05_libero/my_run/20000

Add --default_prompt "task description" when runtime observations omit a prompt.

Step 3: Verify connectivity

uv run examples/simple_client/main.py --env DROID

Step 4: Embed remote client in robot code

Install the lightweight client in your robot environment:

pip install "openpi-client @ git+https://github.com/Physical-Intelligence/openpi.git#subdirectory=packages/openpi-client"

Full integration example:

from openpi_client import websocket_client_policy
import numpy as np

# Connect to remote policy server
client = websocket_client_policy.WebsocketClientPolicy(
    host="gpu-server.local", port=8000
)

# Build observation (keys must match policy transforms)
observation = {
    "image": np.random.rand(224, 224, 3),  # RGB image
    "state": np.zeros(7),                   # Joint positions
    "prompt": "pick up the red block",
}

# Get actions
result = client.infer(observation)
actions = result["actions"]  # shape: (action_chunk_size, action_dim)

# Execute first action on robot
robot.step(actions[0])

---

Common issues

Issue: Missing norm stats error

Fix: run scripts/compute_norm_stats.py --config-name <config_name> before training.

Issue: Out of memory during JAX training

Fix: set XLA_PYTHON_CLIENT_MEM_FRACTION=0.9, lower batch size, or configure fsdp_devices:

# In config: use model-parallel sharding
TrainConfig(
    ...
    fsdp_devices=4,  # Shard across 4 GPUs
)

Issue: OOM while loading PyTorch checkpoints

Fix: export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

Issue: Config not found

Fix: ensure config name exists in src/openpi/training/config.py (exact match from _CONFIGS dict).

Issue: PyTorch training diverges after library changes

Fix: reapply the transformer patch. Run uv cache clean transformers to reset, then reapply.

Issue: `serve_policy.py` crashes with `ModuleNotFoundError`

Fix: resync the public workspace first:

GIT_LFS_SKIP_SMUDGE=1 uv sync
GIT_LFS_SKIP_SMUDGE=1 uv pip install -e .

If the missing module is simulator-related, install the extra runtime dependencies called for by that example:

uv pip install pytest robosuite==1.4.0 gym bddl easydict matplotlib

Issue: `uv sync` fails with `rerun-sdk` wheel mismatch

Fix:

uv sync --no-dev
# or
uv sync --no-dev --no-install-package rerun-sdk

Issue: Checkpoint download times out

Fix: install gsutil and prefetch manually:

pip install gsutil
gsutil -m cp -r gs://openpi-assets/checkpoints/pi05_libero /local/cache/

Remove stale .lock files if a previous download was interrupted.

Issue: Policy server exits with code `137`

Fix: OOM kill. Set JAX memory variables:

export XLA_PYTHON_CLIENT_PREALLOCATE=false
export XLA_PYTHON_CLIENT_ALLOCATOR=platform

---

For HPC/cluster users

On Slurm-managed clusters, wrap commands with resource allocation:

srun --partition=gpu --gpus-per-node=1 --mem=64G --cpus-per-task=8 --pty bash

Route caches to scratch to avoid filling /home:

export HF_HOME=/scratch/$USER/.cache/huggingface
export XDG_CACHE_HOME=/scratch/$USER/.cache
export PIP_CACHE_DIR=/scratch/$USER/.cache/pip
export UV_CACHE_DIR=/scratch/$USER/.cache/uv

Avoid stacking cluster Python modules when using uv-managed environments. Typically module load cuda is sufficient.

---

Advanced topics

Config recipes and baselines: See references/config-recipes.md Training debugging guide: See references/training-debugging.md Checkpoint and environment mapping: See references/checkpoints-and-env-map.md Remote client integration: See references/remote-client-pattern.md PyTorch precision and patching gotchas: See references/pytorch-gotchas.md

Resources

• OpenPI repository: https://github.com/Physical-Intelligence/openpi
• OpenPI client package: https://github.com/Physical-Intelligence/openpi/tree/main/packages/openpi-client
• pi0 paper: https://www.physicalintelligence.company/blog/pi0
• LeRobot dataset format: https://huggingface.co/docs/lerobot

Checkpoints and Environment Map

Use default environment mode for first runs, then switch to explicit checkpoint mode when needed.

Default mapping from scripts/serve_policy.py

Common explicit checkpoint commands

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

# PI 0 FAST DROID
uv run scripts/serve_policy.py policy:checkpoint \
  --policy.config=pi0_fast_droid \
  --policy.dir=gs://openpi-assets/checkpoints/pi0_fast_droid

# PI 0.5 LIBERO
uv run scripts/serve_policy.py policy:checkpoint \
  --policy.config=pi05_libero \
  --policy.dir=gs://openpi-assets/checkpoints/pi05_libero

Local checkpoint command template

uv run scripts/serve_policy.py policy:checkpoint \
  --policy.config=<config_name> \
  --policy.dir=checkpoints/<config_name>/<exp_name>/<step>

Data home and caching

• OpenPI downloads and caches assets under ~/.cache/openpi by default.
• Set OPENPI_DATA_HOME to move download/cache location.

LIBERO checkpoint prefetch on clusters

If policy server startup times out while logs show checkpoint downloading:

# 1) Ensure gsutil exists
pip install gsutil

# 2) Clear stale lock from previous interrupted download
rm -f <OPENPI_DATA_HOME>/openpi-assets/checkpoints/pi05_libero.lock

# 3) Prefetch checkpoint manually
cd <OPENPI_DATA_HOME>/openpi-assets/checkpoints
gsutil -m cp -r gs://openpi-assets/checkpoints/pi05_libero .

Cluster compatibility notes (uv + Slurm)

If uv sync fails with rerun-sdk wheel/platform mismatch:

# 1) Skip dev groups
uv sync --no-dev

# 2) Force skip incompatible package
uv sync --no-dev --no-install-package rerun-sdk

For shared clusters with small /home, point cache roots to scratch:

• HF_HOME, XDG_CACHE_HOME, PIP_CACHE_DIR, UV_CACHE_DIR, TMPDIR

Runtime hotfix dependencies for OpenPI + LIBERO

If server startup fails with ModuleNotFoundError:

uv pip install pytest robosuite==1.4.0 gym bddl easydict matplotlib

Install into both the OpenPI server environment and the LIBERO client environment.

Config Recipes

Use these as starting points when choosing a config to copy or adapt.

Common config baselines

Essential command sequence

# 1) Compute normalization stats
uv run scripts/compute_norm_stats.py --config-name <config_name>

# 2) Train
XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 uv run scripts/train.py <config_name> \
  --exp-name=<run_name> --overwrite

# 3) Serve checkpoint for verification
uv run scripts/serve_policy.py policy:checkpoint \
  --policy.config=<config_name> \
  --policy.dir=checkpoints/<config_name>/<run_name>/<step>

RLDS variant for full DROID

uv run --group rlds scripts/compute_norm_stats.py \
  --config-name pi05_full_droid_finetune --max-frames 10000000

XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 uv run --group rlds scripts/train.py \
  pi05_full_droid_finetune --exp-name=<run_name> --overwrite

High-signal files to inspect while adapting configs

• src/openpi/training/config.py — all config definitions
• src/openpi/policies/libero_policy.py — LIBERO policy transforms
• src/openpi/policies/droid_policy.py — DROID policy transforms
• src/openpi/policies/aloha_policy.py — ALOHA policy transforms

PyTorch Precision and Patching Gotchas

Transformer patch requirement

OpenPI PyTorch requires custom patches applied to the installed transformers package. Training or inference without the patch produces subtle incompatibilities.

Apply patches:

cp -r ./src/openpi/models_pytorch/transformers_replace/* \
  .venv/lib/python3.11/site-packages/transformers/

Verify the patch is active:

Check that modified files in the transformers package directory have recent timestamps matching the patch application.

Patch survives reinstall

If uv sync or pip install reinstalls transformers, the patch is overwritten.

Fix: reapply patches after any dependency reinstall, or run:

uv cache clean transformers

Then reapply the patch.

OOM while loading checkpoints

Set memory allocation strategy before loading large models:

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

Resume mode

• --resume requires --exp_name to match the prior run exactly.
• At least one numeric checkpoint directory must exist under checkpoints/<config_name>/<exp_name>/.
• Do not combine --resume with other conflicting flags.

Precision notes

• Default training precision follows the model config.
• When converting from JAX, ensure the output precision matches expectations (bf16 vs fp32).
• Mixed precision settings in PyTorch should align with the source JAX checkpoint precision.

Remote Client Pattern

Use this pattern when the policy server runs on a GPU machine and control code runs elsewhere.

Server side

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

Default port is 8000.

Robot or eval client side

Install client package:

uv pip install -e packages/openpi-client

Call server from Python:

from openpi_client import websocket_client_policy

client = websocket_client_policy.WebsocketClientPolicy(host="server-ip", port=8000)
result = client.infer(observation)
actions = result["actions"]

Observation contract checks

• Pass observation keys expected by your policy transforms.
• Pass prompt as observation["prompt"] or use server --default_prompt.
• Resize image tensors to the expected model input shape before call (typically 224).
• Keep state values in the policy's expected coordinate and ordering conventions.

Read before integration

• docs/remote_inference.md
• examples/simple_client/README.md
• examples/droid/README.md
• examples/aloha_real/README.md

Training Debugging

Use this quick loop during iteration:

1. Confirm config exists and resolves: src/openpi/training/config.py. 2. Recompute norm stats after transform or dataset changes. 3. Run short training smoke test. 4. Serve a recent checkpoint and run inference sanity check.

Common failures and fixes

Issue: `Config '<name>' not found`

Fix: use exact config name from _CONFIGS in src/openpi/training/config.py.

Issue: Missing normalization stats

Fix: run uv run scripts/compute_norm_stats.py --config-name <name> before training.

Issue: OOM on JAX startup or training

Fix:

• Set XLA_PYTHON_CLIENT_MEM_FRACTION=0.9
• Lower batch size
• Use fsdp_devices for model sharding

Issue: No progress after resume request

Fix: ensure checkpoint directory exists and includes numeric step folders.

Issue: Incompatible resume and overwrite settings

Fix: do not set both simultaneously.

Validation commands

# Quick serve validation
uv run scripts/serve_policy.py policy:checkpoint \
  --policy.config=<config_name> \
  --policy.dir=checkpoints/<config_name>/<exp_name>/<step>

# Quick client test
uv run examples/simple_client/main.py --env DROID