From 80702381659ec54d7c96e6af0939234334dbb271 Mon Sep 17 00:00:00 2001 From: MobKBK <3289288508@QQ.COM> Date: Tue, 14 Apr 2026 04:59:28 +0800 Subject: [PATCH] walk_ok --- go1_walk_these_ways_inference.py | 744 +++++++++++++++++++++++++++++++ 1 file changed, 744 insertions(+) create mode 100644 go1_walk_these_ways_inference.py diff --git a/go1_walk_these_ways_inference.py b/go1_walk_these_ways_inference.py new file mode 100644 index 0000000..b1f97ae --- /dev/null +++ b/go1_walk_these_ways_inference.py @@ -0,0 +1,744 @@ +#!/usr/bin/env python3 +""" +Go1 RL Policy 推理 - Walk-These-Ways 预训练模型 + MuJoCo 可视化 +模型: body_latest.jit + adaptation_module_latest.jit (GRU-based policy with history) +训练环境: IsaacGym (walk-these-ways) + +使用方法: + python3 /home/8x54zj-m/unitree_mujoco/go1_walk_these_ways_inference.py + +键盘控制: + W/S: 前进/后退 + A/D: 左转/右转 + Q/E: 侧向左移/右移 + R/F: 抬腿高度 高/低 + U/J: 身体前倾/后仰 (pitch) + I/K: 身体右倾/左倾 (roll) + 空格: 停止 + ESC: 退出 +""" +import numpy as np +import mujoco +from mujoco import viewer +import os +import sys +import threading +import time +import signal +import collections +import queue + +# ============================================================ +# 全局退出标志 +# ============================================================ +g_exit_requested = False + +def signal_handler(signum, frame): + global g_exit_requested + g_exit_requested = True + +signal.signal(signal.SIGINT, signal_handler) + +# ============================================================ +# 配置 +# ============================================================ +MODEL_DIR = "/home/8x54zj-m/walk-these-ways/runs/gait-conditioned-agility/pretrain-v0/train/025417.456545/checkpoints" +BODY_MODEL_PATH = os.path.join(MODEL_DIR, "body_latest.jit") +ADAPT_MODEL_PATH = os.path.join(MODEL_DIR, "adaptation_module_latest.jit") + +XML_PATH = "/home/8x54zj-m/unitree_mujoco/data/go1/xml/go1.xml" +MESH_PATH = "/home/8x54zj-m/unitree_mujoco/data/go1/meshes" + +# ============================================================ +# Walk-These-Ways 模型参数 (来自 parameters.pkl & deploy配置) +# ============================================================ +NUM_OBS = 70 # 单步观测维度 +NUM_OBS_HISTORY = 30 # 历史观测步数 +OBS_BUFFER_SIZE = NUM_OBS * NUM_OBS_HISTORY # 2100 + +# 观测缩放 (来自 obs_scales) +SCALE_LIN_VEL = 2.0 +SCALE_ANG_VEL = 0.25 +SCALE_DOF_POS = 1.0 +SCALE_DOF_VEL = 0.05 + +# 动作缩放 (来自 control.action_scale) +ACTION_SCALE = 0.25 +HIP_SCALE_REDUCTION = 0.5 # hip关节额外缩放 (来自 control.hip_scale_reduction) + +# 限幅 (来自 normalization) +CLIP_OBSERVATIONS = 100.0 +CLIP_ACTIONS = 10.0 + +# 默认关节角度 (来自 init_state.default_joint_angles, deploy顺序 FL, FR, RL, RR) +# Mujoco 关节顺序是 FR, FL, RR, RL，所以需要映射 +# 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] +# Deploy: [FL_hip, FL_thigh, FL_calf, FR_hip, FR_thigh, FR_calf, RL_hip, RL_thigh, RL_calf, RR_hip, RR_thigh, RR_calf] +DEFAULT_DOF_POS_DEPLOY = np.array([ + 0.1, 0.8, -1.5, # FL: hip, thigh, calf + -0.1, 0.8, -1.5, # FR: hip, thigh, calf + 0.1, 1.0, -1.5, # RL: hip, thigh, calf + -0.1, 1.0, -1.5, # RR: hip, thigh, calf +], dtype=np.float32) + +# Mujoco 默认姿态 (用于初始化Mujoco仿真) +DEFAULT_DOF_POS_MUJOCO = np.array([ + -0.1, 0.8, -1.5, # FR: hip, thigh, calf (Mujoco order) + 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 +], dtype=np.float32) + +# IsaacGym/Deploy -> Mujoco 关节顺序映射 +DEPLOY_TO_MUJOCO_MAPPING = np.array([3, 4, 5, 0, 1, 2, 9, 10, 11, 6, 7, 8]) + +# 速度命令范围 +MAX_LIN_VEL = 1.0 # 最大线速度 m/s +MAX_ANG_VEL = 1.0 # 与 deploy 一致 + +# PD控制参数 +KP = 20.0 +KD = 0.1 + +# MuJoCo 摩擦系数配置 +FRICTION = [0.6, 0.3, 0.3] +FLOOR_FRICTION = [0.0, 0.0, 0.0] # 训练时地面摩擦为0 + +# ============================================================ +# 步态预设 (与 deploy ElegantGaitProfile 一致) +# ============================================================ +GAIT_PRESETS = { + 'trot': { + 'name': 'Trot (小跑)', + 'phase': 0.5, 'offset': 0.0, 'bound': 0.0, 'duration': 0.5, + 'description': '对角腿同步' + }, + 'pace': { + 'name': 'Pace (踱步)', + 'phase': 0.0, 'offset': 0.0, 'bound': 0.5, 'duration': 0.5, + 'description': '同侧腿同步' + }, + 'bound': { + 'name': 'Bound (奔跑)', + 'phase': 0.0, 'offset': 0.5, 'bound': 0.0, 'duration': 0.5, + 'description': '前后腿同步' + }, + 'pronk': { + 'name': 'Pronk (跳跃)', + 'phase': 0.0, 'offset': 0.0, 'bound': 0.0, 'duration': 0.5, + 'description': '四腿同时' + }, +} + +current_gait = 'trot' + + +# ============================================================ +# 键盘输入读取 (独立线程, pynput 事件驱动) +# ============================================================ +from pynput import keyboard + +class KeyboardReader: + def __init__(self): + self._event_queue = queue.Queue() + self.running = True + self.shared_keys_held = set() + self.shared_one_shot = set() + self._reader_thread = None + self._listener = None + + def _normalize_key(self, key): + try: + if hasattr(key, 'char') and key.char is not None: + return key.char.lower() + except: + pass + key_str = str(key) + if key_str == 'Key.esc': + return 'escape' + elif key_str == 'Key.space': + return 'space' + elif key_str.startswith('Key.'): + return key_str.lower() + return key_str.lower() + + def _reader_worker(self): + while self.running: + try: + event_type, key = self._event_queue.get(timeout=0.05) + k = self._normalize_key(key) + if event_type == 'press': + self.shared_keys_held.add(k) + self.shared_one_shot.discard(k) + elif event_type == 'release': + self.shared_keys_held.discard(k) + self.shared_one_shot.discard(k) + except queue.Empty: + pass + + def init(self): + def on_press(key): + self._event_queue.put(('press', key)) + def on_release(key): + self._event_queue.put(('release', key)) + try: + self._listener = keyboard.Listener(on_press=on_press, on_release=on_release) + self._listener.start() + self._reader_thread = threading.Thread(target=self._reader_worker, daemon=True) + self._reader_thread.start() + print("[INFO] 键盘监听已启动 (独立线程)") + except Exception as e: + print(f"[WARN] 无法初始化键盘监听: {e}") + + def is_key_pressed(self, key): + k = self._normalize_key(key) if isinstance(key, str) else self._normalize_key(key) + if k not in self.shared_keys_held or k in self.shared_one_shot: + return False + self.shared_one_shot.add(k) + return True + + def is_key_held(self, key): + k = self._normalize_key(key) if isinstance(key, str) else self._normalize_key(key) + return k in self.shared_keys_held + + def restore(self): + self.running = False + if self._listener: + self._listener.stop() + + +# ============================================================ +# 辅助函数 +# ============================================================ + +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(data, v): + """四元数逆旋转 (世界坐标系 -> 躯干坐标系) + 使用 data.xmat 直接计算，更可靠 + """ + base_rot = data.xmat[1].reshape(3, 3) # body 1 = trunk + return base_rot.T @ np.array(v, dtype=np.float64) + + +def rotation_matrix_from_quat(data): + """从四元数计算旋转矩阵 (躯干坐标系) + 使用 MuJoCo 内置函数 mju_quat2Mat，保证与 MuJoCo 内部一致 + 返回: 3x3 旋转矩阵 R, 使得 R * world_vec = body_vec + """ + quat_raw = data.qpos[3:7] # [x, y, z, w] + quat_mju = np.array([quat_raw[3], quat_raw[0], quat_raw[1], quat_raw[2]], dtype=np.float64) + R_mju = np.zeros(9, dtype=np.float64) + mujoco.mju_quat2Mat(R_mju, quat_mju) + return R_mju.reshape(3, 3) + + +def grav_to_arrow(grav): + """将重力投影向量转为箭头字符串 (躯干坐标系视图) + grav: [gx, gy, gz] 世界重力在躯干坐标系下的投影 + 躯干坐标系: X=前, Y=左, Z=上 + 返回: (world_arrow, body_arrow) + """ + # 世界重力: 始终是 (0, 0, -1) = 纯Z轴负方向 = 向下 ↓ + world_arrow = "↓" + + # 躯干坐标系重力投影 + gx, gy, gz = grav + + # 判断主体方向 (忽略很小分量) + abs_x, abs_y, abs_z = abs(gx), abs(gy), abs(gz) + + # 标准化 + total = abs_x + abs_y + abs_z + nx, ny, nz = gx / total, gy / total, gz / total + + # Z 分量判断上下 + if nz > 0.1: + z_arrow = "↑" + elif nz < -0.1: + z_arrow = "↓" + else: + z_arrow = "·" + + # X 分量判断前后 + if nx > 0.1: + x_arrow = "→" + elif nx < -0.1: + x_arrow = "←" + else: + x_arrow = "·" + + # Y 分量判断左右 (正Y=机器左=世界右→所以箭头反向) + if ny > 0.1: + y_arrow = "←" # 机器左倾 + elif ny < -0.1: + y_arrow = "→" # 机器右倾 + else: + y_arrow = "·" + + body_arrow = f"{x_arrow}{y_arrow}{z_arrow}" + return world_arrow, body_arrow + + +def compute_observations_wtw(data, prev_action, last_action, commands, clock_inputs, default_dof_pos_mujoco, obs_scales): + """ + 计算 Walk-These-Ways 策略的观测向量 (70维) + """ + obs = np.zeros(NUM_OBS, dtype=np.float32) + + # 获取四元数 (MuJoCo 格式: x, y, z, w) -> 转换为 (w, x, y, z) + quat_mujoco = data.qpos[3:7] # (x, y, z, w) + quat = np.concatenate([quat_mujoco[3:4], quat_mujoco[0:3]]) + + # 1. 重力投影到躯干坐标系 (使用单位重力 [0,0,-1]) + projected_gravity = quat_rotate_inverse(data, np.array([0., 0., -1.])) + obs[0:3] = projected_gravity + + # 2. 命令 (已缩放) + obs[3:18] = commands[:15] + + # 3. 相对关节位置 (Mujoco顺序 -> Deploy顺序) + current_joint_pos = data.qpos[7:19] + current_joint_pos = current_joint_pos[DEPLOY_TO_MUJOCO_MAPPING] + default_dof_pos_deploy = default_dof_pos_mujoco[DEPLOY_TO_MUJOCO_MAPPING] + dof_pos_rel = (current_joint_pos - default_dof_pos_deploy) * obs_scales['dof_pos'] + obs[18:30] = dof_pos_rel + + # 4. 关节速度 * scale (Mujoco顺序 -> Deploy顺序) + joint_vel = data.qvel[6:18] + joint_vel = joint_vel[DEPLOY_TO_MUJOCO_MAPPING] + obs[30:42] = joint_vel * obs_scales['dof_vel'] + + # 5. 上一步动作 (clipped) + obs[42:54] = np.clip(prev_action, -CLIP_ACTIONS, CLIP_ACTIONS) + + # 6. 上上步动作 + obs[54:66] = np.clip(last_action, -CLIP_ACTIONS, CLIP_ACTIONS) + + # 7. 时钟输入 + obs[66:70] = clock_inputs + + # 限幅 + obs = np.clip(obs, -CLIP_OBSERVATIONS, CLIP_OBSERVATIONS) + + return obs + + +def compute_commands_scale(): + """计算 commands_scale，与 deploy/lcm_agent.py 一致""" + obs_scales = { + 'lin_vel': 2.0, + 'ang_vel': 0.25, + 'dof_pos': 1.0, + 'dof_vel': 0.05, + 'body_height_cmd': 2.0, + 'footswing_height_cmd': 0.15, + 'body_pitch_cmd': 0.3, + 'body_roll_cmd': 0.3, + 'stance_width_cmd': 1.0, + 'stance_length_cmd': 1.0, + 'aux_reward_cmd': 1.0, + } + + commands_scale = np.array([ + obs_scales['lin_vel'], obs_scales['lin_vel'], + obs_scales['ang_vel'], obs_scales['body_height_cmd'], + 1, 1, 1, 1, 1, # gait params (unscaled) + obs_scales['footswing_height_cmd'], + obs_scales['body_pitch_cmd'], + obs_scales['body_roll_cmd'], + obs_scales['stance_width_cmd'], + obs_scales['stance_length_cmd'], + obs_scales['aux_reward_cmd'], + 1, 1, 1, 1, 1 # padding + ], dtype=np.float32) + return commands_scale[:15], obs_scales + + +# ============================================================ +# 主程序 +# ============================================================ + +def main(): + global current_gait + import torch + + # 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') + + # ---- 添加重力箭头 mocap body (橙色小球) ---- + # mocap body 会自动在viewer中渲染，无需手动管理geom + grav_arrow_body = ''' + + + + ''' + # 在 之前插入 + xml_content = xml_content.replace('', grav_arrow_body + '') + + model = mujoco.MjModel.from_xml_string(xml_content) + data = mujoco.MjData(model) + + # 验证 mocap body 添加成功 + grav_arrow_mocap_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "grav_arrow") + if grav_arrow_mocap_id < 0: + print("[WARN] grav_arrow body not found in model!") + else: + print(f"[INFO] grav_arrow body id: {grav_arrow_mocap_id}, nmocap: {model.nmocap}") + # mocap_pos[0] 对应第一个 mocap body（我们添加的 grav_arrow） + GRAV_MOCAP_IDX = 0 + if model.nmocap < 1: + print("[WARN] No mocap bodies in model! Gravity arrow disabled.") + + print(f"[INFO] MuJoCo Model: {model.nbody} bodies, {model.nq} DoF, {model.nu} actuators") + print(f"[INFO] Mujoco关节顺序: {[mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, i) for i in range(1, 13)]}") + print(f"[INFO] Deploy关节顺序: FL, FR, RL, RR (与Mujoco不同)") + + # 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}") + + for i in range(model.ngeom): + if i != floor_id: + model.geom_friction[i] = FRICTION + print(f"[INFO] 机器人摩擦系数设置为: {FRICTION}") + + # 3. 加载 Walk-These-Ways 模型 + 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}") + print(f"[INFO] Action scale: {ACTION_SCALE}, Hip scale reduction: {HIP_SCALE_REDUCTION}") + + # 4. 计算 commands_scale + commands_scale, obs_scales = compute_commands_scale() + print(f"[INFO] Commands scale: {commands_scale}") + + # 5. 初始化观测历史 buffer + obs_buffer = np.zeros(OBS_BUFFER_SIZE, dtype=np.float32) # 30 * 70 = 2100 + + # 6. 初始化机器人姿态 + data.qpos[7:19] = DEFAULT_DOF_POS_MUJOCO + data.qvel[:] = 0.0 + data.qpos[2] = 0.35 # 抬高躯干 + mujoco.mj_step(model, data) + + # ---- 用 MuJoCo 内置函数验证旋转矩阵 ---- + quat_mj = data.qpos[3:7].copy() # [x,y,z,w] + print(f" [CHECK] MuJoCo qpos[3:7] = {quat_mj}") + # MuJoCo's own rotation matrix from quaternion + R_mj = np.zeros(9) + mujoco.mju_quat2Mat(R_mj, quat_mj) + R_mj = R_mj.reshape(3, 3) + print(f" [CHECK] MuJoCo mju_quat2Mat R (identity if q=[1,0,0,0]≈id): \n{R_mj.round(4)}") + mujoco.mj_step(model, data) + + # 打印初始化后的姿态确认 + quat_mj = data.qpos[3:7] + grav_init = quat_rotate_inverse(data, np.array([0., 0., -1.])) + R_init = rotation_matrix_from_quat(data) + quat_std = np.array([quat_mj[3], quat_mj[0], quat_mj[1], quat_mj[2]]) + print(f" [INIT] quat(MuJoCo)={quat_mj} | quat_as_std={quat_std.round(3)} | grav_world→body={grav_init.round(3)} | trunk_z={data.qpos[2]:.3f}") + + # ---- 3D 可视化: 躯干坐标系 vs 世界坐标系 ---- + try: + import matplotlib.pyplot as plt + from mpl_toolkits.mplot3d import Axes3D + R_init = rotation_matrix_from_quat(data) + fig = plt.figure(figsize=(6, 6)) + ax = fig.add_subplot(111, projection='3d') + # 世界坐标系 (黑色) + ax.quiver(0, 0, 0, 1.2, 0, 0, color='k', linewidth=1.5, arrow_length_ratio=0.1) + ax.quiver(0, 0, 0, 0, 1.2, 0, color='k', linewidth=1.5, arrow_length_ratio=0.1) + ax.quiver(0, 0, 0, 0, 0, 1.2, color='k', linewidth=1.5, arrow_length_ratio=0.1) + ax.text(1.3, 0, 0, "Xw(前)", fontsize=9) + ax.text(0, 1.3, 0, "Yw(左)", fontsize=9) + ax.text(0, 0, 1.3, "Zw(上)", fontsize=9) + # 躯干坐标系 (彩色) + body_x = R_init[:, 0] * 0.8 # R第一列 = body X轴在世界 + body_y = R_init[:, 1] * 0.8 # R第二列 = body Y轴在世界 + body_z = R_init[:, 2] * 0.8 # R第三列 = body Z轴在世界 + ax.quiver(0, 0, 0, *body_x, color='r', linewidth=2, arrow_length_ratio=0.1) + ax.quiver(0, 0, 0, *body_y, color='g', linewidth=2, arrow_length_ratio=0.1) + ax.quiver(0, 0, 0, *body_z, color='b', linewidth=2, arrow_length_ratio=0.1) + ax.text(body_x[0]*1.2, body_x[1]*1.2, body_x[2]*1.2, "Xb(前)", color='r', fontsize=9) + ax.text(body_y[0]*1.2, body_y[1]*1.2, body_y[2]*1.2, "Yb(左)", color='g', fontsize=9) + ax.text(body_z[0]*1.2, body_z[1]*1.2, body_z[2]*1.2, "Zb(上)", color='b', fontsize=9) + # 重力向量 + grav_arrow = grav_init * 0.6 + ax.quiver(0, 0, 0, *grav_arrow, color='orange', linewidth=2.5, arrow_length_ratio=0.1) + ax.text(grav_arrow[0]*1.2, grav_arrow[1]*1.2, grav_arrow[2]*1.2, f"g={grav_init.round(2)}", color='orange', fontsize=9) + ax.set_xlim([-1.5, 1.5]); ax.set_ylim([-1.5, 1.5]); ax.set_zlim([-1.5, 1.5]) + ax.set_xlabel("X (世界)"); ax.set_ylabel("Y (世界)"); ax.set_zlabel("Z (世界)") + ax.set_title("躯干坐标系 (RGB=XYZ轴) vs 世界坐标系 (K) | 橙色=重力投影") + plt.tight_layout() + plt.savefig("/home/8x54zj-m/unitree_mujoco/body_frame_axes.png", dpi=150) + print(f" [INIT] 坐标系可视化已保存: body_frame_axes.png") + print(f" [INIT] 旋转矩阵 R (body←world, 列=body轴在world中):\n{R_init.round(3)}") + plt.close() + except Exception as e: + print(f" [WARN] 可视化失败: {e}") + + # 初始化历史 (填充零) + for _ in range(NUM_OBS_HISTORY): + dummy_obs = np.zeros(NUM_OBS, dtype=np.float32) + obs_buffer = np.concatenate([obs_buffer[NUM_OBS:], dummy_obs]) + + # 7. 初始化键盘控制 + keyboard_reader = KeyboardReader() + keyboard_reader.init() + print("[INFO] 键盘控制已启用!") + print(" W/S: 前进/后退 (max 1.0 m/s)") + print(" A/D: 左转/右转 (max 5.0 rad/s)") + print(" Q/E: 侧向左移/右移") + print(" R/F: 抬腿高度 高/低") + print(" U/J: 身体前倾/后仰") + print(" I/K: 身体右倾/左倾") + print(" 1: Trot (小跑) 2: Pace (踱步)") + print(" 3: Bound (奔跑) 4: Pronk (跳跃)") + print(" 空格: 停止") + print(" ESC: 退出") + print(f"[INFO] 当前步态: {GAIT_PRESETS[current_gait]['name']} - {GAIT_PRESETS[current_gait]['description']}") + + # 8. 启动交互式查看器 + import mujoco.viewer as mv + view = mv.launch_passive(model, data) + print("[INFO] 交互式查看器已启动!") + + # 9. 主循环 + step_count = 0 + last_inference_time = time.time() + inference_interval = 0.02 # 50Hz + + # 默认命令 + x_vel_cmd = 0.0 + y_vel_cmd = 0.0 + yaw_vel_cmd = 0.0 + body_height_cmd = 0.0 + footswing_height_cmd = 0.15 + body_pitch_cmd = 0.0 + body_roll_cmd = 0.0 + gait_frequency_cmd = 3.0 + gait_phase_cmd = GAIT_PRESETS[current_gait]['phase'] + gait_offset_cmd = GAIT_PRESETS[current_gait]['offset'] + gait_bound_cmd = GAIT_PRESETS[current_gait]['bound'] + gait_duration_cmd = GAIT_PRESETS[current_gait]['duration'] + + prev_action = np.zeros(12, dtype=np.float32) + last_action = np.zeros(12, dtype=np.float32) + gait_indices = 0.0 + + clock_inputs = np.zeros(4, dtype=np.float32) + + print("[INFO] 开始 RL 策略推理...") + + while view.is_running() and not g_exit_requested: + current_time = time.time() + + # 空格: 停止 + if keyboard_reader.is_key_pressed(' '): + x_vel_cmd = 0.0 + y_vel_cmd = 0.0 + yaw_vel_cmd = 0.0 + + # 步态切换 + gait_keys = {'1': 'trot', '2': 'pace', '3': 'bound', '4': 'pronk'} + for key, gait_name in gait_keys.items(): + if keyboard_reader.is_key_pressed(key): + if current_gait != gait_name: + current_gait = gait_name + gait_phase_cmd = GAIT_PRESETS[gait_name]['phase'] + gait_offset_cmd = GAIT_PRESETS[gait_name]['offset'] + gait_bound_cmd = GAIT_PRESETS[gait_name]['bound'] + gait_duration_cmd = GAIT_PRESETS[gait_name]['duration'] + print(f"[INFO] 切换步态: {GAIT_PRESETS[gait_name]['name']} - {GAIT_PRESETS[gait_name]['description']}") + + # 连续动作 + if keyboard_reader.is_key_held('w'): + x_vel_cmd = MAX_LIN_VEL + elif keyboard_reader.is_key_held('s'): + x_vel_cmd = -MAX_LIN_VEL + else: + x_vel_cmd = 0.0 + + if keyboard_reader.is_key_held('q'): + y_vel_cmd = MAX_LIN_VEL + elif keyboard_reader.is_key_held('e'): + y_vel_cmd = -MAX_LIN_VEL + else: + y_vel_cmd = 0.0 + + if keyboard_reader.is_key_held('a'): + yaw_vel_cmd = MAX_ANG_VEL + elif keyboard_reader.is_key_held('d'): + yaw_vel_cmd = -MAX_ANG_VEL + else: + yaw_vel_cmd = 0.0 + + # 抬腿高度 + if keyboard_reader.is_key_held('r'): + footswing_height_cmd = min(0.35, footswing_height_cmd + 0.005) + elif keyboard_reader.is_key_held('f'): + footswing_height_cmd = max(0.03, footswing_height_cmd - 0.005) + + # 身体前后倾斜 + if keyboard_reader.is_key_held('u'): + body_pitch_cmd = min(0.4, body_pitch_cmd + 0.01) + elif keyboard_reader.is_key_held('j'): + body_pitch_cmd = max(-0.4, body_pitch_cmd - 0.01) + + # 身体左右倾斜 + if keyboard_reader.is_key_held('i'): + body_roll_cmd = min(0.0, body_roll_cmd + 0.01) + elif keyboard_reader.is_key_held('k'): + body_roll_cmd = max(0.0, body_roll_cmd - 0.01) + + # ESC: 退出 + if keyboard_reader.is_key_pressed('escape'): + print("[INFO] ESC pressed, exiting...") + break + + # 50Hz 推理 + if current_time - last_inference_time >= inference_interval: + # 构建原始命令 (15维) + raw_commands = np.zeros(15, dtype=np.float32) + raw_commands[0] = x_vel_cmd + raw_commands[1] = y_vel_cmd + raw_commands[2] = yaw_vel_cmd + raw_commands[3] = body_height_cmd + raw_commands[4] = gait_frequency_cmd + raw_commands[5] = gait_phase_cmd + raw_commands[6] = gait_offset_cmd + raw_commands[7] = gait_bound_cmd + raw_commands[8] = gait_duration_cmd + raw_commands[9] = footswing_height_cmd + raw_commands[10] = body_pitch_cmd + raw_commands[11] = body_roll_cmd + raw_commands[12] = 0.25 # stance_width + raw_commands[13] = 0.4 # stance_length + + # 应用 commands_scale + commands = raw_commands * commands_scale + + # 更新 gait indices (每 policy step = 4 * sim_dt = 4 * 0.005 = 0.02) + gait_indices += 0.02 * gait_frequency_cmd + if gait_indices > 1.0: + gait_indices -= 1.0 + + # 计算 clock_inputs + phase = gait_phase_cmd + offset = gait_offset_cmd + bound = gait_bound_cmd + foot_indices = [ + gait_indices + phase + offset + bound, # FL + gait_indices + offset, # FR + gait_indices + bound, # RL + gait_indices + phase # RR + ] + clock_inputs[0] = np.sin(2 * np.pi * foot_indices[0]) + clock_inputs[1] = np.sin(2 * np.pi * foot_indices[1]) + clock_inputs[2] = np.sin(2 * np.pi * foot_indices[2]) + clock_inputs[3] = np.sin(2 * np.pi * foot_indices[3]) + + # 计算当前步观测 + obs = compute_observations_wtw( + data, prev_action, last_action, commands, clock_inputs, DEFAULT_DOF_POS_MUJOCO, obs_scales + ) + + # 更新观测历史 buffer + obs_buffer = np.concatenate([obs_buffer[NUM_OBS:], obs]) + + # 准备模型输入 + obs_buffer_tensor = torch.from_numpy(obs_buffer).float().unsqueeze(0) + obs_tensor = torch.from_numpy(obs).float().unsqueeze(0) + + # 推理 + with torch.inference_mode(): + latent = adapt_module(obs_buffer_tensor) + combined_input = torch.cat([obs_buffer_tensor, latent], dim=1) + action = body_module(combined_input).numpy().flatten() + + # 动作后处理 + action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS) + + # 保存动作历史 + last_action = prev_action.copy() + prev_action = action.copy() + + last_inference_time = current_time + + # ---- PD 控制 ---- + action_scaled = prev_action * ACTION_SCALE + + # hip关节额外缩放 + hip_indices = [0, 3, 6, 9] + for i in hip_indices: + action_scaled[i] *= HIP_SCALE_REDUCTION + + joint_targets_mujoco = action_scaled[DEPLOY_TO_MUJOCO_MAPPING] + DEFAULT_DOF_POS_MUJOCO + + current_pos = data.qpos[7:19] + current_vel = data.qvel[6:18] + torque = KP * (joint_targets_mujoco - current_pos) - KD * current_vel + + data.ctrl[:] = torque + + mujoco.mj_step(model, data) + + # grav_proj: 世界重力 [0,0,-1] 变换到躯干坐标系 + grav_proj = quat_rotate_inverse(data, np.array([0., 0., -1.])) + + # ---- 重力箭头可视化: 橙色小球始终指向世界下方 ---- + torso_pos = data.xpos[1] # body 1 = trunk + arrow_pos = torso_pos + np.array([0., 0., -1.]) * 0.15 + + # ---- 画橙色小球: 用 mocap body 更新位置 ---- + if grav_arrow_mocap_id >= 0: + data.mocap_pos[GRAV_MOCAP_IDX] = arrow_pos.astype(np.float64) + + view.sync() + + step_count += 1 + if step_count % 100 == 0: + trunk_z = data.qpos[2] + lin_vel = np.linalg.norm(data.qvel[0:3]) + quat_mj = data.qpos[3:7] + quat = np.concatenate([quat_mj[3:4], quat_mj[0:3]]) + grav_proj = quat_rotate_inverse(data, np.array([0., 0., -1.])) + quat_raw = data.qpos[3:7] + quat_mju = np.array([quat_raw[3], quat_raw[0], quat_raw[1], quat_raw[2]], dtype=np.float64) + R_mju = np.zeros(9, dtype=np.float64) + mujoco.mju_quat2Mat(R_mju, quat_mju) + R = R_mju.reshape(3, 3) + world_a, body_a = grav_to_arrow(grav_proj) + gait_short = {'trot': 'TROT', 'pace': 'PACE', 'bound': 'BOUND', 'pronk': 'PRONK'} + # R 的三列 = body X/Y/Z 在世界坐标系中的方向 + print(f" Step {step_count}: quat={quat_raw.round(3)} grav={grav_proj.round(2)}{body_a} | " + f"R=[{R[0,:].round(2)}, {R[1,:].round(2)}, {R[2,:].round(2)}] | " + f"trunk_z={trunk_z:.3f}m, vel={lin_vel:.3f}m/s") + + keyboard_reader.restore() + view.close() + + +if __name__ == "__main__": + main()