纠错,疑问,交流: 请进入讨论区或请点击进入页面，扫码加入微信群或Q群进行交流

获取最新文章: 扫一扫加入“创客智造”公众号

欢迎加入我们的openduckmini交流群，微信扫描右侧二维码立即进群交流

MuJoCo 推理

理解MuJoCo推理，包含MJInferBase 基类，MjInfer 主类，用户输入处理，实时步态周期，动作平滑滤波器等

概述

MuJoCo 推理系统是 Open Duck Playground 的实时仿真运行环境。它通过 MjInfer 类加载训练好的 ONNX 策略模型，在 MuJoCo 物理引擎中进行实时仿真，并通过 mujoco.viewer 提供 3D 可视化界面。用户可以通过键盘或摇杆实时控制机器人的运动。

MJInferBase 基类

mujoco_infer_base.py 中的 MJInferBase 提供推理所需的基础设施：

模型加载

class MJInferBase:
    def __init__(self, model_path):
        self.model = mujoco.MjModel.from_xml_string(
            epath.Path(model_path).read_text(), assets=base.get_assets()
        )
        self.model.opt.timestep = 0.002  # 仿真步长
        self.data = mujoco.MjData(self.model)
        mujoco.mj_step(self.model, self.data)

关节分类

与训练环境相同，推理基类也对关节进行分类：

self.floating_base_name      # 浮动基座名称
self.actuator_names          # 执行器关节列表
self.joint_names             # 所有关节列表
self.backlash_joint_names    # 回差关节列表

self.num_dofs = self.model.nu  # 执行器数量 (14)

传感器访问

推理基类通过传感器地址直接访问数据，避免了字符串查找的开销：

def get_gyro(self, data):
    return data.sensordata[self.gyro_addr:self.gyro_addr + self.gyro_dimensions]

def get_accelerometer(self, data):
    return data.sensordata[self.accelerometer_addr:self.accelerometer_addr + 3]

def get_gravity(self, data):
    return data.sensordata[self.gravity_id:self.gravity_id + 3]

接触检测

足部接触检测通过 MuJoCo 的碰撞检测实现：

def get_feet_contacts(self, data):
    left_contact = self.check_contact(data, "foot_assembly", "floor")
    right_contact = self.check_contact(data, "foot_assembly_2", "floor")
    return left_contact, right_contact

MjInfer 主类

mujoco_infer.py 中的 MjInfer 是推理的主入口。

初始化

class MjInfer(MJInferBase):
    def __init__(self, model_path, reference_data, onnx_model_path, standing):
        # 策略加载
        self.policy = OnnxInfer(onnx_model_path, awd=True)
        
        # 参考运动（非站立模式）
        if not self.standing:
            self.PRM = PolyReferenceMotion(reference_data)
        
        # 动作滤波器
        self.action_filter = LowPassActionFilter(50, cutoff_frequency=37.5)
        
        # 命令范围
        self.COMMANDS_RANGE_X = [-0.15, 0.15]       # 前进速度
        self.COMMANDS_RANGE_Y = [-0.2, 0.2]         # 横向速度
        self.COMMANDS_RANGE_THETA = [-1.0, 1.0]     # 旋转速度
        
        # 动作历史
        self.last_action = np.zeros(self.num_dofs)
        self.last_last_action = np.zeros(self.num_dofs)
        self.last_last_last_action = np.zeros(self.num_dofs)
        
        # 默认姿态
        self.default_actuator = self.model.keyframe("home").ctrl
        self.motor_targets = self.default_actuator.copy()
        self.prev_motor_targets = self.default_actuator.copy()

观测构建

get_obs() 方法将仿真数据转换为与训练时相同的观测向量：

def get_obs(self, data, command):
    gyro = self.get_gyro(data)
    accelerometer = self.get_accelerometer(data)
    joint_angles = self.get_actuator_joints_qpos(data.qpos)
    joint_vel = self.get_actuator_joints_qvel(data.qvel)
    contacts = self.get_feet_contacts(data)
    
    obs = np.concatenate([
        gyro,           # 3
        accelerometer,  # 3
        command,        # 7
        joint_angles - self.default_actuator,  # 10
        joint_vel * self.dof_vel_scale,        # 10
        self.last_action,       # 10
        self.last_last_action,  # 10
        self.last_last_last_action,  # 10
        self.motor_targets,     # 10
        contacts,               # 2
        self.imitation_phase,   # 2
    ])
    return obs

主循环

def run(self):
    with mujoco.viewer.launch_passive(
        self.model, self.data, key_callback=self.key_callback
    ) as viewer:
        counter = 0
        while viewer.is_running():
            step_start = time.time()
            
            # 物理仿真步进（1 步 = 0.002s）
            mujoco.mj_step(self.model, self.data)
            counter += 1
            
            # 控制频率 (50 Hz, 每 10 仿真步)
            if counter % self.decimation == 0:
                # 1. 更新参考运动相位
                if not self.standing:
                    self.imitation_i += 1.0 * self.phase_frequency_factor
                    self.imitation_i %= self.PRM.nb_steps_in_period
                    self.imitation_phase = np.array([
                        cos(self.imitation_i / N * 2 * pi),
                        sin(self.imitation_i / N * 2 * pi),
                    ])
                
                # 2. 构建观测
                obs = self.get_obs(self.data, self.commands)
                self.saved_obs.append(obs)
                
                # 3. 策略推理
                action = self.policy.infer(obs)
                
                # 4. 更新动作历史
                self.last_last_last_action = self.last_last_action.copy()
                self.last_last_action = self.last_action.copy()
                self.last_action = action.copy()
                
                # 5. 计算电机目标
                self.motor_targets = (
                    self.default_actuator + action * self.action_scale
                )
                
                # 6. 电机速度限制
                if USE_MOTOR_SPEED_LIMITS:
                    self.motor_targets = np.clip(
                        self.motor_targets,
                        self.prev_motor_targets - max_vel * dt,
                        self.prev_motor_targets + max_vel * dt,
                    )
                    self.prev_motor_targets = self.motor_targets.copy()
                
                # 7. 应用动作
                self.data.ctrl = self.motor_targets.copy()
            
            # 同步渲染
            viewer.sync()
            
            # 保持实时帧率
            time_until_next_step = self.model.opt.timestep - (time.time() - step_start)
            if time_until_next_step > 0:
                time.sleep(time_until_next_step)

用户输入处理

键盘回调

key_callback 方法处理用户的键盘输入：

def key_callback(self, keycode):
    if keycode == 72:  # 'H' 键
        self.head_control_mode = not self.head_control_mode
    
    if not self.head_control_mode:
        # 摇杆模式：控制运动
        if keycode == 265:  # ↑ 前进
            lin_vel_x = COMMANDS_RANGE_X[1]
        elif keycode == 264:  # ↓ 后退
            lin_vel_x = COMMANDS_RANGE_X[0]
        # ... 横向 (←/→) 和 旋转 (A/E)
    else:
        # 头部控制模式
        # ↑/↓: 颈部俯仰, ←/→: 头部偏航, A/E: 头部滚动
        ...

控制模式

模式	切换方式	功能
摇杆模式	默认	控制机器人线速度和角速度
头部模式	按 H 切换	控制机器人头部 4 个自由度

实时步态周期

在不使用模仿奖励时，推理系统包含一个内置的步态相位机制：

imitation_phase = [cos(phase), sin(phase)]

其中 phase = (imitation_i / nb_steps_in_period) * 2 * pi

phase 范围 [0, 2π]，通过 P/M 键可以调节步态频率。

动作平滑滤波器

LowPassActionFilter 在推理中可用于平滑动作输出：

self.action_filter = LowPassActionFilter(50, cutoff_frequency=37.5)

# 使用滤波器
self.action_filter.push(action)
action = self.action_filter.get_filtered_action()

目前默认启用了速度限制而非低通滤波，速度限制限制了每步关节位置的变化量（max_motor_velocity = 5.24 rad/s），更接近真实舵机的行为。