Add closed-loop DreamWaQ stair validation
This commit is contained in:
8x54zj-m
@@ -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>
|
||||||
|
|||||||
@@ -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 = None
|
||||||
|
if not args.headless:
|
||||||
|
keyboard = Keyboard()
|
||||||
|
keyboard.start()
|
||||||
view = viewer.launch_passive(model, data)
|
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")
|
||||||
|
if view is not None:
|
||||||
view.cam.lookat = data.body(trunk_id).xpos.copy()
|
view.cam.lookat = data.body(trunk_id).xpos.copy()
|
||||||
view.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
|
view.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
|
||||||
view.cam.trackbodyid = trunk_id
|
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()):
|
||||||
while view.is_running() and not g_exit_requested:
|
keys = keyboard.held if keyboard is not None else set()
|
||||||
keys = kb.held
|
if "escape" in keys:
|
||||||
if 'escape' in keys: break
|
break
|
||||||
if 'r' in keys:
|
if "r" in keys:
|
||||||
reset_state()
|
reset_state()
|
||||||
last_action[:] = 0; history[:] = 0
|
last_action.fill(0.0)
|
||||||
print("[R] Reset")
|
history.fill(0.0)
|
||||||
if ' ' in keys: vx = vy = wz = 0.0
|
reached_top = False
|
||||||
|
fallen_since = None
|
||||||
|
next_log_time = 0.0
|
||||||
|
physics_step = 0
|
||||||
|
wall_start = time.monotonic()
|
||||||
|
print("[Reset]")
|
||||||
|
|
||||||
vx = MAX_VX if 'w' in keys else (-MAX_VX if 's' in keys else 0.0)
|
forward = 0.0 if args.manual else args.forward_speed
|
||||||
vy = MAX_VY if 'q' in keys else (-MAX_VY if 'e' in keys else 0.0)
|
vx = MAX_VX if "w" in keys else (-MAX_VX if "s" in keys else forward)
|
||||||
wz = MAX_WZ if 'a' in keys else (-MAX_WZ if 'd' 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 " " in keys:
|
||||||
|
vx = vy = wz = 0.0
|
||||||
|
|
||||||
if step % decimation == 0:
|
if physics_step % decimation == 0:
|
||||||
cmd = np.array([vx, vy, wz], dtype=np.float32)
|
command = np.array([vx, vy, wz], dtype=np.float32)
|
||||||
obs = compute_obs(model, data, cmd, last_action)
|
obs = compute_obs(model, data, command, last_action)
|
||||||
|
history[:, :-1] = history[:, 1:]
|
||||||
# Shift history + add new obs
|
history[:, -1] = obs
|
||||||
history = np.concatenate([history[:, 1:, :], obs.reshape(1, 1, -1)], axis=1)
|
action = session.run(
|
||||||
|
None,
|
||||||
# ONNX inference
|
{
|
||||||
outputs = session.run(None, {
|
"obs": obs.reshape(1, -1),
|
||||||
'obs': obs.reshape(1, -1).astype(np.float32),
|
"obs_history": history.reshape(1, -1),
|
||||||
'obs_history': history.reshape(1, -1).astype(np.float32),
|
},
|
||||||
})
|
)[0][0]
|
||||||
action = outputs[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)
|
action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS)
|
||||||
last_action = action.copy()
|
last_action = action.copy()
|
||||||
|
|
||||||
# PD control（对齐训练：目标限位 + 力矩裁剪）
|
|
||||||
target = DEFAULT_ANGLES + action * ACTION_SCALE
|
target = DEFAULT_ANGLES + action * ACTION_SCALE
|
||||||
# 关节目标限位（与训练一致）
|
actuator_joints = model.actuator_trnid[:, 0]
|
||||||
jnt_lo = model.jnt_range[:, 0].copy() if hasattr(model, 'jnt_range') else None
|
target = np.clip(
|
||||||
jnt_hi = model.jnt_range[:, 1].copy() if hasattr(model, 'jnt_range') else None
|
target,
|
||||||
# MuJoCo model.actuator_trnid 可能不直接暴露，改用 model.jnt_range
|
model.jnt_range[actuator_joints, 0],
|
||||||
try:
|
model.jnt_range[actuator_joints, 1],
|
||||||
trnid = model.actuator_trnid[:, 0] # transmission joint indices
|
)
|
||||||
lo = model.jnt_range[trnid, 0]
|
torque = KP * (target - data.qpos[7:19]) - KD * data.qvel[6:18]
|
||||||
hi = model.jnt_range[trnid, 1]
|
data.ctrl[:] = np.clip(torque, -CLIP_TORQUES, CLIP_TORQUES)
|
||||||
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)
|
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()
|
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()
|
||||||
|
|
||||||
# Time sync (policy at 50Hz = 0.02s per step)
|
return status
|
||||||
expected = step * 0.02
|
|
||||||
elapsed = time.time() - loop_t0
|
|
||||||
if elapsed < expected:
|
|
||||||
time.sleep(expected - elapsed)
|
|
||||||
|
|
||||||
step += 1
|
|
||||||
|
|
||||||
kb.stop(); view.close()
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
main()
|
sys.exit(main())
|
||||||
|
|||||||
@@ -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}")
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -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[@]}"
|
||||||
|
|||||||
Reference in New Issue
Block a user