diff --git a/scripts/dreamwaq_sim2sim_mujoco.py b/scripts/dreamwaq_sim2sim_mujoco.py index 0e19da5..db2dd0d 100644 --- a/scripts/dreamwaq_sim2sim_mujoco.py +++ b/scripts/dreamwaq_sim2sim_mujoco.py @@ -47,10 +47,10 @@ 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.0, 0.9, -1.8, - 0.0, 0.9, -1.8, - 0.0, 0.9, -1.8, - 0.0, 0.9, -1.8, + -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, ) @@ -272,7 +272,7 @@ def main(): 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 + spawn_z = hfield_z(spawn_x, spawn_y) + 0.34 def reset_state(): mujoco.mj_resetData(model, data) @@ -318,6 +318,7 @@ def main(): 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 @@ -335,6 +336,7 @@ def main(): 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 @@ -352,8 +354,9 @@ def main(): 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 + if not history_initialized: + history[:, -1] = obs + history_initialized = True action = session.run( None, { @@ -367,6 +370,9 @@ def main(): 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] diff --git a/scripts/evaluate_dreamwaq_stairs.py b/scripts/evaluate_dreamwaq_stairs.py new file mode 100644 index 0000000..ea3c2c2 --- /dev/null +++ b/scripts/evaluate_dreamwaq_stairs.py @@ -0,0 +1,401 @@ +#!/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()) diff --git a/scripts/export_dreamwaq_onnx_new.py b/scripts/export_dreamwaq_onnx_new.py index b0cb07d..4250dcc 100644 --- a/scripts/export_dreamwaq_onnx_new.py +++ b/scripts/export_dreamwaq_onnx_new.py @@ -28,7 +28,7 @@ class DwaqInfer(torch.nn.Module): """obs: (N,45), obs_history: (N,225) → action: (N,12)""" code = self.vae.deterministic_code(obs_history) # (N,19) latent = torch.cat([code, obs], dim=-1) # (N,64) - return torch.clamp(self.actor_mlp(latent), -4.0, 4.0) + return self.actor_mlp(latent) def main():