#!/usr/bin/env python3 """ Go1 RL Policy 推理 - 使用 ONNX Runtime + MuJoCo 可视化 训练模型: /opt/IsaacLab/logs/rsl_rl/unitree_go1_flat/2026-03-24_00-25-36/model_1999.pt 使用方法: python3 /home/8x54zj-m/unitree_mujoco/go1_rl_inference.py 运行此脚本需要: pip install mujoco onnxruntime numpy 键盘控制 (无需 root): W/S: 前进/后退 A/D: 左转/右转 Q/E: 侧向左移/右移 空格: 停止 ESC 或 Ctrl+C: 退出查看器 """ import numpy as np import mujoco from mujoco import viewer import os import sys import tty import termios import signal import threading import time # ============================================================ # 全局退出标志 # ============================================================ g_exit_requested = False def signal_handler(signum, frame): """处理 Ctrl+C (SIGINT) 信号""" global g_exit_requested g_exit_requested = True # 注册信号处理器 signal.signal(signal.SIGINT, signal_handler) # ============================================================ # 配置 # ============================================================ ONNX_PATH = "/opt/IsaacLab/logs/rsl_rl/unitree_go1_flat/2026-03-24_00-25-36/exported/policy.onnx" XML_PATH = "/home/8x54zj-m/unitree_mujoco/data/go1/xml/go1.xml" MESH_PATH = "/home/8x54zj-m/unitree_mujoco/data/go1/meshes" MLP_PATH = "/opt/MotrixLab/walk-these-ways/resources/actuator_nets/unitree_go1.pt" # 观测/动作维度 OBS_DIM = 48 ACTION_DIM = 12 # 关节顺序映射 (policy输出 -> MuJoCo关节) # MuJoCo 关节顺序: [FR_hip, FR_thigh, FR_calf, FL_hip, FL_thigh, FL_calf, RR_hip, RR_thigh, RR_calf, RL_hip, RL_thigh, RL_calf] # Policy 输出顺序取决于 IsaacLab/USD 文件中的定义 # 腿部排列共有 4! = 24 种可能性 # # 测试组 (腿部顺序映射): # 0: FR-FL-RR-RL (MuJoCo默认顺序) # 1: FR-FL-RL-RR # 2: FR-RR-FL-RL # 3: FR-RR-RL-FL # 4: FR-RL-FL-RR # 5: FR-RL-RR-FL # 6: FL-FR-RR-RL # 7: FL-FR-RL-RR # 8: FL-RR-FR-RL # 9: FL-RR-RL-FR # 10: FL-RL-FR-RR # 11: FL-RL-RR-FR # 12: RR-FR-FL-RL # 13: RR-FR-RL-FL 。。 # 14: RR-FL-FR-RL # 15: RR-FL-RL-FR # 16: RR-RL-FR-FL # 17: RR-RL-FL-FR # 18: RL-FR-FL-RR 。。 # 19: RL-FR-RR-FL 。。 # 20: RL-FL-FR-RR # 21: RL-FL-RR-FR # 22: RL-RR-FR-FL # 23: RL-RR-FL-FR 。。 ACTION_MAPPING = 0 # 选择测试组 (0-23) # 腿部顺序定义 LEG_MAPPING = [ [0, 1, 2, 3], # 0: FR-FL-RR-RL (MuJoCo顺序) [0, 1, 3, 2], # 1: FR-FL-RL-RR [0, 2, 1, 3], # 2: FR-RR-FL-RL [0, 2, 3, 1], # 3: FR-RR-RL-FL [0, 3, 1, 2], # 4: FR-RL-FL-RR [0, 3, 2, 1], # 5: FR-RL-RR-FL [1, 0, 2, 3], # 6: FL-FR-RR-RL [1, 0, 3, 2], # 7: FL-FR-RL-RR [1, 2, 0, 3], # 8: FL-RR-FR-RL [1, 2, 3, 0], # 9: FL-RR-RL-FR [1, 3, 0, 2], # 10: FL-RL-FR-RR [1, 3, 2, 0], # 11: FL-RL-RR-FR [2, 0, 1, 3], # 12: RR-FR-FL-RL [2, 0, 3, 1], # 13: RR-FR-RL-FL [2, 1, 0, 3], # 14: RR-FL-FR-RL [2, 1, 3, 0], # 15: RR-FL-RL-FR [2, 3, 0, 1], # 16: RR-RL-FR-FL [2, 3, 1, 0], # 17: RR-RL-FL-FR [3, 0, 1, 2], # 18: RL-FR-FL-RR [3, 0, 2, 1], # 19: RL-FR-RR-FL [3, 1, 0, 2], # 20: RL-FL-FR-RR [3, 1, 2, 0], # 21: RL-FL-RR-FR [3, 2, 0, 1], # 22: RL-RR-FR-FL [3, 2, 1, 0], # 23: RL-RR-FL-FR ] # 根据选择生成 DOF_MAPPING # 每个腿有3个关节(hip, thigh, calf) leg_order = LEG_MAPPING[ACTION_MAPPING] DOF_MAPPING = np.array([ leg_order[0]*3 + 0, leg_order[0]*3 + 1, leg_order[0]*3 + 2, # 腿0 leg_order[1]*3 + 0, leg_order[1]*3 + 1, leg_order[1]*3 + 2, # 腿1 leg_order[2]*3 + 0, leg_order[2]*3 + 1, leg_order[2]*3 + 2, # 腿2 leg_order[3]*3 + 0, leg_order[3]*3 + 1, leg_order[3]*3 + 2, # 腿3 ]) # 速度命令范围 MAX_LIN_VEL = 1.0 # 最大线速度 m/s MAX_ANG_VEL = 1.0 # 最大角速度 rad/s # ============================================================ # 缩放参数 # ============================================================ ACTION_SCALE = 0.15 # RL 输出动作的缩放 (映射到关节位置范围) POS_SCALE = -1.0 # MLP 位置误差缩放 VEL_SCALE = 1.0 # MLP 速度缩放 TORQUE_SCALE = 1.0 # MLP torque 输出缩放 TORQUE_LIMIT = 23.7 # torque 限幅 (N·m) # MuJoCo 摩擦系数配置 # friction = [滑动摩擦(X), 滑动摩擦(Y), 扭转摩擦] # - friction[0]: X方向滑动摩擦系数 # - friction[1]: Y方向滑动摩擦系数 (各向异性材料用) # - friction[2]: 扭转摩擦系数 (绕法线的旋转阻力) # # 常见地面摩擦参数参考: # ┌──────────┬───────────────────┐ # │ 地面类型 │ FLOOR_FRICTION │ # ├──────────┼───────────────────┤ # │ 冰面 │ [0.05, 0.02, 0.01]│ # │ 木地板 │ [0.4, 0.3, 0.2 ]│ # │ 瓷砖 │ [0.6, 0.4, 0.3 ]│ # │ 橡胶垫 │ [1.5, 1.0, 0.5 ]│ # │ 粗糙地面 │ [2.0, 1.5, 1.0 ]│ # └──────────┴───────────────────┘ # 机器人 geom 摩擦系数 FRICTION = [0.6, 0.3, 0.3] # 地面摩擦系数 (单独设置) FLOOR_FRICTION = [2.0, 1.5, 1.0] # ============================================================ # 键盘输入读取 (使用后台线程，无需 root 权限) # ============================================================ class KeyboardReader: """使用后台线程读取键盘输入，不需要 root 权限""" def __init__(self): self.old_settings = None self.keys_pressed = set() self.running = True self.thread = None def init(self): """初始化终端并启动键盘读取线程""" self.old_settings = termios.tcgetattr(sys.stdin) tty.setcbreak(sys.stdin.fileno()) self.thread = threading.Thread(target=self._read_loop, daemon=True) self.thread.start() def _read_loop(self): """后台线程：持续读取键盘输入""" try: while self.running: try: import select if select.select([sys.stdin], [], [], 0.1)[0]: ch = sys.stdin.read(1) if ch: if ch == '\x1b': # ESC self.keys_pressed.add('\x1b') elif ch == ' ': self.keys_pressed.add('space') else: self.keys_pressed.add(ch.lower()) except: pass except: pass def restore(self): """恢复终端设置""" self.running = False if self.thread: self.thread.join(timeout=1.0) if self.old_settings: termios.tcsetattr(sys.stdin, termios.TCSADRAIN, self.old_settings) def is_key_pressed(self, key): """检查按键是否被按下""" if key == '\x1b': return '\x1b' in self.keys_pressed elif key == ' ': return 'space' in self.keys_pressed return key in self.keys_pressed # ============================================================ # 辅助函数 # ============================================================ def get_joint_state(model, data, joint_name): """通过关节名获取关节位置""" joint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, joint_name) if joint_id < 0: return None qposadr = model.jnt_qposadr[joint_id] return data.qpos[qposadr] def compute_observations(data, prev_action, command, default_joint_pos): """ 计算 RL policy 观测向量 (48维) 对应 IsaacLab 观测配置: [0:3] base_lin_vel - 躯干坐标系线速度 [3:6] base_ang_vel - 躯干坐标系角速度 [6:9] projected_gravity - 重力投影到躯干坐标系 [9:12] velocity_commands - 期望速度 [12:24] joint_pos_rel - 相对关节位置 (相对于默认位置) [24:36] joint_vel_rel - 相对关节速度 (相对于默认速度) [36:48] prev_actions - 上一步动作 """ obs = np.zeros(OBS_DIM, dtype=np.float32) # 获取躯干四元数 (MuJoCo 格式: qx, qy, qz, qw) quat = data.qpos[3:7] # (x, y, z, w) # 计算旋转矩阵 (从世界坐标系到躯干坐标系) R = quaternion_to_rotation_matrix(quat).T # 转置得到逆变换 (世界->躯干) # 1. 躯干线速度 - 转换到躯干坐标系 # IsaacLab: root_lin_vel_b = R @ root_lin_vel_w global_lin_vel = data.qvel[0:3] obs[0:3] = R @ global_lin_vel # 2. 躯干角速度 - 转换到躯干坐标系 # IsaacLab: root_ang_vel_b = R @ root_ang_vel_w global_ang_vel = data.qvel[3:6] obs[3:6] = R @ global_ang_vel # 3. 重力投影到躯干坐标系 # IsaacLab: projected_gravity_b = R @ gravity_world # 重力向量 (0, 0, -9.81) 在世界坐标系中 gravity_world = np.array([0.0, 0.0, -9.81]) obs[6:9] = R @ gravity_world # 4. 速度命令 obs[9:12] = command # 5. 相对关节位置 (joint_pos - default_joint_pos) # IsaacLab: joint_pos_rel = joint_pos - default_joint_pos current_joint_pos = data.qpos[7:19] obs[12:24] = current_joint_pos - default_joint_pos # 6. 相对关节速度 # IsaacLab: joint_vel_rel = joint_vel - default_joint_vel (default = 0) default_joint_vel = np.zeros(12) # IsaacLab 默认关节速度为 0 obs[24:36] = data.qvel[6:18] - default_joint_vel # 7. 上一步动作 obs[36:48] = prev_action return obs def quaternion_to_rotation_matrix(q): """四元数转旋转矩阵 (MuJoCo 格式: qx, qy, qz, qw) 返回的矩阵将世界坐标系中的向量转换到以该四元数表示的局部坐标系中 """ qx, qy, qz, qw = q # 归一化确保数值稳定性 norm = np.sqrt(qx**2 + qy**2 + qz**2 + qw**2) qx, qy, qz, qw = qx/norm, qy/norm, qz/norm, qw/norm return np.array([ [1-2*(qy**2+qz**2), 2*(qx*qy-qz*qw), 2*(qx*qz+qy*qw)], [2*(qx*qy+qz*qw), 1-2*(qx**2+qz**2), 2*(qy*qz-qx*qw)], [2*(qx*qz-qy*qw), 2*(qy*qz+qx*qw), 1-2*(qx**2+qy**2)] ]) # ============================================================ # MLP Actuator Model (从 IsaacLab 迁移) # ============================================================ class ActuatorMLP: """ActuatorNetMLP 模型 - 用于将 position target 转换为 torque 参考: IsaacLab source/isaaclab/isaaclab/actuators/actuator_net.py 使用批量处理，所有 12 个关节一次前向传播 """ def __init__(self, model_path, num_joints=12): import torch self.num_joints = num_joints # 强制使用 CPU 避免 CUDA 版本问题 self.network = torch.jit.load(model_path, map_location='cpu') self.network.eval() # GO1 MLP 配置 (来自 GO1_ACTUATOR_CFG) self.pos_scale = POS_SCALE self.vel_scale = VEL_SCALE self.torque_scale = TORQUE_SCALE self.input_idx = [0, 1, 2] # 3步历史 self.input_order = "pos_vel" # 历史缓冲区 (每个关节独立) history_len = max(self.input_idx) + 1 # = 3 self.pos_error_history = np.zeros((history_len, num_joints), dtype=np.float32) self.vel_history = np.zeros((history_len, num_joints), dtype=np.float32) # Torque 限幅 (来自 URDF 和 GO1_ACTUATOR_CFG) self.effort_limit = TORQUE_LIMIT # N·m def reset(self): """重置历史缓冲区""" self.pos_error_history.fill(0.0) self.vel_history.fill(0.0) def compute_torque(self, pos_target, current_pos, current_vel): """计算 torque (使用批量处理，匹配 IsaacLab 实现) Args: pos_target: 目标位置 (12,) current_pos: 当前关节位置 (12,) current_vel: 当前关节速度 (12,) Returns: torque: 力矩命令 (12,) """ import torch # 1. 计算 position error pos_error = pos_target - current_pos # (12,) # 2. 更新历史 (移动队列) self.pos_error_history = np.roll(self.pos_error_history, 1, axis=0) self.pos_error_history[0] = pos_error self.vel_history = np.roll(self.vel_history, 1, axis=0) self.vel_history[0] = current_vel # 3. 批量计算 torque (IsaacLab 风格，一次性处理所有关节) # pos_error_history: (history_len, num_joints) = (3, 12) # vel_history: (history_len, num_joints) = (3, 12) # # IsaacLab 构建输入的方式: # pos_input = torch.cat([history[:, i].unsqueeze(2) for i in input_idx], dim=2) # -> (num_envs, num_joints, 3) # pos_input = pos_input.view(num_envs * num_joints, -1) # -> (12, 3) # 然后 concat pos_input 和 vel_input 得到 (12, 6) with torch.inference_mode(): # 构建 pos_input: (num_joints, 3) = (12, 3) # 从 history 矩阵中提取 input_idx 指定的行，然后转置 # pos_error_history[self.input_idx, :] shape = (3, 12) # 转置后变成 (12, 3)，每行是一个关节的 3 步位置误差历史 pos_input = self.pos_error_history[self.input_idx, :].T * self.pos_scale # (12, 3) # 构建 vel_input: (num_joints, 3) = (12, 3) vel_input = self.vel_history[self.input_idx, :].T * self.vel_scale # (12, 3) # 合并为网络输入: (12, 6) if self.input_order == "pos_vel": network_input = np.concatenate([pos_input, vel_input], axis=1) else: network_input = np.concatenate([vel_input, pos_input], axis=1) # 批量推理 (12, 6) -> (12, 1) input_tensor = torch.from_numpy(network_input).float() # (12, 6) torque_output = self.network(input_tensor).numpy().flatten() # (12,) # 限幅 torques = np.clip(torque_output * self.torque_scale, -self.effort_limit, self.effort_limit) return torques.astype(np.float32) # ============================================================ # 主程序 # ============================================================ def main(): # 1. 加载 MuJoCo 模型 os.chdir(MESH_PATH) with open(XML_PATH, 'r') as f: xml_content = f.read() xml_content = xml_content.replace('meshdir="../meshes/"', f'meshdir="{MESH_PATH}"') xml_content = xml_content.replace('\n=', '\n') model = mujoco.MjModel.from_xml_string(xml_content) data = mujoco.MjData(model) print(f"[INFO] MuJoCo Model: {model.nbody} bodies, {model.nq} DoF, {model.nu} actuators") print(f"[INFO] 关节顺序: {[mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, i) for i in range(1, 13)]}") # 2. 设置摩擦系数 floor_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "floor") if floor_id >= 0: model.geom_friction[floor_id] = FLOOR_FRICTION print(f"[INFO] 地面摩擦系数设置为: {FLOOR_FRICTION}") else: print(f"[WARNING] 未找到 floor geom，使用默认摩擦系数") for i in range(model.ngeom): if i != floor_id: model.geom_friction[i] = FRICTION print(f"[INFO] 机器人摩擦系数设置为: {FRICTION}") # 3. 加载 ONNX 策略 import onnxruntime as ort sess = ort.InferenceSession(ONNX_PATH, providers=['CPUExecutionProvider']) print(f"[INFO] ONNX Policy loaded from: {ONNX_PATH}") print(f"[INFO] 观测维度: {OBS_DIM}, 动作维度: {ACTION_DIM}") # 3. 加载 MLP Actuator 模型 (IsaacLab 风格) try: actuator_mlp = ActuatorMLP(MLP_PATH, num_joints=ACTION_DIM) print(f"[INFO] ActuatorMLP loaded from: {MLP_PATH}") use_mlp = True except Exception as e: print(f"[WARNING] Failed to load ActuatorMLP: {e}") print(f"[INFO] Falling back to PD control") use_mlp = False # 4. 初始化 obs = np.zeros(OBS_DIM, dtype=np.float32) prev_action = np.zeros(ACTION_DIM, dtype=np.float32) # IsaacLab UNITREE_GO1_CFG 默认关节位置 (用于计算相对关节位置) # 关节顺序: FR_hip, FR_thigh, FR_calf, FL_hip, FL_thigh, FL_calf, RR_hip, RR_thigh, RR_calf, RL_hip, RL_thigh, RL_calf default_joint_pos = np.array([ -0.1, # FR_hip 0.8, # FR_thigh -1.5, # FR_calf 0.1, # FL_hip 0.8, # FL_thigh -1.5, # FL_calf -0.1, # RR_hip 1.0, # RR_thigh -1.5, # RR_calf 0.1, # RL_hip 1.0, # RL_thigh -1.5, # RL_calf ]) # 重置机器人关节位置 (与 default_joint_pos 一致) # 注意: MuJoCo 关节角度单位是 rad # URDF 关节范围: # Hip: -0.803 ~ 0.803 rad # Thigh: -1.047 ~ 4.189 rad # Calf: -2.697 ~ -0.916 rad crouch_pos = np.array([ -0.1, # FR_hip (URDF: -0.803 ~ 0.803) 0.8, # FR_thigh (URDF: -1.047 ~ 4.189) -1.5, # FR_calf (URDF: -2.697 ~ -0.916) 0.1, # FL_hip 0.8, # FL_thigh -1.5, # FL_calf -0.1, # RR_hip 1.0, # RR_thigh -1.5, # RR_calf 0.1, # RL_hip 1.0, # RL_thigh -1.5, # RL_calf ]) data.qpos[7:19] = crouch_pos # 重置 qvel 为零，避免初始速度导致不稳定 data.qvel[:] = 0.0 # 将躯干抬高一点，避免初始碰撞 data.qpos[2] = 0.35 mujoco.mj_step(model, data) # 执行一步让状态更新 # 重置 MLP actuator 历史 if use_mlp: actuator_mlp.reset() print("[INFO] 机器人初始化完成，开始 RL 策略推理...", flush=True) # 初始化键盘控制 keyboard_reader = KeyboardReader() keyboard_reader.init() print("[INFO] 键盘控制已启用!", flush=True) print(" W/S: 前进/后退", flush=True) print(" A/D: 左转/右转", flush=True) print(" Q/E: 侧向左移/右移", flush=True) print(" 空格: 停止", flush=True) print(" ESC: 退出", flush=True) # 4. 交互式可视化 import mujoco.viewer as mv view = mv.launch_passive(model, data) print("[INFO] 交互式查看器已启动!", flush=True) print(" 鼠标拖拽旋转视角, 滚轮缩放, ESC退出", flush=True) step_count = 0 last_inference_time = time.time() inference_interval = 0.02 # 50Hz = 20ms torque = np.zeros(ACTION_DIM, dtype=np.float32) # 初始化 torque import sys sys.stderr.write(f"[DEBUG] Starting loop\n") sys.stderr.flush() # 命令速度 (使用 numpy array 以便在循环中修改) command = np.array([1.0, 0.0, 0.0], dtype=np.float32) # vx, vy, yaw_rate while view.is_running() and not g_exit_requested: current_time = time.time() # W/S: 前进/后退 (vx) if keyboard_reader.is_key_pressed('w'): command[0] = MAX_LIN_VEL elif keyboard_reader.is_key_pressed('s'): command[0] = -MAX_LIN_VEL else: command[0] = 0.0 # Q/E: 侧向移动 (vy) if keyboard_reader.is_key_pressed('q'): command[1] = MAX_LIN_VEL elif keyboard_reader.is_key_pressed('e'): command[1] = -MAX_LIN_VEL else: command[1] = 0.0 # A/D: 左转/右转 (yaw_rate) if keyboard_reader.is_key_pressed('a'): command[2] = MAX_ANG_VEL elif keyboard_reader.is_key_pressed('d'): command[2] = -MAX_ANG_VEL else: command[2] = 0.0 # 空格: 停止 if keyboard_reader.is_key_pressed(' '): command[:] = 0.0 # ESC: 退出 if keyboard_reader.is_key_pressed('\x1b'): sys.stderr.write("[INFO] ESC pressed, exiting...\n") break # 50Hz 频率限制推理 if current_time - last_inference_time >= inference_interval: # 计算观测 obs = compute_observations(data, prev_action, command, default_joint_pos) # 推理获取动作 action = sess.run(None, {'obs': obs.reshape(1, -1)})[0][0] # 动作缩放 (RL 输出是位置目标, 范围 [-1, 1], 映射到关节范围) # IsaacLab 配置: actions.joint_pos.scale = 0.25 action_scaled = action * ACTION_SCALE # 应用关节顺序映射 (policy输出 -> MuJoCo关节) action_scaled_mapped = action_scaled[DOF_MAPPING] # 获取当前关节状态 current_pos = data.qpos[7:19] current_vel = data.qvel[6:18] # 使用 MLP 计算 torque (与 IsaacLab 一致) if use_mlp: torque = actuator_mlp.compute_torque( action_scaled_mapped, current_pos, current_vel ) else: # Fallback PD 控制 kp = 10.0 kd = 0.0 torque = kp * (action_scaled_mapped - current_pos) - kd * current_vel # 保存上一步动作 prev_action = action.copy() last_inference_time = current_time # 应用控制 (力矩控制) - 每步都应用 data.ctrl[:] = torque # 执行仿真 mujoco.mj_step(model, data) view.sync() step_count += 1 trunk_z = data.qpos[2] lin_vel = np.linalg.norm(data.qvel[0:3]) ang_vel = np.linalg.norm(data.qvel[3:6]) sys.stderr.write(f" Step {step_count}: cmd=[{command[0]:.2f}, {command[1]:.2f}, {command[2]:.2f}] | " f"trunk_z={trunk_z:.3f}m, vel={lin_vel:.3f}m/s\n") sys.stderr.flush() # 恢复终端设置 if g_exit_requested: sys.stderr.write("[INFO] Ctrl+C pressed, exiting...\n") keyboard_reader.restore() view.close() if __name__ == "__main__": main()