#!/usr/bin/env python3 """Run and validate a DreamWaQ ONNX policy in MuJoCo. 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 """ import argparse import os import queue import signal import sys import threading import time import mujoco import numpy as np import onnxruntime as ort from mujoco import viewer g_exit_requested = False signal.signal(signal.SIGINT, lambda *args: globals().update(g_exit_requested=True)) 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") NUM_OBS = 45 NUM_ACTIONS = 12 HISTORY_LEN = 5 ACTION_SCALE = 0.25 KP = 28.0 KD = 0.7 CLIP_ACTIONS = 100.0 CLIP_TORQUES = 80.0 CLIP_OBS = 100.0 MAX_VX, MAX_VY, MAX_WZ = 1.0, 1.0, 1.0 # MuJoCo actuator and qpos order: FR, FL, RR, RL; hip, thigh, calf. DEFAULT_ANGLES = np.array( [ -0.1, 0.8, -1.5, 0.1, 0.8, -1.5, -0.1, 1.0, -1.5, 0.1, 1.0, -1.5, ], dtype=np.float32, ) class Keyboard: def __init__(self): self.events = queue.Queue() self.running = True self.held = set() self.thread = None self.listener = None @staticmethod def _name(key): try: 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: 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() 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): """Build the 45-value observation in the same order used for training.""" obs = np.zeros(NUM_OBS, dtype=np.float32) 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) obs[9:21] = data.qpos[7:19] - DEFAULT_ANGLES obs[21:33] = data.qvel[6:18] * 0.05 obs[33:45] = last_action return np.clip(obs, -CLIP_OBS, CLIP_OBS) 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( "--realtime-factor", type=float, default=1.0, help="viewer playback speed relative to wall time (default: 1.0)", ) 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 if args.realtime_factor <= 0.0: print("[ERROR] --realtime-factor must be positive", file=sys.stderr) return 2 terrain_map = { "flat": "scene_dreamwaq_flat.xml", "rough": "scene_rough_terrain.xml", "stairs": "scene_stairs_terrain.xml", "dreamwaq": "scene_dreamwaq_terrain.xml", "stairs_test": "scene_stairs_test.xml", "stairs_box": "scene_stairs_box.xml", "flat_stairs": "scene_flat_stairs.xml", } xml_file = os.path.join(XML_DIR, terrain_map[args.terrain]) previous_cwd = os.getcwd() os.chdir(XML_DIR) try: model = mujoco.MjModel.from_xml_path(xml_file) finally: os.chdir(previous_cwd) data = mujoco.MjData(model) if args.terrain == "flat_stairs": 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 elif args.terrain == "stairs_box": spawn_x, spawn_y = -2.0, 0.0 elif args.terrain == "dreamwaq": 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 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 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.34 def reset_state(): 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 print(f"[DreamWaQ] {args.onnx}") 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(f"[Viewer] realtime_factor={args.realtime_factor:.2f}x") print("[CTRL] W/S forward/back, Q/E lateral, A/D yaw, Space stop, R reset, Esc quit") keyboard = None view = None if not args.headless: keyboard = Keyboard() keyboard.start() view = viewer.launch_passive(model, data) trunk_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "trunk") 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 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) history_initialized = False decimation = 4 physics_step = 0 reached_top = False fallen_since = None next_log_time = 0.0 status = 0 wall_start = time.monotonic() 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) history_initialized = False reached_top = False fallen_since = None next_log_time = 0.0 physics_step = 0 wall_start = time.monotonic() print("[Reset]") 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 if physics_step % decimation == 0: command = np.array([vx, vy, wz], dtype=np.float32) obs = compute_obs(model, data, command, last_action) if not history_initialized: history[:, -1] = obs history_initialized = True 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 physics_step > 0: history[:, :-1] = history[:, 1:] history[:, -1] = obs 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 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 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 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 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 if view is not None: view.sync() deadline = wall_start + data.time / args.realtime_factor 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() return status if __name__ == "__main__": sys.exit(main())