Backup DreamWaQ rslrl stability fixes
This commit is contained in:
8x54zj-m
458
scripts/diag_terrain_collision.py
Normal file
458
scripts/diag_terrain_collision.py
Normal file
@@ -0,0 +1,458 @@
|
||||
#!/usr/bin/env python3
|
||||
"""地形碰撞诊断工具 — 系统性检测 hfield PNG 地形的物理交互质量。
|
||||
|
||||
检测项：
|
||||
1. 脚部穿透（foot z < terrain z）
|
||||
2. 接触力异常（为 0 但应该接触 / 过大）
|
||||
3. NaN/Inf 传播
|
||||
4. 物理爆炸（关节速度飙升）
|
||||
5. 机器人卡住（速度接近 0 但有命令）
|
||||
6. 躯干沉入地形（base z < terrain z）
|
||||
|
||||
用法:
|
||||
uv run scripts/diag_terrain_collision.py # 全量检测
|
||||
uv run scripts/diag_terrain_collision.py --level 5 # 仅 level 5
|
||||
uv run scripts/diag_terrain_collision.py --type stairs # 仅楼梯地形
|
||||
uv run scripts/diag_terrain_collision.py --render # 渲染可视化
|
||||
"""
|
||||
import argparse, os, sys, time
|
||||
import numpy as np
|
||||
|
||||
os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
|
||||
os.environ.setdefault("JAX_PLATFORMS", "cpu")
|
||||
os.environ.setdefault("DREAMWAQ_TERRAIN", "pyramid")
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
import motrix_envs.locomotion.go1.dreamwaq # noqa: F401
|
||||
from motrix_envs import registry as env_registry
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# 配置
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
NUM_ROWS = 10
|
||||
NUM_COLS = 20
|
||||
CELL_M = 8.0
|
||||
BORDER_M = 5.0
|
||||
SETTLE_STEPS = 100 # 自由落体稳定步数（200Hz）
|
||||
TEST_STEPS = 400 # 测试步数
|
||||
STEPS_PER_CELL = 5 # 每个 cell 测试的 spawn 位置数
|
||||
|
||||
# 地形类型名（与 gen_dreamwaq_terrain.py PROPORTIONS 对应）
|
||||
TYPE_NAMES = ["平滑斜坡", "粗糙斜坡", "下行楼梯", "上行楼梯", "离散障碍"]
|
||||
CUM = [0.1, 0.2, 0.55, 0.9, 1.0]
|
||||
|
||||
|
||||
def _cell_origin(row, col):
|
||||
half_x = BORDER_M + NUM_COLS * CELL_M / 2.0
|
||||
half_y = BORDER_M + NUM_ROWS * CELL_M / 2.0
|
||||
cx = -half_x + BORDER_M + col * CELL_M + CELL_M / 2
|
||||
cy = half_y - BORDER_M - row * CELL_M - CELL_M / 2
|
||||
return cx, cy
|
||||
|
||||
|
||||
def _terrain_type(col):
|
||||
choice = col / NUM_COLS + 0.001
|
||||
for i, c in enumerate(CUM):
|
||||
if choice < c:
|
||||
return i
|
||||
return len(TYPE_NAMES) - 1
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# 诊断核心
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
class CollisionDiag:
|
||||
def __init__(self, env, render=False):
|
||||
self.env = env
|
||||
self.render = render
|
||||
self.renderer = None
|
||||
if render:
|
||||
from motrix_envs.np.renderer import NpRenderer
|
||||
self.renderer = NpRenderer(env)
|
||||
|
||||
# 统计数据
|
||||
self.stats = {
|
||||
"penetrations": [], # (row, col, ttype, foot_z - terrain_z)
|
||||
"contact_anomalies": [], # (row, col, ttype, max_contact_force)
|
||||
"nan_events": [], # (row, col, ttype, step)
|
||||
"explosions": [], # (row, col, ttype, max_dof_vel)
|
||||
"stuck_events": [], # (row, col, ttype, speed)
|
||||
"base_sink": [], # (row, col, ttype, base_z - terrain_z)
|
||||
"ok_count": 0,
|
||||
"total_tested": 0,
|
||||
}
|
||||
|
||||
def _get_terrain_z(self, x, y):
|
||||
"""查询 (x, y) 处的地形高度。"""
|
||||
hm = self.env._hm_cache
|
||||
if hm is None:
|
||||
return 0.0
|
||||
nr, nc = hm.shape
|
||||
xmin = self.env._h_xmin
|
||||
ymax = self.env._h_ymax
|
||||
w = self.env._h_xmax - xmin
|
||||
h = ymax - self.env._h_ymin
|
||||
z_base = self.env._h_zbase
|
||||
col = int(np.clip((x - xmin) / max(w, 1e-6) * (nc - 1), 0, nc - 1))
|
||||
row = int(np.clip((ymax - y) / max(h, 1e-6) * (nr - 1), 0, nr - 1))
|
||||
return float(hm[row, col]) + z_base
|
||||
|
||||
def _get_foot_positions(self, state):
|
||||
"""获取所有足部在世界坐标系中的位置 (N, 4, 3)。"""
|
||||
data = state.data
|
||||
n = data.shape[0]
|
||||
feet = np.zeros((n, 4, 3), dtype=np.float32)
|
||||
foot_names = ["FR_foot", "FL_foot", "RR_foot", "RL_foot"]
|
||||
for i, name in enumerate(foot_names):
|
||||
gidx = self.env._model.get_geom_index(name)
|
||||
if gidx is not None:
|
||||
pose = self.env._model.geoms[gidx].get_pose(data)
|
||||
feet[:, i, :] = pose[:, :3]
|
||||
return feet
|
||||
|
||||
def test_cell(self, row, col, level):
|
||||
"""在单个 cell 上测试多个 spawn 位置。"""
|
||||
cx, cy = _cell_origin(row, col)
|
||||
ttype = _terrain_type(col)
|
||||
tname = TYPE_NAMES[ttype]
|
||||
|
||||
# 在 cell 内取几个 spawn 点（中央 + 四角附近）
|
||||
offsets = [
|
||||
(0, 0),
|
||||
(CELL_M * 0.3, 0), (-CELL_M * 0.3, 0),
|
||||
(0, CELL_M * 0.3), (0, -CELL_M * 0.3),
|
||||
][:STEPS_PER_CELL]
|
||||
|
||||
for ox, oy in offsets:
|
||||
sx, sy = cx + ox, cy + oy
|
||||
terrain_z = float(self.env._sample_terrain_height(
|
||||
np.array([[sx, sy]], dtype=np.float32), radius=0.35)[0])
|
||||
|
||||
self.stats["total_tested"] += 1
|
||||
self._test_single(sx, sy, terrain_z, row, col, level, ttype, tname)
|
||||
|
||||
def _test_single(self, sx, sy, terrain_z, row, col, level, ttype, tname):
|
||||
"""单个 spawn 点的完整测试。"""
|
||||
env = self.env
|
||||
data = env._state.data
|
||||
n = 1
|
||||
|
||||
# Spawn
|
||||
init_pos = env._init_dof_pos.copy().reshape(1, -1)
|
||||
init_pos[0, 0] = sx
|
||||
init_pos[0, 1] = sy
|
||||
spawn_z = terrain_z + 0.45 # 标准 clearance
|
||||
init_pos[0, 2] = spawn_z
|
||||
yaw = np.random.uniform(-np.pi, np.pi)
|
||||
init_pos[0, 3:7] = [0, 0, np.sin(yaw / 2), np.cos(yaw / 2)]
|
||||
|
||||
data.reset(env._model)
|
||||
data.set_dof_pos(init_pos, env._model)
|
||||
env._model.forward_kinematic(data)
|
||||
|
||||
# 给一个随机动作防止完全静止（小幅度 PD 控制）
|
||||
action = 0.1 * (np.random.randn(n, 12).astype(np.float32))
|
||||
|
||||
has_nan = False
|
||||
max_dof_vel = 0.0
|
||||
max_contact_force = 0.0
|
||||
min_foot_z = spawn_z
|
||||
min_base_z = spawn_z
|
||||
total_speed = 0.0
|
||||
|
||||
for step in range(SETTLE_STEPS + TEST_STEPS):
|
||||
# Apply PD
|
||||
actions_scaled = action * env.cfg.control_config.action_scale
|
||||
target = actions_scaled + env.default_angles.reshape(1, -1)
|
||||
lo = env._model.joint_limits[0].reshape(1, -1)
|
||||
hi = env._model.joint_limits[1].reshape(1, -1)
|
||||
target = np.clip(target, lo, hi)
|
||||
torques = env.kps * (target - env.get_dof_pos(data)) - env.kds * env.get_dof_vel(data)
|
||||
data.actuator_ctrls = np.clip(torques, -80.0, 80.0)
|
||||
|
||||
env._model.step(data)
|
||||
|
||||
# ── 检测 ──
|
||||
|
||||
# 1. NaN 检测
|
||||
dv = np.array(data.dof_vel)
|
||||
dp = np.array(data.dof_pos)
|
||||
if np.any(~np.isfinite(dv)) or np.any(~np.isfinite(dp)):
|
||||
self.stats["nan_events"].append((row, col, ttype, step))
|
||||
has_nan = True
|
||||
break
|
||||
|
||||
# 2. 速度爆炸
|
||||
cur_max_vel = float(np.max(np.abs(dv[:, 6:])))
|
||||
max_dof_vel = max(max_dof_vel, cur_max_vel)
|
||||
if cur_max_vel > 100.0:
|
||||
self.stats["explosions"].append((row, col, ttype, cur_max_vel))
|
||||
break
|
||||
|
||||
# 3. 接触力
|
||||
# Read contact forces via MotrixSim sensor API
|
||||
cf_vals = []
|
||||
for foot in ["FR", "FL", "RR", "RL"]:
|
||||
try:
|
||||
v = env._model.get_sensor_value(f"{foot}_foot_contact", data)
|
||||
cf_vals.append(v[0])
|
||||
except Exception:
|
||||
cf_vals.append(np.zeros(3, dtype=np.float32))
|
||||
cf_raw = np.concatenate(cf_vals)
|
||||
cur_cf = float(np.max(np.abs(cf_raw)))
|
||||
max_contact_force = max(max_contact_force, cur_cf)
|
||||
|
||||
# 4. 足部穿透（settle 之后测）
|
||||
if step >= SETTLE_STEPS:
|
||||
feet = self._get_foot_positions(env._state)
|
||||
for fi in range(4):
|
||||
fz = float(feet[0, fi, 2])
|
||||
tz = self._get_terrain_z(float(feet[0, fi, 0]), float(feet[0, fi, 1]))
|
||||
penetration = tz - fz # 正 = 穿透
|
||||
min_foot_z = min(min_foot_z, fz - tz) # 相对地形的高度
|
||||
if penetration > 0.02: # > 2cm 穿透
|
||||
self.stats["penetrations"].append(
|
||||
(row, col, ttype, float(penetration)))
|
||||
|
||||
# 5. 躯干沉入
|
||||
base_pose = env._body.get_pose(data)
|
||||
bz = float(base_pose[0, 2])
|
||||
tz = self._get_terrain_z(float(base_pose[0, 0]), float(base_pose[0, 1]))
|
||||
min_base_z = min(min_base_z, bz - tz)
|
||||
if bz < tz + 0.10 and step >= SETTLE_STEPS: # 躯干离地面 < 10cm
|
||||
self.stats["base_sink"].append((row, col, ttype, float(bz - tz)))
|
||||
|
||||
# 6. 速度（卡住检测）
|
||||
if step >= SETTLE_STEPS:
|
||||
linvel = env.get_local_linvel(data)[0]
|
||||
total_speed += float(np.linalg.norm(linvel[:2]))
|
||||
|
||||
if self.renderer and step % 20 == 0:
|
||||
self.renderer.render()
|
||||
time.sleep(0.005)
|
||||
|
||||
# ── 统计 ──
|
||||
if not has_nan and max_dof_vel < 100.0:
|
||||
avg_speed = total_speed / max(TEST_STEPS, 1)
|
||||
if avg_speed < 0.01 and step == SETTLE_STEPS + TEST_STEPS - 1:
|
||||
self.stats["stuck_events"].append((row, col, ttype, avg_speed))
|
||||
|
||||
if not has_nan and max_dof_vel < 50.0 and min_base_z > 0.05:
|
||||
self.stats["ok_count"] += 1
|
||||
|
||||
# 记录接触力异常（完全无接触力但足部应该着地）
|
||||
if max_contact_force < 1e-3 and step == SETTLE_STEPS + TEST_STEPS - 1:
|
||||
feet = self._get_foot_positions(env._state)
|
||||
for fi in range(4):
|
||||
fz = float(feet[0, fi, 2])
|
||||
tz = self._get_terrain_z(float(feet[0, fi, 0]), float(feet[0, fi, 1]))
|
||||
if abs(fz - tz) < 0.05: # 足部接近地面
|
||||
self.stats["contact_anomalies"].append(
|
||||
(row, col, ttype, 0.0)) # 0 = 无接触力
|
||||
break
|
||||
|
||||
def run(self, target_level=None, target_type=None):
|
||||
"""遍历所有地形 cell 测试。"""
|
||||
total = 0
|
||||
start_t = time.time()
|
||||
|
||||
for row in range(NUM_ROWS):
|
||||
if target_level is not None and row != target_level:
|
||||
continue
|
||||
for col in range(NUM_COLS):
|
||||
ttype = _terrain_type(col)
|
||||
if target_type is not None and ttype != target_type:
|
||||
continue
|
||||
total += 1
|
||||
|
||||
print(f"[Diag] 测试 {NUM_ROWS}×{NUM_COLS} 网格, "
|
||||
f"每 cell {STEPS_PER_CELL} 个 spawn 点 = {total * STEPS_PER_CELL} 次")
|
||||
print(f"[Diag] Settle={SETTLE_STEPS}步 Test={TEST_STEPS}步")
|
||||
|
||||
count = 0
|
||||
for row in range(NUM_ROWS):
|
||||
if target_level is not None and row != target_level:
|
||||
continue
|
||||
for col in range(NUM_COLS):
|
||||
ttype = _terrain_type(col)
|
||||
if target_type is not None and ttype != target_type:
|
||||
continue
|
||||
self.test_cell(row, col, row)
|
||||
count += 1
|
||||
if count % 10 == 0:
|
||||
elapsed = time.time() - start_t
|
||||
print(f" [{count}/{total}] {elapsed:.1f}s "
|
||||
f"OK={self.stats['ok_count']}/{self.stats['total_tested']} "
|
||||
f"穿透={len(self.stats['penetrations'])} "
|
||||
f"NaN={len(self.stats['nan_events'])} "
|
||||
f"爆炸={len(self.stats['explosions'])}")
|
||||
|
||||
elapsed = time.time() - start_t
|
||||
self._report(elapsed)
|
||||
|
||||
def _report(self, elapsed):
|
||||
s = self.stats
|
||||
t = s["total_tested"]
|
||||
print("\n" + "=" * 70)
|
||||
print(f"地形碰撞诊断报告 ({elapsed:.1f}s, {t} 次测试)")
|
||||
print("=" * 70)
|
||||
|
||||
# 总体统计
|
||||
ok_rate = s["ok_count"] / max(t, 1) * 100
|
||||
print(f"\n ✅ 通过: {s['ok_count']}/{t} ({ok_rate:.1f}%)")
|
||||
|
||||
# 穿透
|
||||
if s["penetrations"]:
|
||||
pens = s["penetrations"]
|
||||
pvals = [p[3] for p in pens]
|
||||
print(f"\n 🔴 足部穿透: {len(pens)} 次")
|
||||
print(f" 最大穿透: {max(pvals):.3f}m 平均: {np.mean(pvals):.3f}m")
|
||||
|
||||
# 按地形类型分组
|
||||
by_type = {}
|
||||
for p in pens:
|
||||
tn = TYPE_NAMES[p[2]]
|
||||
by_type.setdefault(tn, []).append(p[3])
|
||||
for tn, vals in sorted(by_type.items()):
|
||||
print(f" {tn}: {len(vals)}次 max={max(vals):.3f}m avg={np.mean(vals):.3f}m")
|
||||
|
||||
# 按 level 分组
|
||||
by_level = {}
|
||||
for p in pens:
|
||||
by_level.setdefault(p[0], []).append(p[3])
|
||||
print(f" 按 level: ", end="")
|
||||
for lv in sorted(by_level.keys()):
|
||||
vals = by_level[lv]
|
||||
print(f"L{lv}({len(vals)}x max={max(vals):.3f}) ", end="")
|
||||
print()
|
||||
else:
|
||||
print(f"\n 🟢 足部穿透: 0 次 ✓")
|
||||
|
||||
# NaN
|
||||
if s["nan_events"]:
|
||||
print(f"\n 🔴 NaN/Inf: {len(s['nan_events'])} 次")
|
||||
for n in s["nan_events"][:5]:
|
||||
print(f" L{n[0]} C{n[1]:02d} {TYPE_NAMES[n[2]]} @step {n[3]}")
|
||||
else:
|
||||
print(f" 🟢 NaN/Inf: 0 次 ✓")
|
||||
|
||||
# 物理爆炸
|
||||
if s["explosions"]:
|
||||
print(f"\n 🟡 速度爆炸 (>{100}rad/s): {len(s['explosions'])} 次")
|
||||
for e in s["explosions"][:5]:
|
||||
print(f" L{e[0]} C{e[1]:02d} {TYPE_NAMES[e[2]]} max_vel={e[3]:.1f}")
|
||||
else:
|
||||
print(f" 🟢 速度爆炸: 0 次 ✓")
|
||||
|
||||
# 接触力异常
|
||||
if s["contact_anomalies"]:
|
||||
print(f"\n 🟡 接触力异常 (无接触但足部近地): {len(s['contact_anomalies'])} 次")
|
||||
else:
|
||||
print(f" 🟢 接触力异常: 0 次 ✓")
|
||||
|
||||
# 卡住
|
||||
if s["stuck_events"]:
|
||||
print(f"\n 🟡 卡住 (速度<0.01m/s): {len(s['stuck_events'])} 次")
|
||||
by_type = {}
|
||||
for e in s["stuck_events"]:
|
||||
tn = TYPE_NAMES[e[2]]
|
||||
by_type.setdefault(tn, []).append(1)
|
||||
for tn, vals in sorted(by_type.items()):
|
||||
print(f" {tn}: {len(vals)}次")
|
||||
else:
|
||||
print(f" 🟢 卡住: 0 次 ✓")
|
||||
|
||||
# 躯干沉入
|
||||
if s["base_sink"]:
|
||||
sinks = s["base_sink"]
|
||||
print(f"\n 🟡 躯干贴地 (<10cm): {len(sinks)} 次")
|
||||
vals = [s[3] for s in sinks]
|
||||
print(f" 最差: {min(vals):.3f}m 平均: {np.mean(vals):.3f}m")
|
||||
else:
|
||||
print(f" 🟢 躯干贴地: 0 次 ✓")
|
||||
|
||||
# 按地形类型汇总
|
||||
print(f"\n── 按地形类型汇总 ──")
|
||||
for ti, tn in enumerate(TYPE_NAMES):
|
||||
pens_t = sum(1 for p in s["penetrations"] if p[2] == ti)
|
||||
nans_t = sum(1 for n in s["nan_events"] if n[2] == ti)
|
||||
exps_t = sum(1 for e in s["explosions"] if e[2] == ti)
|
||||
icon = "🔴" if (pens_t + nans_t + exps_t) > 0 else "🟢"
|
||||
print(f" {icon} {tn}: 穿透{pens_t} NaN{nans_t} 爆炸{exps_t}")
|
||||
|
||||
# 结论
|
||||
print(f"\n── 全局结论 ──")
|
||||
if ok_rate > 95 and len(s["nan_events"]) == 0 and len(s["penetrations"]) < 5:
|
||||
print(f" ✅ 仿真环境正常 — hfield 地形碰撞基本可靠")
|
||||
elif ok_rate > 80:
|
||||
print(f" ⚠ 仿真环境存在一些问题 — 部分地形类型/level 有问题")
|
||||
else:
|
||||
print(f" 🔴 仿真环境问题严重 — 需要排查 MotrixSim hfield 实现")
|
||||
|
||||
if len(s["penetrations"]) > 10:
|
||||
print(f" 💡 穿透较多 → hfield 碰撞面可能偏软或 contact margin 过大")
|
||||
print(f" 可尝试增大 geom margin 或检查 MotrixSim 碰撞参数")
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
# Main
|
||||
# ═══════════════════════════════════════════════════════════════════════
|
||||
|
||||
def main():
|
||||
p = argparse.ArgumentParser(description="DreamWaQ 地形碰撞诊断")
|
||||
p.add_argument("--level", type=int, default=None, help="仅测试指定 level (0-9)")
|
||||
p.add_argument("--type", dest="ttype", default=None,
|
||||
choices=["slope", "rough", "stairs_down", "stairs_up", "obstacles"],
|
||||
help="仅测试指定地形类型")
|
||||
p.add_argument("--render", action="store_true", help="渲染可视化")
|
||||
p.add_argument("--settle", type=int, default=SETTLE_STEPS)
|
||||
p.add_argument("--test", type=int, default=TEST_STEPS)
|
||||
args = p.parse_args()
|
||||
|
||||
type_map = {"slope": 0, "rough": 1, "stairs_down": 2, "stairs_up": 3, "obstacles": 4}
|
||||
target_type = type_map.get(args.ttype)
|
||||
|
||||
print("[Diag] 创建 DreamWaQ 环境...")
|
||||
env = env_registry.make("go1-dreamwaq-walk", num_envs=1)
|
||||
|
||||
# 如果指定 level，强制执行
|
||||
if args.level is not None:
|
||||
env._force_level = args.level
|
||||
print(f"[Diag] 强制 level={args.level}")
|
||||
|
||||
env.init_state()
|
||||
|
||||
# 确保 terrain 初始化（触发 hfield cache）
|
||||
if env._model.num_hfields > 0:
|
||||
hf = env._model.get_hfield(0)
|
||||
env._hm_cache = hf.height_matrix
|
||||
env._h_xmin, env._h_ymin = hf.bound[0], hf.bound[1]
|
||||
env._h_xmax, env._h_ymax = hf.bound[3], hf.bound[4]
|
||||
env._h_nr, env._h_nc = hf.height_matrix.shape
|
||||
cfg_base = getattr(env.cfg, "hfield_z_base", None)
|
||||
env._h_zbase = float(cfg_base) if cfg_base is not None else float(-hf.height_matrix[0, 0])
|
||||
print(f"[Diag] HField: {hf.height_matrix.shape} bound={hf.bound}")
|
||||
print(f"[Diag] z_base={env._h_zbase:.3f}")
|
||||
else:
|
||||
print("[Diag] WARNING: 无 hfield（flat 场景），跳过诊断")
|
||||
return
|
||||
|
||||
diag = CollisionDiag(env, render=args.render)
|
||||
diag.SETTLE_STEPS = args.settle
|
||||
diag.TEST_STEPS = args.test
|
||||
|
||||
try:
|
||||
diag.run(target_level=args.level, target_type=target_type)
|
||||
except KeyboardInterrupt:
|
||||
print("\n[Diag] 中断")
|
||||
finally:
|
||||
if diag.renderer:
|
||||
diag.renderer.close()
|
||||
print("[Diag] 完成")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -31,7 +31,8 @@ HISTORY_LEN = 5
|
||||
ACTION_SCALE = 0.25
|
||||
KP = 28.0
|
||||
KD = 0.7
|
||||
CLIP_ACTIONS = 23.7
|
||||
CLIP_ACTIONS = 4.0
|
||||
CLIP_TORQUES = 80.0
|
||||
CLIP_OBS = 100.0
|
||||
MAX_VX, MAX_VY, MAX_WZ = 1.0, 1.0, 1.0
|
||||
|
||||
@@ -122,7 +123,7 @@ def main():
|
||||
|
||||
# Select XML scene
|
||||
terrain_map = {
|
||||
"flat": "scene_motor_actuator.xml",
|
||||
"flat": "scene_dreamwaq_flat.xml",
|
||||
"rough": "scene_rough_terrain.xml",
|
||||
"stairs": "scene_stairs_terrain.xml",
|
||||
"dreamwaq": "scene_dreamwaq_terrain.xml",
|
||||
@@ -243,10 +244,22 @@ def main():
|
||||
action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS)
|
||||
last_action = action.copy()
|
||||
|
||||
# PD control
|
||||
# PD control（对齐训练：目标限位 + 力矩裁剪）
|
||||
target = DEFAULT_ANGLES + action * ACTION_SCALE
|
||||
# 关节目标限位（与训练一致）
|
||||
jnt_lo = model.jnt_range[:, 0].copy() if hasattr(model, 'jnt_range') else None
|
||||
jnt_hi = model.jnt_range[:, 1].copy() if hasattr(model, 'jnt_range') else None
|
||||
# MuJoCo model.actuator_trnid 可能不直接暴露，改用 model.jnt_range
|
||||
try:
|
||||
trnid = model.actuator_trnid[:, 0] # transmission joint indices
|
||||
lo = model.jnt_range[trnid, 0]
|
||||
hi = model.jnt_range[trnid, 1]
|
||||
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_ACTIONS, CLIP_ACTIONS)
|
||||
data.ctrl[:] = np.clip(torques, -CLIP_TORQUES, CLIP_TORQUES)
|
||||
mujoco.mj_step(model, data)
|
||||
view.sync()
|
||||
|
||||
|
||||
@@ -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 self.actor_mlp(latent)
|
||||
return torch.clamp(self.actor_mlp(latent), -4.0, 4.0)
|
||||
|
||||
|
||||
def main():
|
||||
|
||||
@@ -22,13 +22,20 @@ PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
|
||||
class DwaqInfer(tnn.Module):
|
||||
"""Deterministic inference: CENet mean code + actor."""
|
||||
"""Deterministic inference: CENet mean code + RunningStats normalize + actor."""
|
||||
def __init__(self, ac: ActorCritic_DWAQ):
|
||||
super().__init__()
|
||||
self.encoder = ac.encoder
|
||||
self.encode_mean_vel = ac.encode_mean_vel
|
||||
self.encode_mean_latent = ac.encode_mean_latent
|
||||
self.actor = ac.actor
|
||||
# Bake RunningStats into the exported model
|
||||
if hasattr(ac, 'actor_normalizer') and ac.actor_normalizer.count > 10:
|
||||
self.register_buffer("norm_mean", ac.actor_normalizer.mean.clone())
|
||||
self.register_buffer("norm_std", ac.actor_normalizer.var.sqrt().clone())
|
||||
self.has_norm = True
|
||||
else:
|
||||
self.has_norm = False
|
||||
|
||||
def forward(self, obs, obs_history):
|
||||
h = self.encoder(obs_history.reshape(obs_history.shape[0], -1)) # (B,225)->(B,64)
|
||||
@@ -36,6 +43,8 @@ class DwaqInfer(tnn.Module):
|
||||
latent = self.encode_mean_latent(h) # (B,16) mean latent
|
||||
code = torch.cat([vel, latent], dim=-1) # (B,19) = [vel, latent]
|
||||
x = torch.cat([code, obs], dim=-1) # (B,64) = [code, obs]
|
||||
if self.has_norm:
|
||||
x = torch.clamp((x - self.norm_mean) / (self.norm_std + 1e-8), -5.0, 5.0)
|
||||
return self.actor(x) # (B,12)
|
||||
|
||||
|
||||
|
||||
103
scripts/play_amp_mujoco.py
Normal file
103
scripts/play_amp_mujoco.py
Normal file
@@ -0,0 +1,103 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Deploy Go1 AMP model in MuJoCo."""
|
||||
import mujoco, numpy as np, os, torch, time, sys
|
||||
from mujoco import viewer
|
||||
|
||||
# Use AMP's rsl_rl (NOT MotrixLab's)
|
||||
sys.path.insert(0, '/home/8x54zj-m/amp_go2/rsl_rl')
|
||||
sys.path.insert(0, '/home/8x54zj-m/amp_go2/legged_gym')
|
||||
# Force reload of rsl_rl to pick up AMP version
|
||||
for mod in list(sys.modules.keys()):
|
||||
if 'rsl_rl' in mod:
|
||||
del sys.modules[mod]
|
||||
|
||||
XML_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
|
||||
'motrix_envs/src/motrix_envs/locomotion/go1/xmls')
|
||||
CKPT = '/home/8x54zj-m/amp_go2/legged_gym/logs/go1_amp/Jul05_01-15-41_flat/model_6000.pt'
|
||||
|
||||
import argparse
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--terrain", default="flat", choices=["flat", "flat_stairs"])
|
||||
args = p.parse_args()
|
||||
|
||||
scene = "scene_flat_stairs.xml" if args.terrain == "flat_stairs" else "go1_motor_actuator.xml"
|
||||
os.chdir(XML_DIR)
|
||||
m = mujoco.MjModel.from_xml_string(open(scene).read())
|
||||
d = mujoco.MjData(m)
|
||||
|
||||
# Spawn on level 1 platform if using flat_stairs
|
||||
spawn_x, spawn_y, spawn_z = 0.0, 0.0, 0.45
|
||||
if args.terrain == "flat_stairs":
|
||||
col = np.random.randint(0, 4)
|
||||
spawn_x = -12.0 + col * 8.0
|
||||
spawn_y = -4.0 # level 1 center
|
||||
# Sample terrain height at spawn
|
||||
floor_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_GEOM, "floor")
|
||||
if floor_id >= 0 and m.geom_type[floor_id] == mujoco.mjtGeom.mjGEOM_HFIELD:
|
||||
hf = m.geom_dataid[floor_id]
|
||||
nr, nc = int(m.hfield_nrow[hf]), int(m.hfield_ncol[hf])
|
||||
sx, sy, zt, sb = m.hfield_size[hf]
|
||||
hd = m.hfield_data[m.hfield_adr[hf]:m.hfield_adr[hf]+nr*nc].reshape(nr, nc)
|
||||
gp = m.geom_pos[floor_id]
|
||||
col_idx = int(np.clip(((spawn_x-gp[0])/sx*0.5+0.5)*(nc-1), 0, nc-1))
|
||||
row_idx = int(np.clip(((spawn_y-gp[1])/sy*0.5+0.5)*(nr-1), 0, nr-1))
|
||||
spawn_z = float(gp[2] + sb + hd[row_idx, col_idx] * zt) + 0.45
|
||||
print(f"[Terrain] flat_stairs level=1 spawn=({spawn_x:.1f},{spawn_y:.1f},{spawn_z:.2f})")
|
||||
|
||||
ckpt = torch.load(CKPT, map_location='cpu', weights_only=False)
|
||||
normalizer = ckpt['amp_normalizer']
|
||||
|
||||
actor = torch.nn.Sequential(
|
||||
torch.nn.Linear(45, 512), torch.nn.ELU(),
|
||||
torch.nn.Linear(512, 256), torch.nn.ELU(),
|
||||
torch.nn.Linear(256, 128), torch.nn.ELU(),
|
||||
torch.nn.Linear(128, 12)
|
||||
)
|
||||
sd = {k.replace('actor.',''): v for k,v in ckpt['model_state_dict'].items() if k.startswith('actor.')}
|
||||
actor.load_state_dict(sd)
|
||||
actor.eval()
|
||||
|
||||
DEFAULT = np.array([-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)
|
||||
LIN_VEL_SCALE, ANG_VEL_SCALE = 2.0, 0.25
|
||||
DOF_POS_SCALE, DOF_VEL_SCALE = 1.0, 0.05
|
||||
CMD_SCALE = np.array([2.0, 2.0, 0.25], dtype=np.float32)
|
||||
KP, KD, ACT_SCALE = 20.0, 0.5, 0.25
|
||||
|
||||
def get_gravity(q):
|
||||
qw,qx,qy,qz = q[3],q[0],q[1],q[2]
|
||||
return np.array([2*(-qz*qx+qw*qy), -2*(qz*qy+qw*qx), 1-2*(qw*qw+qz*qz)])
|
||||
|
||||
def compute_obs(d, last_action, cmd):
|
||||
obs = np.zeros(45, dtype=np.float32)
|
||||
obs[0:3] = d.qvel[3:6] * ANG_VEL_SCALE
|
||||
obs[3:6] = get_gravity(d.qpos[3:7])
|
||||
obs[6:9] = cmd * CMD_SCALE
|
||||
obs[9:21] = (d.qpos[7:19] - DEFAULT) * DOF_POS_SCALE
|
||||
obs[21:33] = d.qvel[6:18] * DOF_VEL_SCALE
|
||||
obs[33:45] = last_action
|
||||
return obs
|
||||
|
||||
d.qpos[7:19] = DEFAULT; d.qpos[0:3] = [spawn_x, spawn_y, spawn_z]; d.qvel[:] = 0
|
||||
mujoco.mj_forward(m, d)
|
||||
|
||||
view = viewer.launch_passive(m, d)
|
||||
cmd = np.array([0.8, 0.0, 0.0], dtype=np.float32)
|
||||
last_action = np.zeros(12, dtype=np.float32)
|
||||
|
||||
print("AMP Go1 | W/S前后 Q/E左右 A/D旋转 Space停 R重置 Esc退出")
|
||||
|
||||
while view.is_running():
|
||||
if view.is_running():
|
||||
obs = compute_obs(d, last_action, cmd)
|
||||
obs_t = torch.from_numpy(obs).unsqueeze(0).float()
|
||||
obs_norm = obs_t.clone()
|
||||
obs_norm[:, :43] = normalizer.normalize(obs_t[:, :43])
|
||||
with torch.no_grad():
|
||||
act = actor(obs_norm)[0].numpy()
|
||||
act = np.clip(act, -18, 18)
|
||||
last_action = act.copy()
|
||||
target = DEFAULT + act * ACT_SCALE
|
||||
d.ctrl[:] = KP * (target - d.qpos[7:19]) - KD * d.qvel[6:18]
|
||||
mujoco.mj_step(m, d)
|
||||
view.sync()
|
||||
time.sleep(0.001)
|
||||
107
scripts/test_abrupt_stop.py
Normal file
107
scripts/test_abrupt_stop.py
Normal file
@@ -0,0 +1,107 @@
|
||||
#!/usr/bin/env python3
|
||||
"""无头 MuJoCo 测试：前进 → 突然停止 → 观察姿态变化。"""
|
||||
import numpy as np
|
||||
import mujoco
|
||||
import onnxruntime as ort
|
||||
import os, sys, time
|
||||
|
||||
_PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
XML_DIR = os.path.join(_PROJECT, "motrix_envs", "src", "motrix_envs", "locomotion", "go1", "xmls")
|
||||
ONNX = os.path.join(_PROJECT, "runs/go1-dreamwaq-walk/rslrl/26-07-02_20-03-23-_36054_PPO/policy.onnx")
|
||||
|
||||
NUM_OBS = 45
|
||||
NUM_ACTIONS = 12
|
||||
HISTORY_LEN = 5
|
||||
ACTION_SCALE = 0.25
|
||||
KP = 28.0
|
||||
KD = 0.7
|
||||
DECIMATION = 4
|
||||
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], dtype=np.float32)
|
||||
|
||||
|
||||
def compute_obs(model, data, commands, last_action):
|
||||
obs = np.zeros(NUM_OBS, dtype=np.float32)
|
||||
g = None # no gyro sensor in headless
|
||||
obs[0:3] = (g if g is not None else data.qvel[3:6]) * 0.25
|
||||
grav_world = model.opt.gravity.copy()
|
||||
grav_world = grav_world / np.linalg.norm(grav_world)
|
||||
R = data.xmat[1].reshape(3, 3)
|
||||
obs[3:6] = (R.T @ grav_world).astype(np.float32)
|
||||
obs[6:9] = commands * np.array([2.0, 2.0, 0.25], dtype=np.float32)
|
||||
obs[9:21] = (data.qpos[7:19] - DEFAULT_ANGLES) * 1.0
|
||||
obs[21:33] = data.qvel[6:18] * 0.05
|
||||
obs[33:45] = last_action
|
||||
return obs
|
||||
|
||||
|
||||
def get_base_pose(data):
|
||||
"""返回 base 的 z 高度和 roll/pitch（度）。"""
|
||||
quat = data.xquat[1] # base body quaternion
|
||||
w, x, y, z = quat[0], quat[1], quat[2], quat[3]
|
||||
# roll, pitch from quaternion
|
||||
sinr_cosp = 2 * (w * x + y * z)
|
||||
cosr_cosp = 1 - 2 * (x * x + y * y)
|
||||
roll = np.arctan2(sinr_cosp, cosr_cosp)
|
||||
sinp = 2 * (w * y - z * x)
|
||||
pitch = np.arcsin(np.clip(sinp, -1, 1))
|
||||
return float(data.xpos[1, 2]), np.degrees(roll), np.degrees(pitch)
|
||||
|
||||
|
||||
def main():
|
||||
session = ort.InferenceSession(ONNX, providers=['CPUExecutionProvider'])
|
||||
xml_file = os.path.join(XML_DIR, "scene_dreamwaq_flat.xml")
|
||||
model = mujoco.MjModel.from_xml_path(xml_file)
|
||||
data = mujoco.MjData(model)
|
||||
|
||||
# 初始化
|
||||
data.qpos[7:19] = DEFAULT_ANGLES
|
||||
data.qpos[2] = 0.35
|
||||
mujoco.mj_forward(model, data)
|
||||
|
||||
history = np.zeros((1, HISTORY_LEN, NUM_OBS), dtype=np.float32)
|
||||
last_action = np.zeros(NUM_ACTIONS, dtype=np.float32)
|
||||
|
||||
print(f"{'Step':>5s} {'Cmd_vx':>7s} {'base_z':>8s} {'roll':>8s} {'pitch':>8s} {'speed_xy':>8s}")
|
||||
print("-" * 60)
|
||||
|
||||
for step in range(3000):
|
||||
# 命令：前 1500 步前进，之后突然停止
|
||||
if step < 1500:
|
||||
vx, vy, wz = 0.5, 0.0, 0.0 # 前进
|
||||
else:
|
||||
vx, vy, wz = 0.0, 0.0, 0.0 # 突然停止！
|
||||
|
||||
for _ in range(DECIMATION):
|
||||
mujoco.mj_step(model, data)
|
||||
|
||||
if step % DECIMATION == 0:
|
||||
cmd = np.array([vx, vy, wz], dtype=np.float32)
|
||||
obs = compute_obs(model, data, cmd, last_action)
|
||||
history = np.concatenate([history[:, 1:, :], obs.reshape(1, 1, -1)], axis=1)
|
||||
outputs = session.run(None, {
|
||||
'obs': obs.reshape(1, -1).astype(np.float32),
|
||||
'obs_history': history.reshape(1, -1).astype(np.float32),
|
||||
})
|
||||
action = outputs[0][0]
|
||||
last_action = action
|
||||
|
||||
target = DEFAULT_ANGLES + action * ACTION_SCALE
|
||||
data.ctrl[:] = KP * (target - data.qpos[7:19]) - KD * data.qvel[6:18]
|
||||
|
||||
# 每 100 步打印
|
||||
if step % 100 == 0:
|
||||
bz, roll, pitch = get_base_pose(data)
|
||||
speed_xy = np.linalg.norm(data.qvel[0:2])
|
||||
print(f"{step:5d} {vx:7.1f} {bz:8.3f} {roll:+8.1f} {pitch:+8.1f} {speed_xy:8.3f}")
|
||||
|
||||
# 最终状态
|
||||
bz, roll, pitch = get_base_pose(data)
|
||||
print(f"\n最终: base_z={bz:.3f} roll={roll:.1f}° pitch={pitch:.1f}°")
|
||||
if abs(roll) > 30 or abs(pitch) > 30:
|
||||
print("⚠ 机器人倾覆！突然停止导致翻跟头")
|
||||
else:
|
||||
print("✅ 机器人在突然停止后保持稳定")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user