594 lines
21 KiB
Python
594 lines
21 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Go1 RL Policy 推理demo
|
||
默认加载的Walk-These-Ways 预训练模型 + MuJoCo 可视化
|
||
模型: body_latest.jit + adaptation_module_latest.jit (GRU-based policy with history)
|
||
|
||
使用方法:
|
||
source /opt/mujoco_project/mujoco_env/bin/activate
|
||
python3 /opt/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 threading
|
||
import time
|
||
import signal
|
||
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)
|
||
|
||
# ============================================================
|
||
# 预训练模型 /opt/walk-these-ways/repo/runs/gait-conditioned-agility/pretrain-v0/train/025417.456545/checkpoints
|
||
# ============================================================
|
||
MODEL_DIR = "/opt/walk-these-ways/repo/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 = "/opt/unitree_mujoco/data/go1/xml/go1.xml"
|
||
MESH_PATH = "/opt/unitree_mujoco/data/go1/meshes"
|
||
TERRAIN_PATH = "/opt/unitree_mujoco/data/go1/xml/terrain"
|
||
|
||
# 地形配置: "plane", "stairs", "slope", "rough"
|
||
TERRAIN = "stairs"
|
||
|
||
# ============================================================
|
||
# 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 = 3.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 quat_rotate_inverse(data, v):
|
||
"""
|
||
四元数逆旋转 (世界坐标系 -> 躯干坐标系)
|
||
"""
|
||
base_rot = data.xmat[1].reshape(3, 3) # body 1 = trunk
|
||
return base_rot.T @ np.array(v, dtype=np.float64)
|
||
|
||
|
||
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 模型
|
||
import re
|
||
|
||
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')
|
||
|
||
# ---- 地形替换: 使用单独的 terrain XML 文件 ----
|
||
if TERRAIN != "none":
|
||
terrain_file = os.path.join(TERRAIN_PATH, f"{TERRAIN}.xml")
|
||
if os.path.exists(terrain_file):
|
||
with open(terrain_file, 'r') as f:
|
||
terrain_xml = f.read()
|
||
# 删除原有的 floor geom
|
||
xml_content = re.sub(r'<geom[^>]*name=[^>]*floor[^>]*/>', '', xml_content)
|
||
xml_content = re.sub(r"<geom[^>]*name=[^>]*floor[^>]*/>", '', xml_content)
|
||
# 在 </worldbody> 之前插入新地形
|
||
xml_content = xml_content.replace('</worldbody>', terrain_xml + '</worldbody>')
|
||
print(f"[INFO] 地形已加载: {terrain_file} (TERRAIN={TERRAIN})")
|
||
else:
|
||
print(f"[WARN] 地形文件不存在: {terrain_file}, 使用默认地面")
|
||
|
||
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] 地形: {TERRAIN} ({os.path.join(TERRAIN_PATH, TERRAIN + '.xml')})")
|
||
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}")
|
||
else:
|
||
print(f"[INFO] 未找到 floor geom，跳过地面摩擦设置（地形可能已自定义摩擦系数）")
|
||
|
||
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)
|
||
|
||
# 初始化历史 (填充零)
|
||
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)
|
||
|
||
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}: trunk_z={trunk_z:.3f}m, vel={lin_vel:.3f}m/s")
|
||
|
||
keyboard_reader.restore()
|
||
view.close()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|