pi0(OM)-推理指导

本目录介绍在310P上pi0模型的离线模型转换及推理步骤。

概述

pi0模型在π0: A Vision-Language-Action Flow Model for General Robot Control 论文中提出，原仓库地址为：https://github.com/Physical-Intelligence/openpi

LeRobot 是一套基于 PyTorch 构建的面向真实世界机器人应用的工具集，核心目标是降低机器人领域的研发门槛，让开发者均可参与共享数据集、预训练模型的建设。该指导基于LeRobot框架，采用OM加速路线部署π0，https://huggingface.co/BrunoM42/pi0_aloha_transfer_cube。由于Lerobot在v0.4.0后重构了pi0和pi0.5代码结构。基于老Lerobot版本代码训练的模型在新版本代码推理上会出现诸多config类报错， 因此我们在本样例下，主要适配的基于‘577cd10974b84bea1f06b6472eb9e5e74e07f77a’ commit的Lerbot代码。

插件与驱动准备

• 该模型需要以下插件与驱动

  配套	版本	环境准备指导
  固件与驱动	24.1.RC3	Pytorch框架推理环境准备
  CANN	8.3.RC1	包含kernels包和toolkit包
  Python	3.11	-
  PyTorch	2.7.1	-
  说明：支持Atlas 300I DUO/Atlas 300I Pro，不支持Atlas 800I A2	\	\

获取本仓源码

git clone https://gitcode.com/ascend/ModelZoo-PyTorch.git
cd ModelZoo-PyTorch/ACL_PyTorch/built-in/embodied_ai/vla/pi0
# 记录本仓路径
export PI0=/path/to/ModelZoo-PyTorch/ACL_PyTorch/built-in/embodied_ai/vla/pi0

环境准备

• 获取LeRobot源码，切到指定commit并安装相关依赖：

git clone https://github.com/huggingface/lerobot.git
cd lerobot
git checkout 577cd10974b84bea1f06b6472eb9e5e74e07f77a
# 安装lerobot
pip install -e .
# 安装Aloha仿真依赖
pip install -e ".[aloha]"
# 叠加lerobot patch，只需执行一次
git apply ../patches/lerobot_diff.patch
cd ..

• 安装requirements:

pip3 install -r requirements.txt

查看transformer代码路径: pip show transformers | grep Location

cd /path/to/lib/python3.11/site-packages/transformers/models/gemma
# 叠加patch，只需执行一次
git apply -p1 $PI0/patches/modeling_gemma.patch

• 安装OpenGL/EGL相关依赖: apt install -y libegl1-mesa libegl1-mesa-dev libopengl0 libgl1-mesa-glx libgl1-mesa-dev libosmesa6-dev mesa-utils libglfw3

• 下载模型权重：

cd $PI0
mkdir weight
cd weight

• pi0_aloha_transfer_cube 下载链接
• paligemma-3b-pt-224 tokenizer，下载added_tokens.json, special_tokens_map.json, tokenizer_config.json, tokenizer.json和tokenizer.model这个五个文件即可：下载链接

cd ..

• 运行om时需要ACLLite组件，ACLLite安装指导：https://gitee.com/ascend/ACLLite/tree/master/python

文件目录结构

文件目录结构大致如下：

📁 pi0/
├── 📁 lerobot/
|── 📁 patches
|   |── 📄 lerobot_diff.patch
|   └── 📄 modeling_gemma.patch
|── 📁 weights
|   |── 📁 pi0_aloha_transfer_cube
|   └── 📁 paligemma-3b-pt-224
|── 📄 conver_verify_onnx_action_expert.py
|── 📄 conver_verify_onnx_vlm.py
|── 📄 infer.py
|── 📄 modeling_pi0_om.py
|── 📄 README.md
|── 📄 requirements.txt
|── 📄 utils.py
|── 📄 verify_om_onnx_action_expert.py
└── 📄 verify_om_onnx_vlm.py

导出 ONNX

cd $PI0
mkdir runtime_save # 用于保存中间运行时张量
# 导出vlm
python3 convert_verify_onnx_vlm.py --pretrained-policy-path /path/to/pi0_aloha_transfer_cube --tokenizer-path /path/to/paligemma-3b-pt-224
# 导出action expert
python3 convert_verify_onnx_action_expert.py --pretrained-policy-path /path/to/pi0_aloha_transfer_cube --tokenizer-path /path/to/paligemma-3b-pt-224

参数说明

• --pretrained-policy-path: pi0模型权重路径，默认为"./weight/pi0_aloha_transfer_cube"
• --tokenizer-path: paligemma-3b-pt-224 tokenizer配置文件路径，默认为"./weight/paligemma-3b-pt-224"
• --output: onnx模型输出路径，vlm部分默认为"outputs/onnx/pi0-vlm.onnx", action expert部分默认为"outputs/onnx/pi0-action_expert.onnx"
• --device: 默认为"cpu"
• --opset: onnx opset版本，默认为14
• --seed: 随机种子，默认为42
• --batch-size: batch数，默认为1
• --lang-tokens-len: language token长度，默认为48，与pi0_aloha_transfer_cube中的tokenizer_max_length对齐
• --runtime-save-dir: runtime tensor save路径，默认为"runtime_save"
• --no-validate: 设置该参数可跳过精度验证
• --log-level: 日志等级，默认为"INFO"
• --local-files-only: 设置改参数时从本地加载模型
• --past-kv-path: convert_verify_onnx_action_expert.py独有参数，past_kv_tensor读取路径，默认为"runtime_save/past_kv_tensor.pth"
• --prefix-pad-masks-path: convert_verify_onnx_action_expert.py独有参数，prefix_pad_masks读取路径，默认为"runtime_save/prefix_pad_masks.pth"

ATC将ONNX转为OM

将利用npu-smi info命令获取的芯片型号填入${soc_version}中

atc --model=/path/to/pi0-vlm.onnx --framework=5 --output=outputs/om/pi0-vlm --soc_version=${soc_version} --precision_mode_v2=origin
atc --model=/path/to/pi0-action_expert.onnx --framework=5 --output=outputs/om/pi0-action_expert --soc_version=${soc_version} --precision_mode_v2=origin

模型转换完成后会在指定的output目录下生成pi0-vlm.om和pi0-action_expert.om或者pi0-action_expert_{arch}.om，模型{arch}后缀为当前使用的CPU操作系统。

基于mock的输入数据，比对onnx和om输出的余弦相似度

python3 verify_om_onnx_vlm.py --onnx-model-path /path/to/pi0-vlm.onnx --om-model-path /path/to/pi0-vlm.om
python3 verify_om_onnx_action_expert.py --onnx-model-path /path/to/pi0-action_expert.onnx --om-model-path /path/to/pi0-action_expert.om

端到端Pi0

python3 infer.py --vlm-model-path /path/to/pi0-vlm.om --action-expert-model-path /path/to/pi0-action_expert.om

参数说明：

• --vlm-model-path: pi0 vlm om模型路径，默认为outputs/om/pi0-vlm.om
• --action-expert-model-path: pi0 action expert om模型路径，默认为outputs/om/pi0-action_expert.om
• --device: device id, 默认为0
• --mean-path: 反归一化mean值路径，默认为None，使用默认mean值，使用时必须和--std-path参数一起使用
• --std-path: 反归一化std值路径，默认为None，使用默认std值，使用时必须和--mean-path一起使用

性能数据

infer.py取libero数据集中的一个frame进行推理，性能如下

模型	芯片	数据集	单次推理耗时
π0	300I DUO(单芯)	随机数	220ms

仿真环境测评

export MUJOCO_GL=osmesa
python3 -m lerobot.scripts.eval\
  --policy.path=/path/to/pi0_aloha_transfer_cube\
  --policy.tokenizer_path=/path/to/paligemma-3b-pt-224\
  --env.type=aloha --env.task=AlohaTransferCube-v0\
  --eval.n_episodes=1\
  --eval.batch_size=1\
  --seed=42\
  --vlm-model-path /path/to/pi0-vlm.om\
  --action-expert-model-path /path/to/pi0-action_expert.om

执行后模型在仿真环境里完成任务，并在outputs/eval路径下生成仿真视频。