个人主页链接项目主页链接1. 说明本文档用于补充 LeRobot ACT 模型从 PyTorch 导出为 ONNX并进一步用于 TensorRT FP16 部署时涉及的核心依赖、转换链路、验证方法和常见问题。当前流程目标LeRobot ACT PyTorch 模型 │ ▼ torch.onnx.export │ ▼ model.onnx │ ▼ ONNX 检查 / 数值验证 │ ▼ TensorRT FP16 Engine当前使用的预训练模型lerobot/act_aloha_sim_transfer_cube_human对应验证数据集lerobot/aloha_sim_transfer_cube_human本流程不需要训练环境只需要能够加载预训练模型加载数据集统计信息执行一次 PyTorch 前向推理导出 ONNX使用 ONNX / TensorRT 进行推理验证。2. ONNX 转换相关 Python 包在安装 ONNX 导出环境时常见会安装以下几个包ml_dtypes onnx onnx_ir onnxscript它们分别承担不同作用。2.1 ml_dtypesml_dtypes提供机器学习场景中使用的一些特殊数据类型。典型包括bfloat16 float8 其他扩展浮点类型NumPy 对部分机器学习 dtype 的原生支持并不完整因此 ONNX、JAX、TensorFlow 等库会依赖ml_dtypes。在当前ACT - ONNX - TensorRT FP16流程中通常不需要直接调用ml_dtypes。它主要属于底层 dtype 支持依赖。2.2 onnxonnx是 ONNX 官方核心 Python 包。主要作用定义 ONNX 模型格式加载.onnx保存.onnx检查模型合法性操作 ONNX Graph操作 Tensor、Node、Input、Output 等结构。常用示例importonnx modelonnx.load(model.onnx)onnx.checker.check_model(model)print(onnx.helper.printable_graph(model.graph))也可以重新保存onnx.save(model,model_checked.onnx)TensorRT 最终读取的model.onnx本质上就是 ONNXModelProto序列化后的文件。2.3 onnx_ironnx_ir是较新的 ONNX Python IRIntermediate Representation中间表示库。它主要用于图结构构建图遍历Node / Value 操作Graph RewriteONNX 转换新版 ONNX 导出器内部表示。相比直接修改onnx.ModelProtoonnx_ir更适合现代 ONNX exporter 在 Python 层处理计算图。一般项目代码不会直接使用它但新版torch.onnx和onnxscript可能会间接依赖它。2.4 onnxscriptonnxscript用于通过 Python 表达和生成 ONNX 算子、Function 和计算图。它在较新的 PyTorch ONNX 导出链中非常重要。典型作用PyTorch Operator │ ▼ onnxscript │ ▼ ONNX Operator / ONNX Function例如新版 PyTorch 导出方式torch.onnx.export(model,args,model.onnx,dynamoTrue,)内部可能大量依赖onnxscript onnx_ir onnx完成算子转换和图生成。3. 四个包之间的关系可以简单理解为ml_dtypes │ │ 特殊数据类型支持 ▼ PyTorch Model │ │ torch.onnx.export() ▼ onnxscript │ │ PyTorch 算子 - ONNX 表达 ▼ onnx_ir │ │ 中间图表示 / 图转换 ▼ onnx │ │ ModelProto 序列化 ▼ model.onnx │ ▼ TensorRT │ ▼ FP16 Engine在实际项目中重要程度可以理解为包是否经常直接使用主要作用ml_dtypes否特殊 ML dtypeonnx是ONNX 模型读取、检查、保存onnx_ir否新版 ONNX 图中间表示onnxscript通常间接PyTorch 算子转换为 ONNX4. PyTorch 到 ONNX 的典型转换基础形式torch.onnx.export(model,example_inputs,model.onnx,)推荐明确指定torch.onnx.export(model,example_inputs,model.onnx,opset_version17,input_names[observation.images.top,observation.state,],output_names[action,],)对于新版 PyTorch也可能使用torch.onnx.export(model,example_inputs,model.onnx,opset_version17,dynamoTrue,)是否使用dynamoTrue需要结合当前 LeRobot、PyTorch、onnxscript 和 TensorRT 版本决定。如果旧 exporter 已经能够稳定导出则不一定需要切换到 dynamo exporter。5. 当前 ACT 模型输入输出当前预训练 ACT 模型主要输入包括observation.images.top observation.state其中observation.images.top通常为[B, 3, H, W]例如[1, 3, 480, 640]机器人状态observation.state为 14 维[B, 14]模型最终输出 action一般包含 action chunk[B, chunk_size, action_dim]当前模型action_dim 14 chunk_size 100因此常见原始输出可理解为[1, 100, 14]具体 ONNX 输出是否保留完整 chunk应以实际导出脚本为准。6. 为什么需要 dataset 的 stats.json预训练模型并不只是输入图像 输入 state 输出 actionLeRobot 通常还涉及输入预处理 Normalize 模型推理 Unnormalize 输出后处理数据集meta/stats.json提供模型训练时使用的数据统计信息例如mean std min max可能用于observation.state normalization action normalization因此只有model.safetensors通常不足以完全复现原始推理流程。当前离线导出流程会结合model/config.json dataset/meta/stats.json恢复必要的预处理和后处理逻辑。7. 推荐的 ONNX 导出产物建议最终输出目录act_policy/ ├── model.onnx ├── deployment.json └── golden.ptmodel.onnxTensorRT 的输入模型。deployment.json记录部署相关信息例如输入名称 输出名称 输入 shape action_dim chunk_size 图像尺寸 归一化方式避免 ROS2 / TensorRT runtime 中硬编码过多模型参数。golden.pt用于保存 PyTorch 原始模型的一组基准输入输出。例如{inputs:...,outputs:...}后续用于比较PyTorch ONNX Runtime TensorRT FP32 TensorRT FP16之间的数值误差。8. ONNX 导出后必须进行检查8.1 ONNX 结构检查执行importonnx modelonnx.load(model.onnx)onnx.checker.check_model(model)print(ONNX model check passed.)如果onnx.checker.check_model通过只说明 ONNX 图结构合法。不代表输出数值一定正确 TensorRT 一定支持8.2 查看输入输出可以执行importonnx modelonnx.load(model.onnx)forxinmodel.graph.input:print(INPUT :,x.name)forxinmodel.graph.output:print(OUTPUT:,x.name)确认名称和 TensorRT runtime 中使用的名称一致。9. ONNX Runtime 数值验证推荐在进入 TensorRT 之前先做PyTorch vs ONNX Runtime比较。安装pipinstallonnxruntime如果是 CPUimportonnxruntimeasort sessionort.InferenceSession(model.onnx,providers[CPUExecutionProvider],)运行outputssession.run(None,{observation.images.top:image_numpy,observation.state:state_numpy,},)然后和 PyTorch 输出比较importnumpyasnp diffnp.abs(torch_output-onnx_output)print(max abs diff:,diff.max())print(mean abs diff:,diff.mean())如果 ONNX Runtime 与 PyTorch 差异已经很大则不能直接进入 TensorRT 调试。10. TensorRT 转换建议ONNX 验证通过后再进行 TensorRT 转换。典型model.onnx │ ▼ TensorRT Builder │ ├── FP32 └── FP16 ▼ model.engine建议先验证ONNX - TensorRT FP32再开启FP16这样更容易定位问题。如果 FP32 正常而 FP16 有明显误差则问题主要集中在FP16 精度 算子数值稳定性 TensorRT tactic而不是 ONNX 转换本身。11. 推荐验证顺序完整验证链建议Step 1 PyTorch 原模型推理 │ ▼ 保存 golden.pt Step 2 导出 model.onnx │ ▼ onnx.checker.check_model() Step 3 ONNX Runtime CPU │ ▼ 和 PyTorch golden 对比 Step 4 TensorRT FP32 │ ▼ 和 ONNX / PyTorch 对比 Step 5 TensorRT FP16 │ ▼ 和 FP32 / PyTorch 对比 Step 6 接入 ROS2不要直接从PyTorch跳到ROS2 TensorRT FP16否则出了问题很难判断是哪一层导致。12. 常见 ONNX 导出错误12.1 UnsupportedOperatorError例如UnsupportedOperatorError说明某个 PyTorch operator 没有对应的 ONNX symbolic / translation。处理方向尝试更高 opset检查新版 dynamo exporter替换不支持算子自定义 ONNX operator对模型做 wrapper。12.2 onnxscript conversion error如果出现onnxscript conversion error通常需要检查torch onnx onnxscript onnx_ir版本兼容性。不要只升级单个包。12.3 Dynamic shape 问题机器人部署一般输入尺寸固定例如batch 1 image 480x640 state 14如果没有动态 batch / 动态分辨率需求建议优先导出固定 shape。优点TensorRT optimization profile 更简单 engine 更稳定 构建速度更快 推理性能更容易达到最佳12.4 TensorRT 不支持某个 ONNX 算子即使onnx.checker通过也可能出现 TensorRT parser 错误。例如Unsupported node Unsupported operator No importer registered这时需要检查ONNX opset TensorRT version 算子实现必要时对 ONNX Graph 进行 simplify / rewrite。13. 当前项目环境建议ONNX 导出环境至少应包括python torch torchvision lerobot safetensors numpy onnx onnxscript onnx_ir ml_dtypes验证阶段建议增加onnxruntimeTensorRT 环境另外需要TensorRT CUDA cuDNN取决于具体部署注意模型导出环境和TensorRT 部署环境不一定必须是同一个 Python 环境。当前没有训练环境并不影响 ONNX 导出。只要可以正确加载model.safetensors config.json stats.json并完成一次模型 forward就可以进行导出。14. 当前资产目录当前下载完成后的结构lerobot_assets/ ├── dataset/ │ └── aloha_sim_transfer_cube_human/ │ ├── data/ │ │ └── chunk-000/ │ │ ├── file-000.parquet │ │ ├── file-001.parquet │ │ └── file-002.parquet │ ├── meta/ │ │ ├── episodes/ │ │ │ └── chunk-000/ │ │ │ └── file-000.parquet │ │ ├── info.json │ │ ├── stats.json │ │ └── tasks.parquet │ ├── README.md │ └── videos/ │ └── observation.images.top/ │ └── chunk-000/ │ └── file-000.mp4 │ └── model/ └── act_aloha_sim_transfer_cube_human/ ├── config.json ├── demo.gif ├── eval_info.json ├── model.safetensors ├── README.md ├── train_config.json └── training_curves.png已经具备进行离线 ONNX 导出的主要模型和验证数据。15. 下一步建议下一步执行python export_act_onnx_offline.py\--pretrained./lerobot_assets/model/act_aloha_sim_transfer_cube_human\--dataset-root ./lerobot_assets/dataset/aloha_sim_transfer_cube_human\--output-dir ./act_policy\--devicecpu目标生成act_policy/ ├── model.onnx ├── deployment.json └── golden.pt之后依次完成PyTorch ↓ ONNX Runtime ↓ TensorRT FP32 ↓ TensorRT FP16 ↓ ROS2 Runtime并使用同一组 golden 输入输出进行数值一致性验证。 LeRobot-ONNX 转换补充
个人主页链接项目主页链接1. 说明本文档用于补充 LeRobot ACT 模型从 PyTorch 导出为 ONNX并进一步用于 TensorRT FP16 部署时涉及的核心依赖、转换链路、验证方法和常见问题。当前流程目标LeRobot ACT PyTorch 模型 │ ▼ torch.onnx.export │ ▼ model.onnx │ ▼ ONNX 检查 / 数值验证 │ ▼ TensorRT FP16 Engine当前使用的预训练模型lerobot/act_aloha_sim_transfer_cube_human对应验证数据集lerobot/aloha_sim_transfer_cube_human本流程不需要训练环境只需要能够加载预训练模型加载数据集统计信息执行一次 PyTorch 前向推理导出 ONNX使用 ONNX / TensorRT 进行推理验证。2. ONNX 转换相关 Python 包在安装 ONNX 导出环境时常见会安装以下几个包ml_dtypes onnx onnx_ir onnxscript它们分别承担不同作用。2.1 ml_dtypesml_dtypes提供机器学习场景中使用的一些特殊数据类型。典型包括bfloat16 float8 其他扩展浮点类型NumPy 对部分机器学习 dtype 的原生支持并不完整因此 ONNX、JAX、TensorFlow 等库会依赖ml_dtypes。在当前ACT - ONNX - TensorRT FP16流程中通常不需要直接调用ml_dtypes。它主要属于底层 dtype 支持依赖。2.2 onnxonnx是 ONNX 官方核心 Python 包。主要作用定义 ONNX 模型格式加载.onnx保存.onnx检查模型合法性操作 ONNX Graph操作 Tensor、Node、Input、Output 等结构。常用示例importonnx modelonnx.load(model.onnx)onnx.checker.check_model(model)print(onnx.helper.printable_graph(model.graph))也可以重新保存onnx.save(model,model_checked.onnx)TensorRT 最终读取的model.onnx本质上就是 ONNXModelProto序列化后的文件。2.3 onnx_ironnx_ir是较新的 ONNX Python IRIntermediate Representation中间表示库。它主要用于图结构构建图遍历Node / Value 操作Graph RewriteONNX 转换新版 ONNX 导出器内部表示。相比直接修改onnx.ModelProtoonnx_ir更适合现代 ONNX exporter 在 Python 层处理计算图。一般项目代码不会直接使用它但新版torch.onnx和onnxscript可能会间接依赖它。2.4 onnxscriptonnxscript用于通过 Python 表达和生成 ONNX 算子、Function 和计算图。它在较新的 PyTorch ONNX 导出链中非常重要。典型作用PyTorch Operator │ ▼ onnxscript │ ▼ ONNX Operator / ONNX Function例如新版 PyTorch 导出方式torch.onnx.export(model,args,model.onnx,dynamoTrue,)内部可能大量依赖onnxscript onnx_ir onnx完成算子转换和图生成。3. 四个包之间的关系可以简单理解为ml_dtypes │ │ 特殊数据类型支持 ▼ PyTorch Model │ │ torch.onnx.export() ▼ onnxscript │ │ PyTorch 算子 - ONNX 表达 ▼ onnx_ir │ │ 中间图表示 / 图转换 ▼ onnx │ │ ModelProto 序列化 ▼ model.onnx │ ▼ TensorRT │ ▼ FP16 Engine在实际项目中重要程度可以理解为包是否经常直接使用主要作用ml_dtypes否特殊 ML dtypeonnx是ONNX 模型读取、检查、保存onnx_ir否新版 ONNX 图中间表示onnxscript通常间接PyTorch 算子转换为 ONNX4. PyTorch 到 ONNX 的典型转换基础形式torch.onnx.export(model,example_inputs,model.onnx,)推荐明确指定torch.onnx.export(model,example_inputs,model.onnx,opset_version17,input_names[observation.images.top,observation.state,],output_names[action,],)对于新版 PyTorch也可能使用torch.onnx.export(model,example_inputs,model.onnx,opset_version17,dynamoTrue,)是否使用dynamoTrue需要结合当前 LeRobot、PyTorch、onnxscript 和 TensorRT 版本决定。如果旧 exporter 已经能够稳定导出则不一定需要切换到 dynamo exporter。5. 当前 ACT 模型输入输出当前预训练 ACT 模型主要输入包括observation.images.top observation.state其中observation.images.top通常为[B, 3, H, W]例如[1, 3, 480, 640]机器人状态observation.state为 14 维[B, 14]模型最终输出 action一般包含 action chunk[B, chunk_size, action_dim]当前模型action_dim 14 chunk_size 100因此常见原始输出可理解为[1, 100, 14]具体 ONNX 输出是否保留完整 chunk应以实际导出脚本为准。6. 为什么需要 dataset 的 stats.json预训练模型并不只是输入图像 输入 state 输出 actionLeRobot 通常还涉及输入预处理 Normalize 模型推理 Unnormalize 输出后处理数据集meta/stats.json提供模型训练时使用的数据统计信息例如mean std min max可能用于observation.state normalization action normalization因此只有model.safetensors通常不足以完全复现原始推理流程。当前离线导出流程会结合model/config.json dataset/meta/stats.json恢复必要的预处理和后处理逻辑。7. 推荐的 ONNX 导出产物建议最终输出目录act_policy/ ├── model.onnx ├── deployment.json └── golden.ptmodel.onnxTensorRT 的输入模型。deployment.json记录部署相关信息例如输入名称 输出名称 输入 shape action_dim chunk_size 图像尺寸 归一化方式避免 ROS2 / TensorRT runtime 中硬编码过多模型参数。golden.pt用于保存 PyTorch 原始模型的一组基准输入输出。例如{inputs:...,outputs:...}后续用于比较PyTorch ONNX Runtime TensorRT FP32 TensorRT FP16之间的数值误差。8. ONNX 导出后必须进行检查8.1 ONNX 结构检查执行importonnx modelonnx.load(model.onnx)onnx.checker.check_model(model)print(ONNX model check passed.)如果onnx.checker.check_model通过只说明 ONNX 图结构合法。不代表输出数值一定正确 TensorRT 一定支持8.2 查看输入输出可以执行importonnx modelonnx.load(model.onnx)forxinmodel.graph.input:print(INPUT :,x.name)forxinmodel.graph.output:print(OUTPUT:,x.name)确认名称和 TensorRT runtime 中使用的名称一致。9. ONNX Runtime 数值验证推荐在进入 TensorRT 之前先做PyTorch vs ONNX Runtime比较。安装pipinstallonnxruntime如果是 CPUimportonnxruntimeasort sessionort.InferenceSession(model.onnx,providers[CPUExecutionProvider],)运行outputssession.run(None,{observation.images.top:image_numpy,observation.state:state_numpy,},)然后和 PyTorch 输出比较importnumpyasnp diffnp.abs(torch_output-onnx_output)print(max abs diff:,diff.max())print(mean abs diff:,diff.mean())如果 ONNX Runtime 与 PyTorch 差异已经很大则不能直接进入 TensorRT 调试。10. TensorRT 转换建议ONNX 验证通过后再进行 TensorRT 转换。典型model.onnx │ ▼ TensorRT Builder │ ├── FP32 └── FP16 ▼ model.engine建议先验证ONNX - TensorRT FP32再开启FP16这样更容易定位问题。如果 FP32 正常而 FP16 有明显误差则问题主要集中在FP16 精度 算子数值稳定性 TensorRT tactic而不是 ONNX 转换本身。11. 推荐验证顺序完整验证链建议Step 1 PyTorch 原模型推理 │ ▼ 保存 golden.pt Step 2 导出 model.onnx │ ▼ onnx.checker.check_model() Step 3 ONNX Runtime CPU │ ▼ 和 PyTorch golden 对比 Step 4 TensorRT FP32 │ ▼ 和 ONNX / PyTorch 对比 Step 5 TensorRT FP16 │ ▼ 和 FP32 / PyTorch 对比 Step 6 接入 ROS2不要直接从PyTorch跳到ROS2 TensorRT FP16否则出了问题很难判断是哪一层导致。12. 常见 ONNX 导出错误12.1 UnsupportedOperatorError例如UnsupportedOperatorError说明某个 PyTorch operator 没有对应的 ONNX symbolic / translation。处理方向尝试更高 opset检查新版 dynamo exporter替换不支持算子自定义 ONNX operator对模型做 wrapper。12.2 onnxscript conversion error如果出现onnxscript conversion error通常需要检查torch onnx onnxscript onnx_ir版本兼容性。不要只升级单个包。12.3 Dynamic shape 问题机器人部署一般输入尺寸固定例如batch 1 image 480x640 state 14如果没有动态 batch / 动态分辨率需求建议优先导出固定 shape。优点TensorRT optimization profile 更简单 engine 更稳定 构建速度更快 推理性能更容易达到最佳12.4 TensorRT 不支持某个 ONNX 算子即使onnx.checker通过也可能出现 TensorRT parser 错误。例如Unsupported node Unsupported operator No importer registered这时需要检查ONNX opset TensorRT version 算子实现必要时对 ONNX Graph 进行 simplify / rewrite。13. 当前项目环境建议ONNX 导出环境至少应包括python torch torchvision lerobot safetensors numpy onnx onnxscript onnx_ir ml_dtypes验证阶段建议增加onnxruntimeTensorRT 环境另外需要TensorRT CUDA cuDNN取决于具体部署注意模型导出环境和TensorRT 部署环境不一定必须是同一个 Python 环境。当前没有训练环境并不影响 ONNX 导出。只要可以正确加载model.safetensors config.json stats.json并完成一次模型 forward就可以进行导出。14. 当前资产目录当前下载完成后的结构lerobot_assets/ ├── dataset/ │ └── aloha_sim_transfer_cube_human/ │ ├── data/ │ │ └── chunk-000/ │ │ ├── file-000.parquet │ │ ├── file-001.parquet │ │ └── file-002.parquet │ ├── meta/ │ │ ├── episodes/ │ │ │ └── chunk-000/ │ │ │ └── file-000.parquet │ │ ├── info.json │ │ ├── stats.json │ │ └── tasks.parquet │ ├── README.md │ └── videos/ │ └── observation.images.top/ │ └── chunk-000/ │ └── file-000.mp4 │ └── model/ └── act_aloha_sim_transfer_cube_human/ ├── config.json ├── demo.gif ├── eval_info.json ├── model.safetensors ├── README.md ├── train_config.json └── training_curves.png已经具备进行离线 ONNX 导出的主要模型和验证数据。15. 下一步建议下一步执行python export_act_onnx_offline.py\--pretrained./lerobot_assets/model/act_aloha_sim_transfer_cube_human\--dataset-root ./lerobot_assets/dataset/aloha_sim_transfer_cube_human\--output-dir ./act_policy\--devicecpu目标生成act_policy/ ├── model.onnx ├── deployment.json └── golden.pt之后依次完成PyTorch ↓ ONNX Runtime ↓ TensorRT FP32 ↓ TensorRT FP16 ↓ ROS2 Runtime并使用同一组 golden 输入输出进行数值一致性验证。