feat: DreamWaQ full replication — env, terrain, CENet, PPO

2026-06-30 14:53:12 +08:00
parent c569f4d6d9
commit 422421d263
17 changed files with 2056 additions and 17 deletions
--- a/scripts/test_zscale_limit.py
+++ b/scripts/test_zscale_limit.py
@@ -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}")