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

538
scripts/dreamwaq_sim2sim_mujoco.py
View File

@@ -1,30 +1,38 @@
#!/usr/bin/env python3
"""DreamWaQ MuJoCo sim2sim — VAE encoder + Actor, 5-frame history buffer.
"""Run and validate a DreamWaQ ONNX policy in MuJoCo.
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
By default the robot walks forward through the stairs_box course. The process
returns success only after it climbs the stairs, crosses the top platform,
descends to flat ground, and remains upright.
Controls:
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
from mujoco import viewer
import numpy as np
import onnxruntime as ort
import os, sys, threading, queue, argparse, time, signal
from mujoco import viewer
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 = 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")
PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
XML_DIR = os.path.join(
PROJECT_DIR, "motrix_envs", "src", "motrix_envs", "locomotion", "go1", "xmls"
)
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_ACTIONS = 12
HISTORY_LEN = 5
@@ -36,92 +44,180 @@ 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)
# MuJoCo actuator and qpos order: FR, FL, RR, RL; hip, thigh, calf.
DEFAULT_ANGLES = np.array(
[
0.0, 0.9, -1.8,
0.0, 0.9, -1.8,
0.0, 0.9, -1.8,
0.0, 0.9, -1.8,
],
dtype=np.float32,
)
# ═══════════════════════════════════════════════════════════════════════
# Keyboard
# ═══════════════════════════════════════════════════════════════════════
from pynput import keyboard
class KB:
class Keyboard:
def __init__(self):
self._q = queue.Queue(); self.running = True
self.held = set(); self._t = None; self._l = None
def _n(self, k):
self.events = queue.Queue()
self.running = True
self.held = set()
self.thread = None
self.listener = None
@staticmethod
def _name(key):
try:
if hasattr(k, 'char') and k.char: return k.char.lower()
except: pass
return str(k).lower()
def _w(self):
if hasattr(key, "char") and key.char:
return key.char.lower()
except Exception:
pass
return str(key).lower()
def _worker(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()
event_type, key = self.events.get(timeout=0.05)
name = self._name(key)
if event_type == "press":
self.held.add(name)
else:
self.held.discard(name)
except queue.Empty:
pass
def start(self):
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()
# ═══════════════════════════════════════════════════════════════════════
# 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 get_sensor(model, data, name):
sensor_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, name)
if sensor_id < 0:
return None
address = model.sensor_adr[sensor_id]
dimension = model.sensor_dim[sensor_id]
return data.sensordata[address : address + dimension].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
"""
"""Build the 45-value observation in the same order used for training."""
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]
gyro = get_sensor(model, data, "gyro")
obs[0:3] = (gyro if gyro is not None else data.qvel[3:6]) * 0.25
gravity_world = model.opt.gravity.copy()
gravity_world /= np.linalg.norm(gravity_world)
rotation = data.xmat[1].reshape(3, 3)
obs[3:6] = (rotation.T @ gravity_world).astype(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) * 1.0
# joint_vel [21:33]
obs[9:21] = data.qpos[7:19] - DEFAULT_ANGLES
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()
def body_pose(data, body_id):
rotation = data.xmat[body_id].reshape(3, 3)
roll = np.arctan2(rotation[2, 1], rotation[2, 2])
pitch = np.arcsin(np.clip(-rotation[2, 0], -1.0, 1.0))
yaw = np.arctan2(rotation[1, 0], rotation[0, 0])
return data.xpos[body_id].copy(), np.degrees([roll, pitch, yaw])
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 = {
"flat": "scene_dreamwaq_flat.xml",
"rough": "scene_rough_terrain.xml",
@@ -132,147 +228,213 @@ def main():
"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)
previous_cwd = os.getcwd()
os.chdir(XML_DIR)
with open(xml_file) as f:
model = mujoco.MjModel.from_xml_string(f.read())
try:
model = mujoco.MjModel.from_xml_path(xml_file)
finally:
os.chdir(previous_cwd)
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)
level = max(0, min(1, args.level))
spawn_x, spawn_y = -12.0, 4.0 - level * 8.0
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":
spawn_x, spawn_y = -2.0, 0.0 # flat ground before stairs
spawn_x, spawn_y = -2.0, 0.0
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})")
level = max(0, min(9, args.level))
spawn_x = -8.0
spawn_y = 36.0 - level * 8.0
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
# 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:
def hfield_z(x, y):
floor_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "floor")
if floor_id < 0 or model.geom_type[floor_id] != 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)
hfield_id = model.geom_dataid[floor_id]
rows = int(model.hfield_nrow[hfield_id])
columns = int(model.hfield_ncol[hfield_id])
size_x, size_y, height, base = model.hfield_size[hfield_id]
address = model.hfield_adr[hfield_id]
samples = model.hfield_data[address : address + rows * columns].reshape(rows, columns)
geom_position = model.geom_pos[floor_id]
column = int(np.clip(((x - geom_position[0]) / size_x * 0.5 + 0.5) * (columns - 1), 0, columns - 1))
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():
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_resetData(model, data)
data.qpos[0:3] = [spawn_x, spawn_y, spawn_z]
data.qpos[3:7] = [1.0, 0.0, 0.0, 0.0]
data.qpos[7:19] = DEFAULT_ANGLES
mujoco.mj_forward(model, data)
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"[Terrain] {args.terrain}")
print(f"[CTRL] W/S前后 Q/E左右 A/D旋转 Space停 R重置 Esc退出")
print(f"[Terrain] {args.terrain} spawn=({spawn_x:.2f}, {spawn_y:.2f}, {spawn_z:.2f})")
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()
view = viewer.launch_passive(model, data)
keyboard = None
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")
view.cam.lookat = data.body(trunk_id).xpos.copy()
view.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
view.cam.trackbodyid = trunk_id
if view is not None:
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)
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:
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
forward = 0.0 if args.manual else args.forward_speed
vx = MAX_VX if "w" in keys else (-MAX_VX if "s" in keys else forward)
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 " " 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 physics_step % decimation == 0:
command = np.array([vx, vy, wz], dtype=np.float32)
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:
cmd = np.array([vx, vy, wz], dtype=np.float32)
obs = compute_obs(model, data, cmd, last_action)
target = DEFAULT_ANGLES + action * ACTION_SCALE
actuator_joints = model.actuator_trnid[:, 0]
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
history = np.concatenate([history[:, 1:, :], obs.reshape(1, 1, -1)], axis=1)
position, rpy = body_pose(data, trunk_id)
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
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()
if data.time >= next_log_time:
velocity = data.qvel[0:3]
print(
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}) "
f"rpy=({rpy[0]:6.1f},{rpy[1]:6.1f},{rpy[2]:6.1f})deg "
f"vel=({velocity[0]:5.2f},{velocity[1]:5.2f},{velocity[2]:5.2f})"
)
next_log_time += args.log_interval
# 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()
tilt = max(abs(rpy[0]), abs(rpy[1]))
fallen = position[2] < 0.18 or tilt > 65.0
if data.time > 1.0 and fallen:
fallen_since = data.time if fallen_since is None else fallen_since
else:
fallen_since = None
# 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)
if validation_enabled:
upright_at_finish = position[2] < spawn_z + 0.25 and tilt < 35.0
if reached_top and position[0] >= course["pass_x"] and upright_at_finish:
print(
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__":
main()
sys.exit(main())

59
scripts/gen_stairs_box.py
View File

@@ -1,7 +1,8 @@
#!/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:
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" />
{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>
</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'
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'
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:.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():
@@ -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-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("--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()
h = args.step_height
d = args.step_depth
n = args.num_steps
sz = args.box_thickness # half-height of each box
platform_depth = args.platform_depth
steps_xml = ""
for i in range(n):
x = i * d
z = i * h + sz # center of box = step top surface - sz
steps_xml += STEP_TPL.format(n=i, sx=d/2, sz=sz, x=x, z=z)
top = (i + 1) * h
steps_xml += STEP_TPL.format(
name=f"step_up_{i}", sx=d / 2, sz=top / 2,
x=(i + 0.5) * d, z=top / 2,
)
# 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
fill_sx = total_depth / 2
fill_sz = total_height / 2
fill_x = total_depth / 2
fill_z = -fill_sz
platform_start = n * d
platform_end = platform_start + platform_depth
steps_xml += PLAT_TPL.format(
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(),
fill_sx=fill_sx, fill_sz=fill_sz,
fill_x=fill_x, fill_z=fill_z)
for i in range(n):
top = (n - i - 1) * h
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__), "..",
"motrix_envs", "src", "motrix_envs", "locomotion",
@@ -86,8 +89,10 @@ def main():
f.write(out)
max_h = n * h
print(f"Generated {n} steps × {h*100:.0f}cm = {max_h*100:.0f}cm total")
print(f" step depth: {d*100:.0f}cm box thickness: {sz*200:.0f}cm")
finish_x = platform_end + n * d
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}")

56
scripts/run_dreamwaq_sim2sim.sh
View File

@@ -7,13 +7,18 @@ RUNS_DIR="${PROJECT_DIR}/runs/go1-dreamwaq-walk/rslrl"
checkpoint=""
output=""
terrain="flat"
terrain="stairs_box"
level="0"
forward_speed="0.5"
timeout="30"
export_only=false
manual=false
headless=false
no_validation=false
usage() {
cat <<'EOF'
Export a DreamWaQ checkpoint to ONNX and run MuJoCo sim2sim.
Export a DreamWaQ checkpoint and run the MuJoCo stair-course validation.
Usage:
scripts/run_dreamwaq_sim2sim.sh [options]
@@ -22,15 +27,21 @@ Options:
-c, --checkpoint PATH Checkpoint to export. Defaults to the newest model_*.pt.
-o, --output PATH ONNX output path. Defaults to policy.onnx beside checkpoint.
-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.
--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.
-h, --help Show this help.
Examples:
scripts/run_dreamwaq_sim2sim.sh
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
}
@@ -61,6 +72,28 @@ while (($#)); do
level="$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=true
shift
@@ -80,6 +113,8 @@ case "$terrain" in
*) die "unsupported terrain: ${terrain}" ;;
esac
[[ "$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"
if [[ -z "$checkpoint" ]]; then
@@ -118,7 +153,14 @@ if [[ "$export_only" == true ]]; then
fi
echo "[Pipeline] starting MuJoCo: terrain=${terrain}, level=${level}"
exec uv run scripts/dreamwaq_sim2sim_mujoco.py \
--onnx "$output" \
--terrain "$terrain" \
sim_args=(
--onnx "$output"
--terrain "$terrain"
--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[@]}"