cyy/Motrixlab
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add closed-loop DreamWaQ stair validation

Browse Source
This commit is contained in:
8x54zj-m
2026-07-22 02:36:36 +08:00
parent 9cdf59a99d
commit 660393a9ee
4 changed files with 451 additions and 236 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
motrix_envs/src/motrix_envs/locomotion/go1/xmls/scene_stairs_box.xml
View File

@@ -18,19 +18,25 @@
<geom name="floor" pos="0 0 -0.001" size="0 0 0.001" type="plane" <geom name="floor" pos="0 0 -0.001" size="0 0 0.001" type="plane"
material="motphys-ground" contype="1" conaffinity="0" priority="0" friction="0.6" /> material="motphys-ground" contype="1" conaffinity="0" priority="0" friction="0.6" />
<geom name="step0" type="box" size="0.15 10 0.03" pos="0.0 0 0.03" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/> <geom name="step_up_0" type="box" size="0.15 10 0.03" pos="0.15 0 0.03" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
<geom name="step1" type="box" size="0.15 10 0.03" pos="0.3 0 0.09" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/> <geom name="step_up_1" type="box" size="0.15 10 0.06" pos="0.45 0 0.06" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
<geom name="step2" type="box" size="0.15 10 0.03" pos="0.6 0 0.15" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/> <geom name="step_up_2" type="box" size="0.15 10 0.09" pos="0.75 0 0.09" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
<geom name="step3" type="box" size="0.15 10 0.03" pos="0.8999999999999999 0 0.21" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/> <geom name="step_up_3" type="box" size="0.15 10 0.12" pos="1.05 0 0.12" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
<geom name="step4" type="box" size="0.15 10 0.03" pos="1.2 0 0.27" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/> <geom name="step_up_4" type="box" size="0.15 10 0.15" pos="1.35 0 0.15" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
<geom name="step5" type="box" size="0.15 10 0.03" pos="1.5 0 0.32999999999999996" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/> <geom name="step_up_5" type="box" size="0.15 10 0.18" pos="1.65 0 0.18" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
<geom name="step6" type="box" size="0.15 10 0.03" pos="1.7999999999999998 0 0.39" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/> <geom name="step_up_6" type="box" size="0.15 10 0.21" pos="1.95 0 0.21" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
<geom name="step7" type="box" size="0.15 10 0.03" pos="2.1 0 0.44999999999999996" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/> <geom name="step_up_7" type="box" size="0.15 10 0.24" pos="2.25 0 0.24" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
<geom name="step8" type="box" size="0.15 10 0.03" pos="2.4 0 0.51" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/> <geom name="step_up_8" type="box" size="0.15 10 0.27" pos="2.55 0 0.27" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
<geom name="step9" type="box" size="0.15 10 0.03" pos="2.6999999999999997 0 0.5700000000000001" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/> <geom name="step_up_9" type="box" size="0.15 10 0.3" pos="2.85 0 0.3" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
<geom name="platform" type="box" size="0.5 10 0.03" pos="3.5 0 0.63" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/> <geom name="platform" type="box" size="0.5 10 0.3" pos="3.5 0 0.3" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
<geom name="step_down_0" type="box" size="0.15 10 0.27" pos="4.15 0 0.27" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
<!-- Fill under stairs --> <geom name="step_down_1" type="box" size="0.15 10 0.24" pos="4.45 0 0.24" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
<geom name="fill" type="box" size="1.5 10 0.3" pos="1.5 0 -0.3" rgba="0.5 0.4 0.3 1" friction="0.8 0.3 0.3"/> <geom name="step_down_2" type="box" size="0.15 10 0.21" pos="4.75 0 0.21" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
<geom name="step_down_3" type="box" size="0.15 10 0.18" pos="5.05 0 0.18" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
<geom name="step_down_4" type="box" size="0.15 10 0.15" pos="5.35 0 0.15" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
<geom name="step_down_5" type="box" size="0.15 10 0.12" pos="5.65 0 0.12" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
<geom name="step_down_6" type="box" size="0.15 10 0.09" pos="5.95 0 0.09" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
<geom name="step_down_7" type="box" size="0.15 10 0.06" pos="6.25 0 0.06" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
<geom name="step_down_8" type="box" size="0.15 10 0.03" pos="6.55 0 0.03" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
</worldbody> </worldbody>
</mujoco> </mujoco>

538
scripts/dreamwaq_sim2sim_mujoco.py
View File

@@ -1,30 +1,38 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
"""DreamWaQ MuJoCo sim2sim — VAE encoder + Actor, 5-frame history buffer. """Run and validate a DreamWaQ ONNX policy in MuJoCo.
Usage: By default the robot walks forward through the stairs_box course. The process
uv run scripts/dreamwaq_sim2sim_mujoco.py # flat returns success only after it climbs the stairs, crosses the top platform,
uv run scripts/dreamwaq_sim2sim_mujoco.py --terrain rough descends to flat ground, and remains upright.
uv run scripts/dreamwaq_sim2sim_mujoco.py --onnx path/to/policy.onnx
Controls: Controls:
W/S: forward/back Q/E: left/right A/D: rotate W/S: forward/back Q/E: left/right A/D: rotate
Space: stop R: reset Esc: quit Space: stop R: reset Esc: quit
""" """
import numpy as np
import argparse
import os
import queue
import signal
import sys
import threading
import time
import mujoco import mujoco
from mujoco import viewer import numpy as np
import onnxruntime as ort import onnxruntime as ort
import os, sys, threading, queue, argparse, time, signal from mujoco import viewer
g_exit_requested = False g_exit_requested = False
signal.signal(signal.SIGINT, lambda *a: globals().update(g_exit_requested=True)) signal.signal(signal.SIGINT, lambda *args: globals().update(g_exit_requested=True))
# ═══════════════════════════════════════════════════════════════════════ PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
_PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) XML_DIR = os.path.join(
XML_DIR = os.path.join(_PROJECT, "motrix_envs", "src", "motrix_envs", "locomotion", "go1", "xmls") PROJECT_DIR, "motrix_envs", "src", "motrix_envs", "locomotion", "go1", "xmls"
DEFAULT_ONNX = os.path.join(_PROJECT, "exports_go1_dreamwaq", "policy.onnx") )
DEFAULT_ONNX = os.path.join(PROJECT_DIR, "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_OBS = 45
NUM_ACTIONS = 12 NUM_ACTIONS = 12
HISTORY_LEN = 5 HISTORY_LEN = 5
@@ -36,92 +44,180 @@ CLIP_TORQUES = 80.0
CLIP_OBS = 100.0 CLIP_OBS = 100.0
MAX_VX, MAX_VY, MAX_WZ = 1.0, 1.0, 1.0 MAX_VX, MAX_VY, MAX_WZ = 1.0, 1.0, 1.0
# DreamWaQ default joint angles — MUST match MuJoCo XML joint order: # MuJoCo actuator and qpos order: FR, FL, RR, RL; hip, thigh, calf.
# 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 DEFAULT_ANGLES = np.array(
# 必须与 MotrixSim 训练的 default_angles 完全一致！ [
DEFAULT_ANGLES = np.array([ 0.0, 0.9, -1.8,
0.0, 0.9, -1.8, # FR 0.0, 0.9, -1.8,
0.0, 0.9, -1.8, # FL 0.0, 0.9, -1.8,
0.0, 0.9, -1.8, # RR 0.0, 0.9, -1.8,
0.0, 0.9, -1.8, # RL ],
], dtype=np.float32) dtype=np.float32,
)
# ═══════════════════════════════════════════════════════════════════════ class Keyboard:
# Keyboard
# ═══════════════════════════════════════════════════════════════════════
from pynput import keyboard
class KB:
def __init__(self): def __init__(self):
self._q = queue.Queue(); self.running = True self.events = queue.Queue()
self.held = set(); self._t = None; self._l = None self.running = True
def _n(self, k): self.held = set()
self.thread = None
self.listener = None
@staticmethod
def _name(key):
try: try:
if hasattr(k, 'char') and k.char: return k.char.lower() if hasattr(key, "char") and key.char:
except: pass return key.char.lower()
return str(k).lower() except Exception:
def _w(self): pass
return str(key).lower()
def _worker(self):
while self.running: while self.running:
try: try:
et, k = self._q.get(timeout=0.05) event_type, key = self.events.get(timeout=0.05)
n = self._n(k) name = self._name(key)
if et == 'press': self.held.add(n) if event_type == "press":
elif et == 'release': self.held.discard(n) self.held.add(name)
except queue.Empty: pass else:
def init(self): self.held.discard(name)
self._l = keyboard.Listener(on_press=lambda k: self._q.put(('press', k)), except queue.Empty:
on_release=lambda k: self._q.put(('release', k))) pass
self._l.start()
self._t = threading.Thread(target=self._w, daemon=True); self._t.start() def start(self):
def stop(self): self.running = False; self._l.stop() from pynput import keyboard
self.listener = keyboard.Listener(
on_press=lambda key: self.events.put(("press", key)),
on_release=lambda key: self.events.put(("release", key)),
)
self.listener.start()
self.thread = threading.Thread(target=self._worker, daemon=True)
self.thread.start()
def stop(self):
self.running = False
if self.listener is not None:
self.listener.stop()
# ═══════════════════════════════════════════════════════════════════════ def get_sensor(model, data, name):
# Sensor sensor_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, name)
# ═══════════════════════════════════════════════════════════════════════ if sensor_id < 0:
def get_sensor(m, d, name): return None
sid = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_SENSOR, name) address = model.sensor_adr[sensor_id]
if sid < 0: return None dimension = model.sensor_dim[sensor_id]
adr = m.sensor_adr[sid]; dim = m.sensor_dim[sid] return data.sensordata[address : address + dimension].copy()
return d.sensordata[adr:adr+dim].copy()
def compute_obs(model, data, commands, last_action): def compute_obs(model, data, commands, last_action):
"""DreamWaQ observation (Manaro-Alpha order): """Build the 45-value observation in the same order used for training."""
ang_vel(3) + gravity(3) + commands(3) + joint_pos(12) + joint_vel(12) + actions(12) = 45
"""
obs = np.zeros(NUM_OBS, dtype=np.float32) obs = np.zeros(NUM_OBS, dtype=np.float32)
# ang_vel [0:3] gyro = get_sensor(model, data, "gyro")
g = get_sensor(model, data, "gyro") obs[0:3] = (gyro if gyro is not None else data.qvel[3:6]) * 0.25
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) gravity_world = model.opt.gravity.copy()
grav_world = model.opt.gravity.copy() gravity_world /= np.linalg.norm(gravity_world)
grav_world = grav_world / np.linalg.norm(grav_world) # normalize rotation = data.xmat[1].reshape(3, 3)
R = data.xmat[1].reshape(3, 3) obs[3:6] = (rotation.T @ gravity_world).astype(np.float32)
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) 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
obs[9:21] = (data.qpos[7:19] - DEFAULT_ANGLES) * 1.0
# joint_vel [21:33]
obs[21:33] = data.qvel[6:18] * 0.05 obs[21:33] = data.qvel[6:18] * 0.05
# last_action [33:45]
obs[33:45] = last_action obs[33:45] = last_action
return np.clip(obs, -CLIP_OBS, CLIP_OBS) return np.clip(obs, -CLIP_OBS, CLIP_OBS)
# ═══════════════════════════════════════════════════════════════════════ def body_pose(data, body_id):
# Main rotation = data.xmat[body_id].reshape(3, 3)
# ═══════════════════════════════════════════════════════════════════════ roll = np.arctan2(rotation[2, 1], rotation[2, 2])
def main(): pitch = np.arcsin(np.clip(-rotation[2, 0], -1.0, 1.0))
p = argparse.ArgumentParser() yaw = np.arctan2(rotation[1, 0], rotation[0, 0])
p.add_argument("--onnx", default=DEFAULT_ONNX) return data.xpos[body_id].copy(), np.degrees([roll, pitch, yaw])
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() def geom_x_range(model, prefix):
ranges = []
for geom_id in range(model.ngeom):
name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_GEOM, geom_id) or ""
if name.startswith(prefix):
center = model.geom_pos[geom_id, 0]
half_size = model.geom_size[geom_id, 0]
ranges.append((center - half_size, center + half_size))
if not ranges:
return None
return min(start for start, _ in ranges), max(end for _, end in ranges)
def course_geometry(model):
up_range = geom_x_range(model, "step_up_")
down_range = geom_x_range(model, "step_down_")
platform_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "platform")
if up_range is None or down_range is None or platform_id < 0:
return None
platform_x = model.geom_pos[platform_id, 0]
platform_half = model.geom_size[platform_id, 0]
platform_top = model.geom_pos[platform_id, 2] + model.geom_size[platform_id, 2]
return {
"up_start": up_range[0],
"platform_start": platform_x - platform_half,
"platform_end": platform_x + platform_half,
"down_end": down_range[1],
"platform_top": platform_top,
"pass_x": down_range[1] + 0.6,
}
def stage_for_x(x, course):
if course is None:
return "terrain"
if x < course["up_start"]:
return "approach"
if x < course["platform_start"]:
return "ascending"
if x < course["platform_end"]:
return "platform"
if x < course["pass_x"]:
return "descending"
return "finish"
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--onnx", default=DEFAULT_ONNX)
parser.add_argument(
"--terrain",
default="stairs_box",
choices=[
"flat", "rough", "stairs", "dreamwaq", "stairs_test",
"stairs_box", "flat_stairs",
],
)
parser.add_argument("--level", type=int, default=0)
parser.add_argument(
"--forward-speed", type=float, default=0.5,
help="default autonomous forward command in m/s",
)
parser.add_argument("--timeout", type=float, default=30.0, help="validation timeout in simulation seconds")
parser.add_argument("--log-interval", type=float, default=0.5, help="pose log interval in simulation seconds")
parser.add_argument("--manual", action="store_true", help="start with zero velocity and use keyboard commands")
parser.add_argument("--headless", action="store_true", help="run validation without viewer or real-time delay")
parser.add_argument("--no-validation", action="store_true", help="do not stop with automatic PASS/FAIL")
return parser.parse_args()
def main():
args = parse_args()
if hasattr(sys.stdout, "reconfigure"):
sys.stdout.reconfigure(line_buffering=True)
if not os.path.isfile(args.onnx):
print(f"[ERROR] ONNX not found: {args.onnx}", file=sys.stderr)
return 2
if args.headless and args.no_validation:
print("[ERROR] --headless requires automatic validation", file=sys.stderr)
return 2
# Select XML scene
terrain_map = { terrain_map = {
"flat": "scene_dreamwaq_flat.xml", "flat": "scene_dreamwaq_flat.xml",
"rough": "scene_rough_terrain.xml", "rough": "scene_rough_terrain.xml",
@@ -132,147 +228,213 @@ def main():
"flat_stairs": "scene_flat_stairs.xml", "flat_stairs": "scene_flat_stairs.xml",
} }
xml_file = os.path.join(XML_DIR, terrain_map[args.terrain]) xml_file = os.path.join(XML_DIR, terrain_map[args.terrain])
previous_cwd = os.getcwd()
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) os.chdir(XML_DIR)
with open(xml_file) as f: try:
model = mujoco.MjModel.from_xml_string(f.read()) model = mujoco.MjModel.from_xml_path(xml_file)
finally:
os.chdir(previous_cwd)
data = mujoco.MjData(model) 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": if args.terrain == "flat_stairs":
lvl = max(0, min(1, args.level)) level = max(0, min(1, args.level))
col = np.random.randint(0, 4) spawn_x, spawn_y = -12.0, 4.0 - level * 8.0
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": elif args.terrain == "stairs_test":
spawn_x, spawn_y = -7.5, -4.0 # flat approach before first step (1m zone) spawn_x, spawn_y = -7.5, -4.0
elif args.terrain == "stairs_box": elif args.terrain == "stairs_box":
spawn_x, spawn_y = -2.0, 0.0 # flat ground before stairs spawn_x, spawn_y = -2.0, 0.0
elif args.terrain == "dreamwaq": elif args.terrain == "dreamwaq":
lvl = max(0, min(9, args.level)) level = max(0, min(9, args.level))
col = np.random.randint(0, 4) # NUM_COLS=4 spawn_x = -8.0
spawn_y = 36.0 - lvl * 8.0 # level 0 flat at y=+36, level 9 stairs at y=-36 spawn_y = 36.0 - level * 8.0
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: else:
spawn_x, spawn_y = (0.0, 0.0) spawn_x, spawn_y = 0.0, 0.0
# When DISPLAY is a virtual framebuffer (Xvfb), MuJoCo headless rendering is def hfield_z(x, y):
# handled transparently; on a real display this opens a normal GUI window. floor_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "floor")
# No explicit headless flag needed — MuJoCo glfw detects the display type. if floor_id < 0 or model.geom_type[floor_id] != mujoco.mjtGeom.mjGEOM_HFIELD:
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 return 0.0
hf = model.geom_dataid[gid] hfield_id = model.geom_dataid[floor_id]
nrow, ncol = int(model.hfield_nrow[hf]), int(model.hfield_ncol[hf]) rows = int(model.hfield_nrow[hfield_id])
sx, sy, ztop, sbase = model.hfield_size[hf] columns = int(model.hfield_ncol[hfield_id])
adr = model.hfield_adr[hf] size_x, size_y, height, base = model.hfield_size[hfield_id]
hd = model.hfield_data[adr:adr + nrow * ncol].reshape(nrow, ncol) address = model.hfield_adr[hfield_id]
gp = model.geom_pos[gid] samples = model.hfield_data[address : address + rows * columns].reshape(rows, columns)
col = int(np.clip(((mx - gp[0]) / sx * 0.5 + 0.5) * (ncol - 1), 0, ncol - 1)) geom_position = model.geom_pos[floor_id]
row = int(np.clip(((my - gp[1]) / sy * 0.5 + 0.5) * (nrow - 1), 0, nrow - 1)) column = int(np.clip(((x - geom_position[0]) / size_x * 0.5 + 0.5) * (columns - 1), 0, columns - 1))
return float(gp[2] + sbase + hd[row, col] * ztop) row = int(np.clip(((y - geom_position[1]) / size_y * 0.5 + 0.5) * (rows - 1), 0, rows - 1))
return float(geom_position[2] + base + samples[row, column] * height)
spawn_z = hfield_z(spawn_x, spawn_y) + 0.45 # standing clearance above terrain spawn_z = hfield_z(spawn_x, spawn_y) + 0.45
def reset_state(): def reset_state():
data.qpos[:] = 0 mujoco.mj_resetData(model, data)
data.qpos[0:3] = [spawn_x, spawn_y, spawn_z]; data.qpos[3:7] = [1, 0, 0, 0] data.qpos[0:3] = [spawn_x, spawn_y, spawn_z]
data.qpos[7:19] = DEFAULT_ANGLES; data.qvel[:] = 0 data.qpos[3:7] = [1.0, 0.0, 0.0, 0.0]
data.qpos[7:19] = DEFAULT_ANGLES
mujoco.mj_forward(model, data) mujoco.mj_forward(model, data)
reset_state() reset_state()
session = ort.InferenceSession(args.onnx, providers=["CPUExecutionProvider"])
course = course_geometry(model) if args.terrain == "stairs_box" else None
validation_enabled = not args.no_validation
if validation_enabled and course is None:
print("[ERROR] automatic validation requires --terrain stairs_box", file=sys.stderr)
return 2
# ONNX (2 inputs: observations + obs_history)
session = ort.InferenceSession(args.onnx, providers=['CPUExecutionProvider'])
print(f"[DreamWaQ] {args.onnx}") print(f"[DreamWaQ] {args.onnx}")
print(f"[Terrain] {args.terrain}") print(f"[Terrain] {args.terrain} spawn=({spawn_x:.2f}, {spawn_y:.2f}, {spawn_z:.2f})")
print(f"[CTRL] W/S前后 Q/E左右 A/D旋转 Space停 R重置 Esc退出") print(f"[Command] vx={0.0 if args.manual else args.forward_speed:.2f} m/s")
if course is not None:
print(
"[Course] up={up_start:.2f}->{platform_start:.2f}m "
"platform={platform_start:.2f}->{platform_end:.2f}m "
"down={platform_end:.2f}->{down_end:.2f}m pass_x={pass_x:.2f}m".format(**course)
)
if not args.headless:
print("[CTRL] W/S forward/back, Q/E lateral, A/D yaw, Space stop, R reset, Esc quit")
kb = KB(); kb.init() keyboard = None
view = viewer.launch_passive(model, data) view = None
if not args.headless:
keyboard = Keyboard()
keyboard.start()
view = viewer.launch_passive(model, data)
# Camera tracking: follow the trunk body
trunk_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "trunk") trunk_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "trunk")
view.cam.lookat = data.body(trunk_id).xpos.copy() if view is not None:
view.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING view.cam.lookat = data.body(trunk_id).xpos.copy()
view.cam.trackbodyid = trunk_id 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) last_action = np.zeros(NUM_ACTIONS, dtype=np.float32)
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) history = np.zeros((1, HISTORY_LEN, NUM_OBS), dtype=np.float32)
decimation = 4 # MuJoCo dt=0.005, policy dt=0.02 (DreamWaQ aligned) decimation = 4
physics_step = 0
reached_top = False
fallen_since = None
next_log_time = 0.0
status = 0
wall_start = time.monotonic()
loop_t0 = time.time() try:
while not g_exit_requested and (view is None or view.is_running()):
keys = keyboard.held if keyboard is not None else set()
if "escape" in keys:
break
if "r" in keys:
reset_state()
last_action.fill(0.0)
history.fill(0.0)
reached_top = False
fallen_since = None
next_log_time = 0.0
physics_step = 0
wall_start = time.monotonic()
print("[Reset]")
while view.is_running() and not g_exit_requested: forward = 0.0 if args.manual else args.forward_speed
keys = kb.held vx = MAX_VX if "w" in keys else (-MAX_VX if "s" in keys else forward)
if 'escape' in keys: break vy = MAX_VY if "q" in keys else (-MAX_VY if "e" in keys else 0.0)
if 'r' in keys: wz = MAX_WZ if "a" in keys else (-MAX_WZ if "d" in keys else 0.0)
reset_state() if " " in keys:
last_action[:] = 0; history[:] = 0 vx = vy = wz = 0.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) if physics_step % decimation == 0:
vy = MAX_VY if 'q' in keys else (-MAX_VY if 'e' in keys else 0.0) command = np.array([vx, vy, wz], dtype=np.float32)
wz = MAX_WZ if 'a' in keys else (-MAX_WZ if 'd' in keys else 0.0) obs = compute_obs(model, data, command, last_action)
history[:, :-1] = history[:, 1:]
history[:, -1] = obs
action = session.run(
None,
{
"obs": obs.reshape(1, -1),
"obs_history": history.reshape(1, -1),
},
)[0][0]
if not np.all(np.isfinite(action)):
print("[FAIL] policy produced a non-finite action", file=sys.stderr)
status = 1
break
action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS)
last_action = action.copy()
if step % decimation == 0: target = DEFAULT_ANGLES + action * ACTION_SCALE
cmd = np.array([vx, vy, wz], dtype=np.float32) actuator_joints = model.actuator_trnid[:, 0]
obs = compute_obs(model, data, cmd, last_action) target = np.clip(
target,
model.jnt_range[actuator_joints, 0],
model.jnt_range[actuator_joints, 1],
)
torque = KP * (target - data.qpos[7:19]) - KD * data.qvel[6:18]
data.ctrl[:] = np.clip(torque, -CLIP_TORQUES, CLIP_TORQUES)
mujoco.mj_step(model, data)
physics_step += 1
# Shift history + add new obs position, rpy = body_pose(data, trunk_id)
history = np.concatenate([history[:, 1:, :], obs.reshape(1, 1, -1)], axis=1) stage = stage_for_x(position[0], course)
if course is not None and position[2] >= spawn_z + 0.35 and position[0] >= course["platform_start"] - 0.5:
reached_top = True
# ONNX inference if data.time >= next_log_time:
outputs = session.run(None, { velocity = data.qvel[0:3]
'obs': obs.reshape(1, -1).astype(np.float32), print(
'obs_history': history.reshape(1, -1).astype(np.float32), f"[Pose] t={data.time:5.1f}s stage={stage:10s} "
}) f"pos=({position[0]:6.2f},{position[1]:6.2f},{position[2]:5.2f}) "
action = outputs[0][0] f"rpy=({rpy[0]:6.1f},{rpy[1]:6.1f},{rpy[2]:6.1f})deg "
action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS) f"vel=({velocity[0]:5.2f},{velocity[1]:5.2f},{velocity[2]:5.2f})"
last_action = action.copy() )
next_log_time += args.log_interval
# PD control（对齐训练：目标限位 + 力矩裁剪） tilt = max(abs(rpy[0]), abs(rpy[1]))
target = DEFAULT_ANGLES + action * ACTION_SCALE fallen = position[2] < 0.18 or tilt > 65.0
# 关节目标限位（与训练一致） if data.time > 1.0 and fallen:
jnt_lo = model.jnt_range[:, 0].copy() if hasattr(model, 'jnt_range') else None fallen_since = data.time if fallen_since is None else fallen_since
jnt_hi = model.jnt_range[:, 1].copy() if hasattr(model, 'jnt_range') else None else:
# MuJoCo model.actuator_trnid 可能不直接暴露，改用 model.jnt_range fallen_since = None
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) if validation_enabled:
expected = step * 0.02 upright_at_finish = position[2] < spawn_z + 0.25 and tilt < 35.0
elapsed = time.time() - loop_t0 if reached_top and position[0] >= course["pass_x"] and upright_at_finish:
if elapsed < expected: print(
time.sleep(expected - elapsed) f"[PASS] climbed and descended stairs in {data.time:.2f}s; "
f"final_pos=({position[0]:.2f}, {position[1]:.2f}, {position[2]:.2f}), "
f"final_rpy=({rpy[0]:.1f}, {rpy[1]:.1f}, {rpy[2]:.1f})deg"
)
break
if fallen_since is not None and data.time - fallen_since >= 0.5:
print(
f"[FAIL] robot fell at stage={stage}, t={data.time:.2f}s, "
f"pos=({position[0]:.2f}, {position[1]:.2f}, {position[2]:.2f}), "
f"rpy=({rpy[0]:.1f}, {rpy[1]:.1f}, {rpy[2]:.1f})deg",
file=sys.stderr,
)
status = 1
break
if data.time >= args.timeout:
print(
f"[FAIL] timeout after {data.time:.1f}s at stage={stage}; "
f"reached_top={reached_top}, x={position[0]:.2f}m",
file=sys.stderr,
)
status = 1
break
step += 1 if view is not None:
view.sync()
deadline = wall_start + data.time
remaining = deadline - time.monotonic()
if remaining > 0:
time.sleep(remaining)
finally:
if keyboard is not None:
keyboard.stop()
if view is not None:
view.close()
kb.stop(); view.close() return status
if __name__ == "__main__": if __name__ == "__main__":
main() sys.exit(main())

59
scripts/gen_stairs_box.py
View File

@@ -1,7 +1,8 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
"""Generate box-geom stairs XML for MuJoCo sim2sim. """Generate symmetric box-geom stairs XML for MuJoCo sim2sim.
Each step is a separate box with vertical rises — much steeper than hfield. The course is a flat approach, ascending stairs, a top platform, and matching
descending stairs. Each tread is a solid box rising from the ground.
Usage: Usage:
uv run scripts/gen_stairs_box.py # default: 10 steps × 6cm = 60cm uv run scripts/gen_stairs_box.py # default: 10 steps × 6cm = 60cm
@@ -31,15 +32,12 @@ TPL = '''<mujoco model="go1 box stairs scene">
material="motphys-ground" contype="1" conaffinity="0" priority="0" friction="0.6" /> material="motphys-ground" contype="1" conaffinity="0" priority="0" friction="0.6" />
{steps} {steps}
<!-- Fill under stairs -->
<geom name="fill" type="box" size="{fill_sx} 10 {fill_sz}" pos="{fill_x} 0 {fill_z}" rgba="0.5 0.4 0.3 1" friction="0.8 0.3 0.3"/>
</worldbody> </worldbody>
</mujoco> </mujoco>
''' '''
STEP_TPL = ' <geom name="step{n}" type="box" size="{sx} 10 {sz}" pos="{x} 0 {z}" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>\n' STEP_TPL = ' <geom name="{name}" type="box" size="{sx:.6g} 10 {sz:.6g}" pos="{x:.6g} 0 {z:.6g}" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>\n'
PLAT_TPL = ' <geom name="platform" type="box" size="{sx} 10 {sz}" pos="{x} 0 {z}" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>\n' PLAT_TPL = ' <geom name="platform" type="box" size="{sx:.6g} 10 {sz:.6g}" pos="{x:.6g} 0 {z:.6g}" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>\n'
def main(): def main():
@@ -47,36 +45,41 @@ def main():
p.add_argument("--step-height", type=float, default=0.06, help="rise per step [m]") p.add_argument("--step-height", type=float, default=0.06, help="rise per step [m]")
p.add_argument("--step-depth", type=float, default=0.30, help="tread depth per step [m]") p.add_argument("--step-depth", type=float, default=0.30, help="tread depth per step [m]")
p.add_argument("--num-steps", type=int, default=10, help="number of steps") p.add_argument("--num-steps", type=int, default=10, help="number of steps")
p.add_argument("--box-thickness", type=float, default=0.03, help="box half-height [m]") p.add_argument("--platform-depth", type=float, default=1.0, help="top platform depth [m]")
args = p.parse_args() args = p.parse_args()
h = args.step_height h = args.step_height
d = args.step_depth d = args.step_depth
n = args.num_steps n = args.num_steps
sz = args.box_thickness # half-height of each box platform_depth = args.platform_depth
steps_xml = "" steps_xml = ""
for i in range(n): for i in range(n):
x = i * d top = (i + 1) * h
z = i * h + sz # center of box = step top surface - sz steps_xml += STEP_TPL.format(
steps_xml += STEP_TPL.format(n=i, sx=d/2, sz=sz, x=x, z=z) name=f"step_up_{i}", sx=d / 2, sz=top / 2,
x=(i + 0.5) * d, z=top / 2,
)
# Platform at top # Platform at top
plat_x = n * d + 0.5
plat_z = n * h + sz
steps_xml += PLAT_TPL.format(sx=0.5, sz=sz, x=plat_x, z=plat_z)
# Fill box under stairs
total_depth = n * d
total_height = n * h total_height = n * h
fill_sx = total_depth / 2 platform_start = n * d
fill_sz = total_height / 2 platform_end = platform_start + platform_depth
fill_x = total_depth / 2 steps_xml += PLAT_TPL.format(
fill_z = -fill_sz sx=platform_depth / 2, sz=total_height / 2,
x=platform_start + platform_depth / 2, z=total_height / 2,
)
out = TPL.format(steps=steps_xml.rstrip(), for i in range(n):
fill_sx=fill_sx, fill_sz=fill_sz, top = (n - i - 1) * h
fill_x=fill_x, fill_z=fill_z) if top <= 0:
continue
steps_xml += STEP_TPL.format(
name=f"step_down_{i}", sx=d / 2, sz=top / 2,
x=platform_end + (i + 0.5) * d, z=top / 2,
)
out = TPL.format(steps=steps_xml.rstrip())
out_dir = os.path.join(os.path.dirname(__file__), "..", out_dir = os.path.join(os.path.dirname(__file__), "..",
"motrix_envs", "src", "motrix_envs", "locomotion", "motrix_envs", "src", "motrix_envs", "locomotion",
@@ -86,8 +89,10 @@ def main():
f.write(out) f.write(out)
max_h = n * h max_h = n * h
print(f"Generated {n} steps × {h*100:.0f}cm = {max_h*100:.0f}cm total") finish_x = platform_end + n * d
print(f" step depth: {d*100:.0f}cm box thickness: {sz*200:.0f}cm") print(f"Generated {n} steps up/down x {h*100:.0f}cm = {max_h*100:.0f}cm total")
print(f" step depth: {d*100:.0f}cm platform: {platform_depth:.2f}m")
print(f" course: x=0.00m to x={finish_x:.2f}m")
print(f" saved: {out_path}") print(f" saved: {out_path}")

56
scripts/run_dreamwaq_sim2sim.sh
View File

@@ -7,13 +7,18 @@ RUNS_DIR="${PROJECT_DIR}/runs/go1-dreamwaq-walk/rslrl"
checkpoint="" checkpoint=""
output="" output=""
terrain="flat" terrain="stairs_box"
level="0" level="0"
forward_speed="0.5"
timeout="30"
export_only=false export_only=false
manual=false
headless=false
no_validation=false
usage() { usage() {
cat <<'EOF' cat <<'EOF'
Export a DreamWaQ checkpoint to ONNX and run MuJoCo sim2sim. Export a DreamWaQ checkpoint and run the MuJoCo stair-course validation.
Usage: Usage:
scripts/run_dreamwaq_sim2sim.sh [options] scripts/run_dreamwaq_sim2sim.sh [options]
@@ -22,15 +27,21 @@ Options:
-c, --checkpoint PATH Checkpoint to export. Defaults to the newest model_*.pt. -c, --checkpoint PATH Checkpoint to export. Defaults to the newest model_*.pt.
-o, --output PATH ONNX output path. Defaults to policy.onnx beside checkpoint. -o, --output PATH ONNX output path. Defaults to policy.onnx beside checkpoint.
-t, --terrain NAME flat, rough, stairs, dreamwaq, stairs_test, -t, --terrain NAME flat, rough, stairs, dreamwaq, stairs_test,
stairs_box, or flat_stairs. Default: flat. stairs_box, or flat_stairs. Default: stairs_box.
-l, --level N Terrain difficulty level passed to sim2sim. Default: 0. -l, --level N Terrain difficulty level passed to sim2sim. Default: 0.
--speed MPS Autonomous forward command. Default: 0.5.
--timeout SEC Validation timeout in simulation seconds. Default: 30.
--manual Start at zero velocity and use keyboard commands.
--headless Run validation without a viewer or real-time delay.
--no-validation Run until the viewer closes instead of returning PASS/FAIL.
--export-only Export and validate ONNX without starting MuJoCo. --export-only Export and validate ONNX without starting MuJoCo.
-h, --help Show this help. -h, --help Show this help.
Examples: Examples:
scripts/run_dreamwaq_sim2sim.sh scripts/run_dreamwaq_sim2sim.sh
scripts/run_dreamwaq_sim2sim.sh -c runs/.../model_1200.pt scripts/run_dreamwaq_sim2sim.sh -c runs/.../model_1200.pt
scripts/run_dreamwaq_sim2sim.sh -t dreamwaq -l 3 scripts/run_dreamwaq_sim2sim.sh --headless
scripts/run_dreamwaq_sim2sim.sh -c runs/.../model_1200.pt --speed 0.6
EOF EOF
} }
@@ -61,6 +72,28 @@ while (($#)); do
level="$2" level="$2"
shift 2 shift 2
;; ;;
--speed)
(($# >= 2)) || die "$1 requires a number"
forward_speed="$2"
shift 2
;;
--timeout)
(($# >= 2)) || die "$1 requires a number"
timeout="$2"
shift 2
;;
--manual)
manual=true
shift
;;
--headless)
headless=true
shift
;;
--no-validation)
no_validation=true
shift
;;
--export-only) --export-only)
export_only=true export_only=true
shift shift
@@ -80,6 +113,8 @@ case "$terrain" in
*) die "unsupported terrain: ${terrain}" ;; *) die "unsupported terrain: ${terrain}" ;;
esac esac
[[ "$level" =~ ^[0-9]+$ ]] || die "level must be a non-negative integer: ${level}" [[ "$level" =~ ^[0-9]+$ ]] || die "level must be a non-negative integer: ${level}"
[[ "$forward_speed" =~ ^[0-9]+([.][0-9]+)?$ ]] || die "speed must be a non-negative number: ${forward_speed}"
[[ "$timeout" =~ ^[0-9]+([.][0-9]+)?$ ]] || die "timeout must be a non-negative number: ${timeout}"
command -v uv >/dev/null 2>&1 || die "uv is not available in PATH" command -v uv >/dev/null 2>&1 || die "uv is not available in PATH"
if [[ -z "$checkpoint" ]]; then if [[ -z "$checkpoint" ]]; then
@@ -118,7 +153,14 @@ if [[ "$export_only" == true ]]; then
fi fi
echo "[Pipeline] starting MuJoCo: terrain=${terrain}, level=${level}" echo "[Pipeline] starting MuJoCo: terrain=${terrain}, level=${level}"
exec uv run scripts/dreamwaq_sim2sim_mujoco.py \ sim_args=(
--onnx "$output" \ --onnx "$output"
--terrain "$terrain" \ --terrain "$terrain"
--level "$level" --level "$level"
--forward-speed "$forward_speed"
--timeout "$timeout"
)
[[ "$manual" == true ]] && sim_args+=(--manual)
[[ "$headless" == true ]] && sim_args+=(--headless)
[[ "$no_validation" == true ]] && sim_args+=(--no-validation)
exec uv run scripts/dreamwaq_sim2sim_mujoco.py "${sim_args[@]}"