This commit is contained in:
2026-04-05 07:37:08 +08:00
parent f3300ff1bf
commit d988ed941b
3 changed files with 1319 additions and 0 deletions

618
go1_rl_inference.py Normal file
View File

@@ -0,0 +1,618 @@
#!/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('<default>\n=', '<default>\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()