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 检查点时发生 OOM

解决方法：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/集群用户

详情、代码示例和参数表：请参阅 references/for-hpc-cluster-users.md。适用时请阅读。

高级主题

配置示例和基线：请参阅 references/config-recipes.md 训练调试指南：请参阅 references/training-debugging.md 检查点和环境映射：请参阅 references/checkpoints-and-env-map.md 远程客户端集成：请参阅 references/remote-client-pattern.md PyTorch 精度和补丁注意事项：请参阅 references/pytorch-gotchas.md

资源

• OpenPI 仓库：https://github.com/Physical-Intelligence/openpi
• OpenPI 客户端包：https://github.com/Physical-Intelligence/openpi/tree/main/packages/openpi-client
• pi0 论文：https://www.physicalintelligence.company/blog/pi0
• LeRobot 数据集格式：https://huggingface.co/docs/lerobot

Agent 操作流程

1. 检查环境。 确认 GPU 内存、模型检查点和许可证、框架版本及输入格式。
2. 明确输入。 根据数据或用户提供的信息确认格式、标识符和参数。任何会改变结果的值都应询问，不要猜测。
3. 先运行小规模版本。 先以低分辨率或较短长度对单个样本运行推理。
4. 使用上述说明和参考资料执行完整任务。
5. 验证结果。 通过定性检查和任务指标检查输出；检查生成内容的许可证限制。
6. 报告。 说明运行了什么（版本、命令、参数）、检查了什么，以及还有哪些不确定之处。

如果发生以下情况	请这样做
大型输入导致内存不足	降低分辨率或序列长度，使用半精度，或对输入进行分块处理。
已安装版本中缺少这些说明里的函数、标志或端点	查阅已安装版本自己的文档（help()、--help、官方文档），进行调整并告知用户。切勿编造 API。
必需的输入、标识符或参数不明确	询问用户，或在运行前明确说明所采用的假设。

诚信规则

• 切勿捏造结果、参数、标识符、引用或统计数据。如果无法运行或验证某项内容，请明确说明。
• 商用或发布前，请检查模型和数据集的许可证。
• 此处与版本相关的详细信息可能已经过时：请对照已安装版本的官方文档进行确认。
• 对于需要付费、占用共享 GPU 或云配额、涉及个人或患者数据，或无法撤销的操作，请先征求同意。

相关技能

• evaluating-cosmos-policy：在 LIBERO 和 RoboCasa 仿真环境中评估 NVIDIA Cosmos Policy。
• fine-tuning-openvla-oft：微调并评估 OpenVLA-OFT 和 OpenVLA-OFT+ 策略，用于通过连续动作头和 LoRA 适配生成机器人动作，……
• pufferlib：针对 PufferLib 强化学习环境、向量化、策略、PuffeRL 训练与评估以及安全的 c… 提供版本感知型指导。