459 lines
19 KiB
Python
459 lines
19 KiB
Python
#!/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()
|