[sim] 整理后期Sim2Sim与比赛Rough策略

2026-07-27 13:20:41 +08:00
parent c05c1cb162
commit 4ee4af028c
20 changed files with 5161 additions and 15 deletions
@@ -0,0 +1,58 @@
+# 后期 Sim2Sim 工具
+
+本目录保存比赛训练架构之后形成的 MuJoCo 策略验证工具。`v0.8.0` 在早期 Sim2Sim 基础上增加 ONNX 策略加载、IK 参数扫描、纯 IK 绕桩验证和批量路线检查；训练任务与 MJCF 不在本阶段修改。
+
+## 主要入口
+
+- `nav_sim2sim.py`：Pygame 面板与 MuJoCo 多任务导航，Rough 策略优先加载根目录的 `model_6800.onnx`。
+- `sim2sim.py`：较轻量的键盘控制与策略回放入口，优先加载 `model_6800.onnx`，缺失时回退到早期 `model_rough.pt`。
+- `ik_slalom_sim2sim.py`：不依赖 RL 策略的 IK、差速轮、路径跟踪和绕桩测试。
+- `ik_compensation_sweep.py`：批量扫描 IK 补偿参数并输出排序结果。
+- `nav_route_sim2sim_check.py`：使用 ONNX 策略批量检查内置任务或外部航点路线。
+- `export_onnx.py`：将兼容的 PyTorch actor checkpoint 导出并核对为 ONNX。
+- `interface/mujoco_io.py`：MuJoCo 模型、传感器和执行器接口。
+- `policy/policy_runner.py`：PT/ONNX 策略加载与历史观测缓存。
+
+## 环境
+
+主训练环境继续由根目录的 `uv.lock` 管理。后期 Sim2Sim 新增的 Pygame 与 ONNX Runtime 单独记录在 `sim2sim/requirements.txt`，运行时叠加，避免重新解析时改变已归档的 MuJoCo nightly 版本：
+
+```powershell
+uv run --with-requirements .\sim2sim\requirements.txt python .\sim2sim\nav_sim2sim.py
+```
+
+训练工程提供 MuJoCo、NumPy、PyTorch、Matplotlib 和 `pynput`；专用 requirements 显式补充 Pygame 与 ONNX Runtime。下面其他命令同样使用 `--with-requirements .\sim2sim\requirements.txt`。
+
+## 常用命令
+
+```powershell
+# 比赛 Rough 策略交互回放
+uv run --with-requirements .\sim2sim\requirements.txt python .\sim2sim\nav_sim2sim.py
+
+# 轻量策略回放
+uv run --with-requirements .\sim2sim\requirements.txt python .\sim2sim\sim2sim.py
+
+# 纯 IK 绕桩验证
+uv run --with-requirements .\sim2sim\requirements.txt python .\sim2sim\ik_slalom_sim2sim.py --test slalom
+
+# IK 补偿参数扫描
+uv run --with-requirements .\sim2sim\requirements.txt python .\sim2sim\ik_compensation_sweep.py --top 12
+
+# 使用内置绕桩任务做批量 Sim2Sim 路线检查
+uv run --with-requirements .\sim2sim\requirements.txt python .\sim2sim\nav_route_sim2sim_check.py `
+  --terrain-xml .\sim2sim\terrain\scene_terrain.xml `
+  --mission slalom `
+  --onnx .\model_6800.onnx
+
+# 导出早期参考 PT 权重；也可用 --pt-path 指定其他 checkpoint
+uv run --with-requirements .\sim2sim\requirements.txt python .\sim2sim\export_onnx.py
+```
+
+## 模型与边界
+
+- `../model_6800.onnx` 是 `last_not_slalom_1050` 最终真机工程使用的比赛 Rough 策略，SHA-256 为 `3C994BDD3434AD15770A52AC0E8D229F502F00D6511CDD42C2E2C742301AEF13`。
+- `../model_rough.pt` 是较早阶段的参考 checkpoint，两者不是同一版本的权重。
+- Crawl 模型未在本阶段归档；需要 Crawl 策略的入口会查找 `model_crawl.onnx` 或 `model_crawl.pt`。
+- `nav_route_sim2sim_check.py` 依赖 `../tools/nav_tools/route_safety_check.py` 的航点和避障几何定义。
+
+运行时生成的日志、临时 XML 和 `route_check_runs/` 不纳入版本库。源目录中的大量路线试验结果也未复制；它们包含重复轨迹和本机绝对路径，不属于可复用程序源码。
@@ -0,0 +1,110 @@
+import argparse
+import torch
+import torch.nn as nn
+import numpy as np
+from pathlib import Path
+import onnxruntime as ort
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+
+class PolicyMLP(nn.Module):
+    def __init__(self, obs_dim=53, action_dim=16):
+        super().__init__()
+        self.register_buffer("obs_mean", torch.zeros(obs_dim))
+        self.register_buffer("obs_std", torch.ones(obs_dim))
+        self.net = nn.Sequential(
+            nn.Linear(obs_dim, 512), nn.ELU(),
+            nn.Linear(512, 256), nn.ELU(),
+            nn.Linear(256, 128), nn.ELU(),
+            nn.Linear(128, action_dim),
+        )
+
+    def forward(self, x):
+        x = (x - self.obs_mean) / torch.clamp(self.obs_std, min=1e-6)
+        return self.net(x)
+
+def load_policy(model_path, device):
+    ckpt = torch.load(model_path, map_location=device, weights_only=False)
+    state_dict = ckpt["actor_state_dict"]
+    weight_key = "mlp.0.weight" if "mlp.0.weight" in state_dict else "net.0.weight"
+    obs_dim = state_dict[weight_key].shape[1]
+
+    output_key = "mlp.6.weight" if "mlp.6.weight" in state_dict else "net.6.weight"
+    action_dim = state_dict[output_key].shape[0]
+
+    model = PolicyMLP(obs_dim=obs_dim, action_dim=action_dim)
+    my_sd = {}
+    for k, v in state_dict.items():
+        if k.startswith("mlp."):
+            my_sd[k.replace("mlp.", "net.")] = v
+        elif k.startswith("net."):
+            my_sd[k] = v
+        elif k == "obs_normalizer._mean":
+            my_sd["obs_mean"] = v.squeeze()
+        elif k == "obs_normalizer._var":
+            my_sd["obs_std"] = torch.sqrt(v.squeeze() + 1e-5)
+
+    model.load_state_dict(my_sd, strict=False)
+    model.eval()
+    model.to(device)
+    return model, obs_dim
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--pt-path",
+        "--pt_path",
+        dest="pt_path",
+        type=Path,
+        default=PROJECT_ROOT / "model_rough.pt",
+        help="PyTorch checkpoint to export (default: ../model_rough.pt).",
+    )
+    args = parser.parse_args()
+
+    pt_path = args.pt_path.expanduser().resolve()
+    if not pt_path.exists():
+        print(f"File not found: {pt_path}")
+        return
+
+    device = torch.device("cpu")
+    print(f"Loading {pt_path}...")
+    model, obs_dim = load_policy(pt_path, device)
+
+    onnx_path = pt_path.with_suffix(".onnx")
+
+    dummy_input = torch.randn(1, obs_dim, device=device)
+
+    print(f"Exporting to {onnx_path}...")
+    torch.onnx.export(
+        model,
+        dummy_input,
+        str(onnx_path),
+        export_params=True,
+        opset_version=14,
+        do_constant_folding=True,
+        input_names=["obs"],
+        output_names=["action"],
+        dynamic_axes={"obs": {0: "batch_size"}, "action": {0: "batch_size"}}
+    )
+
+    print("Verifying ONNX export...")
+    try:
+        session = ort.InferenceSession(str(onnx_path))
+        with torch.no_grad():
+            pt_out = model(dummy_input).numpy()
+        onnx_out = session.run(["action"], {"obs": dummy_input.numpy()})[0]
+
+        max_diff = np.max(np.abs(pt_out - onnx_out))
+        mean_diff = np.mean(np.abs(pt_out - onnx_out))
+        print(f"ONNX vs PyTorch - max_diff: {max_diff:.6f}, mean_diff: {mean_diff:.6f}")
+
+        if max_diff < 1e-4:
+            print("ONNX export verified OK.")
+        else:
+            print("WARNING: ONNX export has significant divergence from PyTorch model.")
+    except ImportError:
+        print("onnxruntime not installed. Skipping verification. Install with: pip install onnxruntime")
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,238 @@
+#!/usr/bin/env python3
+"""Sweep IK compensation parameters in the standalone sim2sim scene."""
+
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import json
+import sys
+from pathlib import Path
+from typing import Any
+
+THIS_DIR = Path(__file__).resolve().parent
+PROJECT_ROOT = THIS_DIR.parent
+if str(THIS_DIR) not in sys.path:
+    sys.path.insert(0, str(THIS_DIR))
+
+SIM_PATH = THIS_DIR / "ik_slalom_sim2sim.py"
+spec = importlib.util.spec_from_file_location("ik_slalom_sim2sim", SIM_PATH)
+if spec is None or spec.loader is None:
+    raise RuntimeError(f"Cannot load {SIM_PATH}")
+sim = importlib.util.module_from_spec(spec)
+sys.modules["ik_slalom_sim2sim"] = sim
+spec.loader.exec_module(sim)
+
+
+TRIALS = [
+    {"name": "forward", "speed": 1.0, "yaw": 0.0, "target_vx": 1.0, "target_yaw": 0.0},
+    {"name": "yaw", "speed": 0.0, "yaw": 1.0, "target_vx": 0.0, "target_yaw": 1.0},
+    {"name": "arc", "speed": 1.0, "yaw": 1.0, "target_vx": 1.0, "target_yaw": 1.0},
+]
+
+
+def parse_float_list(text: str) -> list[float]:
+    return [float(x.strip()) for x in text.split(",") if x.strip()]
+
+
+def parse_bool_list(text: str) -> list[bool]:
+    out: list[bool] = []
+    for item in text.split(","):
+        key = item.strip().lower()
+        if not key:
+            continue
+        if key in {"1", "true", "on", "yes"}:
+            out.append(True)
+        elif key in {"0", "false", "off", "no"}:
+            out.append(False)
+        else:
+            raise argparse.ArgumentTypeError(f"Invalid bool item: {item}")
+    return out
+
+
+def make_sim_args(args: argparse.Namespace, trial: dict[str, float | str], cfg: dict[str, Any]) -> argparse.Namespace:
+    argv = [
+        "ik_slalom_sim2sim.py",
+        "--test",
+        str(trial["name"]),
+        "--duration",
+        str(args.duration),
+        "--settle",
+        str(args.settle),
+        "--speed",
+        str(trial["speed"]),
+        "--yaw-rate",
+        str(trial["yaw"]),
+        "--posture",
+        "custom",
+        "--custom-abduction",
+        str(args.custom_abduction),
+        "--custom-hip",
+        str(args.custom_hip),
+        "--custom-knee",
+        str(args.custom_knee),
+        "--wheel-model",
+        "direct",
+        "--linear-wheel-gain",
+        str(args.linear_wheel_gain),
+        "--direct-yaw-wheel-gain",
+        str(args.direct_yaw_wheel_gain),
+        "--max-wheel-speed",
+        str(args.max_wheel_speed),
+        "--wheel-accel-limit",
+        str(args.wheel_accel_limit),
+        "--yaw-rate-kp",
+        str(cfg["yaw_rate_kp"]),
+        "--encoder-posture-kp",
+        str(cfg["encoder_posture_kp"]),
+        "--encoder-posture-max",
+        str(cfg["encoder_posture_max"]),
+        "--roll-comp-gain",
+        str(cfg["roll_comp_gain"]),
+        "--pitch-comp-gain",
+        str(cfg["pitch_comp_gain"]),
+        "--no-realtime",
+    ]
+    argv.append("--imu-posture" if cfg["imu_posture"] else "--no-imu-posture")
+    argv.append("--encoder-guard" if cfg["encoder_guard"] else "--no-encoder-guard")
+    argv.append("--imu-guard" if cfg["imu_guard"] else "--no-imu-guard")
+    old_argv = sys.argv
+    try:
+        sys.argv = argv
+        return sim.parse_args()
+    finally:
+        sys.argv = old_argv
+
+
+def score_trial(out: dict[str, Any], trial: dict[str, float | str]) -> dict[str, float]:
+    vx = float(out["mean_body_vx_mps"])
+    yaw = float(out["mean_yaw_rate_rad_s"])
+    vx_err = abs(vx - float(trial["target_vx"]))
+    yaw_err = abs(yaw - float(trial["target_yaw"]))
+    return {
+        "vx": vx,
+        "yaw": yaw,
+        "imu_gyro_z": float(out["mean_imu_gyro_z_rad_s"]),
+        "vx_err": vx_err,
+        "yaw_err": yaw_err,
+        "err": vx_err + yaw_err,
+    }
+
+
+def run_sweep(args: argparse.Namespace) -> list[dict[str, Any]]:
+    rows: list[dict[str, Any]] = []
+    for imu_posture in args.imu_posture_values:
+        for encoder_guard in args.encoder_guard_values:
+            for imu_guard in args.imu_guard_values:
+                for encoder_posture_kp in args.encoder_posture_kps:
+                    for encoder_posture_max in args.encoder_posture_maxs:
+                        for yaw_rate_kp in args.yaw_rate_kps:
+                            for roll_comp_gain in args.roll_comp_gains:
+                                for pitch_comp_gain in args.pitch_comp_gains:
+                                    cfg = {
+                                        "imu_posture": imu_posture,
+                                        "encoder_guard": encoder_guard,
+                                        "imu_guard": imu_guard,
+                                        "encoder_posture_kp": encoder_posture_kp,
+                                        "encoder_posture_max": encoder_posture_max,
+                                        "yaw_rate_kp": yaw_rate_kp,
+                                        "roll_comp_gain": roll_comp_gain,
+                                        "pitch_comp_gain": pitch_comp_gain,
+                                    }
+                                    detail: list[dict[str, Any]] = []
+                                    speed_error = 0.0
+                                    max_tilt = 0.0
+                                    max_leg = 0.0
+                                    mean_wheel_err = 0.0
+                                    stable_all = True
+                                    for trial in TRIALS:
+                                        sim_args = make_sim_args(args, trial, cfg)
+                                        out = sim.run_one(str(trial["name"]), sim_args)
+                                        trial_score = score_trial(out, trial)
+                                        trial_score["test"] = str(trial["name"])
+                                        detail.append(trial_score)
+                                        speed_error += trial_score["err"]
+                                        max_tilt = max(max_tilt, float(out["max_tilt_deg"]))
+                                        max_leg = max(max_leg, float(out["max_leg_encoder_error_rad"]))
+                                        mean_wheel_err += float(out["mean_wheel_speed_error_rad_s"])
+                                        stable_all = stable_all and bool(out["stable"])
+
+                                    score = (
+                                        speed_error
+                                        + args.tilt_weight * max_tilt
+                                        + args.leg_error_weight * max_leg
+                                        + args.wheel_error_weight * (mean_wheel_err / len(TRIALS))
+                                    )
+                                    row = {
+                                        **cfg,
+                                        "score": round(score, 6),
+                                        "speed_error_sum": round(speed_error, 6),
+                                        "max_tilt_deg": round(max_tilt, 5),
+                                        "max_leg_encoder_error_rad": round(max_leg, 6),
+                                        "mean_wheel_speed_error_rad_s": round(mean_wheel_err / len(TRIALS), 6),
+                                        "stable_all": stable_all,
+                                        "detail": detail,
+                                    }
+                                    rows.append(row)
+                                    print(
+                                        "DONE "
+                                        + json.dumps(
+                                            {k: v for k, v in row.items() if k != "detail"},
+                                            ensure_ascii=False,
+                                        ),
+                                        flush=True,
+                                    )
+    rows.sort(key=lambda r: float(r["score"]))
+    return rows
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--duration", type=float, default=3.0)
+    parser.add_argument("--settle", type=float, default=1.5)
+    parser.add_argument("--custom-abduction", type=float, default=0.2)
+    parser.add_argument("--custom-hip", type=float, default=1.697)
+    parser.add_argument("--custom-knee", type=float, default=-2.650)
+    parser.add_argument("--linear-wheel-gain", type=float, default=12.5)
+    parser.add_argument("--direct-yaw-wheel-gain", type=float, default=8.0)
+    parser.add_argument("--max-wheel-speed", type=float, default=12.0)
+    parser.add_argument("--wheel-accel-limit", type=float, default=35.0)
+    parser.add_argument("--imu-posture-values", type=parse_bool_list, default=[True, False])
+    parser.add_argument("--encoder-guard-values", type=parse_bool_list, default=[True])
+    parser.add_argument("--imu-guard-values", type=parse_bool_list, default=[True])
+    parser.add_argument("--encoder-posture-kps", type=parse_float_list, default=[0.0, 0.05, 0.15, 0.30])
+    parser.add_argument("--encoder-posture-maxs", type=parse_float_list, default=[0.03])
+    parser.add_argument("--yaw-rate-kps", type=parse_float_list, default=[0.0, 0.4, 0.8])
+    parser.add_argument("--roll-comp-gains", type=parse_float_list, default=[0.35])
+    parser.add_argument("--pitch-comp-gains", type=parse_float_list, default=[0.35])
+    parser.add_argument("--tilt-weight", type=float, default=0.02)
+    parser.add_argument("--leg-error-weight", type=float, default=0.5)
+    parser.add_argument("--wheel-error-weight", type=float, default=0.0)
+    parser.add_argument("--top", type=int, default=12)
+    parser.add_argument("--json", type=Path, default=None)
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    rows = run_sweep(args)
+    if args.json:
+        args.json.write_text(json.dumps(rows, indent=2), encoding="utf-8")
+
+    print("\nTop compensation parameter sets")
+    print("rank score speed_err tilt leg_err wheel_err imu enc_kp yaw_kp enc_guard imu_guard")
+    for i, row in enumerate(rows[: args.top], 1):
+        print(
+            f"{i:2d} {row['score']:7.4f} {row['speed_error_sum']:7.4f} "
+            f"{row['max_tilt_deg']:5.2f} {row['max_leg_encoder_error_rad']:7.4f} "
+            f"{row['mean_wheel_speed_error_rad_s']:7.4f} "
+            f"{int(row['imu_posture'])} {row['encoder_posture_kp']:6.3f} "
+            f"{row['yaw_rate_kp']:6.3f} {int(row['encoder_guard'])} {int(row['imu_guard'])}"
+        )
+        for d in row["detail"]:
+            print(f"   {d['test']:<7} vx={d['vx']:+.3f} yaw={d['yaw']:+.3f} err={d['err']:.3f}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
@@ -87,9 +87,14 @@ class MuJoCoIO:
        return out_xml_path

    def _rebuild_actuators(self, spec):
-        actuators_to_delete = list(spec.actuators)
-        for act in actuators_to_delete:
-            spec.delete(act)
+        if hasattr(spec, "delete"):
+            actuators_to_delete = list(spec.actuators)
+            for act in actuators_to_delete:
+                spec.delete(act)
+        else:
+            actuators_to_delete = list(spec.actuators)
+            for act in actuators_to_delete:
+                act.delete()
        
        # Keep sim2sim aligned with the training robot config and sim2real runtime:
        # leg position PD = (50.0, 1.5), wheel velocity damping = 1.0.
@@ -324,9 +324,11 @@ def main():
    terrain_dir = Path(__file__).parent / "terrain"
    terrain_xml = terrain_dir / "scene_terrain.xml"
    robot_xml = project_root / "mjcf" / "wheelleg.xml"
+    rough_onnx = project_root / "model_6800.onnx"
+    crawl_onnx = project_root / "model_crawl.onnx"
    policy_path = {
-        "rough": project_root / "model_rough.pt",
-        "crawl": project_root / "model_crawl.pt"
+        "rough": rough_onnx if rough_onnx.exists() else project_root / "model_rough.pt",
+        "crawl": crawl_onnx if crawl_onnx.exists() else project_root / "model_crawl.pt"
    }

    # 1. 解析 XML 地图障碍物，实现 100% 可视化精准对应
@@ -3,7 +3,10 @@ import torch.nn as nn
 import numpy as np
 from collections import deque
 from pathlib import Path
-from pynput import keyboard
+try:
+    from pynput import keyboard
+except ImportError:
+    keyboard = None

 # ============================================================
 # Policy Model
@@ -26,6 +29,28 @@ class PolicyMLP(nn.Module):


 def load_policy(model_path, device):
+    if str(model_path).endswith('.onnx'):
+        import onnxruntime as ort
+        session = ort.InferenceSession(str(model_path))
+        class OnnxWrapper:
+            def __init__(self, session):
+                self.session = session
+                self.obs_dim = session.get_inputs()[0].shape[1]
+                if isinstance(self.obs_dim, str):
+                    self.obs_dim = 53
+                class MockMean:
+                    def __init__(self, d):
+                        self.d = d
+                    def numel(self):
+                        return self.d
+                self.obs_mean = MockMean(self.obs_dim)
+
+            def __call__(self, x):
+                inputs = {self.session.get_inputs()[0].name: x.cpu().numpy()}
+                out = self.session.run(None, inputs)[0]
+                return torch.tensor(out, device=x.device)
+        return OnnxWrapper(session)
+
    ckpt = torch.load(model_path, map_location=device, weights_only=False)
    state_dict = ckpt["actor_state_dict"]
    
@@ -129,9 +154,13 @@ class PolicyRunner:
        ], dtype=np.float32)

        # Background keyboard listener for seamless switcher keys ('1' and '2')
-        self.listener = keyboard.Listener(on_press=self._on_press)
-        self.listener.start()
-        print("[PolicyRunner] Background Keyboard Switcher active: Press '1' for ROUGH, '2' for CRAWL")
+        self.listener = None
+        if keyboard is not None:
+            self.listener = keyboard.Listener(on_press=self._on_press)
+            self.listener.start()
+            print("[PolicyRunner] Background Keyboard Switcher active: Press '1' for ROUGH, '2' for CRAWL")
+        else:
+            print("[PolicyRunner] pynput not installed; background keyboard switcher disabled.")

    def _on_press(self, key):
        try:
@@ -0,0 +1,3 @@
+# Additional runtime dependencies for the post-training Sim2Sim tools.
+onnxruntime>=1.19.0
+pygame>=2.6.1
@@ -41,6 +41,19 @@ class PolicyMLP(nn.Module):


 def load_policy(model_path, device):
+    if str(model_path).endswith('.onnx'):
+        import onnxruntime as ort
+        session = ort.InferenceSession(str(model_path))
+        class OnnxWrapper:
+            def __init__(self, session):
+                self.session = session
+            def __call__(self, x):
+                inputs = {self.session.get_inputs()[0].name: x.cpu().numpy()}
+                out = self.session.run(None, inputs)[0]
+                return torch.tensor(out, device=x.device)
+        return OnnxWrapper(session)
+
+
    ckpt = torch.load(model_path, map_location=device, weights_only=False)
    state_dict = ckpt["actor_state_dict"]
    model = PolicyMLP()
@@ -131,7 +144,9 @@ def main():
    terrain_dir = Path(__file__).parent / "terrain"
    terrain_xml = terrain_dir / "scene_terrain.xml"
    robot_xml = Path(__file__).parent.parent / "mjcf" / "wheelleg.xml"
-    policy_path = Path(__file__).parent.parent / "model_1700.pt"
+    policy_path = Path(__file__).parent.parent / "model_6800.onnx"
+    if not policy_path.exists():
+        policy_path = Path(__file__).parent.parent / "model_rough.pt"
    hfield_dir = terrain_dir

    temp_xml = project_root / "mjcf" / "sim2sim_temp.xml"