Pi0 具身智能 VLA 大模型在昇腾 310P 上的离线模型转换与推理实战指南

Pi0 具身智能 VLA 大模型在昇腾 310P 上的离线模型转换与推理实战指南

【免费下载链接】cann-recipes-embodied-ai本项目针对具身智能业务中的典型模型、加速算法，提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-embodied-ai

导读：本文以 CANN 社区仓库 manipulation/pi0/infer_with_om 为蓝本，系统讲解如何将基于 Diffusion 的 VLA（Vision-Language-Action）大模型 Pi0 在昇腾 310P 上进行「PyTorch → ONNX → OM」离线转换、精度校验与端到端推理，并接入 MuJoCo Aloha 仿真环境完成评测。读者学完后，可完整复现从环境搭建、模型导出、ATC 离线编译到仿真评测的全流程，并掌握双模块（PaliGemma + Gemma 动作专家）拆分为两个 OM 图的设计思路。

背景与模型介绍

Pi0 是论文π0: A Vision-Language-Action Flow Model for General Robot Control（arXiv:2410.24164）中提出的机器人策略模型，它把机器人控制策略建模为条件去噪扩散过程（conditional denoising diffusion）：给定当前观测（摄像头图像 + 机器人状态）与语言指令，模型从一个噪声动作序列出发，通过多轮迭代去噪逐步逼近真实动作分布，输出未来一段时域内的动作轨迹。这种基于 Diffusion 的模仿学习控制策略在叠衣服、整理桌面等具有挑战性的精细操作任务上表现良好，是当前主流 VLA 模型之一。

在代码层面，Pi0 最早由 Physical Intelligence 以 JAX 实现（openpi 仓库），HuggingFace LeRobot 将其移植到 PyTorch。关键兼容性前提：LeRobot 在 v0.4.0 之后重构了 pi0 与 pi0.5 的代码结构，基于老版本 LeRobot 训练的模型在新版本代码上推理会出现大量 config 类报错。因此本样例固定适配577cd10974b84bea1f06b6472eb9e5e74e07f77a这个 commit 的 LeRobot 代码，配套使用的修改文件也在仓库的 lerobot_modify 目录中，使用其他 commit 可能导致类名、配置字段不匹配。

模型输入输出定义

本样例使用的示例模型为 HuggingFace 上的BrunoM42/pi0_aloha_transfer_cube（一个在 MuJoCo 仿真环境中微调、用于双臂递方块任务的模型）。其输入输出规格如下：

参数符号说明：

• $N_c$：摄像头个数（图片个数），训练时决定，通常为一到三路摄像头；
• $N_k$：机器人自由度，由构型决定（示例 Aloha 双臂为 14）；
• $H, W$：摄像头 RGB 图像分辨率（示例为 480×640）；
• $chunk_size$：机器人执行步数个数（示例为 50）。

输入输出格式定义在模型目录的config.json中。示例模型为一路摄像头，其关键内容如下（input_features声明输入张量类型与形状，output_features声明动作维度，chunk_size声明预测的步数）：

"input_features": { "observation.images.top": { "type": "VISUAL", "shape": [3, 480, 640] }, "observation.state": { "type": "STATE", "shape": [14] } }, "output_features": { "action": { "type": "ACTION", "shape": [14] } }, "chunk_size": 50,

环境搭建：昇腾 310P 运行配置

与昇腾平台相关的环境配置

OM 模型转化及运行需要安装 CANN 软件包。本样例的编译执行依赖 CANN 开发套件包（cann-toolkit）与 CANN 二进制算子包（cann-kernels），支持版本为CANN 8.0.0 ~ 8.2.RC1。请从昇腾社区软件包下载地址下载对应架构软件包，例如Ascend-cann-toolkit_8.2.RC1_linux-x86_64.run与Ascend-cann-kernels-310p_8.2.RC1_linux-x86_64.run，并参考 CANN 安装文档依次安装。

安装完成后，每个新终端都需要先 source 环境变量（${cann_install_path} 为 CANN 包的实际安装目录）：

# 方式1：默认路径安装，以 root 用户为例 source /usr/local/Ascend/ascend-toolkit/set_env.sh # 方式2：指定路径进行安装 source ${cann_install_path}/ascend-toolkit/set_env.sh

与昇腾服务器无关的环境配置

# 1) 拉取本 Pi0 使用仓库 git clone https://gitcode.com/CANN/cann-recipes-embodied-ai.git # 2) 拉取 lerobot 仓库并切到指定 commit cd cann-recipes-embodied-ai/manipulation git clone https://github.com/huggingface/lerobot.git # 如果 contrib 下没有 lerobot 目录 cd lerobot git checkout 577cd10974b84bea1f06b6472eb9e5e74e07f77a # 3) 创建运行环境 conda create -y -n lerobot python=3.10 # 如果之前没有创建过该环境 conda activate lerobot # 4) 安装 lerobot（本地可编辑安装） pip install -e . # 如果使用不同版本的 lerobot，请重新执行 pip install -e . 进行更新 # 5) 仿真依赖（Aloha 仿真需要 gym_aloha / gym-aloha，推荐用 extra 一次性安装） pip install -e ".[aloha]" # 6) 固定关键依赖版本 pip install numpy==1.26.0 # 使用 atc 工具时推荐，避免潜在兼容性问题 pip install transformers==4.55.4 # 按模型训练时所用 transformers 版本安装，本样例模型使用该版本 pip install decorator==5.2.1 # 7) 库文件替换（Gemma 精度相关补丁） export LEROBOT=/path/to/cann-recipes-embodied-ai/manipulation/lerobot export INFER_WITH_OM=/path/to/cann-recipes-embodied-ai/manipulation/pi0/infer_with_om cd /path/to/conda/envs/lerobot/lib/python3.10/site-packages/transformers/models/gemma # 可通过 pip show transformers 查看 transformers 安装路径 git apply --check -p1 $INFER_WITH_OM/lib/modeling_gemma.patch # 无报错表示可以应用补丁 git apply -p1 $INFER_WITH_OM/lib/modeling_gemma.patch # 8) 模型分部文件添加（将仓库中适配后的代码复制进 lerobot 源码） cd "$INFER_WITH_OM" cp $INFER_WITH_OM/lerobot_modify/modeling_pi0_vlm.py $LEROBOT/src/lerobot/policies/pi0/ cp $INFER_WITH_OM/lerobot_modify/modeling_pi0_action_expert.py $LEROBOT/src/lerobot/policies/pi0/ cp $INFER_WITH_OM/lerobot_modify/paligemma_with_expert_fp16.py $LEROBOT/src/lerobot/policies/pi0/ cp $INFER_WITH_OM/lerobot_modify/normalize.py $LEROBOT/src/lerobot/policies/

为什么需要这些补丁：从源码看，lib/modeling_gemma.patch 将 Gemma 的 RMSNorm 归一化改为在 float32 下先计算均值平方（x_fp32.pow(2).sum(...)后除以维度），再对原输入做rsqrt缩放，避免半精度下累积误差导致精度下降；而 lerobot_modify/paligemma_with_expert_fp16.py 则在模型内部关键路径上将中间张量统一为 float16（见其中hidden_states.to(dtype=torch.float16)等代码），保证与 310P 上的 OM 图计算精度一致。

ONNX 转换辅助依赖

转化 ONNX 的机器需要额外安装：

pip install onnx pytest onnxscript # 基于 Host CPU 转换 onnx（310P 宿主机执行）： pip install onnxruntime # 基于 Host GPU 转换 onnx： pip install onnxruntime-gpu

双模块拆分与整体推理链路设计

Pi0 由两个解耦模块组成：VLM 部分的PaliGemma（负责视觉-语言理解）和动作专家的Gemma（负责条件扩散去噪生成动作）。二者在推理流上相互独立，且动作专家是扩散模型，需要迭代执行 10 轮去噪。如果把它整体编译成一张离线图，会占据大量内存和编译时间，因此样例将其拆分为两个 ONNX/OM 文件：

• part1（PaliGemma）：输入摄像头图像、机器人状态、语言 token 与 mask，输出 KV 缓存张量past_kv_tensor与prefix_pad_masks中间张量；
• part2（Gemma 动作专家）：依赖 part1 保存的中间张量，输入状态、语言信息、KV 缓存、扩散时间步time与噪声noise，输出动作序列，需被循环调用 10 次。

推理时两步加载执行：先调用 part1 的 VLM 模块并保存中间张量，再多次调用 part2 的动作专家模块，最后对动作输出做反归一化映射回原始动作空间。推荐的单机链路为：

1. 在 310P 宿主机用 Host CPU 导出 ONNX；
2. 使用 ATC 将 ONNX 转为 OM（310P 上）；
3. 使用 run_om_e2e.py 完成端到端 OM 推理校验，或使用 eval_pi0_ascend.py 接入 Aloha 仿真环境评测。

第一步：导出 ONNX

在 Host（CPU 或 GPU）上执行（以本地目录pi0_model/为例，其中包含 config.json 等文件；也可以先用huggingface-cli download BrunoM42/pi0_aloha_transfer_cube --local-dir pi0_model下载）：

cd $INFER_WITH_OM mkdir runtime_save # 用于保存中间运行时张量 export PYTHONPATH=/path/to/manipulation/lerobot/src:$PYTHONPATH # 指向 lerobot 路径保证成功 import ./run_pi0_export.sh --pretrained-policy-path ./pi0_model # 替换为你的模型目录路径

脚本行为说明（依据 run_pi0_export.sh 的源码）：

• 脚本按顺序调用 convert_verify_onnx_vlm.py 与 convert_verify_onnx_action_expert.py，分别导出outputs/onnx/pi0-vlm.onnx与outputs/onnx/pi0-action_expert.onnx；
• part1 导出时会通过validate_onnx把 PyTorch 计算的past_kv_tensor与prefix_pad_masks保存到runtime_save/目录（past_kv_tensor.pth、prefix_pad_masks.pth），part2 导出时再从该目录加载这些运行时张量作为输入；
• 默认会用 ONNXRuntime CPU 对比 PyTorch 输出，打印max/mean abs diff，如需跳过可加--no-validate（对应脚本--skip-verify）；
• 输入 schema（包含多个摄像头 key 时也支持）严格来自config.json.input_features。

该脚本还暴露了若干可选参数：--vlm-output/--action_expert-output自定义 ONNX 输出路径，--runtime-save-dir指定中间张量目录，--device指定导出设备（默认 cpu），--extra-vlm "--opset 14"与--extra-action_expert "--lang-len 48"可向两个子脚本透传额外参数（动作专家导出默认 opset 为 14、语言长度 48）。

第二步：ATC 将 ONNX 转为 OM

在 310P 上（已安装并 source CANN 环境）执行 ATC：

# 推荐直接用 atc（路径按你的实际文件位置修改） atc --model=outputs/onnx/pi0-vlm.onnx \ --framework=5 \ --output=outputs/om/pi0_vlm \ --soc_version=Ascend310P1 \ --precision_mode_v2=origin atc --model=outputs/onnx/pi0-action_expert.onnx \ --framework=5 \ --output=outputs/om/pi0_action_expert \ --soc_version=Ascend310P1 \ --precision_mode_v2=origin

注意事项：

• --framework=5表示 ONNX 输入；--precision_mode_v2=origin保持原始精度；
• soc_version需要根据npu-smi info得到的 Name Device 中芯片型号填写。例如显示为 "310P1"，则soc_version填写Ascend310P1；
• 若报错[ERROR] Execute model failed for acl.mdl.execute error 507011，请检查npu-smi info显示的信息是否与 ATC 转化时使用的--soc_version相同；
• 模型转换完成后，当前目录（或--output指定目录）下应存在pi0_vlm.om和pi0_action_expert.om（也可能是pi0_action_expert_linux_x86_64.om等包含适配系统信息的格式），终端输出 "ATC run success, welcome to the next use" 即代表成功，后续使用请基于生成的 om 对应文件名。

基于 mock 数据的 CPU/GPU 与原始 PyTorch 输出相似度对比

构造确定性输入，对比 Pytorch CPU/GPU 和 OM 310P NPU 的输出余弦相似度（需要在 310P 上执行，具备 ACL/AclLite Python 依赖；安装后可参考以下方式设置 PYTHONPATH，具体路径按实际安装路径修改，可通过sudo find / -name "acllite_utils.py"查找，设置为该文件所在目录的上一级目录）：

# 如 export PYTHONPATH=/path/to/Ascend/ascend-toolkit/8.2.RC1/thirdpart/python/:$PYTHONPATH python3 verify_om_onnx_vlm.py \ --onnx-model-path outputs/onnx/pi0-vlm.onnx \ --om-model-path outputs/om/pi0_vlm.om python3 verify_om_onnx_action_expert.py \ --onnx-model-path outputs/onnx/pi0-action_expert.onnx \ --om-model-path outputs/om/pi0_action_expert.om

脚本会根据config.json生成确定性的 dummy 输入（支持多摄像头输入），并对比 ONNXRuntime(CPU) vs OM(NPU)。从源码看，verify_om_onnx_vlm.py 的 VLM 输入为observation.state（float32, B×14）、observation.images.top（float32, B×3×H×W）、lang_tokens（int64, B×L）、lang_masks（bool, B×L），输出past_kv_tensor与prefix_pad_masks；verify_om_onnx_action_expert.py 的动作专家输入为observation.state（float16）、lang_tokens、lang_masks、past_kv_tensor（float16,[LAYER, 2, B, S, H, D]）、prefix_pad_masks、time（float16, B）、noise（float16, B×50×32）。两者都会报告 abs/rel 误差以及最后一维上的余弦相似度。

端到端 Pi0 模型推理

Pi0 模型推理需要分两步调用：先调用 part1 的 VLM 模块并保存中间张量，再多次调用 part2（10 次）的动作专家模块，最后将动作输出反归一化映射到原动作空间。

cd ../pi0/infer_with_om python ./run_om_e2e.py \ --vlm-model-path ./outputs/om/pi0_vlm.om \ --action-expert-model-path ./outputs/om/pi0_action_expert.om # 反归一化参数默认值为示例模型的参数；若使用自定义模型，需先运行以下命令获取 mean.pt 和 std.pt： # --policy.path 为你的 safetensors 模型路径 # 注意：仅运行示例模型时无需获取 mean.pt/std.pt，直接使用默认值即可 # 按你的实际渲染环境配置 MUJOCO_GL，若无 GPU 渲染需求，推荐使用 osmesa 渲染（需先安装 osmesa 库）： # sudo apt-get install libosmesa6 libosmesa6-dev # 若未配置代理则使用 huggingface 镜像网站 cd ../../lerobot export HF_ENDPOINT=https://hf-mirror.com export MUJOCO_GL=osmesa python ./src/lerobot/scripts/eval.py --policy.path=/path/to/your/safetensors/file --env.type=aloha --env.task=AlohaTransferCube-v0 --env.episode_length=10 cd ../pi0/infer_with_om # 获取到 mean.pt 和 std.pt 后，运行端到端 OM 推理： python ./run_om_e2e.py \ --mean-path ../../lerobot/mean.pt \ --std-path ../../lerobot/std.pt \ --vlm-model-path ./outputs/om/pi0_vlm.om \ --action-expert-model-path ./outputs/om/pi0_action_expert.om

推理细节（依据 run_om_e2e.py 源码）：

• Pi0类通过AclLiteModel分别加载 VLM 与动作专家两个 OM 模型，interface()方法先组装 part1 输入（state、image 均为 float32，lang_tokens 为 int64，lang_masks 为 bool），执行后从输出中取出kv_tensor与prefix_pad_masks；
• part2 使用 float16 的 state、int64 的 lang_tokens、bool 的 lang_masks、float16 的 kv_tensor 与 prefix_pad_masks，外加初始化为 1.0 的扩散时间步dummy_time和全零噪声dummy_noise（shape 为(1, 50, 32)）；
• 在 10 次循环中每次把上一次的output[0]作为新的noise输入，同时dummy_time -= 0.1实现时间步从 1.0 到 0.0 的退火去噪调度；
• 最终动作通过_to_action_14()从(1,50,32)截取到 14 维，再经Unnormalize（结合 config 中的normalization_mapping，VISUAL 为 IDENTITY、STATE 与 ACTION 为 MEAN_STD）反归一化得到最终动作；
• 反归一化统计量优先级：用户传入--mean-path/--std-path（二者必须同时提供）> 示例模型内置默认值；select_action()会把整段 50 步动作存入队列，按执行步逐条弹出，与 LeRobot 的策略执行语义一致。

仿真环境 Pi0 评测

eval_pi0_ascend.py 基于 LeRobot eval 流程接入 VLM OM 与 action expert OM，用于 Aloha 仿真环境评测。运行前需确认已安装 MuJoCo/Aloha、ACL/AclLite 依赖，并已通过前文步骤生成 OM 文件：

# 此处 policy.path 替换为你的 safetensors 模型路径；脚本会使用该目录中的 config 信息构造环境和 policy export MUJOCO_GL=osmesa python ./eval_pi0_ascend.py \ --policy.path=/path/to/models/pi0_aloha_model \ --env.type=aloha --env.task=AlohaTransferCube-v0 \ --env.episode_length=600 \ --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

从源码看，该脚本复用了 LeRobot 的EvalPipelineConfig配置体系（lerobot.configs.eval），通过make_env/make_policy构造 Aloha 环境与策略，并在内部维护一个动作队列（deque），逐帧将 OM 模型输出的动作反归一化后交给仿真环境执行、收集评测指标。如果使用自定义模型并需要自定义动作归一化统计量，可额外传入--mean-path与--std-path，这两个参数需要同时提供。

此外，Pi0.5 目录中的 lerobot_eval_om.py 可作为其他 OM backend 接入 LeRobot eval 流程的参考实现，但不能直接作为 Pi0 脚本使用。

可能遇到的问题

1. 运行lerobot/scripts/eval.py时，若使用网络环境下载google/paligemma-3b-pt-224模型，需提前到模型对应的 HuggingFace 页面请求访问权限（参考 HuggingFace 快速入门与 gated models 文档）。
2. 若网络环境下载 HuggingFace 模型较慢，遇到google/paligemma-3b-pt-224下载卡顿，可手动下载模型到本地路径，再修改以下文件中对应行处的google/paligemma-3b-pt-224为本地路径：
   • lerobot/src/lerobot/policies/pi0/modeling_pi0.py中对应第 247 行；
   • lerobot/src/lerobot/policies/pi0/modeling_pi0_vlm.py中对应第 306 行、第 350 行。

附录：lerobot 根目录相关代码目录树

经过上述操作，pi0 适配昇腾的 lerobot 根目录最终相关代码目录树如下：

|-- manipulation/pi0/infer_with_om/ # 本目录 ├── README.md # 使用指南（本文件） ├── convert_verify_onnx_action_expert.py # PyTorch -> ONNX（动作专家）转换与验证 ├── convert_verify_onnx_vlm.py # PyTorch -> ONNX（VLM 部分）转换与验证 ├── eval_pi0_ascend.py # Pi0 昇腾仿真评测脚本 ├── run_om_e2e.py # Pi0 昇腾端到端推理脚本 ├── verify_om_onnx_action_expert.py # ONNXRuntime(CPU) vs OM(NPU) 误差对比（动作专家） ├── verify_om_onnx_vlm.py # ONNXRuntime(CPU) vs OM(NPU) 误差对比（VLM） ├── run_pi0_export.sh # 示例导出脚本 ├── lerobot # lerobot 代码文件夹 └── lib/ └── └── modeling_gemma.patch # Gemma 相关补丁 └── outputs/ # 输出结果 └── onnx/ ├── pi0-vlm.onnx └── pi0-action_expert.onnx └── om/ ├── pi0_vlm.om └── pi0_action_expert.om

参考资料

本文所述 Pi0 模型及 LeRobot 框架相关引用：

@misc{black2024pi0visionlanguageactionflowmodel, title={$\pi_0$: A Vision-Language-Action Flow Model for General Robot Control}, author={Kevin Black and Noah Brown and Danny Driess and Adnan Esmail and Michael Equi and Chelsea Finn and Niccolo Fusai and Lachy Groom and Karol Hausman and Brian Ichter and Szymon Jakubczak and Tim Jones and Liyiming Ke and Sergey Levine and Adrian Li-Bell and Mohith Mothukuri and Suraj Nair and Karl Pertsch and Lucy Xiaoyang Shi and James Tanner and Quan Vuong and Anna Walling and Haohuan Wang and Ury Zhilinsky}, year={2024}, eprint={2410.24164}, archivePrefix={arXiv}, primaryClass={cs.LG}, } @misc{cadene2024lerobot, author = {Cadene, Remi and Alibert, Simon and Soare, Alexander and Gallouedec, Quentin and Zouitine, Adil and Palma, Steven and Kooijmans, Pepijn and Aractingi, Michel and Shukor, Mustafa and Aubakirova, Dana and Russi, Martino and Capuano, Francesco and Choghari, Jade and Moss, Jess and Wolf, Thomas}, title = {LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch}, howpublished = "\url{https://github.com/huggingface/lerobot}", year = {2024} }

【免费下载链接】cann-recipes-embodied-ai本项目针对具身智能业务中的典型模型、加速算法，提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-embodied-ai