149 lines
5.6 KiB
Python
149 lines
5.6 KiB
Python
#!/usr/bin/env python3
|
||
"""测试 MotrixSim hfield 的 z_scale 上限——自己动手测，不信文档。"""
|
||
import os, sys, numpy as np, time
|
||
|
||
os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
|
||
os.environ.setdefault("JAX_PLATFORMS", "cpu")
|
||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
|
||
import motrix_envs.locomotion.go1.dreamwaq # noqa
|
||
from motrix_envs import registry as env_registry
|
||
|
||
NUM_ENVS = 256
|
||
TEST_STEPS = 100
|
||
os.environ["DREAMWAQ_TERRAIN"] = "flat" # 用 flat 场景，手动覆盖 hfield 参数
|
||
|
||
|
||
def _build_custom_hfield(x_radius, y_radius, z_scale, z_base=0.001):
|
||
"""构建自定义 hfield 描述字符串，用于覆盖 XML 中的 hfield 参数。"""
|
||
import tempfile, cv2
|
||
# 生成一个纯斜坡的 hfield PNG 用于测试
|
||
nx, ny = 200, 200
|
||
canvas = np.zeros((ny, nx), dtype=np.uint16)
|
||
# 从左上到右下的斜坡：高度从 0 到 z_scale
|
||
for i in range(ny):
|
||
for j in range(nx):
|
||
# 对角线斜坡，最高点在右下角
|
||
h = int((i + j) / (nx + ny) * 65535)
|
||
canvas[i, j] = h
|
||
# 中间 1/3 区域做平台（平坦）
|
||
cx, cy = nx // 2, ny // 2
|
||
p = nx // 6
|
||
canvas[cy - p:cy + p, cx - p:cx + p] = 0
|
||
|
||
out_d = os.path.join(os.path.dirname(__file__), "..",
|
||
"motrix_envs", "src", "motrix_envs",
|
||
"locomotion", "go1", "xmls", "assets")
|
||
os.makedirs(out_d, exist_ok=True)
|
||
png_path = os.path.join(out_d, "zscale_test.png")
|
||
cv2.imwrite(png_path, canvas)
|
||
|
||
return png_path, (x_radius, y_radius, z_scale, z_base)
|
||
|
||
|
||
def test_z_scale(z_scale, n_envs=NUM_ENVS, n_steps=TEST_STEPS):
|
||
"""测试给定 z_scale 下的物理稳定性。"""
|
||
import tempfile, cv2
|
||
|
||
# 生成测试 hfield
|
||
out_d = os.path.join(os.path.dirname(__file__), "..",
|
||
"motrix_envs", "src", "motrix_envs",
|
||
"locomotion", "go1", "xmls", "assets")
|
||
os.makedirs(out_d, exist_ok=True)
|
||
|
||
# 简单斜坡地形
|
||
nx, ny = 40, 40 # 小尺寸快速生成
|
||
canvas = np.zeros((ny, nx), dtype=np.uint16)
|
||
for i in range(ny):
|
||
for j in range(nx):
|
||
h = int((i / ny) * 65535) # y 方向斜坡
|
||
canvas[i, j] = h
|
||
# 中央平台
|
||
p = nx // 6
|
||
canvas[ny // 2 - p:ny // 2 + p, nx // 2 - p:nx // 2 + p] = 0
|
||
|
||
png_path = os.path.join(out_d, "zscale_test.png")
|
||
cv2.imwrite(png_path, canvas)
|
||
|
||
total_x = nx * 0.1 # HS=0.1, 粗略
|
||
total_y = ny * 0.1
|
||
|
||
# 构建 scene XML
|
||
xml_path = os.path.join(os.path.dirname(__file__), "..",
|
||
"motrix_envs", "src", "motrix_envs",
|
||
"locomotion", "go1", "xmls", "scene_zscale_test.xml")
|
||
xml = f"""<mujoco model="zscale test">
|
||
<include file="go1_motor_actuator.xml" />
|
||
<include file="materials.xml" />
|
||
<statistic center="0 0 0.3" extent="2" />
|
||
<visual>
|
||
<headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0" />
|
||
<rgba haze="0.15 0.25 0.35 1" />
|
||
<global azimuth="120" elevation="-20" />
|
||
</visual>
|
||
<asset>
|
||
<hfield name="test_hf" file="assets/zscale_test.png"
|
||
size="{total_x/2:.1f} {total_y/2:.1f} {z_scale:.3f} 0.001" />
|
||
</asset>
|
||
<worldbody>
|
||
<light pos="0 0 2" dir="0 0 -1" directional="true" />
|
||
<geom name="floor" pos="0 0 0" type="hfield" hfield="test_hf"
|
||
material="motphys-ground" contype="1" conaffinity="0" priority="1" friction="0.6" />
|
||
</worldbody>
|
||
<sensor>
|
||
<contact name="FR_foot_contact" geom2="FR_foot" geom1="floor" data="force" num="1" />
|
||
<contact name="FL_foot_contact" geom2="FL_foot" geom1="floor" data="force" num="1" />
|
||
<contact name="RR_foot_contact" geom2="RR_foot" geom1="floor" data="force" num="1" />
|
||
<contact name="RL_foot_contact" geom2="RL_foot" geom1="floor" data="force" num="1" />
|
||
</sensor>
|
||
</mujoco>"""
|
||
with open(xml_path, "w") as f:
|
||
f.write(xml)
|
||
|
||
import motrixsim as mtx
|
||
t0 = time.time()
|
||
|
||
try:
|
||
model = mtx.load_model(xml_path)
|
||
data = mtx.SceneData(model, batch=[n_envs])
|
||
data.reset(model)
|
||
body = model.get_body(0)
|
||
|
||
# 随机初始位置（在平台上）
|
||
init_pos = model.compute_init_dof_pos()
|
||
init_pos = np.tile(init_pos, (n_envs, 1))
|
||
init_pos[:, 0] += np.random.uniform(-0.5, 0.5, n_envs)
|
||
init_pos[:, 1] += np.random.uniform(-0.5, 0.5, n_envs)
|
||
init_pos[:, 2] = 0.5 # 从 0.5m 掉落
|
||
data.set_dof_pos(init_pos.astype(np.float32), model)
|
||
model.forward_kinematic(data)
|
||
|
||
heights = np.zeros((n_steps, n_envs), dtype=np.float32)
|
||
fall_count = 0
|
||
for step in range(n_steps):
|
||
model.step(data)
|
||
h = body.get_pose(data)[:, 2]
|
||
heights[step] = h
|
||
# 摔倒检测：base_z < 0.15 (趴了)
|
||
fall_count += np.sum(h < 0.15)
|
||
|
||
mean_h = float(np.mean(heights[-20:])) # 最后 20 步平均
|
||
total_falls = fall_count
|
||
dt = time.time() - t0
|
||
return mean_h, total_falls, dt
|
||
except Exception as e:
|
||
return None, str(e), 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
print(f"{'z_scale':>8s} {'mean_base_z':>12s} {'falls':>8s} {'time':>8s} verdict")
|
||
print("-" * 65)
|
||
|
||
for zs in [0.3, 0.5, 0.54, 0.8, 1.0, 1.5, 2.0, 3.0]:
|
||
mean_h, falls, dt = test_z_scale(zs, n_envs=64, n_steps=50)
|
||
if mean_h is None:
|
||
print(f"{zs:8.3f} {'ERROR':>12s} {str(falls)[:20]:>8s}")
|
||
continue
|
||
ok = "✅ 稳定" if mean_h > 0.25 and falls < 10 else "⚠ 不稳" if mean_h > 0.15 else "❌ 崩溃"
|
||
print(f"{zs:8.3f} {mean_h:12.4f} {falls:8d} {dt:7.1f}s {ok}")
|