fine-tuning-serving-openpi

用途

本技能面向 Physical Intelligence 的 OpenPI 模型（pi0、pi0-fast、pi0.5），提供在机器人操作任务上微调与部署的端到端工作流，基于公开的 openpi 仓库。覆盖空白机器环境搭建、JAX 训练、PyTorch 训练、检查点转换以及策略推理服务。

模型家族

OpenPI 实现三种模型变体：

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

关键设计：双后端（JAX 为主，官方训练；PyTorch 为社区版，便于部署）；配置驱动，所有训练/服务参数定义在 src/openpi/training/config.py；每个配置在训练前都需要预计算归一化统计量；策略服务器通过 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	快速，可在登录节点运行

使用方式

空白机器环境搭建

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 python -c "from openpi.training import config as _config; print(_config.get_config('pi05_droid').name)"
uv run scripts/serve_policy.py --help

JAX 微调（LeRobot 数据）

从 src/openpi/training/config.py 复制最接近的配置（如 pi05_libero、pi0_libero、pi0_fast_libero、pi0_aloha_pen_uncap、pi05_droid_finetune、pi05_full_droid_finetune）。
更新数据集映射与基础检查点，设置 repo_id，确保 weight_loader 与模型类型匹配。
计算归一化统计量：uv run scripts/compute_norm_stats.py --config-name <config_name>。
启动训练：XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 uv run scripts/train.py <config_name> --exp-name=<run_name> --overwrite。
服务并验证：uv run scripts/serve_policy.py policy:checkpoint --policy.config=<config_name> --policy.dir=checkpoints/<config_name>/<run_name>/<step>。

PyTorch 训练与检查点转换

同步依赖并确认 transformers 版本。
应用 OpenPI 对 transformers 的补丁：cp -r ./src/openpi/models_pytorch/transformers_replace/* .venv/lib/python3.11/site-packages/transformers/。
转换 JAX 检查点：uv run examples/convert_jax_model_to_pytorch.py --checkpoint_dir <jax_checkpoint_dir> --config_name <config_name> --output_path <pytorch_checkpoint_dir>。
训练或服务：单 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>。

策略推理服务

默认环境预设：ALOHA（pi05_aloha）、ALOHA_SIM（pi0_aloha_sim）、DROID（pi05_droid）、LIBERO（pi05_libero）。

启动服务：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。

客户端集成：

from openpi_client import websocket_client_policy
import numpy as np

client = websocket_client_policy.WebsocketClientPolicy(host="gpu-server.local", port=8000)
observation = {
    "image": np.random.rand(224, 224, 3),
    "state": np.zeros(7),
    "prompt": "pick up the red block",
}
result = client.infer(observation)
actions = result["actions"]  # shape: (action_chunk_size, action_dim)
robot.step(actions[0])

适用场景

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

限制与注意事项

需要相应 GPU 与显存，训练前必须计算归一化统计量。
PyTorch 训练依赖对 transformers 的补丁，库变更后需重新应用。
检查点下载可能超时，可安装 gsutil 手动预取。
商业或公开发布前需检查模型与数据集许可证。
版本相关细节可能过时，应以所安装版本的官方文档为准。
涉及费用、共享 GPU、云配额、个人或患者数据以及不可撤销操作前应先询问。

技能介绍

用途