SuperAPI AI 工具集
合作ENAI 行业资讯收藏
Skill 广场fine-tuning-serving-openpi

fine-tuning-serving-openpi

智能体与模型

查看项目

更新于 10-15 14:01

技能介绍

用途

本技能面向 Physical Intelligence 的 OpenPI 模型（pi0、pi0-fast、pi0.5），提供从环境搭建、微调到策略推理服务的完整工作流，用于机器人操作任务。

主要能力

  • 支持三种模型变体：pi0（流匹配 VLA，质量最高）、pi0-fast（自回归动作 token，速度提升 2-5 倍）、pi0.5（pi0 加改进视觉编码器，最新默认）。
  • 双后端：JAX（官方主训练后端）与 PyTorch（社区、部署友好）。
  • 配置驱动：所有训练与服务参数定义在 src/openpi/training/config.py。
  • 归一化统计：每个配置在训练前都需要预计算归一化统计量。
  • WebSocket 服务：策略服务器暴露 WebSocket API 用于低延迟推理。

使用方式

空白机器环境搭建

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 数据）

  1. 从 src/openpi/training/config.py 复制最接近的配置（如 pi05_libero、pi0_libero、pi0_fast_libero、pi0_aloha_pen_uncap、pi05_droid_finetune、pi05_full_droid_finetune）。
  2. 更新数据集映射与基础检查点，确保 weight_loader 与模型类型匹配。
  3. 计算归一化统计：uv run scripts/compute_norm_stats.py --config-name <config_name>。
  4. 启动训练：XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 uv run scripts/train.py <config_name> --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>。

PyTorch 训练与检查点转换

  1. 同步依赖并确认 transformers 版本。
  2. 应用 OpenPI 的 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>。

策略推理服务

  1. 选择环境预设（ALOHA、ALOHA_SIM、DROID、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。
  3. 验证连通性：uv run examples/simple_client/main.py --env DROID。
  4. 在机器人代码中嵌入远程客户端：
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"]
robot.step(actions[0])

适用场景

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

限制与注意事项

  • 训练循环不变式：每次配置或数据集变更后，必须重新执行“计算归一化统计 → 训练 → 服务检查点 → 验证推理”循环。
  • 计算资源要求：服务 pi0.5 推理约需 1x A100/H100 与约 24 GB 显存；JAX 微调 pi0.5 约需 1x A100 80GB 与约 60 GB 显存；PyTorch DDP 微调约需 1-8x A100 与约 40 GB/GPU 显存；计算归一化统计约需 8 GB 显存。
  • 常见问题包括缺失归一化统计、JAX 训练内存不足、PyTorch 检查点加载 OOM、配置未找到、PyTorch 训练发散、serve_policy.py 模块缺失、uv sync 失败、检查点下载超时、策略服务器退出码 137 等，资料中提供了对应修复命令。
  • HPC/集群用户需使用 Slurm 资源分配并重定向缓存目录。
  • 高级主题参考文档：config-recipes.md、training-debugging.md、checkpoints-and-env-map.md、remote-client-pattern.md、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