[train] 更新新MJCF与第一版完整训练框架

2026-07-27 12:31:21 +08:00
parent dfd646cf47
commit b08956aec7
116 changed files with 7649 additions and 6760 deletions
@@ -0,0 +1,115 @@
+"""Interactive single-patch terrain explorer (Viser + MuJoCo MjSpec).
+
+Run with:
+  uv run python scripts/tools/terrain_explorer.py
+  uv run python scripts/tools/terrain_explorer.py --port 8081
+
+Then open the printed URL (default http://localhost:8080).
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+
+import mujoco
+import numpy as np
+import viser
+from mjviser.conversions import merge_geoms
+
+from mjlab.terrains.config import ALL_TERRAIN_PRESETS
+from mjlab.terrains.terrain_generator import TerrainGenerator, TerrainGeneratorCfg
+
+PATCH_SIZE = (8.0, 8.0)
+
+
+# Per-preset overrides applied when building in the explorer (e.g. to surface
+# difficulty-driven behavior that is off by default).
+_PRESET_OVERRIDES: dict[str, dict] = {
+  "random_rough": {"scale_with_difficulty": True},
+}
+
+
+def _build_terrain_mesh(preset_name: str, difficulty: float, seed: int):
+  """Generate a single terrain patch and return a merged trimesh (or raise)."""
+  preset_fn = ALL_TERRAIN_PRESETS[preset_name]
+  overrides = _PRESET_OVERRIDES.get(preset_name, {})
+  generator_cfg = TerrainGeneratorCfg(
+    seed=seed,
+    size=PATCH_SIZE,
+    num_rows=1,
+    num_cols=1,
+    border_width=0.0,
+    curriculum=False,
+    # A degenerate range pins the single patch to exactly this difficulty.
+    difficulty_range=(difficulty, difficulty),
+    color_scheme="height",
+    sub_terrains={preset_name: preset_fn(proportion=1.0, **overrides)},
+  )
+  generator = TerrainGenerator(generator_cfg)
+  spec = mujoco.MjSpec()
+  generator.compile(spec)
+  model = spec.compile()
+
+  terrain_body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "terrain")
+  geom_ids = [i for i in range(model.ngeom) if model.geom_bodyid[i] == terrain_body_id]
+  return merge_geoms(model, geom_ids)
+
+
+def main() -> None:
+  parser = argparse.ArgumentParser(description=__doc__)
+  parser.add_argument(
+    "--port", type=int, default=8080, help="Port for the viser server."
+  )
+  args = parser.parse_args()
+
+  server = viser.ViserServer(port=args.port)
+  preset_names = sorted(ALL_TERRAIN_PRESETS)
+
+  terrain_dropdown = server.gui.add_dropdown(
+    "Terrain", options=preset_names, initial_value=preset_names[0]
+  )
+  difficulty_slider = server.gui.add_slider(
+    "Difficulty", min=0.0, max=1.0, step=0.01, initial_value=0.0
+  )
+  seed_input = server.gui.add_number("Seed", initial_value=42, step=1)
+  status = server.gui.add_markdown("**Status:** ready")
+
+  handle: viser.SceneNodeHandle | None = None
+
+  def update() -> None:
+    nonlocal handle
+    name = terrain_dropdown.value
+    difficulty = float(difficulty_slider.value)
+    seed = int(seed_input.value)
+    status.content = f"**Status:** building `{name}` at difficulty {difficulty:.2f}..."
+    try:
+      mesh = _build_terrain_mesh(name, difficulty, seed)
+    except Exception as e:  # noqa: BLE001 - surface any generation failure in the UI.
+      status.content = f"**Error:** {type(e).__name__}: {e}"
+      print(f"Failed to build {name} at difficulty {difficulty}: {e}")
+      return
+    if handle is not None:
+      handle.remove()
+    handle = server.scene.add_mesh_trimesh("/terrain", mesh)
+    status.content = (
+      f"**Loaded** `{name}` at difficulty {difficulty:.2f} ({len(mesh.faces):,} faces)"
+    )
+
+  terrain_dropdown.on_update(lambda _: update())
+  difficulty_slider.on_update(lambda _: update())
+  seed_input.on_update(lambda _: update())
+
+  # Top-down-ish initial camera.
+  @server.on_client_connect
+  def _(client: viser.ClientHandle) -> None:
+    client.camera.position = np.array([10.0, 10.0, 8.0])
+    client.camera.look_at = np.array([0.0, 0.0, 0.0])
+
+  update()
+  while True:
+    time.sleep(1.0)
+
+
+if __name__ == "__main__":
+  main()