Add independent DreamWaQ stair capability benchmark

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:]
+                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]

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())

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():