From 660393a9ee1bf323b8d7f8767211c5be6556ce54 Mon Sep 17 00:00:00 2001
From: 8x54zj-m <8x54zj-m@motrixlab.local>
Date: Wed, 22 Jul 2026 02:36:36 +0800
Subject: [PATCH] Add closed-loop DreamWaQ stair validation
---
.../locomotion/go1/xmls/scene_stairs_box.xml | 34 +-
scripts/dreamwaq_sim2sim_mujoco.py | 538 ++++++++++++------
scripts/gen_stairs_box.py | 59 +-
scripts/run_dreamwaq_sim2sim.sh | 56 +-
4 files changed, 451 insertions(+), 236 deletions(-)
diff --git a/motrix_envs/src/motrix_envs/locomotion/go1/xmls/scene_stairs_box.xml b/motrix_envs/src/motrix_envs/locomotion/go1/xmls/scene_stairs_box.xml
index 2bfc8d1..d6ddcff 100644
--- a/motrix_envs/src/motrix_envs/locomotion/go1/xmls/scene_stairs_box.xml
+++ b/motrix_envs/src/motrix_envs/locomotion/go1/xmls/scene_stairs_box.xml
@@ -18,19 +18,25 @@
-
-
-
-
-
-
-
-
-
-
-
-
-
-
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
diff --git a/scripts/dreamwaq_sim2sim_mujoco.py b/scripts/dreamwaq_sim2sim_mujoco.py
index 3565408..da77da5 100644
--- a/scripts/dreamwaq_sim2sim_mujoco.py
+++ b/scripts/dreamwaq_sim2sim_mujoco.py
@@ -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())
diff --git a/scripts/gen_stairs_box.py b/scripts/gen_stairs_box.py
index 3f1bba1..8bc6ade 100644
--- a/scripts/gen_stairs_box.py
+++ b/scripts/gen_stairs_box.py
@@ -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 = '''
material="motphys-ground" contype="1" conaffinity="0" priority="0" friction="0.6" />
{steps}
-
-
-
'''
-STEP_TPL = ' \n'
-PLAT_TPL = ' \n'
+STEP_TPL = ' \n'
+PLAT_TPL = ' \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}")
diff --git a/scripts/run_dreamwaq_sim2sim.sh b/scripts/run_dreamwaq_sim2sim.sh
index c63aa85..74711a9 100755
--- a/scripts/run_dreamwaq_sim2sim.sh
+++ b/scripts/run_dreamwaq_sim2sim.sh
@@ -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[@]}"