279 lines
12 KiB
Python
279 lines
12 KiB
Python
#!/usr/bin/env python3
|
||
"""DreamWaQ MuJoCo sim2sim — VAE encoder + Actor, 5-frame history buffer.
|
||
|
||
Usage:
|
||
uv run scripts/dreamwaq_sim2sim_mujoco.py # flat
|
||
uv run scripts/dreamwaq_sim2sim_mujoco.py --terrain rough
|
||
uv run scripts/dreamwaq_sim2sim_mujoco.py --onnx path/to/policy.onnx
|
||
|
||
Controls:
|
||
W/S: forward/back Q/E: left/right A/D: rotate
|
||
Space: stop R: reset Esc: quit
|
||
"""
|
||
import numpy as np
|
||
import mujoco
|
||
from mujoco import viewer
|
||
import onnxruntime as ort
|
||
import os, sys, threading, queue, argparse, time, signal
|
||
|
||
g_exit_requested = False
|
||
signal.signal(signal.SIGINT, lambda *a: globals().update(g_exit_requested=True))
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════
|
||
_PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
XML_DIR = os.path.join(_PROJECT, "motrix_envs", "src", "motrix_envs", "locomotion", "go1", "xmls")
|
||
DEFAULT_ONNX = os.path.join(_PROJECT, "exports_go1_dreamwaq", "policy.onnx")
|
||
|
||
# ── DreamWaQ params (matching training: PD 28/0.7, action_scale 0.25, ctrl_dt=0.02) ──
|
||
NUM_OBS = 45
|
||
NUM_ACTIONS = 12
|
||
HISTORY_LEN = 5
|
||
ACTION_SCALE = 0.25
|
||
KP = 28.0
|
||
KD = 0.7
|
||
CLIP_ACTIONS = 4.0
|
||
CLIP_TORQUES = 80.0
|
||
CLIP_OBS = 100.0
|
||
MAX_VX, MAX_VY, MAX_WZ = 1.0, 1.0, 1.0
|
||
|
||
# DreamWaQ default joint angles — MUST match MuJoCo XML joint order:
|
||
# qpos[7:19] = FR_hip,FR_thigh,FR_calf, FL_hip,FL_thigh,FL_calf, RR_hip,RR_thigh,RR_calf, RL_hip,RL_thigh,RL_calf
|
||
# 必须与 MotrixSim 训练的 default_angles 完全一致！
|
||
DEFAULT_ANGLES = np.array([
|
||
0.0, 0.9, -1.8, # FR
|
||
0.0, 0.9, -1.8, # FL
|
||
0.0, 0.9, -1.8, # RR
|
||
0.0, 0.9, -1.8, # RL
|
||
], dtype=np.float32)
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════
|
||
# Keyboard
|
||
# ═══════════════════════════════════════════════════════════════════════
|
||
from pynput import keyboard
|
||
|
||
class KB:
|
||
def __init__(self):
|
||
self._q = queue.Queue(); self.running = True
|
||
self.held = set(); self._t = None; self._l = None
|
||
def _n(self, k):
|
||
try:
|
||
if hasattr(k, 'char') and k.char: return k.char.lower()
|
||
except: pass
|
||
return str(k).lower()
|
||
def _w(self):
|
||
while self.running:
|
||
try:
|
||
et, k = self._q.get(timeout=0.05)
|
||
n = self._n(k)
|
||
if et == 'press': self.held.add(n)
|
||
elif et == 'release': self.held.discard(n)
|
||
except queue.Empty: pass
|
||
def init(self):
|
||
self._l = keyboard.Listener(on_press=lambda k: self._q.put(('press', k)),
|
||
on_release=lambda k: self._q.put(('release', k)))
|
||
self._l.start()
|
||
self._t = threading.Thread(target=self._w, daemon=True); self._t.start()
|
||
def stop(self): self.running = False; self._l.stop()
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════
|
||
# Sensor
|
||
# ═══════════════════════════════════════════════════════════════════════
|
||
def get_sensor(m, d, name):
|
||
sid = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_SENSOR, name)
|
||
if sid < 0: return None
|
||
adr = m.sensor_adr[sid]; dim = m.sensor_dim[sid]
|
||
return d.sensordata[adr:adr+dim].copy()
|
||
|
||
|
||
def compute_obs(model, data, commands, last_action):
|
||
"""DreamWaQ observation (Manaro-Alpha order):
|
||
ang_vel(3) + gravity(3) + commands(3) + joint_pos(12) + joint_vel(12) + actions(12) = 45
|
||
"""
|
||
obs = np.zeros(NUM_OBS, dtype=np.float32)
|
||
# ang_vel [0:3]
|
||
g = get_sensor(model, data, "gyro")
|
||
obs[0:3] = (g if g is not None else data.qvel[3:6]) * 0.25
|
||
# gravity [3:6] (read from MuJoCo model, matching training)
|
||
grav_world = model.opt.gravity.copy()
|
||
grav_world = grav_world / np.linalg.norm(grav_world) # normalize
|
||
R = data.xmat[1].reshape(3, 3)
|
||
obs[3:6] = (R.T @ grav_world).astype(np.float32)
|
||
# commands [6:9]
|
||
obs[6:9] = commands * np.array([2.0, 2.0, 0.25], dtype=np.float32)
|
||
# joint_pos [9:21]
|
||
obs[9:21] = (data.qpos[7:19] - DEFAULT_ANGLES) * 1.0
|
||
# joint_vel [21:33]
|
||
obs[21:33] = data.qvel[6:18] * 0.05
|
||
# last_action [33:45]
|
||
obs[33:45] = last_action
|
||
return np.clip(obs, -CLIP_OBS, CLIP_OBS)
|
||
|
||
|
||
# ═══════════════════════════════════════════════════════════════════════
|
||
# Main
|
||
# ═══════════════════════════════════════════════════════════════════════
|
||
def main():
|
||
p = argparse.ArgumentParser()
|
||
p.add_argument("--onnx", default=DEFAULT_ONNX)
|
||
p.add_argument("--terrain", default="flat", choices=["flat", "rough", "stairs", "dreamwaq", "stairs_test", "stairs_box", "flat_stairs"])
|
||
p.add_argument("--level", type=int, default=0, help="terrain difficulty level 0-9 (0=flat, 9=hardest)")
|
||
args = p.parse_args()
|
||
|
||
# Select XML scene
|
||
terrain_map = {
|
||
"flat": "scene_dreamwaq_flat.xml",
|
||
"rough": "scene_rough_terrain.xml",
|
||
"stairs": "scene_stairs_terrain.xml",
|
||
"dreamwaq": "scene_dreamwaq_terrain.xml",
|
||
"stairs_test": "scene_stairs_test.xml",
|
||
"stairs_box": "scene_stairs_box.xml",
|
||
"flat_stairs": "scene_flat_stairs.xml",
|
||
}
|
||
xml_file = os.path.join(XML_DIR, terrain_map[args.terrain])
|
||
|
||
if not os.path.exists(args.onnx):
|
||
print(f"[ERROR] ONNX not found: {args.onnx}")
|
||
print("Run: uv run scripts/export_dreamwaq_onnx.py (after training completes)")
|
||
sys.exit(1)
|
||
|
||
os.chdir(XML_DIR)
|
||
with open(xml_file) as f:
|
||
model = mujoco.MjModel.from_xml_string(f.read())
|
||
data = mujoco.MjData(model)
|
||
|
||
# Spawn pose. Hfield heights: MuJoCo z = gp[2] + sbase + (hd * ztop).
|
||
# The stairs_test terrain has sbase=0, flat platform z=0; just lift by clearance.
|
||
if args.terrain == "flat_stairs":
|
||
lvl = max(0, min(1, args.level))
|
||
col = np.random.randint(0, 4)
|
||
spawn_y = 4.0 - lvl * 8.0 # level 0 flat at y=+4, level 1 stairs at y=-4
|
||
spawn_x = -12.0 + col * 8.0 # platform center (cell center x)
|
||
elif args.terrain == "stairs_test":
|
||
spawn_x, spawn_y = -7.5, -4.0 # flat approach before first step (1m zone)
|
||
elif args.terrain == "stairs_box":
|
||
spawn_x, spawn_y = -2.0, 0.0 # flat ground before stairs
|
||
elif args.terrain == "dreamwaq":
|
||
lvl = max(0, min(9, args.level))
|
||
col = np.random.randint(0, 4) # NUM_COLS=4
|
||
spawn_y = 36.0 - lvl * 8.0 # level 0 flat at y=+36, level 9 stairs at y=-36
|
||
spawn_x = -12.0 + col * 8.0 + 4.0 # centre of cell
|
||
print(f"[Level {lvl}] type={col} spawn=({spawn_x:.1f}, {spawn_y:.1f})")
|
||
else:
|
||
spawn_x, spawn_y = (0.0, 0.0)
|
||
|
||
# When DISPLAY is a virtual framebuffer (Xvfb), MuJoCo headless rendering is
|
||
# handled transparently; on a real display this opens a normal GUI window.
|
||
# No explicit headless flag needed — MuJoCo glfw detects the display type.
|
||
|
||
def hfield_z(mx, my):
|
||
gid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "floor")
|
||
if gid < 0 or model.geom_type[gid] != mujoco.mjtGeom.mjGEOM_HFIELD:
|
||
return 0.0
|
||
hf = model.geom_dataid[gid]
|
||
nrow, ncol = int(model.hfield_nrow[hf]), int(model.hfield_ncol[hf])
|
||
sx, sy, ztop, sbase = model.hfield_size[hf]
|
||
adr = model.hfield_adr[hf]
|
||
hd = model.hfield_data[adr:adr + nrow * ncol].reshape(nrow, ncol)
|
||
gp = model.geom_pos[gid]
|
||
col = int(np.clip(((mx - gp[0]) / sx * 0.5 + 0.5) * (ncol - 1), 0, ncol - 1))
|
||
row = int(np.clip(((my - gp[1]) / sy * 0.5 + 0.5) * (nrow - 1), 0, nrow - 1))
|
||
return float(gp[2] + sbase + hd[row, col] * ztop)
|
||
|
||
spawn_z = hfield_z(spawn_x, spawn_y) + 0.45 # standing clearance above terrain
|
||
|
||
def reset_state():
|
||
data.qpos[:] = 0
|
||
data.qpos[0:3] = [spawn_x, spawn_y, spawn_z]; data.qpos[3:7] = [1, 0, 0, 0]
|
||
data.qpos[7:19] = DEFAULT_ANGLES; data.qvel[:] = 0
|
||
mujoco.mj_forward(model, data)
|
||
|
||
reset_state()
|
||
|
||
# ONNX (2 inputs: observations + obs_history)
|
||
session = ort.InferenceSession(args.onnx, providers=['CPUExecutionProvider'])
|
||
print(f"[DreamWaQ] {args.onnx}")
|
||
print(f"[Terrain] {args.terrain}")
|
||
print(f"[CTRL] W/S前后 Q/E左右 A/D旋转 Space停 R重置 Esc退出")
|
||
|
||
kb = KB(); kb.init()
|
||
view = viewer.launch_passive(model, data)
|
||
|
||
# Camera tracking: follow the trunk body
|
||
trunk_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "trunk")
|
||
view.cam.lookat = data.body(trunk_id).xpos.copy()
|
||
view.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
|
||
view.cam.trackbodyid = trunk_id
|
||
|
||
step = 0
|
||
vx, vy, wz = 0.0, 0.0, 0.0
|
||
last_action = np.zeros(NUM_ACTIONS, dtype=np.float32)
|
||
action = np.zeros(NUM_ACTIONS, dtype=np.float32)
|
||
history = np.zeros((1, HISTORY_LEN, NUM_OBS), dtype=np.float32)
|
||
decimation = 4 # MuJoCo dt=0.005, policy dt=0.02 (DreamWaQ aligned)
|
||
|
||
loop_t0 = time.time()
|
||
|
||
while view.is_running() and not g_exit_requested:
|
||
keys = kb.held
|
||
if 'escape' in keys: break
|
||
if 'r' in keys:
|
||
reset_state()
|
||
last_action[:] = 0; history[:] = 0
|
||
print("[R] Reset")
|
||
if ' ' in keys: vx = vy = wz = 0.0
|
||
|
||
vx = MAX_VX if 'w' in keys else (-MAX_VX if 's' in keys else 0.0)
|
||
vy = MAX_VY if 'q' in keys else (-MAX_VY if 'e' in keys else 0.0)
|
||
wz = MAX_WZ if 'a' in keys else (-MAX_WZ if 'd' in keys else 0.0)
|
||
|
||
if step % decimation == 0:
|
||
cmd = np.array([vx, vy, wz], dtype=np.float32)
|
||
obs = compute_obs(model, data, cmd, last_action)
|
||
|
||
# Shift history + add new obs
|
||
history = np.concatenate([history[:, 1:, :], obs.reshape(1, 1, -1)], axis=1)
|
||
|
||
# ONNX inference
|
||
outputs = session.run(None, {
|
||
'obs': obs.reshape(1, -1).astype(np.float32),
|
||
'obs_history': history.reshape(1, -1).astype(np.float32),
|
||
})
|
||
action = outputs[0][0]
|
||
action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS)
|
||
last_action = action.copy()
|
||
|
||
# PD control（对齐训练：目标限位 + 力矩裁剪）
|
||
target = DEFAULT_ANGLES + action * ACTION_SCALE
|
||
# 关节目标限位（与训练一致）
|
||
jnt_lo = model.jnt_range[:, 0].copy() if hasattr(model, 'jnt_range') else None
|
||
jnt_hi = model.jnt_range[:, 1].copy() if hasattr(model, 'jnt_range') else None
|
||
# MuJoCo model.actuator_trnid 可能不直接暴露，改用 model.jnt_range
|
||
try:
|
||
trnid = model.actuator_trnid[:, 0] # transmission joint indices
|
||
lo = model.jnt_range[trnid, 0]
|
||
hi = model.jnt_range[trnid, 1]
|
||
except Exception:
|
||
lo = np.full(12, -12.0)
|
||
hi = np.full(12, 12.0)
|
||
target = np.clip(target, lo, hi)
|
||
torques = KP * (target - data.qpos[7:19]) - KD * data.qvel[6:18]
|
||
data.ctrl[:] = np.clip(torques, -CLIP_TORQUES, CLIP_TORQUES)
|
||
mujoco.mj_step(model, data)
|
||
view.sync()
|
||
|
||
# Time sync (policy at 50Hz = 0.02s per step)
|
||
expected = step * 0.02
|
||
elapsed = time.time() - loop_t0
|
||
if elapsed < expected:
|
||
time.sleep(expected - elapsed)
|
||
|
||
step += 1
|
||
|
||
kb.stop(); view.close()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|