walk_ok

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('<default>\n=', '<default>\n')
+
+    # ---- 添加重力箭头 mocap body (橙色小球) ----
+    # mocap body 会自动在viewer中渲染，无需手动管理geom
+    grav_arrow_body = '''
+    <body name="grav_arrow" pos="0 0 0" mocap="true">
+        <geom type="sphere" size="0.04 0 0" rgba="1.0 0.5 0.0 0.8" contype="0" conaffinity="0" friction="0 0 0"/>
+    </body>
+    '''
+    # 在 </worldbody> 之前插入
+    xml_content = xml_content.replace('</worldbody>', grav_arrow_body + '</worldbody>')
+
+    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()