Compare commits

...

10 Commits

Author SHA1 Message Date
8x54zj-m
3963773b6f Sample symmetric DreamWaQ speed commands 2026-07-22 20:08:07 +08:00
8x54zj-m
3fa2849123 Sample forward DreamWaQ speed commands 2026-07-22 20:03:02 +08:00
8x54zj-m
509b3d635c Cap DreamWaQ adaptive learning rate 2026-07-22 19:58:19 +08:00
8x54zj-m
944e91e28a Require consecutive failures for terrain downgrade 2026-07-22 19:48:48 +08:00
8x54zj-m
d70b084318 Revert "Limit DreamWaQ commands to stand or walk"
This reverts commit 6d6a805091.
2026-07-22 19:42:40 +08:00
8x54zj-m
6d6a805091 Limit DreamWaQ commands to stand or walk 2026-07-22 19:39:13 +08:00
8x54zj-m
d22292d726 Revert "Use integrated progress for terrain curriculum"
This reverts commit 7358b31de1.
2026-07-22 19:37:26 +08:00
8x54zj-m
7358b31de1 Use integrated progress for terrain curriculum 2026-07-22 19:10:48 +08:00
8x54zj-m
913d5061ff Fix terrain curriculum traversal and hfield roughness 2026-07-22 16:46:32 +08:00
8x54zj-m
a0e8501d11 Add independent DreamWaQ stair capability benchmark 2026-07-22 16:31:09 +08:00
9 changed files with 561 additions and 55 deletions

View File

@@ -43,6 +43,11 @@ def _sanitize_dof_pos(dof_pos: np.ndarray) -> np.ndarray:
return clean
def _advance_streak(streak: np.ndarray, condition: np.ndarray) -> np.ndarray:
"""Increment streaks where condition holds and reset the rest."""
return np.where(condition, streak + 1, 0)
def _scene_file():
"""选择地形场景。DREAMWAQ_TERRAIN=flat|pyramid|flat_stairs。默认 FLAT。"""
global _SCENE_PRINTED
@@ -94,10 +99,11 @@ class DreamWaQCfg(Go1WalkNpEnvCfg):
cell_size: float = 8.0
border_size: float = 5.0
# 命令范围(上游: [-1,1] 对称)
# Symmetric walking commands with no low-speed samples except standing.
@dataclass
class Commands:
vel_limit = [[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]]
forward_speed_range = (0.5, 1.0)
stand_probability = 0.2
commands: Commands = field(default_factory=Commands)
@@ -191,13 +197,20 @@ class DreamWaQTask(Go1WalkTask):
self._hf_cache = None
self._hm_cache = None
# ── 命令:全范围,无课程(匹配上游)──
# ── 命令:站立或随机速度直行 ──
def resample_commands(self, num_envs: int) -> np.ndarray:
lim = np.array(self.cfg.commands.vel_limit, dtype=np.float32)
cmds = np.random.uniform(lim[0], lim[1], size=(num_envs, 3)).astype(np.float32)
small = np.linalg.norm(cmds[:, :2], axis=1) < 0.2
cmds[small, :2] = 0.0
cmds = np.zeros((num_envs, 3), dtype=np.float32)
speed_min, speed_max = self.cfg.commands.forward_speed_range
cmds[:, 0] = np.random.uniform(
speed_min, speed_max, size=num_envs
).astype(np.float32)
signs = np.random.choice(
np.array([-1.0, 1.0], dtype=np.float32), size=num_envs
)
cmds[:, 0] *= signs
stand = np.random.random(num_envs) < self.cfg.commands.stand_probability
cmds[stand, 0] = 0.0
return cmds
# ── 动作裁剪 + 力矩计算 ──
@@ -279,6 +292,25 @@ class DreamWaQTask(Go1WalkTask):
# ── 高度测量 ──
def _ensure_hfield_cache(self) -> bool:
"""Load hfield metadata before reset-time terrain height queries."""
if self._hm_cache is not None:
return True
try:
hf = self._model.get_hfield(0)
self._hf_cache = hf
self._hm_cache = hf.height_matrix
nr, nc = self._hm_cache.shape
b = hf.bound
self._h_xmin, self._h_ymin = b[0], b[1]
self._h_xmax, self._h_ymax = b[3], b[4]
self._h_nr, self._h_nc = nr, nc
cfg_base = getattr(self.cfg, "hfield_z_base", None)
self._h_zbase = float(cfg_base) if cfg_base is not None else float(-self._hm_cache[0, 0])
return True
except Exception:
return False
def _get_heights(self, data: mtx.SceneData) -> np.ndarray:
n = data.shape[0]
nx, ny = len(self._hx), len(self._hy)
@@ -286,25 +318,12 @@ class DreamWaQTask(Go1WalkTask):
base_pos = pose[:, 0:3]
yaw = quaternion.get_yaw(pose[:, 3:7])
cos_yaw, sin_yaw = np.cos(yaw), np.sin(yaw)
try:
if self._hm_cache is None:
hf = self._model.get_hfield(0)
self._hf_cache = hf
hm = hf.height_matrix
self._hm_cache = hm
nr, nc = hm.shape
b = hf.bound
self._h_xmin, self._h_ymin = b[0], b[1]
self._h_xmax, self._h_ymax = b[3], b[4]
self._h_nr, self._h_nc = nr, nc
cfg_base = getattr(self.cfg, "hfield_z_base", None)
self._h_zbase = float(cfg_base) if cfg_base is not None else float(-hm[0, 0])
if not self._ensure_hfield_cache():
return np.zeros((n, nx * ny), dtype=np.float32)
hm, nr, nc = self._hm_cache, self._h_nr, self._h_nc
xmin, ymin, xmax, ymax = self._h_xmin, self._h_ymin, self._h_xmax, self._h_ymax
z_base = self._h_zbase
w, h = xmax - xmin, ymax - ymin
except Exception:
return np.zeros((n, nx * ny), dtype=np.float32)
heights = np.zeros((n, nx * ny), dtype=np.float32)
idx = 0
@@ -324,13 +343,12 @@ class DreamWaQTask(Go1WalkTask):
def _sample_terrain_height(self, xy: np.ndarray, radius: float = 0.0) -> np.ndarray:
n = xy.shape[0]
try:
if not self._ensure_hfield_cache():
return np.zeros(n, dtype=np.float32)
hm, nr, nc = self._hm_cache, self._h_nr, self._h_nc
xmin, ymin, xmax, ymax = self._h_xmin, self._h_ymin, self._h_xmax, self._h_ymax
z_base = self._h_zbase
w, h = xmax - xmin, ymax - ymin
except Exception:
return np.zeros(n, dtype=np.float32)
col = np.clip(np.nan_to_num((xy[:, 0] - xmin) / max(w, 1e-6) * (nc - 1), nan=0).astype(np.int32), 0, nc - 1)
row = np.clip(np.nan_to_num((ymax - xy[:, 1]) / max(h, 1e-6) * (nr - 1), nan=0).astype(np.int32), 0, nr - 1)
if radius <= 0:
@@ -523,22 +541,28 @@ class DreamWaQTask(Go1WalkTask):
ep_bh = old_info.get("ep_base_height", np.zeros(self._num_envs, dtype=np.float32))[done]
avg_bh = np.zeros(num_reset, dtype=np.float32)
avg_bh[mask] = ep_bh[mask] / ep_steps[mask].astype(np.float32)
# orient: -0.5*dt*sin²(theta); bh: -10*dt*z_err²
# Require actual traversal of most of the current terrain cell.
# Tracking velocity while staying near the spawn point must not
# promote the curriculum.
move_up_raw = (
(avg_tracking > 0.5)
& not_fallen
& (avg_orient > -0.0005)
& (avg_bh > -0.001)
& (distance > self._cell_size / 2.0)
)
# ── 渐进升级:连续通过 3 次 + 冷却期 + 每次只升 1 级 ──
if not hasattr(self, '_consecutive_pass_count'):
self._consecutive_pass_count = np.zeros(self._num_envs, dtype=np.int32)
self._pending_upgrade = np.zeros(self._num_envs, dtype=bool)
self._upgrade_cooldown = np.zeros(self._num_envs, dtype=np.int32)
if not hasattr(self, '_consecutive_fail_count'):
self._consecutive_fail_count = np.zeros(self._num_envs, dtype=np.int32)
# 每 env: 连续通过 +1, 失败重置(仅更新被 reset 的 env)
# Cooldown episodes do not contribute to a new validation streak.
cooldown_active = self._upgrade_cooldown[done_idx] > 0
self._consecutive_pass_count[done_idx] = np.where(
move_up_raw,
move_up_raw & ~cooldown_active,
self._consecutive_pass_count[done_idx] + 1, 0)
# 冷却期以该 env 自己完成的 episode 为单位,而不是全局 reset 批次数。
self._upgrade_cooldown[done_idx] = np.maximum(0, self._upgrade_cooldown[done_idx] - 1)
@@ -559,7 +583,14 @@ class DreamWaQTask(Go1WalkTask):
| (avg_orient <= -0.001)
| (avg_bh <= -0.002)
)
move_down = move_down_raw & ~move_up
self._consecutive_fail_count[done_idx] = _advance_streak(
self._consecutive_fail_count[done_idx], move_down_raw
)
move_down_ready = self._consecutive_fail_count[done_idx] >= 2
move_down = move_down_ready & ~move_up
self._consecutive_fail_count[done_idx] = np.where(
move_down, 0, self._consecutive_fail_count[done_idx]
)
new_levels = np.where(move_up, old_levels + 1, old_levels)
new_levels = np.where(move_down, new_levels - 1, new_levels)
new_levels = np.clip(new_levels, 0, self._num_rows - 1)
@@ -578,7 +609,9 @@ class DreamWaQTask(Go1WalkTask):
self._curric_raw_all_ok = 0; self._curric_streak_ok = 0
self._curric_upgrade_ok = 0; self._curric_cooldown_block = 0
self._curric_down_ok = 0; self._curric_down_track = 0
self._curric_down_stability = 0; self._curric_total = 0
self._curric_down_stability = 0; self._curric_progress_ok = 0
self._curric_down_raw = 0
self._curric_total = 0
self._curric_log_counter += 1
self._curric_track_ok += int((avg_tracking > 0.5).sum())
self._curric_survive_ok += int(not_fallen.sum())
@@ -589,8 +622,10 @@ class DreamWaQTask(Go1WalkTask):
self._curric_upgrade_ok += int(move_up.sum())
self._curric_cooldown_block += int(cooldown_blocked.sum())
self._curric_down_ok += int(move_down.sum())
self._curric_down_raw += int(move_down_raw.sum())
self._curric_down_track += int((avg_tracking < 0.35).sum())
self._curric_down_stability += int((~not_fallen | (avg_orient <= -0.001) | (avg_bh <= -0.002)).sum())
self._curric_progress_ok += int((distance > self._cell_size / 2.0).sum())
self._curric_total += num_reset
if self._curric_log_counter % 50 == 0:
def pct(n): return 100*n/max(self._curric_total,1)
@@ -603,9 +638,11 @@ class DreamWaQTask(Go1WalkTask):
f"streak3={pct(self._curric_streak_ok):.0f}% "
f"up={pct(self._curric_upgrade_ok):.0f}% "
f"cool={pct(self._curric_cooldown_block):.0f}% "
f"down={pct(self._curric_down_ok):.0f}% "
f"downRaw={pct(self._curric_down_raw):.0f}% "
f"down2={pct(self._curric_down_ok):.0f}% "
f"downT={pct(self._curric_down_track):.0f}% "
f"downS={pct(self._curric_down_stability):.0f}% "
f"progress={pct(self._curric_progress_ok):.0f}% "
f"max_init={self._max_init_level}")
else:
new_levels = np.random.randint(0, self._max_init_level + 1, size=num_reset, dtype=np.int32)

View File

@@ -1,3 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:5fe0435e385f736b46a910b53d30135d7f6280c0d8daa2bd2f64b7df7962d998
size 1439313
oid sha256:a8289226407a014f04e4d5e9d7f526c2c154d50d1fe91865fce3451c212b31cb
size 912324

View File

@@ -2,7 +2,12 @@ from types import SimpleNamespace
import numpy as np
from motrix_envs.locomotion.go1.dreamwaq import DreamWaQCfg, DreamWaQTask, _sanitize_dof_pos
from motrix_envs.locomotion.go1.dreamwaq import (
DreamWaQCfg,
DreamWaQTask,
_advance_streak,
_sanitize_dof_pos,
)
from motrix_envs.locomotion.go1.walk_np import Go1WalkTask
@@ -22,6 +27,16 @@ def test_sanitize_dof_pos_handles_zero_quaternion_with_nonfinite_joints():
assert clean[1, 8] == 0.0
def test_advance_streak_requires_consecutive_failures():
streak = np.array([0, 1, 3], dtype=np.int32)
streak = _advance_streak(streak, np.array([True, True, False]))
np.testing.assert_array_equal(streak, [1, 2, 0])
streak = _advance_streak(streak, np.array([False, True, True]))
np.testing.assert_array_equal(streak, [0, 3, 1])
def test_apply_action_advances_three_frame_action_history():
task = Go1WalkTask.__new__(Go1WalkTask)
task.get_dof_vel = lambda data: np.ones((1, 12), dtype=np.float32)
@@ -44,6 +59,35 @@ def test_apply_action_advances_three_frame_action_history():
np.testing.assert_array_equal(state.info["current_actions"], 3.0)
def test_dreamwaq_commands_are_zero_or_symmetric_without_low_speeds(monkeypatch):
task = DreamWaQTask.__new__(DreamWaQTask)
task._cfg = SimpleNamespace(
commands=SimpleNamespace(
forward_speed_range=(0.5, 1.0), stand_probability=0.2
)
)
monkeypatch.setattr(
np.random,
"uniform",
lambda low, high, size: np.array([0.5, 0.75, 1.0], dtype=np.float32),
)
monkeypatch.setattr(
np.random,
"choice",
lambda values, size: np.array([-1.0, -1.0, 1.0], dtype=np.float32),
)
monkeypatch.setattr(
np.random,
"random",
lambda size: np.array([0.1, 0.2, 0.9], dtype=np.float32),
)
commands = task.resample_commands(3)
np.testing.assert_array_equal(commands[:, 0], [0.0, -0.75, 1.0])
np.testing.assert_array_equal(commands[:, 1:], 0.0)
def test_dreamwaq_feet_air_time_matches_upstream_contact_filtering():
task = DreamWaQTask.__new__(DreamWaQTask)
task._cfg = SimpleNamespace(ctrl_dt=0.02)

View File

@@ -23,6 +23,7 @@ class DreamWaQPPO(PPO):
# VAE 损失权重(与上游 beta=1.0 一致)
vae_beta: float = 1.0
max_learning_rate: float = 1e-3
@staticmethod
def construct_algorithm(obs: TensorDict, env, cfg: dict, device: str) -> "DreamWaQPPO":
@@ -133,7 +134,9 @@ class DreamWaQPPO(PPO):
if kl_mean > self.desired_kl * 2.0:
self.learning_rate = max(1e-5, self.learning_rate / 1.5)
elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0:
self.learning_rate = min(1e-2, self.learning_rate * 1.5)
self.learning_rate = min(
self.max_learning_rate, self.learning_rate * 1.5
)
for param_group in self.optimizer.param_groups:
param_group['lr'] = self.learning_rate

View File

@@ -3,9 +3,14 @@ from types import SimpleNamespace
import numpy as np
import torch
from motrix_rl.rslrl.torch.train.dreamwaq_ppo import DreamWaQPPO
from motrix_rl.rslrl.torch.wrap_vec_env import RslrlNpEnvWrap
def test_dreamwaq_adaptive_learning_rate_is_capped_at_initial_rate():
assert DreamWaQPPO.max_learning_rate == 1e-3
def test_step_preserves_rewards_and_marks_only_truncations_as_timeouts():
state = SimpleNamespace(
obs=np.zeros((2, 3), dtype=np.float32),

View File

@@ -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]

View File

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

View File

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

View File

@@ -22,7 +22,12 @@ BORDER_M = 5.0
PROPORTIONS = [0.1, 0.1, 0.35, 0.35, 0.1]
CUM = [sum(PROPORTIONS[:i + 1]) for i in range(len(PROPORTIONS))]
PLATFORM_M = 3.0
_SLOPE_SCALE = 0.4 # 上游原值(已验证 z_scale 上限远超 0.54)
_SLOPE_SCALE = 0.25 # MotrixSim-friendly slope range
ROUGH_GRID_M = 0.20 # Correlate roughness over 20 cm, not each 5 cm pixel.
ROUGH_BASE_M = 0.01
ROUGH_GAIN_M = 0.03
OBSTACLE_BASE_M = 0.03
OBSTACLE_GAIN_M = 0.10
CELL_PX = int(CELL_M / HS) # 160
BORDER_PX = int(BORDER_M / HS) # 100
@@ -61,11 +66,16 @@ def draw_slope(canvas, x0, y0, difficulty, noise=False):
hi = max(edge_h, 0)
hf = np.clip(hf, lo, hi).astype(np.uint16)
if noise:
na = int(0.05 / VS)
n = np.random.randint(-na, na + 1, (CELL_PX, CELL_PX), dtype=np.int16)
# 噪声也只在平台外
# Independent 5 cm samples created 10 cm jumps between adjacent
# pixels. Interpolate a 20 cm grid for physically coherent roughness.
coarse_px = max(2, int(round(ROUGH_GRID_M / HS)))
amp = ROUGH_BASE_M + ROUGH_GAIN_M * difficulty
coarse = np.random.uniform(
-amp, amp, (coarse_px, coarse_px)).astype(np.float32)
n = cv2.resize(coarse, (CELL_PX, CELL_PX), interpolation=cv2.INTER_LINEAR)
n[cx - p2:cx + p2, cy - p2:cy + p2] = 0
hf = np.clip(hf.astype(np.int32) + n, 0, 65535).astype(np.uint16)
hf = np.clip(hf.astype(np.float32) * VS + n, 0, None)
hf = np.rint(hf / VS).astype(np.uint16)
canvas[y0:y0 + CELL_PX, x0:x0 + CELL_PX] += hf
@@ -105,7 +115,7 @@ def draw_obstacles(canvas, x0, y0, difficulty):
"""离散障碍物(随机矩形块)。"""
if difficulty <= 0:
return
max_h = int((0.05 + 0.2 * difficulty) / VS)
max_h = int((OBSTACLE_BASE_M + OBSTACLE_GAIN_M * difficulty) / VS)
if max_h <= 0:
return
p2 = PLATFORM_PX // 2