机器人控制器开发(训练到Jetson本地部署)

训练部署遵循 “在仿真中训练 -> 优化转换 -> 在边缘设备上部署” 的 Sim2Real 流程。

以下是详细的步骤和说明：

整体流程概览

1. 训练阶段：在 Isaac Lab 中训练强化学习策略，并导出模型（通常是 ONNX 格式）。
2. 转换阶段：在 x86 工作站上，将模型转换为 TensorRT 引擎（.plan 或 .engine 文件）。
3. 部署阶段：将 TensorRT 引擎文件和推理代码移植到 Jetson 设备。
4. 集成阶段：在 Jetson 上编写应用程序，获取真实传感器数据，用 TensorRT 引擎进行推理，并将输出动作发送给真实的机器人执行器。

第一阶段：在 Isaac Lab 中训练和导出模型

Isaac Lab 通常使用 PyTorch 作为后端深度学习框架。训练完成后，你需要将训练好的策略网络（Policy Network）导出为一个可用于部署的格式。

最关键的一步：导出为 ONNX 格式。

ONNX (Open Neural Network Exchange) 是一种开放的模型格式，是 PyTorch/TensorFlow 等训练框架与 TensorRT 等推理引擎之间的桥梁。

1. 在你的训练代码中加入导出逻辑：
   通常是在训练结束后，或者找到一个好的检查点（checkpoint）时。假设你的策略网络是一个 torch.nn.Module 实例 policy_net。

   python
   在你的 isaac lab 训练脚本中，训练完成后添加类似代码
   import torch

   1. 切换到评估模式
      policy_net.eval()

   2. 创建一个示例输入张量（ dummy input ）
      这个输入必须和你的网络在实际接收的输入维度、类型完全一致。
      例如，你的输入可能是 batch_size, observation_dim
      对于部署，batch_size 通常为 1（处理单帧数据），但也可以使用动态批次。
      example_obs = torch.randn(1, observation_dim, device=“cuda”) 注意维度要和你的环境观察值一致

   3. 导出模型为 ONNX
      torch.onnx.export(
      policy_net, 要导出的模型
      example_obs, 模型输入示例
      “policy_net.onnx”, 导出的 ONNX 文件名
      input_names=“observations”, 输入节点名称
      output_names=“actions”, 输出节点名称
      dynamic_axes={ 指定动态维度（非常重要！）
      “observations”: {0: “batch_size”}, 第0维（批次维）是动态的
      “actions”: {0: “batch_size”}
      },
      opset_version=17, 使用较高的 ONNX 算子集版本，兼容性更好
      verbose=True
      )

   • 关键点：dynamic_axes 允许你定义动态维度（如可变的批次大小），这使模型更加灵活。
   • 确保你的网络在推理模式下不会产生随机行为（例如，关闭 Dropout，使用确定的行动采样方式）。

第二阶段：转换 ONNX 模型为 TensorRT 引擎

你可以在 x86 工作站上完成转换，也可以直接在 Jetson 上转换。推荐在 x86 工作站上完成，因为转换过程计算密集，工作站速度更快。

1. 在 x86 工作站上安装 TensorRT
   从 NVIDIA 官网下载 TensorRT 的 tar 包安装文件，或者使用 pip install tensorrt。

2. 使用 trtexec 命令行工具进行转换
   trtexec 是 TensorRT 自带的一个非常强大的命令行工具。
   bash
   trtexec --onnx=policy_net.onnx
   –saveEngine=policy_net.plan
   –workspace=1024 \ 分配显存用于优化过程
   –fp16 启用 FP16 精度，极大提升速度，精度损失通常很小

   • 常用参数：
     • --fp16： 生成 FP16 精度的引擎，在 Jetson 上强烈推荐，速度提升显著。
     • --int8： 生成 INT8 精度的引擎，需要校准数据，速度最快，但需要量化校准。
     • --bestGraph： 尝试进一步优化网络。
     • --workspace： 设置优化过程中可使用的最大显存（MB）。
   • 生成的 policy_net.plan 文件就是优化后的 TensorRT 引擎。

第三阶段：在 Jetson 上部署和推理

1. 设置 Jetson 环境

   • 确保你的 Jetson 设备已刷入最新的 JetPack SDK，它包含了 TensorRT、CUDA、cuDNN 等所有必要组件。
   • 将转换好的 policy_net.plan 文件传输到 Jetson 上。

2. 编写 C++ 或 Python 推理代码

   • Python 示例更简洁，适合快速原型开发。
   • C++ 示例性能更优，是最终部署的首选。

   Python 推理示例代码：

   python
   import tensorrt as trt
   import pycuda.driver as cuda
   import pycuda.autoinit 初始化 CUDA 上下文
   import numpy as np

   class TensorRTInference:
   def init(self, engine_path):
   1. 加载 TensorRT 引擎
   self.logger = trt.Logger(trt.Logger.WARNING)
   with open(engine_path, “rb”) as f, trt.Runtime(self.logger) as runtime:
   self.engine = runtime.deserialize_cuda_engine(f.read())

        2. 创建执行上下文
       self.context = self.engine.create_execution_context()
   
        3. 分配输入输出内存（Host and Device）
       self.inputs, self.outputs, self.bindings = , , 
       self.stream = cuda.Stream()
   
       for binding in self.engine:
           size = trt.volume(self.engine.get_binding_shape(binding))  计算字节数
           dtype = trt.nptype(self.engine.get_binding_dtype(binding))
            在 GPU 上分配内存
           device_mem = cuda.mem_alloc(size * dtype.itemsize)
           self.bindings.append(int(device_mem))
            在 CPU 上分配页锁定内存（pinned memory）以加快数据传输
           host_mem = cuda.pagelocked_empty(size, dtype)
            根据是输入还是输出，添加到相应的列表
           if self.engine.binding_is_input(binding):
               self.inputs.append({'host': host_mem, 'device': device_mem})
           else:
               self.outputs.append({'host': host_mem, 'device': device_mem})
   
   def infer(self, input_obs):
       """
       input_obs: numpy array 包含当前的观测值
       returns: numpy array 包含网络输出的动作
       """
        1. 将输入数据复制到 CPU 的输入内存中
       np.copyto(self.inputs0'host', input_obs.ravel())  确保输入是一维的
   
        2. 将数据从 CPU (Host) 拷贝到 GPU (Device)
       cuda.memcpy_htod_async(self.inputs0'device', self.inputs0'host', self.stream)
   
        3. 执行推理
       self.context.execute_async_v2(bindings=self.bindings, stream_handle=self.stream.handle)
   
        4. 将结果从 GPU 拷贝回 CPU
       cuda.memcpy_dtoh_async(self.outputs0'host', self.outputs0'device', self.stream)
   
        5. 同步 stream，等待所有操作完成
       self.stream.synchronize()
   
        6. 返回输出数据
       return self.outputs0'host'.copy().reshape(-1)   reshape 成你需要的动作维度
   

   使用示例
   trt_model = TensorRTInference(“policy_net.plan”)
   假设你从相机/传感器获取了新的观测值
   current_observation = np.random.randn(1, observation_dim).astype(np.float32)
   action = trt_model.infer(current_observation)
   print(f"Predicted action: {action}")
   将这个 action 发送给你的机器人电机/执行器

重要注意事项和最佳实践

1. 数据预处理/后处理：确保你在 Jetson 上对传感器数据的预处理（归一化、缩放等）与 Isaac Lab 训练时的处理完全一致。同样，对网络输出的后处理也要一致。
2. 性能监控：使用 sudo tegrastats 命令监控 Jetson 的 CPU、GPU、RAM 使用情况，确保推理速度满足你的机器人实时控制要求。
3. 精度：fp16 模式在大多数情况下工作良好。如果遇到问题，可以回退到 fp32 模式进行比较。
4. 动态形状：如果你的导出步骤正确配置了 dynamic_axes，你可以在推理时使用不同的批次大小，但需要在执行推理前调用 context.set_binding_shape(...) 来设置具体的输入维度。

通过以上步骤，你就可以将在虚拟世界中训练出的“大脑”，高效地部署到真实的机器人“身体”（Jetson）中，并利用 TensorRT 发挥出极致的性能。