Pi0 具身智能 VLA 大模型在昇腾 310P 上的离线模型转换与推理实战指南【免费下载链接】cann-recipes-embodied-ai本项目针对具身智能业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-embodied-ai导读本文以 CANN 社区仓库 manipulation/pi0/infer_with_om 为蓝本系统讲解如何将基于 Diffusion 的 VLAVision-Language-Action大模型 Pi0 在昇腾 310P 上进行「PyTorch → ONNX → OM」离线转换、精度校验与端到端推理并接入 MuJoCo Aloha 仿真环境完成评测。读者学完后可完整复现从环境搭建、模型导出、ATC 离线编译到仿真评测的全流程并掌握双模块PaliGemma Gemma 动作专家拆分为两个 OM 图的设计思路。背景与模型介绍Pi0 是论文π0: A Vision-Language-Action Flow Model for General Robot ControlarXiv: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 仿真环境中微调、用于双臂递方块任务的模型。其输入输出规格如下输入数据名数据类型dtype数据大小shapeobservation.images_xxx摄像头图像Float32[1, N_c, 3, H, W]observation.state机器人位姿Float32[1, N_k]输出数据名数据大小shapeactions[1, chunk_size, N_k]参数符号说明$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 python3.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 numpy1.26.0 # 使用 atc 工具时推荐避免潜在兼容性问题 pip install transformers4.55.4 # 按模型训练时所用 transformers 版本安装本样例模型使用该版本 pip install decorator5.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(dtypetorch.float16)等代码保证与 310P 上的 OM 图计算精度一致。ONNX 转换辅助依赖转化 ONNX 的机器需要额外安装pip install onnx pytest onnxscript # 基于 Host CPU 转换 onnx310P 宿主机执行 pip install onnxruntime # 基于 Host GPU 转换 onnx pip install onnxruntime-gpu双模块拆分与整体推理链路设计Pi0 由两个解耦模块组成VLM 部分的PaliGemma负责视觉-语言理解和动作专家的Gemma负责条件扩散去噪生成动作。二者在推理流上相互独立且动作专家是扩散模型需要迭代执行 10 轮去噪。如果把它整体编译成一张离线图会占据大量内存和编译时间因此样例将其拆分为两个 ONNX/OM 文件part1PaliGemma输入摄像头图像、机器人状态、语言 token 与 mask输出 KV 缓存张量past_kv_tensor与prefix_pad_masks中间张量part2Gemma 动作专家依赖 part1 保存的中间张量输入状态、语言信息、KV 缓存、扩散时间步time与噪声noise输出动作序列需被循环调用 10 次。推理时两步加载执行先调用 part1 的 VLM 模块并保存中间张量再多次调用 part2 的动作专家模块最后对动作输出做反归一化映射回原始动作空间。推荐的单机链路为在 310P 宿主机用 Host CPU 导出 ONNX使用 ATC 将 ONNX 转为 OM310P 上使用 run_om_e2e.py 完成端到端 OM 推理校验或使用 eval_pi0_ascend.py 接入 Aloha 仿真环境评测。第一步导出 ONNX在 HostCPU 或 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.onnxpart1 导出时会通过validate_onnx把 PyTorch 计算的past_kv_tensor与prefix_pad_masks保存到runtime_save/目录past_kv_tensor.pth、prefix_pad_masks.pthpart2 导出时再从该目录加载这些运行时张量作为输入默认会用 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 --modeloutputs/onnx/pi0-vlm.onnx \ --framework5 \ --outputoutputs/om/pi0_vlm \ --soc_versionAscend310P1 \ --precision_mode_v2origin atc --modeloutputs/onnx/pi0-action_expert.onnx \ --framework5 \ --outputoutputs/om/pi0_action_expert \ --soc_versionAscend310P1 \ --precision_mode_v2origin注意事项--framework5表示 ONNX 输入--precision_mode_v2origin保持原始精度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.statefloat32, B×14、observation.images.topfloat32, B×3×H×W、lang_tokensint64, B×L、lang_masksbool, B×L输出past_kv_tensor与prefix_pad_masksverify_om_onnx_action_expert.py 的动作专家输入为observation.statefloat16、lang_tokens、lang_masks、past_kv_tensorfloat16,[LAYER, 2, B, S, H, D]、prefix_pad_masks、timefloat16, B、noisefloat16, B×50×32。两者都会报告 abs/rel 误差以及最后一维上的余弦相似度。端到端 Pi0 模型推理Pi0 模型推理需要分两步调用先调用 part1 的 VLM 模块并保存中间张量再多次调用 part210 次的动作专家模块最后将动作输出反归一化映射到原动作空间。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_ENDPOINThttps://hf-mirror.com export MUJOCO_GLosmesa python ./src/lerobot/scripts/eval.py --policy.path/path/to/your/safetensors/file --env.typealoha --env.taskAlohaTransferCube-v0 --env.episode_length10 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 均为 float32lang_tokens 为 int64lang_masks 为 bool执行后从输出中取出kv_tensor与prefix_pad_maskspart2 使用 float16 的 state、int64 的 lang_tokens、bool 的 lang_masks、float16 的 kv_tensor 与 prefix_pad_masks外加初始化为 1.0 的扩散时间步dummy_time和全零噪声dummy_noiseshape 为(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_mappingVISUAL 为 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_GLosmesa python ./eval_pi0_ascend.py \ --policy.path/path/to/models/pi0_aloha_model \ --env.typealoha --env.taskAlohaTransferCube-v0 \ --env.episode_length600 \ --eval.n_episodes1 \ --eval.batch_size1 \ --seed42 \ --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 脚本使用。可能遇到的问题运行lerobot/scripts/eval.py时若使用网络环境下载google/paligemma-3b-pt-224模型需提前到模型对应的 HuggingFace 页面请求访问权限参考 HuggingFace 快速入门与 gated models 文档。若网络环境下载 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 - ONNXVLM 部分转换与验证 ├── 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考