fine-tuning-serving-openpi

用途

本技能用于在公开的 openpi 仓库上，针对机器人操作任务对 Physical Intelligence 的 OpenPI 模型（pi0、pi0-fast、pi0.5）进行微调与推理部署，支持 JAX 与 PyTorch 两种后端，并覆盖从空白机器环境搭建到策略服务、客户端集成的完整流程。

主要能力

支持三种模型变体：pi0（流匹配 VLA，基线速度、最高质量）、pi0-fast（自回归动作 token，速度提升 2-5 倍）、pi0.5（pi0 加改进视觉编码器，最新默认）。
双后端：JAX 为官方主训练后端，PyTorch 面向部署且支持 DDP 多卡训练。
配置驱动：训练与推理参数集中在 src/openpi/training/config.py。
归一化统计：每个配置在训练前需预计算 norm stats。
WebSocket 服务：策略服务器暴露 WebSocket API 供低延迟推理。
检查点转换：可将 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 python -c "from openpi.training import config as _config; print(_config.get_config('pi05_droid').name)"
uv run scripts/serve_policy.py --help

JAX 微调流程

从 src/openpi/training/config.py 复制最接近的配置（如 pi05_libero、pi0_aloha_pen_uncap、pi05_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 版本，应用 transformers_replace 补丁。
转换 JAX 检查点：

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

单卡或多卡训练：

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

推理服务与客户端

默认环境预设（ALOHA、ALOHA_SIM、DROID、LIBERO）可直接启动：

uv run scripts/serve_policy.py --env ALOHA
uv run scripts/serve_policy.py --env DROID

自定义检查点：

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

客户端集成：

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

from openpi_client import websocket_client_policy
client = websocket_client_policy.WebsocketClientPolicy(host="gpu-server.local", port=8000)
observation = {"image": img, "state": state, "prompt": "pick up the red block"}
result = client.infer(observation)
actions = result["actions"]
robot.step(actions[0])

适用场景

在 LeRobot 或 RLDS 数据集上微调 pi0、pi0-fast、pi0.5。
为 ALOHA、DROID、LIBERO 等环境部署 OpenPI 策略。
在 JAX 与 PyTorch 之间转换检查点。
排查 OpenPI 训练问题（归一化统计、显存、配置错误等）。

限制与注意事项

训练前必须计算 norm stats，配置、数据集或变换变化后需重新计算。
显存需求：推理 pi0.5 约 24 GB；JAX 微调 pi0.5 约 60 GB（建议 A100 80GB）；PyTorch DDP 约 40 GB/GPU。
PyTorch 训练依赖对 transformers 包应用自定义补丁，库变更后需重新应用。
检查点下载超时需手动使用 gsutil 预取；服务进程退出码 137 表示 OOM，可调整 JAX 显存变量。
在 Slurm 集群上需使用 srun 分配资源，并将缓存目录指向 scratch。

资源

OpenPI 仓库：https://github.com/Physical-Intelligence/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