feat: DreamWaQ full replication — env, terrain, CENet, PPO
This commit is contained in:
223
scripts/gen_dreamwaq_terrain.py
Normal file
223
scripts/gen_dreamwaq_terrain.py
Normal file
@@ -0,0 +1,223 @@
|
||||
#!/usr/bin/env python3
|
||||
"""生成 DreamWaQ 10×20 纯 hfield 地形——OpenCV 绘制。
|
||||
|
||||
5 种地形类型 × 10 难度,全部在单张 PNG 高度图中。
|
||||
楼梯用 1px riser 近垂直面(HS=0.05 时每像素 5cm)。
|
||||
|
||||
用法:
|
||||
uv run python3 scripts/gen_dreamwaq_terrain.py
|
||||
"""
|
||||
import cv2
|
||||
import numpy as np
|
||||
import os
|
||||
import argparse
|
||||
|
||||
# ═══ 参数 ═══
|
||||
HS = 0.05 # 水平分辨率 [m/px]
|
||||
VS = 0.005 # 垂直分辨率 [m/unit]
|
||||
CELL_M = 8.0
|
||||
NUM_ROWS = 10
|
||||
NUM_COLS = 20
|
||||
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)
|
||||
|
||||
CELL_PX = int(CELL_M / HS) # 160
|
||||
BORDER_PX = int(BORDER_M / HS) # 100
|
||||
PLATFORM_PX = int(PLATFORM_M / HS) # 60
|
||||
TOT_ROWS_PX = NUM_ROWS * CELL_PX + 2 * BORDER_PX # 1800
|
||||
TOT_COLS_PX = NUM_COLS * CELL_PX + 2 * BORDER_PX # 3400
|
||||
TOTAL_X = TOT_COLS_PX * HS
|
||||
TOTAL_Y = TOT_ROWS_PX * HS
|
||||
|
||||
|
||||
# ═══ 地形绘制 ═══
|
||||
|
||||
def draw_slope(canvas, x0, y0, difficulty, noise=False):
|
||||
"""平滑/粗糙斜坡——与上游 pyramid_sloped_terrain 对齐。
|
||||
|
||||
上游逻辑:先建金字塔(中心高→边缘低),再用平台边缘高度 clip 整个 terrain,
|
||||
形成与周围地形齐平的平台(而非硬清零到 0)。
|
||||
"""
|
||||
if difficulty <= 0:
|
||||
return
|
||||
slope = difficulty * _SLOPE_SCALE
|
||||
max_h = int(slope * (1.0 / VS) * (CELL_M / 2.0))
|
||||
if max_h <= 0:
|
||||
return
|
||||
cx, cy = CELL_PX // 2, CELL_PX // 2
|
||||
x = np.arange(0, CELL_PX)
|
||||
y = np.arange(0, CELL_PX)
|
||||
xx, yy = np.meshgrid(x, y, sparse=True)
|
||||
xx = (cx - np.abs(cx - xx)) / cx
|
||||
yy = (cy - np.abs(cy - yy)) / cy
|
||||
hf = (max_h * xx.reshape(CELL_PX, 1) * yy.reshape(1, CELL_PX)).astype(np.int32)
|
||||
p2 = PLATFORM_PX // 2
|
||||
# 上游 clip: 取平台边缘高度作为上下界
|
||||
edge_h = int(hf[cx - p2, cy - p2])
|
||||
lo = min(edge_h, 0)
|
||||
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)
|
||||
# 噪声也只在平台外
|
||||
n[cx - p2:cx + p2, cy - p2:cy + p2] = 0
|
||||
hf = np.clip(hf.astype(np.int32) + n, 0, 65535).astype(np.uint16)
|
||||
canvas[y0:y0 + CELL_PX, x0:x0 + CELL_PX] += hf
|
||||
|
||||
|
||||
def draw_pyramid_stairs(canvas, x0, y0, difficulty, concave=False):
|
||||
"""金字塔楼梯——OpenCV 同心矩形(近垂直 riser)。
|
||||
|
||||
每级台阶 2px 宽(10cm tread),高度缩放保持 z_scale < 0.54。
|
||||
"""
|
||||
if difficulty <= 0:
|
||||
return
|
||||
# 上游公式:step_height = 0.05 + 0.18 * difficulty [m]
|
||||
step_h_m = 0.05 + 0.18 * difficulty
|
||||
step_h = max(1, int(step_h_m / VS))
|
||||
cx = x0 + CELL_PX // 2
|
||||
cy = y0 + CELL_PX // 2
|
||||
p2 = PLATFORM_PX // 2
|
||||
|
||||
# 上游踏面 31cm → 6px (HS=0.05), 最多约 8 级
|
||||
tread_px = max(1, int(0.31 / HS))
|
||||
n_steps = min(8, (CELL_PX // 2 - p2) // tread_px)
|
||||
|
||||
if concave:
|
||||
base_h = step_h * n_steps
|
||||
cv2.rectangle(canvas, (x0, y0), (x0 + CELL_PX, y0 + CELL_PX), int(base_h), -1)
|
||||
for i in range(n_steps + 1):
|
||||
half = p2 + (n_steps - i) * tread_px
|
||||
h = int(base_h - step_h * i)
|
||||
cv2.rectangle(canvas, (cx - half, cy - half), (cx + half, cy + half), h, -1)
|
||||
else:
|
||||
for i in range(n_steps + 1):
|
||||
half = p2 + (n_steps - i) * tread_px
|
||||
h = int(step_h * i)
|
||||
cv2.rectangle(canvas, (cx - half, cy - half), (cx + half, cy + half), h, -1)
|
||||
|
||||
|
||||
def draw_obstacles(canvas, x0, y0, difficulty):
|
||||
"""离散障碍物(随机矩形块)。"""
|
||||
if difficulty <= 0:
|
||||
return
|
||||
max_h = int((0.05 + 0.2 * difficulty) / VS)
|
||||
if max_h <= 0:
|
||||
return
|
||||
p2 = PLATFORM_PX // 2
|
||||
# 上游: min_size=1.0m, max_size=2.0m, 20 个矩形
|
||||
min_sz = int(1.0 / HS); max_sz = int(2.0 / HS)
|
||||
for _ in range(20):
|
||||
w = np.random.randint(min_sz, max_sz + 1)
|
||||
ln = np.random.randint(min_sz, max_sz + 1)
|
||||
si = np.random.randint(0, CELL_PX - w)
|
||||
sj = np.random.randint(0, CELL_PX - ln)
|
||||
cv2.rectangle(canvas, (x0 + si, y0 + sj),
|
||||
(x0 + si + w, y0 + sj + ln),
|
||||
int(np.random.choice([max_h // 2, max_h])), -1)
|
||||
cx, cy = x0 + CELL_PX // 2, y0 + CELL_PX // 2
|
||||
cv2.rectangle(canvas, (cx - p2, cy - p2), (cx + p2, cy + p2), 0, -1)
|
||||
|
||||
|
||||
# ═══ 主流程 ═══
|
||||
|
||||
def main():
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--flat-only", action="store_true")
|
||||
p.add_argument("--max-level", type=int, default=None)
|
||||
args = p.parse_args()
|
||||
|
||||
max_row = NUM_ROWS if args.max_level is None else min(args.max_level + 1, NUM_ROWS)
|
||||
print(f"DreamWaQ 纯 hfield ({max_row}×{NUM_COLS}) {TOT_COLS_PX}×{TOT_ROWS_PX}px")
|
||||
|
||||
canvas = np.zeros((TOT_ROWS_PX, TOT_COLS_PX), dtype=np.uint16)
|
||||
|
||||
for row in range(max_row):
|
||||
difficulty = row / NUM_ROWS
|
||||
for col in range(NUM_COLS):
|
||||
if args.flat_only or difficulty == 0:
|
||||
continue
|
||||
x0 = BORDER_PX + col * CELL_PX
|
||||
y0 = BORDER_PX + row * CELL_PX
|
||||
choice = col / NUM_COLS + 0.001
|
||||
if choice < CUM[0]:
|
||||
draw_slope(canvas, x0, y0, difficulty)
|
||||
elif choice < CUM[1]:
|
||||
draw_slope(canvas, x0, y0, difficulty, noise=True)
|
||||
elif choice < CUM[2]:
|
||||
draw_pyramid_stairs(canvas, x0, y0, difficulty, concave=True)
|
||||
elif choice < CUM[3]:
|
||||
draw_pyramid_stairs(canvas, x0, y0, difficulty, concave=False)
|
||||
else:
|
||||
draw_obstacles(canvas, x0, y0, difficulty)
|
||||
|
||||
hf_m = canvas.astype(np.float32) * VS
|
||||
z_min, z_max = float(hf_m.min()), float(hf_m.max())
|
||||
z_range = max(z_max - z_min, 0.001)
|
||||
print(f" 高度范围: [{z_min:.3f}, {z_max:.3f}]m z_scale={z_range:.3f}")
|
||||
|
||||
if z_range > 0.54:
|
||||
print(f" ⚠ z_scale={z_range:.3f} > 0.54!")
|
||||
|
||||
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 = ((hf_m - z_min) / z_range * 65535.0).astype(np.uint16)
|
||||
cv2.imwrite(os.path.join(out_d, "dreamwaq_terrain.png"), png)
|
||||
|
||||
# XML
|
||||
xml = f"""<mujoco model="go1 dreamwaq terrain scene">
|
||||
<include file="go1_motor_actuator.xml" />
|
||||
<include file="materials.xml" />
|
||||
<statistic center="0 0 0.2" extent="5" meansize="0.04" />
|
||||
|
||||
<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" />
|
||||
<map force="0.01" />
|
||||
<scale forcewidth="0.3" contactwidth="0.5" contactheight="0.2" />
|
||||
<quality shadowsize="8192" />
|
||||
</visual>
|
||||
|
||||
<asset>
|
||||
<hfield name="dreamwaq_terrain"
|
||||
file="assets/dreamwaq_terrain.png"
|
||||
size="{TOTAL_X / 2:.1f} {TOTAL_Y / 2:.1f} {z_range:.3f} {max(z_min, 0.001):.3f}" />
|
||||
</asset>
|
||||
|
||||
<worldbody>
|
||||
<light pos="0 0 4" dir="0 0 -1" directional="true" />
|
||||
<geom name="floor" pos="0 0 0" type="hfield" hfield="dreamwaq_terrain"
|
||||
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>
|
||||
"""
|
||||
xml_dir = os.path.join(os.path.dirname(__file__), "..",
|
||||
"motrix_envs", "src", "motrix_envs",
|
||||
"locomotion", "go1", "xmls")
|
||||
with open(os.path.join(xml_dir, "scene_dreamwaq_terrain.xml"), "w") as f:
|
||||
f.write(xml)
|
||||
|
||||
half_x = TOTAL_X / 2
|
||||
half_y = TOTAL_Y / 2
|
||||
print(f" XML: size=\"{half_x:.1f} {half_y:.1f} {z_range:.3f} {max(z_min, 0.001):.3f}\"")
|
||||
print(f" 楼梯: 1px tread (5cm), 1px riser → 近垂直面")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
np.random.seed(42)
|
||||
main()
|
||||
86
scripts/test_mesh_collision.py
Normal file
86
scripts/test_mesh_collision.py
Normal file
@@ -0,0 +1,86 @@
|
||||
#!/usr/bin/env python3
|
||||
"""最小 mesh 碰撞测试——验证 MotrixSim 的 OBJ mesh 是否支持碰撞。"""
|
||||
|
||||
import os, numpy as np
|
||||
|
||||
xml_dir = '/home/8x54zj-m/MotrixLab/motrix_envs/src/motrix_envs/locomotion/go1/xmls'
|
||||
assets_dir = os.path.join(xml_dir, 'assets', 'tmp_test')
|
||||
os.makedirs(assets_dir, exist_ok=True)
|
||||
obj_path = os.path.join(assets_dir, 'test_box.obj')
|
||||
|
||||
# 封闭 box OBJ (1x1x0.1m),带法线
|
||||
with open(obj_path, 'w') as f:
|
||||
f.write("""# closed box
|
||||
v -0.5 -0.5 0.0
|
||||
v 0.5 -0.5 0.0
|
||||
v 0.5 0.5 0.0
|
||||
v -0.5 0.5 0.0
|
||||
v -0.5 -0.5 0.1
|
||||
v 0.5 -0.5 0.1
|
||||
v 0.5 0.5 0.1
|
||||
v -0.5 0.5 0.1
|
||||
f 1 3 2
|
||||
f 1 4 3
|
||||
f 5 6 7
|
||||
f 5 7 8
|
||||
f 1 5 6
|
||||
f 1 6 2
|
||||
f 2 6 7
|
||||
f 2 7 3
|
||||
f 3 7 8
|
||||
f 3 8 4
|
||||
f 4 8 5
|
||||
f 4 5 1
|
||||
""")
|
||||
|
||||
# 生成测试 XML
|
||||
xml_path = os.path.join(xml_dir, 'scene_test_mesh.xml')
|
||||
with open(xml_path, 'w') as f:
|
||||
f.write("""<mujoco model="test mesh">
|
||||
<include file="go1_motor_actuator.xml" />
|
||||
<include file="materials.xml" />
|
||||
<statistic center="0 0 0.3" extent="1" />
|
||||
<visual>
|
||||
<headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0" />
|
||||
</visual>
|
||||
<asset>
|
||||
<mesh name="test_box" file="assets/tmp_test/test_box.obj" />
|
||||
</asset>
|
||||
<worldbody>
|
||||
<light pos="0 0 2" dir="0 0 -1" directional="true" />
|
||||
<!-- 无 plane floor -- 纯 mesh 碰撞测试 -->
|
||||
<geom name="box_mesh" type="mesh" mesh="test_box" pos="0 0 0.15"
|
||||
contype="1" conaffinity="1" rgba="0.8 0.3 0.3 1" friction="0.8 0.3 0.3"/>
|
||||
</worldbody>
|
||||
</mujoco>
|
||||
""")
|
||||
|
||||
import motrixsim as mtx
|
||||
model = mtx.load_model(xml_path)
|
||||
print('加载成功')
|
||||
|
||||
data = mtx.SceneData(model, batch=[1])
|
||||
data.reset(model)
|
||||
body = model.get_body(0)
|
||||
|
||||
init_pos = model.compute_init_dof_pos().reshape(1, -1)
|
||||
init_pos[0, 0:2] = 0.0 # 在 box 正上方
|
||||
init_pos[0, 2] = 0.8 # 从 0.8m 自由落体 (box 顶面在 z=0.25)
|
||||
data.set_dof_pos(init_pos, model)
|
||||
model.forward_kinematic(data)
|
||||
|
||||
print('自由落体到 box mesh (顶部 z=0.25):')
|
||||
for i in range(80):
|
||||
model.step(data)
|
||||
bz = body.get_pose(data)[0, 2]
|
||||
if i < 15 or i % 15 == 0:
|
||||
print(f' 步{i+1}: base_z={bz:.4f}')
|
||||
|
||||
bz_final = body.get_pose(data)[0, 2]
|
||||
print(f'\n最终 base_z={bz_final:.4f}')
|
||||
if bz_final > 0.45:
|
||||
print('✅ mesh 碰撞正常!机器人站在 box 上')
|
||||
elif bz_final < 0.10:
|
||||
print('❌ mesh 碰撞不工作!机器人穿透 box 坠入深渊')
|
||||
else:
|
||||
print(f'⚠ 不确定: base_z={bz_final:.4f}')
|
||||
148
scripts/test_zscale_limit.py
Normal file
148
scripts/test_zscale_limit.py
Normal file
@@ -0,0 +1,148 @@
|
||||
#!/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}")
|
||||
269
scripts/view_dreamwaq.py
Normal file
269
scripts/view_dreamwaq.py
Normal file
@@ -0,0 +1,269 @@
|
||||
#!/usr/bin/env python3
|
||||
"""DreamWaQ 地形可视化。
|
||||
|
||||
键盘:
|
||||
R=重置 H=高度采样点 T=遍历出生点调试 Esc=退出
|
||||
|
||||
用法:
|
||||
uv run scripts/view_dreamwaq.py # 金字塔地形
|
||||
uv run scripts/view_dreamwaq.py --flat --num-envs 1
|
||||
"""
|
||||
import argparse, os, sys, time
|
||||
import numpy as np
|
||||
|
||||
os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
|
||||
os.environ.setdefault("JAX_PLATFORMS", "cpu")
|
||||
|
||||
if "--flat" in sys.argv:
|
||||
os.environ["DREAMWAQ_TERRAIN"] = "flat"
|
||||
elif "--flat-stairs" in sys.argv:
|
||||
os.environ["DREAMWAQ_TERRAIN"] = "flat_stairs"
|
||||
elif "--stairs" in sys.argv:
|
||||
os.environ["DREAMWAQ_TERRAIN"] = "stairs"
|
||||
else:
|
||||
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
|
||||
from motrix_envs.np.renderer import NpRenderer
|
||||
from motrix_envs.math import quaternion
|
||||
from motrixsim.render import RenderClosedError
|
||||
|
||||
# 地形类型名称(与 gen_dreamwaq_terrain.py 的 PROPORTIONS 对应)
|
||||
PROPORTIONS = [0.1, 0.1, 0.35, 0.35, 0.1]
|
||||
CUM = [sum(PROPORTIONS[:i + 1]) for i in range(len(PROPORTIONS))]
|
||||
TYPE_NAMES = ["平滑斜坡", "粗糙斜坡", "下行楼梯", "上行楼梯", "离散障碍"]
|
||||
|
||||
|
||||
def _cell_origin(row, col, border_m=5.0, cell_m=8.0, num_rows=10, num_cols=20):
|
||||
"""计算 cell (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 _type_name(col):
|
||||
"""根据列索引返回地形类型名称。"""
|
||||
choice = col / 20 + 0.001
|
||||
for i, cum in enumerate(CUM):
|
||||
if choice < cum:
|
||||
return TYPE_NAMES[i]
|
||||
return TYPE_NAMES[-1]
|
||||
|
||||
|
||||
def _spawn_at(env, cx, cy, spawn_z=None):
|
||||
"""在指定世界坐标 spawn 单个机器人。"""
|
||||
state = env._state
|
||||
data = state.data
|
||||
init_pos = env._init_dof_pos.copy().reshape(1, -1)
|
||||
init_pos[0, 0] = cx
|
||||
init_pos[0, 1] = cy
|
||||
# 计算地形高度
|
||||
terrain_z = float(env._sample_terrain_height(
|
||||
np.array([[cx, cy]], dtype=np.float32), radius=0.35)[0])
|
||||
if spawn_z is None:
|
||||
spawn_z = terrain_z + 0.45 # 默认 clearance
|
||||
init_pos[0, 2] = spawn_z
|
||||
data.reset(env._model)
|
||||
data.set_dof_pos(init_pos, env._model)
|
||||
env._model.forward_kinematic(data)
|
||||
# 重置 info
|
||||
state.info["commands"][0] = np.array([0.0, 0.0, 0.0], dtype=np.float32)
|
||||
state.info["steps"][0] = 0
|
||||
state.info["obs_history"][0] = 0.0
|
||||
return terrain_z, spawn_z
|
||||
|
||||
|
||||
def main():
|
||||
p = argparse.ArgumentParser(description="DreamWaQ 地形可视化")
|
||||
p.add_argument("--num-envs", type=int, default=1)
|
||||
p.add_argument("--flat", action="store_true")
|
||||
p.add_argument("--flat-stairs", action="store_true")
|
||||
p.add_argument("--stairs", action="store_true")
|
||||
p.add_argument("--level", type=int, default=None)
|
||||
p.add_argument("--no-stand", action="store_true")
|
||||
p.add_argument("--vx", type=float, default=0.5)
|
||||
args = p.parse_args()
|
||||
|
||||
env = env_registry.make("go1-dreamwaq-walk", num_envs=max(args.num_envs, 1))
|
||||
if args.level is not None:
|
||||
env._force_level = args.level
|
||||
|
||||
env.init_state()
|
||||
n = env._num_envs
|
||||
cmd = np.array([args.vx, 0.0, 0.0], dtype=np.float32)
|
||||
|
||||
try:
|
||||
renderer = NpRenderer(env)
|
||||
except Exception as e:
|
||||
print(f"[ERROR] 渲染器创建失败: {e}")
|
||||
renderer = None
|
||||
|
||||
terrain_name = os.environ.get("DREAMWAQ_TERRAIN", "pyramid")
|
||||
print(f"[View] {n} 机器人 | 地形={terrain_name}")
|
||||
print(f"[View] R=重置 H=高度点 T=遍历出生点 Esc=退出")
|
||||
|
||||
show_heights = False
|
||||
traverse_mode = False
|
||||
traverse_row = 0
|
||||
traverse_col = 0
|
||||
traverse_pending = False # 刚切换 cell, 等待稳定
|
||||
traverse_settle = 0
|
||||
step_count = 0
|
||||
|
||||
# 地形信息
|
||||
num_rows = env._num_rows
|
||||
num_cols = env._num_cols
|
||||
cell_m = env._cell_size
|
||||
|
||||
def enter_traverse():
|
||||
nonlocal traverse_mode, traverse_row, traverse_col, traverse_pending, traverse_settle
|
||||
traverse_mode = True
|
||||
traverse_row = 0
|
||||
traverse_col = 0
|
||||
traverse_pending = True
|
||||
traverse_settle = 0
|
||||
print(f"\n[T] 遍历模式: {num_rows}行 × {num_cols}列")
|
||||
print(f"[T] 按 T 前进, R 退出遍历\n")
|
||||
|
||||
def exit_traverse():
|
||||
nonlocal traverse_mode, traverse_pending
|
||||
traverse_mode = False
|
||||
traverse_pending = False
|
||||
env.init_state()
|
||||
print("[T] 退出遍历模式\n")
|
||||
|
||||
def advance_traverse():
|
||||
nonlocal traverse_row, traverse_col, traverse_pending, traverse_settle
|
||||
traverse_col += 1
|
||||
if traverse_col >= num_cols:
|
||||
traverse_col = 0
|
||||
traverse_row += 1
|
||||
if traverse_row >= num_rows:
|
||||
print("[T] 遍历完成! 按 R 退出")
|
||||
traverse_row = num_rows - 1
|
||||
traverse_col = num_cols - 1
|
||||
return False
|
||||
traverse_pending = True
|
||||
traverse_settle = 0
|
||||
return True
|
||||
|
||||
def do_traverse_spawn():
|
||||
"""在当前位置 spawn 并打印信息。"""
|
||||
nonlocal traverse_settle
|
||||
cx, cy = _cell_origin(traverse_row, traverse_col)
|
||||
tname = _type_name(traverse_col)
|
||||
terrain_z, spawn_z = _spawn_at(env, cx, cy)
|
||||
|
||||
# 标记 spawn 点(绿色球)
|
||||
if renderer is not None:
|
||||
g = renderer._render.gizmos
|
||||
g.draw_sphere(0.15, (np.float32(cx), np.float32(cy),
|
||||
np.float32(spawn_z)))
|
||||
|
||||
# 让机器人稳定几步
|
||||
for _ in range(30):
|
||||
env.step(np.zeros((1, 12), dtype=np.float32))
|
||||
if renderer is not None:
|
||||
renderer.render()
|
||||
time.sleep(0.005)
|
||||
|
||||
base_z = env._body.get_pose(env._state.data)[0, 2]
|
||||
contacts = env._state.info.get("contacts", np.zeros(4))
|
||||
cf = env._state.info.get("privileged_obs",
|
||||
np.zeros((1, 247)))[0, 45:57]
|
||||
total_cf = np.sum(np.abs(cf))
|
||||
|
||||
print(f" r{traverse_row}c{traverse_col:02d} {tname:6s} "
|
||||
f"origin=({cx:+.0f},{cy:+.0f}) "
|
||||
f"terrain_z={terrain_z:.3f} spawn_z={spawn_z:.3f} "
|
||||
f"base_z={base_z:.3f} cf={total_cf:.0f}N "
|
||||
f"feet={contacts.astype(int).tolist()}")
|
||||
|
||||
traverse_settle = 30
|
||||
|
||||
try:
|
||||
while True:
|
||||
# ── 键盘 ──
|
||||
if renderer is not None:
|
||||
try:
|
||||
inp = renderer._render.input
|
||||
if inp.is_key_just_pressed("r"):
|
||||
if traverse_mode:
|
||||
exit_traverse()
|
||||
else:
|
||||
env.init_state()
|
||||
step_count = 0
|
||||
print("[R] 重置")
|
||||
if inp.is_key_just_pressed("h"):
|
||||
show_heights = not show_heights
|
||||
print(f"[H] 高度点: {'开' if show_heights else '关'}")
|
||||
if inp.is_key_just_pressed("t"):
|
||||
if traverse_mode:
|
||||
advance_traverse()
|
||||
else:
|
||||
enter_traverse()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ── 遍历模式 ──
|
||||
if traverse_mode and traverse_pending:
|
||||
do_traverse_spawn()
|
||||
traverse_pending = False
|
||||
|
||||
# ── 正常模式动作 ──
|
||||
if not traverse_mode or traverse_settle > 0:
|
||||
if args.no_stand:
|
||||
env._state.info["commands"][:] = cmd
|
||||
else:
|
||||
act = 0.3 * (np.random.rand(n, 12).astype(np.float32) - 0.5)
|
||||
env.step(act)
|
||||
if traverse_settle > 0:
|
||||
traverse_settle -= 1
|
||||
|
||||
if step_count % 100 == 0 and step_count > 0:
|
||||
bz = env._body.get_pose(env._state.data)[0, 2]
|
||||
print(f"[{step_count}] base_z={bz:.3f}")
|
||||
|
||||
# ── 高度点 ──
|
||||
if show_heights and renderer is not None:
|
||||
state = env._state
|
||||
pose = env._body.get_pose(state.data)
|
||||
bp = pose[0, :3]
|
||||
yaw = quaternion.get_yaw(pose[0:1, 3:7])[0]
|
||||
cos_y, sin_y = np.cos(yaw), np.sin(yaw)
|
||||
for gy in env._hy:
|
||||
for gx in env._hx:
|
||||
wx = bp[0] + cos_y * gx - sin_y * gy
|
||||
wy = bp[1] + sin_y * gx + cos_y * gy
|
||||
try:
|
||||
wz = float(env._sample_terrain_height(
|
||||
np.array([[wx, wy]]))[0])
|
||||
g = renderer._render.gizmos
|
||||
g.draw_sphere(0.02, (np.float32(wx), np.float32(wy),
|
||||
np.float32(wz)))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if renderer is not None:
|
||||
renderer.render()
|
||||
time.sleep(0.01)
|
||||
|
||||
step_count += 1
|
||||
|
||||
except (KeyboardInterrupt, RenderClosedError):
|
||||
pass
|
||||
try:
|
||||
if renderer is not None:
|
||||
renderer.close()
|
||||
except Exception:
|
||||
pass
|
||||
print("[View] 结束")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user