#!/usr/bin/env python3 """Evaluate DreamWaQ checkpoints on a fixed MuJoCo stair benchmark. The benchmark is independent from the training curriculum. Each trial must climb ten box steps, cross the platform, descend ten steps, and finish upright. Results are written to CSV and TensorBoard using the checkpoint iteration as the global step. """ import argparse import csv import glob import json import math import os import re import subprocess import sys from dataclasses import asdict, dataclass import mujoco import numpy as np import onnxruntime as ort from torch.utils.tensorboard import SummaryWriter from dreamwaq_sim2sim_mujoco import ( ACTION_SCALE, CLIP_ACTIONS, CLIP_TORQUES, DEFAULT_ANGLES, HISTORY_LEN, KD, KP, NUM_ACTIONS, NUM_OBS, body_pose, compute_obs, course_geometry, stage_for_x, ) 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_RUNS_DIR = os.path.join(PROJECT_DIR, "runs", "go1-dreamwaq-walk", "rslrl") CHECKPOINT_RE = re.compile(r"model_(\d+)\.pt$") @dataclass class TrialResult: success: bool reached_top: bool reached_finish: bool fell: bool timeout: bool completion_time_s: float max_x_m: float final_x_m: float final_tilt_deg: float final_stage: str def checkpoint_iteration(path: str) -> int: match = CHECKPOINT_RE.search(path) if match is None: raise ValueError(f"checkpoint name must match model_N.pt: {path}") return int(match.group(1)) def latest_run_dir() -> str: runs = [path for path in glob.glob(os.path.join(DEFAULT_RUNS_DIR, "*")) if os.path.isdir(path)] if not runs: raise FileNotFoundError(f"no runs found under {DEFAULT_RUNS_DIR}") return max(runs, key=os.path.getmtime) def discover_checkpoints(args) -> list[str]: if args.checkpoint: paths = [os.path.abspath(path) for path in args.checkpoint] else: run_dir = os.path.abspath(args.run_dir or latest_run_dir()) paths = glob.glob(os.path.join(run_dir, "model_*.pt")) paths = [path for path in paths if os.path.isfile(path)] paths.sort(key=checkpoint_iteration) selected = [] for path in paths: iteration = checkpoint_iteration(path) if iteration < args.min_iteration or iteration > args.max_iteration: continue if args.every > 0 and iteration % args.every != 0: continue selected.append(path) if not selected: raise FileNotFoundError("no checkpoints matched the requested iteration range") return selected def export_checkpoint(checkpoint: str, output: str) -> None: command = [ sys.executable, os.path.join(PROJECT_DIR, "scripts", "export_dreamwaq_onnx_new.py"), "--checkpoint", checkpoint, "--output", output, ] subprocess.run(command, cwd=PROJECT_DIR, check=True) def load_stair_model(step_height_m: float) -> tuple[mujoco.MjModel, dict]: previous_cwd = os.getcwd() os.chdir(XML_DIR) try: model = mujoco.MjModel.from_xml_path(os.path.join(XML_DIR, "scene_stairs_box.xml")) finally: os.chdir(previous_cwd) up_ids = [] down_ids = [] for geom_id in range(model.ngeom): name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_GEOM, geom_id) or "" if name.startswith("step_up_"): up_ids.append((int(name.rsplit("_", 1)[1]), geom_id)) elif name.startswith("step_down_"): down_ids.append((int(name.rsplit("_", 1)[1]), geom_id)) up_ids.sort() down_ids.sort() if not up_ids: raise RuntimeError("stairs_box has no ascending step geoms") num_steps = len(up_ids) for index, geom_id in up_ids: top = (index + 1) * step_height_m model.geom_pos[geom_id, 2] = top / 2.0 model.geom_size[geom_id, 2] = top / 2.0 for index, geom_id in down_ids: top = (num_steps - index - 1) * step_height_m model.geom_pos[geom_id, 2] = top / 2.0 model.geom_size[geom_id, 2] = top / 2.0 platform_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "platform") platform_top = num_steps * step_height_m model.geom_pos[platform_id, 2] = platform_top / 2.0 model.geom_size[platform_id, 2] = platform_top / 2.0 mujoco.mj_setConst(model, mujoco.MjData(model)) course = course_geometry(model) if course is None: raise RuntimeError("failed to read stairs_box course geometry") return model, course def run_trial( model: mujoco.MjModel, course: dict, session: ort.InferenceSession, speed_mps: float, timeout_s: float, lateral_offset_m: float, yaw_offset_rad: float, ) -> TrialResult: data = mujoco.MjData(model) mujoco.mj_resetData(model, data) spawn_x = -2.0 spawn_z = 0.34 data.qpos[0:3] = [spawn_x, lateral_offset_m, spawn_z] data.qpos[3:7] = [math.cos(yaw_offset_rad / 2.0), 0.0, 0.0, math.sin(yaw_offset_rad / 2.0)] data.qpos[7:19] = DEFAULT_ANGLES mujoco.mj_forward(model, data) trunk_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "trunk") history = np.zeros((1, HISTORY_LEN, NUM_OBS), dtype=np.float32) history_initialized = False last_action = np.zeros(NUM_ACTIONS, dtype=np.float32) action = np.zeros(NUM_ACTIONS, dtype=np.float32) command = np.array([speed_mps, 0.0, 0.0], dtype=np.float32) actuator_joints = model.actuator_trnid[:, 0] decimation = 4 physics_step = 0 reached_top = False reached_finish = False fallen_since = None fell = False max_x = spawn_x final_stage = "approach" final_tilt = 0.0 while data.time < timeout_s: if physics_step % decimation == 0: 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)): fell = True 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 = np.clip( DEFAULT_ANGLES + action * ACTION_SCALE, 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) max_x = max(max_x, float(position[0])) final_stage = stage_for_x(position[0], course) final_tilt = float(max(abs(rpy[0]), abs(rpy[1]))) if ( position[0] >= course["platform_start"] - 0.15 and position[2] >= course["platform_top"] + 0.15 ): reached_top = True if position[0] >= course["pass_x"]: reached_finish = True fallen = position[2] < 0.18 or final_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 fallen_since is not None and data.time - fallen_since >= 0.5: fell = True break if reached_top and reached_finish and position[2] < spawn_z + 0.25 and final_tilt < 35.0: return TrialResult( True, True, True, False, False, float(data.time), max_x, float(position[0]), final_tilt, final_stage, ) position, _ = body_pose(data, trunk_id) return TrialResult( False, reached_top, reached_finish, fell, not fell and data.time >= timeout_s, float(data.time), max_x, float(position[0]), final_tilt, final_stage, ) def mastered_height(heights_cm: list[float], success_rates: list[float], threshold: float) -> float: mastered = 0.0 for height, rate in sorted(zip(heights_cm, success_rates)): if rate < threshold: break mastered = height return mastered def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--checkpoint", action="append", help="model_N.pt; may be repeated") parser.add_argument("--run-dir", help="run directory; defaults to newest RSLRL run") parser.add_argument("--min-iteration", type=int, default=0) parser.add_argument("--max-iteration", type=int, default=10**9) parser.add_argument("--every", type=int, default=100, help="checkpoint interval") parser.add_argument("--heights-cm", default="4,6,8,10,12") parser.add_argument("--trials", type=int, default=5) parser.add_argument("--speed", type=float, default=0.5) parser.add_argument("--timeout", type=float, default=30.0) parser.add_argument("--success-threshold", type=float, default=0.8) parser.add_argument("--seed", type=int, default=5) parser.add_argument("--output-dir") parser.add_argument("--keep-onnx", action="store_true") return parser.parse_args() def main() -> int: args = parse_args() if args.trials < 1: raise ValueError("--trials must be at least 1") heights_cm = sorted({float(value) for value in args.heights_cm.split(",")}) if not heights_cm or heights_cm[0] <= 0: raise ValueError("--heights-cm must contain positive values") checkpoints = discover_checkpoints(args) run_dir = os.path.dirname(checkpoints[-1]) output_dir = os.path.abspath(args.output_dir or os.path.join(run_dir, "stairs_eval")) os.makedirs(output_dir, exist_ok=True) csv_path = os.path.join(output_dir, "results.csv") jsonl_path = os.path.join(output_dir, "trials.jsonl") writer = SummaryWriter(log_dir=os.path.join(output_dir, "tensorboard")) csv_exists = os.path.isfile(csv_path) and os.path.getsize(csv_path) > 0 with open(csv_path, "a", newline="") as csv_file, open(jsonl_path, "a") as jsonl_file: fieldnames = [ "iteration", "checkpoint", "step_height_cm", "trials", "success_rate", "reached_top_rate", "reached_finish_rate", "fall_rate", "median_completion_time_s", "mean_max_x_m", ] csv_writer = csv.DictWriter(csv_file, fieldnames=fieldnames) if not csv_exists: csv_writer.writeheader() for checkpoint in checkpoints: iteration = checkpoint_iteration(checkpoint) onnx_path = os.path.join(output_dir, f"policy_{iteration}.onnx") export_checkpoint(checkpoint, onnx_path) session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"]) rng = np.random.default_rng(args.seed + iteration) success_rates = [] print(f"[Eval] iteration={iteration} checkpoint={checkpoint}") for height_cm in heights_cm: model, course = load_stair_model(height_cm / 100.0) results = [] for trial in range(args.trials): lateral = float(rng.uniform(-0.10, 0.10)) yaw = float(rng.uniform(-math.radians(5.0), math.radians(5.0))) speed = float(args.speed * rng.uniform(0.9, 1.1)) result = run_trial(model, course, session, speed, args.timeout, lateral, yaw) results.append(result) jsonl_file.write(json.dumps({ "iteration": iteration, "checkpoint": checkpoint, "step_height_cm": height_cm, "trial": trial, "speed_mps": speed, "lateral_offset_m": lateral, "yaw_offset_rad": yaw, **asdict(result), }) + "\n") success_rate = float(np.mean([result.success for result in results])) top_rate = float(np.mean([result.reached_top for result in results])) finish_rate = float(np.mean([result.reached_finish for result in results])) fall_rate = float(np.mean([result.fell for result in results])) completed = [result.completion_time_s for result in results if result.success] median_time = float(np.median(completed)) if completed else float("nan") mean_max_x = float(np.mean([result.max_x_m for result in results])) success_rates.append(success_rate) suffix = f"{height_cm:g}cm" writer.add_scalar(f"Stairs/success_rate_{suffix}", success_rate, iteration) writer.add_scalar(f"Stairs/reached_top_rate_{suffix}", top_rate, iteration) writer.add_scalar(f"Stairs/reached_finish_rate_{suffix}", finish_rate, iteration) writer.add_scalar(f"Stairs/fall_rate_{suffix}", fall_rate, iteration) if completed: writer.add_scalar(f"Stairs/completion_time_s_{suffix}", median_time, iteration) csv_writer.writerow({ "iteration": iteration, "checkpoint": checkpoint, "step_height_cm": height_cm, "trials": args.trials, "success_rate": success_rate, "reached_top_rate": top_rate, "reached_finish_rate": finish_rate, "fall_rate": fall_rate, "median_completion_time_s": median_time, "mean_max_x_m": mean_max_x, }) print( f" {height_cm:4.1f}cm success={success_rate:5.1%} " f"top={top_rate:5.1%} finish={finish_rate:5.1%} " f"fall={fall_rate:5.1%} max_x={mean_max_x:.2f}m" ) mastered = mastered_height(heights_cm, success_rates, args.success_threshold) if len(heights_cm) == 1: auc = success_rates[0] else: auc = float( np.trapezoid(success_rates, heights_cm) / (heights_cm[-1] - heights_cm[0]) ) writer.add_scalar("Stairs/mastered_step_height_cm", mastered, iteration) writer.add_scalar("Stairs/success_auc", auc, iteration) writer.flush() csv_file.flush() jsonl_file.flush() print(f" mastered={mastered:.1f}cm success_auc={auc:.3f}") if not args.keep_onnx: os.remove(onnx_path) writer.close() print(f"[Done] CSV: {csv_path}") print(f"[Done] TensorBoard: {os.path.join(output_dir, 'tensorboard')}") return 0 if __name__ == "__main__": raise SystemExit(main())