619 lines
22 KiB
Python
619 lines
22 KiB
Python
#!/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()
|