471 lines
16 KiB
Python
471 lines
16 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Go1 RL Policy 推理 - 移植自 unitree_rl 项目
|
||
基于 https://github.com/dstx123/unitree_rl
|
||
|
||
模型: body.jit + adapt.jit (GRU-based policy with history)
|
||
训练环境: IsaacGym
|
||
|
||
使用方法:
|
||
python3 /home/8x54zj-m/unitree_mujoco/go1_unitree_rl_inference.py
|
||
|
||
键盘控制:
|
||
W/S: 前进/后退
|
||
A/D: 左转/右转
|
||
Q/E: 侧向左移/右移
|
||
空格: 停止
|
||
ESC: 退出
|
||
"""
|
||
import numpy as np
|
||
import mujoco
|
||
from mujoco import viewer
|
||
import os
|
||
import sys
|
||
import tty
|
||
import termios
|
||
import signal
|
||
import threading
|
||
import time
|
||
import torch
|
||
|
||
# ============================================================
|
||
# 全局退出标志
|
||
# ============================================================
|
||
g_exit_requested = False
|
||
|
||
def signal_handler(signum, frame):
|
||
global g_exit_requested
|
||
g_exit_requested = True
|
||
|
||
signal.signal(signal.SIGINT, signal_handler)
|
||
|
||
# ============================================================
|
||
# 配置
|
||
# ============================================================
|
||
BODY_MODEL_PATH = "/tmp/unitree_rl/src/unitree_guide/unitree_guide/model/body.jit"
|
||
ADAPT_MODEL_PATH = "/tmp/unitree_rl/src/unitree_guide/unitree_guide/model/adapt.jit"
|
||
XML_PATH = "/home/8x54zj-m/unitree_mujoco/data/go1/xml/go1.xml"
|
||
MESH_PATH = "/home/8x54zj-m/unitree_mujoco/data/go1/meshes"
|
||
|
||
# ============================================================
|
||
# 模型参数 (来自 unitree_rl/State_RL.h)
|
||
# ============================================================
|
||
NUM_OBS = 45 # 观测维度
|
||
NUM_OBS_HISTORY = 10 # 历史观测步数
|
||
OBS_BUFFER_SIZE = NUM_OBS * NUM_OBS_HISTORY # 450
|
||
|
||
# 观测缩放
|
||
SCALE_LIN_VEL = 2.0
|
||
SCALE_ANG_VEL = 0.25
|
||
SCALE_COMMANDS = np.array([SCALE_LIN_VEL, SCALE_LIN_VEL, SCALE_ANG_VEL])
|
||
SCALE_DOF_POS = 1.0
|
||
SCALE_DOF_VEL = 0.05
|
||
|
||
# 动作缩放
|
||
ACTION_SCALE = 0.25
|
||
HIP_SCALE_REDUCTION = 0.5 # hip关节额外缩放
|
||
|
||
# 限幅
|
||
CLIP_OBSERVATIONS = 100.0
|
||
CLIP_ACTIONS = 100.0
|
||
|
||
# 默认关节角度 (与 unitree_rl 一致)
|
||
DEFAULT_DOF_POS = np.array([
|
||
-0.1, 0.8, -1.5, # FR: hip, thigh, calf
|
||
0.1, 0.8, -1.5, # FL: hip, thigh, calf
|
||
-0.1, 1.0, -1.5, # RR: hip, thigh, calf
|
||
0.1, 1.0, -1.5 # RL: hip, thigh, calf
|
||
])
|
||
|
||
# Motor order in code: FR(3,4,5), FL(0,1,2), RR(9,10,11), RL(6,7,8)
|
||
# Model order: FL(0,1,2), FR(3,4,5), RL(6,7,8), RR(9,10,11)
|
||
# dof_mapping = {3, 4, 5, 0, 1, 2, 9, 10, 11, 6, 7, 8}
|
||
# Inverse mapping (model idx -> motor idx): 3->0, 4->1, 5->2, 0->3, 1->4, 2->5, 9->6, 10->7, 11->8, 6->9, 7->10, 8->11
|
||
# This maps: model_output[0..11] -> motor_q[0..11]
|
||
DOF_MAPPING = np.array([3, 4, 5, 0, 1, 2, 9, 10, 11, 6, 7, 8])
|
||
|
||
# 速度命令范围
|
||
MAX_LIN_VEL = 1.0
|
||
MAX_ANG_VEL = 1.0
|
||
|
||
# PD控制参数 (来自 unitree_rl)
|
||
KP = 7.0
|
||
KD = 1.0
|
||
|
||
# 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]
|
||
|
||
|
||
# ============================================================
|
||
# 键盘输入读取 (termios 非阻塞方式，无 root 依赖)
|
||
# ============================================================
|
||
import select
|
||
|
||
class KeyboardReader:
|
||
def __init__(self):
|
||
self.keys_pressed = set()
|
||
self.last_key_time = 0
|
||
self.timeout = 0.1 # 无按键超时时间（秒），超时后清除所有按键
|
||
|
||
def init(self):
|
||
import termios
|
||
self.old_settings = termios.tcgetattr(sys.stdin)
|
||
# 非canonical模式，关闭回显
|
||
new_settings = termios.tcgetattr(sys.stdin)
|
||
new_settings[3] = new_settings[3] & ~termios.ICANON & ~termios.ECHO
|
||
new_settings[6][termios.VMIN] = 0 # 非阻塞
|
||
new_settings[6][termios.VTIME] = 0
|
||
termios.tcsetattr(sys.stdin, termios.TCSADRAIN, new_settings)
|
||
print("[INFO] Keyboard reader initialized (termios mode, no root required)")
|
||
|
||
def _read_key(self):
|
||
"""尝试读取一个按键，返回按键字符或None"""
|
||
try:
|
||
if select.select([sys.stdin], [], [], 0)[0]:
|
||
ch = sys.stdin.read(1)
|
||
return ch
|
||
except:
|
||
pass
|
||
return None
|
||
|
||
def update(self):
|
||
"""在主循环中调用，更新按键状态"""
|
||
import time
|
||
current_time = time.time()
|
||
|
||
# 尝试读取按键
|
||
ch = self._read_key()
|
||
if ch:
|
||
if ch == '\x1b': # ESC
|
||
self.keys_pressed.add('escape')
|
||
elif ch == ' ':
|
||
self.keys_pressed.add('space')
|
||
elif ch == '\n' or ch == '\r':
|
||
pass
|
||
else:
|
||
self.keys_pressed.add(ch.lower())
|
||
self.last_key_time = current_time
|
||
elif current_time - self.last_key_time > self.timeout:
|
||
# 超时后清除所有按键（模拟按键释放）
|
||
self.keys_pressed.clear()
|
||
|
||
def restore(self):
|
||
import termios
|
||
if hasattr(self, 'old_settings'):
|
||
termios.tcsetattr(sys.stdin, termios.TCSADRAIN, self.old_settings)
|
||
|
||
def is_key_pressed(self, key):
|
||
if key == '\x1b' or key == 'escape':
|
||
return 'escape' in self.keys_pressed
|
||
elif key == ' ':
|
||
return 'space' in self.keys_pressed
|
||
return key.lower() in self.keys_pressed
|
||
|
||
def clear(self):
|
||
self.keys_pressed.clear()
|
||
|
||
|
||
# ============================================================
|
||
# 辅助函数
|
||
# ============================================================
|
||
|
||
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)]
|
||
])
|
||
|
||
|
||
def quat_rotate_inverse(q, v):
|
||
"""四元数逆旋转 (世界坐标系 -> 躯干坐标系)
|
||
|
||
对应 IsaacLab 的 quat_apply_inverse
|
||
q: (w, x, y, z) 格式四元数
|
||
v: (x, y, z) 向量
|
||
"""
|
||
q_w, q_x, q_y, q_z = q[3], q[0], q[1], q[2]
|
||
return np.array([
|
||
v[0] * (2*q_w**2 - 1) + 2*q_x*(q_y*v[2] - q_z*v[1]) + 2*q_w*(q_z*v[0] - q_x*v[2]),
|
||
v[1] * (2*q_w**2 - 1) + 2*q_y*(q_z*v[0] - q_x*v[2]) + 2*q_w*(q_x*v[1] - q_y*v[0]),
|
||
v[2] * (2*q_w**2 - 1) + 2*q_z*(q_x*v[1] - q_y*v[0]) + 2*q_w*(q_y*v[0] - q_x*v[1])
|
||
])
|
||
|
||
|
||
def compute_observations(data, prev_action, commands, default_dof_pos):
|
||
"""
|
||
计算 RL policy 观测向量 (45维)
|
||
|
||
对应 unitree_rl 的观测计算:
|
||
[0:3] body_ang_vel - 躯干角速度 (body frame)
|
||
[3:6] projected_gravity - 重力投影到躯干坐标系
|
||
[6:9] commands - 速度命令 (已缩放)
|
||
[9:21] dof_pos_rel - 相对关节位置
|
||
[21:33] dof_vel - 关节速度
|
||
[33:45] prev_action - 上一步动作
|
||
"""
|
||
obs = np.zeros(NUM_OBS, dtype=np.float32)
|
||
|
||
# 获取四元数 (MuJoCo 格式: x, y, z, w)
|
||
quat_mujoco = data.qpos[3:7] # (x, y, z, w)
|
||
# 转换为 unitree_rl 格式 (x, y, z, w) for quat_rotate_inverse
|
||
|
||
# 1. 躯干角速度 (body frame)
|
||
# unitree_rl 直接使用 imu.gyroscope (已是在 body frame)
|
||
# MuJoCo 的 qvel[3:6] 是世界坐标系，需要转换
|
||
ang_vel_world = data.qvel[3:6]
|
||
ang_vel_body = quat_rotate_inverse(np.concatenate([quat_mujoco[3:4], quat_mujoco[0:3]]), ang_vel_world)
|
||
obs[0:3] = ang_vel_body * SCALE_ANG_VEL
|
||
|
||
# 2. 重力投影 (body frame)
|
||
# gravity_vec = (0, 0, -1) in unitree_rl
|
||
gravity_world = np.array([0.0, 0.0, -1.0])
|
||
projected_gravity = quat_rotate_inverse(np.concatenate([quat_mujoco[3:4], quat_mujoco[0:3]]), gravity_world)
|
||
obs[3:6] = projected_gravity
|
||
|
||
# 3. 速度命令 (已缩放)
|
||
obs[6:9] = commands * SCALE_COMMANDS
|
||
|
||
# 4. 相对关节位置
|
||
# Motor order: FR(3,4,5), FL(0,1,2), RR(9,10,11), RL(6,7,8)
|
||
# Model order: FL(0,1,2), FR(3,4,5), RL(6,7,8), RR(9,10,11)
|
||
motor_pos = data.qpos[7:19] # MuJoCo motor order
|
||
# Map to model order
|
||
model_order_pos = motor_pos[DOF_MAPPING] # (12,)
|
||
dof_pos_rel = model_order_pos - default_dof_pos
|
||
obs[9:21] = dof_pos_rel * SCALE_DOF_POS
|
||
|
||
# 5. 关节速度
|
||
motor_vel = data.qvel[6:18]
|
||
model_order_vel = motor_vel[DOF_MAPPING]
|
||
obs[21:33] = model_order_vel * SCALE_DOF_VEL
|
||
|
||
# 6. 上一步动作
|
||
obs[33:45] = prev_action
|
||
|
||
# 限幅
|
||
obs = np.clip(obs, -CLIP_OBSERVATIONS, CLIP_OBSERVATIONS)
|
||
|
||
return obs
|
||
|
||
|
||
# ============================================================
|
||
# 主程序
|
||
# ============================================================
|
||
|
||
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. 设置摩擦系数
|
||
# friction = [滑动摩擦, 扭转摩擦, 滚动摩擦]
|
||
# 设置地面摩擦系数 (通过 geom 名称查找)
|
||
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，使用默认摩擦系数")
|
||
|
||
# 设置机器人 geom 的摩擦系数 (排除地面)
|
||
for i in range(model.ngeom):
|
||
if i != floor_id:
|
||
model.geom_friction[i] = FRICTION
|
||
print(f"[INFO] 机器人摩擦系数设置为: {FRICTION}")
|
||
|
||
# 3. 加载 RL 模型
|
||
body_module = torch.jit.load(BODY_MODEL_PATH, map_location='cpu')
|
||
adapt_module = torch.jit.load(ADAPT_MODEL_PATH, map_location='cpu')
|
||
body_module.eval()
|
||
adapt_module.eval()
|
||
print(f"[INFO] Body model loaded from: {BODY_MODEL_PATH}")
|
||
print(f"[INFO] Adapt model loaded from: {ADAPT_MODEL_PATH}")
|
||
print(f"[INFO] 观测维度: {NUM_OBS}, 历史步数: {NUM_OBS_HISTORY}")
|
||
|
||
# 3. 初始化
|
||
obs = np.zeros(NUM_OBS, dtype=np.float32)
|
||
prev_action = np.zeros(12, dtype=np.float32)
|
||
obs_buffer = np.zeros(OBS_BUFFER_SIZE, dtype=np.float32) # 10 * 45 = 450
|
||
|
||
# 初始化机器人姿态
|
||
data.qpos[7:19] = DEFAULT_DOF_POS
|
||
data.qvel[:] = 0.0
|
||
data.qpos[2] = 0.35 # 抬高躯干
|
||
mujoco.mj_step(model, data)
|
||
|
||
print("[INFO] 机器人初始化完成，开始 RL 策略推理...")
|
||
|
||
# 4. 初始化键盘控制
|
||
keyboard_reader = KeyboardReader()
|
||
keyboard_reader.init()
|
||
print("[INFO] 键盘控制已启用!")
|
||
print(" W/S: 前进/后退")
|
||
print(" A/D: 左转/右转")
|
||
print(" Q/E: 侧向左移/右移")
|
||
print(" 空格: 停止")
|
||
print(" ESC: 退出")
|
||
|
||
# 5. 启动交互式查看器
|
||
import mujoco.viewer as mv
|
||
view = mv.launch_passive(model, data)
|
||
print("[INFO] 交互式查看器已启动!")
|
||
|
||
# 初始化观测 buffer (填充历史)
|
||
for _ in range(NUM_OBS_HISTORY):
|
||
obs = compute_observations(data, prev_action, np.array([0.0, 0.0, 0.0]), DEFAULT_DOF_POS)
|
||
obs_buffer = np.concatenate([obs_buffer[NUM_OBS:], obs])
|
||
|
||
step_count = 0
|
||
last_inference_time = time.time()
|
||
inference_interval = 0.02 # 50Hz
|
||
|
||
command = np.array([0.0, 0.0, 0.0], dtype=np.float32)
|
||
joint_targets = DEFAULT_DOF_POS.copy() # 初始化关节目标
|
||
|
||
while view.is_running() and not g_exit_requested:
|
||
current_time = time.time()
|
||
|
||
# 更新键盘状态
|
||
keyboard_reader.update()
|
||
|
||
# 读取键盘输入 (切换模式: 按下切换状态)
|
||
# 空格: 停止所有
|
||
if keyboard_reader.is_key_pressed(' '):
|
||
command[:] = 0.0
|
||
keyboard_reader.clear()
|
||
|
||
# W/S: 前进/后退 (vx)
|
||
if keyboard_reader.is_key_pressed('w'):
|
||
command[0] = MAX_LIN_VEL
|
||
command[1] = 0.0 # 清除侧向
|
||
command[2] = 0.0 # 清除转向
|
||
keyboard_reader.clear()
|
||
|
||
if keyboard_reader.is_key_pressed('s'):
|
||
command[0] = -MAX_LIN_VEL
|
||
command[1] = 0.0
|
||
command[2] = 0.0
|
||
keyboard_reader.clear()
|
||
|
||
# Q/E: 侧向移动 (vy)
|
||
if keyboard_reader.is_key_pressed('q'):
|
||
command[1] = MAX_LIN_VEL
|
||
command[0] = 0.0
|
||
command[2] = 0.0
|
||
keyboard_reader.clear()
|
||
|
||
if keyboard_reader.is_key_pressed('e'):
|
||
command[1] = -MAX_LIN_VEL
|
||
command[0] = 0.0
|
||
command[2] = 0.0
|
||
keyboard_reader.clear()
|
||
|
||
# A/D: 左转/右转 (yaw_rate)
|
||
if keyboard_reader.is_key_pressed('a'):
|
||
command[2] = MAX_ANG_VEL
|
||
command[0] = 0.0
|
||
command[1] = 0.0
|
||
keyboard_reader.clear()
|
||
|
||
if keyboard_reader.is_key_pressed('d'):
|
||
command[2] = -MAX_LIN_VEL
|
||
command[0] = 0.0
|
||
command[1] = 0.0
|
||
keyboard_reader.clear()
|
||
|
||
if keyboard_reader.is_key_pressed('\x1b'):
|
||
print("[INFO] ESC pressed, exiting...")
|
||
break
|
||
|
||
# 50Hz 推理
|
||
if current_time - last_inference_time >= inference_interval:
|
||
# 计算观测
|
||
obs = compute_observations(data, prev_action, command, DEFAULT_DOF_POS)
|
||
|
||
# 更新观测历史 buffer
|
||
obs_buffer = np.concatenate([obs_buffer[NUM_OBS:], obs])
|
||
|
||
# 模型推理
|
||
with torch.inference_mode():
|
||
obs_buffer_tensor = torch.from_numpy(obs_buffer).float().unsqueeze(0) # (1, 450)
|
||
obs_tensor = torch.from_numpy(obs).float().unsqueeze(0) # (1, 45)
|
||
|
||
# Adapt 模块: 编码历史观测为 latent
|
||
latent = adapt_module(obs_buffer_tensor) # (1, 21)
|
||
|
||
# Body 模块: 结合当前观测和 latent 生成动作
|
||
combined_input = torch.cat([obs_tensor, latent], dim=1) # (1, 66)
|
||
action = body_module(combined_input).numpy().flatten() # (12,)
|
||
|
||
# 动作后处理
|
||
action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS)
|
||
actions_scaled = action * ACTION_SCALE
|
||
|
||
# Hip关节额外缩放
|
||
hip_indices = [0, 3, 6, 9]
|
||
for i in hip_indices:
|
||
actions_scaled[i] *= HIP_SCALE_REDUCTION
|
||
|
||
# 转换为电机角度目标
|
||
joint_targets = actions_scaled[DOF_MAPPING] + DEFAULT_DOF_POS
|
||
|
||
# 保存上一步动作
|
||
prev_action = action.copy()
|
||
last_inference_time = current_time
|
||
|
||
# PD 控制 (每步执行)
|
||
current_pos = data.qpos[7:19]
|
||
current_vel = data.qvel[6:18]
|
||
torque = KP * (joint_targets - current_pos) - KD * current_vel
|
||
|
||
# 应用控制
|
||
data.ctrl[:] = torque
|
||
mujoco.mj_step(model, data)
|
||
view.sync()
|
||
|
||
step_count += 1
|
||
if step_count % 100 == 0:
|
||
trunk_z = data.qpos[2]
|
||
lin_vel = np.linalg.norm(data.qvel[0:3])
|
||
print(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")
|
||
|
||
keyboard_reader.restore()
|
||
view.close()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|