fine-tuning-serving-openpi

OpenPI 微调与服务部署

提供端到端工作流，用于在公共 openpi 仓库中的机器人操作任务上微调和部署 Physical Intelligence 的 OpenPI 模型（pi0、pi0-fast、pi0.5）。涵盖空白机器配置、JAX 训练、PyTorch 训练、检查点转换和策略推理服务部署。

快速开始

克隆公共仓库，安装工作区，然后部署预训练策略：

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)

核心概念

模型系列：OpenPI 实现了 Physical Intelligence 的三个模型变体：

模型	架构	速度	质量	典型用途
pi0	流匹配 VLA	基准	最高	研究、复杂任务
pi0-fast	自回归动作 Token	快 2-5 倍	良好	实时控制
pi0.5	pi0 + 改进的视觉编码器	基准	最佳	最新默认模型

关键设计选择：

• 双后端：JAX（主要后端，官方训练）和 PyTorch（社区后端，便于部署）
• 配置驱动：所有训练和服务部署参数都在 src/openpi/training/config.py 中定义
• 归一化统计数据：每个配置都需要在训练前预先计算归一化统计数据
• WebSocket 服务部署：策略服务通过 WebSocket API 提供低延迟推理

训练循环不变原则：每次更改配置或数据集后，都要重新运行以下流程：

1. 计算归一化统计数据 → 2. 训练 → 3. 部署检查点 → 4. 验证推理

计算资源要求

任务	GPU	显存	备注
部署 pi0.5（推理）	1x A100/H100	~24 GB	单个 GPU 即可
微调 pi0.5（JAX）	1x A100 80GB	~60 GB	多 GPU 时使用 fsdp_devices
微调 pi0（JAX）	1x A100 80GB	~40 GB	模型占用空间更小
微调（PyTorch DDP）	1-8x A100	~40 GB/GPU	使用 torchrun 启动器
计算归一化统计数据	CPU 或 1x GPU	~8 GB	速度快，可在登录节点运行

工作流 0：空白机器配置

复制此清单并跟踪进度：

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

步骤 1：克隆仓库

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

如果克隆时没有包含子模块：

git submodule update --init --recursive

步骤 2：同步依赖项

GIT_LFS_SKIP_SMUDGE=1 uv sync

步骤 3：安装可编辑软件包

GIT_LFS_SKIP_SMUDGE=1 uv pip install -e .

步骤 4：验证安装

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

何时使用及替代方案

以下情况请使用此技能：

• 在 LeRobot 或 RLDS 数据集上微调 pi0、pi0-fast 或 pi0.5
• 为 ALOHA、DROID 或 LIBERO 评估部署 OpenPI 策略
• 将 JAX 检查点转换为 PyTorch 格式
• 调试 OpenPI 训练问题（归一化统计数据、内存、配置）

以下情况请改用 fine-tuning-openvla-oft：

• 使用连续动作头和 LoRA 微调 OpenVLA
• 在 LIBERO 或 ALOHA 上复现 OpenVLA-OFT 论文结果

以下情况请改用 evaluating-cosmos-policy：

• 在仿真基准上评估 NVIDIA Cosmos Policy

---

工作流 1：在 LeRobot 数据上进行 JAX 微调

复制此清单并跟踪进度：

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

步骤 1：选择配置

从 src/openpi/training/config.py 复制最接近的配置：

配置	用途
pi05_libero	在 LIBERO 上微调 pi0.5
pi0_libero	在 LIBERO 上对 pi0 进行完整微调
pi0_fast_libero	在 LIBERO 上使用 pi0-fast
pi0_aloha_pen_uncap	ALOHA 自定义数据
pi05_droid_finetune	小型自定义 DROID 数据集（LeRobot 格式）
pi05_full_droid_finetune	完整 DROID RLDS 大规模训练

步骤 2：更新数据集和变换

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

为你的数据集设置 repo_id，并确保 weight_loader 与模型类型（pi0 或 pi0.5）匹配。

步骤 3：计算归一化统计数据

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

每次更改配置、数据集或变换后，都必须在启动训练前运行此步骤。

步骤 4：启动 JAX 训练

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

对于完整的 DROID RLDS 训练，请添加 rlds 依赖组：

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

步骤 5：部署并验证

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

使用测试客户端进行验证：

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)

---

工作流 2：PyTorch 训练与检查点转换

复制此清单并跟踪进度：

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

步骤 1：同步依赖项

uv sync
uv pip show transformers

步骤 2：应用所需补丁

OpenPI PyTorch 需要对已安装的 transformers 包进行自定义修改：

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

步骤 3：转换 JAX 检查点

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

步骤 4：训练或提供服务

单 GPU 训练：

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

多 GPU 分布式训练：

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

使用转换后的检查点进行程序化推理：

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/<config_name>/<exp_name>/<step>/。

---

工作流 3：策略推理服务

复制此清单并跟踪进度：

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

步骤 1：选择环境

默认环境预设：

环境	配置	默认检查点
ALOHA	pi05_aloha	gs://openpi-assets/checkpoints/pi05_base
ALOHA_SIM	pi0_aloha_sim	gs://openpi-assets/checkpoints/pi0_aloha_sim
DROID	pi05_droid	gs://openpi-assets/checkpoints/pi05_droid
LIBERO	pi05_libero	gs://openpi-assets/checkpoints/pi05_libero

步骤 2：启动服务端

默认模式（使用预设检查点）：

uv run scripts/serve_policy.py --env ALOHA

显式指定检查点模式（自定义或本地模型）：

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

当运行时观测中未包含提示词时，添加 --default_prompt "task description"。

步骤 3：验证连接

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

步骤 4：在机器人代码中嵌入远程客户端

在机器人环境中安装轻量级客户端：

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

完整集成示例：

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

---

常见问题

问题：缺少 norm stats 错误

解决方法：在训练前运行 scripts/compute_norm_stats.py --config-name <config_name>。

问题：JAX 训练期间内存不足

解决方法：设置 XLA_PYTHON_CLIENT_MEM_FRACTION=0.9、降低批次大小，或配置 fsdp_devices：

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

问题：加载 PyTorch 检查点时内存不足

解决方法：export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

问题：找不到配置

解决方法：确保配置名称存在于 src/openpi/training/config.py 中（必须与 _CONFIGS 字典中的名称完全一致）。

问题：库变更后 PyTorch 训练出现发散

解决方法：重新应用 transformer 补丁。运行 uv cache clean transformers 重置，然后重新应用补丁。

问题：serve_policy.py 因 ModuleNotFoundError 崩溃

解决方法：先重新同步公共工作区：

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

如果缺失的模块与模拟器有关，请安装该示例所需的额外运行时依赖项：

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

问题：uv sync 因 rerun-sdk wheel 不匹配而失败

解决方法：

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

问题：检查点下载超时

解决方法：安装 gsutil 并手动预先下载：

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

如果之前的下载中断，请删除过期的 .lock 文件。

问题：策略服务以代码 137 退出

解决方法：这是 OOM 终止。设置 JAX 内存变量：

export XLA_PYTHON_CLIENT_PREALLOCATE=false
export XLA_PYTHON_CLIENT_ALLOCATOR=platform

---

面向 HPC/集群用户

在由 Slurm 管理的集群上，使用资源分配命令包装命令：

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

将缓存目录指向临时存储，避免占满 /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

使用由 uv 管理的环境时，避免叠加集群 Python 模块。通常只需 module load cuda。

---

高级主题

配置方案与基线：参见 references/config-recipes.md 训练调试指南：参见 references/training-debugging.md 检查点与环境映射：参见 references/checkpoints-and-env-map.md 远程客户端集成：参见 references/remote-client-pattern.md PyTorch 精度与补丁注意事项：参见 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