推理系统

• 理解推理系统，包含推理管线，核心组件，推理流程，与训练环境的区别，用户输入处理，实时性能，观测保存，参考运动可视化和部署选项等

概述

推理系统（Inference System）负责将训练好的策略在 MuJoCo 仿真中实时运行。它实现了从训练到部署的完整管线，允许用户在可视化环境中观察训练好的策略行为，并通过键盘或摇杆与仿真交互。

推理管线总览

训练阶段                         推理阶段
─────────                       ─────────

┌──────────────┐               ┌──────────────┐
│  PPO 训练      │               │  ONNX 模型     │
│  (runner.py)  │               │  (.onnx)      │
└──────┬───────┘               └──────┬───────┘
       │                              │
       ▼                              ▼
┌──────────────┐               ┌──────────────┐
│  JAX 参数     │    export_    │  OnnxInfer    │
│  (Flax)      │──────onnx───→│  (onnx_infer  │
└──────────────┘               │   .py)        │
                               └──────┬───────┘
                                       │
                                       ▼
                               ┌──────────────┐
                               │  MjInfer      │
                               │  (mujoco_     │
                               │   infer.py)   │
                               └──────┬───────┘
                                       │
                          ┌────────────┴────────────┐
                          │                         │
                          ▼                         ▼
                  ┌──────────────┐         ┌──────────────┐
                  │  MuJoCo 仿真  │         │  用户输入      │
                  │  (实时渲染)    │←────────│  (键盘/摇杆)   │
                  └──────────────┘         └──────────────┘

核心组件

OnnxInfer

负责加载 ONNX 模型并执行前向推理：

class OnnxInfer:
    def __init__(self, onnx_model_path, input_name="obs", awd=False):
        self.ort_session = onnxruntime.InferenceSession(
            onnx_model_path, providers=["CPUExecutionProvider"])
    
    def infer(self, inputs):
        outputs = self.ort_session.run(None, {self.input_name: [inputs]})
        return outputs[0][0]

MJInferBase

MuJoCo 推理的基类，加载机器人模型并处理关节/传感器查询：

class MJInferBase:
    def __init__(self, model_path):
        self.model = mujoco.MjModel.from_xml_string(...)
        self.data = mujoco.MjData(self.model)
        # ... 关节管理、传感器地址计算等

MjInfer

主要的推理主类，集成策略网络和仿真循环：

class MjInfer(MJInferBase):
    def __init__(self, model_path, reference_data, onnx_model_path, standing):
        self.policy = OnnxInfer(onnx_model_path, awd=True)
        self.PRM = PolyReferenceMotion(reference_data)  # 可选
        # ... 命令范围、动作历史等初始化

推理流程

1. 加载 ONNX 模型和 MuJoCo 场景
          │
2. 初始化仿真状态 (keyframe "home")
          │
3. 进入主循环
          │
    ┌─────┴─────┐ (循环, ~50 Hz 控制频率)
    │           │
    ▼           │
  用户输入处理   │
  (键盘方向键/  │
   摇杆/头部控制)│
    │           │
    ▼           │
  MuJoCo 仿真步  │
  (mj_step,     │
   decimation=10)│
    │           │
    ▼           │
  控制步 (每     │
  decimation 步) │
    │           │
    ▼           │
  构建观测 →    │
  策略推理 →    │
  合成动作 →    │──→ 继续
  应用动作      │

与训练环境的区别

推理系统的观测构建与训练时略有不同：

用户输入处理

键盘控制

推理系统通过 key_callback 函数处理键盘输入，支持两种模式：

摇杆模式（默认）：

头部控制模式（H 键切换）：

摇杆支持

ref_motion_viewer.py 支持通过 pygame 使用实体游戏手柄：

左摇杆：前进/后退 + 横向
右摇杆（可选）：旋转

实时性能

推理系统以约 50 Hz 的控制频率运行：

self.sim_dt = 0.002        # 仿真步长 (500 Hz)
self.decimation = 10       # 每 10 个仿真步执行一次控制 (50 Hz)

每帧循环：

1. mj_step()：10 次仿真步进 (500 Hz)
2. 策略推理：ONNX Runtime CPU 推理 (~0.1 ms)
3. viewer.sync()：渲染更新

观测保存

推理过程中，观测数据会自动保存，方便后续分析：

self.saved_obs.append(obs)
# 按 Ctrl+C 退出时保存到文件
pickle.dump(self.saved_obs, open("mujoco_saved_obs.pkl", "wb"))

参考运动可视化

ref_motion_viewer.py 是一个独立的参考运动可视化工具，允许用户在三维环境中观察参考运动数据生成的步态：

uv run playground/open_duck_mini_v2/ref_motion_viewer.py \
    --command 0.0 -0.05 -0.1 \
    --reference_data playground/open_duck_mini_v2/data/polynomial_coefficients.pkl \
    --scene flat_terrain

它支持键盘控制和 pygame 摇杆控制，同时展示参考运动的关节位置。

部署选项

当前推理系统主要用于桌面仿真验证，部署到真实机器人需要：

1. 将 ONNX 模型加载到机器人上的轻量级推理引擎（如 ONNX Runtime）
2. 实现传感器数据读取（IMU、关节编码器）
3. 实现执行器控制输出（PWM 或串行指令）
4. 确保实时性能（通常需要 50 Hz 控制频率）