Initial commit
This commit is contained in:
cyy
371
demo/go1_sim2sim_mujoco.py
Normal file
371
demo/go1_sim2sim_mujoco.py
Normal file
@@ -0,0 +1,371 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
MOTRIXLAB_UNTRIEE_GO1_SIM2SIM
|
||||
source /opt/mujoco/venv/bin/activate
|
||||
cd /opt/unitree_mujoco
|
||||
python demo/go1_sim2sim_mujoco.py
|
||||
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import mujoco
|
||||
from mujoco import viewer
|
||||
import os
|
||||
import threading
|
||||
import signal
|
||||
import queue
|
||||
import argparse
|
||||
import time
|
||||
|
||||
g_exit_requested = False
|
||||
|
||||
def signal_handler(signum, frame):
|
||||
global g_exit_requested
|
||||
g_exit_requested = True
|
||||
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
|
||||
# ============================================================
|
||||
# 配置
|
||||
# ============================================================
|
||||
DEFAULT_ONNX_PATH = "/opt/motrixlab/repo/export_flat/policy.onnx"
|
||||
MOTRIX_XML_DIR = "/opt/motrixlab/repo/motrix_envs/src/motrix_envs/locomotion/go1/xmls"
|
||||
XML_PATH = f"{MOTRIX_XML_DIR}/go1_motor_actuator.xml"
|
||||
|
||||
TERRAIN = "none"
|
||||
|
||||
# ============================================================
|
||||
# MotrixLab 参数 (来自 cfg.py)
|
||||
# ============================================================
|
||||
NUM_OBS = 48
|
||||
NUM_ACTIONS = 12
|
||||
OBS_SCALES = {'lin_vel': 2.0, 'ang_vel': 0.25, 'dof_pos': 1.0, 'dof_vel': 0.05}
|
||||
ACTION_SCALE = 0.05
|
||||
KP = 80.0
|
||||
KD = 1.0
|
||||
CLIP_ACTIONS = 23.7
|
||||
CLIP_OBSERVATIONS = 100.0
|
||||
MAX_LIN_VEL_X = 1.0
|
||||
MAX_LIN_VEL_Y = 1.0
|
||||
MAX_ANG_VEL = 0.5
|
||||
|
||||
# ============================================================
|
||||
# 关节名称和顺序
|
||||
# ============================================================
|
||||
POLICY_JOINT_NAMES = [
|
||||
"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_ANGLES = np.array([
|
||||
-0.0, 0.9, -1.8, # FR_hip, FR_thigh, FR_calf
|
||||
0.0, 0.9, -1.8, # FL_hip, FL_thigh, FL_calf
|
||||
-0.0, 0.9, -1.8, # RR_hip, RR_thigh, RR_calf
|
||||
0.0, 0.9, -1.8, # RL_hip, RL_thigh, RL_calf
|
||||
], dtype=np.float32)
|
||||
|
||||
MUJOCO_TO_POLICY = np.arange(12, dtype=np.int64)
|
||||
POLICY_TO_MUJOCO = np.arange(12, dtype=np.int64)
|
||||
|
||||
# ============================================================
|
||||
# 键盘输入
|
||||
# ============================================================
|
||||
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()
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Sensor 读取
|
||||
# ============================================================
|
||||
def get_sensor(model, data, name):
|
||||
sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, name)
|
||||
if sid < 0:
|
||||
return None
|
||||
adr = model.sensor_adr[sid]
|
||||
dim = model.sensor_dim[sid]
|
||||
return data.sensordata[adr:adr + dim].copy()
|
||||
|
||||
|
||||
def compute_observations_motrix(model, data, commands, last_actions):
|
||||
"""计算 48 维观测"""
|
||||
obs = np.zeros(NUM_OBS, dtype=np.float32)
|
||||
|
||||
# 局部线速度 (来自 sensor)
|
||||
local_vel = get_sensor(model, data, "local_linvel")
|
||||
if local_vel is not None:
|
||||
obs[0:3] = local_vel * OBS_SCALES['lin_vel']
|
||||
else:
|
||||
base_rot = data.xmat[1].reshape(3, 3)
|
||||
world_linvel = data.qvel[0:3]
|
||||
local_linvel = base_rot.T @ world_linvel
|
||||
obs[0:3] = local_linvel * OBS_SCALES['lin_vel']
|
||||
|
||||
# 陀螺仪
|
||||
gyro = get_sensor(model, data, "gyro")
|
||||
if gyro is not None:
|
||||
obs[3:6] = gyro * OBS_SCALES['ang_vel']
|
||||
else:
|
||||
obs[3:6] = data.qvel[3:6] * OBS_SCALES['ang_vel']
|
||||
|
||||
# 重力向量 (躯干坐标系
|
||||
base_rot = data.xmat[1].reshape(3, 3)
|
||||
gravity_world = np.array([0., 0., -1.], dtype=np.float64)
|
||||
local_gravity = base_rot.T @ gravity_world
|
||||
obs[6:9] = local_gravity.astype(np.float32)
|
||||
|
||||
# 关节位置偏差
|
||||
joint_pos = data.qpos[7:19]
|
||||
dof_pos_rel = (joint_pos - DEFAULT_JOINT_ANGLES) * OBS_SCALES['dof_pos']
|
||||
obs[9:21] = dof_pos_rel
|
||||
|
||||
# 关节速度
|
||||
joint_vel = data.qvel[6:18]
|
||||
obs[21:33] = joint_vel * OBS_SCALES['dof_vel']
|
||||
|
||||
# 上一步动作
|
||||
obs[33:45] = last_actions
|
||||
|
||||
# 命令
|
||||
obs[45:48] = commands * np.array([OBS_SCALES['lin_vel'], OBS_SCALES['lin_vel'], OBS_SCALES['ang_vel']], dtype=np.float32)
|
||||
|
||||
# 限幅
|
||||
obs = np.clip(obs, -CLIP_OBSERVATIONS, CLIP_OBSERVATIONS)
|
||||
return obs
|
||||
|
||||
|
||||
def main():
|
||||
import re
|
||||
import onnxruntime as ort
|
||||
|
||||
parser = argparse.ArgumentParser(description="MotrixLab Go1 Policy Inference in MuJoCo")
|
||||
parser.add_argument("--onnx", type=str, default=DEFAULT_ONNX_PATH)
|
||||
parser.add_argument("--terrain", type=str, default=TERRAIN, choices=["none", "rough", "stairs"])
|
||||
args = parser.parse_args()
|
||||
|
||||
os.chdir(MOTRIX_XML_DIR)
|
||||
|
||||
if args.terrain == "rough":
|
||||
xml_file = f"{MOTRIX_XML_DIR}/scene_rough_terrain.xml"
|
||||
elif args.terrain == "stairs":
|
||||
xml_file = f"{MOTRIX_XML_DIR}/scene_stairs_terrain.xml"
|
||||
else:
|
||||
xml_file = f"{MOTRIX_XML_DIR}/scene_motor_actuator.xml" # flat floor
|
||||
|
||||
with open(xml_file, 'r') as f:
|
||||
xml_content = f.read()
|
||||
|
||||
model = mujoco.MjModel.from_xml_string(xml_content)
|
||||
data = mujoco.MjData(model)
|
||||
|
||||
print(f"[INFO] Model: {model.nbody} bodies, {model.nq} DoF, {model.nu} actuators")
|
||||
print(f"[INFO] MuJoCo timestep: {model.opt.timestep}")
|
||||
print(f"[INFO] 关节顺序 (qpos[7:19]): {[mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, i) for i in range(1, 13)]}")
|
||||
|
||||
for sensor_name in ["gyro", "local_linvel"]:
|
||||
sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, sensor_name)
|
||||
print(f"[SENSOR] {sensor_name}: {'存在' if sid >= 0 else '不存在'}")
|
||||
|
||||
# 初始化
|
||||
data.qpos[0:3] = np.array([0.0, 0.0, 0.42])
|
||||
data.qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0]) # 默认四元数
|
||||
data.qpos[7:19] = DEFAULT_JOINT_ANGLES
|
||||
data.qvel[:] = 0.0
|
||||
data.ctrl[:] = 0.0
|
||||
mujoco.mj_forward(model, data)
|
||||
|
||||
print(f"[INIT] qpos[2]={data.qpos[2]:.3f}")
|
||||
print(f"[INIT] qpos[7:19]={data.qpos[7:19]}")
|
||||
print(f"[INIT] DEFAULT_JOINT_ANGLES={DEFAULT_JOINT_ANGLES}")
|
||||
|
||||
# 加载onnx
|
||||
session = ort.InferenceSession(args.onnx, providers=['CPUExecutionProvider'])
|
||||
print(f"[INFO] loaded")
|
||||
|
||||
# 主循环
|
||||
ctrl_dt = 0.01 # 100Hz
|
||||
num_steps_per_inference = int(ctrl_dt / model.opt.timestep)
|
||||
print(f"[INFO] 每 {num_steps_per_inference} 步推理一次 ")
|
||||
|
||||
step_count = 0
|
||||
inference_step = 0
|
||||
|
||||
x_vel_cmd = 0.0
|
||||
y_vel_cmd = 0.0
|
||||
yaw_vel_cmd = 0.0
|
||||
commands = np.array([x_vel_cmd, y_vel_cmd, yaw_vel_cmd], dtype=np.float32)
|
||||
last_actions = np.zeros(NUM_ACTIONS, dtype=np.float32)
|
||||
action = np.zeros(NUM_ACTIONS, dtype=np.float32)
|
||||
|
||||
# 键盘
|
||||
keyboard_reader = KeyboardReader()
|
||||
keyboard_reader.init()
|
||||
|
||||
view = viewer.launch_passive(model, data)
|
||||
print("[INFO] 已启动!")
|
||||
|
||||
loop_start_time = time.time()
|
||||
|
||||
while view.is_running() and not g_exit_requested:
|
||||
# 键盘命令
|
||||
if keyboard_reader.is_key_pressed(' '):
|
||||
x_vel_cmd = 0.0
|
||||
y_vel_cmd = 0.0
|
||||
yaw_vel_cmd = 0.0
|
||||
|
||||
if keyboard_reader.is_key_held('w'):
|
||||
x_vel_cmd = MAX_LIN_VEL_X
|
||||
elif keyboard_reader.is_key_held('s'):
|
||||
x_vel_cmd = -MAX_LIN_VEL_X
|
||||
else:
|
||||
x_vel_cmd = 0.0
|
||||
|
||||
if keyboard_reader.is_key_held('q'):
|
||||
y_vel_cmd = MAX_LIN_VEL_Y
|
||||
elif keyboard_reader.is_key_held('e'):
|
||||
y_vel_cmd = -MAX_LIN_VEL_Y
|
||||
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_pressed('r'):
|
||||
# 重置机器人到初始位置
|
||||
data.qpos[0:3] = np.array([0.0, 0.0, 0.42])
|
||||
data.qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0])
|
||||
data.qpos[7:19] = DEFAULT_JOINT_ANGLES
|
||||
data.qvel[:] = 0.0
|
||||
data.ctrl[:] = 0.0
|
||||
last_actions = np.zeros(NUM_ACTIONS, dtype=np.float32)
|
||||
mujoco.mj_forward(model, data)
|
||||
print("[RESET] 机器人已重置")
|
||||
|
||||
if keyboard_reader.is_key_pressed('escape'):
|
||||
break
|
||||
|
||||
# 推理 (每 N 步一次)
|
||||
if inference_step == 0:
|
||||
commands[0] = x_vel_cmd
|
||||
commands[1] = y_vel_cmd
|
||||
commands[2] = yaw_vel_cmd
|
||||
|
||||
obs = compute_observations_motrix(model, data, commands, last_actions)
|
||||
|
||||
# 推理
|
||||
action = session.run(None, {'observations': obs.reshape(1, -1).astype(np.float32)})[0][0]
|
||||
action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS)
|
||||
last_actions = action.copy()
|
||||
|
||||
# PD 控制
|
||||
# joint_targets = action * action_scale + default_angles
|
||||
joint_targets = DEFAULT_JOINT_ANGLES + action * ACTION_SCALE
|
||||
|
||||
current_pos = data.qpos[7:19]
|
||||
current_vel = data.qvel[6:18]
|
||||
torques = KP * (joint_targets - current_pos) - KD * current_vel
|
||||
torques = np.clip(torques, -CLIP_ACTIONS, CLIP_ACTIONS)
|
||||
data.ctrl[:] = torques
|
||||
|
||||
mujoco.mj_step(model, data)
|
||||
view.sync()
|
||||
|
||||
expected_time = step_count * ctrl_dt
|
||||
elapsed = time.time() - loop_start_time
|
||||
sleep_time = expected_time - elapsed
|
||||
if sleep_time > 0:
|
||||
time.sleep(sleep_time)
|
||||
|
||||
step_count += 1
|
||||
inference_step = (inference_step + 1) % num_steps_per_inference
|
||||
|
||||
if step_count % 200 == 0:
|
||||
trunk_z = data.qpos[2]
|
||||
lin_vel = np.linalg.norm(data.qvel[0:3])
|
||||
print(f"\n========== Step {step_count} ==========")
|
||||
print(f"[CMD] x={x_vel_cmd:.2f}, y={y_vel_cmd:.2f}, yaw={yaw_vel_cmd:.2f}")
|
||||
print(f"[OBS] lin_vel={obs[0:3]}, ang_vel={obs[3:6]}, grav={obs[6:9]}")
|
||||
print(f"[ACTION] raw={action[:4]}... scaled={action[:4]*ACTION_SCALE}...")
|
||||
print(f"[TARGET] {joint_targets[:4]}...")
|
||||
print(f"[TORQUE] {torques[:4]}...")
|
||||
print(f"[STATE] z={trunk_z:.3f}m, vel={lin_vel:.3f}m/s")
|
||||
print(f"==========================================\n")
|
||||
|
||||
keyboard_reader.restore()
|
||||
view.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user