diff --git a/.gitignore b/.gitignore index 29a3009..ec9d871 100644 --- a/.gitignore +++ b/.gitignore @@ -36,6 +36,7 @@ checkpoints/ wandb/ sim2sim_log_*.txt **/sim2sim_temp.xml +**/route_check_runs/ # IDE and operating system files .idea/ diff --git a/01_doc/version_history.md b/01_doc/version_history.md index e5a5a06..6970060 100644 --- a/01_doc/version_history.md +++ b/01_doc/version_history.md @@ -12,6 +12,7 @@ | `v0.5.0` | 随机化增强 | 扩大观测、延迟和动力学随机化，加入持续外力扰动 | | `v0.6.0` | 比赛训练架构 | 分轴奖励、自适应指令课程、障碍释放课程和比赛站姿 | | `v0.7.0` | MuJoCo 工具 | 姿态优化、IK 扫描、动力学、MPC 和 GUI 调试工具 | +| `v0.8.0` | 后期 Sim2Sim | ONNX 回放、IK/路线检查工具和比赛最终 Rough 策略 | ## `v0.4.0` 的模型变化 @@ -51,3 +52,12 @@ - 增加 IK/差速轮参数扫描，可导出 JSON 结果。 - 增加 Robot、Controller、Dynamics、MPCController 和 GUI 调试链路。 - 记录历史工具常量与新版 MJCF 质量、比赛默认站姿之间的参数边界，避免将分析结果直接当作已校准真机参数。 + +## `v0.8.0` 的后期 Sim2Sim + +- 比赛训练任务、MJCF 和 `v0.7.0` 的 MuJoCo 工具保持不变。 +- 策略运行器增加 ONNX 加载，并允许在缺少 `pynput` 时关闭后台键盘监听继续运行。 +- MuJoCo 执行器重建同时兼容新旧 Spec 删除接口。 +- 增加 PT→ONNX 导出、IK 补偿扫描、纯 IK 绕桩和 ONNX 批量路线检查入口。 +- 归档比赛最终 Rough 策略 `model_6800.onnx`；其 SHA-256 为 `3C994BDD3434AD15770A52AC0E8D229F502F00D6511CDD42C2E2C742301AEF13`。 +- Crawl 权重、运行日志、临时 XML 和大量重复路线实验不在本阶段归档。 diff --git a/05_software/README.md b/05_software/README.md index 3ed6053..157bf58 100644 --- a/05_software/README.md +++ b/05_software/README.md @@ -28,7 +28,7 @@ MJCF + mjlab task IK real --------------------------------> 电机 ``` -`rc_mjlab` 是自包含工程。训练、MJCF、MuJoCo、Sim2Sim 和策略权重通过相对路径绑定，因此保留其内部布局，没有为了目录外观拆散。第一代完整闭环见 `v0.3.0`，第一份新版 MJCF 与训练框架见 `v0.4.0`，随机化增强版见 `v0.5.0`，比赛最终训练架构见 `v0.6.0`，后期 MuJoCo 工具集见 `v0.7.0`。 +`rc_mjlab` 是自包含工程。训练、MJCF、MuJoCo、Sim2Sim 和策略权重通过相对路径绑定，因此保留其内部布局，没有为了目录外观拆散。第一代完整闭环见 `v0.3.0`，第一份新版 MJCF 与训练框架见 `v0.4.0`，随机化增强版见 `v0.5.0`，比赛最终训练架构见 `v0.6.0`，后期 MuJoCo 工具集见 `v0.7.0`，后期 Sim2Sim 与比赛 Rough 策略见 `v0.8.0`。 详细说明见： diff --git a/05_software/train/README.md b/05_software/train/README.md index fbb3db4..4fb7831 100644 --- a/05_software/train/README.md +++ b/05_software/train/README.md @@ -2,7 +2,7 @@ `rc_mjlab/` 保存 16DOF 轮足机器人的当前训练与 Sim2Sim 工程。历史快照由 Git Tag 保留，不在目录中复制 `old`、`new` 或 `final` 版本。 -当前内容对应 `v0.7.0`：训练代码保持 `v0.6.0` 的比赛架构，新增后期 MuJoCo 姿态、IK、动力学和 MPC 工具。训练过程可能先获得基模，再调整奖励、课程和环境参数继续训练；模型 checkpoint 的变化不等同于软件架构变化。 +当前内容对应 `v0.8.0`：训练代码保持 `v0.6.0` 的比赛架构，包含 `v0.7.0` 的 MuJoCo 独立工具，并新增后期 Sim2Sim、路线检查和比赛最终 Rough ONNX 策略。训练过程可能先获得基模，再调整奖励、课程和环境参数继续训练；模型 checkpoint 的变化不等同于软件架构变化。 ## 内容 @@ -12,6 +12,7 @@ - `mujoco_sim`：不依赖训练循环的姿态、IK、动力学和 MPC 分析 - `mjlab`：固定版本的本地训练框架依赖 - `model_rough.pt`：本阶段 Rough 策略权重 +- `model_6800.onnx`：比赛最终使用的 Rough 策略 - `pyproject.toml`、`uv.lock`：Python 环境与依赖锁定 与 `v0.3.0` 相比，本版本更新了 MJCF 质量和惯性参数，并将 mjlab 上游基准从 `00409797` 更新到 `40f8d93e`。机械 CAD 未发生变化。 @@ -22,4 +23,6 @@ `v0.7.0` 不修改比赛训练架构，增加独立 MuJoCo 工具；入口和参数边界见 [`rc_mjlab/mujoco_sim/README.md`](rc_mjlab/mujoco_sim/README.md)。 +`v0.8.0` 继续保持训练架构和 MJCF 不变，归档后期 Sim2Sim 增量与比赛 Rough ONNX 策略；入口和归档边界见 [`rc_mjlab/sim2sim/README.md`](rc_mjlab/sim2sim/README.md)。 + 工程命令和任务说明见 [`rc_mjlab/README.md`](rc_mjlab/README.md)，本地依赖来源见 [`rc_mjlab/DEPENDENCIES.md`](rc_mjlab/DEPENDENCIES.md)。 diff --git a/05_software/train/rc_mjlab/DEPENDENCIES.md b/05_software/train/rc_mjlab/DEPENDENCIES.md index b3fb65e..768e85d 100644 --- a/05_software/train/rc_mjlab/DEPENDENCIES.md +++ b/05_software/train/rc_mjlab/DEPENDENCIES.md @@ -8,6 +8,7 @@ - `mjlab[cu128]` - PyTorch CUDA 12.8 环境 - `pynput` +- 后期 Sim2Sim 可选依赖：Pygame、ONNX Runtime 精确解析结果保存在 `uv.lock`。项目使用本地可编辑的 `mjlab`： @@ -39,4 +40,10 @@ uv run train Robot-Flat-v0 uv run play Robot-Rough-v0 ``` +根 `uv.lock` 保留比赛训练环境的历史解析结果。后期 Sim2Sim 新增依赖单独保存在 `sim2sim/requirements.txt`，运行时叠加，避免重新锁定时升级历史 MuJoCo nightly： + +```bash +uv run --with-requirements sim2sim/requirements.txt python sim2sim/nav_sim2sim.py +``` + GPU、CUDA、MuJoCo development wheel 和驱动版本必须满足 `pyproject.toml` 与 `uv.lock` 的约束。 diff --git a/05_software/train/rc_mjlab/README.md b/05_software/train/rc_mjlab/README.md index 283d13f..71d71b2 100644 --- a/05_software/train/rc_mjlab/README.md +++ b/05_software/train/rc_mjlab/README.md @@ -2,7 +2,7 @@ 基于 [mjlab](https://github.com/google-deepmind/mjlab) 框架的四轮腿混合机器人强化学习训练与部署部署项目，面向机器人竞赛场景（如越障、匍匐、斜坡、台阶等复合任务）。 -> 当前目录对应 `v0.7.0`：保留 `v0.6.0` 的比赛训练架构，并加入后期 MuJoCo 独立工具集。当前目录中的 `model_rough.pt` 是早期参考权重；比赛最终使用的 `model_6800.onnx` 将随最终部署版本归档。 +> 当前目录对应 `v0.8.0`：保留 `v0.6.0` 的比赛训练架构与 `v0.7.0` 的 MuJoCo 独立工具集，加入后期 Sim2Sim 工具和比赛最终 Rough 策略 `model_6800.onnx`。`model_rough.pt` 仍作为早期参考权重保留。 --- @@ -41,6 +41,10 @@ rc_mjlab/ │ └── competition_terrains.py # 竞赛自定义地形（高墙障碍、低杆障碍） ├── sim2sim/ # Sim2Sim 物理部署与高精度交互回放工具 │ ├── nav_sim2sim.py # 主程序：2D Pygame 交互面板 + 全自动多地形导航追踪 +│ ├── nav_route_sim2sim_check.py # ONNX 策略批量路线检查 +│ ├── ik_slalom_sim2sim.py # 纯 IK、路径跟踪与绕桩验证 +│ ├── ik_compensation_sweep.py # IK 补偿参数扫描 +│ ├── export_onnx.py # PT actor 导出与 ONNX 一致性检查 │ ├── sim2sim.py # 简易版键盘调试工具 │ ├── interface/ │ │ └── mujoco_io.py # MuJoCo 输入输出与传感器、低通滤波器接口 @@ -54,7 +58,9 @@ rc_mjlab/ │ ├── scene.xml # mjlab 场景入口文件 │ └── meshes/ # STL/OBJ 碰撞与外观网格 ├── mujoco_sim/ # 姿态、IK、动力学和 MPC 独立工具 -├── model_rough.pt # 本阶段用于回放和 Sim2Sim 的 Rough 策略 +├── tools/nav_tools/ # 路线安全检查公共模块 +├── model_rough.pt # 早期 Rough 参考 checkpoint +├── model_6800.onnx # 比赛最终 Rough 策略 ├── pyproject.toml # 项目依赖（uv 管理，含清华镜像源加速） └── uv.lock # 精确依赖锁定文件 ``` @@ -89,6 +95,8 @@ cd sim2sim uv run python nav_sim2sim.py ``` +后期 Sim2Sim 的入口、模型边界和批量检查命令见 [`sim2sim/README.md`](sim2sim/README.md)。 + --- ## 🖥️ 交互式自动导航平台 (sim2sim/nav_sim2sim.py) diff --git a/05_software/train/rc_mjlab/model_6800.onnx b/05_software/train/rc_mjlab/model_6800.onnx new file mode 100644 index 0000000..49ea593 Binary files /dev/null and b/05_software/train/rc_mjlab/model_6800.onnx differ diff --git a/05_software/train/rc_mjlab/sim2sim/README.md b/05_software/train/rc_mjlab/sim2sim/README.md new file mode 100644 index 0000000..8e43e94 --- /dev/null +++ b/05_software/train/rc_mjlab/sim2sim/README.md @@ -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/` 不纳入版本库。源目录中的大量路线试验结果也未复制；它们包含重复轨迹和本机绝对路径，不属于可复用程序源码。 diff --git a/05_software/train/rc_mjlab/sim2sim/export_onnx.py b/05_software/train/rc_mjlab/sim2sim/export_onnx.py new file mode 100644 index 0000000..d1be128 --- /dev/null +++ b/05_software/train/rc_mjlab/sim2sim/export_onnx.py @@ -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() diff --git a/05_software/train/rc_mjlab/sim2sim/ik_compensation_sweep.py b/05_software/train/rc_mjlab/sim2sim/ik_compensation_sweep.py new file mode 100644 index 0000000..f013f7b --- /dev/null +++ b/05_software/train/rc_mjlab/sim2sim/ik_compensation_sweep.py @@ -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()) diff --git a/05_software/train/rc_mjlab/sim2sim/ik_slalom_sim2sim.py b/05_software/train/rc_mjlab/sim2sim/ik_slalom_sim2sim.py new file mode 100644 index 0000000..5107681 --- /dev/null +++ b/05_software/train/rc_mjlab/sim2sim/ik_slalom_sim2sim.py @@ -0,0 +1,1235 @@ +#!/usr/bin/env python3 +""" +Standalone sim2sim check for wheel-leg IK navigation. + +This file is intentionally isolated from nav_sim2sim.py and the ROS2 runtime. +It keeps the 12 leg joints at the RL rough-terrain default pose and drives the +four wheel velocity actuators with a differential-drive inverse kinematics law. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import sys +import tempfile +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Iterable + +import mujoco +import mujoco.viewer +import numpy as np + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +SIM2SIM_DIR = PROJECT_ROOT / "sim2sim" +if str(SIM2SIM_DIR) not in sys.path: + sys.path.insert(0, str(SIM2SIM_DIR)) + +from tools.math_utils import get_gravity_orientation # noqa: E402 + + +LEG_JOINT_NAMES = [ + "fl_hip_abduction_joint", + "fl_hip_pitch_joint", + "fl_knee_joint", + "fr_hip_abduction_joint", + "fr_hip_pitch_joint", + "fr_knee_joint", + "rl_hip_abduction_joint", + "rl_hip_pitch_joint", + "rl_knee_joint", + "rr_hip_abduction_joint", + "rr_hip_pitch_joint", + "rr_knee_joint", +] +WHEEL_JOINT_NAMES = [ + "fl_wheel_joint", + "fr_wheel_joint", + "rl_wheel_joint", + "rr_wheel_joint", +] +ALL_JOINT_NAMES = LEG_JOINT_NAMES + WHEEL_JOINT_NAMES +WHEEL_BODY_NAMES = [ + "fl_wheel_Link", + "fr_wheel_Link", + "rl_wheel_Link", + "rr_wheel_Link", +] + +# Must match src/robot/robot_cfg.py rough policy initial state. +ROUGH_DEFAULT_DOF_POS = np.array( + [0.0, 0.550, -1.125] * 4 + [0.0] * 4, + dtype=np.float64, +) + +# From mujoco_sim/config.py. These poses keep the wheel contact close to the +# hip in sagittal X while changing body height. +HEIGHT_TABLE = np.array( + [ + (0.17, 0.914, -2.628), + (0.19, 0.926, -2.528), + (0.21, 0.924, -2.428), + (0.23, 0.912, -2.328), + (0.25, 0.892, -2.226), + (0.27, 0.864, -2.122), + (0.29, 0.834, -2.014), + (0.31, 0.798, -1.906), + (0.33, 0.758, -1.792), + (0.35, 0.714, -1.672), + (0.37, 0.666, -1.546), + (0.39, 0.612, -1.412), + (0.41, 0.552, -1.266), + (0.43, 0.484, -1.104), + (0.45, 0.404, -0.918), + ], + dtype=np.float64, +) + + +def wrap_pi(x: float) -> float: + return (x + math.pi) % (2.0 * math.pi) - math.pi + + +def quat_to_euler_wxyz(q: np.ndarray) -> tuple[float, float, float]: + w, x, y, z = q + sinr_cosp = 2.0 * (w * x + y * z) + cosr_cosp = 1.0 - 2.0 * (x * x + y * y) + roll = math.atan2(sinr_cosp, cosr_cosp) + + sinp = 2.0 * (w * y - z * x) + pitch = math.asin(float(np.clip(sinp, -1.0, 1.0))) + + siny_cosp = 2.0 * (w * z + x * y) + cosy_cosp = 1.0 - 2.0 * (y * y + z * z) + yaw = math.atan2(siny_cosp, cosy_cosp) + return roll, pitch, yaw + + +def yaw_to_quat_wxyz(yaw: float) -> np.ndarray: + half = 0.5 * yaw + return np.array([math.cos(half), 0.0, 0.0, math.sin(half)], dtype=np.float64) + + +def quat_to_mat_wxyz(q: np.ndarray) -> np.ndarray: + w, x, y, z = q + return np.array( + [ + [1.0 - 2.0 * (y * y + z * z), 2.0 * (x * y - z * w), 2.0 * (x * z + y * w)], + [2.0 * (x * y + z * w), 1.0 - 2.0 * (x * x + z * z), 2.0 * (y * z - x * w)], + [2.0 * (x * z - y * w), 2.0 * (y * z + x * w), 1.0 - 2.0 * (x * x + y * y)], + ], + dtype=np.float64, + ) + + +def table_pose_for_height(height: float) -> np.ndarray: + h = HEIGHT_TABLE[:, 0] + hip = float(np.interp(np.clip(height, h[0], h[-1]), h, HEIGHT_TABLE[:, 1])) + knee = float(np.interp(np.clip(height, h[0], h[-1]), h, HEIGHT_TABLE[:, 2])) + return np.array([0.0, hip, knee] * 4 + [0.0] * 4, dtype=np.float64) + + +def nominal_dof_pos(args: argparse.Namespace) -> np.ndarray: + if args.posture == "custom": + ab = args.custom_abduction + return np.array( + [ + ab, + args.custom_hip, + args.custom_knee, + -ab, + args.custom_hip, + args.custom_knee, + ab, + args.custom_hip, + args.custom_knee, + -ab, + args.custom_hip, + args.custom_knee, + 0.0, + 0.0, + 0.0, + 0.0, + ], + dtype=np.float64, + ) + if args.posture == "table": + return table_pose_for_height(args.body_height) + return ROUGH_DEFAULT_DOF_POS.copy() + + +def make_flat_scene_xml(robot_abs: str, tmp_dir: Path) -> Path: + out = tmp_dir / "ik_flat_scene.xml" + out.write_text( + f""" + + + +""", + encoding="utf-8", + ) + return out + + +def make_sim_xml(terrain_xml: Path | None, robot_xml: Path, tmp_dir: Path) -> Path: + robot_content = robot_xml.read_text(encoding="utf-8") + mesh_dir = str((robot_xml.parent / "meshes").resolve()).replace("\\", "/") + robot_content = robot_content.replace( + '', + f'', + ) + tmp_dir.mkdir(parents=True, exist_ok=True) + robot_out = tmp_dir / "wheelleg_abs_mesh.xml" + robot_out.write_text(robot_content, encoding="utf-8") + + robot_abs = str(robot_out.resolve()).replace("\\", "/") + if terrain_xml is None: + return make_flat_scene_xml(robot_abs, tmp_dir) + + content = terrain_xml.read_text(encoding="utf-8") + content = content.replace('', f'') + content = content.replace('', f'') + + terrain_dir = terrain_xml.parent + replacements = { + "../height_field.png": str((terrain_dir / "height_field.png").resolve()).replace("\\", "/"), + "../unitree_hfield.png": str((terrain_dir / "unitree_hfield.png").resolve()).replace("\\", "/"), + } + for old, new in replacements.items(): + content = content.replace(old, new) + + out = tmp_dir / "ik_slalom_scene.xml" + out.write_text(content, encoding="utf-8") + return out + + +def rebuild_actuators(spec: mujoco.MjSpec) -> None: + for actuator in list(spec.actuators): + spec.delete(actuator) + + kp_leg = 50.0 + kd_leg = 1.5 + kd_wheel = 1.0 + effort_limit = 17.0 + + for jname in LEG_JOINT_NAMES: + act = spec.add_actuator(name=jname, target=jname) + act.trntype = mujoco.mjtTrn.mjTRN_JOINT + act.dyntype = mujoco.mjtDyn.mjDYN_NONE + act.gaintype = mujoco.mjtGain.mjGAIN_FIXED + act.biastype = mujoco.mjtBias.mjBIAS_AFFINE + act.gainprm[0] = kp_leg + act.biasprm[1] = -kp_leg + act.biasprm[2] = -kd_leg + act.forcelimited = True + act.forcerange[:] = [-effort_limit, effort_limit] + act.ctrllimited = False + + for jname in WHEEL_JOINT_NAMES: + act = spec.add_actuator(name=jname, target=jname) + act.trntype = mujoco.mjtTrn.mjTRN_JOINT + act.dyntype = mujoco.mjtDyn.mjDYN_NONE + act.gaintype = mujoco.mjtGain.mjGAIN_FIXED + act.biastype = mujoco.mjtBias.mjBIAS_AFFINE + act.gainprm[0] = kd_wheel + act.biasprm[2] = -kd_wheel + act.forcelimited = True + act.forcerange[:] = [-effort_limit, effort_limit] + act.ctrllimited = False + + +@dataclass(frozen=True) +class ModelIds: + qpos: np.ndarray + qvel: np.ndarray + ctrl: np.ndarray + wheel_body: np.ndarray + wheel_geom: np.ndarray + + +@dataclass +class SensorFrame: + time: float + base_pos: np.ndarray + quat_wxyz: np.ndarray + rpy: np.ndarray + gyro: np.ndarray + projected_gravity: np.ndarray + joint_pos: np.ndarray + joint_vel: np.ndarray + wheel_vel: np.ndarray + + +def name_id(model: mujoco.MjModel, obj_type: mujoco.mjtObj, name: str) -> int: + idx = mujoco.mj_name2id(model, obj_type, name) + if idx < 0: + raise RuntimeError(f"MuJoCo object not found: {name}") + return idx + + +def build_model(args: argparse.Namespace, tmp_dir: Path) -> tuple[mujoco.MjModel, mujoco.MjData, ModelIds]: + scene_xml = make_sim_xml(args.terrain_xml, args.robot_xml, tmp_dir) + + old_cwd = Path.cwd() + os.chdir(str(args.robot_xml.parent)) + try: + spec = mujoco.MjSpec.from_file(str(scene_xml)) + rebuild_actuators(spec) + model = spec.compile() + finally: + os.chdir(str(old_cwd)) + + data = mujoco.MjData(model) + qpos_ids = np.array( + [model.jnt_qposadr[name_id(model, mujoco.mjtObj.mjOBJ_JOINT, n)] for n in ALL_JOINT_NAMES], + dtype=np.int32, + ) + qvel_ids = np.array( + [model.jnt_dofadr[name_id(model, mujoco.mjtObj.mjOBJ_JOINT, n)] for n in ALL_JOINT_NAMES], + dtype=np.int32, + ) + ctrl_ids = np.array( + [name_id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, n) for n in ALL_JOINT_NAMES], + dtype=np.int32, + ) + wheel_body_ids = np.array( + [name_id(model, mujoco.mjtObj.mjOBJ_BODY, n) for n in WHEEL_BODY_NAMES], + dtype=np.int32, + ) + + wheel_geom_ids: list[int] = [] + for body_id in wheel_body_ids: + for geom_id in range(model.ngeom): + if model.geom_bodyid[geom_id] == body_id and model.geom_type[geom_id] == mujoco.mjtGeom.mjGEOM_CYLINDER: + wheel_geom_ids.append(geom_id) + break + else: + raise RuntimeError(f"No cylinder collision geom found for body id {body_id}") + + return model, data, ModelIds(qpos_ids, qvel_ids, ctrl_ids, wheel_body_ids, np.array(wheel_geom_ids)) + + +def reset_robot(model: mujoco.MjModel, data: mujoco.MjData, ids: ModelIds, args: argparse.Namespace) -> np.ndarray: + default_dof_pos = nominal_dof_pos(args) + data.qpos[:3] = [args.start_x, args.start_y, args.start_z] + data.qpos[3:7] = [1.0, 0.0, 0.0, 0.0] + data.qpos[ids.qpos] = default_dof_pos + data.qvel[:] = 0.0 + data.ctrl[ids.ctrl[:12]] = default_dof_pos[:12] + data.ctrl[ids.ctrl[12:]] = 0.0 + mujoco.mj_forward(model, data) + + settle_steps = int(round(args.settle / model.opt.timestep)) + for _ in range(settle_steps): + data.ctrl[ids.ctrl[:12]] = default_dof_pos[:12] + data.ctrl[ids.ctrl[12:]] = 0.0 + mujoco.mj_step(model, data) + return default_dof_pos + + +def read_sensors(model: mujoco.MjModel, data: mujoco.MjData, ids: ModelIds) -> SensorFrame: + del model + quat = data.qpos[3:7].copy() + rpy = np.array(quat_to_euler_wxyz(quat), dtype=np.float64) + joint_pos = data.qpos[ids.qpos].copy() + joint_vel = data.qvel[ids.qvel].copy() + return SensorFrame( + time=float(data.time), + base_pos=data.qpos[:3].copy(), + quat_wxyz=quat, + rpy=rpy, + gyro=data.qvel[3:6].copy(), + projected_gravity=get_gravity_orientation(quat), + joint_pos=joint_pos, + joint_vel=joint_vel, + wheel_vel=joint_vel[12:].copy(), + ) + + +@dataclass(frozen=True) +class WheelGeometry: + radius: float + track_width: float + wheel_y: np.ndarray + + +def infer_wheel_geometry(model: mujoco.MjModel, data: mujoco.MjData, ids: ModelIds) -> WheelGeometry: + radius_values = model.geom_size[ids.wheel_geom, 0] + radius = float(np.mean(radius_values)) + wheel_y = data.xipos[ids.wheel_body, 1].copy() + left_y = wheel_y[[0, 2]] + right_y = wheel_y[[1, 3]] + track = float(np.mean(left_y) - np.mean(right_y)) + if radius <= 0.0 or track <= 0.0: + raise RuntimeError(f"Invalid wheel geometry: radius={radius}, track={track}") + return WheelGeometry(radius=radius, track_width=track, wheel_y=wheel_y) + + +@dataclass +class IkCommand: + linear_x: float + yaw_rate: float + target_index: int = -1 + target_distance: float = 0.0 + + +class DifferentialWheelIk: + def __init__( + self, + geometry: WheelGeometry, + max_wheel_speed: float, + wheel_signs: np.ndarray, + yaw_wheel_gain: float, + args: argparse.Namespace, + ) -> None: + self.geometry = geometry + self.max_wheel_speed = max_wheel_speed + self.wheel_signs = wheel_signs.astype(np.float64) + self.yaw_wheel_gain = yaw_wheel_gain + self.args = args + + def solve(self, cmd: IkCommand, sensor: SensorFrame | None = None) -> np.ndarray: + half_track = 0.5 * self.geometry.track_width + yaw_rate = cmd.yaw_rate + if sensor is not None: + yaw_rate += self.args.yaw_rate_kp * (cmd.yaw_rate - float(sensor.gyro[2])) + + if self.args.wheel_model == "direct": + linear_wheel = self.args.linear_wheel_gain * cmd.linear_x + yaw_wheel = self.args.direct_yaw_wheel_gain * yaw_rate + else: + linear_wheel = cmd.linear_x / self.geometry.radius + yaw_wheel = self.yaw_wheel_gain * yaw_rate * half_track / self.geometry.radius + yaw_wheel = float(np.clip(yaw_wheel, -self.args.max_yaw_wheel_speed, self.args.max_yaw_wheel_speed)) + left = linear_wheel - yaw_wheel + right = linear_wheel + yaw_wheel + wheel = np.array([left, right, left, right], dtype=np.float64) + wheel = np.clip(wheel, -self.max_wheel_speed, self.max_wheel_speed) + return wheel * self.wheel_signs + + +class LegStanceIk: + """Numerical wheel-center IK for keeping the rough RL stance level.""" + + def __init__(self, model: mujoco.MjModel, data: mujoco.MjData, ids: ModelIds, args: argparse.Namespace) -> None: + self.model = model + self.ids = ids + self.args = args + self.shadow = mujoco.MjData(model) + self.jacp = np.zeros((3, model.nv), dtype=np.float64) + self.jacr = np.zeros((3, model.nv), dtype=np.float64) + self.last_target = ROUGH_DEFAULT_DOF_POS[:12].copy() + + self.leg_qpos = [ids.qpos[i * 3 : i * 3 + 3] for i in range(4)] + self.leg_qvel = [ids.qvel[i * 3 : i * 3 + 3] for i in range(4)] + self.leg_joint_ids = [ + [name_id(model, mujoco.mjtObj.mjOBJ_JOINT, LEG_JOINT_NAMES[i * 3 + j]) for j in range(3)] + for i in range(4) + ] + self.joint_ranges = [ + np.array([model.jnt_range[jid].copy() for jid in leg], dtype=np.float64) + for leg in self.leg_joint_ids + ] + + mujoco.mj_forward(model, data) + base_pos = data.qpos[:3].copy() + base_rot = quat_to_mat_wxyz(data.qpos[3:7]) + self.nominal_body_wheel = (base_rot.T @ (data.xipos[ids.wheel_body] - base_pos).T).T + + def compute(self, data: mujoco.MjData, sensor: SensorFrame) -> np.ndarray: + current_rot = quat_to_mat_wxyz(sensor.quat_wxyz) + level_rot = quat_to_mat_wxyz(yaw_to_quat_wxyz(float(sensor.rpy[2]))) + base_targets = (current_rot @ self.nominal_body_wheel.T).T + level_targets = (level_rot @ self.nominal_body_wheel.T).T + target_world = sensor.base_pos + base_targets + self.args.posture_gain * (level_targets - base_targets) + + self.shadow.qpos[:] = data.qpos + self.shadow.qvel[:] = 0.0 + mujoco.mj_forward(self.model, self.shadow) + + damping = self.args.leg_ik_damping + for _ in range(self.args.leg_ik_iterations): + for leg_idx, body_id in enumerate(self.ids.wheel_body): + err = target_world[leg_idx] - self.shadow.xipos[body_id] + if float(np.linalg.norm(err)) < 1e-4: + continue + + self.jacp.fill(0.0) + self.jacr.fill(0.0) + mujoco.mj_jacBody(self.model, self.shadow, self.jacp, self.jacr, int(body_id)) + cols = self.leg_qvel[leg_idx] + j_leg = self.jacp[:, cols] + lhs = j_leg @ j_leg.T + (damping * damping) * np.eye(3) + dq = j_leg.T @ np.linalg.solve(lhs, err) + dq = np.clip(dq, -self.args.leg_ik_step, self.args.leg_ik_step) + + qpos_ids = self.leg_qpos[leg_idx] + self.shadow.qpos[qpos_ids] += dq + ranges = self.joint_ranges[leg_idx] + self.shadow.qpos[qpos_ids] = np.clip(self.shadow.qpos[qpos_ids], ranges[:, 0], ranges[:, 1]) + mujoco.mj_forward(self.model, self.shadow) + + raw_target = self.shadow.qpos[self.ids.qpos[:12]].copy() + offset = np.clip( + raw_target - ROUGH_DEFAULT_DOF_POS[:12], + -self.args.max_leg_offset, + self.args.max_leg_offset, + ) + raw_target = ROUGH_DEFAULT_DOF_POS[:12] + offset + + max_delta = self.args.leg_target_rate + filtered = self.last_target + np.clip(raw_target - self.last_target, -max_delta, max_delta) + alpha = self.args.leg_target_filter + filtered = (1.0 - alpha) * self.last_target + alpha * filtered + self.last_target = filtered.copy() + return filtered + + +class TablePostureController: + def __init__(self, args: argparse.Namespace, initial_target: np.ndarray) -> None: + self.args = args + self.target = initial_target[:12].copy() + + def compute(self, sensor: SensorFrame) -> np.ndarray: + if self.args.posture == "rough": + raw = ROUGH_DEFAULT_DOF_POS[:12].copy() + elif self.args.posture == "custom": + raw = nominal_dof_pos(self.args)[:12].copy() + else: + pose = table_pose_for_height(self.args.body_height) + raw = pose[:12].copy() + + if self.args.imu_posture: + roll_corr = -self.args.roll_comp_gain * float(sensor.rpy[0]) + pitch_corr = -self.args.pitch_comp_gain * float(sensor.rpy[1]) + for i, side in enumerate((1.0, -1.0, 1.0, -1.0)): + raw[i * 3] = np.clip(raw[i * 3] + side * roll_corr, -0.5, 0.5) + raw[i * 3 + 1] = np.clip(raw[i * 3 + 1] + pitch_corr, -1.0, 2.5) + raw[i * 3 + 2] = np.clip(raw[i * 3 + 2], -2.65, -0.3) + + if self.args.encoder_posture_kp > 0.0: + encoder_err = self.target - sensor.joint_pos[:12] + raw += np.clip( + self.args.encoder_posture_kp * encoder_err, + -self.args.encoder_posture_max, + self.args.encoder_posture_max, + ) + + max_delta = self.args.leg_target_rate + raw = self.target + np.clip(raw - self.target, -max_delta, max_delta) + alpha = self.args.leg_target_filter + self.target = (1.0 - alpha) * self.target + alpha * raw + return self.target.copy() + + +def sensor_command_scale(sensor: SensorFrame, leg_target: np.ndarray, args: argparse.Namespace) -> float: + scale = 1.0 + if args.encoder_guard: + leg_error = float(np.max(np.abs(sensor.joint_pos[:12] - leg_target))) + if leg_error >= args.encoder_guard_stop: + scale = 0.0 + elif leg_error > args.encoder_guard_start: + span = max(1e-6, args.encoder_guard_stop - args.encoder_guard_start) + scale *= 1.0 - (leg_error - args.encoder_guard_start) / span + + if args.imu_guard: + tilt = math.hypot(float(sensor.rpy[0]), float(sensor.rpy[1])) + start = math.radians(args.imu_guard_start_deg) + stop = math.radians(args.imu_guard_stop_deg) + if tilt >= stop: + scale = 0.0 + elif tilt > start: + scale *= 1.0 - (tilt - start) / max(1e-6, stop - start) + + return float(np.clip(scale, 0.0, 1.0)) + + +class PathFollower: + def __init__( + self, + waypoints: Iterable[tuple[float, float]], + max_speed: float, + max_yaw_rate: float, + arrive_radius: float, + allow_reverse: bool, + ) -> None: + self.waypoints = list(waypoints) + self.idx = 0 + self.max_speed = max_speed + self.max_yaw_rate = max_yaw_rate + self.arrive_radius = arrive_radius + self.allow_reverse = allow_reverse + + @property + def done(self) -> bool: + return self.idx >= len(self.waypoints) + + def update(self, sensor: SensorFrame) -> IkCommand: + if self.done: + return IkCommand(0.0, 0.0, self.idx, 0.0) + + x, y = sensor.base_pos[:2] + target = self.waypoints[self.idx] + dx = target[0] - float(x) + dy = target[1] - float(y) + dist = math.hypot(dx, dy) + while dist < self.arrive_radius and self.idx < len(self.waypoints) - 1: + self.idx += 1 + target = self.waypoints[self.idx] + dx = target[0] - float(x) + dy = target[1] - float(y) + dist = math.hypot(dx, dy) + if dist < self.arrive_radius and self.idx == len(self.waypoints) - 1: + self.idx = len(self.waypoints) + return IkCommand(0.0, 0.0, self.idx, dist) + + desired_yaw = math.atan2(dy, dx) + yaw_error = wrap_pi(desired_yaw - float(sensor.rpy[2])) + direction = 1.0 + if self.allow_reverse and abs(yaw_error) > math.pi * 0.5: + desired_yaw = wrap_pi(desired_yaw + math.pi) + yaw_error = wrap_pi(desired_yaw - float(sensor.rpy[2])) + direction = -1.0 + yaw_rate = float(np.clip(2.8 * yaw_error, -self.max_yaw_rate, self.max_yaw_rate)) + + speed = min(self.max_speed, 1.2 * dist) + speed *= max(0.15, math.cos(min(abs(yaw_error), math.pi * 0.5))) + speed *= direction + return IkCommand(speed, yaw_rate, self.idx, dist) + + +class PurePursuitFollower: + def __init__(self, waypoints: Iterable[tuple[float, float]], args: argparse.Namespace) -> None: + self.points = np.array(list(waypoints), dtype=np.float64) + self.args = args + self.progress = 0 + self.closest_distance = float("inf") + self.finished = False + if len(self.points) < 2: + raise ValueError("PurePursuitFollower needs at least two points") + seg = self.points[1:] - self.points[:-1] + self.seg_len = np.linalg.norm(seg, axis=1) + self.cum = np.concatenate([[0.0], np.cumsum(self.seg_len)]) + self.total_len = float(self.cum[-1]) + + @property + def done(self) -> bool: + return self.finished + + @property + def idx(self) -> int: + return int(self.progress) + + def _project_s(self, p: np.ndarray) -> tuple[float, float]: + best_s = 0.0 + best_d = float("inf") + start_i = max(0, self.progress - 1) + end_i = min(len(self.points) - 1, self.progress + self.args.pure_search_segments) + for i in range(start_i, end_i): + a = self.points[i] + b = self.points[i + 1] + ab = b - a + denom = float(np.dot(ab, ab)) + if denom < 1e-9: + continue + t = float(np.clip(np.dot(p - a, ab) / denom, 0.0, 1.0)) + proj = a + t * ab + d = float(np.linalg.norm(p - proj)) + if d < best_d: + best_d = d + best_s = float(self.cum[i] + t * self.seg_len[i]) + self.progress = max(self.progress, i) + return best_s, best_d + + def _point_at_s(self, s: float) -> np.ndarray: + s = float(np.clip(s, 0.0, self.total_len)) + i = int(np.searchsorted(self.cum, s, side="right") - 1) + i = int(np.clip(i, 0, len(self.seg_len) - 1)) + local = 0.0 if self.seg_len[i] < 1e-9 else (s - self.cum[i]) / self.seg_len[i] + return self.points[i] + local * (self.points[i + 1] - self.points[i]) + + def update(self, sensor: SensorFrame) -> IkCommand: + p = sensor.base_pos[:2].astype(np.float64) + s, dist_path = self._project_s(p) + self.closest_distance = min(self.closest_distance, dist_path) + goal_dist = float(np.linalg.norm(self.points[-1] - p)) + if self.total_len - s < self.args.arrive_radius and goal_dist < self.args.arrive_radius: + self.finished = True + return IkCommand(0.0, 0.0, len(self.points) - 1, goal_dist) + + lookahead = np.clip( + self.args.lookahead_base + self.args.lookahead_time * abs(self.args.speed), + self.args.lookahead_min, + self.args.lookahead_max, + ) + target = self._point_at_s(s + lookahead) + dx, dy = target - p + yaw = float(sensor.rpy[2]) + c, sn = math.cos(yaw), math.sin(yaw) + x_body = c * dx + sn * dy + y_body = -sn * dx + c * dy + ld2 = max(lookahead * lookahead, 1e-6) + curvature = 2.0 * y_body / ld2 + + abs_curv = abs(curvature) + if abs_curv > 1e-6: + v_curve = math.sqrt(max(self.args.max_lat_acc, 1e-6) / abs_curv) + else: + v_curve = self.args.speed + speed = min(self.args.speed, v_curve) + speed = max(self.args.min_path_speed, speed) + if x_body < -0.05 and self.args.allow_reverse: + speed = -speed + yaw_rate = float(np.clip(speed * curvature, -self.args.max_yaw_rate, self.args.max_yaw_rate)) + return IkCommand(float(speed), yaw_rate, self.progress, dist_path) + + +def args_sign(x: float) -> float: + return -1.0 if x < -0.05 else 1.0 + + +def test_command(test_name: str, elapsed: float, sensor: SensorFrame, follower: PathFollower | None, args: argparse.Namespace) -> IkCommand: + if test_name == "forward": + return IkCommand(args.speed, 0.0) + if test_name == "yaw": + return IkCommand(0.0, args.yaw_rate) + if test_name == "arc": + return IkCommand(args.speed, args.yaw_rate) + if follower is None: + raise RuntimeError(f"Test {test_name} needs a path follower") + del elapsed + return follower.update(sensor) + + +def waypoints_for_test(test_name: str, args: argparse.Namespace) -> list[tuple[float, float]]: + sx = float(args.start_x) + sy = float(args.start_y) + if test_name == "waypoint": + return [ + (sx + 0.3, sy), + (sx + 0.65, sy), + (sx + 1.0, sy), + ] + if test_name == "slalom": + local = [ + (0.0, 0.0), + (0.8, -0.5), + (1.9, -0.95), + (2.8, -1.0), + (2.85, -1.8), + (2.75, -2.05), + (1.05, -2.0), + (0.9, -2.8), + (1.05, -3.0), + (2.75, -3.0), + (2.75, -3.72), + (2.25, -3.85), + (1.4, -4.25), + (1.4, -2.55), + (0.45, -3.42), + (0.0, 0.0), + ] + return [(sx + x, sy + y) for x, y in local] + return [] + + +@dataclass +class Metrics: + start_pos: np.ndarray + start_yaw: float + start_time: float = 0.0 + max_abs_roll: float = 0.0 + max_abs_pitch: float = 0.0 + max_tilt: float = 0.0 + max_leg_error: float = 0.0 + max_wheel_air_height: float = -1e9 + min_wheel_air_height: float = 1e9 + wheel_speed_error_sum: float = 0.0 + wheel_speed_error_count: int = 0 + max_wheel_speed: float = 0.0 + max_gyro_z: float = 0.0 + body_vx_sum: float = 0.0 + gyro_z_sum: float = 0.0 + yaw_unwrapped: float = 0.0 + velocity_count: int = 0 + last_yaw: float = field(init=False) + + def __post_init__(self) -> None: + self.last_yaw = float(self.start_yaw) + + def update( + self, + sensor: SensorFrame, + model: mujoco.MjModel, + data: mujoco.MjData, + ids: ModelIds, + geometry: WheelGeometry, + wheel_target: np.ndarray, + leg_target: np.ndarray, + ) -> None: + roll, pitch, _ = sensor.rpy + self.max_abs_roll = max(self.max_abs_roll, abs(float(roll))) + self.max_abs_pitch = max(self.max_abs_pitch, abs(float(pitch))) + self.max_tilt = max(self.max_tilt, math.hypot(float(roll), float(pitch))) + self.max_leg_error = max( + self.max_leg_error, + float(np.max(np.abs(sensor.joint_pos[:12] - leg_target))), + ) + + wheel_center_z = data.xipos[ids.wheel_body, 2] + floor_air_height = wheel_center_z - geometry.radius + self.max_wheel_air_height = max(self.max_wheel_air_height, float(np.max(floor_air_height))) + self.min_wheel_air_height = min(self.min_wheel_air_height, float(np.min(floor_air_height))) + + actual = sensor.wheel_vel + self.wheel_speed_error_sum += float(np.mean(np.abs(actual - wheel_target))) + self.wheel_speed_error_count += 1 + self.max_wheel_speed = max(self.max_wheel_speed, float(np.max(np.abs(actual)))) + self.max_gyro_z = max(self.max_gyro_z, abs(float(sensor.gyro[2]))) + rot = quat_to_mat_wxyz(sensor.quat_wxyz) + body_vel = rot.T @ data.qvel[:3] + yaw_now = float(sensor.rpy[2]) + self.yaw_unwrapped += wrap_pi(yaw_now - self.last_yaw) + self.last_yaw = yaw_now + self.body_vx_sum += float(body_vel[0]) + self.gyro_z_sum += float(sensor.gyro[2]) + self.velocity_count += 1 + del model + + def as_dict(self, sensor: SensorFrame, reached: int, total: int, reason: str, success: bool) -> dict: + yaw_change = wrap_pi(float(sensor.rpy[2]) - self.start_yaw) + delta = sensor.base_pos[:2] - self.start_pos[:2] + mean_wheel_err = self.wheel_speed_error_sum / max(1, self.wheel_speed_error_count) + mean_body_vx = self.body_vx_sum / max(1, self.velocity_count) + elapsed = max(1e-6, float(sensor.time) - float(self.start_time)) + mean_yaw_rate = self.yaw_unwrapped / elapsed + mean_gyro_z = self.gyro_z_sum / max(1, self.velocity_count) + return { + "success": bool(success), + "reason": reason, + "sim_time": round(sensor.time, 4), + "start_xy": [round(float(v), 4) for v in self.start_pos[:2]], + "final_xy": [round(float(v), 4) for v in sensor.base_pos[:2]], + "delta_xy": [round(float(v), 4) for v in delta], + "final_z": round(float(sensor.base_pos[2]), 4), + "yaw_change_deg": round(math.degrees(yaw_change), 3), + "final_yaw_deg": round(math.degrees(float(sensor.rpy[2])), 3), + "max_roll_deg": round(math.degrees(self.max_abs_roll), 3), + "max_pitch_deg": round(math.degrees(self.max_abs_pitch), 3), + "max_tilt_deg": round(math.degrees(self.max_tilt), 3), + "max_leg_encoder_error_rad": round(self.max_leg_error, 5), + "mean_wheel_speed_error_rad_s": round(mean_wheel_err, 5), + "max_wheel_speed_rad_s": round(self.max_wheel_speed, 5), + "mean_body_vx_mps": round(mean_body_vx, 5), + "mean_yaw_rate_rad_s": round(mean_yaw_rate, 5), + "mean_imu_gyro_z_rad_s": round(mean_gyro_z, 5), + "min_wheel_air_height_m": round(self.min_wheel_air_height, 5), + "max_wheel_air_height_m": round(self.max_wheel_air_height, 5), + "max_imu_gyro_z_rad_s": round(self.max_gyro_z, 5), + "reached_waypoints": int(reached), + "total_waypoints": int(total), + } + + +class InteractivePanel: + def __init__(self, args: argparse.Namespace) -> None: + try: + import tkinter as tk + from tkinter import ttk + except ImportError as exc: + raise RuntimeError("Interactive mode needs tkinter, but it is not available in this Python.") from exc + + self.tk = tk + self.alive = True + self.last_update = 0.0 + self.root = tk.Tk() + self.root.title("Wheel-Leg IK Control") + self.root.protocol("WM_DELETE_WINDOW", self.close) + + self.vx = tk.DoubleVar(value=float(args.speed)) + self.yaw = tk.DoubleVar(value=float(args.yaw_rate)) + self.status_vars: dict[str, object] = {} + + main = ttk.Frame(self.root, padding=10) + main.grid(row=0, column=0, sticky="nsew") + self.root.columnconfigure(0, weight=1) + main.columnconfigure(1, weight=1) + + self._add_slider(main, 0, "linear x m/s", self.vx, args.interactive_vx_limit) + self._add_slider(main, 1, "yaw z rad/s", self.yaw, args.interactive_yaw_limit) + + buttons = ttk.Frame(main) + buttons.grid(row=2, column=0, columnspan=3, sticky="ew", pady=(6, 10)) + ttk.Button(buttons, text="Zero", command=self.zero).pack(side="left", padx=(0, 6)) + ttk.Button(buttons, text="Quit", command=self.close).pack(side="left") + + metrics = ttk.LabelFrame(main, text="State", padding=8) + metrics.grid(row=3, column=0, columnspan=3, sticky="nsew") + metrics.columnconfigure(1, weight=1) + for row, key in enumerate( + [ + "cmd", + "pos", + "rpy", + "body_vel", + "imu", + "wheel", + "wheel_air", + "leg", + "limits", + ] + ): + ttk.Label(metrics, text=key).grid(row=row, column=0, sticky="w", padx=(0, 8)) + var = tk.StringVar(value="-") + self.status_vars[key] = var + ttk.Label(metrics, textvariable=var, width=48).grid(row=row, column=1, sticky="w") + + def _add_slider(self, parent: object, row: int, label: str, var: object, limit: float) -> None: + ttk = __import__("tkinter.ttk").ttk + ttk.Label(parent, text=label).grid(row=row, column=0, sticky="w") + scale = ttk.Scale(parent, from_=-limit, to=limit, variable=var, orient="horizontal") + scale.grid(row=row, column=1, sticky="ew", padx=8) + value = ttk.Label(parent, width=8) + value.grid(row=row, column=2, sticky="e") + + def refresh_value(*_: object) -> None: + value.configure(text=f"{float(var.get()): .2f}") + + var.trace_add("write", refresh_value) + refresh_value() + + def zero(self) -> None: + self.vx.set(0.0) + self.yaw.set(0.0) + + def close(self) -> None: + self.alive = False + + def command(self) -> IkCommand: + return IkCommand(float(self.vx.get()), float(self.yaw.get())) + + def update_events(self) -> bool: + if not self.alive: + return False + try: + self.root.update_idletasks() + self.root.update() + except self.tk.TclError: + self.alive = False + return self.alive + + def update_metrics( + self, + sensor: SensorFrame, + data: mujoco.MjData, + ids: ModelIds, + geometry: WheelGeometry, + wheel_target: np.ndarray, + leg_target: np.ndarray, + metrics: Metrics, + args: argparse.Namespace, + ) -> None: + if sensor.time - self.last_update < args.ui_update_dt: + return + self.last_update = sensor.time + + rot = quat_to_mat_wxyz(sensor.quat_wxyz) + body_vel = rot.T @ data.qvel[:3] + wheel_air = data.xipos[ids.wheel_body, 2] - geometry.radius + wheel_err = float(np.mean(np.abs(sensor.wheel_vel - wheel_target))) + leg_err = float(np.max(np.abs(sensor.joint_pos[:12] - leg_target))) + tilt = math.hypot(float(sensor.rpy[0]), float(sensor.rpy[1])) + + values = { + "cmd": f"x={self.vx.get(): .2f} m/s, yaw={self.yaw.get(): .2f} rad/s", + "pos": f"x={sensor.base_pos[0]: .3f}, y={sensor.base_pos[1]: .3f}, z={sensor.base_pos[2]: .3f}", + "rpy": ( + f"roll={math.degrees(sensor.rpy[0]): .2f}, " + f"pitch={math.degrees(sensor.rpy[1]): .2f}, " + f"yaw={math.degrees(sensor.rpy[2]): .2f} deg" + ), + "body_vel": f"vx={body_vel[0]: .3f}, vy={body_vel[1]: .3f}, vz={body_vel[2]: .3f} m/s", + "imu": f"gyro_z={sensor.gyro[2]: .3f} rad/s, tilt={math.degrees(tilt): .2f} deg", + "wheel": ( + f"target=[{', '.join(f'{v: .1f}' for v in wheel_target)}], " + f"err={wheel_err: .3f} rad/s" + ), + "wheel_air": f"min={np.min(wheel_air): .4f}, max={np.max(wheel_air): .4f} m", + "leg": f"encoder_err={leg_err: .4f} rad, max={metrics.max_leg_error: .4f} rad", + "limits": ( + f"max_tilt={math.degrees(metrics.max_tilt): .2f} deg, " + f"max_wheel_speed={metrics.max_wheel_speed: .2f} rad/s" + ), + } + for key, value in values.items(): + self.status_vars[key].set(value) + + def destroy(self) -> None: + try: + self.root.destroy() + except self.tk.TclError: + pass + + +def run_one(test_name: str, args: argparse.Namespace) -> dict: + tmp_root = Path(args.tmp_dir) if args.tmp_dir else Path(tempfile.gettempdir()) + with tempfile.TemporaryDirectory(prefix="ik_slalom_", dir=str(tmp_root)) as tmp: + model, data, ids = build_model(args, Path(tmp)) + default_dof_pos = reset_robot(model, data, ids, args) + geometry = infer_wheel_geometry(model, data, ids) + wheel_signs = np.array(args.wheel_signs, dtype=np.float64) + ik = DifferentialWheelIk(geometry, args.max_wheel_speed, wheel_signs, args.yaw_wheel_gain, args) + leg_ik = LegStanceIk(model, data, ids, args) if args.leg_ik else None + posture = TablePostureController(args, default_dof_pos) if leg_ik is None else None + path = waypoints_for_test(test_name, args) + if path and args.path_follower == "pure": + follower = PurePursuitFollower(path, args) + else: + follower = PathFollower(path, args.speed, args.max_yaw_rate, args.arrive_radius, args.allow_reverse) if path else None + + sensor0 = read_sensors(model, data, ids) + metrics = Metrics( + start_pos=sensor0.base_pos.copy(), + start_yaw=float(sensor0.rpy[2]), + start_time=float(sensor0.time), + ) + + control_steps = max(1, int(round(args.control_dt / model.opt.timestep))) + total_steps = int(round(args.duration / model.opt.timestep)) + wheel_target = np.zeros(4, dtype=np.float64) + leg_target = default_dof_pos[:12].copy() + reason = "timeout" + success = False + + viewer = mujoco.viewer.launch_passive(model, data) if args.viewer else None + panel = InteractivePanel(args) if args.interactive else None + try: + for step in range(total_steps): + if panel is not None and not panel.update_events(): + reason = "interactive panel closed" + break + if viewer is not None and not viewer.is_running(): + reason = "viewer closed" + break + + if step % control_steps == 0: + sensor = read_sensors(model, data, ids) + cmd = panel.command() if panel is not None else test_command(test_name, step * model.opt.timestep, sensor, follower, args) + raw_wheel_target = ik.solve(cmd, sensor) + if leg_ik is not None: + leg_target = leg_ik.compute(data, sensor) + elif posture is not None: + leg_target = posture.compute(sensor) + raw_wheel_target *= sensor_command_scale(sensor, leg_target, args) + max_wheel_delta = args.wheel_accel_limit * args.control_dt + wheel_target = wheel_target + np.clip(raw_wheel_target - wheel_target, -max_wheel_delta, max_wheel_delta) + data.ctrl[ids.ctrl[:12]] = leg_target + data.ctrl[ids.ctrl[12:]] = wheel_target + + if follower is not None and follower.done: + reason = "path complete" + success = True + break + + mujoco.mj_step(model, data) + if step % control_steps == 0: + sensor = read_sensors(model, data, ids) + metrics.update(sensor, model, data, ids, geometry, wheel_target, leg_target) + if panel is not None: + panel.update_metrics(sensor, data, ids, geometry, wheel_target, leg_target, metrics, args) + if viewer is not None: + viewer.sync() + + if viewer is not None and args.realtime: + time.sleep(model.opt.timestep * args.realtime_scale) + finally: + if panel is not None: + panel.destroy() + if viewer is not None: + viewer.close() + + final_sensor = read_sensors(model, data, ids) + reached = follower.idx if follower is not None else 0 + total = len(path) + + stable = ( + metrics.max_tilt < math.radians(args.max_tilt_deg) + and metrics.max_leg_error < args.max_leg_error + and metrics.max_wheel_air_height < args.max_wheel_air_height + ) + + if args.interactive: + success = stable + reason = reason if reason != "timeout" else "interactive duration elapsed" + elif test_name == "forward": + success = final_sensor.base_pos[0] - metrics.start_pos[0] > max(0.25, 0.35 * args.speed * args.duration) and stable + reason = "forward displacement ok" if success else reason + elif test_name == "yaw": + success = abs(wrap_pi(final_sensor.rpy[2] - metrics.start_yaw)) > max(0.15, 0.14 * abs(args.yaw_rate) * args.duration) and stable + reason = "yaw change ok" if success else reason + elif test_name == "arc": + delta = np.linalg.norm(final_sensor.base_pos[:2] - metrics.start_pos[:2]) + success = ( + delta > max(0.25, 0.25 * args.speed * args.duration) + and abs(wrap_pi(final_sensor.rpy[2] - metrics.start_yaw)) > 0.15 + and stable + ) + reason = "arc displacement and yaw ok" if success else reason + elif test_name in {"waypoint", "slalom"}: + success = follower is not None and follower.done and stable + reason = "path complete" if success else reason + + out = metrics.as_dict(final_sensor, reached, total, reason, success) + out.update( + { + "test": test_name, + "wheel_radius_m": round(geometry.radius, 5), + "track_width_m": round(geometry.track_width, 5), + "wheel_body_y_m": [round(float(v), 5) for v in geometry.wheel_y], + "wheel_signs": [float(v) for v in wheel_signs], + "wheel_model": args.wheel_model, + "linear_wheel_gain": float(args.linear_wheel_gain), + "direct_yaw_wheel_gain": float(args.direct_yaw_wheel_gain), + "yaw_wheel_gain": float(args.yaw_wheel_gain), + "stable": bool(stable), + "leg_ik": bool(args.leg_ik), + "posture": args.posture, + "body_height": float(args.body_height), + "imu_posture": bool(args.imu_posture), + "yaw_rate_kp": float(args.yaw_rate_kp), + "encoder_posture_kp": float(args.encoder_posture_kp), + "encoder_guard": bool(args.encoder_guard), + "imu_guard": bool(args.imu_guard), + "posture_gain": float(args.posture_gain), + "path_follower": args.path_follower, + "default_pose": ( + "rough_rl" + if args.posture == "rough" + else "custom" + if args.posture == "custom" + else "mujoco_sim_height_table" + ), + } + ) + return out + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--test", choices=("forward", "yaw", "arc", "waypoint", "slalom", "all"), default="forward") + parser.add_argument("--duration", type=float, default=8.0) + parser.add_argument("--settle", type=float, default=1.5) + parser.add_argument("--control-dt", type=float, default=0.02) + parser.add_argument("--speed", type=float, default=0.35) + parser.add_argument("--yaw-rate", type=float, default=0.25) + parser.add_argument("--max-yaw-rate", type=float, default=0.45) + parser.add_argument("--max-wheel-speed", type=float, default=10.0) + parser.add_argument("--wheel-model", choices=("diff", "direct"), default="diff") + 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("--yaw-wheel-gain", type=float, default=2.2) + parser.add_argument("--yaw-rate-kp", type=float, default=0.0) + parser.add_argument("--max-yaw-wheel-speed", type=float, default=1.65) + parser.add_argument("--wheel-accel-limit", type=float, default=28.0) + parser.add_argument("--max-tilt-deg", type=float, default=25.0) + parser.add_argument("--max-leg-error", type=float, default=0.65) + parser.add_argument("--max-wheel-air-height", type=float, default=0.08) + parser.add_argument("--posture", choices=("rough", "table", "custom"), default="rough") + parser.add_argument("--body-height", type=float, default=0.37) + parser.add_argument("--custom-abduction", type=float, default=0.0) + parser.add_argument("--custom-hip", type=float, default=0.55) + parser.add_argument("--custom-knee", type=float, default=-1.30) + parser.add_argument("--imu-posture", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--roll-comp-gain", type=float, default=0.35) + parser.add_argument("--pitch-comp-gain", type=float, default=0.35) + parser.add_argument("--encoder-posture-kp", type=float, default=0.15) + parser.add_argument("--encoder-posture-max", type=float, default=0.03) + parser.add_argument("--encoder-guard", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--encoder-guard-start", type=float, default=0.28) + parser.add_argument("--encoder-guard-stop", type=float, default=0.65) + parser.add_argument("--imu-guard", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--imu-guard-start-deg", type=float, default=12.0) + parser.add_argument("--imu-guard-stop-deg", type=float, default=28.0) + parser.add_argument("--leg-ik", action=argparse.BooleanOptionalAction, default=False) + parser.add_argument("--posture-gain", type=float, default=0.65) + parser.add_argument("--leg-ik-iterations", type=int, default=3) + parser.add_argument("--leg-ik-damping", type=float, default=0.02) + parser.add_argument("--leg-ik-step", type=float, default=0.035) + parser.add_argument("--max-leg-offset", type=float, default=0.28) + parser.add_argument("--leg-target-rate", type=float, default=0.035) + parser.add_argument("--leg-target-filter", type=float, default=0.45) + parser.add_argument("--arrive-radius", type=float, default=0.20) + parser.add_argument("--allow-reverse", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--path-follower", choices=("waypoint", "pure"), default="pure") + parser.add_argument("--lookahead-base", type=float, default=0.45) + parser.add_argument("--lookahead-time", type=float, default=0.35) + parser.add_argument("--lookahead-min", type=float, default=0.35) + parser.add_argument("--lookahead-max", type=float, default=0.90) + parser.add_argument("--max-lat-acc", type=float, default=0.45) + parser.add_argument("--min-path-speed", type=float, default=0.25) + parser.add_argument("--pure-search-segments", type=int, default=4) + parser.add_argument("--start-x", type=float, default=0.0) + parser.add_argument("--start-y", type=float, default=0.0) + parser.add_argument("--start-z", type=float, default=0.6) + parser.add_argument( + "--terrain-xml", + type=Path, + default=None, + help="Optional terrain scene XML. Omit it to use a generated flat ground scene.", + ) + parser.add_argument("--robot-xml", type=Path, default=PROJECT_ROOT / "mjcf" / "wheelleg.xml") + parser.add_argument("--tmp-dir", type=Path, default=Path("D:/tmp") if Path("D:/tmp").exists() else None) + parser.add_argument("--interactive", action="store_true", help="Open a live command/metrics panel and drive manually.") + parser.add_argument("--interactive-vx-limit", type=float, default=1.0) + parser.add_argument("--interactive-yaw-limit", type=float, default=1.0) + parser.add_argument("--ui-update-dt", type=float, default=0.10) + parser.add_argument("--viewer", action="store_true", help="Open MuJoCo passive viewer while running the test.") + parser.add_argument("--realtime", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--realtime-scale", type=float, default=1.0) + parser.add_argument( + "--wheel-signs", + type=float, + nargs=4, + default=[1.0, 1.0, 1.0, 1.0], + metavar=("FL", "FR", "RL", "RR"), + help="Per-wheel velocity sign multipliers in joint order.", + ) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + if args.interactive: + args.viewer = True + args.duration = max(args.duration, 3600.0) + tests = ["forward", "yaw", "arc", "waypoint", "slalom"] if args.test == "all" else [args.test] + results = [run_one(name, args) for name in tests] + print(json.dumps(results[0] if len(results) == 1 else results, indent=2, ensure_ascii=False)) + return 0 if all(r["success"] for r in results) else 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/05_software/train/rc_mjlab/sim2sim/interface/mujoco_io.py b/05_software/train/rc_mjlab/sim2sim/interface/mujoco_io.py index 276aa64..a7447f7 100644 --- a/05_software/train/rc_mjlab/sim2sim/interface/mujoco_io.py +++ b/05_software/train/rc_mjlab/sim2sim/interface/mujoco_io.py @@ -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. diff --git a/05_software/train/rc_mjlab/sim2sim/nav_route_sim2sim_check.py b/05_software/train/rc_mjlab/sim2sim/nav_route_sim2sim_check.py new file mode 100644 index 0000000..4425482 --- /dev/null +++ b/05_software/train/rc_mjlab/sim2sim/nav_route_sim2sim_check.py @@ -0,0 +1,2776 @@ +#!/usr/bin/env python3 +"""Batch sim2sim route validation for nav_tools waypoints. + +This script runs the MuJoCo robot with an ONNX locomotion policy and a simple +waypoint follower, then reports whether the route can be completed without +falling or entering avoid-region clearance. + +Example: + uv run python sim2sim/nav_route_sim2sim_check.py \ + --terrain-xml tools/nav_tools/xml/A.xml \ + --points tools/nav_tools/points/points_20260705_174627.json \ + --onnx model_6800.onnx \ + --start-yaw-offset-deg -180 \ + --heading-offset-deg 180 +""" + +from __future__ import annotations + +import argparse +import csv +import json +import math +import os +import pkgutil +import random +import sys +import time +import xml.etree.ElementTree as ET +from collections import deque +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import mujoco +import mujoco.viewer +import numpy as np +import onnxruntime as ort +import torch + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +SIM2SIM_DIR = Path(__file__).resolve().parent +if str(SIM2SIM_DIR) not in sys.path: + sys.path.insert(0, str(SIM2SIM_DIR)) +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from interface.mujoco_io import MuJoCoIO # noqa: E402 +from policy.policy_runner import PolicyRunner # noqa: E402 +from tools.nav_tools.route_safety_check import ( # noqa: E402 + AvoidRegion, + Waypoint, + default_lateral_footprint_radius, + load_regions, + load_waypoints, + point_in_polygon, + point_segment_distance, +) + + +ACTION_SCALE = np.array( + [0.125, 0.25, 0.25] * 4 + [5.0, 5.0, 5.0, 5.0], + dtype=np.float32, +) +DEPLOY_DEFAULT_DOF_POS = np.array([0.0, 0.550, -1.125] * 4 + [0.0] * 4, dtype=np.float32) +ROBOT_BODY_LENGTH = 0.356 +ROBOT_BODY_WIDTH = 0.235 +ROBOT_BODY_CENTER_X = 0.1518 +ROBOT_ORIGIN_FROM_FRONT = 0.105 +ROBOT_BODY_CENTER_OFFSET_X = ROBOT_ORIGIN_FROM_FRONT - ROBOT_BODY_LENGTH * 0.5 +ROBOT_WHEEL_VIS_LENGTH = 0.16 +ROBOT_WHEEL_VIS_WIDTH = 0.055 +ROBOT_WHEEL_POSITIONS = { + "fl": ((0.32826 + 0.06389) - ROBOT_BODY_CENTER_X, 0.066172 - 0.027344, 0.1035, 0.014699, 0.04074, 0.0), + "fr": ((0.32826 + 0.06389) - ROBOT_BODY_CENTER_X, -0.065853 + 0.027311, -0.1035, -0.018447, -0.040735, -0.00075079), + "rl": ((-0.024743 - 0.06389) - ROBOT_BODY_CENTER_X, 0.066141 - 0.027309, 0.099459, 0.012475, 0.040737, 0.0), + "rr": ((-0.024743 - 0.06389) - ROBOT_BODY_CENTER_X, -0.065884 + 0.027341, -0.099408, -0.012435, -0.040737, -0.00075079), +} + + +@dataclass +class SimConfig: + control_hz: float = 50.0 + max_vx: float = 1.2 + max_vy: float = 0.0 + max_wz: float = 0.8 + kp_dist: float = 0.8 + kp_yaw: float = 1.8 + yaw_stop_threshold_deg: float = 45.0 + turn_in_place_enter_deg: float = 70.0 + turn_in_place_exit_deg: float = 18.0 + turn_in_place_max_wz: float = 0.8 + final_align_max_wz: float = 0.45 + final_align_kp_scale: float = 0.6 + waypoint_timeout_s: float = 30.0 + max_total_time_s: float = 420.0 + stable_cycles: int = 2 + lookahead_m: float = 0.45 + min_cmd_vx: float = 0.08 + creep_cmd_vx: float = 0.04 + cmd_vx_scale: float = 1.0 + stuck_timeout_s: float = 4.0 + stuck_progress_m: float = 0.08 + recovery_duration_s: float = 2.0 + max_recoveries: int = 4 + slalom_script_enabled: bool = True + slalom_script_start_tolerance: float = 0.10 + slalom_script_pos_tolerance: float = 0.08 + slalom_script_yaw_tolerance_deg: float = 5.0 + slalom_script_stable_cycles: int = 1 + slalom_script_rotate_steps_enabled: bool = False + slalom_script_final_rotate_enabled: bool = False + slalom_script_require_yaw_at_step: bool = False + slalom_script_kp_dist: float = 1.0 + slalom_script_kp_yaw: float = 1.6 + slalom_script_max_vx: float = 0.60 + slalom_script_max_vy: float = 0.50 + slalom_script_max_wz: float = 0.50 + slalom_script_min_cmd_linear: float = 0.2 + slalom_script_min_cmd_angular: float = 0.0 + slalom_script_min_cmd_epsilon: float = 0.05 + slalom_script_min_step_distance: float = 0.02 + slalom_script_yaw_gate_deg: float = 25.0 + slalom_script_drive_yaw_source: str = "segment" + slalom_script_curvature_speed_enabled: bool = True + slalom_script_curvature_min_scale: float = 0.55 + + +@dataclass(frozen=True) +class SlalomScriptStep: + kind: str + start_index: int + end_index: int + target_x: float + target_y: float + target_yaw: float | None + pos_tolerance: float | None = None + forward: float = 0.0 + left: float = 0.0 + + +@dataclass +class SimResult: + success: bool + reason: str + sim_time: float + reached_count: int + waypoint_count: int + min_clearance: float + min_margin: float + min_clearance_region: str + min_clearance_wp: str + max_roll_deg: float + max_pitch_deg: float + max_tilt_deg: float + samples: int + wall_time: float + real_time_factor: float + + +@dataclass +class RouteSnapshot: + time: float + qpos: np.ndarray + qvel: np.ndarray + ctrl: np.ndarray + follower_index: int + follower_turn_in_place: bool + follower_stable_count: int + follower_wp_start_time: float + follower_best_dist: float | None + follower_last_progress_time: float | None + follower_recovery_until: float | None + follower_recovery_count: int | None + follower_recovery_turn_sign: float | None + + +class OnnxPolicy: + def __init__(self, path: Path) -> None: + providers = ["CPUExecutionProvider"] + self.session = ort.InferenceSession(str(path), providers=providers) + self.input_name = self.session.get_inputs()[0].name + self.output_name = self.session.get_outputs()[0].name + shape = self.session.get_inputs()[0].shape + dim = shape[1] if len(shape) >= 2 else 53 + self.obs_dim = 53 if isinstance(dim, str) else int(dim) + self.last_actions = np.zeros(16, dtype=np.float32) + + def reset(self) -> None: + self.last_actions[:] = 0.0 + + def step(self, obs_53d: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + if self.obs_dim == obs_53d.shape[0]: + model_input = obs_53d + else: + raise RuntimeError( + f"ONNX expects obs dim {self.obs_dim}, but this checker provides {obs_53d.shape[0]}" + ) + raw = self.session.run( + [self.output_name], + {self.input_name: model_input.astype(np.float32)[None, :]}, + )[0][0].astype(np.float32) + raw = np.clip(raw, -100.0, 100.0) + self.last_actions = raw.copy() + return raw * ACTION_SCALE, raw + + +class ExistingPolicyBackend: + """Adapter around sim2sim/policy/PolicyRunner to match nav_sim2sim.py.""" + + def __init__(self, path: Path, crawl_path: Path | None = None) -> None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + policy_path: Path | dict[str, Path] + if crawl_path is not None: + policy_path = {"rough": path, "crawl": crawl_path} + else: + policy_path = path + self.runner = PolicyRunner(policy_path, device) + self.runner.reset() + + @property + def default_dof_pos(self) -> np.ndarray: + return self.runner.default_dof_pos + + @property + def last_actions(self) -> np.ndarray: + return self.runner.last_actions + + def reset(self) -> None: + self.runner.reset() + + def maybe_switch_policy(self, policy_name: str | None) -> None: + if not policy_name: + return + if policy_name == self.runner.current_policy_name: + return + if self.runner.transition_in_progress and policy_name == self.runner.current_policy_name: + return + if policy_name in self.runner.policies: + self.runner.trigger_transition(policy_name) + + def step(self, obs_53d: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + return self.runner.step(obs_53d) + + def close(self) -> None: + listener = getattr(self.runner, "listener", None) + if listener is not None: + listener.stop() + + +class RoughPolicyIkCrawlBackend(ExistingPolicyBackend): + def __init__(self, path: Path) -> None: + super().__init__(path, None) + self.mode = "rough" + self.requested_mode = "rough" + self.crawl_default_dof_pos = np.array( + [ + 0.2, 1.697, -2.650, + -0.2, 1.697, -2.650, + 0.2, 1.697, -2.650, + -0.2, 1.697, -2.650, + 0.0, 0.0, 0.0, 0.0, + ], + dtype=np.float32, + ) + self.slalom_ik_default_dof_pos = np.array( + [ + 0.0, 1.697, -2.650, + 0.0, 1.697, -2.650, + 0.0, 1.697, -2.650, + 0.0, 1.697, -2.650, + 0.0, 0.0, 0.0, 0.0, + ], + dtype=np.float32, + ) + self._last_actions = np.zeros(16, dtype=np.float32) + self.crawl_ik_wheel_linear_gain = 6.25 + self.crawl_ik_wheel_yaw_gain = 4.0 + self.crawl_ik_max_wheel_speed = 6.0 + self.transition_steps = 60 + self.transition_step = self.transition_steps + self.transition_start_pose = self.runner.default_dof_pos.copy() + self.transition_target_pose = self.runner.default_dof_pos.copy() + + def reset(self) -> None: + self.runner.reset() + self.mode = "rough" + self.requested_mode = "rough" + self.transition_step = self.transition_steps + self.transition_start_pose = self.runner.default_dof_pos.copy() + self.transition_target_pose = self.runner.default_dof_pos.copy() + self._last_actions[:] = 0.0 + + @property + def default_dof_pos(self) -> np.ndarray: + if self.transition_step < self.transition_steps: + alpha = self.transition_step / max(1, self.transition_steps) + return (1.0 - alpha) * self.transition_start_pose + alpha * self.transition_target_pose + if self.requested_mode == "crawl": + return self.crawl_default_dof_pos + if self.requested_mode == "slalom_ik": + return self.slalom_ik_default_dof_pos + return self.runner.default_dof_pos + + @property + def last_actions(self) -> np.ndarray: + if self.requested_mode in {"crawl", "slalom_ik"}: + return self._last_actions + return self.runner.last_actions + + def maybe_switch_policy(self, policy_name: str | None) -> None: + normalized = (policy_name or "").lower() + if normalized in {"slalom", "slalom_ik", "slalomik"}: + next_mode = "slalom_ik" + elif normalized in {"crawl", "ik"}: + next_mode = "crawl" + else: + next_mode = "rough" + if next_mode == self.requested_mode: + return + self.transition_start_pose = self.default_dof_pos.copy() + if next_mode == "crawl": + self.transition_target_pose = self.crawl_default_dof_pos.copy() + elif next_mode == "slalom_ik": + self.transition_target_pose = self.slalom_ik_default_dof_pos.copy() + else: + self.transition_target_pose = self.runner.default_dof_pos.copy() + self.transition_step = 0 + self.requested_mode = next_mode + + def step(self, obs_53d: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + in_transition = self.transition_step < self.transition_steps + if in_transition: + self.transition_step += 1 + if self.requested_mode not in {"crawl", "slalom_ik"}: + return self.runner.step(obs_53d) + cmd = obs_53d[6:9].astype(np.float32) + target = ( + self.slalom_ik_default_dof_pos.copy() + if self.requested_mode == "slalom_ik" + else self.crawl_default_dof_pos.copy() + ) + left_wheel = float(np.clip( + cmd[0] * self.crawl_ik_wheel_linear_gain - cmd[2] * self.crawl_ik_wheel_yaw_gain, + -self.crawl_ik_max_wheel_speed, + self.crawl_ik_max_wheel_speed, + )) + right_wheel = float(np.clip( + cmd[0] * self.crawl_ik_wheel_linear_gain + cmd[2] * self.crawl_ik_wheel_yaw_gain, + -self.crawl_ik_max_wheel_speed, + self.crawl_ik_max_wheel_speed, + )) + alpha = 1.0 + if in_transition: + alpha = min(1.0, self.transition_step / max(1, self.transition_steps)) + scaled = target.copy() + scaled[12:] = [left_wheel, right_wheel, left_wheel, right_wheel] + scaled[12:] *= alpha + raw = np.zeros(16, dtype=np.float32) + raw[12:] = (np.array([left_wheel, right_wheel, left_wheel, right_wheel], dtype=np.float32) * alpha) / 5.0 + self._last_actions = raw + # send_actions adds default_dof_pos to scaled actions, so return leg offsets. + scaled[:12] = 0.0 + return scaled, raw + + +class DirectOnnxBackend: + """Lightweight ONNX adapter kept for isolating PolicyRunner effects.""" + + def __init__(self, path: Path, default_dof_pos: np.ndarray) -> None: + self.policy = OnnxPolicy(path) + self._default_dof_pos = default_dof_pos.astype(np.float32) + + @property + def default_dof_pos(self) -> np.ndarray: + return self._default_dof_pos + + @property + def last_actions(self) -> np.ndarray: + return self.policy.last_actions + + def reset(self) -> None: + self.policy.reset() + + def maybe_switch_policy(self, policy_name: str | None) -> None: + _ = policy_name + + def step(self, obs_53d: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + return self.policy.step(obs_53d) + + def close(self) -> None: + pass + + +class WaypointFollower: + def __init__(self, waypoints: list[Waypoint], cfg: SimConfig) -> None: + self.waypoints = waypoints + self.cfg = cfg + self.index = 0 + self.turn_in_place = False + self.stable_count = 0 + self.wp_start_time = 0.0 + + @property + def active(self) -> Waypoint | None: + if self.index >= len(self.waypoints): + return None + return self.waypoints[self.index] + + def reset_timing(self, sim_time: float) -> None: + self.wp_start_time = sim_time + + @staticmethod + def normalize_angle(value: float) -> float: + while value > math.pi: + value -= 2.0 * math.pi + while value < -math.pi: + value += 2.0 * math.pi + return value + + @staticmethod + def clamp(value: float, lo: float, hi: float) -> float: + return max(lo, min(hi, value)) + + def update(self, x: float, y: float, yaw: float, sim_time: float) -> tuple[np.ndarray, str | None]: + wp = self.active + if wp is None: + return np.zeros(3, dtype=np.float32), "complete" + + if sim_time - self.wp_start_time > self.cfg.waypoint_timeout_s: + return np.zeros(3, dtype=np.float32), f"timeout at waypoint {wp.id}" + + dx = wp.x - x + dy = wp.y - y + dist = math.hypot(dx, dy) + tolerance = wp.tolerance if wp.tolerance is not None else 0.15 + if dist <= tolerance: + self.stable_count += 1 + if self.stable_count >= self.cfg.stable_cycles: + self.index += 1 + self.turn_in_place = False + self.stable_count = 0 + self.wp_start_time = sim_time + if self.index >= len(self.waypoints): + return np.zeros(3, dtype=np.float32), "complete" + return np.zeros(3, dtype=np.float32), None + return np.zeros(3, dtype=np.float32), None + self.stable_count = 0 + + target_yaw = math.atan2(dy, dx) + yaw_err = self.normalize_angle(target_yaw - yaw) + yaw_stop_threshold = math.radians(self.cfg.yaw_stop_threshold_deg) + + if self.turn_in_place: + if abs(yaw_err) <= math.radians(self.cfg.turn_in_place_exit_deg): + self.turn_in_place = False + elif abs(yaw_err) >= math.radians(self.cfg.turn_in_place_enter_deg): + self.turn_in_place = True + + if self.turn_in_place: + return np.array( + [ + 0.0, + 0.0, + self.clamp( + self.cfg.kp_yaw * yaw_err, + -self.cfg.turn_in_place_max_wz, + self.cfg.turn_in_place_max_wz, + ), + ], + dtype=np.float32, + ), None + + speed_limit = self.cfg.max_vx if wp.speed is None else min(self.cfg.max_vx, max(0.0, wp.speed)) + cmd = np.zeros(3, dtype=np.float32) + cmd[2] = self.clamp(self.cfg.kp_yaw * yaw_err, -self.cfg.max_wz, self.cfg.max_wz) + if abs(yaw_err) <= yaw_stop_threshold: + cmd[0] = self.clamp(self.cfg.kp_dist * dist * math.cos(yaw_err), 0.0, speed_limit) + cmd[0] *= self.cfg.cmd_vx_scale + return cmd, None + + +class PurePursuitFollower(WaypointFollower): + def __init__(self, waypoints: list[Waypoint], cfg: SimConfig) -> None: + super().__init__(waypoints, cfg) + self.best_dist = float("inf") + self.last_progress_time = 0.0 + self.recovery_until = -1.0 + self.recovery_count = 0 + self.recovery_turn_sign = 1.0 + + def reset_timing(self, sim_time: float) -> None: + super().reset_timing(sim_time) + self.best_dist = float("inf") + self.last_progress_time = sim_time + self.recovery_until = -1.0 + self.recovery_count = 0 + + def _advance_reached(self, x: float, y: float, sim_time: float) -> str | None: + while self.active is not None: + wp = self.active + if wp.require_yaw: + break + tol = wp.tolerance if wp.tolerance is not None else 0.15 + if wp.mandatory_cross: + tol = min(tol, 0.15 if wp.mandatory_radius is None else max(0.01, wp.mandatory_radius)) + completion_x = wp.x if wp.mandatory_center_x is None else wp.mandatory_center_x + completion_y = wp.y if wp.mandatory_center_y is None else wp.mandatory_center_y + if math.hypot(completion_x - x, completion_y - y) > tol: + break + self.index += 1 + self.turn_in_place = False + self.stable_count = 0 + self.wp_start_time = sim_time + self.best_dist = float("inf") + self.last_progress_time = sim_time + if self.index >= len(self.waypoints): + return "complete" + return None + + def _lookahead_target(self, x: float, y: float) -> Waypoint: + active = self.active + if active is None: + return self.waypoints[-1] + if active.exact_reach: + return active + target = active + for wp in self.waypoints[self.index :]: + target = wp + if wp.mandatory_cross: + mandatory_radius = 0.15 if wp.mandatory_radius is None else max(0.01, wp.mandatory_radius) + center_x = wp.x if wp.mandatory_center_x is None else wp.mandatory_center_x + center_y = wp.y if wp.mandatory_center_y is None else wp.mandatory_center_y + if math.hypot(center_x - x, center_y - y) > mandatory_radius: + break + if math.hypot(wp.x - x, wp.y - y) >= self.cfg.lookahead_m: + break + return target + + def _current_speed_limit(self) -> float: + speed = self.cfg.max_vx + end = min(len(self.waypoints), self.index + 4) + for wp in self.waypoints[self.index : end]: + if wp.speed is not None: + speed = min(speed, max(0.0, wp.speed)) + return speed + + def update(self, x: float, y: float, yaw: float, sim_time: float) -> tuple[np.ndarray, str | None]: + nav_state = self._advance_reached(x, y, sim_time) + if nav_state is not None: + return np.zeros(3, dtype=np.float32), nav_state + + active = self.active + if active is None: + return np.zeros(3, dtype=np.float32), "complete" + active_dist = math.hypot(active.x - x, active.y - y) + if active.require_yaw and active_dist <= (active.tolerance if active.tolerance is not None else 0.15): + yaw_tolerance_deg = 5.0 if active.yaw_tolerance_deg is None else max(0.0, active.yaw_tolerance_deg) + yaw_err = self.normalize_angle(math.radians(active.yaw_deg) - yaw) + if abs(yaw_err) > math.radians(yaw_tolerance_deg): + self.stable_count = 0 + return np.array( + [ + 0.0, + 0.0, + self.clamp( + self.cfg.kp_yaw * self.cfg.final_align_kp_scale * yaw_err, + -self.cfg.final_align_max_wz, + self.cfg.final_align_max_wz, + ), + ], + dtype=np.float32, + ), None + required_stable_cycles = ( + self.cfg.stable_cycles if active.stable_cycles is None else max(1, active.stable_cycles) + ) + self.stable_count += 1 + if self.stable_count < required_stable_cycles: + return np.zeros(3, dtype=np.float32), None + self.index += 1 + self.turn_in_place = False + self.stable_count = 0 + self.wp_start_time = sim_time + self.best_dist = float("inf") + self.last_progress_time = sim_time + if self.index >= len(self.waypoints): + return np.zeros(3, dtype=np.float32), "complete" + return np.zeros(3, dtype=np.float32), None + if active_dist + self.cfg.stuck_progress_m < self.best_dist: + self.best_dist = active_dist + self.last_progress_time = sim_time + + if self.recovery_until > sim_time: + return np.array( + [ + -0.10 * abs(self.cfg.cmd_vx_scale), + 0.0, + 0.45 * self.recovery_turn_sign, + ], + dtype=np.float32, + ), None + + if sim_time - self.wp_start_time > self.cfg.waypoint_timeout_s: + return np.zeros(3, dtype=np.float32), f"timeout at waypoint {active.id}" + + if sim_time - self.last_progress_time > self.cfg.stuck_timeout_s: + if active.require_yaw: + self.last_progress_time = sim_time + elif self.recovery_count >= self.cfg.max_recoveries: + return np.zeros(3, dtype=np.float32), f"stuck near waypoint {active.id}" + else: + self.recovery_count += 1 + self.recovery_turn_sign *= -1.0 + self.recovery_until = sim_time + self.cfg.recovery_duration_s + self.last_progress_time = sim_time + return np.array([-0.10 * abs(self.cfg.cmd_vx_scale), 0.0, 0.45 * self.recovery_turn_sign], dtype=np.float32), None + + target = self._lookahead_target(x, y) + dx = target.x - x + dy = target.y - y + target_dist = math.hypot(dx, dy) + target_yaw = math.atan2(dy, dx) + yaw_err = self.normalize_angle(target_yaw - yaw) + + cmd = np.zeros(3, dtype=np.float32) + cmd[2] = self.clamp(self.cfg.kp_yaw * yaw_err, -self.cfg.max_wz, self.cfg.max_wz) + if abs(yaw_err) > math.radians(self.cfg.turn_in_place_enter_deg): + cmd[0] = 0.0 + return cmd, None + + speed_limit = self._current_speed_limit() + curvature_slow = max(0.35, 1.0 - abs(yaw_err) / math.radians(80.0)) + cmd_vx = self.cfg.kp_dist * target_dist * math.cos(yaw_err) + cmd_vx = self.clamp(cmd_vx, self.cfg.min_cmd_vx, speed_limit * curvature_slow) + if abs(yaw_err) > math.radians(self.cfg.yaw_stop_threshold_deg): + if abs(yaw_err) < math.radians(self.cfg.turn_in_place_enter_deg): + cmd_vx = min(self.cfg.creep_cmd_vx, speed_limit) + else: + cmd_vx = 0.0 + cmd[0] = cmd_vx * self.cfg.cmd_vx_scale + return cmd, None + + +class NavGoodFollower(WaypointFollower): + """Follower shaped after sim2real nav_good/simple_nav_node.py. + + This keeps the route as the active path, advances through near path points, + selects a lookahead target, then applies the same yaw gating and + turn-in-place structure as the ROS2 simple_nav node. + """ + + def __init__(self, waypoints: list[Waypoint], cfg: SimConfig) -> None: + super().__init__(waypoints, cfg) + self.goal_entered_tolerance = False + self.goal_complete_stable_count = 0 + self.goal_exit_tolerance_margin = 0.05 + self.final_align_kp_yaw_scale = 0.6 + self.final_align_creep_speed = 0.03 + self.path_reach_dist = 0.18 + self.best_dist = float("inf") + self.last_progress_time = 0.0 + + def reset_timing(self, sim_time: float) -> None: + super().reset_timing(sim_time) + self.goal_entered_tolerance = False + self.goal_complete_stable_count = 0 + self.turn_in_place = False + self.best_dist = float("inf") + self.last_progress_time = sim_time + + def _current_goal(self) -> Waypoint | None: + return self.active + + def _is_charge_segment(self, x: float, y: float, goal: Waypoint) -> bool: + policy = (goal.policy or "").lower() + if policy in {"charge", "wall", "obstacle", "climb"}: + return True + return -0.45 <= x <= 0.35 and -4.35 <= y <= -2.85 + + def _point_is_behind_segment( + self, + x: float, + y: float, + point: Waypoint, + next_point: Waypoint, + ) -> bool: + seg_x = next_point.x - point.x + seg_y = next_point.y - point.y + seg_len_sq = seg_x * seg_x + seg_y * seg_y + if seg_len_sq <= 1.0e-9: + return False + proj = ((x - point.x) * seg_x + (y - point.y) * seg_y) / seg_len_sq + return proj > 0.65 and math.hypot(next_point.x - x, next_point.y - y) < math.hypot(point.x - x, point.y - y) + + def _advance_path_nodes(self, x: float, y: float, sim_time: float) -> str | None: + while self.index < len(self.waypoints) - 1: + goal = self.waypoints[self.index] + if goal.require_yaw: + break + next_goal = self.waypoints[self.index + 1] + tol = max(self.path_reach_dist, goal.tolerance if goal.tolerance is not None else 0.15) + dist = math.hypot(goal.x - x, goal.y - y) + if dist > tol and not self._point_is_behind_segment(x, y, goal, next_goal): + break + self.index += 1 + self.wp_start_time = sim_time + self.turn_in_place = False + self.goal_entered_tolerance = False + self.goal_complete_stable_count = 0 + self.best_dist = float("inf") + self.last_progress_time = sim_time + if self.index >= len(self.waypoints): + return "complete" + return None + + def _path_follow_target(self, x: float, y: float, goal: Waypoint) -> Waypoint: + target_index = self.index + while target_index < len(self.waypoints) - 1: + candidate = self.waypoints[target_index + 1] + if math.hypot(candidate.x - x, candidate.y - y) > self.cfg.lookahead_m: + break + target_index += 1 + return self.waypoints[target_index] + + def update(self, x: float, y: float, yaw: float, sim_time: float) -> tuple[np.ndarray, str | None]: + nav_state = self._advance_path_nodes(x, y, sim_time) + if nav_state is not None: + return np.zeros(3, dtype=np.float32), nav_state + + goal = self._current_goal() + if goal is None: + return np.zeros(3, dtype=np.float32), "complete" + if sim_time - self.wp_start_time > self.cfg.waypoint_timeout_s: + return np.zeros(3, dtype=np.float32), f"timeout at waypoint {goal.id}" + + goal_dx = goal.x - x + goal_dy = goal.y - y + dist = math.hypot(goal_dx, goal_dy) + tol_enter = goal.tolerance if goal.tolerance is not None else 0.15 + tol_exit = tol_enter + self.goal_exit_tolerance_margin + position_ready = dist < (tol_exit if self.goal_entered_tolerance else tol_enter) + max_vx = self.cfg.max_vx if goal.speed is None else min(self.cfg.max_vx, max(0.0, goal.speed)) + + if goal.require_yaw and position_ready: + self.goal_entered_tolerance = True + yaw_tolerance_deg = 5.0 if goal.yaw_tolerance_deg is None else max(0.0, goal.yaw_tolerance_deg) + yaw_err = self.normalize_angle(math.radians(goal.yaw_deg) - yaw) + if abs(yaw_err) > math.radians(yaw_tolerance_deg): + self.goal_complete_stable_count = 0 + return np.array( + [ + 0.0, + 0.0, + self.clamp( + self.cfg.kp_yaw * self.final_align_kp_yaw_scale * yaw_err, + -self.cfg.final_align_max_wz, + self.cfg.final_align_max_wz, + ), + ], + dtype=np.float32, + ), None + required_stable_cycles = ( + self.cfg.stable_cycles if goal.stable_cycles is None else max(1, goal.stable_cycles) + ) + self.goal_complete_stable_count += 1 + if self.goal_complete_stable_count < required_stable_cycles: + return np.zeros(3, dtype=np.float32), None + self.index += 1 + self.wp_start_time = sim_time + self.turn_in_place = False + self.goal_entered_tolerance = False + self.goal_complete_stable_count = 0 + self.best_dist = float("inf") + self.last_progress_time = sim_time + if self.index >= len(self.waypoints): + return np.zeros(3, dtype=np.float32), "complete" + return np.zeros(3, dtype=np.float32), None + + is_final_goal = self.index >= len(self.waypoints) - 1 + if is_final_goal and position_ready: + self.goal_entered_tolerance = True + self.goal_complete_stable_count += 1 + if self.goal_complete_stable_count < self.cfg.stable_cycles: + return np.zeros(3, dtype=np.float32), None + self.index += 1 + self.wp_start_time = sim_time + self.turn_in_place = False + self.goal_entered_tolerance = False + self.goal_complete_stable_count = 0 + if self.index >= len(self.waypoints): + return np.zeros(3, dtype=np.float32), "complete" + return np.zeros(3, dtype=np.float32), None + + self.goal_entered_tolerance = False + self.goal_complete_stable_count = 0 + if dist + self.cfg.stuck_progress_m < self.best_dist: + self.best_dist = dist + self.last_progress_time = sim_time + if sim_time - self.last_progress_time > self.cfg.waypoint_timeout_s: + return np.zeros(3, dtype=np.float32), f"no progress near waypoint {goal.id}" + + target = goal if goal.require_yaw else self._path_follow_target(x, y, goal) + dx = target.x - x + dy = target.y - y + target_yaw = math.atan2(dy, dx) + yaw_err = self.normalize_angle(target_yaw - yaw) + + if self.turn_in_place: + if abs(yaw_err) <= math.radians(self.cfg.turn_in_place_exit_deg): + self.turn_in_place = False + elif abs(yaw_err) >= math.radians(self.cfg.turn_in_place_enter_deg): + self.turn_in_place = True + + if self.turn_in_place: + return np.array( + [ + 0.0, + 0.0, + self.clamp(self.cfg.kp_yaw * yaw_err, -self.cfg.turn_in_place_max_wz, self.cfg.turn_in_place_max_wz), + ], + dtype=np.float32, + ), None + + cmd = np.zeros(3, dtype=np.float32) + cmd[2] = self.clamp(self.cfg.kp_yaw * yaw_err, -self.cfg.max_wz, self.cfg.max_wz) + if abs(yaw_err) <= math.radians(self.cfg.yaw_stop_threshold_deg): + path_dist = math.hypot(dx, dy) + if self._is_charge_segment(x, y, goal): + cmd[0] = max_vx + else: + cmd[0] = self.clamp(self.cfg.kp_dist * path_dist * math.cos(yaw_err), 0.0, max_vx) + cmd[0] *= self.cfg.cmd_vx_scale + return cmd, None + + +class NavSim2SimFollower(WaypointFollower): + """Follower matched to sim2sim/nav_sim2sim.py NavController. + + It tracks only the current waypoint, switches at 0.15 m, stops forward + motion above ~45 deg yaw error, and can keep full waypoint speed through + obstacle-charge segments instead of scaling speed down by distance. + """ + + def __init__(self, waypoints: list[Waypoint], cfg: SimConfig) -> None: + super().__init__(waypoints, cfg) + self.reach_dist = 0.15 + + def _advance_reached(self, x: float, y: float, sim_time: float) -> str | None: + while self.active is not None: + wp = self.active + tol = wp.tolerance if wp.tolerance is not None else self.reach_dist + tol = max(self.reach_dist, tol) + if math.hypot(wp.x - x, wp.y - y) >= tol: + break + if self.index < len(self.waypoints) - 1: + self.index += 1 + self.turn_in_place = False + self.stable_count = 0 + self.wp_start_time = sim_time + return None + self.index += 1 + return "complete" + if self.index >= len(self.waypoints): + return "complete" + return None + + def _is_charge_segment(self, x: float, y: float, wp: Waypoint) -> bool: + policy = (wp.policy or "").lower() + if policy in {"charge", "wall", "obstacle", "climb"}: + return True + # A.xml has a low wall/step around x=0, y=-3.5. nav_sim2sim keeps + # speed through comparable obstacle sections instead of distance taper. + return -0.45 <= x <= 0.35 and -4.35 <= y <= -2.85 + + def update(self, x: float, y: float, yaw: float, sim_time: float) -> tuple[np.ndarray, str | None]: + nav_state = self._advance_reached(x, y, sim_time) + if nav_state is not None: + return np.zeros(3, dtype=np.float32), nav_state + wp = self.active + if wp is None: + return np.zeros(3, dtype=np.float32), "complete" + if sim_time - self.wp_start_time > self.cfg.waypoint_timeout_s: + return np.zeros(3, dtype=np.float32), f"timeout at waypoint {wp.id}" + + dx = wp.x - x + dy = wp.y - y + dist = math.hypot(dx, dy) + target_yaw = math.atan2(dy, dx) + yaw_err = self.normalize_angle(target_yaw - yaw) + speed_limit = self.cfg.max_vx if wp.speed is None else min(self.cfg.max_vx, max(0.0, wp.speed)) + + cmd = np.zeros(3, dtype=np.float32) + cmd[2] = self.clamp(self.cfg.kp_yaw * yaw_err, -self.cfg.max_wz, self.cfg.max_wz) + if abs(yaw_err) > 0.8: + cmd[0] = 0.0 + elif self._is_charge_segment(x, y, wp): + cmd[0] = speed_limit + else: + cmd[0] = self.clamp(self.cfg.kp_dist * dist * math.cos(yaw_err), -0.2, speed_limit) + cmd[0] *= self.cfg.cmd_vx_scale + return cmd, None + + +class SlalomScriptFollower(NavGoodFollower): + """NavGood follower plus the sim2real slalomStraight odometry script. + + The script starts only after the robot reaches the first slalomStraight + point, then drives consecutive slalomStraight points with body-frame + vx/vy/wz feedback. This mirrors simple_nav_node.py instead of the older + pure waypoint followers. + """ + + def __init__(self, waypoints: list[Waypoint], cfg: SimConfig) -> None: + super().__init__(waypoints, cfg) + self.slalom_script_active = False + self.slalom_script_start_index = 0 + self.slalom_script_end_index = 0 + self.slalom_script_steps: list[SlalomScriptStep] = [] + self.slalom_script_step_index = 0 + self.slalom_script_step_stable_count = 0 + + def reset_timing(self, sim_time: float) -> None: + super().reset_timing(sim_time) + self.reset_slalom_script() + + def reset_slalom_script(self) -> None: + self.slalom_script_active = False + self.slalom_script_start_index = 0 + self.slalom_script_end_index = 0 + self.slalom_script_steps = [] + self.slalom_script_step_index = 0 + self.slalom_script_step_stable_count = 0 + + @staticmethod + def _is_slalom_script_waypoint(wp: Waypoint | None) -> bool: + return bool(wp is not None and getattr(wp, "slalom_straight", False)) + + def _advance_path_nodes(self, x: float, y: float, sim_time: float) -> str | None: + active = self.active + if ( + self.cfg.slalom_script_enabled + and self._is_slalom_script_waypoint(active) + and not self.slalom_script_active + and self.index + 1 < len(self.waypoints) + and self._is_slalom_script_waypoint(self.waypoints[self.index + 1]) + ): + # Do not let the normal path follower skip the script start point at + # the wider path tolerance; the real node gates this with + # nav_slalom_script_start_tolerance. + return None + return super()._advance_path_nodes(x, y, sim_time) + + def get_slalom_script_bounds(self, start_index: int) -> tuple[int, int]: + if start_index < 0 or start_index >= len(self.waypoints): + return start_index, start_index + if not self._is_slalom_script_waypoint(self.waypoints[start_index]): + return start_index, start_index + end_index = start_index + while end_index + 1 < len(self.waypoints): + if getattr(self.waypoints[end_index], "slalom_script_break", False): + break + if not self._is_slalom_script_waypoint(self.waypoints[end_index + 1]): + break + end_index += 1 + return start_index, end_index + + def build_slalom_script_steps(self, start_index: int, end_index: int) -> list[SlalomScriptStep]: + if end_index <= start_index: + return [] + + steps: list[SlalomScriptStep] = [] + for index in range(start_index, end_index): + goal = self.waypoints[index] + next_goal = self.waypoints[index + 1] + route_dx = next_goal.x - goal.x + route_dy = next_goal.y - goal.y + route_dist = math.hypot(route_dx, route_dy) + route_yaw = math.atan2(route_dy, route_dx) if route_dist > 1.0e-6 else math.radians(goal.yaw_deg) + yaw_source = self.cfg.slalom_script_drive_yaw_source + if yaw_source == "segment": + target_yaw = route_yaw + elif yaw_source == "blend": + target_yaw = route_yaw + if index + 2 <= end_index: + following_goal = self.waypoints[index + 2] + next_dx = following_goal.x - next_goal.x + next_dy = following_goal.y - next_goal.y + if math.hypot(next_dx, next_dy) > 1.0e-6: + next_yaw = math.atan2(next_dy, next_dx) + delta = self.normalize_angle(next_yaw - route_yaw) + target_yaw = route_yaw + 0.15 * delta + elif yaw_source == "next": + target_yaw = math.radians(next_goal.yaw_deg) + else: + target_yaw = math.radians(goal.yaw_deg) + target_yaw = self.normalize_angle(target_yaw) + if self.cfg.slalom_script_rotate_steps_enabled: + steps.append( + SlalomScriptStep( + kind="rotate", + start_index=index, + end_index=index, + target_x=goal.x, + target_y=goal.y, + target_yaw=target_yaw, + pos_tolerance=getattr(goal, "slalom_script_pos_tolerance", None), + ) + ) + if route_dist >= self.cfg.slalom_script_min_step_distance: + forward = math.cos(target_yaw) * route_dx + math.sin(target_yaw) * route_dy + left = -math.sin(target_yaw) * route_dx + math.cos(target_yaw) * route_dy + steps.append( + SlalomScriptStep( + kind="drive", + start_index=index, + end_index=index + 1, + target_x=next_goal.x, + target_y=next_goal.y, + target_yaw=target_yaw, + pos_tolerance=getattr(next_goal, "slalom_script_pos_tolerance", None), + forward=forward, + left=left, + ) + ) + + final_goal = self.waypoints[end_index] + if self.cfg.slalom_script_final_rotate_enabled: + steps.append( + SlalomScriptStep( + kind="rotate", + start_index=end_index, + end_index=end_index, + target_x=final_goal.x, + target_y=final_goal.y, + target_yaw=math.radians(final_goal.yaw_deg), + pos_tolerance=getattr(final_goal, "slalom_script_pos_tolerance", None), + ) + ) + return steps + + def maybe_start_slalom_script(self, x: float, y: float, sim_time: float) -> bool: + if not self.cfg.slalom_script_enabled or self.slalom_script_active: + return self.slalom_script_active + active = self.active + if not self._is_slalom_script_waypoint(active): + return False + if math.hypot(active.x - x, active.y - y) > self.cfg.slalom_script_start_tolerance: + return False + + start_index, end_index = self.get_slalom_script_bounds(self.index) + if end_index <= start_index: + return False + steps = self.build_slalom_script_steps(start_index, end_index) + if not steps: + return False + + self.slalom_script_active = True + self.slalom_script_start_index = start_index + self.slalom_script_end_index = end_index + self.slalom_script_steps = steps + self.slalom_script_step_index = 0 + self.slalom_script_step_stable_count = 0 + self.turn_in_place = False + self.goal_entered_tolerance = False + self.goal_complete_stable_count = 0 + self.wp_start_time = sim_time + return True + + def compute_slalom_script_command( + self, + step: SlalomScriptStep, + x: float, + y: float, + yaw: float, + ) -> tuple[np.ndarray, bool]: + cmd = np.zeros(3, dtype=np.float32) + target_yaw = step.target_yaw + yaw_tolerance = math.radians(self.cfg.slalom_script_yaw_tolerance_deg) + + if step.kind == "rotate": + if target_yaw is None: + return cmd, True + yaw_err = self.normalize_angle(target_yaw - yaw) + if abs(yaw_err) <= yaw_tolerance: + return cmd, True + cmd[2] = self.clamp( + self.cfg.slalom_script_kp_yaw * yaw_err, + -self.cfg.slalom_script_max_wz, + self.cfg.slalom_script_max_wz, + ) + cmd[2] = self.apply_min_command( + float(cmd[2]), + self.cfg.slalom_script_min_cmd_angular, + self.cfg.slalom_script_max_wz, + self.cfg.slalom_script_min_cmd_epsilon, + ) + return cmd, False + + dx = step.target_x - x + dy = step.target_y - y + if target_yaw is None: + target_yaw = math.atan2(dy, dx) if math.hypot(dx, dy) > 1.0e-6 else yaw + yaw_err = self.normalize_angle(target_yaw - yaw) + pos_err = math.hypot(dx, dy) + pos_tolerance = ( + self.cfg.slalom_script_pos_tolerance + if step.pos_tolerance is None + else max(0.01, float(step.pos_tolerance)) + ) + complete = ( + pos_err <= pos_tolerance + and ( + not self.cfg.slalom_script_require_yaw_at_step + or abs(yaw_err) <= yaw_tolerance + ) + ) + if complete: + return cmd, True + + cmd[2] = self.clamp( + self.cfg.slalom_script_kp_yaw * yaw_err, + -self.cfg.slalom_script_max_wz, + self.cfg.slalom_script_max_wz, + ) + cmd[2] = self.apply_min_command( + float(cmd[2]), + self.cfg.slalom_script_min_cmd_angular, + self.cfg.slalom_script_max_wz, + self.cfg.slalom_script_min_cmd_epsilon, + ) + if abs(yaw_err) <= math.radians(self.cfg.slalom_script_yaw_gate_deg): + err_forward = math.cos(yaw) * dx + math.sin(yaw) * dy + err_left = -math.sin(yaw) * dx + math.cos(yaw) * dy + yaw_scale = max(0.35, math.cos(yaw_err) ** 2) + curvature_scale = 1.0 + if self.cfg.slalom_script_curvature_speed_enabled: + gate = max(1.0e-6, math.radians(self.cfg.slalom_script_yaw_gate_deg)) + curvature_scale = max( + self.cfg.slalom_script_curvature_min_scale, + 1.0 - abs(yaw_err) / gate, + ) + cmd[0] = self.clamp( + self.cfg.slalom_script_kp_dist * err_forward, + -self.cfg.slalom_script_max_vx * yaw_scale * curvature_scale, + self.cfg.slalom_script_max_vx * yaw_scale * curvature_scale, + ) + cmd[1] = self.clamp( + self.cfg.slalom_script_kp_dist * err_left, + -self.cfg.slalom_script_max_vy * curvature_scale, + self.cfg.slalom_script_max_vy * curvature_scale, + ) + cmd[0] = self.apply_min_command( + float(cmd[0]), + self.cfg.slalom_script_min_cmd_linear, + self.cfg.slalom_script_max_vx, + self.cfg.slalom_script_min_cmd_epsilon, + ) + cmd[1] = self.apply_min_command( + float(cmd[1]), + self.cfg.slalom_script_min_cmd_linear, + self.cfg.slalom_script_max_vy, + self.cfg.slalom_script_min_cmd_epsilon, + ) + return cmd, False + + @classmethod + def apply_min_command( + cls, + value: float, + min_abs: float, + max_abs: float, + epsilon: float, + ) -> float: + if max_abs <= 0.0 or min_abs <= 0.0: + return value + abs_value = abs(value) + if abs_value < max(0.0, epsilon) or abs_value >= min_abs: + return value + return cls.clamp(math.copysign(min_abs, value), -max_abs, max_abs) + + def finish_slalom_script(self, sim_time: float) -> tuple[np.ndarray, str | None]: + end_index = self.slalom_script_end_index + self.reset_slalom_script() + self.index = end_index + 1 + self.wp_start_time = sim_time + self.turn_in_place = False + self.goal_entered_tolerance = False + self.goal_complete_stable_count = 0 + self.best_dist = float("inf") + self.last_progress_time = sim_time + if self.index >= len(self.waypoints): + return np.zeros(3, dtype=np.float32), "complete" + return np.zeros(3, dtype=np.float32), None + + def handle_slalom_script( + self, + x: float, + y: float, + yaw: float, + sim_time: float, + ) -> tuple[np.ndarray, str | None] | None: + if not self.slalom_script_active and not self.maybe_start_slalom_script(x, y, sim_time): + return None + if sim_time - self.wp_start_time > self.cfg.waypoint_timeout_s: + active = self.active + active_id = active.id if active is not None else "unknown" + return np.zeros(3, dtype=np.float32), ( + f"timeout at slalom step {self.slalom_script_step_index + 1}/" + f"{len(self.slalom_script_steps)} near waypoint {active_id}" + ) + + while self.slalom_script_step_index < len(self.slalom_script_steps): + step = self.slalom_script_steps[self.slalom_script_step_index] + cmd, complete = self.compute_slalom_script_command(step, x, y, yaw) + if complete: + self.slalom_script_step_stable_count += 1 + if self.slalom_script_step_stable_count >= self.cfg.slalom_script_stable_cycles: + self.slalom_script_step_index += 1 + self.index = max(self.index, min(step.end_index, self.slalom_script_end_index)) + self.wp_start_time = sim_time + self.slalom_script_step_stable_count = 0 + continue + else: + self.slalom_script_step_stable_count = 0 + return cmd, None + + return self.finish_slalom_script(sim_time) + + def update(self, x: float, y: float, yaw: float, sim_time: float) -> tuple[np.ndarray, str | None]: + script_result = self.handle_slalom_script(x, y, yaw, sim_time) + if script_result is not None: + return script_result + return super().update(x, y, yaw, sim_time) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Run MuJoCo+ONNX route validation.") + parser.add_argument("--terrain-xml", type=Path, default=PROJECT_ROOT / "tools/nav_tools/xml/A.xml") + parser.add_argument("--points", type=Path, default=PROJECT_ROOT / "tools/nav_tools/points/points_20260705_174627.json") + parser.add_argument("--onnx", type=Path, default=PROJECT_ROOT / "model_6800.onnx") + parser.add_argument("--crawl-onnx", type=Path, default=PROJECT_ROOT / "model_crawl.onnx") + parser.add_argument( + "--policy-backend", + choices=("policy-runner", "onnx", "rough-ik-crawl"), + default="policy-runner", + help="policy-runner matches nav_sim2sim.py; rough-ik-crawl matches sim2real nav_good crawl_backend=ik.", + ) + parser.add_argument("--enable-policy-switch", action="store_true") + parser.add_argument( + "--mission", + choices=("json", "slalom", "grand"), + default="json", + help="json uses --points; slalom/grand use the hardcoded nav_sim2sim coordinate frame.", + ) + parser.add_argument("--robot-xml", type=Path, default=PROJECT_ROOT / "mjcf/wheelleg.xml") + parser.add_argument("--hfield-dir", type=Path, default=PROJECT_ROOT / "sim2sim/terrain") + parser.add_argument("--out-dir", type=Path, default=PROJECT_ROOT / "sim2sim/route_check_runs") + parser.add_argument("--avoid-margin", type=float, default=0.05) + parser.add_argument("--footprint-radius", type=float, default=None) + parser.add_argument("--max-time", type=float, default=420.0) + parser.add_argument("--waypoint-timeout", type=float, default=30.0) + parser.add_argument("--max-vx", type=float, default=1.2) + parser.add_argument("--max-vy", type=float, default=0.0) + parser.add_argument("--max-wz", type=float, default=0.8) + parser.add_argument("--kp-dist", type=float, default=0.8) + parser.add_argument("--kp-yaw", type=float, default=1.8) + parser.add_argument("--yaw-stop-threshold-deg", type=float, default=45.0) + parser.add_argument("--turn-in-place-enter-deg", type=float, default=70.0) + parser.add_argument("--turn-in-place-exit-deg", type=float, default=18.0) + parser.add_argument("--follower", choices=("waypoint", "pure-pursuit", "nav-good", "nav-sim2sim", "slalom-script", "nav-script"), default="pure-pursuit") + parser.add_argument("--lookahead", type=float, default=0.45) + parser.add_argument("--min-cmd-vx", type=float, default=0.08) + parser.add_argument("--creep-cmd-vx", type=float, default=0.04) + parser.add_argument("--cmd-vx-scale", type=float, default=1.0) + parser.add_argument("--stuck-timeout", type=float, default=4.0) + parser.add_argument("--recovery-duration", type=float, default=2.0) + parser.add_argument("--max-recoveries", type=int, default=4) + parser.add_argument("--slalom-script", action="store_true", default=True) + parser.add_argument("--no-slalom-script", dest="slalom_script", action="store_false") + parser.add_argument("--slalom-script-start-tolerance", type=float, default=0.10) + parser.add_argument("--slalom-script-pos-tolerance", type=float, default=0.08) + parser.add_argument("--slalom-script-yaw-tolerance-deg", type=float, default=5.0) + parser.add_argument("--slalom-script-stable-cycles", type=int, default=1) + parser.add_argument("--slalom-script-rotate-steps", action="store_true", default=False) + parser.add_argument("--slalom-script-final-rotate", action="store_true", default=False) + parser.add_argument("--slalom-script-require-yaw-at-step", action="store_true", default=False) + parser.add_argument("--slalom-script-kp-dist", type=float, default=1.0) + parser.add_argument("--slalom-script-kp-yaw", type=float, default=1.6) + parser.add_argument("--slalom-script-max-vx", type=float, default=0.60) + parser.add_argument("--slalom-script-max-vy", type=float, default=0.50) + parser.add_argument("--slalom-script-max-wz", type=float, default=0.50) + parser.add_argument("--slalom-script-min-cmd-linear", type=float, default=0.2) + parser.add_argument("--slalom-script-min-cmd-angular", type=float, default=0.0) + parser.add_argument("--slalom-script-min-cmd-epsilon", type=float, default=0.05) + parser.add_argument("--slalom-script-min-step-distance", type=float, default=0.02) + parser.add_argument("--slalom-script-yaw-gate-deg", type=float, default=25.0) + parser.add_argument("--slalom-script-curvature-speed", action="store_true", default=True) + parser.add_argument( + "--no-slalom-script-curvature-speed", + dest="slalom_script_curvature_speed", + action="store_false", + ) + parser.add_argument("--slalom-script-curvature-min-scale", type=float, default=0.55) + parser.add_argument( + "--slalom-script-drive-yaw-source", + choices=("current", "next", "segment", "blend"), + default="segment", + help="Yaw used while driving each slalom segment.", + ) + parser.add_argument("--local-safety", action="store_true", default=False) + parser.add_argument("--no-local-safety", dest="local_safety", action="store_false") + parser.add_argument("--safety-horizon", type=float, default=1.0) + parser.add_argument("--safety-dt", type=float, default=0.1) + parser.add_argument( + "--start-yaw-offset-deg", + type=float, + default=0.0, + help="Added to route yawDeg when placing the robot.", + ) + parser.add_argument( + "--heading-offset-deg", + type=float, + default=0.0, + help="Added to MuJoCo base yaw before waypoint heading control.", + ) + parser.add_argument("--sample-every", type=int, default=5) + parser.add_argument("--sim-speed", type=float, default=1.0, help="Simulation wall-clock speed factor, matching nav_sim2sim.py timing semantics.") + parser.add_argument("--viewer", action="store_true", help="Open MuJoCo passive viewer while the check runs.") + parser.add_argument("--viewer-sync-every", type=int, default=2) + parser.add_argument("--dashboard", action="store_true", help="Open a pygame telemetry dashboard with route, avoid regions, and trail.") + parser.add_argument("--dashboard-fps", type=float, default=30.0) + parser.add_argument("--dashboard-sim-speed", type=float, default=1.0, help="Initial dashboard simulation speed factor; can be changed in the UI up to 10x.") + parser.add_argument("--dashboard-width", type=int, default=1180) + parser.add_argument("--dashboard-height", type=int, default=760) + parser.add_argument("--dashboard-trail", type=int, default=6000) + parser.add_argument("--start-index", type=int, default=1, help="1-based waypoint index to start testing from.") + parser.add_argument("--end-index", type=int, default=None, help="1-based waypoint index to stop testing at.") + parser.add_argument("--start-id", type=str, default=None, help="Waypoint id to start testing from; overrides --start-index.") + parser.add_argument("--end-id", type=str, default=None, help="Waypoint id to stop testing at; overrides --end-index.") + parser.add_argument( + "--skip-initial-waypoint", + action="store_true", + help="Use the first sliced waypoint as spawn pose, then navigate to the next waypoint.", + ) + parser.add_argument( + "--auto-skip-slice-start", + action="store_true", + default=True, + help="When --start-index > 1, skip the sliced start waypoint by default.", + ) + parser.add_argument("--no-auto-skip-slice-start", dest="auto_skip_slice_start", action="store_false") + parser.add_argument("--start-z", type=float, default=0.75, help="Initial base z before settling.") + parser.add_argument("--start-x-offset", type=float, default=0.0, help="World x offset added to the initial spawn pose.") + parser.add_argument("--start-y-offset", type=float, default=0.0, help="World y offset added to the initial spawn pose.") + parser.add_argument("--settle-steps", type=int, default=500) + parser.add_argument("--min-base-z", type=float, default=0.12) + parser.add_argument("--stop-on-clearance-violation", action="store_true") + parser.add_argument("--deployment-randomization", action="store_true") + parser.add_argument("--random-seed", type=int, default=1) + parser.add_argument("--pose-noise-xy", type=float, default=0.015) + parser.add_argument("--pose-noise-yaw-deg", type=float, default=1.0) + parser.add_argument("--command-delay-ms", type=float, default=40.0) + parser.add_argument("--action-delay-ms", type=float, default=20.0) + parser.add_argument("--control-jitter-ms", type=float, default=4.0) + parser.add_argument("--command-dropout-prob", type=float, default=0.01) + parser.add_argument("--push-force", type=float, default=25.0) + parser.add_argument("--push-interval", type=float, default=8.0) + parser.add_argument("--push-duration", type=float, default=0.08) + parser.add_argument("--no-csv", action="store_true") + return parser.parse_args() + + +def make_combined_terrain_xml(terrain_xml: Path, work_dir: Path) -> Path: + text = terrain_xml.read_text(encoding="utf-8") + if '' not in text and '' not in text: + end = text.find(">") + if end < 0: + raise ValueError(f"Invalid terrain XML: {terrain_xml}") + text = text[: end + 1] + '\n' + text[end + 1 :] + out_path = work_dir / f"{terrain_xml.stem}_with_robot_include.xml" + out_path.write_text(text, encoding="utf-8") + return out_path + + +def builtin_mission_waypoints(name: str) -> list[Waypoint]: + slalom = [ + (2.9, -9.5, 0.4, "rough"), + (1.8, -9.95, 0.4, "rough"), + (0.9, -10.0, 0.4, "rough"), + (0.8, -10.8, 0.4, "rough"), + (0.9, -11.0, 0.4, "rough"), + (2.7, -11.0, 0.4, "rough"), + (2.8, -11.8, 0.4, "rough"), + (2.7, -12.0, 0.4, "rough"), + (0.9, -12.0, 0.4, "rough"), + (0.9, -12.75, 0.4, "rough"), + (1.42, -12.82, 0.4, "rough"), + (2.3, -13.3, 0.4, "rough"), + (2.3, -11.5, 0.4, "rough"), + (3.32, -12.44, 0.4, "rough"), + (3.7, -9.0, 0.6, "rough"), + ] + grand_tail = [ + (5.7, -10.2, 0.5, "rough"), + (5.7, -9.0, 0.3, "crawl"), + (5.7, -7.8, 0.3, "rough"), + (5.7, -7.0, 0.5, "rough"), + (3.7, -9.0, 0.5, "rough"), + (3.7, -12.5, 0.5, "rough"), + (4.84, -12.5, 0.4, "rough"), + (5.84, -12.0, 0.4, "rough"), + (5.84, -10.2, 0.4, "rough"), + (3.3, -10.2, 0.5, "rough"), + (3.3, -9.0, 0.5, "rough"), + (3.7, -9.0, 0.5, "rough"), + (3.3, -9.0, 0.5, "rough"), + (1.8, -9.0, 0.5, "rough"), + (1.8, -8.0, 0.4, "rough"), + (1.8, -7.5, 0.4, "rough"), + (1.8, -5.8, 1.2, "rough"), + (1.8, -3.5, 0.4, "rough"), + (1.8, -1.6, 0.4, "rough"), + (1.8, -1.38, 0.4, "rough"), + (1.8, 0.0, 0.4, "rough"), + (3.2, 0.0, 0.4, "rough"), + (4.85, 0.0, 0.4, "rough"), + (5.7, -0.5, 0.4, "rough"), + (5.7, -2.25, 0.4, "rough"), + (5.7, -3.5, 0.4, "rough"), + (5.7, -5.0, 0.4, "rough"), + (5.7, -7.0, 0.4, "rough"), + (3.3, -7.0, 0.5, "rough"), + (3.3, -9.0, 0.5, "rough"), + (3.7, -9.0, 0.6, "rough"), + ] + rows = slalom if name == "slalom" else slalom + grand_tail + return [ + Waypoint( + index=i, + id=f"{name}_{i}", + x=x, + y=y, + yaw_deg=0.0, + speed=speed, + policy=policy, + tolerance=0.15, + ) + for i, (x, y, speed, policy) in enumerate(rows, start=1) + ] + + +def load_route(path: Path) -> tuple[list[Waypoint], list[AvoidRegion]]: + with path.open("r", encoding="utf-8") as f: + payload = json.load(f) + if not isinstance(payload, dict): + raise ValueError(f"Route JSON root must be object: {path}") + return load_waypoints(payload), load_regions(payload) + + +def load_requested_route(args: argparse.Namespace) -> tuple[list[Waypoint], list[AvoidRegion]]: + if args.mission == "json": + waypoints, regions = load_route(args.points) + else: + waypoints, regions = builtin_mission_waypoints(args.mission), [] + + id_to_pos = {wp.id: i + 1 for i, wp in enumerate(waypoints)} + start = id_to_pos.get(str(args.start_id), max(1, int(args.start_index))) if args.start_id is not None else max(1, int(args.start_index)) + end_default = int(args.end_index) if args.end_index is not None else len(waypoints) + end = id_to_pos.get(str(args.end_id), end_default) if args.end_id is not None else end_default + if start > end or start > len(waypoints): + raise ValueError(f"Invalid waypoint slice: start={start} end={end} total={len(waypoints)}") + sliced = waypoints[start - 1 : min(end, len(waypoints))] + if len(sliced) < 1: + raise ValueError("Waypoint slice is empty") + return sliced, regions + + +def initial_pose_for_route(args: argparse.Namespace, waypoints: list[Waypoint]) -> tuple[float, float, float]: + if args.mission == "json": + start = waypoints[0] + return start.x, start.y, start.yaw_deg + float(args.start_yaw_offset_deg) + return 3.7, -9.0, 0.0 + + +def set_robot_pose( + io: MuJoCoIO, + x: float, + y: float, + yaw_deg: float, + z: float, + settle_steps: int, + default_dof_pos: np.ndarray, +) -> None: + yaw = math.radians(yaw_deg) + io.d.qpos[0] = x + io.d.qpos[1] = y + io.d.qpos[2] = z + io.d.qpos[3:7] = [math.cos(yaw * 0.5), 0.0, 0.0, math.sin(yaw * 0.5)] + io.d.qpos[io.qpos_ids] = default_dof_pos + io.d.qvel[:] = 0.0 + io.d.ctrl[io.ctrl_ids[:12]] = default_dof_pos[:12] + io.d.ctrl[io.ctrl_ids[12:]] = 0.0 + mujoco.mj_forward(io.m, io.d) + for _ in range(max(0, int(settle_steps))): + mujoco.mj_step(io.m, io.d) + + +def clearance_to_regions( + x: float, + y: float, + regions: list[AvoidRegion], +) -> tuple[float, str]: + best = float("inf") + best_name = "" + for region in regions: + polygon = region.polygon + if len(polygon) < 3: + continue + if point_in_polygon(x, y, polygon): + return 0.0, region.name + dist = min( + point_segment_distance(x, y, ax, ay, bx, by) + for (ax, ay), (bx, by) in zip(polygon, polygon[1:] + polygon[:1]) + ) + if dist < best: + best = dist + best_name = region.name + return best, best_name + + +def safety_filter_command( + desired: np.ndarray, + x: float, + y: float, + yaw: float, + target: Waypoint | None, + regions: list[AvoidRegion], + required_clearance: float, + cfg: SimConfig, + horizon_s: float, + dt_s: float, +) -> np.ndarray: + if not regions or target is None: + return desired + + base_vx = float(desired[0]) + base_vy = float(desired[1]) + base_wz = float(desired[2]) + vx_abs = abs(base_vx) + vy_abs = abs(base_vy) + vx_sign = -1.0 if base_vx < 0.0 else 1.0 + vy_sign = -1.0 if base_vy < 0.0 else 1.0 + vx_samples = [0.0, 0.35 * vx_abs, 0.65 * vx_abs, vx_abs] + vy_samples = [0.0, 0.5 * vy_abs, vy_abs] + wz_span = max(0.35, min(cfg.max_wz, abs(base_wz) + 0.35)) + wz_samples = [ + base_wz - wz_span, + base_wz - 0.5 * wz_span, + base_wz, + base_wz + 0.5 * wz_span, + base_wz + wz_span, + ] + best_cmd = desired.copy() + best_score = -float("inf") + initial_dist = math.hypot(target.x - x, target.y - y) + steps = max(1, int(round(horizon_s / dt_s))) + + for vx_mag in vx_samples: + for vy_mag in vy_samples: + for wz in wz_samples: + px, py, pyaw = x, y, yaw + min_margin = float("inf") + for _ in range(steps): + vx = vx_sign * vx_mag + vy = vy_sign * vy_mag + px += (vx * math.cos(pyaw) - vy * math.sin(pyaw)) * dt_s + py += (vx * math.sin(pyaw) + vy * math.cos(pyaw)) * dt_s + pyaw = WaypointFollower.normalize_angle(pyaw + wz * dt_s) + clearance, _ = clearance_to_regions(px, py, regions) + min_margin = min(min_margin, clearance - required_clearance) + final_dist = math.hypot(target.x - px, target.y - py) + progress = initial_dist - final_dist + yaw_cost = abs(wz - base_wz) + speed_cost = abs(vx_mag - vx_abs) + abs(vy_mag - vy_abs) + risk_penalty = 20.0 * max(0.0, -min_margin) + margin_reward = min(0.4, min_margin) + score = 3.0 * progress + margin_reward - risk_penalty - 0.12 * yaw_cost - 0.08 * speed_cost + if score > best_score: + best_score = score + best_cmd = np.array( + [ + vx_sign * vx_mag, + vy_sign * vy_mag, + WaypointFollower.clamp(wz, -cfg.max_wz, cfg.max_wz), + ], + dtype=np.float32, + ) + return best_cmd + + +def robot_yaw_from_qpos(qpos: np.ndarray) -> float: + qw, qx, qy, qz = qpos[3:7] + siny_cosp = 2.0 * (qw * qz + qx * qy) + cosy_cosp = 1.0 - 2.0 * (qy * qy + qz * qz) + return math.atan2(siny_cosp, cosy_cosp) + + +def quat_to_euler_wxyz(quat: np.ndarray) -> tuple[float, float, float]: + qw, qx, qy, qz = [float(v) for v in quat] + sinr_cosp = 2.0 * (qw * qx + qy * qz) + cosr_cosp = 1.0 - 2.0 * (qx * qx + qy * qy) + roll = math.atan2(sinr_cosp, cosr_cosp) + sinp = 2.0 * (qw * qy - qz * qx) + pitch = math.asin(max(-1.0, min(1.0, sinp))) + yaw = robot_yaw_from_qpos(np.array([0.0, 0.0, 0.0, qw, qx, qy, qz], dtype=np.float64)) + return roll, pitch, yaw + + +@dataclass(frozen=True) +class DashboardGeom: + name: str + kind: str + pos: tuple[float, ...] + size: tuple[float, ...] + quat: tuple[float, float, float, float] + rgba: tuple[float, float, float, float] + collidable: bool + + +def parse_dashboard_geoms(xml_path: Path) -> list[DashboardGeom]: + geoms: list[DashboardGeom] = [] + if not xml_path.exists(): + return geoms + try: + root = ET.parse(xml_path).getroot() + except ET.ParseError: + return geoms + for index, geom in enumerate(root.iter("geom"), start=1): + name = geom.get("name", f"geom_{index}") + if name == "floor": + continue + kind = geom.get("type", "box") + pos_text = geom.get("pos") + size_text = geom.get("size") + if not pos_text or not size_text: + continue + try: + pos = tuple(float(v) for v in pos_text.split()) + size = tuple(float(v) for v in size_text.split()) + quat_values = tuple(float(v) for v in geom.get("quat", "1 0 0 0").split()) + rgba_values = tuple(float(v) for v in geom.get("rgba", "0.65 0.65 0.65 1").split()) + except ValueError: + continue + if len(pos) < 2 or len(size) < 1 or len(quat_values) != 4: + continue + rgba = rgba_values if len(rgba_values) == 4 else (0.65, 0.65, 0.65, 1.0) + collidable = geom.get("contype", "1") != "0" and geom.get("conaffinity", "1") != "0" + geoms.append( + DashboardGeom( + name=name, + kind=kind, + pos=pos, + size=size, + quat=quat_values, # type: ignore[arg-type] + rgba=rgba, # type: ignore[arg-type] + collidable=collidable, + ) + ) + return geoms + + +def quat_yaw_wxyz(quat: tuple[float, float, float, float]) -> float: + qw, qx, qy, qz = quat + siny_cosp = 2.0 * (qw * qz + qx * qy) + cosy_cosp = 1.0 - 2.0 * (qy * qy + qz * qz) + return math.atan2(siny_cosp, cosy_cosp) + + +def local_rect_polygon( + origin_x: float, + origin_y: float, + yaw: float, + center_x: float, + center_y: float, + length: float, + width: float, +) -> list[tuple[float, float]]: + c = math.cos(yaw) + s = math.sin(yaw) + hx = length * 0.5 + hy = width * 0.5 + points = [] + for lx, ly in [(-hx, -hy), (hx, -hy), (hx, hy), (-hx, hy)]: + px = center_x + lx + py = center_y + ly + points.append((origin_x + px * c - py * s, origin_y + px * s + py * c)) + return points + + +def robot_wheel_local_points() -> list[tuple[str, float, float]]: + wheels = [] + for name, (pitch_x, pitch_y, knee_y, wheel_y, wheel_geom_y, knee_x) in ROBOT_WHEEL_POSITIONS.items(): + x = ROBOT_BODY_CENTER_OFFSET_X + pitch_x + knee_x + y = pitch_y + knee_y + wheel_y + wheel_geom_y + wheels.append((name, x, y)) + return wheels + + +class DashboardButton: + def __init__(self, rect: tuple[int, int, int, int], label: str, action: str) -> None: + self.rect_tuple = rect + self.label = label + self.action = action + + def rect(self, pygame: Any) -> Any: + return pygame.Rect(*self.rect_tuple) + + +class PygameRouteDashboard: + def __init__( + self, + waypoints: list[Waypoint], + regions: list[AvoidRegion], + terrain_xml: Path, + width: int, + height: int, + fps: float, + max_trail: int, + control_hz: float, + initial_sim_speed: float, + ) -> None: + os.environ.setdefault("PYGAME_HIDE_SUPPORT_PROMPT", "1") + # Older pkg_resources versions used by some pygame installs still + # reference pkgutil.ImpImporter, which was removed in Python 3.12. + if not hasattr(pkgutil, "ImpImporter"): + import zipimport + pkgutil.ImpImporter = zipimport.zipimporter # type: ignore[attr-defined] + import pygame + + self.pygame = pygame + pygame.init() + pygame.font.init() + self.screen = pygame.display.set_mode((width, height), pygame.RESIZABLE) + pygame.display.set_caption("nav_route_sim2sim_check dashboard") + self.clock = pygame.time.Clock() + self.font = pygame.font.Font(None, 24) + self.small_font = pygame.font.Font(None, 19) + self.title_font = pygame.font.Font(None, 28) + self.waypoints = waypoints + self.regions = regions + self.geoms = parse_dashboard_geoms(terrain_xml) + self.width = width + self.height = height + self.panel_width = 400 + self.fps = max(1.0, float(fps)) + self.control_hz = max(1.0, float(control_hz)) + self.max_trail = max(50, int(max_trail)) + self.trail: list[tuple[float, float]] = [] + self.zoom = 60.0 + self.pan_x = 0.0 + self.pan_y = 0.0 + self.speed_scale = 1.0 + self.sim_speed = max(0.1, min(10.0, float(initial_sim_speed))) + self.rewind_seconds = 0.0 + self.paused = False + self.buttons: list[DashboardButton] = [] + self._dragging = False + self._drag_start: tuple[int, int] | None = None + self._last_draw_time = 0.0 + self._fit_view() + + @property + def map_width(self) -> int: + return max(240, self.width - self.panel_width) + + def _fit_view(self) -> None: + points: list[tuple[float, float]] = [(wp.x, wp.y) for wp in self.waypoints] + for region in self.regions: + points.extend(region.polygon) + for geom in self.geoms: + if len(geom.pos) >= 2: + points.append((geom.pos[0], geom.pos[1])) + if not points: + return + xs = [p[0] for p in points] + ys = [p[1] for p in points] + span_x = max(1.0, max(xs) - min(xs)) + span_y = max(1.0, max(ys) - min(ys)) + margin = 70 + self.zoom = max( + 12.0, + min( + 120.0, + min((self.map_width - margin * 2) / span_x, (self.height - margin * 2) / span_y), + ), + ) + center_x = (min(xs) + max(xs)) * 0.5 + center_y = (min(ys) + max(ys)) * 0.5 + self.pan_x = -center_x * self.zoom + self.pan_y = center_y * self.zoom + + def _screen(self, x: float, y: float) -> tuple[int, int]: + sx = int(self.map_width / 2.0 + x * self.zoom + self.pan_x) + sy = int(self.height / 2.0 - y * self.zoom + self.pan_y) + return sx, sy + + def _draw_text(self, text: str, x: int, y: int, color: tuple[int, int, int] = (226, 232, 240), small: bool = False) -> None: + font = self.small_font if small else self.font + self.screen.blit(font.render(text, True, color), (x, y)) + + def _draw_card(self, rect: tuple[int, int, int, int], title: str | None = None) -> None: + pygame = self.pygame + pygame.draw.rect(self.screen, (20, 28, 48), rect, border_radius=7) + pygame.draw.rect(self.screen, (47, 61, 90), rect, width=1, border_radius=7) + if title: + self._draw_text(title, rect[0] + 12, rect[1] + 9, (203, 213, 225), small=True) + + def _draw_bar(self, x: int, y: int, w: int, value: float, lo: float, hi: float, color: tuple[int, int, int]) -> None: + pygame = self.pygame + frac = 0.0 if hi <= lo else max(0.0, min(1.0, (value - lo) / (hi - lo))) + pygame.draw.rect(self.screen, (8, 13, 25), (x, y, w, 8), border_radius=4) + pygame.draw.rect(self.screen, color, (x, y, int(w * frac), 8), border_radius=4) + + def _rgba_to_color(self, rgba: tuple[float, float, float, float], alpha: int | None = None) -> tuple[int, int, int] | tuple[int, int, int, int]: + rgb = tuple(max(0, min(255, int(v * 255))) for v in rgba[:3]) + if alpha is None: + return rgb + return (*rgb, alpha) + + def _draw_grid(self) -> None: + pygame = self.pygame + bg = (9, 14, 26) + grid = (28, 36, 58) + axis = (58, 72, 106) + self.screen.fill(bg) + step = 1.0 + left_world = -(self.map_width / 2.0 + self.pan_x) / self.zoom + right_world = (self.map_width / 2.0 - self.pan_x) / self.zoom + bottom_world = -(self.height / 2.0 - self.pan_y) / self.zoom + top_world = (self.height / 2.0 + self.pan_y) / self.zoom + gx = math.floor(left_world / step) * step + while gx <= right_world: + sx, _ = self._screen(gx, 0.0) + pygame.draw.line(self.screen, axis if abs(gx) < 1e-6 else grid, (sx, 0), (sx, self.height), 2 if abs(gx) < 1e-6 else 1) + gx += step + gy = math.floor(bottom_world / step) * step + while gy <= top_world: + _, sy = self._screen(0.0, gy) + pygame.draw.line(self.screen, axis if abs(gy) < 1e-6 else grid, (0, sy), (self.map_width, sy), 2 if abs(gy) < 1e-6 else 1) + gy += step + + def _draw_map(self, telemetry: dict[str, Any]) -> None: + pygame = self.pygame + self._draw_grid() + self._draw_xml_geoms() + overlay = pygame.Surface((self.map_width, self.height), pygame.SRCALPHA) + for region in self.regions: + pts = [self._screen(px, py) for px, py in region.polygon] + if len(pts) >= 3: + pygame.draw.polygon(overlay, (244, 63, 94, 72), pts) + pygame.draw.polygon(self.screen, (251, 113, 133), pts, width=2) + self.screen.blit(overlay, (0, 0)) + + route = [self._screen(wp.x, wp.y) for wp in self.waypoints] + if len(route) >= 2: + pygame.draw.lines(self.screen, (14, 165, 233), False, route, width=2) + active_index = int(telemetry.get("wp_index", 0)) + for index, wp in enumerate(self.waypoints): + color = (148, 163, 184) + radius = 4 + if index < active_index: + color = (34, 197, 94) + elif index == active_index: + color = (250, 204, 21) + radius = 7 + pygame.draw.circle(self.screen, color, self._screen(wp.x, wp.y), radius) + + if len(self.trail) >= 2: + pygame.draw.lines(self.screen, (34, 197, 94), False, [self._screen(x, y) for x, y in self.trail], width=3) + + x = float(telemetry.get("x", 0.0)) + y = float(telemetry.get("y", 0.0)) + yaw = float(telemetry.get("yaw", 0.0)) + self._draw_robot_pose(x, y, yaw) + + def _draw_robot_pose(self, x: float, y: float, yaw: float) -> None: + pygame = self.pygame + body = local_rect_polygon( + x, + y, + yaw, + ROBOT_BODY_CENTER_OFFSET_X, + 0.0, + ROBOT_BODY_LENGTH, + ROBOT_BODY_WIDTH, + ) + body_screen = [self._screen(px, py) for px, py in body] + pygame.draw.polygon(self.screen, (30, 41, 59), body_screen) + pygame.draw.polygon(self.screen, (226, 232, 240), body_screen, width=2) + + c = math.cos(yaw) + s = math.sin(yaw) + for _, lx, ly in robot_wheel_local_points(): + wx = x + lx * c - ly * s + wy = y + lx * s + ly * c + wheel = local_rect_polygon( + wx, + wy, + yaw, + 0.0, + 0.0, + ROBOT_WHEEL_VIS_LENGTH, + ROBOT_WHEEL_VIS_WIDTH, + ) + pygame.draw.polygon(self.screen, (2, 6, 23), [self._screen(px, py) for px, py in wheel]) + + origin = self._screen(x, y) + nose = self._screen(x + 0.32 * math.cos(yaw), y + 0.32 * math.sin(yaw)) + radius = max(3, int(self.zoom * 0.025)) + pygame.draw.circle(self.screen, (45, 212, 191), origin, radius) + pygame.draw.line(self.screen, (45, 212, 191), (origin[0] - radius - 2, origin[1]), (origin[0] + radius + 2, origin[1]), 1) + pygame.draw.line(self.screen, (45, 212, 191), (origin[0], origin[1] - radius - 2), (origin[0], origin[1] + radius + 2), 1) + pygame.draw.line(self.screen, (240, 253, 250), origin, nose, 2) + + def _draw_xml_geoms(self) -> None: + pygame = self.pygame + geom_layer = pygame.Surface((self.map_width, self.height), pygame.SRCALPHA) + for geom in self.geoms: + if len(geom.pos) < 2 or not geom.size: + continue + color = self._rgba_to_color(geom.rgba, 130 if geom.collidable else 55) + outline = self._rgba_to_color(geom.rgba) + if geom.kind == "box" and len(geom.size) >= 2: + cx, cy = geom.pos[0], geom.pos[1] + sx, sy = geom.size[0], geom.size[1] + yaw = quat_yaw_wxyz(geom.quat) + points = [] + for lx, ly in [(-sx, -sy), (sx, -sy), (sx, sy), (-sx, sy)]: + wx = cx + lx * math.cos(yaw) - ly * math.sin(yaw) + wy = cy + lx * math.sin(yaw) + ly * math.cos(yaw) + points.append(self._screen(wx, wy)) + pygame.draw.polygon(geom_layer, color, points) + pygame.draw.polygon(self.screen, outline, points, width=1) + elif geom.kind == "cylinder": + cx, cy = self._screen(geom.pos[0], geom.pos[1]) + radius = max(2, int(abs(geom.size[0]) * self.zoom)) + pygame.draw.circle(geom_layer, color, (cx, cy), radius) + pygame.draw.circle(self.screen, outline, (cx, cy), radius, width=1) + elif geom.kind == "plane": + continue + self.screen.blit(geom_layer, (0, 0)) + + def _draw_panel(self, telemetry: dict[str, Any]) -> None: + pygame = self.pygame + x0 = self.map_width + pygame.draw.rect(self.screen, (12, 18, 32), (x0, 0, self.panel_width, self.height)) + pygame.draw.line(self.screen, (55, 65, 88), (x0, 0), (x0, self.height), 2) + x = x0 + 16 + y = 14 + self.screen.blit(self.title_font.render("Route Monitor", True, (248, 250, 252)), (x, y)) + self._draw_text("MuJoCo sim2sim telemetry", x, y + 25, (148, 163, 184), small=True) + + y += 58 + card_w = self.panel_width - 32 + self._draw_card((x, y, card_w, 112), "state") + mode_color = (250, 204, 21) if str(telemetry.get("mode", "")).lower().startswith("crawl") else (45, 212, 191) + self._draw_text(str(telemetry.get("mode", "rough")), x + 12, y + 31, mode_color) + self._draw_text(f"t {float(telemetry.get('time', 0.0)):.2f}s", x + 130, y + 31, (226, 232, 240)) + self._draw_text( + f"wp {telemetry.get('wp_id', '')} {int(telemetry.get('wp_index', 0))}/{int(telemetry.get('wp_count', 0))}", + x + 12, + y + 60, + (226, 232, 240), + ) + self._draw_text( + f"x {float(telemetry.get('x', 0.0)):.2f} y {float(telemetry.get('y', 0.0)):.2f} yaw {math.degrees(float(telemetry.get('yaw', 0.0))):.1f}", + x + 12, + y + 86, + (148, 163, 184), + small=True, + ) + + y += 124 + self._draw_card((x, y, card_w, 128), "command limits") + cmd_vx = float(telemetry.get("cmd_vx", 0.0)) + cmd_vy = float(telemetry.get("cmd_vy", 0.0)) + cmd_wz = float(telemetry.get("cmd_wz", 0.0)) + for row, (name, value, color) in enumerate( + [ + ("x", cmd_vx, (34, 197, 94)), + ("y", cmd_vy, (14, 165, 233)), + ("yaw", cmd_wz, (250, 204, 21)), + ] + ): + yy = y + 32 + row * 30 + self._draw_text(f"{name} {value:+.3f}", x + 12, yy - 7, (226, 232, 240), small=True) + self._draw_bar(x + 90, yy, 245, value, -1.0, 1.0, color) + cx = x + 90 + 245 // 2 + pygame.draw.line(self.screen, (71, 85, 105), (cx, yy - 3), (cx, yy + 11), 1) + + y += 140 + self._draw_card((x, y, card_w, 112), "safety") + margin = float(telemetry.get("margin", 0.0)) + margin_color = (251, 113, 133) if margin < 0.0 else (34, 197, 94) + self._draw_text(f"margin {margin:+.3f} m", x + 12, y + 32, margin_color) + self._draw_text(f"clearance {float(telemetry.get('clearance', 0.0)):.3f} m", x + 190, y + 32, (226, 232, 240)) + self._draw_text(f"region {telemetry.get('region', '')}", x + 12, y + 62, (148, 163, 184), small=True) + self._draw_text( + f"roll {float(telemetry.get('roll_deg', 0.0)):.1f} pitch {float(telemetry.get('pitch_deg', 0.0)):.1f} tilt {float(telemetry.get('max_tilt_deg', 0.0)):.1f}", + x + 12, + y + 84, + (148, 163, 184), + small=True, + ) + + y += 124 + self._draw_controls(x, y, card_w) + + def _draw_controls(self, x: int, y: int, card_w: int) -> None: + pygame = self.pygame + self._draw_card((x, y, card_w, min(245, self.height - y - 12)), "controls") + x += 12 + y += 32 + self._draw_text(f"nav {self.speed_scale:.2f}x", x, y, (226, 232, 240)) + self._draw_text(f"sim {self.sim_speed:.1f}x", x + 132, y, (226, 232, 240)) + self._draw_text(f"rtf {float(self._last_rtf):.1f}x" if hasattr(self, "_last_rtf") else "rtf --", x + 250, y, (148, 163, 184), small=True) + y += 30 + labels = [ + ("nav -", "speed_down"), + ("nav +", "speed_up"), + ("nav 1x", "speed_normal"), + ("sim -", "sim_down"), + ("sim +", "sim_up"), + ("sim 1x", "sim_normal"), + ("sim 5x", "sim_5"), + ("sim 10x", "sim_10"), + ("rew -1s", "rewind_1"), + ("rew -5s", "rewind_5"), + ("pause", "pause"), + ("fit", "fit"), + ] + self.buttons = [] + bx = x + by = y + for i, (label, action) in enumerate(labels): + if i and i % 3 == 0: + bx = x + by += 36 + button = DashboardButton((bx, by, 104, 28), label, action) + rect = button.rect(pygame) + fill = (31, 41, 55) + if action == "pause" and self.paused: + fill = (245, 158, 11) + elif action == "speed_normal" and abs(self.speed_scale - 1.0) < 1e-6: + fill = (16, 185, 129) + elif action == "sim_normal" and abs(self.sim_speed - 1.0) < 1e-6: + fill = (14, 165, 233) + pygame.draw.rect(self.screen, fill, rect, border_radius=5) + pygame.draw.rect(self.screen, (71, 85, 105), rect, width=1, border_radius=5) + text = self.small_font.render(label, True, (248, 250, 252)) + self.screen.blit(text, (rect.centerx - text.get_width() // 2, rect.centery - text.get_height() // 2)) + self.buttons.append(button) + bx += 112 + by += 48 + self._draw_text("keys: +/- nav, [/ ] sim, backspace rewind, space pause", x, by, (148, 163, 184), small=True) + self._draw_text("wheel zoom, right drag pan, f fit", x, by + 20, (148, 163, 184), small=True) + + def pump_events(self) -> bool: + pygame = self.pygame + for event in pygame.event.get(): + if event.type == pygame.QUIT: + return False + if event.type == pygame.KEYDOWN: + if event.key in (pygame.K_EQUALS, pygame.K_PLUS, pygame.K_KP_PLUS): + self.speed_scale = min(10.0, self.speed_scale + 0.10) + elif event.key in (pygame.K_MINUS, pygame.K_KP_MINUS): + self.speed_scale = max(0.05, self.speed_scale - 0.10) + elif event.key == pygame.K_0: + self.speed_scale = 1.0 + elif event.key == pygame.K_RIGHTBRACKET: + self.sim_speed = self._next_sim_speed(1) + elif event.key == pygame.K_LEFTBRACKET: + self.sim_speed = self._next_sim_speed(-1) + elif event.key == pygame.K_BACKSPACE: + self.rewind_seconds = max(self.rewind_seconds, 1.0) + elif event.key == pygame.K_SPACE: + self.paused = not self.paused + elif event.key == pygame.K_f: + self._fit_view() + if event.type == pygame.VIDEORESIZE: + self.width = max(700, int(event.w)) + self.height = max(480, int(event.h)) + self.screen = pygame.display.set_mode((self.width, self.height), pygame.RESIZABLE) + self._fit_view() + if event.type == pygame.MOUSEBUTTONDOWN: + if event.button == 1: + for button in self.buttons: + if button.rect(pygame).collidepoint(event.pos): + self._handle_button(button.action) + break + elif event.button == 3: + self._dragging = True + self._drag_start = event.pos + elif event.button == 4 and event.pos[0] < self.map_width: + self.zoom = min(240.0, self.zoom * 1.12) + elif event.button == 5 and event.pos[0] < self.map_width: + self.zoom = max(8.0, self.zoom / 1.12) + if event.type == pygame.MOUSEBUTTONUP and event.button == 3: + self._dragging = False + self._drag_start = None + if event.type == pygame.MOUSEMOTION and self._dragging and self._drag_start is not None: + dx = event.pos[0] - self._drag_start[0] + dy = event.pos[1] - self._drag_start[1] + self.pan_x += dx + self.pan_y += dy + self._drag_start = event.pos + return True + + def draw_if_due(self, telemetry: dict[str, Any], force: bool = False) -> None: + now = time.perf_counter() + draw_interval = 1.0 / max(1.0, self.fps) + if force or now - self._last_draw_time >= draw_interval: + self._last_rtf = float(telemetry.get("rtf", 0.0)) + self.trail.append((float(telemetry.get("x", 0.0)), float(telemetry.get("y", 0.0)))) + if len(self.trail) > self.max_trail: + self.trail = self.trail[-self.max_trail :] + self._draw_map(telemetry) + self._draw_panel(telemetry) + self.pygame.display.flip() + self._last_draw_time = now + + def _handle_button(self, action: str) -> None: + if action == "speed_down": + self.speed_scale = max(0.05, self.speed_scale - 0.10) + elif action == "speed_up": + self.speed_scale = min(10.0, self.speed_scale + 0.10) + elif action == "speed_normal": + self.speed_scale = 1.0 + elif action == "sim_down": + self.sim_speed = self._next_sim_speed(-1) + elif action == "sim_up": + self.sim_speed = self._next_sim_speed(1) + elif action == "sim_normal": + self.sim_speed = 1.0 + elif action == "sim_5": + self.sim_speed = 5.0 + elif action == "sim_10": + self.sim_speed = 10.0 + elif action == "rewind_1": + self.rewind_seconds = max(self.rewind_seconds, 1.0) + elif action == "rewind_5": + self.rewind_seconds = max(self.rewind_seconds, 5.0) + elif action == "pause": + self.paused = not self.paused + elif action == "fit": + self._fit_view() + + def _next_sim_speed(self, direction: int) -> float: + steps = [0.2, 0.5, 0.8, 1.0, 1.5, 2.0, 3.0, 5.0, 10.0] + current = float(self.sim_speed) + if direction > 0: + for value in steps: + if value > current + 1e-6: + return value + return steps[-1] + for value in reversed(steps): + if value < current - 1e-6: + return value + return steps[0] + + def consume_rewind_seconds(self) -> float: + value = self.rewind_seconds + self.rewind_seconds = 0.0 + return value + + def close(self) -> None: + self.pygame.quit() + + +def capture_route_snapshot(io: MuJoCoIO, follower: WaypointFollower) -> RouteSnapshot: + return RouteSnapshot( + time=float(io.d.time), + qpos=io.d.qpos.copy(), + qvel=io.d.qvel.copy(), + ctrl=io.d.ctrl.copy(), + follower_index=int(follower.index), + follower_turn_in_place=bool(follower.turn_in_place), + follower_stable_count=int(follower.stable_count), + follower_wp_start_time=float(follower.wp_start_time), + follower_best_dist=float(getattr(follower, "best_dist", 0.0)) if hasattr(follower, "best_dist") else None, + follower_last_progress_time=float(getattr(follower, "last_progress_time", 0.0)) if hasattr(follower, "last_progress_time") else None, + follower_recovery_until=float(getattr(follower, "recovery_until", 0.0)) if hasattr(follower, "recovery_until") else None, + follower_recovery_count=int(getattr(follower, "recovery_count", 0)) if hasattr(follower, "recovery_count") else None, + follower_recovery_turn_sign=float(getattr(follower, "recovery_turn_sign", 1.0)) if hasattr(follower, "recovery_turn_sign") else None, + ) + + +def restore_route_snapshot(io: MuJoCoIO, follower: WaypointFollower, snapshot: RouteSnapshot) -> None: + io.d.time = snapshot.time + io.d.qpos[:] = snapshot.qpos + io.d.qvel[:] = snapshot.qvel + io.d.ctrl[:] = snapshot.ctrl + follower.index = snapshot.follower_index + follower.turn_in_place = snapshot.follower_turn_in_place + follower.stable_count = snapshot.follower_stable_count + follower.wp_start_time = snapshot.follower_wp_start_time + if snapshot.follower_best_dist is not None and hasattr(follower, "best_dist"): + setattr(follower, "best_dist", snapshot.follower_best_dist) + if snapshot.follower_last_progress_time is not None and hasattr(follower, "last_progress_time"): + setattr(follower, "last_progress_time", snapshot.follower_last_progress_time) + if snapshot.follower_recovery_until is not None and hasattr(follower, "recovery_until"): + setattr(follower, "recovery_until", snapshot.follower_recovery_until) + if snapshot.follower_recovery_count is not None and hasattr(follower, "recovery_count"): + setattr(follower, "recovery_count", snapshot.follower_recovery_count) + if snapshot.follower_recovery_turn_sign is not None and hasattr(follower, "recovery_turn_sign"): + setattr(follower, "recovery_turn_sign", snapshot.follower_recovery_turn_sign) + mujoco.mj_forward(io.m, io.d) + + +def run_sim(args: argparse.Namespace) -> tuple[SimResult, Path | None, Path]: + args.out_dir.mkdir(parents=True, exist_ok=True) + stamp = time.strftime("%Y%m%d_%H%M%S") + f"_{time.time_ns() % 1_000_000_000:09d}" + run_dir = args.out_dir / f"route_check_{stamp}" + run_dir.mkdir(parents=True, exist_ok=True) + + waypoints, regions = load_requested_route(args) + footprint_radius = ( + float(args.footprint_radius) + if args.footprint_radius is not None + else default_lateral_footprint_radius() + ) + required_clearance = footprint_radius + float(args.avoid_margin) + cfg = SimConfig( + max_vx=float(args.max_vx), + max_vy=float(args.max_vy), + max_wz=float(args.max_wz), + kp_dist=float(args.kp_dist), + kp_yaw=float(args.kp_yaw), + yaw_stop_threshold_deg=float(args.yaw_stop_threshold_deg), + turn_in_place_enter_deg=float(args.turn_in_place_enter_deg), + turn_in_place_exit_deg=float(args.turn_in_place_exit_deg), + max_total_time_s=float(args.max_time), + waypoint_timeout_s=float(args.waypoint_timeout), + lookahead_m=float(args.lookahead), + min_cmd_vx=float(args.min_cmd_vx), + creep_cmd_vx=float(args.creep_cmd_vx), + cmd_vx_scale=float(args.cmd_vx_scale), + stuck_timeout_s=float(args.stuck_timeout), + recovery_duration_s=float(args.recovery_duration), + max_recoveries=int(args.max_recoveries), + slalom_script_enabled=bool(args.slalom_script), + slalom_script_start_tolerance=float(args.slalom_script_start_tolerance), + slalom_script_pos_tolerance=float(args.slalom_script_pos_tolerance), + slalom_script_yaw_tolerance_deg=float(args.slalom_script_yaw_tolerance_deg), + slalom_script_stable_cycles=int(args.slalom_script_stable_cycles), + slalom_script_rotate_steps_enabled=bool(args.slalom_script_rotate_steps), + slalom_script_final_rotate_enabled=bool(args.slalom_script_final_rotate), + slalom_script_require_yaw_at_step=bool(args.slalom_script_require_yaw_at_step), + slalom_script_kp_dist=float(args.slalom_script_kp_dist), + slalom_script_kp_yaw=float(args.slalom_script_kp_yaw), + slalom_script_max_vx=float(args.slalom_script_max_vx), + slalom_script_max_vy=float(args.slalom_script_max_vy), + slalom_script_max_wz=float(args.slalom_script_max_wz), + slalom_script_min_cmd_linear=float(args.slalom_script_min_cmd_linear), + slalom_script_min_cmd_angular=float(args.slalom_script_min_cmd_angular), + slalom_script_min_cmd_epsilon=float(args.slalom_script_min_cmd_epsilon), + slalom_script_min_step_distance=float(args.slalom_script_min_step_distance), + slalom_script_yaw_gate_deg=float(args.slalom_script_yaw_gate_deg), + slalom_script_drive_yaw_source=str(args.slalom_script_drive_yaw_source), + slalom_script_curvature_speed_enabled=bool(args.slalom_script_curvature_speed), + slalom_script_curvature_min_scale=max( + 0.1, min(1.0, float(args.slalom_script_curvature_min_scale)) + ), + ) + + combined_xml = make_combined_terrain_xml(args.terrain_xml.resolve(), run_dir) + old_cwd = Path.cwd() + try: + io = MuJoCoIO(combined_xml, args.robot_xml.resolve(), args.hfield_dir.resolve()) + finally: + os.chdir(old_cwd) + if args.policy_backend == "rough-ik-crawl": + policy = RoughPolicyIkCrawlBackend(args.onnx.resolve()) + elif args.policy_backend == "policy-runner": + crawl_path = args.crawl_onnx.resolve() if args.crawl_onnx and args.crawl_onnx.exists() else None + policy = ExistingPolicyBackend(args.onnx.resolve(), crawl_path) + else: + policy = DirectOnnxBackend(args.onnx.resolve(), DEPLOY_DEFAULT_DOF_POS) + follower: WaypointFollower + if args.follower in {"slalom-script", "nav-script"}: + follower = SlalomScriptFollower(waypoints, cfg) + elif args.follower == "nav-good": + follower = NavGoodFollower(waypoints, cfg) + elif args.follower == "nav-sim2sim": + follower = NavSim2SimFollower(waypoints, cfg) + elif args.follower == "pure-pursuit": + follower = PurePursuitFollower(waypoints, cfg) + else: + follower = WaypointFollower(waypoints, cfg) + if ( + args.skip_initial_waypoint + or (args.auto_skip_slice_start and int(args.start_index) > 1) + ) and len(waypoints) > 1: + follower.index = 1 + + start_x, start_y, start_yaw_deg = initial_pose_for_route(args, waypoints) + start_x += float(args.start_x_offset) + start_y += float(args.start_y_offset) + set_robot_pose( + io, + start_x, + start_y, + start_yaw_deg, + float(args.start_z), + int(args.settle_steps), + policy.default_dof_pos, + ) + policy.reset() + follower.reset_timing(io.d.time) + + sim_steps_per_control = max(1, int(round((1.0 / cfg.control_hz) / io.m.opt.timestep))) + csv_path = None if args.no_csv else run_dir / "trajectory.csv" + csv_file = None + writer = None + if csv_path is not None: + csv_file = csv_path.open("w", newline="", encoding="utf-8") + writer = csv.writer(csv_file) + writer.writerow( + [ + "time", + "wp_index", + "wp_id", + "script_active", + "script_step", + "script_step_count", + "x", + "y", + "z", + "yaw_deg", + "roll_deg", + "pitch_deg", + "cmd_vx", + "cmd_vy", + "cmd_wz", + "clearance", + "margin", + "region", + "raw_wheel_fl", + "raw_wheel_fr", + "raw_wheel_rl", + "raw_wheel_rr", + "qvel_wheel_fl", + "qvel_wheel_fr", + "qvel_wheel_rl", + "qvel_wheel_rr", + ] + ) + + min_clearance = float("inf") + min_clearance_region = "" + min_clearance_wp = "" + max_roll = 0.0 + max_pitch = 0.0 + max_tilt = 0.0 + samples = 0 + reason = "max time reached" + success = False + + viewer = None + if args.viewer: + viewer = mujoco.viewer.launch_passive(io.m, io.d) + viewer.cam.distance = 5.0 + viewer.cam.elevation = -35.0 + viewer.cam.azimuth = 135.0 + sim_speed_factor = max(0.1, min(10.0, float(args.sim_speed))) + if args.dashboard: + sim_speed_factor = max(0.1, min(10.0, float(args.dashboard_sim_speed))) + dashboard = None + if args.dashboard: + dashboard = PygameRouteDashboard( + waypoints=waypoints, + regions=regions, + terrain_xml=args.terrain_xml.resolve(), + width=int(args.dashboard_width), + height=int(args.dashboard_height), + fps=float(args.dashboard_fps), + max_trail=int(args.dashboard_trail), + control_hz=cfg.control_hz, + initial_sim_speed=sim_speed_factor, + ) + snapshots: list[RouteSnapshot] = [] + snapshot_keep_s = 45.0 + control_dt = 1.0 / cfg.control_hz + next_exec_time = time.perf_counter() + wall_start_time = next_exec_time + sim_start_time = float(io.d.time) + randomizer = random.Random(int(args.random_seed)) + command_delay_steps = max(0, int(round(float(args.command_delay_ms) * 0.001 * cfg.control_hz))) + action_delay_steps = max(0, int(round(float(args.action_delay_ms) * 0.001 * cfg.control_hz))) + command_queue: deque[np.ndarray] = deque() + action_queue: deque[np.ndarray] = deque() + last_deployed_command = np.zeros(3, dtype=np.float32) + base_body_id = mujoco.mj_name2id(io.m, mujoco.mjtObj.mjOBJ_BODY, "base_link") + + try: + while float(io.d.time) - sim_start_time < cfg.max_total_time_s: + if viewer is not None and not viewer.is_running(): + reason = "viewer closed" + break + if dashboard is not None and not dashboard.pump_events(): + reason = "dashboard closed" + break + if dashboard is not None: + sim_speed_factor = max(0.1, min(10.0, float(dashboard.sim_speed))) + x = float(io.d.qpos[0]) + y = float(io.d.qpos[1]) + z = float(io.d.qpos[2]) + yaw = robot_yaw_from_qpos(io.d.qpos) + sensed_x = x + sensed_y = y + sensed_yaw = yaw + if args.deployment_randomization: + sensed_x += randomizer.gauss(0.0, max(0.0, float(args.pose_noise_xy))) + sensed_y += randomizer.gauss(0.0, max(0.0, float(args.pose_noise_xy))) + sensed_yaw += math.radians( + randomizer.gauss(0.0, max(0.0, float(args.pose_noise_yaw_deg))) + ) + heading_yaw = follower.normalize_angle(sensed_yaw + math.radians(float(args.heading_offset_deg))) + roll, pitch, _ = quat_to_euler_wxyz(io.d.qpos[3:7].copy()) + roll_deg = math.degrees(roll) + pitch_deg = math.degrees(pitch) + max_roll = max(max_roll, abs(roll_deg)) + max_pitch = max(max_pitch, abs(pitch_deg)) + max_tilt = max(max_tilt, math.hypot(roll_deg, pitch_deg)) + + active = follower.active + active_id = active.id if active is not None else "done" + clearance, region_name = clearance_to_regions(x, y, regions) + margin = clearance - required_clearance + if clearance < min_clearance: + min_clearance = clearance + min_clearance_region = region_name + min_clearance_wp = active_id + + if z < float(args.min_base_z) or abs(roll_deg) > 70.0 or abs(pitch_deg) > 70.0: + reason = f"robot fell or tipped at t={io.d.time:.2f}s" + break + if args.stop_on_clearance_violation and margin < 0.0: + reason = ( + f"clearance violation at t={io.d.time:.2f}s: " + f"clearance={clearance:.3f} required={required_clearance:.3f} region={region_name}" + ) + break + + command, nav_state = follower.update(sensed_x, sensed_y, heading_yaw, io.d.time) + if nav_state == "complete": + success = True + reason = "route complete" + break + if nav_state and nav_state.startswith("timeout"): + reason = nav_state + break + if nav_state: + reason = nav_state + break + skip_local_safety = False + if active is not None and (active.policy or "").lower() == "crawl": + skip_local_safety = True + if hasattr(follower, "_is_charge_segment") and active is not None: + skip_local_safety = bool(getattr(follower, "_is_charge_segment")(x, y, active)) or skip_local_safety + if args.local_safety and not skip_local_safety: + command = safety_filter_command( + command, + x, + y, + heading_yaw, + follower.active, + regions, + required_clearance, + cfg, + float(args.safety_horizon), + float(args.safety_dt), + ) + if dashboard is not None: + command[0] = float(command[0]) * dashboard.speed_scale + if dashboard.paused: + command[:] = 0.0 + command = np.clip(command, -1.0, 1.0).astype(np.float32) + if args.deployment_randomization: + if randomizer.random() < max(0.0, min(1.0, float(args.command_dropout_prob))): + command = last_deployed_command.copy() + command_queue.append(command.copy()) + command = command_queue.popleft() if len(command_queue) > command_delay_steps else np.zeros(3, dtype=np.float32) + last_deployed_command = command.copy() + + use_route_policy = args.enable_policy_switch or args.policy_backend == "rough-ik-crawl" + active_policy = active.policy if active is not None and use_route_policy else "rough" + policy.maybe_switch_policy(active_policy) + policy_mode = str(getattr(policy, "requested_mode", getattr(getattr(policy, "runner", None), "current_policy_name", active_policy))) + default_dof_pos = policy.default_dof_pos + obs = io.get_obs_53d(command, default_dof_pos, policy.last_actions) + scaled_actions, raw_actions = policy.step(obs) + if args.deployment_randomization: + action_queue.append(scaled_actions.copy()) + scaled_actions = action_queue.popleft() if len(action_queue) > action_delay_steps else np.zeros_like(scaled_actions) + io.send_actions(scaled_actions, default_dof_pos) + for _ in range(sim_steps_per_control): + if args.deployment_randomization and base_body_id != -1: + io.d.xfrc_applied[base_body_id] = 0.0 + push_phase = (float(io.d.time) - sim_start_time) % max(0.1, float(args.push_interval)) + if push_phase < max(0.0, float(args.push_duration)): + angle = randomizer.uniform(-math.pi, math.pi) + force = max(0.0, float(args.push_force)) + io.d.xfrc_applied[base_body_id, 0] = force * math.cos(angle) + io.d.xfrc_applied[base_body_id, 1] = force * math.sin(angle) + mujoco.mj_step(io.m, io.d) + pole_contact = None + pole_contact_robot_geom = None + pole_contact_robot_body = None + pole_contact_position = None + border_contact = None + for contact_index in range(io.d.ncon): + contact = io.d.contact[contact_index] + geom1 = mujoco.mj_id2name(io.m, mujoco.mjtObj.mjOBJ_GEOM, contact.geom1) or "" + geom2 = mujoco.mj_id2name(io.m, mujoco.mjtObj.mjOBJ_GEOM, contact.geom2) or "" + if geom1.startswith("slalom_pole_") or geom2.startswith("slalom_pole_"): + pole_contact = geom1 if geom1.startswith("slalom_pole_") else geom2 + robot_geom_id = contact.geom2 if geom1.startswith("slalom_pole_") else contact.geom1 + pole_contact_robot_geom = geom2 if geom1.startswith("slalom_pole_") else geom1 + robot_body_id = int(io.m.geom_bodyid[robot_geom_id]) + pole_contact_robot_body = ( + mujoco.mj_id2name(io.m, mujoco.mjtObj.mjOBJ_BODY, robot_body_id) or "unknown" + ) + pole_contact_position = tuple(float(value) for value in contact.pos) + break + if geom1.startswith("border_") or geom2.startswith("border_"): + border_contact = geom1 if geom1.startswith("border_") else geom2 + break + if pole_contact is not None: + contact_xyz = "" + if pole_contact_position is not None: + contact_xyz = ( + f" at xyz=({pole_contact_position[0]:.3f}," + f"{pole_contact_position[1]:.3f},{pole_contact_position[2]:.3f})" + ) + reason = ( + f"physical contact: robot body {pole_contact_robot_body or 'unknown'} " + f"geom {pole_contact_robot_geom or 'unnamed'} with " + f"{pole_contact}{contact_xyz} at t={io.d.time:.2f}s near waypoint {active_id}" + ) + break + if border_contact is not None: + reason = f"physical contact with field boundary {border_contact} at t={io.d.time:.2f}s near waypoint {active_id}" + break + if pole_contact is not None or border_contact is not None: + break + if dashboard is not None: + snapshots.append(capture_route_snapshot(io, follower)) + cutoff_time = float(io.d.time) - snapshot_keep_s + while len(snapshots) > 2 and snapshots[0].time < cutoff_time: + snapshots.pop(0) + viewer_sync_stride = max(1, int(args.viewer_sync_every), int(math.ceil(max(1.0, sim_speed_factor) * 2.0))) + if viewer is not None and samples % viewer_sync_stride == 0: + base_id = mujoco.mj_name2id(io.m, mujoco.mjtObj.mjOBJ_BODY, "base_link") + if base_id != -1: + viewer.cam.lookat[:] = io.d.xpos[base_id] + viewer.sync() + + if dashboard is not None: + wall_elapsed = max(1e-6, time.perf_counter() - wall_start_time) + rtf = max(0.0, (float(io.d.time) - sim_start_time) / wall_elapsed) + telemetry = { + "time": float(io.d.time), + "rtf": rtf, + "status": "running", + "mode": policy_mode, + "wp_index": follower.index, + "wp_count": len(waypoints), + "wp_id": active_id, + "x": x, + "y": y, + "z": z, + "yaw": yaw, + "roll_deg": roll_deg, + "pitch_deg": pitch_deg, + "cmd_vx": float(command[0]), + "cmd_vy": float(command[1]), + "cmd_wz": float(command[2]), + "clearance": clearance, + "margin": margin, + "region": region_name, + "min_margin": min_clearance - required_clearance, + "max_tilt_deg": max_tilt, + } + dashboard.draw_if_due(telemetry) + rewind_seconds = dashboard.consume_rewind_seconds() + if rewind_seconds > 0.0 and snapshots: + target_time = max(0.0, float(io.d.time) - rewind_seconds) + rewind_snapshot = snapshots[0] + for snapshot in snapshots: + if snapshot.time <= target_time: + rewind_snapshot = snapshot + else: + break + restore_route_snapshot(io, follower, rewind_snapshot) + snapshots = [snapshot for snapshot in snapshots if snapshot.time <= rewind_snapshot.time] + if hasattr(policy, "reset"): + policy.reset() + reason = "max time reached" + next_exec_time = time.perf_counter() + + if writer is not None and samples % max(1, int(args.sample_every)) == 0: + writer.writerow( + [ + f"{io.d.time:.4f}", + follower.index, + active_id, + int(bool(getattr(follower, "slalom_script_active", False))), + int(getattr(follower, "slalom_script_step_index", -1)), + len(getattr(follower, "slalom_script_steps", [])), + f"{x:.5f}", + f"{y:.5f}", + f"{z:.5f}", + f"{math.degrees(yaw):.3f}", + f"{roll_deg:.3f}", + f"{pitch_deg:.3f}", + f"{float(command[0]):.4f}", + f"{float(command[1]):.4f}", + f"{float(command[2]):.4f}", + f"{clearance:.5f}", + f"{margin:.5f}", + region_name, + f"{raw_actions[12]:.5f}", + f"{raw_actions[13]:.5f}", + f"{raw_actions[14]:.5f}", + f"{raw_actions[15]:.5f}", + f"{io.d.qvel[io.qvel_ids[12]]:.5f}", + f"{io.d.qvel[io.qvel_ids[13]]:.5f}", + f"{io.d.qvel[io.qvel_ids[14]]:.5f}", + f"{io.d.qvel[io.qvel_ids[15]]:.5f}", + ] + ) + samples += 1 + jitter_s = 0.0 + if args.deployment_randomization: + jitter_s = randomizer.uniform( + -max(0.0, float(args.control_jitter_ms)) * 0.001, + max(0.0, float(args.control_jitter_ms)) * 0.001, + ) + dt_real = max(0.001, control_dt + jitter_s) / sim_speed_factor + next_exec_time += dt_real + now = time.perf_counter() + sleep_time = next_exec_time - now + if sleep_time > 0.0: + time.sleep(sleep_time) + elif sleep_time < -dt_real: + next_exec_time = now + finally: + wall_time = max(1e-6, time.perf_counter() - wall_start_time) + elapsed_sim_time = max(0.0, float(io.d.time) - sim_start_time) + real_time_factor = max(0.0, elapsed_sim_time / wall_time) + if csv_file is not None: + csv_file.close() + if viewer is not None: + viewer.close() + if dashboard is not None: + dashboard.close() + policy.close() + + result = SimResult( + success=success, + reason=reason, + sim_time=elapsed_sim_time, + reached_count=follower.index, + waypoint_count=len(waypoints), + min_clearance=float(min_clearance), + min_margin=float(min_clearance - required_clearance), + min_clearance_region=min_clearance_region, + min_clearance_wp=min_clearance_wp, + max_roll_deg=max_roll, + max_pitch_deg=max_pitch, + max_tilt_deg=max_tilt, + samples=samples, + wall_time=wall_time, + real_time_factor=real_time_factor, + ) + report = { + "success": result.success, + "reason": result.reason, + "sim_time": round(result.sim_time, 4), + "reached_count": result.reached_count, + "waypoint_count": result.waypoint_count, + "min_clearance": round(result.min_clearance, 6), + "required_clearance": round(required_clearance, 6), + "min_margin": round(result.min_margin, 6), + "min_clearance_region": result.min_clearance_region, + "min_clearance_wp": result.min_clearance_wp, + "max_roll_deg": round(result.max_roll_deg, 3), + "max_pitch_deg": round(result.max_pitch_deg, 3), + "max_tilt_deg": round(result.max_tilt_deg, 3), + "samples": result.samples, + "wall_time": round(result.wall_time, 4), + "real_time_factor": round(result.real_time_factor, 4), + "sim_speed": round(sim_speed_factor, 4), + "points": str(args.points.resolve()), + "mission": args.mission, + "terrain_xml": str(args.terrain_xml.resolve()), + "onnx": str(args.onnx.resolve()), + "crawl_onnx": str(args.crawl_onnx.resolve()) if args.crawl_onnx else None, + "policy_backend": args.policy_backend, + "enable_policy_switch": args.enable_policy_switch, + "follower": args.follower, + "max_vx": args.max_vx, + "max_vy": args.max_vy, + "max_wz": args.max_wz, + "lookahead": args.lookahead, + "cmd_vx_scale": args.cmd_vx_scale, + "slalom_script": args.slalom_script, + "slalom_script_start_tolerance": args.slalom_script_start_tolerance, + "slalom_script_pos_tolerance": args.slalom_script_pos_tolerance, + "slalom_script_kp_dist": args.slalom_script_kp_dist, + "slalom_script_kp_yaw": args.slalom_script_kp_yaw, + "slalom_script_max_vx": args.slalom_script_max_vx, + "slalom_script_max_vy": args.slalom_script_max_vy, + "slalom_script_max_wz": args.slalom_script_max_wz, + "slalom_script_min_cmd_linear": args.slalom_script_min_cmd_linear, + "slalom_script_min_cmd_angular": args.slalom_script_min_cmd_angular, + "slalom_script_min_cmd_epsilon": args.slalom_script_min_cmd_epsilon, + "slalom_script_yaw_gate_deg": args.slalom_script_yaw_gate_deg, + "slalom_script_drive_yaw_source": args.slalom_script_drive_yaw_source, + "local_safety": args.local_safety, + "safety_horizon": args.safety_horizon, + "start_yaw_offset_deg": args.start_yaw_offset_deg, + "start_x_offset": args.start_x_offset, + "start_y_offset": args.start_y_offset, + "heading_offset_deg": args.heading_offset_deg, + "start_z": args.start_z, + "settle_steps": args.settle_steps, + "viewer": args.viewer, + "start_index": args.start_index, + "end_index": args.end_index, + "skip_initial_waypoint": args.skip_initial_waypoint, + "auto_skip_slice_start": args.auto_skip_slice_start, + "trajectory_csv": str(csv_path) if csv_path else None, + } + (run_dir / "report.json").write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8") + return result, csv_path, run_dir + + +def main() -> int: + args = parse_args() + result, csv_path, run_dir = run_sim(args) + print("Route sim2sim check") + print(f" success: {result.success}") + print(f" reason: {result.reason}") + print(f" reached: {result.reached_count}/{result.waypoint_count}") + print(f" sim_time: {result.sim_time:.2f}s") + print(f" wall_time: {result.wall_time:.2f}s rtf={result.real_time_factor:.2f}x") + print( + " clearance: " + f"min={result.min_clearance:.3f}m margin={result.min_margin:.3f}m " + f"region={result.min_clearance_region} wp={result.min_clearance_wp}" + ) + print( + " attitude: " + f"max_roll={result.max_roll_deg:.1f}deg " + f"max_pitch={result.max_pitch_deg:.1f}deg max_tilt={result.max_tilt_deg:.1f}deg" + ) + print(f" report_dir: {run_dir}") + if csv_path: + print(f" trajectory: {csv_path}") + return 0 if result.success and result.min_margin >= 0.0 else 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/05_software/train/rc_mjlab/sim2sim/nav_sim2sim.py b/05_software/train/rc_mjlab/sim2sim/nav_sim2sim.py index 9e1a740..b798dac 100644 --- a/05_software/train/rc_mjlab/sim2sim/nav_sim2sim.py +++ b/05_software/train/rc_mjlab/sim2sim/nav_sim2sim.py @@ -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% 可视化精准对应 diff --git a/05_software/train/rc_mjlab/sim2sim/policy/policy_runner.py b/05_software/train/rc_mjlab/sim2sim/policy/policy_runner.py index 3b15db4..b5a3134 100644 --- a/05_software/train/rc_mjlab/sim2sim/policy/policy_runner.py +++ b/05_software/train/rc_mjlab/sim2sim/policy/policy_runner.py @@ -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: diff --git a/05_software/train/rc_mjlab/sim2sim/requirements.txt b/05_software/train/rc_mjlab/sim2sim/requirements.txt new file mode 100644 index 0000000..4da9002 --- /dev/null +++ b/05_software/train/rc_mjlab/sim2sim/requirements.txt @@ -0,0 +1,3 @@ +# Additional runtime dependencies for the post-training Sim2Sim tools. +onnxruntime>=1.19.0 +pygame>=2.6.1 diff --git a/05_software/train/rc_mjlab/sim2sim/sim2sim.py b/05_software/train/rc_mjlab/sim2sim/sim2sim.py index 91560c8..c6adf7c 100644 --- a/05_software/train/rc_mjlab/sim2sim/sim2sim.py +++ b/05_software/train/rc_mjlab/sim2sim/sim2sim.py @@ -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" diff --git a/05_software/train/rc_mjlab/tools/nav_tools/README.md b/05_software/train/rc_mjlab/tools/nav_tools/README.md new file mode 100644 index 0000000..2b672c1 --- /dev/null +++ b/05_software/train/rc_mjlab/tools/nav_tools/README.md @@ -0,0 +1,7 @@ +# 路线检查公共模块 + +`route_safety_check.py` 提供航点、避障区域、机器人平面包络和几何距离计算，供 `sim2sim/nav_route_sim2sim_check.py` 复用。 + +该模块只依赖 Python 标准库。原开发目录中的地图编辑器、PCD、比赛路线 JSON、备份和批量实验结果不属于本次后期 Sim2Sim 里程碑，未在这里复制。 + +外部路线文件需要包含 `waypoints`（或 `segments[].waypoints`）以及 `regions` / `avoid_regions`。也可以不提供路线文件，直接使用 Sim2Sim 检查器的内置任务。 diff --git a/05_software/train/rc_mjlab/tools/nav_tools/route_safety_check.py b/05_software/train/rc_mjlab/tools/nav_tools/route_safety_check.py new file mode 100644 index 0000000..1e3dca3 --- /dev/null +++ b/05_software/train/rc_mjlab/tools/nav_tools/route_safety_check.py @@ -0,0 +1,638 @@ +#!/usr/bin/env python3 +"""Offline route safety checker for nav_tools waypoint JSON files. + +The checker treats avoid regions as hard no-go polygons and validates the +route centerline with a circular robot footprint. It is intentionally light on +dependencies so it can run on the robot laptop without ROS, pygame, or shapely. +""" + +from __future__ import annotations + +import argparse +import json +import math +import sys +import xml.etree.ElementTree as ET +from dataclasses import dataclass +from pathlib import Path +from typing import Any + + +ROBOT_BODY_LENGTH = 0.356 +ROBOT_BODY_WIDTH = 0.235 +ROBOT_BODY_CENTER_X = 0.1518 +ROBOT_ORIGIN_FROM_FRONT = 0.105 +ROBOT_WHEEL_VIS_LENGTH = 0.16 +ROBOT_WHEEL_VIS_WIDTH = 0.055 +ROBOT_POSE_HIP = 0.550 +ROBOT_POSE_KNEE = -1.125 +PCD_ROBOT_RADIUS = 0.18 +ROBOT_FOOTPRINT_PADDING = 0.03 + + +@dataclass(frozen=True) +class Waypoint: + index: int + id: str + x: float + y: float + yaw_deg: float + speed: float | None + policy: str + tolerance: float | None + slalom_straight: bool = False + slalom_script_break: bool = False + slalom_script_pos_tolerance: float | None = None + exact_reach: bool = False + precision_follow: bool = False + require_yaw: bool = False + yaw_tolerance_deg: float | None = None + stable_cycles: int | None = None + mandatory_cross: bool = False + mandatory_radius: float | None = None + mandatory_center_x: float | None = None + mandatory_center_y: float | None = None + + +@dataclass(frozen=True) +class AvoidRegion: + name: str + kind: str + polygon: tuple[tuple[float, float], ...] + + +@dataclass(frozen=True) +class SegmentRisk: + start_id: str + end_id: str + region: str + clearance_m: float + required_m: float + margin_m: float + length_m: float + centerline_intersects: bool + + @property + def status(self) -> str: + if self.centerline_intersects: + return "INTERSECT" + if self.margin_m < 0.0: + return "VIOLATION" + if self.margin_m < 0.05: + return "TIGHT" + return "OK" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Check nav_tools waypoint routes against avoid/no-go polygons." + ) + parser.add_argument( + "--points", + type=Path, + default=Path("tools/nav_tools/points/points_20260705_174627.json"), + help="Route JSON exported by nav_map_viewer.", + ) + parser.add_argument( + "--xml", + type=Path, + default=Path("tools/nav_tools/xml/A.xml"), + help="Optional MuJoCo terrain XML used for metadata checks.", + ) + parser.add_argument( + "--onnx", + type=Path, + default=Path("model_6800.onnx"), + help="Optional ONNX policy path used for input/output shape reporting.", + ) + parser.add_argument( + "--footprint-radius", + type=float, + default=None, + help="Robot circular footprint radius in meters. Defaults to sim2real lateral footprint.", + ) + parser.add_argument( + "--avoid-margin", + type=float, + default=0.05, + help="Extra clearance added outside the robot footprint.", + ) + parser.add_argument( + "--warn-margin", + type=float, + default=0.05, + help="Report a TIGHT warning when spare margin is below this value.", + ) + parser.add_argument( + "--top", + type=int, + default=20, + help="Number of closest segment-region pairs to print.", + ) + parser.add_argument( + "--json-out", + type=Path, + default=None, + help="Optional machine-readable report path.", + ) + parser.add_argument( + "--allow-violations", + action="store_true", + help="Exit with code 0 even when violations are detected.", + ) + return parser.parse_args() + + +def load_json(path: Path) -> dict[str, Any]: + with path.open("r", encoding="utf-8") as f: + data = json.load(f) + if not isinstance(data, dict): + raise ValueError(f"JSON root must be an object: {path}") + return data + + +def _optional_float(value: Any) -> float | None: + if value is None: + return None + try: + return float(value) + except (TypeError, ValueError): + return None + + +def load_waypoints(payload: dict[str, Any]) -> list[Waypoint]: + rows = None + if isinstance(payload.get("segments"), list) and payload["segments"]: + rows = [] + for segment in payload["segments"]: + if isinstance(segment, dict) and isinstance(segment.get("waypoints"), list): + rows.extend(segment["waypoints"]) + if rows is None: + rows = payload.get("waypoints") + if not isinstance(rows, list): + raise ValueError("Route JSON has no top-level waypoints or segments[].waypoints") + + waypoints: list[Waypoint] = [] + for index, row in enumerate(rows, start=1): + if not isinstance(row, dict): + continue + x = float(row.get("world_x", row.get("x", 0.0))) + y = float(row.get("world_y", row.get("y", 0.0))) + yaw = float(row.get("yawDeg", row.get("yaw_deg", row.get("yaw", 0.0)))) + speed = row.get("speed") + tolerance = row.get("tolerance") + waypoints.append( + Waypoint( + index=index, + id=str(row.get("id", index)), + x=x, + y=y, + yaw_deg=yaw, + speed=float(speed) if speed is not None else None, + policy=str(row.get("policy", "")), + tolerance=float(tolerance) if tolerance is not None else None, + slalom_straight=bool(row.get("slalom_straight", row.get("slalomStraight", False))), + slalom_script_break=bool( + row.get("slalom_script_break", row.get("slalomScriptBreak", False)) + ), + slalom_script_pos_tolerance=_optional_float( + row.get( + "slalom_script_pos_tolerance", + row.get("slalomScriptPosTolerance", row.get("scriptTolerance")), + ) + ), + exact_reach=bool(row.get("exact_reach", row.get("exactReach", False))), + precision_follow=bool(row.get("precision_follow", row.get("precisionFollow", False))), + require_yaw=bool(row.get("require_yaw", row.get("requireYaw", False))), + yaw_tolerance_deg=_optional_float( + row.get("yaw_tolerance_deg", row.get("yawToleranceDeg")) + ), + stable_cycles=( + int(row.get("stable_cycles", row.get("stableCycles"))) + if row.get("stable_cycles", row.get("stableCycles")) is not None + else None + ), + mandatory_cross=bool(row.get("mandatory_cross", row.get("mandatoryCross", False))), + mandatory_radius=_optional_float( + row.get("mandatory_radius", row.get("mandatoryRadius")) + ), + mandatory_center_x=_optional_float( + row.get("mandatory_center_x", row.get("mandatoryCenterX")) + ), + mandatory_center_y=_optional_float( + row.get("mandatory_center_y", row.get("mandatoryCenterY")) + ), + ) + ) + if len(waypoints) < 2: + raise ValueError("Route must contain at least two waypoints") + return waypoints + + +def load_regions(payload: dict[str, Any]) -> list[AvoidRegion]: + rows = payload.get("regions", payload.get("avoid_regions", [])) + if not isinstance(rows, list): + return [] + + regions: list[AvoidRegion] = [] + for index, row in enumerate(rows, start=1): + if not isinstance(row, dict): + continue + polygon_rows = row.get("polygon", row.get("points", [])) + if not isinstance(polygon_rows, list): + continue + polygon: list[tuple[float, float]] = [] + for point in polygon_rows: + if isinstance(point, dict): + polygon.append((float(point.get("x", 0.0)), float(point.get("y", 0.0)))) + elif isinstance(point, (list, tuple)) and len(point) >= 2: + polygon.append((float(point[0]), float(point[1]))) + if len(polygon) >= 3: + regions.append( + AvoidRegion( + name=str(row.get("name", f"avoid_{index}")), + kind=str(row.get("kind", "avoid")), + polygon=tuple(polygon), + ) + ) + return regions + + +def robot_wheel_local_points(body_center_offset_x: float) -> list[tuple[float, float]]: + thigh_dx = -0.25 * math.sin(ROBOT_POSE_HIP) + shank_dx = -0.2 * math.sin(ROBOT_POSE_HIP + ROBOT_POSE_KNEE) + wheel_positions = ( + ((0.32826 + 0.06389) - ROBOT_BODY_CENTER_X, 0.066172 - 0.027344, 0.1035, 0.014699, 0.04074, 0.0), + ((0.32826 + 0.06389) - ROBOT_BODY_CENTER_X, -0.065853 + 0.027311, -0.1035, -0.018447, -0.040735, -0.00075079), + ((-0.024743 - 0.06389) - ROBOT_BODY_CENTER_X, 0.066141 - 0.027309, 0.099459, 0.012475, 0.040737, 0.0), + ((-0.024743 - 0.06389) - ROBOT_BODY_CENTER_X, -0.065884 + 0.027341, -0.099408, -0.012435, -0.040737, -0.00075079), + ) + return [ + ( + body_center_offset_x + pitch_x + knee_x + thigh_dx + shank_dx, + pitch_y + knee_y + wheel_y + wheel_geom_y, + ) + for pitch_x, pitch_y, knee_y, wheel_y, wheel_geom_y, knee_x in wheel_positions + ] + + +def default_lateral_footprint_radius() -> float: + half_width = ROBOT_BODY_WIDTH * 0.5 + radius = max(PCD_ROBOT_RADIUS, half_width) + body_center_offset_x = ROBOT_ORIGIN_FROM_FRONT - ROBOT_BODY_LENGTH * 0.5 + for _, wheel_y in robot_wheel_local_points(body_center_offset_x): + radius = max(radius, abs(wheel_y) + ROBOT_WHEEL_VIS_WIDTH * 0.5) + return radius + ROBOT_FOOTPRINT_PADDING + + +def point_segment_distance( + px: float, + py: float, + ax: float, + ay: float, + bx: float, + by: float, +) -> float: + dx = bx - ax + dy = by - ay + length_sq = dx * dx + dy * dy + if length_sq <= 1.0e-12: + return math.hypot(px - ax, py - ay) + t = ((px - ax) * dx + (py - ay) * dy) / length_sq + t = max(0.0, min(1.0, t)) + qx = ax + t * dx + qy = ay + t * dy + return math.hypot(px - qx, py - qy) + + +def orientation( + ax: float, + ay: float, + bx: float, + by: float, + cx: float, + cy: float, +) -> float: + return (bx - ax) * (cy - ay) - (by - ay) * (cx - ax) + + +def on_segment( + ax: float, + ay: float, + bx: float, + by: float, + cx: float, + cy: float, +) -> bool: + return ( + min(ax, bx) - 1.0e-9 <= cx <= max(ax, bx) + 1.0e-9 + and min(ay, by) - 1.0e-9 <= cy <= max(ay, by) + 1.0e-9 + and abs(orientation(ax, ay, bx, by, cx, cy)) <= 1.0e-9 + ) + + +def segments_intersect( + a: tuple[float, float], + b: tuple[float, float], + c: tuple[float, float], + d: tuple[float, float], +) -> bool: + ax, ay = a + bx, by = b + cx, cy = c + dx, dy = d + o1 = orientation(ax, ay, bx, by, cx, cy) + o2 = orientation(ax, ay, bx, by, dx, dy) + o3 = orientation(cx, cy, dx, dy, ax, ay) + o4 = orientation(cx, cy, dx, dy, bx, by) + if o1 * o2 < 0.0 and o3 * o4 < 0.0: + return True + return ( + on_segment(ax, ay, bx, by, cx, cy) + or on_segment(ax, ay, bx, by, dx, dy) + or on_segment(cx, cy, dx, dy, ax, ay) + or on_segment(cx, cy, dx, dy, bx, by) + ) + + +def point_in_polygon(x: float, y: float, polygon: tuple[tuple[float, float], ...]) -> bool: + inside = False + for index, (ax, ay) in enumerate(polygon): + bx, by = polygon[(index + 1) % len(polygon)] + if point_segment_distance(x, y, ax, ay, bx, by) <= 1.0e-9: + return True + if (ay > y) != (by > y): + x_cross = (bx - ax) * (y - ay) / (by - ay) + ax + if x < x_cross: + inside = not inside + return inside + + +def segment_polygon_intersects( + a: tuple[float, float], + b: tuple[float, float], + polygon: tuple[tuple[float, float], ...], +) -> bool: + if point_in_polygon(a[0], a[1], polygon) or point_in_polygon(b[0], b[1], polygon): + return True + return any( + segments_intersect(a, b, polygon[index], polygon[(index + 1) % len(polygon)]) + for index in range(len(polygon)) + ) + + +def segment_polygon_distance( + a: tuple[float, float], + b: tuple[float, float], + polygon: tuple[tuple[float, float], ...], +) -> float: + if segment_polygon_intersects(a, b, polygon): + return 0.0 + distances = [point_segment_distance(px, py, a[0], a[1], b[0], b[1]) for px, py in polygon] + for index, (ax, ay) in enumerate(polygon): + bx, by = polygon[(index + 1) % len(polygon)] + distances.append(point_segment_distance(a[0], a[1], ax, ay, bx, by)) + distances.append(point_segment_distance(b[0], b[1], ax, ay, bx, by)) + return min(distances) + + +def analyze_route( + waypoints: list[Waypoint], + regions: list[AvoidRegion], + required_clearance: float, +) -> list[SegmentRisk]: + risks: list[SegmentRisk] = [] + for start, end in zip(waypoints, waypoints[1:]): + a = (start.x, start.y) + b = (end.x, end.y) + length = math.hypot(end.x - start.x, end.y - start.y) + for region in regions: + intersects = segment_polygon_intersects(a, b, region.polygon) + clearance = 0.0 if intersects else segment_polygon_distance(a, b, region.polygon) + risks.append( + SegmentRisk( + start_id=start.id, + end_id=end.id, + region=region.name, + clearance_m=clearance, + required_m=required_clearance, + margin_m=clearance - required_clearance, + length_m=length, + centerline_intersects=intersects, + ) + ) + risks.sort(key=lambda item: (item.margin_m, item.clearance_m)) + return risks + + +def parse_xml_summary(path: Path) -> dict[str, Any]: + if not path.exists(): + return {"path": str(path), "exists": False} + root = ET.parse(path).getroot() + geoms = [geom for geom in root.iter("geom")] + collidable = [ + geom for geom in geoms + if geom.get("name") != "floor" + and geom.get("contype", "1") != "0" + and geom.get("conaffinity", "1") != "0" + ] + return { + "path": str(path), + "exists": True, + "model": root.get("model", ""), + "geom_count": len(geoms), + "collidable_geom_count": len(collidable), + } + + +def parse_onnx_summary(path: Path) -> dict[str, Any]: + if not path.exists(): + return {"path": str(path), "exists": False} + try: + import onnx # type: ignore + except Exception as exc: # pragma: no cover - depends on local env + return {"path": str(path), "exists": True, "error": f"onnx import failed: {exc}"} + + model = onnx.load(str(path)) + inputs = [ + { + "name": item.name, + "shape": [ + dim.dim_value if dim.dim_value else dim.dim_param + for dim in item.type.tensor_type.shape.dim + ], + } + for item in model.graph.input + ] + outputs = [ + { + "name": item.name, + "shape": [ + dim.dim_value if dim.dim_value else dim.dim_param + for dim in item.type.tensor_type.shape.dim + ], + } + for item in model.graph.output + ] + return { + "path": str(path), + "exists": True, + "inputs": inputs, + "outputs": outputs, + "metadata_keys": [prop.key for prop in model.metadata_props], + } + + +def risk_to_dict(risk: SegmentRisk) -> dict[str, Any]: + return { + "start_id": risk.start_id, + "end_id": risk.end_id, + "region": risk.region, + "clearance_m": round(risk.clearance_m, 6), + "required_m": round(risk.required_m, 6), + "margin_m": round(risk.margin_m, 6), + "length_m": round(risk.length_m, 6), + "centerline_intersects": risk.centerline_intersects, + "status": risk.status, + } + + +def print_report( + points_path: Path, + xml_summary: dict[str, Any], + onnx_summary: dict[str, Any], + waypoints: list[Waypoint], + regions: list[AvoidRegion], + footprint_radius: float, + avoid_margin: float, + warn_margin: float, + risks: list[SegmentRisk], + top: int, +) -> None: + required_clearance = footprint_radius + avoid_margin + violations = [risk for risk in risks if risk.margin_m < 0.0 or risk.centerline_intersects] + tight = [ + risk for risk in risks + if risk.margin_m >= 0.0 and risk.margin_m < warn_margin + ] + route_len = sum( + math.hypot(b.x - a.x, b.y - a.y) + for a, b in zip(waypoints, waypoints[1:]) + ) + + print("Route safety check") + print(f" points: {points_path}") + print(f" waypoints: {len(waypoints)}, regions: {len(regions)}, path_length: {route_len:.3f} m") + print( + " clearance: " + f"footprint={footprint_radius:.3f} m + avoid_margin={avoid_margin:.3f} m " + f"=> required={required_clearance:.3f} m" + ) + if xml_summary.get("exists"): + print( + " xml: " + f"{xml_summary.get('path')} " + f"(model={xml_summary.get('model')}, geoms={xml_summary.get('geom_count')}, " + f"collidable={xml_summary.get('collidable_geom_count')})" + ) + else: + print(f" xml: missing ({xml_summary.get('path')})") + if onnx_summary.get("exists") and not onnx_summary.get("error"): + print(f" onnx: {onnx_summary.get('path')}") + print(f" inputs: {onnx_summary.get('inputs')}") + print(f" outputs: {onnx_summary.get('outputs')}") + elif onnx_summary.get("exists"): + print(f" onnx: {onnx_summary.get('error')}") + else: + print(f" onnx: missing ({onnx_summary.get('path')})") + + print("") + if violations: + print(f"FAIL: {len(violations)} segment-region pairs are inside required clearance.") + elif tight: + print(f"WARN: no violations, but {len(tight)} segment-region pairs are tight.") + else: + print("PASS: all segment-region pairs satisfy the requested clearance.") + + print("") + print(f"Closest {min(top, len(risks))} segment-region pairs:") + print(" status wp_start->wp_end region clear req spare") + for risk in risks[:top]: + print( + f" {risk.status:<10} " + f"{risk.start_id:>4}->{risk.end_id:<4} " + f"{risk.region:<10} " + f"{risk.clearance_m:>6.3f} " + f"{risk.required_m:>6.3f} " + f"{risk.margin_m:>7.3f}" + ) + + +def main() -> int: + args = parse_args() + points_path = args.points.resolve() + xml_path = args.xml.resolve() + onnx_path = args.onnx.resolve() + + payload = load_json(points_path) + waypoints = load_waypoints(payload) + regions = load_regions(payload) + if not regions: + raise ValueError(f"No avoid regions found in {points_path}") + + footprint_radius = ( + float(args.footprint_radius) + if args.footprint_radius is not None + else default_lateral_footprint_radius() + ) + required_clearance = footprint_radius + float(args.avoid_margin) + risks = analyze_route(waypoints, regions, required_clearance) + xml_summary = parse_xml_summary(xml_path) + onnx_summary = parse_onnx_summary(onnx_path) + + print_report( + points_path, + xml_summary, + onnx_summary, + waypoints, + regions, + footprint_radius, + float(args.avoid_margin), + float(args.warn_margin), + risks, + max(0, int(args.top)), + ) + + violations = [risk for risk in risks if risk.margin_m < 0.0 or risk.centerline_intersects] + tight = [ + risk for risk in risks + if risk.margin_m >= 0.0 and risk.margin_m < float(args.warn_margin) + ] + report = { + "points": str(points_path), + "waypoint_count": len(waypoints), + "region_count": len(regions), + "footprint_radius_m": round(footprint_radius, 6), + "avoid_margin_m": round(float(args.avoid_margin), 6), + "required_clearance_m": round(required_clearance, 6), + "violations": [risk_to_dict(risk) for risk in violations], + "tight": [risk_to_dict(risk) for risk in tight], + "closest": [risk_to_dict(risk) for risk in risks[: max(0, int(args.top))]], + "xml": xml_summary, + "onnx": onnx_summary, + } + if args.json_out: + args.json_out.parent.mkdir(parents=True, exist_ok=True) + args.json_out.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8") + + if violations and not args.allow_violations: + return 2 + return 0 + + +if __name__ == "__main__": + try: + raise SystemExit(main()) + except Exception as exc: + print(f"ERROR: {exc}", file=sys.stderr) + raise SystemExit(1) diff --git a/README.md b/README.md index e6b0c26..dd76f47 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ RC_WheelLeg 是山东华宇工学院 16DOF 串联轮足机器人项目。 -当前 `16dof` 分支用于整理 16DOF 机械、强化学习训练、Sim2Sim、Sim2Real、ROS 2 部署和比赛版本。机械资料、比赛训练架构和后期 MuJoCo 工具集已经完成整理。 +当前 `16dof` 分支用于整理 16DOF 机械、强化学习训练、Sim2Sim、Sim2Real、ROS 2 部署和比赛版本。机械资料、比赛训练架构、后期 MuJoCo 工具集和后期 Sim2Sim 已经完成整理。 ## 平台概览 @@ -37,6 +37,7 @@ RC_WheelLeg/ - [x] 整理第二版 Sim2Real 随机化训练配置 - [x] 整理比赛最终训练代码架构 - [x] 整理后期 MuJoCo 姿态、IK、动力学和 MPC 工具 +- [x] 整理后期 Sim2Sim、路线检查与比赛 Rough ONNX 策略 - [ ] 核对比赛机械与仿真模型参数 - [ ] 整理 URDF/MJCF 机器人描述 - [ ] 整理后续统一训练、ROS 2 和比赛版本