[tools] 整理导航打点工具与抽样PCD

[sim] 整理后期Sim2Sim与比赛Rough策略
[sim] 添加后期MuJoCo姿态与MPC工具集
2026-07-27 13:39:52 +08:00 · 2026-07-27 13:20:41 +08:00 · 2026-07-27 12:57:03 +08:00 · 2026-07-27 12:48:12 +08:00 · 2026-07-27 12:39:27 +08:00 · 2026-07-27 12:31:21 +08:00
171 changed files with 444259 additions and 5431 deletions
@@ -35,6 +35,10 @@ logs/
 checkpoints/
 wandb/
 sim2sim_log_*.txt
+**/sim2sim_temp.xml
+**/route_check_runs/
+**/route_experiments/suite_*/
+**/tools/nav_tools/points/auto_candidates/

 # IDE and operating system files
 .idea/
@@ -0,0 +1,37 @@
+# 训练代码演进
+
+本项目将“训练代码架构”和“训练产生的模型 checkpoint”分别管理。代码、奖励、课程或观测契约发生变化时形成新的 Git 版本；同一架构下继续训练产生的模型编号作为实验和部署工件记录。
+
+| 版本 | 来源快照 | 主要目的 |
+| --- | --- | --- |
+| `v0.4.0` | `uni_mjlab(1)` | 第一份完整的新 MJCF 与新 mjlab 训练工程 |
+| `v0.5.0` | `uni_mjlab_new` | 扩大观测、延迟和动力学随机化，探索更强 Sim2Real 鲁棒性 |
+| `v0.6.0` | `best` | 面向比赛越障重新设计奖励、课程、站姿和诊断体系 |
+
+## 从随机化增强到比赛训练
+
+`v0.6.0` 不是简单地继续增大 `v0.5.0` 的随机范围。实践中重新降低了部分噪声、延迟和动力学随机化强度，并把训练重点转向可控的比赛任务课程：
+
+- 投影重力噪声由 `±0.08` 调回 `±0.05`。
+- 腿和轮动作最大随机延迟由 4 步调回 2 步。
+- 摩擦、刚度和阻尼随机化调回较窄范围。
+- 移除该阶段的连续机身外力扰动与腿部质量随机化。
+- 将 x、y 和 yaw 跟踪拆成独立奖励和独立指令课程。
+- 障碍由课程逐步释放，高墙训练地形改为五道重复横墙。
+- 增加楼梯侧向漂移和偏航漂移惩罚。
+- 增加大量只用于训练评估的误差、轮速、姿态和接触指标。
+
+这种调整反映的是从“广泛鲁棒性探索”转向“比赛场景定向优化”，不表示 `v0.5.0` 被删除；它仍由对应 Tag 完整保留。
+
+## 模型与比赛部署
+
+比赛训练可能先获得基模，再修改参数继续训练和筛选 checkpoint。最终关系为：
+
+```text
+训练代码架构：v0.6.0 / best
+比赛 Rough 策略：model_6800.onnx
+比赛真机工程：last_not_slalom_1050
+比赛得分：1050
+```
+
+`model_6800.onnx` 是比赛最终部署工件，不用模型编号替代训练代码版本号。它将在最终比赛部署版本中与运行配置一起归档。
@@ -0,0 +1,73 @@
+# 版本演进
+
+本项目使用同一条 `16dof` 主线和里程碑 Tag 保存线性演进，不在源码目录中复制历史版本。
+
+| Tag | 阶段 | 核心内容 |
+| --- | --- | --- |
+| `v0.1.0` | 8DOF 中期检查 | 8DOF 串联足机械与大疆 A 板实机版本 |
+| `v0.2.0` | 16DOF 机械 | 16DOF 串联轮足机械 CAD 与 STEP |
+| `v0.3.0` | 第一代软件闭环 | 早期训练、MJCF、MuJoCo、Sim2Sim、IK 与 Python Sim2Real |
+| `v0.3.1` | 实机记录 | 补充第一代 Sim2Real 实机视频 |
+| `v0.4.0` | 新训练基线 | 第一份完整的新 MJCF、新 mjlab 框架和 Rough 策略工程 |
+| `v0.5.0` | 随机化增强 | 扩大观测、延迟和动力学随机化，加入持续外力扰动 |
+| `v0.6.0` | 比赛训练架构 | 分轴奖励、自适应指令课程、障碍释放课程和比赛站姿 |
+| `v0.7.0` | MuJoCo 工具 | 姿态优化、IK 扫描、动力学、MPC 和 GUI 调试工具 |
+| `v0.8.0` | 后期 Sim2Sim | ONNX 回放、IK/路线检查工具和比赛最终 Rough 策略 |
+| `v0.8.1` | 导航打点工具 | 地图/航点编辑、路线迭代和抽样 PCD 补充包 |
+
+## `v0.4.0` 的模型变化
+
+- 机械 CAD 不变。
+- MJCF 更新整机质量和惯性参数，旧、新 `wheelleg.xml` 的 SHA-256 不同。
+- mjlab 上游基准从 `00409797` 更新到 `40f8d93e`。
+- 保留轮腿分组执行器随机化所需的本地补丁。
+- 本阶段归档 `model_rough.pt`，不将生成日志、缓存和临时 XML 纳入版本库。
+
+## `v0.5.0` 的训练变化
+
+- MJCF、mjlab 基准和已有模型文件保持不变。
+- 投影重力噪声由 `±0.05` 扩大到 `±0.08`。
+- 腿与轮动作的最大随机延迟由 2 步增加到 4 步。
+- 地面摩擦随机范围由 `0.3–1.0` 扩大到 `0.15–1.25`。
+- 执行器刚度和阻尼缩放由 `0.9–1.1` 扩大到 `0.5–1.5`。
+- 增加膝部和轮部质量的 `0.7–1.3` 随机缩放。
+- 增加作用于机身的连续随机外力和力矩扰动。
+
+## `v0.6.0` 的比赛训练架构
+
+- 保持 `v0.4.0` 引入的新 MJCF 和 mjlab 框架不变。
+- 将线速度奖励拆分为 x/y 两轴，并独立配置偏航角速度奖励。
+- 增加自适应 x/y/yaw 指令范围课程。
+- 增加障碍地形逐步释放与更严格的地形晋级逻辑。
+- 增加楼梯横向速度和偏航漂移约束。
+- 默认站姿调整为髋俯仰 `0.550`、膝关节 `-1.125`，初始机身高度为 `0.42 m`。
+- 增加速度误差、轮速跟踪、动作和姿态等训练诊断指标。
+- 该 Tag 保存比赛训练代码架构，不把每次继续训练产生的 checkpoint 误记为新的软件版本。
+
+训练阶段的详细关系见 [`training_evolution.md`](training_evolution.md)。
+
+## `v0.7.0` 的独立 MuJoCo 工具
+
+- 比赛训练架构、MJCF、模型和依赖锁文件保持 `v0.6.0` 状态不变。
+- 增加解析姿态表、RL 友好姿态筛选和 MuJoCo 静态姿态优化。
+- 增加 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 和大量重复路线实验不在本阶段归档。
+
+## `v0.8.1` 的导航打点工具
+
+- 补充 Pygame 地图/PCD/航点综合编辑器、避障区域编辑器和坐标变换工具。
+- 补充路线安全检查、候选航点优化、XML/航点镜像和批量 Sim2Sim 实验入口。
+- 按源文件时间保留 14 份比赛路线 JSON，不将开发期文件名误解释为正式版本号。
+- 补充 `1hao.xml`、`2hao.xml`、`A_C.xml`，并为 `1B_FF.json` 补齐其引用的 `B_C.xml`。
+- 将两份约 915 MiB 的原始 ASCII PCD 确定性抽样为各小于 10 MB 的预览点云；抽样参数、点数和哈希记录在工具 README。
+- 训练代码、MJCF、比赛策略和历史依赖锁保持 `v0.8.0` 状态不变。
@@ -1,6 +1,6 @@
 # 软件

-本目录当前保存 16DOF 轮足机器人的第一代完整软件闭环。
+本目录保存 16DOF 轮足机器人的训练、仿真和真机软件演进。

 ```text
 05_software/
@@ -19,7 +19,7 @@ MJCF + mjlab task
        v
   PPO 训练策略
        |
-        +----> MuJoCo 独立模型调试
+        +----> MuJoCo 姿态 / IK / MPC 调试
        |
        +----> Sim2Sim 策略验证
        |
@@ -28,7 +28,7 @@ MJCF + mjlab task
 IK real --------------------------------> 电机
 ```

-`rc_mjlab` 在早期版本中是自包含工程。训练、MJCF、独立 MuJoCo、Sim2Sim 和策略权重通过相对路径绑定，因此本次保留其原始内部布局，没有为了目录外观拆散。
+`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`，完整导航打点工具见 `v0.8.1`。

 详细说明见：

@@ -1,15 +1,31 @@
-# 第一代强化学习与仿真工程
+# 强化学习与仿真工程

-`rc_mjlab/` 是 16DOF 轮足机器人的第一代自包含训练与仿真工程。
+`rc_mjlab/` 保存 16DOF 轮足机器人的当前训练与 Sim2Sim 工程。历史快照由 Git Tag 保留，不在目录中复制 `old`、`new` 或 `final` 版本。
+
+当前内容对应 `v0.8.1`：训练代码保持 `v0.6.0` 的比赛架构，包含后期 MuJoCo、Sim2Sim、比赛最终 Rough ONNX 策略，并补充完整导航打点工具、路线迭代和抽样 PCD。训练过程可能先获得基模，再调整奖励、课程和环境参数继续训练；模型 checkpoint 的变化不等同于软件架构变化。

 ## 内容

 - `src/robot`：Flat、Rough、Crawl 训练任务和自定义 MDP
 - `mjcf`：轮足机器人 MuJoCo 模型和网格
- `mujoco_sim`：不依赖策略的独立 MuJoCo/MPC 调试工具
 - `sim2sim`：策略加载、交互控制和比赛地形验证
+- `mujoco_sim`：不依赖训练循环的姿态、IK、动力学和 MPC 分析
+- `tools/nav_tools`：地图/PCD 查看、航点编辑、路线检查和比赛路线数据
 - `mjlab`：固定版本的本地训练框架依赖
- `model_rough.pt`、`model_crawl.pt`：对应的早期策略权重
+- `model_rough.pt`：本阶段 Rough 策略权重
+- `model_6800.onnx`：比赛最终使用的 Rough 策略
 - `pyproject.toml`、`uv.lock`：Python 环境与依赖锁定

+与 `v0.3.0` 相比，本版本更新了 MJCF 质量和惯性参数，并将 mjlab 上游基准从 `00409797` 更新到 `40f8d93e`。机械 CAD 未发生变化。
+
+与 `v0.4.0` 相比，本版本没有再次修改 MJCF 和训练框架，只调整训练环境配置：投影重力噪声从 `±0.05` 扩大到 `±0.08`，最大动作延迟从 2 步增加到 4 步，扩大摩擦、刚度和阻尼随机化，增加腿部质量随机化与连续外力/力矩扰动。
+
+`v0.6.0` 在 `v0.5.0` 之后转向比赛任务优化：降低部分过强随机化，加入分轴速度跟踪奖励、自适应指令课程、障碍地形释放课程、楼梯横向/偏航约束以及更完整的训练诊断。详细对比见 [`../../01_doc/training_evolution.md`](../../01_doc/training_evolution.md)。
+
+`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)。
+
+`v0.8.1` 补充完整导航打点工具和小体积预览点云；入口、路线清单和抽样边界见 [`rc_mjlab/tools/nav_tools/README.md`](rc_mjlab/tools/nav_tools/README.md)。
+
 工程命令和任务说明见 [`rc_mjlab/README.md`](rc_mjlab/README.md)，本地依赖来源见 [`rc_mjlab/DEPENDENCIES.md`](rc_mjlab/DEPENDENCIES.md)。
@@ -2,14 +2,16 @@

 ## Python 环境

- Python `>=3.10`
+- Python `>=3.10,<3.14`
 - `uv` 依赖管理
- MuJoCo development wheel
+- MuJoCo `3.8` 系列
 - `mjlab[cu128]`
 - PyTorch CUDA 12.8 环境
 - `pynput`
+- 后期 Sim2Sim 可选依赖：Pygame、ONNX Runtime
+- 导航打点可选依赖：Pygame、Pillow

-精确解析结果保存在 `uv.lock`。项目使用本地可编辑 `mjlab`：
+精确解析结果保存在 `uv.lock`。项目使用本地可编辑的 `mjlab`：

 ```toml
 [tool.uv.sources]
@@ -19,17 +21,15 @@ mjlab = { path = "mjlab", editable = true }
 ## mjlab 来源

 - 上游仓库：`https://github.com/mujocolab/mjlab.git`
- 基准提交：`0040979763ab43bc1220812c9de4bc74e2631f42`
- 基准日期：`2026-04-28`
+- 基准提交：`40f8d93e31b589dccae78ba6aadfc4b74cd1e3fd`
+- 基准日期：`2026-06-02`
 - 上游许可证：Apache-2.0，许可证文件保留在 `mjlab/LICENSE`

-早期工程在该基准上保留了 3 处本地修改：
+本版本在该基准上保留 1 处本地修改：

-1. `mjlab/pyproject.toml`：增加清华 PyPI 镜像。
-2. `mjlab/src/mjlab/envs/mdp/dr/actuator.py`：让 effort limit 随机化支持轮子使用的 velocity/motor actuator。
-3. `mjlab/src/mjlab/scene/scene.py`：通过 XML 字符串加载场景，以适配当时的场景组合方式。
+1. `mjlab/src/mjlab/envs/mdp/dr/actuator.py`：为分组执行器补充名称到运行时执行器对象的解析，使 PD 增益和力矩限制随机化能够正确作用于轮腿机器人的执行器组。

-本次归档保留修改后的完整工作树，但不包含上游 `.git`、本地 `.venv`、缓存和生成日志。
+本次归档保留修改后的完整工作树，但不包含上游 `.git`、本地缓存、生成日志和运行时临时文件。

 ## 基本入口

@@ -41,4 +41,16 @@ 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
+```
+
+导航打点工具同样不修改历史锁文件：
+
+```bash
+uv run --with-requirements tools/nav_tools/requirements.txt python tools/nav_tools/nav_map_viewer.py
+```
+
 GPU、CUDA、MuJoCo development wheel 和驱动版本必须满足 `pyproject.toml` 与 `uv.lock` 的约束。
@@ -2,6 +2,8 @@

 基于 [mjlab](https://github.com/google-deepmind/mjlab) 框架的四轮腿混合机器人强化学习训练与部署部署项目，面向机器人竞赛场景（如越障、匍匐、斜坡、台阶等复合任务）。

+> 当前目录对应 `v0.8.1`：保留比赛训练架构、后期 MuJoCo/Sim2Sim 和比赛最终 Rough 策略，并补充导航地图、打点工具、路线迭代与抽样 PCD。`model_rough.pt` 仍作为早期参考权重保留。
+
 ---

 ## 🛠️ 项目简介
@@ -39,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 输入输出与传感器、低通滤波器接口
@@ -51,8 +57,10 @@ rc_mjlab/
 │   ├── wheelleg.xml              # 机器人 MuJoCo 模型（含网格引用）
 │   ├── scene.xml                 # mjlab 场景入口文件
 │   └── meshes/                   # STL/OBJ 碰撞与外观网格
-├── mujoco_sim/                   # 独立 MPC 仿真调试工具（不依赖 RL 训练）
-├── logs/                         # 训练日志（rsl_rl 格式，按任务名/日期/checkpoint 归档）
+├── mujoco_sim/                   # 姿态、IK、动力学和 MPC 独立工具
+├── tools/nav_tools/              # 地图/PCD/航点编辑、路线检查与比赛路线数据
+├── model_rough.pt                # 早期 Rough 参考 checkpoint
+├── model_6800.onnx               # 比赛最终 Rough 策略
 ├── pyproject.toml                # 项目依赖（uv 管理，含清华镜像源加速）
 └── uv.lock                       # 精确依赖锁定文件
 ```
@@ -87,6 +95,10 @@ cd sim2sim
 uv run python nav_sim2sim.py
 ```

+后期 Sim2Sim 的入口、模型边界和批量检查命令见 [`sim2sim/README.md`](sim2sim/README.md)。
+
+导航打点工具、路线快照和抽样点云说明见 [`tools/nav_tools/README.md`](tools/nav_tools/README.md)。
+
 ---

 ## 🖥️ 交互式自动导航平台 (sim2sim/nav_sim2sim.py)
@@ -1,327 +0,0 @@
-<mujoco model="go2w scene">
-<include file="C:/Users/31560/Documents/00_legged/new_rl/rc_mjlab/mjcf/wheelleg.xml"/>
-<statistic center="3.7 -9.0 0.4" extent="5.0"/>
-<visual>
-<headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0"/>
-<rgba haze="0.15 0.25 0.35 1"/>
-<global azimuth="90" elevation="-20"/>
-</visual>
-<asset>
-<texture type="skybox" builtin="gradient" rgb1="0.3 0.5 0.7" rgb2="0 0 0" width="512" height="3072"/>
-<texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3" markrgb="0.8 0.8 0.8" width="300" height="300"/>
-<material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance="0.2"/>
-<hfield name="perlin_hfield" size="1.0 0.75 0.2 0.2" file="C:/Users/31560/Documents/00_legged/new_rl/rc_mjlab/sim2sim/terrain/height_field.png"/>
-<hfield name="image_hfield" size="1.0 1.0 0.02 0.1" file="C:/Users/31560/Documents/00_legged/new_rl/rc_mjlab/sim2sim/terrain/unitree_hfield.png"/>
-</asset>
-<worldbody>
-<light pos="0 0 1.5" dir="0 0 -1" directional="true" />
-<geom name="floor" size="0 0 0.05" type="plane" material="groundplane" />
-
-<!-- 30cm高墙：旋转90度，沿x轴方向放置，并与T型楼梯中心线 y=-3.50 对齐 -->
-<geom pos="1.8 -7.0 0.15"
-      type="box"
-      size="0.025 0.5 0.15"
-      quat="0.7071068 0.0 0.0 0.7071068" 
-      rgba="1.0 0.9 0.4 1.0"/>
-
-<!-- 沙砾碎木坑：x正方向边界与10度斜坡+x边界对齐，y正边界距斜坡y负边界4m -->
-
-<geom pos="4.8361 -12.5 0.075"
-      type="box"
-      size="0.5 0.5 0.075"
-      quat="0.0 0.0 0.0 1.0" 
-      rgba="0.75 0.72 0.55 1.0"/>
-
-<geom pos="5.8361 -12.0 0.075"
-      type="box"
-      size="0.5 1.0 0.075"
-      quat="0.0 0.0 0.0 1.0" 
-      rgba="0.75 0.72 0.55 1.0"/>
-
-
-
-
-<!-- 限高杆 -->
-<geom pos="6.2 -9.0 0.155"
-      type="cylinder"
-      size="0.025 0.155"
-      quat="1.0 0.0 0.0 0.0"
-      rgba="0.8 0.1 0.1 1.0" />
-
-<geom pos="5.2 -9.0 0.155"
-      type="cylinder"
-      size="0.025 0.155"
-      quat="1.0 0.0 0.0 0.0"
-      rgba="0.8 0.1 0.1 1.0" />
-<geom pos="5.7 -9.0 0.325"
-      type="cylinder"
-      size="0.015 0.5"
-      quat="0.7071068 0.0 0.7071068 0.0" 
-      rgba="1.0 0.9 0.4 1.0"/>
-
-
-<!-- 1m × 1m 正方形颜色块，出发区-->
-<geom pos="3.7 -9.0 0.0"
-      type="box"
-      size="0.5 0.5 0.001"
-      rgba="1.0 0.0 0.0 0.35"
-      contype="0"
-      conaffinity="0" />
-
-
-
-<!-- 10cm梯形台阶：T型楼梯，最高平台与 x=5.7 y=-3.5 平台在y轴方向对齐 -->
-
-<geom pos="1.80 -4.75 0.05" type="box" size="0.15 0.5 0.05" quat="0.7071068 0.0 0.0 0.7071068" rgba="0.75 0.72 0.55 1.0" />
-<geom pos="1.80 -4.45 0.15" type="box" size="0.15 0.5 0.05" quat="0.7071068 0.0 0.0 0.7071068" rgba="0.75 0.72 0.55 1.0"/>
-<geom pos="1.80 -4.15 0.25" type="box" size="0.15 0.5 0.05" quat="0.7071068 0.0 0.0 0.7071068" rgba="0.75 0.72 0.55 1.0"/>
-
-<!-- 最高平台：y = -3.50，与目标平台y轴对齐 -->
-<geom pos="1.80 -3.50 0.35" type="box" size="0.5 0.5 0.05" quat="0.7071068 0.0 0.0 0.7071068" rgba="0.75 0.72 0.55 1.0"/>
-
-<geom pos="1.80 -2.85 0.25" type="box" size="0.15 0.5 0.05" quat="0.7071068 0.0 0.0 0.7071068" rgba="0.75 0.72 0.55 1.0"/>
-<geom pos="1.80 -2.55 0.15" type="box" size="0.15 0.5 0.05" quat="0.7071068 0.0 0.0 0.7071068" rgba="0.75 0.72 0.55 1.0"/>
-<geom pos="1.80 -2.25 0.05" type="box" size="0.15 0.5 0.05" quat="0.7071068 0.0 0.0 0.7071068" rgba="0.75 0.72 0.55 1.0"/>
-
-
-<!-- 顶部平台向 +x 方向连接地面的10cm台阶，同样y轴移动到 -3.50 -->
-
-<geom pos="2.45 -3.50 0.25" type="box" size="0.15 0.5 0.05" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
-<geom pos="2.75 -3.50 0.15" type="box" size="0.15 0.5 0.05" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
-<geom pos="3.05 -3.50 0.05" type="box" size="0.15 0.5 0.05" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
-
-
-
-
-<!-- 斜坡木桥A木桥B -->
-<geom pos="1.8 -0.88 0.1" type="box" size="0.40 0.5 0.005" quat="0.701836 0.086175 -0.086175 0.701836" rgba="0.75 0.72 0.55 1.0"/>
-<geom pos="1.8 0.0 0.10" type="box" size="0.5 0.5 0.10" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
-<geom pos="2.65 0.0 0.10" type="box" size="0.2 0.5 0.10" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
-<geom pos="3.2 0.0 0.10" type="box" size="0.2 0.5 0.10" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
-<geom pos="3.75 0.0 0.10" type="box" size="0.2 0.5 0.10" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
-<geom pos="4.3 0.0 0.10" type="box" size="0.2 0.5 0.10" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
-<geom pos="4.85 0.0 0.10" type="box" size="0.2 0.5 0.10" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
-<geom pos="5.7 -0.5 0.10" type="box" size="0.5 1.0 0.10" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
-<geom pos="4.8002 -1.0 0.0951"
-      type="box"
-      size="0.4133 0.5 0.005"
-      quat="0.992546 0.0 -0.121869 0.0" 
-      rgba="0.75 0.72 0.55 1.0"/>
-<geom pos="5.35 -2.25 0.10" type="box" size="0.1 0.75 0.10" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
-<geom pos="5.65 -2.25 0.10" type="box" size="0.1 0.75 0.10" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
-<geom pos="5.95 -2.25 0.10" type="box" size="0.1 0.75 0.10" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
-<geom pos="5.7 -3.5 0.10" type="box" size="0.5 0.5 0.10" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
-<!-- 10度斜坡：宽4m，斜坡 y正方向边缘 与平台 y正方向边缘 对齐 -->
-<geom pos="4.6338 -5.0 0.0951"
-      type="box"
-      size="0.5759 2.0 0.005"
-      quat="0.9961947 0.0 -0.0871557 0.0" 
-      rgba="0.75 0.72 0.55 1.0"/>
-<!-- 新建10度斜坡：宽3m，高端与前一个10度斜坡高端衔接，向+x方向下坡 -->
-<geom pos="5.7681 -5.5 0.0951"
-      type="box"
-      size="0.5759 1.5 0.005"
-      quat="0.9961947 0.0 0.0871557 0.0" 
-      rgba="0.75 0.72 0.55 1.0"/>
-
-
-
-
-
-<!--绕杆-->
-<!-- 直径1m圆形颜色块，仅显示，不碰撞 -->
-<geom pos="1.8 -10.1 0.0"
-      type="cylinder"
-      size="0.1 0.001"
-      rgba="1.0 0.0 0.0 0.35"
-      contype="0"
-      conaffinity="0" />
-
-<geom pos="3.2 -12.5 0.0"
-      type="cylinder"
-      size="0.1 0.001"
-      rgba="1.0 0.0 0.0 0.35"
-      contype="0"
-      conaffinity="0" />
-
-<geom pos="1.55 -12.75 0.0"
-      type="cylinder"
-      size="0.1 0.001"
-      rgba="1.0 0.0 0.0 0.35"
-      contype="0"
-      conaffinity="0" />
-
-<!-- 原杆 -->
-<geom pos="1.8 -10.5 0.02"
-      type="cylinder"
-      size="0.05 0.02"
-      quat="1.0 0.0 0.0 0.0"
-      rgba="0.75 0.72 0.55 1.0" />
-
-<geom pos="1.8 -10.5 0.37"
-      type="cylinder"
-      size="0.015 0.33"
-      quat="1.0 0.0 0.0 0.0"
-      rgba="0.75 0.72 0.55 1.0" />
-
-
-<!-- y轴负方向第1根：间隔1m -->
-<geom pos="1.8 -11.5 0.02"
-      type="cylinder"
-      size="0.05 0.02"
-      quat="1.0 0.0 0.0 0.0"
-      rgba="0.75 0.72 0.55 1.0" />
-
-<geom pos="1.8 -11.5 0.37"
-      type="cylinder"
-      size="0.015 0.33"
-      quat="1.0 0.0 0.0 0.0"
-      rgba="0.75 0.72 0.55 1.0" />
-
-
-<!-- y轴负方向第2根：继续间隔1m -->
-<geom pos="1.8 -12.5 0.02"
-      type="cylinder"
-      size="0.05 0.02"
-      quat="1.0 0.0 0.0 0.0"
-      rgba="0.75 0.72 0.55 1.0" />
-
-<geom pos="1.8 -12.5 0.37"
-      type="cylinder"
-      size="0.015 0.33"
-      quat="1.0 0.0 0.0 0.0"
-      rgba="0.75 0.72 0.55 1.0" />
-
-
-<!-- x轴正方向第3根：继续间隔1m -->
-<geom pos="2.8 -12.5 0.02"
-      type="cylinder"
-      size="0.05 0.02"
-      quat="1.0 0.0 0.0 0.0"
-      rgba="0.75 0.72 0.55 1.0" />
-
-<geom pos="2.8 -12.5 0.37"
-      type="cylinder"
-      size="0.015 0.33"
-      quat="1.0 0.0 0.0 0.0"
-      rgba="0.75 0.72 0.55 1.0" />
-
-
-<!--===================================================================其他障碍=====================================================================================-->
-
-<!-- 5cm台阶 -->
-<geom pos="1.0 2.0 0.025" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
-<geom pos="1.3 2.0 0.075" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
-<geom pos="1.6 2.0 0.125" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
-<geom pos="1.9 2.0 0.175" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
-<geom pos="2.2 2.0 0.225" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
-<geom pos="2.5 2.0 0.275" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
-<geom pos="2.8 2.0 0.325" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
-<geom pos="3.45 2.0 0.375" type="box" size="0.5 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
-<geom pos="4.1 2.0 0.325" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
-<geom pos="4.4 2.0 0.275" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
-<geom pos="4.7 2.0 0.225" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
-<geom pos="5.0 2.0 0.175" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
-<geom pos="5.3 2.0 0.125" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
-<geom pos="5.6 2.0 0.075" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
-<geom pos="5.9 2.0 0.025" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
-
-<!-- 斜坡 -->
-<geom pos="2.0 4.0 0.1" type="box" size="1.5 0.75 0.005" quat="0.9950041652780258 0.0 -0.09983341664682815 0.0" />
-
-
-
-
-<geom pos="1.4 6.0 0.165" type="box" size="0.1 0.75 0.0049999999999999975" quat="1.0 0.0 0.0 0.0"/>
-<geom pos="1.6 6.0 0.275" type="box" size="0.1 0.75 0.0049999999999999975" quat="1.0 0.0 0.0 0.0"/>
-<geom pos="1.8 6.0 0.385" type="box" size="0.1 0.75 0.0049999999999999975" quat="1.0 0.0 0.0 0.0"/>
-<geom pos="2.0 6.0 0.495" type="box" size="0.1 0.75 0.0049999999999999975" quat="1.0 0.0 0.0 0.0"/>
-<geom pos="2.2 6.0 0.605" type="box" size="0.1 0.75 0.0049999999999999975" quat="1.0 0.0 0.0 0.0"/>
-<geom pos="2.4 6.0 0.715" type="box" size="0.1 0.75 0.0049999999999999975" quat="1.0 0.0 0.0 0.0"/>
-<geom pos="2.5999999999999996 6.0 0.825" type="box" size="0.1 0.75 0.0049999999999999975" quat="1.0 0.0 0.0 0.0"/>
-<geom pos="2.8 6.0 0.9349999999999999" type="box" size="0.1 0.75 0.0049999999999999975" quat="1.0 0.0 0.0 0.0"/>
-<geom pos="3.0 6.0 1.045" type="box" size="0.1 0.75 0.0049999999999999975" quat="1.0 0.0 0.0 0.0"/>
-<geom pos="-2.3179973398407565 5.173660321080885 -0.25" type="box" size="0.2568216778785459 0.2608020098770089 0.2619541072037832" quat="0.9930658271270357 -0.05910360995856133 -0.06731065544310916 0.0761334482746007"/>
-<geom pos="-2.3179973398407565 5.3620612607436735 -0.25" type="box" size="0.23778028947050467 0.26051225137569556 0.27137457425982286" quat="0.9959802827127009 0.07435995796871331 0.0004862731479656538 -0.04993632582444846"/>
-<geom pos="-2.3179973398407565 5.545897059602109 -0.25" type="box" size="0.23841994611378936 0.2717839884518381 0.22585827399286504" quat="0.9983504429181294 -0.004811901627841528 0.03471877670768104 -0.045473566737405234"/>
-<geom pos="-2.3179973398407565 5.7436240471772795 -0.25" type="box" size="0.2552179019048769 0.2548992578792955 0.22547735976326444" quat="0.9968270877924189 -0.029673908987198697 0.06777847718526858 -0.029347814214192768"/>
-<geom pos="-2.3179973398407565 5.940214647011584 -0.25" type="box" size="0.24313116620329878 0.2372064979204117 0.26079933745117434" quat="0.9952106364954509 0.05088987714749159 -0.07605843245051555 -0.03436748846516202"/>
-<geom pos="-2.3179973398407565 6.165430585901471 -0.25" type="box" size="0.24786042386990592 0.2322559052231109 0.2644037606269708" quat="0.9936075351397807 -0.05017393314139304 0.06986162224674641 -0.07311632022760194"/>
-<geom pos="-2.3179973398407565 6.315657865031069 -0.25" type="box" size="0.23704265198840277 0.24982080672772003 0.2530694373586838" quat="0.9981716459301547 0.036179123385437884 0.04523497974247521 0.017269420947002005"/>
-<geom pos="-2.3179973398407565 6.489372835072359 -0.25" type="box" size="0.2647428927965494 0.2716292502682415 0.23725049444938928" quat="0.9954041516289313 0.019099981466976248 0.07663366510923347 0.05415761257454444"/>
-<geom pos="-2.094119617957536 5.194943631570213 -0.25" type="box" size="0.23176840693038148 0.23782936054799508 0.2282032657053922" quat="0.9973366591969123 0.030692664684553082 -0.030015392527106853 0.058962910104125923"/>
-<geom pos="-2.094119617957536 5.441090326561234 -0.25" type="box" size="0.2602142310926322 0.27213502289176367 0.2574009440402366" quat="0.9948915480473169 0.023650461407535205 0.08508706956405496 -0.04890453856469255"/>
-<geom pos="-2.094119617957536 5.642958951230403 -0.25" type="box" size="0.24800055056479955 0.24050676557282252 0.23522489807277194" quat="0.9912055235117067 0.06188914817453858 -0.09720856678264424 -0.0650525790586179"/>
-<geom pos="-2.094119617957536 5.884810142659838 -0.25" type="box" size="0.24637516954806898 0.24364583504893206 0.2682443460752295" quat="0.9964495083946477 -0.07467561549070643 -0.038754439691442995 0.003165924091257204"/>
-<geom pos="-2.094119617957536 6.129893093768145 -0.25" type="box" size="0.2378958269703999 0.25045408022075055 0.24760364656411665" quat="0.9946240462872651 -0.03963038800955788 -0.08917648605658504 0.03464091840526007"/>
-<geom pos="-2.094119617957536 6.31135792935651 -0.25" type="box" size="0.2612154064649683 0.2252849660353503 0.26933214617336004" quat="0.997742809172301 0.019590451180868298 0.06391648942914147 -0.006339033565912353"/>
-<geom pos="-2.094119617957536 6.4905285510291915 -0.25" type="box" size="0.22768983323309258 0.23025468157122184 0.2748062344121069" quat="0.9940843963100783 0.05105793020576152 0.009164090402224958 -0.0954218016127878"/>
-<geom pos="-2.094119617957536 6.650993287072798 -0.25" type="box" size="0.25443787917913946 0.24575387105878618 0.22934424641420587" quat="0.9953832160779188 -0.0075142647566912866 0.08305501342257929 0.04751477371219154"/>
-<geom pos="-1.8951777534978893 5.178742263102535 -0.25" type="box" size="0.256127371544264 0.24353861032614957 0.2714578153598507" quat="0.996547165679354 -0.008558846106514282 0.009384988866828401 0.08205129318749771"/>
-<geom pos="-1.8951777534978893 5.353646654030374 -0.25" type="box" size="0.23858897994497189 0.2426691342857623 0.26961659613969635" quat="0.997689480058865 0.014145078555906384 0.009115651940537827 -0.06582190381793727"/>
-<geom pos="-1.8951777534978893 5.575640096199311 -0.25" type="box" size="0.24430646334884096 0.2676093509240798 0.23693271670520474" quat="0.9976019172675967 -0.06794745094616256 0.006187887696065785 -0.011630503849579829"/>
-<geom pos="-1.8951777534978893 5.728196634759373 -0.25" type="box" size="0.25668697148716296 0.2743986770827004 0.23696403861156407" quat="0.9946311966964959 -0.0774649076913635 0.05866511796477466 0.03558615698054331"/>
-<geom pos="-1.8951777534978893 5.954288148947323 -0.25" type="box" size="0.2748179458028994 0.25407956175122554 0.25142548243710355" quat="0.9927720663552536 0.03885622734718657 0.06180632563910765 0.09525647469886914"/>
-<geom pos="-1.8951777534978893 6.198076908897641 -0.25" type="box" size="0.23522918501552104 0.2714895927340788 0.23659922178360912" quat="0.9944408196408429 0.027722532803612636 -0.06883505448003607 -0.07470376618171105"/>
-<geom pos="-1.8951777534978893 6.356335079330304 -0.25" type="box" size="0.253243646653926 0.26648639763264487 0.22751627090926196" quat="0.9924345646356405 0.06351743627133029 0.09661714031088077 -0.041283149154730255"/>
-<geom pos="-1.8951777534978893 6.6037178705715744 -0.25" type="box" size="0.2664884535468762 0.26472237049442093 0.2545826559482188" quat="0.996208812433608 0.04792667434016892 0.03088022810200046 -0.06570728596343135"/>
-<geom pos="-1.7450143557797366 5.181798049538029 -0.25" type="box" size="0.22764974830145798 0.2314500042225232 0.26635647118774647" quat="0.996420562587967 -0.02530598413067966 -0.01712092616039989 -0.07881968983996092"/>
-<geom pos="-1.7450143557797366 5.396539066742657 -0.25" type="box" size="0.24515338305934362 0.25502436912192245 0.23509532716059323" quat="0.9974817825235419 0.033698559354477776 -0.06057332552735373 -0.01501242370995589"/>
-<geom pos="-1.7450143557797366 5.5477605493550115 -0.25" type="box" size="0.2368415768982088 0.2653984778068547 0.25193186806340717" quat="0.9933816957092091 -0.04426535864216467 0.0892264769079807 0.057201577537394625"/>
-<geom pos="-1.7450143557797366 5.764238738853998 -0.25" type="box" size="0.257475541143157 0.25587521442146555 0.2684267956125346" quat="0.9955849335110497 0.004076118120639651 -0.09299366924131539 0.012091439447074191"/>
-<geom pos="-1.7450143557797366 5.942956213887727 -0.25" type="box" size="0.23647345784635188 0.22779489919605103 0.2690566454457882" quat="0.9997372307105926 0.020286738832439623 -0.010113822112874779 0.003410038259163467"/>
-<geom pos="-1.7450143557797366 6.162981335139796 -0.25" type="box" size="0.2336152227884161 0.23785626414299832 0.26272786991330355" quat="0.9954569549676405 0.08638675499005252 -0.03150236191317218 -0.024705881136540285"/>
-<geom pos="-1.7450143557797366 6.344025407207907 -0.25" type="box" size="0.2364006101773331 0.23674709170116234 0.2660167000427503" quat="0.9941020005048972 -0.09350614404282417 -0.05493059192638792 0.0006660998579442658"/>
-<geom pos="-1.7450143557797366 6.5791872438688825 -0.25" type="box" size="0.259653062502181 0.26359758888480966 0.27170867851854713" quat="0.9955751902075182 0.06572529509454274 0.04059253213564211 -0.05350207998588654"/>
-<geom pos="-1.4950537649717406 5.196975242498044 -0.25" type="box" size="0.24902209995429173 0.24604186796843594 0.26555264385759036" quat="0.9981753741369879 0.027059282981497578 0.02078151255996191 0.049818133312398136"/>
-<geom pos="-1.4950537649717406 5.415093649392617 -0.25" type="box" size="0.22561317519118573 0.23246623591498758 0.2516053992906602" quat="0.9961431033663982 -0.019679813255757937 0.07777937635153075 0.035524515199325105"/>
-<geom pos="-1.4950537649717406 5.596072881787104 -0.25" type="box" size="0.2545347271371952 0.2527292932516396 0.272364707011277" quat="0.9987990019474599 -0.0468132050914785 0.002764373695484487 -0.01419280718860587"/>
-<geom pos="-1.4950537649717406 5.788444207457658 -0.25" type="box" size="0.257823578756761 0.22815201323013437 0.2506904868770564" quat="0.9907800223935689 0.08207700037219429 -0.06679273748567588 -0.08459931119619979"/>
-<geom pos="-1.4950537649717406 5.962479724512656 -0.25" type="box" size="0.23447002037921025 0.260091883647859 0.2613547781123637" quat="0.994185747243896 -0.045779429335849345 -0.09537527521663794 -0.020062420196245392"/>
-<geom pos="-1.4950537649717406 6.137620949716146 -0.25" type="box" size="0.25502378303486756 0.24137626830945555 0.26521755821448284" quat="0.9954511067004865 0.04827241542514884 0.023378972724717166 -0.07874193109224191"/>
-<geom pos="-1.4950537649717406 6.345742058941403 -0.25" type="box" size="0.23763376291279575 0.27259418014745557 0.24184880568666417" quat="0.9979730938550316 -0.0521706328108408 -0.03256129013590577 0.01636127740178204"/>
-<geom pos="-1.4950537649717406 6.548966807280673 -0.25" type="box" size="0.23004532593424165 0.24736965888987583 0.22917624237245732" quat="0.993938322382385 -0.041286377558295506 -0.09233464220936707 0.04308549844056845"/>
-<geom pos="-1.2872521554407157 5.194072124237239 -0.25" type="box" size="0.2725857334366114 0.23611730648841156 0.25109418723334265" quat="0.9870880287314421 -0.09538067591152324 -0.0888119213223197 -0.09312460914696315"/>
-<geom pos="-1.2872521554407157 5.418639818976418 -0.25" type="box" size="0.2326397456179607 0.2609646699674687 0.2717115772948157" quat="0.9928206111485406 -0.07383059897166898 -0.08473752136579614 0.040936892980589765"/>
-<geom pos="-1.2872521554407157 5.655974569843163 -0.25" type="box" size="0.22633109461398734 0.25911291311168594 0.23532484499883452" quat="0.9947769304653071 -0.023834507399340135 0.09176246279588879 -0.037820963666801724"/>
-<geom pos="-1.2872521554407157 5.8981628648303595 -0.25" type="box" size="0.27278934520816167 0.2559269445001904 0.26076472835929454" quat="0.9936318869346414 0.0561019165840355 0.09234219275279033 0.0319557140415977"/>
-<geom pos="-1.2872521554407157 6.064939027275534 -0.25" type="box" size="0.2481756437845515 0.2613413397088905 0.24788858471207542" quat="0.9992310163557966 0.006919915638721734 0.03652449975321466 -0.012468024618680441"/>
-<geom pos="-1.2872521554407157 6.288670079370122 -0.25" type="box" size="0.2574127585385611 0.27445220033632356 0.22507618952620437" quat="0.9984337612791753 0.04947883993440395 -0.0010074967830211658 -0.026093173185657285"/>
-<geom pos="-1.2872521554407157 6.4987760234638605 -0.25" type="box" size="0.2442462188069663 0.2639082925274208 0.24918893213917415" quat="0.9910480688126146 0.08639971284060838 -0.0903343762232188 -0.04688832899783648"/>
-<geom pos="-1.2872521554407157 6.675611985491267 -0.25" type="box" size="0.2507168814441121 0.26699708557208374 0.26588306060638556" quat="0.9968882348111167 0.009466759168649483 -0.06920917552433178 0.03652831489763918"/>
-<geom pos="-1.0678149575697586 5.238251070535694 -0.25" type="box" size="0.2331754586513582 0.22873009754409884 0.2593258638743009" quat="0.9962563091989393 0.014506022500180732 0.069646748226122 0.049114887295595266"/>
-<geom pos="-1.0678149575697586 5.471732496077581 -0.25" type="box" size="0.2653495693182789 0.26581370557074685 0.2509273010188512" quat="0.9948019811445705 -0.05898118120008214 0.04034485064897968 -0.07254330845220838"/>
-<geom pos="-1.0678149575697586 5.691153717662407 -0.25" type="box" size="0.26309077142424103 0.26536949948987093 0.26566703066149744" quat="0.998331608401392 -0.025703385975953945 -0.05139460706256987 0.005650661991725779"/>
-<geom pos="-1.0678149575697586 5.938910564562482 -0.25" type="box" size="0.23082606038737274 0.23890770539441533 0.25941695887199245" quat="0.99040745217393 -0.09494180013051996 0.05600240658814475 -0.08332384846283197"/>
-<geom pos="-1.0678149575697586 6.13746533389903 -0.25" type="box" size="0.2676431215498913 0.2569308659994288 0.24597694927356487" quat="0.9955894165044442 0.04511524606234168 -0.003741035154326279 -0.08217258042101841"/>
-<geom pos="-1.0678149575697586 6.328213254628996 -0.25" type="box" size="0.24188540919954635 0.25556145122229207 0.2605619987765001" quat="0.9929731934387934 -0.06598945456828725 -0.06809928840246085 0.07079629875087447"/>
-<geom pos="-1.0678149575697586 6.554891603524267 -0.25" type="box" size="0.23121916641107804 0.25266417867731916 0.25489063309084503" quat="0.9984565221200774 -0.02714563973423854 0.029423903667849492 0.03849573446816443"/>
-<geom pos="-1.0678149575697586 6.718482966003368 -0.25" type="box" size="0.27418971733780156 0.2623437838864593 0.23694037285314332" quat="0.9968347573500278 0.015128289848683302 0.06990516881209438 0.03471121949050888"/>
-<geom pos="-0.8851136474992356 5.233272057743225 -0.25" type="box" size="0.2587161160845902 0.2542459313914242 0.25268742624288776" quat="0.9956323623605605 0.07455213252171443 -0.05585558926184348 0.006191260373025539"/>
-<geom pos="-0.8851136474992356 5.410759300563648 -0.25" type="box" size="0.24019625793101548 0.2509280955260936 0.26698317271101046" quat="0.9960056830073374 0.025487333399660517 0.08516393482003458 -0.008377318140719903"/>
-<geom pos="-0.8851136474992356 5.622546965631826 -0.25" type="box" size="0.22527085609750916 0.22924847380626232 0.23073331588883172" quat="0.9982201065088221 0.016292210993739946 0.013620740418279952 0.05572843307423328"/>
-<geom pos="-0.8851136474992356 5.856118124892678 -0.25" type="box" size="0.26833255247166926 0.2512767990265972 0.2502231376336179" quat="0.9914910532543448 0.0657880493114159 0.05073139628229043 -0.10021850784978792"/>
-<geom pos="-0.8851136474992356 6.010914595891683 -0.25" type="box" size="0.24926441451407527 0.22868800964152894 0.26501630221174116" quat="0.990736605756697 -0.0765632176593298 -0.09876375004658954 0.05314859727296924"/>
-<geom pos="-0.8851136474992356 6.1657945570092965 -0.25" type="box" size="0.26823427325895977 0.263134568634285 0.23692064318485426" quat="0.9903439387899988 0.07219524819688558 0.09310840228148876 0.07305856872598579"/>
-<geom pos="-0.8851136474992356 6.383066094410735 -0.25" type="box" size="0.23772929143507357 0.2619329708548053 0.23258884134606547" quat="0.9928463525385446 0.07348810370848206 0.09252710448645944 -0.017156742103069993"/>
-<geom pos="-0.8851136474992356 6.551482176174038 -0.25" type="box" size="0.25127741834221506 0.25307864976328337 0.234931895271364" quat="0.994062479132312 -0.084918495271965 -0.05565932062022146 -0.03912386445844099"/>
-<geom pos="-0.6441452943469552 5.193190826541208 -0.25" type="box" size="0.27016649011158844 0.23778885128164629 0.25859862297032477" quat="0.9957774962333289 0.08893738517887846 -0.021654911548168992 0.006955883744530942"/>
-<geom pos="-0.6441452943469552 5.407266358883468 -0.25" type="box" size="0.24982888074520473 0.26972381569065607 0.2275629646713632" quat="0.9942571418988088 0.02444442288698024 0.050336521935401994 -0.09122193010664256"/>
-<geom pos="-0.6441452943469552 5.56363948710079 -0.25" type="box" size="0.24725848381417467 0.2326432801330426 0.2476341019968084" quat="0.9935769557551848 -0.05688836927692267 -0.06858315894908604 0.0697488117593134"/>
-<geom pos="-0.6441452943469552 5.716461923308758 -0.25" type="box" size="0.226746163574222 0.25188961955216527 0.24650452053954758" quat="0.9981029078160676 0.03471313574303335 -0.0493522060742233 -0.012245136651074443"/>
-<geom pos="-0.6441452943469552 5.896430073001983 -0.25" type="box" size="0.25284949298017617 0.23421620066432108 0.2621382648463894" quat="0.9996413412447853 0.009523771195346772 -0.02374945051322109 -0.007902547492127014"/>
-<geom pos="-0.6441452943469552 6.110845613879558 -0.25" type="box" size="0.23421403581041014 0.2504556848552827 0.24096555776925194" quat="0.9982371379292613 0.05906199571247059 0.004184466367765714 0.004097238396071534"/>
-<geom pos="-0.6441452943469552 6.301410327752269 -0.25" type="box" size="0.2296817145699627 0.2701766916372966 0.22803599234835198" quat="0.9993879202059675 0.00903361239651759 0.0071600522792900426 -0.03302896372607247"/>
-<geom pos="-0.6441452943469552 6.470340592771366 -0.25" type="box" size="0.2742166750011587 0.23990439594655139 0.2609404931878803" quat="0.9954980192986956 -0.06477558432537366 -0.03478863652617152 0.059812774691783824"/>
-<geom pos="-0.43033316974190594 5.176946043237525 -0.25" type="box" size="0.2511344405622137 0.26383180387822514 0.2729516287367074" quat="0.9939152287364423 -0.054017256979398534 -0.08418383902399898 -0.0461273810376057"/>
-<geom pos="-0.43033316974190594 5.409693792302493 -0.25" type="box" size="0.23878438774624536 0.22858057493504688 0.24808981791323623" quat="0.9942721168135948 -0.06322741578652484 -0.032958930763236916 -0.0796175891553822"/>
-<geom pos="-0.43033316974190594 5.573593446642387 -0.25" type="box" size="0.2647545698926724 0.25215736713350106 0.25522885339490087" quat="0.9971918839448468 -0.004732654353889047 0.07222000940536273 0.019241071144380294"/>
-<geom pos="-0.43033316974190594 5.819531161619798 -0.25" type="box" size="0.2538318010250847 0.23661725487197 0.26323696729639623" quat="0.9929867679849618 0.07588268800385974 -0.046355180482866284 -0.07791208834634974"/>
-<geom pos="-0.43033316974190594 6.047888701789067 -0.25" type="box" size="0.25684720236482583 0.25568474221031684 0.24397094296107352" quat="0.9925256546500537 -0.08214715587137877 0.06615741576234473 0.06138294537877935"/>
-<geom pos="-0.43033316974190594 6.283676056613722 -0.25" type="box" size="0.2432692787269645 0.2742601402399693 0.2689467974776609" quat="0.9989897410981501 0.0047037010179340685 0.0014590493232283363 -0.04466814919444852"/>
-<geom pos="-0.43033316974190594 6.458239009480634 -0.25" type="box" size="0.2474330649412728 0.2551435998627988 0.23002773988805483" quat="0.9982311862818315 0.0506459597703481 0.025989046975463247 -0.01714802993392417"/>
-<geom pos="-0.43033316974190594 6.703309377175586 -0.25" type="box" size="0.2301029552870502 0.2556223065028505 0.2349557527965884" quat="0.9948672058648923 -0.034182171957127715 -0.0030330068097560517 -0.09519255582537542"/>
-<geom type="hfield" hfield="perlin_hfield" pos="-1.5 4.0 0.0" quat="1.0 0.0 0.0 0.0"/>
-<geom type="hfield" hfield="image_hfield" pos="-1.5 2.0 0.0" quat="0.7073882691671998 0.0 0.0 -0.706825181105366"/>
-</worldbody>
-</mujoco>
@@ -26,28 +26,28 @@

  <worldbody>
    <body name="base_link">
-      <inertial pos="0.1517 0.0002 0.0542" mass="3.5" diaginertia="0.0215 0.0904 0.0985"/>
+      <inertial pos="0.1517 0.0002 0.0542" mass="6.5377" fullinertia="0.0402 0.1689 0.1840 -0.0001 0.0000 0.0000"/>
      <joint type="free"/>
      <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="base_link"/>
      <geom size="0.178 0.1175 0.073" pos="0.1518 0 0.054" type="box" rgba="0.75294 0.75294 0.75294 1"/>
      <body name="fl_hip_abduction_Link" pos="0.32826 0.066172 0.053981">
-        <inertial pos="0.0488 -0.0026 0.0007" mass="0.5" diaginertia="0.0003 0.0006 0.0005"/>
+        <inertial pos="0.0488 -0.0026 0.0007" mass="0.63" fullinertia="0.0004 0.0007 0.0006 0.0000 0.0000 0.0000"/>
        <joint name="fl_hip_abduction_joint" pos="0 0 0" axis="1 0 0" range="-0.436 0.611" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
        <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="fl_hip_abduction_Link"/>
        <body name="fl_hip_pitch_Link" pos="0.06389 -0.027344 0.00010727" quat="0.999997 -0.0025023 0 0">
-          <inertial pos="0.0019 0.1119 -0.048" mass="0.935" diaginertia="0.0062 0.0064 0.001"/>
+          <inertial pos="0.0019 0.1119 -0.0480" mass="0.998" fullinertia="0.0066 0.0068 0.0011 0.0000 0.0000 0.0004"/>
          <joint name="fl_hip_pitch_joint" pos="0 0 0" axis="0 1 0" range="-2.58 2.58" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
          <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="fl_hip_pitch_Link"/>
          <geom size="0.046 0.048" pos="0 0.048 0" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
          <geom size="0.0435 0.0115 0.06" pos="0 0.1155 -0.06" type="box" rgba="0.75294 0.75294 0.75294 1"/>
          <body name="fl_knee_Link" pos="0 0.1035 -0.25" quat="0.999997 0.0025023 0 0">
-            <inertial pos="0.0002 0.0242 -0.1539" mass="0.651" fullinertia="0.0042 0.0045 0.0005 0 0 0.0002"/>
+            <inertial pos="0.0002 0.0242 -0.1539" mass="0.6965" fullinertia="0.0045 0.0048 0.0006 0.0000 0.0000 0.0002"/>
            <joint name="fl_knee_joint" pos="0 0 0" axis="0 1 0" range="-2.65 2.65" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
            <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="fl_knee_Link"/>
            <geom size="0.0475 0.015" pos="0 0.025 -0.20011" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
            <geom size="0.015 0.0125 0.06" pos="0 0.0125 -0.09" type="box" rgba="0.75294 0.75294 0.75294 1"/>
            <body name="fl_wheel_Link" pos="0 0.014699 -0.20011">
-              <inertial pos="-0.0002 0.0407 -0.0001" mass="0.53" diaginertia="0.0017 0.0032 0.0017"/>
+              <inertial pos="-0.0002 0.0407 -0.0001" mass="0.5505" fullinertia="0.0017 0.0034 0.0017 0.0000 0.0000 0.0000"/>
              <joint name="fl_wheel_joint" pos="0 0 0" axis="0 1 0" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
              <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="fl_wheel_Link"/>
              <geom size="0.1 0.015" pos="0 0.04074 0" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
@@ -56,23 +56,23 @@
        </body>
      </body>
      <body name="fr_hip_abduction_Link" pos="0.32826 -0.065853 0.054034">
-        <inertial pos="0.0488 0.0026 0.0008" mass="0.5" diaginertia="0.0003 0.0006 0.0005"/>
+        <inertial pos="0.0488 0.0026 0.0008" mass="0.63" fullinertia="0.0004 0.0007 0.0006 0.0000 0.0000 0.0000"/>
        <joint name="fr_hip_abduction_joint" pos="0 0 0" axis="1 0 0" range="-0.611 0.436" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
        <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="fr_hip_abduction_Link"/>
        <body name="fr_hip_pitch_Link" pos="0.06389 0.027311 -0.00036027" quat="0.999976 -0.00686995 0 0">
-          <inertial pos="-0.0019 -0.1119 -0.048" mass="0.935" diaginertia="0.0062 0.0064 0.001"/>
+          <inertial pos="-0.0019 -0.1119 -0.048" mass="0.998" fullinertia="0.0066 0.0068 0.0011 0.0000 0.0000 0.0004"/>
          <joint name="fr_hip_pitch_joint" pos="0 0 0" axis="0 1 0" range="-2.58 2.58" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
          <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="fr_hip_pitch_Link"/>
          <geom size="0.046 0.048" pos="0 -0.048 0" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
          <geom size="0.0435 0.0115 0.06" pos="0 -0.1155 -0.06" type="box" rgba="0.75294 0.75294 0.75294 1"/>
          <body name="fr_knee_Link" pos="-0.00075079 -0.1035 -0.25" quat="0.999976 0.00686995 0 0">
-            <inertial pos="-0.0002 -0.0242 -0.1539" mass="0.651" fullinertia="0.0042 0.0045 0.0005 0 0 0.0001"/>
+            <inertial pos="-0.0002 -0.0242 -0.1539" mass="0.6965" fullinertia="0.0045 0.0048 0.0006 0.0000 0.0000 0.0002"/>
            <joint name="fr_knee_joint" pos="0 0 0" axis="0 1 0" range="-2.65 2.65" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
            <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="fr_knee_Link"/>
            <geom size="0.0475 0.015" pos="0 -0.025 -0.1998" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
            <geom size="0.015 0.0125 0.06" pos="0 -0.0125 -0.09" type="box" rgba="0.75294 0.75294 0.75294 1"/>
            <body name="fr_wheel_Link" pos="0 -0.018447 -0.1998">
-              <inertial pos="0.0002 -0.0407 -0.0001" mass="0.53" diaginertia="0.0017 0.0032 0.0017"/>
+              <inertial pos="0.0002 -0.0407 -0.0001" mass="0.5505" fullinertia="0.0017 0.0034 0.0017 0.0000 0.0000 0.0000"/>
              <joint name="fr_wheel_joint" pos="0 0 0" axis="0 1 0" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
              <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="fr_wheel_Link"/>
              <geom size="0.1 0.015" pos="0 -0.040735 0" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
@@ -81,23 +81,23 @@
        </body>
      </body>
      <body name="rl_hip_abduction_Link" pos="-0.024743 0.066141 0.054034">
-        <inertial pos="-0.0488 -0.0026 -0.0008" mass="0.5" diaginertia="0.0003 0.0006 0.0005"/>
+        <inertial pos="-0.0488 -0.0026 -0.0007" mass="0.63" fullinertia="0.0004 0.0007 0.0006 0.0000 0.0000 0.0000"/>
        <joint name="rl_hip_abduction_joint" pos="0 0 0" axis="1 0 0" range="-0.436 0.611" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
        <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="rl_hip_abduction_Link"/>
        <body name="rl_hip_pitch_Link" pos="-0.06389 -0.027309 0.00045509">
-          <inertial pos="0.0019 0.1119 -0.048" mass="0.935" diaginertia="0.0062 0.0064 0.001"/>
+          <inertial pos="0.0019 0.1119 -0.048" mass="0.998" fullinertia="0.0066 0.0068 0.0011 0.0000 -0.0001 -0.0005"/>
          <joint name="rl_hip_pitch_joint" pos="0 0 0" axis="0 1 0" range="-2.58 2.58" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
          <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="rl_hip_pitch_Link"/>
          <geom size="0.046 0.048" pos="0 0.048 0" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
          <geom size="0.0435 0.0115 0.06" pos="0 0.1155 -0.06" type="box" rgba="0.75294 0.75294 0.75294 1"/>
          <body name="rl_knee_Link" pos="0 0.099459 -0.25163">
-            <inertial pos="0.0002 0.0242 -0.1539" mass="0.651" fullinertia="0.0042 0.0045 0.0005 0 0 -0.0003"/>
+            <inertial pos="0.0002 0.0242 -0.1539" mass="0.6965" fullinertia="0.0045 0.0048 0.0006 0.0000 0.0000 -0.0003"/>
            <joint name="rl_knee_joint" pos="0 0 0" axis="0 1 0" range="-2.65 2.65" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
            <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="rl_knee_Link"/>
            <geom size="0.0475 0.015" pos="0 0.025 -0.20027" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
            <geom size="0.015 0.0125 0.06" pos="0 0.0125 -0.09" type="box" rgba="0.75294 0.75294 0.75294 1"/>
            <body name="rl_wheel_Link" pos="0 0.012475 -0.20027">
-              <inertial pos="-0.0002 0.0407 -0.0001" mass="0.53" diaginertia="0.0017 0.0032 0.0017"/>
+              <inertial pos="-0.0002 0.0407 -0.0001" mass="0.5505" fullinertia="0.0017 0.0034 0.0017 0.0000 0.0000 0.0000"/>
              <joint name="rl_wheel_joint" pos="0 0 0" axis="0 1 0" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
              <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="rl_wheel_Link"/>
              <geom size="0.1 0.015" pos="0 0.040737 0" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
@@ -106,23 +106,23 @@
        </body>
      </body>
      <body name="rr_hip_abduction_Link" pos="-0.024743 -0.065884 0.053981">
-        <inertial pos="-0.0488 0.0026 0.0008" mass="0.5" diaginertia="0.0003 0.0006 0.0005"/>
+        <inertial pos="-0.0488 0.0026 0.0008" mass="0.63" fullinertia="0.0004 0.0007 0.0006 0.0000 0.0000 0.0000"/>
        <joint name="rr_hip_abduction_joint" pos="0 0 0" axis="1 0 0" range="-0.611 0.436" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
        <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="rr_hip_abduction_Link"/>
        <body name="rr_hip_pitch_Link" pos="-0.06389 0.027341 0.00041625">
-          <inertial pos="-0.002 -0.1111 -0.0498" mass="0.935" diaginertia="0.0062 0.0064 0.001"/>
+          <inertial pos="-0.0020 -0.1111 -0.0498" mass="0.998" fullinertia="0.0066 0.0068 0.0011 0.0000 -0.0001 -0.0003"/>
          <joint name="rr_hip_pitch_joint" pos="0 0 0" axis="0 1 0" range="-2.58 2.58" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
          <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="rr_hip_pitch_Link"/>
          <geom size="0.046 0.048" pos="0 -0.048 0" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
          <geom size="0.0435 0.0115 0.06" pos="0 -0.1155 -0.06" type="box" rgba="0.75294 0.75294 0.75294 1"/>
          <body name="rr_knee_Link" pos="-0.00075079 -0.099408 -0.25165">
-            <inertial pos="-0.0002 -0.0225 -0.1541" mass="0.651" fullinertia="0.0042 0.0045 0.0005 0 0 -0.0001"/>
+            <inertial pos="-0.0002 -0.0225 -0.1541" mass="0.6965" fullinertia="0.0045 0.0048 0.0006 0.0000 0.0000 -0.0002"/>
            <joint name="rr_knee_joint" pos="0 0 0" axis="0 1 0" range="-2.65 2.65" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
            <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="rr_knee_Link"/>
            <geom size="0.0475 0.015" pos="0 -0.025 -0.20027" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
            <geom size="0.015 0.0125 0.06" pos="0 -0.0125 -0.09" type="box" rgba="0.75294 0.75294 0.75294 1"/>
            <body name="rr_wheel_Link" pos="0 -0.012435 -0.20027">
-              <inertial pos="0.0002 -0.0407 -0.0005" mass="0.53" diaginertia="0.0017 0.0032 0.0017"/>
+              <inertial pos="0.0002 -0.0407 -0.0005" mass="0.5505" fullinertia="0.0017 0.0034 0.0017 0.0000 0.0000 0.0000"/>
              <joint name="rr_wheel_joint" pos="0 0 0" axis="0 1 0" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
              <geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="rr_wheel_Link"/>
              <geom size="0.1 0.015" pos="0 -0.040737 0" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
@@ -28,9 +28,9 @@ jobs:
    runs-on: ubuntu-latest
    steps:
      - name: Checkout repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
      - name: Setup uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
        with:
          enable-cache: true
          version: "0.9.27"
@@ -45,9 +45,9 @@ jobs:
      matrix:
        python-version: ["3.10", "3.11", "3.12", "3.13"]
    steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
      - name: Setup uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
        with:
          python-version: ${{ matrix.python-version }}
          enable-cache: true
@@ -68,9 +68,9 @@ jobs:
      matrix:
        python-version: ["3.10", "3.11", "3.12", "3.13"]
    steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
      - name: Setup uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
        with:
          python-version: ${{ matrix.python-version }}
          enable-cache: true
@@ -84,9 +84,9 @@ jobs:
      matrix:
        python-version: ["3.10", "3.11", "3.12", "3.13"]
    steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
      - name: Setup uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
        with:
          python-version: ${{ matrix.python-version }}
          enable-cache: true
@@ -0,0 +1,44 @@
+name: Claude Code Review
+
+on:
+  pull_request:
+    types: [opened, synchronize, ready_for_review, reopened]
+    # Optional: Only run on specific file changes
+    # paths:
+    #   - "src/**/*.ts"
+    #   - "src/**/*.tsx"
+    #   - "src/**/*.js"
+    #   - "src/**/*.jsx"
+
+jobs:
+  claude-review:
+    # Optional: Filter by PR author
+    # if: |
+    #   github.event.pull_request.user.login == 'external-contributor' ||
+    #   github.event.pull_request.user.login == 'new-developer' ||
+    #   github.event.pull_request.author_association == 'FIRST_TIME_CONTRIBUTOR'
+
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      pull-requests: read
+      issues: read
+      id-token: write
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 1
+
+      - name: Run Claude Code Review
+        id: claude-review
+        uses: anthropics/claude-code-action@v1
+        with:
+          claude_code_oauth_token: ${{ secrets.CLAUDE_CODE_OAUTH_TOKEN }}
+          plugin_marketplaces: 'https://github.com/anthropics/claude-code.git'
+          plugins: 'code-review@claude-code-plugins'
+          prompt: '/code-review:code-review ${{ github.repository }}/pull/${{ github.event.pull_request.number }}'
+          # See https://github.com/anthropics/claude-code-action/blob/main/docs/usage.md
+          # or https://code.claude.com/docs/en/cli-reference for available options
+
@@ -0,0 +1,50 @@
+name: Claude Code
+
+on:
+  issue_comment:
+    types: [created]
+  pull_request_review_comment:
+    types: [created]
+  issues:
+    types: [opened, assigned]
+  pull_request_review:
+    types: [submitted]
+
+jobs:
+  claude:
+    if: |
+      (github.event_name == 'issue_comment' && contains(github.event.comment.body, '@claude')) ||
+      (github.event_name == 'pull_request_review_comment' && contains(github.event.comment.body, '@claude')) ||
+      (github.event_name == 'pull_request_review' && contains(github.event.review.body, '@claude')) ||
+      (github.event_name == 'issues' && (contains(github.event.issue.body, '@claude') || contains(github.event.issue.title, '@claude')))
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      pull-requests: read
+      issues: read
+      id-token: write
+      actions: read # Required for Claude to read CI results on PRs
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 1
+
+      - name: Run Claude Code
+        id: claude
+        uses: anthropics/claude-code-action@v1
+        with:
+          claude_code_oauth_token: ${{ secrets.CLAUDE_CODE_OAUTH_TOKEN }}
+
+          # This is an optional setting that allows Claude to read CI results on PRs
+          additional_permissions: |
+            actions: read
+
+          # Optional: Give a custom prompt to Claude. If this is not specified, Claude will perform the instructions specified in the comment that tagged it.
+          # prompt: 'Update the pull request description to include a summary of changes.'
+
+          # Optional: Add claude_args to customize behavior and configuration
+          # See https://github.com/anthropics/claude-code-action/blob/main/docs/usage.md
+          # or https://code.claude.com/docs/en/cli-reference for available options
+          # claude_args: '--allowed-tools Bash(gh pr *)'
+
@@ -17,7 +17,7 @@ jobs:
  build:
    runs-on: ubuntu-latest
    steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
        with:
          fetch-depth: 0

@@ -26,7 +26,7 @@ jobs:
          python-version: '3.13'

      - name: Install uv
-        uses: astral-sh/setup-uv@v4
+        uses: astral-sh/setup-uv@v7

      - name: Build Sphinx Documentation
        run: uv run --group docs sphinx-multiversion docs docs/_build
@@ -15,9 +15,9 @@ jobs:
      contents: read
    steps:
      - name: Checkout
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
      - name: Install uv
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@v7
      - name: Install Python 3.13
        run: uv python install 3.13
      - name: Build
@@ -33,9 +33,9 @@ keywords:
  - reinforcement-learning
  - robotics
 license: Apache-2.0
-commit: e2f33c6fb49caa26ec11f7b2de3c0c9aba71e9fd
-version: 1.3.0
-date-released: '2026-04-14'
+commit: 3cc461cd15e7155a8998b75ad767fae6dd448072
+version: 1.4.0
+date-released: '2026-05-26'
 preferred-citation:
  type: article
  title: >-
@@ -1,6 +1,10 @@
 .PHONY: sync
 sync:
-	uv sync --all-extras --all-packages --group dev
+	uv sync --all-packages --extra cu128 --group dev
+
+.PHONY: sync-cpu
+sync-cpu:
+	uv sync --all-packages --extra cpu --group dev

 .PHONY: format
 format:
@@ -61,6 +61,7 @@ MuJoCo's integrator handles velocity-dependent forces.

 **Built-in actuators** (``BuiltinPositionActuator``,
 ``BuiltinVelocityActuator``, ``BuiltinMotorActuator``,
+``BuiltinPdActuator``, ``BuiltinDcMotorActuator``,
 ``BuiltinMuscleActuator``) create native MuJoCo actuator elements in the
 MjSpec. The physics engine computes the control law and integrates
 velocity-dependent damping forces implicitly. This provides the best
@@ -119,6 +120,31 @@ control.
 **BuiltinMotorActuator**: Creates ``<motor>`` actuators for direct torque
 control.

+**BuiltinPdActuator**: Native PD that closes on both a position and a
+velocity target, implemented as paired ``<position>`` + ``<velocity>``
+actuators summing to ``kp * (p_target - q) + kd * (v_target - qdot)``.
+``BuiltinPositionActuator`` puts kd on the ``<position>`` element and
+implicitly assumes a zero velocity reference; use this when the policy
+emits a non-zero velocity target. Native delivery lets
+``implicit`` / ``implicitfast`` see the kd term in their velocity update,
+unlike ``IdealPdActuator`` which forwards Python-computed torque through
+an opaque ``<motor>``.
+
+**BuiltinDcMotorActuator**: Wraps MuJoCo's native
+`<dcmotor> <https://mujoco.readthedocs.io/en/stable/XMLreference.html#actuator-dcmotor>`_
+element. Torque is ``tau = K * (V - K * omega) / R``; the back-EMF runs
+through the native bias path, so ``implicit`` / ``implicitfast`` pick up
+its velocity derivative as effective damping. Three input modes pick what
+``ctrl`` carries: VOLTAGE drives the motor directly; POSITION / VELOCITY
+close an internal PID (with anti-windup and slew limiting) against a
+single setpoint, whose Vmax-clamped output becomes torque. POSITION mode
+pins v_target = 0 (the kd term acts on raw velocity). Optional physics:
+inductance,
+thermal model with I^2R heating, cogging ripple, LuGre friction.
+``DcMotorActuator`` (the explicit version) is a software PD with a
+velocity-dependent torque clamp on top of a ``<motor>``; this is the real
+electrical model.
+
 **BuiltinMuscleActuator**: Creates ``<muscle>`` actuators for
 biologically-inspired muscle dynamics with force-length-velocity
 characteristics.
@@ -18,6 +18,13 @@ mjlab.actuator
  - :class:`BuiltinPositionActuatorCfg`
  - :class:`BuiltinVelocityActuator`
  - :class:`BuiltinVelocityActuatorCfg`
+  - :class:`BuiltinPdActuator`
+  - :class:`BuiltinPdActuatorCfg`
+  - :class:`BuiltinDcMotorActuator`
+  - :class:`BuiltinDcMotorActuatorCfg`
+  - :class:`DcMotorInputMode`
+  - :class:`DcMotorDatasheetParams`
+  - :class:`DcMotorPhysicalParams`
  - :class:`BuiltinMuscleActuator`
  - :class:`BuiltinMuscleActuatorCfg`
  - :class:`XmlActuator`
@@ -84,6 +91,40 @@ Builtin Actuators
  :undoc-members:


+.. autoclass:: BuiltinPdActuator
+  :members:
+  :show-inheritance:
+
+.. autoclass:: BuiltinPdActuatorCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
+.. autoclass:: BuiltinDcMotorActuator
+  :members:
+  :show-inheritance:
+
+.. autoclass:: BuiltinDcMotorActuatorCfg
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+.. autoclass:: DcMotorInputMode
+  :members:
+  :show-inheritance:
+
+.. autoclass:: DcMotorDatasheetParams
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+.. autoclass:: DcMotorPhysicalParams
+  :members:
+  :exclude-members: __init__
+  :undoc-members:
+
+
 .. autoclass:: BuiltinMuscleActuator
  :members:
  :show-inheritance:
@@ -8,6 +8,89 @@ Upcoming version (not yet released)
 Added
 ^^^^^

+- Added ``BuiltinDcMotorActuator``, a native MuJoCo ``<dcmotor>`` wrapper.
+  Supports voltage / position / velocity input modes with back-EMF,
+  configurable motor constants, and optional integral, slew, inductance,
+  thermal, LuGre, and cogging extensions.
+- Added ``scale_with_difficulty`` to ``HfRandomUniformTerrainCfg``. When
+  enabled, the noise amplitude scales with difficulty (flat at 0, full
+  ``noise_range`` at 1) so the terrain progresses in a curriculum. Defaults to
+  ``False``, preserving the previous difficulty-independent behavior.
+
+Changed
+^^^^^^^
+
+- Bumped ``rsl-rl-lib`` from 5.2.0 to 5.4.0.
+- Curriculum-mode terrain difficulty is now deterministic across rows
+  and reaches the configured ``difficulty_range`` endpoints
+  (:issue:`1027`).
+- Heightfield terrains now color by absolute height with a diverging palette
+  (cool below the ground plane, green at ground level, warm above) on a fixed
+  scale, replacing the per-patch normalization. Color is now consistent across
+  terrains, and low-amplitude terrain such as ``random_rough`` reads as gently
+  tinted ground instead of high-contrast noise.
+- ``BoxNestedRingsTerrainCfg`` now builds uniform-height concentric ridges
+  whose separating gaps widen with difficulty, replacing the random per-ring
+  heights. Rings are colored by height (like the other terrains) and the outer
+  border matches the ring height.
+- Terrain generation no longer prints timing information to stdout.
+
+Fixed
+^^^^^
+
+- Fixed ``select_gpus`` crashing when ``CUDA_VISIBLE_DEVICES`` contains MIG UUIDs instead of numeric indices.
+- Fixed pyramid-stairs terrains (``BoxPyramidStairsTerrainCfg``,
+  ``BoxInvertedPyramidStairsTerrainCfg``, and ``BoxOpenStairsTerrainCfg``)
+  leaving an empty, geometry-free border at difficulty 0, where the step
+  height collapses to zero. The flat border frame is now always generated as
+  solid geometry flush with the ground (:issue:`1033`).
+- Fixed ``HfPerlinNoiseTerrainCfg`` failing to compile at difficulty 0, where
+  the target height collapses to zero and MuJoCo rejects the non-positive
+  heightfield size.
+- Fixed ``BoxRandomGridTerrainCfg`` producing NaN colors (and failing to build)
+  at difficulty 0, where the grid height is zero and the color normalization
+  divided by zero.
+- Fixed the center platform z-fighting with surrounding geometry in
+  ``BoxRandomGridTerrainCfg`` (grid cells were left underneath the platform) and
+  ``BoxRandomSpreadTerrainCfg`` (the platform duplicated the floor surface).
+- Fixed ``BoxNarrowBeamsTerrainCfg`` square platform corners protruding between
+  the beams at high difficulty; the platform now shrinks to stay within the
+  beams' angular coverage.
+- Fixed ``BoxSteppingStonesTerrainCfg`` reconfiguring abruptly at a difficulty
+  threshold, where the stone grid re-tiled as its spacing crossed an integer
+  boundary, and leaving an oversized gap around the center platform. The grid is
+  now difficulty-independent and the platform snaps to it as a clean island.
+- Fixed ``train --video``, ``play``, and ``demo`` crashing with ``OpenGL
+  platform library not loaded`` on headless Linux hosts that don't pre-set
+  ``MUJOCO_GL``. The default is now applied in ``mjlab/__init__.py`` (Linux
+  only) so it takes effect before mujoco's GL backend selection runs.
+
+Version 1.4.0 (May 26, 2026)
+----------------------------
+
+Added
+^^^^^
+
+- Added ``BuiltinPdActuator``, the implicit-integration version of
+  ``IdealPdActuator``. Same interface (position + velocity targets,
+  kp/kd gains), but expresses the PD as native MuJoCo ``<position>``
+  and ``<velocity>`` elements so the ``implicit`` / ``implicitfast``
+  integrators include the kp/kd derivatives in their velocity update.
+  The actuator stays stable at gain/timestep combinations where
+  explicit Python PD would diverge, which matters when you want to
+  run a real motor's stiff on-board PD gains in sim. ``effort_limit``
+  is enforced as a sum-clamp on the two PD terms via
+  ``jnt_actfrcrange`` (or ``tendon_actfrcrange``). Supported by
+  ``dr.pd_gains`` and ``dr.effort_limits``.
+- Added ``mdp.projected_gravity_from_sensor``, an observation that derives
+  projected gravity from a ``framezaxis`` up-vector sensor (negated) rather
+  than from the root body orientation. Unlike ``mdp.projected_gravity``, it
+  reflects the sensor's site frame, so it can observe IMU mounting domain
+  randomization (e.g. via ``dr.site_quat``). Go1 and G1 ship an
+  ``imu_upvector`` sensor for this.
+- Added ``DebugVisualizer.add_box`` for drawing an axis-oriented box
+  primitive, mirroring ``add_ellipsoid``. Supported by both the native
+  and Viser viewers. ``size`` is the box half-extents (:issue:`992`).
 - Added ``--log-root`` CLI option to ``train``, ``play``, and ``evaluate``
  scripts for choosing where training logs are stored. Defaults to
  ``logs/rsl_rl`` (unchanged behavior). Useful for directing outputs to a
@@ -21,22 +104,41 @@ Added
  primary names in the order they appear along the per-contact axis of the
  output tensors. This makes it possible to map a contact-data column back
  to the primary it belongs to (:issue:`914`).
- Added per-world mesh variant support via ``VariantEntityCfg`` and
-  ``VariantCfg``. Each world in a batched simulation can now use a
-  different mesh asset for the same logical entity (e.g. world 0 holds a
-  cube, world 1 a sphere), with weights controlling the proportion of
-  worlds assigned to each variant. Mesh-derived constants (collision
-  bounds, body inertials, subtree mass, inverse weights) are compiled
-  per-variant and stored as per-world arrays in the Warp model, so domain
-  randomization, the native viewer, the offscreen renderer, and the Viser
-  viewer all pick up the variant assignment automatically. Variants must
-  share the same kinematic structure (same bodies, joints, joint types);
-  only mesh geoms may differ. Assignment is fixed at simulation init.
-  See :ref:`per_world_mesh` for usage. With help from @XiangruiJiang.
+- Added per-world mesh variant support via ``VariantEntityCfg``. Each
+  world in a batched simulation can now use a different mesh asset for
+  the same logical entity (e.g. world 0 holds a cube, world 1 a
+  sphere). Variants are passed as a ``dict[str, Callable]`` of named
+  spec callables; the optional ``assignment`` field controls how worlds
+  map to variants and accepts ``None`` (uniform), a ``dict[str, float]``
+  of per-variant weights, or a custom ``Callable[[int], Sequence[int]]``.
+  Mesh-derived constants (collision bounds, body inertials, subtree
+  mass, inverse weights) are compiled per-variant and stored as
+  per-world arrays in the Warp model, so domain randomization, the
+  native viewer, the offscreen renderer, and the Viser viewer all pick
+  up the variant assignment automatically. Variants must share the
+  same kinematic structure (same bodies, joints, joint types); only
+  mesh geoms may differ. Assignment is fixed at simulation init. See
+  :ref:`heterogeneous_worlds` for usage. With help from @XiangruiJiang.
+- Per-world mesh variants now support per-variant materials and textures.
+  Each variant can reference its own named material, which is automatically
+  prefixed and scattered via ``geom_matid`` alongside the existing
+  ``geom_dataid`` table. Variants without a material get ``matid = -1``.
+  Contribution by @omarrayyann.

 Changed
 ^^^^^^^

+- ``Entity`` now raises a clear error at construction when its spec contains
+  more than one freejoint. An entity models a single system rooted at one
+  body, so it has at most one freejoint; a second one was previously accepted
+  silently and only surfaced later as a cryptic shape mismatch when writing
+  root state. Model each detached floating body as its own entry in
+  ``SceneCfg.entities`` instead.
+- Changed ``compute_root_relative_mpkpe`` to re-anchor the reference to the
+  robot's root each step, removing yaw drift as well as translation so it
+  measures intrinsic body pose error.
+- Changed ``compute_joint_velocity_error`` from an L2 norm to a per-joint
+  RMS, so it no longer scales with the number of joints.
 - Bumped ``mujoco`` to 3.8 and ``mujoco-warp`` to 3.8.0. The ``multiccd``
  enable flag was removed in mujoco 3.8 (it became default-on), so configs
  that listed ``"multiccd"`` in ``MujocoCfg.enableflags`` need to drop it.
@@ -68,15 +170,63 @@ Changed
  air-time fields (``current_air_time``, ``last_air_time``,
  ``current_contact_time``, ``last_contact_time``) have shape ``[B, P]``,
  where ``P`` is the number of resolved primaries (:issue:`914`).
+- Event functions now share a single ``resolve_env_ids`` helper to expand
+  ``env_ids=None`` to all environments, replacing five copies of the same
+  guard. ``push_by_setting_velocity`` and ``apply_external_force_torque``
+  accept ``env_ids=None`` too, so they work as global-time interval terms.
+  Documented when to use ``apply_external_force_torque`` (a constant,
+  self-managed wrench) versus ``apply_body_impulse`` (transient, automatic
+  impulses) versus ``push_by_setting_velocity`` (an instantaneous velocity
+  kick).

 Fixed
 ^^^^^

- Fixed ``ManagerBasedRlEnv`` initializing Warp on all visible CUDA devices
-  even when constructed with ``device="cpu"``. ``seed_rng`` now accepts a
-  ``device`` argument and skips ``wp.rand_init`` on CPU devices, so a
-  CPU-only env no longer claims a CUDA context on machines with a visible
-  GPU (:issue:`949`).
+- Removed use of deprecated ``warp-lang`` symbols (``wp.context.runtime``
+  and ``wp.context.Device``) that were dropped in newer ``warp-lang``
+  releases, causing ``AttributeError: module 'warp' has no attribute
+  'context'`` at import/runtime. mjlab now uses
+  ``wp.get_cuda_driver_version()`` and ``wp.Device`` instead
+  (:issue:`967`). Contribution by @rdeits.
+- Fixed the tracking ``evaluate`` script scoring each metric against the
+  next motion frame; the reference is now snapshotted before each step to
+  match the reward.
+- Fixed the tracking end-effector metrics silently scoring zero for an
+  unknown body name; they now raise ``ValueError``.
+- Fixed ``compute_mpkpe`` measuring root-relative instead of global error;
+  it now uses the global reference ``body_pos_w`` (:issue:`1006`).
+- Fixed heavy flicker in offscreen training videos on rough-terrain tasks.
+  The renderer recomputed its context "neighbor" robots every frame from
+  ``env_origins``, which the terrain curriculum mutates on reset, so the
+  neighbor set kept changing and robots popped in and out. The neighbor
+  set is now computed once and cached (:issue:`979`).
+- Fixed command delay only applying to an actuator's position target.
+  ``IdealPdActuator`` and ``DcMotorActuator`` also use velocity and effort, which
+  arrived undelayed and out of sync; all command targets now share one delay.
+  Zero-reference setups are unaffected.
+- Fixed duplicate random seeds across nodes in multi-node training. The
+  per-process seed offset in ``scripts/train.py`` now uses the global
+  ``RANK`` instead of ``LOCAL_RANK``. Contribution by @bd-pdomanico.
+- Fixed ``apply_body_impulse`` firing an impulse on the very first step (and
+  the first step after every reset) instead of starting with a cooldown as
+  documented. The cooldown is now sampled lazily on the first call so impulse
+  timing is decorrelated from episode resets (:issue:`973`).
+- Fixed ``dr.pd_gains`` and ``dr.effort_limits`` silently no-oping when
+  passed an ``Operation`` object (e.g. ``dr.scale``) instead of a string.
+  Both functions now accept ``Operation | str`` like every other DR event
+  and raise ``ValueError`` for unsupported operations (:issue:`971`).
+- Fixed ``ContactSensor`` with ``global_frame=True`` and
+  ``reduce`` ∈ {``"none"``, ``"mindist"``, ``"maxforce"``} producing forces
+  rotated onto the wrong axis. The contact-frame→world rotation matrix had
+  its columns ordered ``[tangent, tangent2, normal]`` instead of
+  ``[normal, tangent, tangent2]``, projecting the normal-force component
+  onto a tangent direction. Contribution by @bd-pdomanico.
+- Fixed ``extras["log"]`` entries written by reward terms (e.g. ``Metrics/*``
+  values in velocity tasks) being silently discarded on any step where at
+  least one environment resets. ``_reset_idx`` was clearing the dict after
+  ``reward_manager.compute()`` had already populated it. The clear now
+  happens at the top of ``step()`` and ``reset()`` so that all entries
+  survive (:issue:`957`).
 - Fixed ``ContactSensor.compute_first_contact`` and ``compute_first_air``
  occasionally missing events when a contact began or ended right at the
  last physics substep of a control step. ``current_contact_time`` /
@@ -197,8 +197,8 @@ example, a ``CollisionCfg`` with ``geom_names_expr=(".*_foot.*",)``
 sets contact parameters only on foot geoms. See the asset zoo
 (``mjlab.asset_zoo.robots``) for complete examples.

-Per-world mesh variants
-^^^^^^^^^^^^^^^^^^^^^^^
+Heterogeneous worlds
+^^^^^^^^^^^^^^^^^^^^

 For scenes that need different mesh assets in different parallel worlds
 (for example, training a manipulation policy that generalizes across
@@ -206,7 +206,7 @@ object shapes), use ``VariantEntityCfg`` instead of ``EntityCfg``. Each
 world is assigned a variant proportional to a configurable weight, and
 mesh-dependent compiled constants (collision bounds, body inertials,
 subtree mass) are stored as per-world arrays so domain randomization and
-viewers stay consistent. See :ref:`per_world_mesh`.
+viewers stay consistent. See :ref:`heterogeneous_worlds`.

 Subclassing Entity
 ^^^^^^^^^^^^^^^^^^
@@ -1,50 +1,35 @@
-.. _per_world_mesh:
+.. _heterogeneous_worlds:

-Mesh Variants
-=============
+Heterogeneous Worlds
+====================

-Mesh variants let a single batched simulation run with different mesh
-assets in different parallel worlds. World 0 may simulate a cube, world
-1 a sphere, and world 2 a bowl, all sharing the same compiled scene
-and the same kinematic structure. The result is a heterogeneous batch
-in which the mesh and its derived constants vary across worlds while
-everything else (the body tree, the joint structure, the contact and
-solver setup) is fixed.
-
-Mesh variants are configured at the entity level through
-``VariantEntityCfg`` and ``VariantCfg``. Once configured,
-domain randomization, the native viewer, the offscreen renderer, and
-the Viser viewer all pick up the variant assignment automatically.
+mjlab can run a single batched simulation in which different parallel
+worlds use different mesh assets for the same logical entity. World 0
+may simulate a cube, world 1 a sphere, world 2 a bowl. All worlds
+share the same compiled scene and the same body and joint structure;
+only the meshes and the per-geom attributes that travel with them
+(friction, contact bits, mass, density, and a few more) differ across
+worlds. Articulated props work too (you can have a hinge or slide
+below the variant's root), as long as the joint topology matches
+across variants. The feature is exposed through ``VariantEntityCfg``.
+The full breakdown of what can and cannot vary across variants is in
+the next section.


-How it works
------------
+Quickstart
+----------

-A standard ``EntityCfg`` provides a single ``spec_fn`` that returns one
-``MjSpec``. A ``VariantEntityCfg`` provides a dictionary of named
-variants, each with its own ``spec_fn`` and a weight controlling the
-proportion of worlds that use it.
-
-**All variants must declare the same kinematic structure.** The batched
-simulator assumes a single topology across worlds; per-world variation
-is confined to mesh assets and the constants derived from them. mjlab
-uses the first variant's body tree as the template and copies mesh
-assets and explicit body inertials from the others. Geom-level
-properties on later variants such as ``rgba``, friction, and material
-assignments are not propagated; control per-world appearance through
-domain randomization on ``geom_rgba`` or ``mat_rgba``. The structural
-check is enforced at construction time and raises a ``ValueError``
-describing the first mismatch. Variants must also be floating-base
-(declare a free joint on the root body); fixed-base variants are
-rejected.
-
-A minimal two-variant config:
+Say you want some parallel worlds to hold a sphere and others to hold
+a cone, with a single shared scene running both at once. Define each
+variant as a function that returns an ``MjSpec``, then group them
+under one ``VariantEntityCfg``:

 .. code-block:: python

    import mujoco

-    from mjlab.entity import EntityCfg, VariantCfg, VariantEntityCfg
+    from mjlab.entity import EntityCfg, VariantEntityCfg
+

    def make_sphere_spec() -> mujoco.MjSpec:
        spec = mujoco.MjSpec()
@@ -56,40 +41,267 @@ A minimal two-variant config:
        body.add_geom(type=mujoco.mjtGeom.mjGEOM_MESH, meshname="visual")
        return spec

-    # ``make_cone_spec`` follows the same shape with
-    # ``mesh.make_cone(nedge=16, radius=0.04)`` in place of the sphere call.
+
+    def make_cone_spec() -> mujoco.MjSpec:
+        spec = mujoco.MjSpec()
+        mesh = spec.add_mesh(name="visual")
+        mesh.make_cone(nedge=16, radius=0.04)
+        body = spec.worldbody.add_body(name="prop")
+        body.add_freejoint()
+        body.add_geom(type=mujoco.mjtGeom.mjGEOM_MESH, meshname="visual")
+        return spec
+

    object_cfg = VariantEntityCfg(
        variants={
-            "sphere": VariantCfg(spec_fn=make_sphere_spec, weight=1.0),
-            "cone": VariantCfg(spec_fn=make_cone_spec, weight=2.0),
+            "sphere": make_sphere_spec,
+            "cone":   make_cone_spec,
        },
+        assignment={"cone": 2.0},  # twice as many cones as spheres
        init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
    )

-During scene construction mjlab merges the per-variant specs into a
-single ``MjSpec`` whose mesh slots are padded to the maximum count any
-variant uses, then writes a per-world ``geom_dataid`` table that
-selects the right mesh for each world. In the merged scene
-``geom_dataid`` is no longer a flat ``(ngeom,)`` vector but a
-``(num_envs, ngeom)`` table whose rows differ by variant. A value of
-``-1`` marks a disabled mesh slot, used for variants with fewer mesh
-geoms than the maximum.
+Plug the variant entity into a :ref:`scene` exactly like a regular
+``EntityCfg``:

-Mesh choice is entangled with several other compiled-model constants:
-geom collision bounds, geom local frames, body inertials, subtree mass,
-and inverse weights. mjlab compiles each unique row of the
-``geom_dataid`` table on the host and copies the relevant compiled
-fields into per-world arrays on the GPU, so each world's compiled
-constants stay consistent with that world's mesh selection. The full
-list of fields handled this way is in
-``mjlab.sim.mesh_variants.VARIANT_DEPENDENT_FIELDS``.
+.. code-block:: python
+
+    from mjlab.scene import SceneCfg
+
+    scene_cfg = SceneCfg(
+        num_envs=4096,
+        entities={"object": object_cfg},
+    )
+
+Twice as many worlds will hold a cone as a sphere. Variants not listed
+in the ``assignment`` dict default to weight 1.0; omit ``assignment``
+entirely for uniform allocation across all variants.
+
+
+What variants can differ in
+---------------------------
+
+**Free to vary across variants:** the mesh asset assigned to each
+slot, the number of mesh geoms per ``(body, role)`` bucket on the
+variant body (one variant can have more collision meshes than
+another), the per-mesh-geom attributes that travel with the mesh
+(friction, contact bits, mass, density, ``condim``, and a handful of
+others), and explicit body inertial values within whichever single
+inertial mode the variants agree on per body.
+
+**Must match across variants:** the body tree, joint topology,
+primitive (non-mesh) geoms, and any actuators / sensors / tendons /
+equalities. Variants must also agree on the inertial representation
+per body (mesh-derived, diagonal, or fullinertia), and may not use the
+reserved ``mjlab/pad/`` name prefix on any element. Variant entities
+must also be floating-base: the root body declares a freejoint.
+
+The validator runs at entity build time and raises ``ValueError``
+naming the offending variant and the exact mismatch.
+
+
+How variants are assembled
+--------------------------
+
+mjlab merges every variant's mesh assets into a single ``MjSpec`` and
+gives the variant body enough mesh-geom *slots* to cover the maximum
+mesh count any variant uses for each ``(body, role)`` bucket. A slot
+is identified by ``(body_path, role, ordinal)``. ``role`` is "visual"
+or "collision", derived from ``contype``/``conaffinity``;
+mujoco_warp's ``geom_contype``/``geom_conaffinity`` are 1D shared
+(not per-world), so a slot's role is fixed across worlds by
+construction.
+
+A worked example
+~~~~~~~~~~~~~~~~
+
+Say variant ``sphere`` has 1 visual mesh geom and 2 collision mesh
+geoms on the prop body, and variant ``cone`` has 1 visual mesh geom
+and 4 collision mesh geoms on the same body.
+
+.. code-block:: text
+
+    sphere variant body              cone variant body
+    -------------------              -------------------
+    prop body                        prop body
+      [visual] sphere_vis              [visual] cone_vis
+      [coll]   sphere_col_0            [coll]   cone_col_0
+      [coll]   sphere_col_1            [coll]   cone_col_1
+                                       [coll]   cone_col_2
+                                       [coll]   cone_col_3
+
+mjlab walks each variant's body tree, buckets mesh geoms by
+``(body_path, role)``, and lays the union out as slots:
+
+.. list-table::
+   :header-rows: 1
+   :widths: 8 18 8 12 27 27
+
+   * - Slot
+     - body_path
+     - role
+     - ordinal
+     - sphere fills with
+     - cone fills with
+   * - 0
+     - /prop
+     - visual
+     - 0
+     - sphere_vis
+     - cone_vis
+   * - 1
+     - /prop
+     - collision
+     - 0
+     - sphere_col_0
+     - cone_col_0
+   * - 2
+     - /prop
+     - collision
+     - 1
+     - sphere_col_1
+     - cone_col_1
+   * - 3
+     - /prop
+     - collision
+     - 2
+     - *(unfilled)*
+     - cone_col_2
+   * - 4
+     - /prop
+     - collision
+     - 3
+     - *(unfilled)*
+     - cone_col_3
+
+Five slots total. The merged scene's prop body has five mesh geoms:
+slot 0 plus four collision slots (the union of sphere's two and
+cone's four). At merge time, every variant's mesh asset is added to
+the merged spec under a unique name (e.g.
+``sphere/sphere_vis``, ``cone/cone_col_2``).
+
+The merged scene compiles once into a single canonical ``MjModel``
+that every world in the batch agrees on layout-wise: same nbody,
+ngeom, same body and geom IDs. mjlab's per-world overrides on top of
+that one model are what make worlds heterogeneous.
+
+What each world sees at runtime
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Worlds where ``sphere`` is active see only its three meshes; the two
+extra collision slots are disabled via per-world ``geom_dataid = -1``,
+and mujoco_warp skips them. Worlds where ``cone`` is active see all
+five meshes wired up.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 14 14 12 12 12 12 12
+
+   * - World
+     - variant
+     - slot 0
+     - slot 1
+     - slot 2
+     - slot 3
+     - slot 4
+   * - 0
+     - sphere
+     - sphere_vis
+     - sphere_col_0
+     - sphere_col_1
+     - **off (-1)**
+     - **off (-1)**
+   * - 1
+     - cone
+     - cone_vis
+     - cone_col_0
+     - cone_col_1
+     - cone_col_2
+     - cone_col_3
+
+Three categories of per-world override carry the variation:
+
+* **geom_dataid** is a ``(num_envs, ngeom)`` table. Its row for
+  world W picks which compiled mesh each slot points at. ``-1`` is
+  the "skip me" sentinel mujoco_warp already understands.
+* **Mesh-derived fields** (``geom_size``, ``geom_rbound``,
+  ``geom_aabb``, ``geom_pos``, ``geom_quat``, ``body_mass``,
+  ``body_subtreemass``, ``body_inertia``, ``body_invweight0``,
+  ``body_ipos``, ``body_iquat``) are stored as ``(num_envs, ...)``
+  arrays. The values for sphere worlds reflect a sphere-shaped
+  inertia tensor and sphere-sized AABBs; the values for cone worlds
+  reflect the cone. The full list is in
+  ``mjlab.entity.variants.VARIANT_DEPENDENT_FIELDS``.
+* **Per-mesh-geom attributes** (contact bits, friction, mass,
+  density, condim, group, priority, rgba, solref, solimp, margin,
+  gap) are captured per variant in ``VariantGeomSpec`` at merge time
+  and restored verbatim on the slot geom during the per-variant
+  reference compile. So if sphere's collision geoms have
+  ``friction=0.5`` and cone's have ``friction=1.2``, world W's
+  per-step friction reflects the assigned variant's source value.
+  The one exception is ``material``, which is not propagated across
+  variants; if you need per-world appearance variation use DR on
+  ``geom_rgba`` / ``mat_rgba``.
+
+If ``sphere`` adds a body that ``cone`` lacks (or vice versa), the
+validator rejects the configuration before any of the merge logic
+runs. The slot mechanism only flexes mesh geom counts within
+matching bodies; everything structural above the geom level must
+agree.
+
+.. note::
+
+   **Doesn't compiling the merged scene ruin the prop body's
+   inertia?**
+
+   No, but it's worth understanding why, because the naive intuition
+   says it should. If you stuck every variant's mesh geoms on the
+   prop body and called ``spec.compile()``, MuJoCo would sum each
+   geom's inertial contribution, and you would get a body whose mass
+   and inertia tensor are a meaningless mix of every variant's shape.
+
+   mjlab avoids this in two layers:
+
+   * **The merged scene does not stick every variant's geoms on the
+     body.** The prop body in the merged spec carries variant 0's
+     mesh geoms (with their original mass and density) plus, for any
+     slot variant 0 doesn't fill, a synthesized padding geom that has
+     ``mass = 0`` and ``density = 0``. Padding contributes nothing to
+     body inertia. Other variants' meshes are present in the merged
+     spec only as **mesh assets** (in the assets section, not as geoms
+     on any body). They get wired in at runtime via per-world
+     ``geom_dataid`` and never affect the host compile's inertial
+     sums.
+   * **Per-world overrides come from per-variant source compiles.**
+     Even with the above, the merged-scene compile's prop body inertia
+     is only correct for variant 0. For every other variant, mjlab
+     compiles that variant's original source spec in isolation (one
+     body, one variant's worth of meshes), reads the resulting
+     ``body_mass``, ``body_inertia``, ``body_ipos``, ``body_iquat``,
+     ``body_invweight0``, and ``body_subtreemass``, and writes them
+     into the per-world arrays at the prop body's index.
+
+   Net result: world W's prop body inertia is byte-equal to what you
+   would get by compiling variant W's source spec on its own. There
+   is a regression test
+   (``test_visual_collision_split_inertia_matches_independent_compile``
+   in ``tests/test_variants.py``) that asserts exactly this against
+   independent per-variant compiles.


 World assignment
 ----------------

-mjlab assigns variants to worlds proportionally by weight using the
+How worlds get mapped to variants is controlled by the ``assignment``
+field on ``VariantEntityCfg``. It accepts three shapes:
+
+* ``None`` (default): uniform allocation across variants.
+* ``dict[str, float]``: per-variant weights. Variants not listed
+  default to weight 1.0.
+* ``Callable[[int], Sequence[int]]``: an explicit assignment function
+  called with ``num_envs`` at simulation init.
+
+Both the ``None`` and dict cases use the
 `largest remainder method
 <https://en.wikipedia.org/wiki/Largest_remainder_method>`_. Each
 variant's quota is ``q_i = (w_i / sum(w)) * num_envs``; each variant
@@ -98,14 +310,28 @@ first receives ``floor(q_i)`` worlds, and the remaining
 fractional remainders, with ties broken by declaration order. For
 ``num_envs = 10`` and weights ``(1.0, 2.0, 1.0)`` this gives
 ``(3, 5, 2)`` worlds per variant. Weights are normalized internally,
-so ``(1, 2, 1)`` and ``(0.25, 0.5, 0.25)`` produce identical
-assignments. A weight of zero is allowed and produces zero worlds for
-that variant; at least one variant must have a positive weight.
+so ``{"a": 1, "b": 2, "c": 1}`` and ``{"a": 0.25, "b": 0.5, "c": 0.25}``
+produce identical assignments. A weight of zero is allowed and
+produces zero worlds for that variant; at least one variant must end
+up with positive weight.

-Variant assignment is fixed at simulation initialization and does not
-resample on episode reset. The intended use is heterogeneous training
-across the batch, not per-episode mesh randomization. To inspect the
-assignment from user code, read ``env.sim.world_to_variant``:
+The default and dict paths are purely deterministic given
+``(assignment, num_envs)``. With ``assignment={"a": 1, "b": 1}`` and
+``num_envs = 8`` you always get ``[0, 0, 0, 0, 1, 1, 1, 1]``. There is
+no seed involved; rerunning the same config produces the same
+partition every time. Note that the partition's *boundaries* depend
+on ``num_envs``, so world W's variant is not necessarily stable when
+you change ``num_envs``. If you need explicit per-world stability
+across batch sizes (e.g. "world 0 is always variant 0, world 1 is
+always variant 1, regardless of how many envs I launch"), use a
+callable assignment as below.
+
+Variant assignment is fixed at ``Simulation`` initialization and does
+not resample on episode reset. The intended use is heterogeneous
+training across the batch, not per-episode mesh randomization.
+
+Read the resolved assignment from user code via
+``env.sim.world_to_variant``:

 .. code-block:: python

@@ -118,23 +344,53 @@ variants were declared in ``VariantEntityCfg.variants``. The dict is
 empty for non-variant scenes.


+Custom assignment with a callable
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+When the weighted default is not what you want, pass a callable to
+``assignment``. The callable receives ``num_envs`` and must return a
+length-``num_envs`` sequence of variant indices in
+``[0, len(variants))``. The returned sequence's length and bounds are
+validated at sim init; mismatches raise a ``ValueError`` naming the
+offending entity.
+
+A few patterns:
+
+**Round-robin** - cycle through variants by world index.
+
+.. code-block:: python
+
+    cfg = VariantEntityCfg(
+        variants={"a": make_a, "b": make_b, "c": make_c},
+        assignment=lambda n: [w % 3 for w in range(n)],
+    )
+
+**Stratified halves** - first half is variant 0, second half is
+variant 1.
+
+.. code-block:: python
+
+    cfg = VariantEntityCfg(
+        variants={"easy": make_easy, "hard": make_hard},
+        assignment=lambda n: [0] * (n // 2) + [1] * (n - n // 2),
+    )
+
 Domain randomization
 --------------------

 Domain randomization on variant scenes preserves per-variant baselines
 automatically. When the simulation initializes, mjlab snapshots the
-variant-dependent fields (``body_mass``, ``body_inertia``,
-``geom_size``, and others listed in ``VARIANT_DEPENDENT_FIELDS``) as
-``(num_envs, ...)`` tensors and registers them in
-``sim.per_world_default_fields``. Domain randomization operations that
-read defaults (scale, additive offsets) detect this registration and
-index the per-world default array by environment, so a 10% mass scale
+variant-dependent fields as ``(num_envs, ...)`` tensors and registers
+them in ``sim.per_world_default_fields``. DR operations that read
+defaults (scale, additive offsets) detect this registration and index
+the per-world default array by environment, so a 10% mass scale
 applied across a batch containing a 100 g sphere variant and a 1 kg
-cube variant produces 10% perturbations around each variant's own
-mass, not 10% of a shared template mass. Fields that are not
-variant-dependent (``geom_friction``, ``dof_armature``,
-``dof_damping``, and so on) behave identically on variant and
-non-variant scenes.
+cube variant produces 10% perturbations *around each variant's own
+mass*, not 10% of a shared template mass.
+
+Fields that are not variant-dependent (``geom_friction``,
+``dof_armature``, ``dof_damping``, and so on) behave identically on
+variant and non-variant scenes.

 For inertial randomization the recommended path is
 ``dr.pseudo_inertia``, which jointly randomizes mass, COM offset,
@@ -169,32 +425,66 @@ Convex hull visualization is computed per variant from the variant's
 mesh vertices.


-Performance considerations
--------------------------
+Performance
+-----------

-Mesh variants do not add per-step overhead in the GPU kernels.
-Variant-dependent fields are stored as per-world arrays accessed by
-world index in the existing kernels, with no branching or dispatch
-on variant.
+**Per-step cost is unaffected by variant count.** Variant-dependent
+fields are stored as per-world arrays accessed by world index in the
+existing kernels, with no branching or dispatch on variant.

-Initialization is the main consideration. mjlab compiles each unique
-row of the ``geom_dataid`` table by taking a fresh ``MjSpec.copy()``,
-editing the mesh selection and (if applicable) the explicit body
-inertials, and calling ``spec.compile()``. This work scales with the
-number of unique variant combinations rather than with ``num_envs``.
-For a scene with one variant entity declaring k variants, this is k
-host compiles regardless of how many worlds use each variant. With
-multiple variant entities the unique-row count is bounded by the
-product of their variant counts in the worst case, so a scene with
-two variant entities of 5 variants each could trigger up to 25 host
-compiles at init.
+**Construction cost is linear in the total variant count.** mjlab
+compiles the merged scene once to produce the canonical ``MjModel``,
+then compiles each variant's original (un-merged) source spec in
+isolation to recover that variant's per-body and per-geom mesh-derived
+fields. Each per-variant compile sees only that variant's single body
+and mesh, so its cost is independent of the total number of variants
+in the scene.

-``MjSpec.copy()`` and ``spec.compile()`` are non-trivial operations,
-and their cost grows with scene size. For a scene with many variant
-entities or many variants per entity, the cumulative initialization
-cost can be measured in seconds. This cost is paid once at startup
-and does not affect training throughput.
+For a scene with one variant entity declaring k variants, construction
+runs ``1 + k`` compiles. With multiple variant entities, compiles
+decouple across entities: two variant entities of 5 variants each cost
+``1 + 5 + 5 = 11`` compiles, not ``1 + 5 * 5 = 26``. As an order of
+magnitude on CPU with typical procedural meshes, each per-variant
+compile takes around 1-2 ms, so a scene with 100 variants pays a few
+hundred milliseconds at startup and a scene with 1000 variants pays
+roughly two seconds.

-The merged spec contains every variant's mesh assets simultaneously.
-Memory footprint at scene-build time scales with the total number of
-mesh vertices and faces across all declared variants.
+The merged spec contains every variant's mesh assets simultaneously,
+so memory at scene-build time scales with the total mesh vertex /
+face count across all variants. This is paid once at startup and does
+not affect training throughput.
+
+
+Limitations
+-----------
+
+**Floating-base only.** Each variant's root body must declare a free
+joint. Fixed-base variants are rejected; mocap auto-wrapping that
+applies to non-variant entities is not applied here.
+
+**Material assets are not propagated.** Each variant's ``contype``,
+``conaffinity``, ``condim``, ``friction``, ``mass``, ``density``,
+``group``, ``priority``, ``rgba``, ``solref``, ``solimp``, ``margin``,
+and ``gap`` are restored per-world during compile, but the
+``material`` reference on slot geoms inherits whichever material the
+template variant set. Use DR on ``geom_rgba`` / ``mat_rgba`` for
+per-world appearance variation.
+
+**Assignment is fixed at sim init.** There is no API to swap a world
+to a different variant on episode reset. World W's mesh asset is
+whatever it was assigned at init for the lifetime of the simulation.
+Per-episode mesh randomization is not supported today; DR can vary
+scalar properties (mass, friction, color, scale) on a fixed variant
+but cannot swap one mesh for another.
+
+**No support for per-world differing kinematic topology.** Variants
+must share the same body tree, joints, and actuator/sensor counts,
+so you cannot configure things like:
+
+* a different number of objects per world (world 0 has two props on
+  the table, world 1 has three);
+* different articulation per world (world 0's prop is an articulated
+  drawer with a slider joint, world 1's prop is a rigid block).
+
+True heterogeneous topology requires upstream support in mujoco_warp
+that does not currently exist.
@@ -42,6 +42,27 @@ Not all CUDA versions are supported by MuJoCo Warp.
 - **Recommended**: CUDA **12.4+** (for conditional execution support in CUDA
  graphs).

+How do I run on CPU without touching the GPU?
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Passing ``device="cpu"`` puts all mjlab computation on the CPU, but it does
+**not** stop Warp from initializing the GPU. The first time Warp's runtime
+comes up, it eagerly enumerates and creates a CUDA context on **every**
+visible device, regardless of which device you requested. So on a machine
+with a visible GPU, a ``device="cpu"`` run still claims VRAM.
+
+This happens inside Warp and cannot be prevented from Python once the
+package is imported. To keep the process entirely off the GPU, hide the
+devices from CUDA before launching:
+
+.. code-block:: bash
+
+   CUDA_VISIBLE_DEVICES="" uv run train.py ...
+
+With no visible CUDA devices, Warp initializes CPU-only and never allocates
+on the GPU. See `issue #949
+<https://github.com/mujocolab/mjlab/issues/949>`_ for background.
+
 Performance
 -----------

@@ -49,7 +49,6 @@ the geometry and how it scales with difficulty.
        terrain_generator=TerrainGeneratorCfg(
            size=(8.0, 8.0),
            num_rows=10,
-            num_cols=20,
            border_width=20.0,
            curriculum=True,
            sub_terrains={
@@ -70,10 +69,12 @@ the geometry and how it scales with difficulty.
        max_init_terrain_level=5,
    )

-The generator creates a ``num_rows x num_cols`` grid of patches. The
-``sub_terrains`` dictionary maps names to ``SubTerrainCfg`` instances,
-and each sub-terrain's ``proportion`` weight controls how many columns
-(curriculum mode) or sampling probability (random mode) it receives.
+The generator creates a grid of patches sized ``num_rows`` by either
+``num_cols`` (random mode) or ``len(sub_terrains)`` (curriculum mode,
+where ``num_cols`` is ignored). The ``sub_terrains`` dictionary maps
+names to ``SubTerrainCfg`` instances; each sub-terrain's ``proportion``
+controls robot spawning distribution across columns in curriculum mode,
+or per-patch sampling probability in random mode.


 Grid layout
@@ -82,30 +83,48 @@ Grid layout
 Two generation modes control how terrain types are distributed across
 the grid:

-**Curriculum mode** (``curriculum=True``). Columns are deterministically
-assigned to terrain types based on their ``proportion`` weights. A type
-with proportion 0.4 in a 20-column grid gets 8 columns. All patches in
-a column share the same terrain type, and difficulty increases from row 0
-(easiest) to row ``num_rows - 1`` (hardest). This structured layout is
-what enables the curriculum system to advance environments to harder rows
-as performance improves.
+**Curriculum mode** (``curriculum=True``). Each terrain type gets exactly
+one column; the generator uses ``len(sub_terrains)`` columns regardless of
+``num_cols``. All patches in a column share the same terrain type, and
+difficulty increases from row 0 (easiest) to row ``num_rows - 1``
+(hardest). The ``proportion`` field controls how robots are distributed
+across columns at spawn time, not column count. This structured layout
+is what enables the curriculum system to advance environments to harder
+rows as performance improves.

 **Random mode** (``curriculum=False``). Every patch independently samples
 a terrain type weighted by ``proportion`` and a difficulty from
-``difficulty_range``. This provides maximum variety but no structured
-difficulty progression.
+``difficulty_range``. ``num_cols`` is honored. This provides maximum
+variety but no structured difficulty progression.


 The difficulty parameter
 ^^^^^^^^^^^^^^^^^^^^^^^^

 Each sub-terrain's generation function receives a ``difficulty`` value
-in ``[0, 1]``. This value linearly interpolates the terrain's
-configurable ranges. For example, a ``BoxPyramidStairsTerrainCfg`` with
+that linearly interpolates the terrain's configurable ranges. For
+example, a ``BoxPyramidStairsTerrainCfg`` with
 ``step_height_range=(0.0, 0.2)`` produces flat ground at difficulty 0
-and 20 cm steps at difficulty 1. In curriculum mode, difficulty is
-determined by the row: row 0 gets the minimum, row ``num_rows - 1`` gets
-the maximum.
+and 20 cm steps at difficulty 1.
+
+In curriculum mode, difficulty is determined by the row:
+``difficulty = lower + (upper - lower) * row / max(num_rows - 1, 1)``,
+where ``(lower, upper) = difficulty_range``. Row 0 is exactly
+``lower``, row ``num_rows - 1`` is exactly ``upper``, and intermediate
+rows are evenly spaced between them. All columns in a given row share
+the same difficulty scalar; the visible variation across columns comes
+from each sub-terrain type generating different geometry at the same
+difficulty.
+
+.. note::
+
+   With ``num_rows=1`` and ``curriculum=True``, every patch is generated
+   at ``difficulty = lower`` (the easiest configured difficulty). Use
+   ``curriculum=False`` if you want a single grid of randomly sampled
+   difficulties instead.
+
+In random mode, difficulty is sampled uniformly from
+``difficulty_range`` independently for every patch.


 Sub-terrain types
@@ -244,17 +263,23 @@ and undulating ground that box geoms cannot represent.
 Preset configurations
 ---------------------

-mjlab ships two ready-made ``TerrainGeneratorCfg`` presets in
+mjlab ships three ready-made ``TerrainGeneratorCfg`` presets in
 ``mjlab.terrains.config``:

 ``ROUGH_TERRAINS_CFG``
-    A 10x20 grid with seven terrain types (flat, stairs, inverted
-    stairs, slopes, inverted slopes, random rough, waves). Designed for
-    locomotion training with a moderate difficulty range.
+    A 10x20 random-mode grid with seven terrain types (flat, stairs,
+    inverted stairs, slopes, inverted slopes, random rough, waves).
+    Designed for locomotion training with a moderate difficulty range.
+    Set ``curriculum=True`` via ``dataclasses.replace`` to use it as a
+    curriculum grid (one column per terrain type).
+
+``STAIRS_TERRAINS_CFG``
+    A 10-row curriculum grid focused on stair traversal: flat plus
+    three pyramid-stair variants of increasing difficulty.

 ``ALL_TERRAINS_CFG``
-    A 10x16 grid with all sixteen terrain types at equal proportion.
-    Useful for training on maximum terrain variety.
+    A 10-row random-mode grid covering all available terrain types at
+    equal proportion. Useful for training on maximum terrain variety.

 Both can be used directly or customized with ``dataclasses.replace()``:

@@ -285,9 +310,9 @@ The key concepts:
 - The built-in ``terrain_levels_vel`` curriculum term promotes
  environments that track commanded velocity well and demotes
  environments that fall or fail to make progress.
- When an environment reaches the maximum row, it is randomly reassigned
-  to a lower row to prevent the policy from collapsing to a single
-  difficulty level.
+- When an environment is promoted past the hardest row, it is randomly
+  reassigned to any row in ``[0, num_rows)`` to prevent the policy from
+  collapsing to a single difficulty level.


 Flat patch detection
@@ -4,7 +4,7 @@ build-backend = "uv_build"

 [project]
 name = "mjlab"
-version = "1.3.0"
+version = "1.4.0"
 license = "Apache-2.0"
 license-files = ["LICENSE"]
 readme = { file = "README.md", content-type = "text/markdown" }
@@ -37,15 +37,16 @@ dependencies = [
  "torch>=2.7.0",
  "torchrunx>=0.3.4",
  "warp-lang>=1.12.0",
-  "mujoco-warp>=3.8.0",
-  "mujoco>=3.8.0",
+  "mujoco-warp>=3.8.0.3,~=3.8.0",
+  "mujoco~=3.8.0",
  "trimesh>=4.8.3",
-  "viser>=1.0.26",
-  "mjviser>=0.0.13",
+  "scipy>=1.15",
+  "viser>=1.0.27",
+  "mjviser>=0.0.14",
  "mediapy>=1.2.6",
  "imageio-ffmpeg",
  "tensordict",
-  "rsl-rl-lib==5.2.0",
+  "rsl-rl-lib==5.4.0",
  "tensorboard>=2.20.0",
  "onnxscript>=0.5.4",
  "wandb>=0.22.3",
@@ -98,11 +99,11 @@ conflicts = [
  [{extra = "cu128"}, {extra = "cpu"}],
 ]
 # The nightly index (py.mujoco.org) only has dev builds, and PEP 440 ranks
-# 3.7.0.devN < 3.7.0, so the >=3.7.0 floor in [project.dependencies] would
+# 3.8.0.devN < 3.8.0, so the ~=3.8.0 floor in [project.dependencies] would
 # reject them. This override loosens the constraint for uv resolution only.
 override-dependencies = ["mujoco>=3.8.0.dev0"]
 constraint-dependencies = [
-  "GitPython>=3.1.47",
+  "GitPython>=3.1.49",
  "lxml>=6.1.0",
 ]
 required-environments = [
@@ -111,12 +112,6 @@ required-environments = [
 ]

 [[tool.uv.index]]
-name = "tsinghua"
-url = "https://pypi.tuna.tsinghua.edu.cn/simple"
-default = true
-
-[[tool.uv.index]]
-name = "pypi"
 url = "https://pypi.org/simple"

 [[tool.uv.index]]
@@ -146,8 +141,7 @@ torch = [
  { index = "pytorch-cpu", extra = "cpu", marker = "sys_platform != 'darwin'" },
 ]
 mujoco = { index = "mujoco" }
-mujoco-warp = { git = "https://github.com/google-deepmind/mujoco_warp", rev = "6f235d4" }
-mjviser = { git = "https://github.com/mujocolab/mjviser", rev = "1bdfd6fe79066b847a5f430000fcfbb53ec31a6f" }
+mujoco-warp = { git = "https://github.com/google-deepmind/mujoco_warp", rev = "88b55fc2696960b927bc12584994bb8412b36558" }

 [tool.ruff]
 src = ["src"]  # Helpful for recognizing first-party imports.
@@ -215,6 +215,27 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
            border-color: var(--accent);
            color: white;
        }}
+        .range-selector {{
+            display: flex;
+            gap: 0.4rem;
+            margin-bottom: 1rem;
+        }}
+        .range-btn {{
+            background: var(--bg-card);
+            border: 1px solid var(--border);
+            border-radius: 6px;
+            padding: 0.3rem 0.75rem;
+            cursor: pointer;
+            color: var(--text);
+            font-size: 0.8rem;
+            font-weight: 500;
+        }}
+        .range-btn:hover {{ border-color: var(--accent); }}
+        .range-btn.active {{
+            background: var(--accent);
+            border-color: var(--accent);
+            color: white;
+        }}
        .tab-content {{ display: none; }}
        .tab-content.active {{ display: block; }}
        .tab-description {{
@@ -303,12 +324,24 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st

    <div id="tracking" class="tab-content active">
        <p class="tab-description">Nightly motion imitation training and evaluation on Unitree G1 (1024 trials per run).</p>
+        <div class="range-selector" id="range-selector">
+            <button class="range-btn" data-days="30">30d</button>
+            <button class="range-btn active" data-days="90">90d</button>
+            <button class="range-btn" data-days="180">180d</button>
+            <button class="range-btn" data-days="0">All</button>
+        </div>
        <div class="charts" id="charts"></div>
    </div>

    <div id="throughput" class="tab-content">
        <p class="tab-description">Physics simulation throughput across tasks (4096 parallel envs, NVIDIA RTX 5090).</p>
        <div class="task-grid" id="task-grid"></div>
+        <div class="range-selector" id="range-selector-tp">
+            <button class="range-btn" data-days="30">30d</button>
+            <button class="range-btn active" data-days="90">90d</button>
+            <button class="range-btn" data-days="180">180d</button>
+            <button class="range-btn" data-days="0">All</button>
+        </div>
        <div id="task-chart-panels"></div>
    </div>

@@ -393,6 +426,8 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
        }};

        let charts = [];
+        let trackingCharts = [];
+        let throughputCharts = [];

        function updateChartColors() {{
            const style = getComputedStyle(root);
@@ -449,7 +484,7 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
            `;
            chartsContainer.appendChild(card);

-            charts.push(new Chart(card.querySelector('canvas'), {{
+            const chart = new Chart(card.querySelector('canvas'), {{
                type: 'line',
                data: {{
                    datasets: [
@@ -459,7 +494,8 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
                            borderColor: color,
                            backgroundColor: color + '20',
                            borderWidth: 2,
-                            pointRadius: 4,
+                            pointRadius: 2,
+                            pointHoverRadius: 5,
                            tension: 0.1,
                            fill: true
                        }},
@@ -531,7 +567,9 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
                        }}
                    }}
                }}
-            }}));
+            }});
+            charts.push(chart);
+            trackingCharts.push(chart);
        }});

        // Tab switching
@@ -621,7 +659,8 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
                                borderColor: '#58a6ff',
                                backgroundColor: '#58a6ff20',
                                borderWidth: 2,
-                                pointRadius: 4,
+                                pointRadius: 2,
+                                pointHoverRadius: 5,
                                tension: 0.1,
                                fill: true
                            }},
@@ -631,7 +670,8 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
                                borderColor: '#3fb950',
                                backgroundColor: '#3fb95020',
                                borderWidth: 2,
-                                pointRadius: 4,
+                                pointRadius: 2,
+                                pointHoverRadius: 5,
                                tension: 0.1,
                                fill: true
                            }}
@@ -697,6 +737,7 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
                    }}
                }});
                charts.push(chart);
+                throughputCharts.push(chart);
                throughputChartInstances[task] = {{ chart, panelId: `task-panel-${{i}}` }};

                // Card click handler
@@ -712,6 +753,23 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
        }} else {{
            taskGrid.innerHTML = '<p style="color: var(--text-dim)">No throughput data available. Run measure_throughput.py to generate data.</p>';
        }}
+
+        // Date-range windowing across both tracking and throughput charts.
+        // Setting min and clearing max also resets any zoom/pan.
+        function setRange(days) {{
+            const min = days > 0 ? Date.now() - days * 86400000 : undefined;
+            [...trackingCharts, ...throughputCharts].forEach(c => {{
+                c.options.scales.x.min = min;
+                c.options.scales.x.max = undefined;
+                c.update();
+            }});
+            document.querySelectorAll('.range-btn').forEach(b =>
+                b.classList.toggle('active', parseInt(b.dataset.days) === days));
+        }}
+        document.querySelectorAll('.range-btn').forEach(btn => {{
+            btn.addEventListener('click', () => setRange(parseInt(btn.dataset.days)));
+        }});
+        setRange(90);
    </script>
 </body>
 </html>
@@ -764,7 +822,11 @@ def main(
      if run_id in eval_results_by_id:
        print(f"Using cached result for {run_id}")
      else:
-        result = evaluate_run(run_path, num_envs)
+        try:
+          result = evaluate_run(run_path, num_envs)
+        except RuntimeError as e:
+          print(f"Skipping {run_path}: {e}")
+          continue
        eval_results_by_id[run_id] = result
        new_evals += 1
  else:
@@ -783,7 +845,11 @@ def main(
          print(f"Reached eval limit ({eval_limit}), skipping remaining new runs")
          break
        run_path = f"{entity}/{project}/{run.id}"
-        result = evaluate_run(run_path, num_envs)
+        try:
+          result = evaluate_run(run_path, num_envs)
+        except RuntimeError as e:
+          print(f"Skipping {run.name} ({run.id}): {e}")
+          continue
        eval_results_by_id[run.id] = result
        new_evals += 1

@@ -0,0 +1,115 @@
+"""Interactive single-patch terrain explorer (Viser + MuJoCo MjSpec).
+
+Run with:
+  uv run python scripts/tools/terrain_explorer.py
+  uv run python scripts/tools/terrain_explorer.py --port 8081
+
+Then open the printed URL (default http://localhost:8080).
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+
+import mujoco
+import numpy as np
+import viser
+from mjviser.conversions import merge_geoms
+
+from mjlab.terrains.config import ALL_TERRAIN_PRESETS
+from mjlab.terrains.terrain_generator import TerrainGenerator, TerrainGeneratorCfg
+
+PATCH_SIZE = (8.0, 8.0)
+
+
+# Per-preset overrides applied when building in the explorer (e.g. to surface
+# difficulty-driven behavior that is off by default).
+_PRESET_OVERRIDES: dict[str, dict] = {
+  "random_rough": {"scale_with_difficulty": True},
+}
+
+
+def _build_terrain_mesh(preset_name: str, difficulty: float, seed: int):
+  """Generate a single terrain patch and return a merged trimesh (or raise)."""
+  preset_fn = ALL_TERRAIN_PRESETS[preset_name]
+  overrides = _PRESET_OVERRIDES.get(preset_name, {})
+  generator_cfg = TerrainGeneratorCfg(
+    seed=seed,
+    size=PATCH_SIZE,
+    num_rows=1,
+    num_cols=1,
+    border_width=0.0,
+    curriculum=False,
+    # A degenerate range pins the single patch to exactly this difficulty.
+    difficulty_range=(difficulty, difficulty),
+    color_scheme="height",
+    sub_terrains={preset_name: preset_fn(proportion=1.0, **overrides)},
+  )
+  generator = TerrainGenerator(generator_cfg)
+  spec = mujoco.MjSpec()
+  generator.compile(spec)
+  model = spec.compile()
+
+  terrain_body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "terrain")
+  geom_ids = [i for i in range(model.ngeom) if model.geom_bodyid[i] == terrain_body_id]
+  return merge_geoms(model, geom_ids)
+
+
+def main() -> None:
+  parser = argparse.ArgumentParser(description=__doc__)
+  parser.add_argument(
+    "--port", type=int, default=8080, help="Port for the viser server."
+  )
+  args = parser.parse_args()
+
+  server = viser.ViserServer(port=args.port)
+  preset_names = sorted(ALL_TERRAIN_PRESETS)
+
+  terrain_dropdown = server.gui.add_dropdown(
+    "Terrain", options=preset_names, initial_value=preset_names[0]
+  )
+  difficulty_slider = server.gui.add_slider(
+    "Difficulty", min=0.0, max=1.0, step=0.01, initial_value=0.0
+  )
+  seed_input = server.gui.add_number("Seed", initial_value=42, step=1)
+  status = server.gui.add_markdown("**Status:** ready")
+
+  handle: viser.SceneNodeHandle | None = None
+
+  def update() -> None:
+    nonlocal handle
+    name = terrain_dropdown.value
+    difficulty = float(difficulty_slider.value)
+    seed = int(seed_input.value)
+    status.content = f"**Status:** building `{name}` at difficulty {difficulty:.2f}..."
+    try:
+      mesh = _build_terrain_mesh(name, difficulty, seed)
+    except Exception as e:  # noqa: BLE001 - surface any generation failure in the UI.
+      status.content = f"**Error:** {type(e).__name__}: {e}"
+      print(f"Failed to build {name} at difficulty {difficulty}: {e}")
+      return
+    if handle is not None:
+      handle.remove()
+    handle = server.scene.add_mesh_trimesh("/terrain", mesh)
+    status.content = (
+      f"**Loaded** `{name}` at difficulty {difficulty:.2f} ({len(mesh.faces):,} faces)"
+    )
+
+  terrain_dropdown.on_update(lambda _: update())
+  difficulty_slider.on_update(lambda _: update())
+  seed_input.on_update(lambda _: update())
+
+  # Top-down-ish initial camera.
+  @server.on_client_connect
+  def _(client: viser.ClientHandle) -> None:
+    client.camera.position = np.array([10.0, 10.0, 8.0])
+    client.camera.look_at = np.array([0.0, 0.0, 0.0])
+
+  update()
+  while True:
+    time.sleep(1.0)
+
+
+if __name__ == "__main__":
+  main()
@@ -1,5 +1,15 @@
 import os
 import sys
+
+# Default to EGL for GPU-accelerated offscreen rendering on Linux. Must be set
+# before any mujoco import: mujoco's gl_context module captures MUJOCO_GL once
+# at load time. Override with e.g. MUJOCO_GL=osmesa on clusters without EGL.
+# Linux-only because mujoco's gl_context rejects "egl" on macOS/Windows and
+# raises at import. On those platforms we leave MUJOCO_GL alone so mujoco
+# defaults to GLFW.
+if sys.platform.startswith("linux"):
+  os.environ.setdefault("MUJOCO_GL", "egl")
+
 import traceback
 from importlib.metadata import entry_points
 from pathlib import Path
@@ -4,6 +4,12 @@ from mjlab.actuator.actuator import Actuator as Actuator
 from mjlab.actuator.actuator import ActuatorCfg as ActuatorCfg
 from mjlab.actuator.actuator import ActuatorCmd as ActuatorCmd
 from mjlab.actuator.actuator import CommandField as CommandField
+from mjlab.actuator.builtin_actuator import (
+  BuiltinDcMotorActuator as BuiltinDcMotorActuator,
+)
+from mjlab.actuator.builtin_actuator import (
+  BuiltinDcMotorActuatorCfg as BuiltinDcMotorActuatorCfg,
+)
 from mjlab.actuator.builtin_actuator import (
  BuiltinMotorActuator as BuiltinMotorActuator,
 )
@@ -16,6 +22,12 @@ from mjlab.actuator.builtin_actuator import (
 from mjlab.actuator.builtin_actuator import (
  BuiltinMuscleActuatorCfg as BuiltinMuscleActuatorCfg,
 )
+from mjlab.actuator.builtin_actuator import (
+  BuiltinPdActuator as BuiltinPdActuator,
+)
+from mjlab.actuator.builtin_actuator import (
+  BuiltinPdActuatorCfg as BuiltinPdActuatorCfg,
+)
 from mjlab.actuator.builtin_actuator import (
  BuiltinPositionActuator as BuiltinPositionActuator,
 )
@@ -28,6 +40,15 @@ from mjlab.actuator.builtin_actuator import (
 from mjlab.actuator.builtin_actuator import (
  BuiltinVelocityActuatorCfg as BuiltinVelocityActuatorCfg,
 )
+from mjlab.actuator.builtin_actuator import (
+  DcMotorDatasheetParams as DcMotorDatasheetParams,
+)
+from mjlab.actuator.builtin_actuator import (
+  DcMotorInputMode as DcMotorInputMode,
+)
+from mjlab.actuator.builtin_actuator import (
+  DcMotorPhysicalParams as DcMotorPhysicalParams,
+)
 from mjlab.actuator.builtin_group import BuiltinActuatorGroup as BuiltinActuatorGroup
 from mjlab.actuator.dc_actuator import DcMotorActuator as DcMotorActuator
 from mjlab.actuator.dc_actuator import DcMotorActuatorCfg as DcMotorActuatorCfg
@@ -174,15 +174,6 @@ class Actuator(ABC, Generic[ActuatorCfgT]):
    """Whether this actuator has delay configured."""
    return self.cfg.delay_max_lag > 0

-  @property
-  def command_field(self) -> CommandField | None:
-    """The primary command field this actuator consumes.
-
-    Returns None by default. Subclasses should override to return the
-    appropriate field.
-    """
-    return None
-
  @property
  def target_ids(self) -> torch.Tensor:
    """Local indices of targets controlled by this actuator."""
@@ -271,11 +262,6 @@ class Actuator(ABC, Generic[ActuatorCfgT]):
    """Create delay buffer. Called during initialize()."""
    if not self.has_delay:
      return
-    if self.command_field is None:
-      raise ValueError(
-        f"{self.__class__.__name__}: delay is configured (delay_max_lag="
-        f"{self.cfg.delay_max_lag}) but command_field is not defined."
-      )
    self._delay_buffer = DelayBuffer(
      min_lag=self.cfg.delay_min_lag,
      max_lag=self.cfg.delay_max_lag,
@@ -287,19 +273,25 @@ class Actuator(ABC, Generic[ActuatorCfgT]):
    )

  def apply_delay(self, cmd: ActuatorCmd) -> ActuatorCmd:
-    """Apply delay to the command_field target. No-op without delay."""
+    """Delay all command targets with one shared lag. No-op without delay.
+
+    Every target the policy issues (position, velocity, effort) travels the same
+    command channel and experiences the same latency, so they are stacked and
+    delayed together. Feedback fields (``pos``, ``vel``) are never delayed.
+    """
    if self._delay_buffer is None:
      return cmd
-    cf = self.command_field
-    if cf == "position":
-      self._delay_buffer.append(cmd.position_target)
-      return dataclasses.replace(cmd, position_target=self._delay_buffer.compute())
-    elif cf == "velocity":
-      self._delay_buffer.append(cmd.velocity_target)
-      return dataclasses.replace(cmd, velocity_target=self._delay_buffer.compute())
-    else:
-      self._delay_buffer.append(cmd.effort_target)
-      return dataclasses.replace(cmd, effort_target=self._delay_buffer.compute())
+    targets = torch.stack(
+      (cmd.position_target, cmd.velocity_target, cmd.effort_target), dim=-1
+    )
+    self._delay_buffer.append(targets)
+    delayed = self._delay_buffer.compute()
+    return dataclasses.replace(
+      cmd,
+      position_target=delayed[..., 0],
+      velocity_target=delayed[..., 1],
+      effort_target=delayed[..., 2],
+    )

  def set_lags(
    self,
@@ -7,19 +7,21 @@ created programmatically via the MjSpec API.
 from __future__ import annotations

 from dataclasses import dataclass
+from enum import IntEnum
 from typing import TYPE_CHECKING

 import mujoco
+import numpy as np
 import torch

 from mjlab.actuator.actuator import (
  Actuator,
  ActuatorCfg,
  ActuatorCmd,
-  CommandField,
  TransmissionType,
 )
 from mjlab.utils.spec import (
+  apply_target_overrides,
  create_motor_actuator,
  create_muscle_actuator,
  create_position_actuator,
@@ -63,10 +65,6 @@ class BuiltinPositionActuatorCfg(ActuatorCfg):
 class BuiltinPositionActuator(Actuator[BuiltinPositionActuatorCfg]):
  """MuJoCo built-in position actuator."""

-  @property
-  def command_field(self) -> CommandField:
-    return "position"
-
  def __init__(
    self,
    cfg: BuiltinPositionActuatorCfg,
@@ -96,6 +94,102 @@ class BuiltinPositionActuator(Actuator[BuiltinPositionActuatorCfg]):
    return cmd.position_target


+@dataclass(kw_only=True)
+class BuiltinPdActuatorCfg(ActuatorCfg):
+  """Implicit-integration version of IdealPdActuator.
+
+  Both consume a position target and a velocity target with kp/kd gains. The
+  difference is in how the PD is delivered to MuJoCo: IdealPdActuator computes
+  the PD force in Python and feeds it to a ``<motor>`` element, which MuJoCo
+  sees as an opaque external force. This actuator expresses the PD as native
+  MuJoCo elements (a ``<position>`` carrying kp, a ``<velocity>`` carrying kd),
+  so the implicit and implicitfast integrators include the kp/kd derivatives
+  in their velocity update. That makes the actuator numerically stable at
+  gain/timestep combinations where explicit Python PD would diverge, which
+  matters when you want to run a real motor's stiff on-board PD gains in sim.
+  """
+
+  stiffness: float
+  """Proportional gain (kp)."""
+  damping: float
+  """Derivative gain (kd)."""
+  effort_limit: float | None = None
+  """Maximum total torque applied to the joint or tendon. Enforced as a
+  sum-clamp on the two PD terms via jnt_actfrcrange (JOINT) or
+  tendon_actfrcrange (TENDON). None leaves the limit unset."""
+
+  def __post_init__(self) -> None:
+    super().__post_init__()
+    if self.transmission_type == TransmissionType.SITE:
+      raise ValueError(
+        "BuiltinPdActuatorCfg does not support SITE transmission. "
+        "Use BuiltinMotorActuatorCfg for site transmission."
+      )
+
+  def build(
+    self, entity: Entity, target_ids: list[int], target_names: list[str]
+  ) -> BuiltinPdActuator:
+    return BuiltinPdActuator(self, entity, target_ids, target_names)
+
+
+class BuiltinPdActuator(Actuator[BuiltinPdActuatorCfg]):
+  """MuJoCo native PD: paired <position> + <velocity> elements per target."""
+
+  def __init__(
+    self,
+    cfg: BuiltinPdActuatorCfg,
+    entity: Entity,
+    target_ids: list[int],
+    target_names: list[str],
+  ) -> None:
+    super().__init__(cfg, entity, target_ids, target_names)
+
+  @property
+  def num_targets(self) -> int:
+    """Number of targets. ``ctrl_ids`` is laid out as ``[pos..., vel...]``,
+    each block of length ``num_targets``."""
+    return len(self._target_ids_list)
+
+  def edit_spec(self, spec: mujoco.MjSpec, target_names: list[str]) -> None:
+    # Position elements first, then velocity elements, so ctrl_ids is laid out
+    # as [pos_0..pos_{N-1}, vel_0..vel_{N-1}].
+    for target_name in target_names:
+      pos_act = create_position_actuator(
+        spec,
+        target_name,
+        actuator_name=f"{target_name}_pd_pos",
+        stiffness=self.cfg.stiffness,
+        damping=0.0,  # damping lives on the <velocity> element.
+        armature=self.cfg.armature,
+        frictionloss=self.cfg.frictionloss,
+        viscous_damping=self.cfg.viscous_damping,
+        transmission_type=self.cfg.transmission_type,
+      )
+      self._mjs_actuators.append(pos_act)
+    for target_name in target_names:
+      vel_act = create_velocity_actuator(
+        spec,
+        target_name,
+        actuator_name=f"{target_name}_pd_vel",
+        damping=self.cfg.damping,
+        transmission_type=self.cfg.transmission_type,
+      )
+      self._mjs_actuators.append(vel_act)
+    # Effort limit: sum-clamp on the joint/tendon, not on each element.
+    if self.cfg.effort_limit is not None:
+      lim = self.cfg.effort_limit
+      for target_name in target_names:
+        if self.cfg.transmission_type == TransmissionType.JOINT:
+          target = spec.joint(target_name)
+        else:
+          target = spec.tendon(target_name)
+        target.actfrclimited = mujoco.mjtLimited.mjLIMITED_TRUE
+        target.actfrcrange[:] = np.array([-lim, lim])
+
+  def compute(self, cmd: ActuatorCmd) -> torch.Tensor:
+    return torch.cat((cmd.position_target, cmd.velocity_target), dim=1)
+
+
@dataclass(kw_only=True)
 class BuiltinMotorActuatorCfg(ActuatorCfg):
  """Configuration for MuJoCo built-in motor actuator.
@@ -119,10 +213,6 @@ class BuiltinMotorActuatorCfg(ActuatorCfg):
 class BuiltinMotorActuator(Actuator[BuiltinMotorActuatorCfg]):
  """MuJoCo built-in motor actuator."""

-  @property
-  def command_field(self) -> CommandField:
-    return "effort"
-
  def __init__(
    self,
    cfg: BuiltinMotorActuatorCfg,
@@ -151,6 +241,248 @@ class BuiltinMotorActuator(Actuator[BuiltinMotorActuatorCfg]):
    return cmd.effort_target


+def _or_zeros(t: tuple[float, ...] | None, n: int) -> list[float]:
+  return list(t) if t is not None else [0.0] * n
+
+
+class DcMotorInputMode(IntEnum):
+  """What the ``ctrl`` signal of a ``<dcmotor>`` represents.
+
+  Values match MuJoCo's enum, consumed by mjs_setToDCMotor and read as gainprm[8].
+  """
+
+  VOLTAGE = 0
+  POSITION = 1
+  VELOCITY = 2
+
+
+@dataclass(frozen=True)
+class DcMotorDatasheetParams:
+  """Datasheet characterization of a DC motor."""
+
+  nominal_voltage: float
+  """Nominal (rated) voltage V_n [V]."""
+  stall_torque: float
+  """Stall torque tau_stall at V_n [N*m]."""
+  no_load_speed: float
+  """No-load angular velocity omega_no_load at V_n [rad/s]."""
+
+  def _pack(self) -> tuple[list[float], float, list[float]]:
+    """Returns (motorconst, resistance, nominal) for set_to_dcmotor."""
+    return (
+      [0.0, 0.0],
+      0.0,
+      [self.nominal_voltage, self.stall_torque, self.no_load_speed],
+    )
+
+
+@dataclass(frozen=True)
+class DcMotorPhysicalParams:
+  """Physical characterization of a DC motor."""
+
+  kt: float
+  """Torque constant [N*m/A]."""
+  ke: float
+  """Back-EMF constant [V*s/rad]."""
+  resistance: float
+  """Terminal resistance R [Ohm]."""
+
+  def _pack(self) -> tuple[list[float], float, list[float]]:
+    """Returns (motorconst, resistance, nominal) for set_to_dcmotor."""
+    return [self.kt, self.ke], self.resistance, [0.0, 0.0, 0.0]
+
+
+@dataclass(kw_only=True)
+class BuiltinDcMotorActuatorCfg(ActuatorCfg):
+  """Native MuJoCo ``<dcmotor>`` wrapper.
+
+  Models a DC motor: torque is derived from voltage via the motor constant K and
+  back-EMF, tau = K * (V - K * omega) / R. The back-EMF term lives in biasprm, so
+  MuJoCo's implicit / implicitfast integrators pick up its velocity derivative as
+  effective damping.
+
+  Three input modes select what ctrl carries:
+
+  * VOLTAGE: ctrl is the drive voltage. cmd.effort_target carries volts, not torque.
+  * POSITION / VELOCITY: an internal PID closes on the setpoint and the motor produces
+    torque from its (Vmax-clamped) voltage output.
+
+  Motor characterization: pass either DcMotorDatasheetParams or DcMotorPhysicalParams
+  as motor_params. mjs_setToDCMotor derives K and R (including the viscous-damping
+  correction) and packs the generic gainprm / biasprm / dynprm slots.
+
+  Optional extensions, off by default: integral_gain / integral_limit, slew_rate,
+  inductance / electrical_time_constant, thermal, lugre, cogging.
+
+  dr.pd_gains randomizes only kp and kd; for DR over the extensions, write directly to
+  actuator_gainprm or actuator_dynprm.
+  """
+
+  motor_params: DcMotorDatasheetParams | DcMotorPhysicalParams
+  """Motor characterization. Datasheet form: (V_n, tau_stall, omega_no_load).
+  Physical form: (Kt, Ke, R)."""
+
+  mode: DcMotorInputMode = DcMotorInputMode.POSITION
+  """ctrl input semantics. See class docstring."""
+
+  stiffness: float = 0.0
+  """PID proportional gain kp. Required in POSITION / VELOCITY mode; must be
+  0 in VOLTAGE mode."""
+
+  damping: float = 0.0
+  """PID derivative gain kd. Used in POSITION / VELOCITY mode; must be 0 in
+  VOLTAGE mode."""
+
+  voltage_limit: float = 0.0
+  """Max drive voltage Vmax. Required in POSITION / VELOCITY mode (clamps the
+  PID output). In VOLTAGE mode it is an optional clamp on ctrl; 0 disables."""
+
+  integral_gain: float = 0.0
+  """PID integral gain ki. In position mode the integrator tracks
+  ki * integral(target - q); in velocity mode, ki * (integral(target) - q).
+  Must be 0 in VOLTAGE mode."""
+
+  integral_limit: float = 0.0
+  """Anti-windup clamp Imax on the integrator state. 0 disables (the
+  integrator can run away)."""
+
+  slew_rate: float = 0.0
+  """Max rate of change of ctrl per second. 0 disables."""
+
+  effort_limit: float | None = None
+  """Continuous torque cap [N*m]. Sets actuator_forcerange. None leaves the
+  per-element forcerange unset."""
+
+  gear: float = 1.0
+  """Mechanical gear ratio."""
+
+  inductance: float = 0.0
+  """Winding inductance L [H]. Enables first-order electrical dynamics on the
+  motor current. MuJoCo internally uses te = L / R; pass
+  electrical_time_constant directly to skip the divide. 0 disables."""
+
+  electrical_time_constant: float = 0.0
+  """Alternative to inductance: specify te [s] directly. Ignored if
+  inductance > 0. 0 disables."""
+
+  thermal: tuple[float, float, float, float, float, float] | None = None
+  """Thermal model (R_thermal, C_thermal, tau_thermal, alpha, T0, T_ambient).
+  See MuJoCo's ``<dcmotor thermal=...>`` reference for units and which of the
+  first three may be underspecified. Effective resistance becomes
+  R * (1 + alpha * (T + T_ambient - T0)). None disables."""
+
+  cogging: tuple[float, float, float] | None = None
+  """Cogging torque (amplitude, periodicity, phase) in (N*m, cycles per unit
+  length, rad). Models magnetic torque ripple from rotor-stator interaction;
+  at joint angle q the contribution is amplitude * sin(periodicity * q + phase).
+
+  Added *after* effort_limit is enforced, matching MuJoCo's physical model:
+  effort_limit bounds the electromagnetic torque (the current limit), not the
+  mechanical torque. Total joint torque can exceed effort_limit by up to
+  amplitude. None disables."""
+
+  lugre: tuple[float, float, float, float, float] | None = None
+  """LuGre friction (sigma0, sigma1, F_Coulomb, F_Stribeck, v_Stribeck).
+  Stick-slip friction with bristle-deflection state. Subtracted from joint
+  torque after the effort_limit clamp (mechanical, like cogging). None
+  disables."""
+
+  def __post_init__(self) -> None:
+    super().__post_init__()
+    if self.transmission_type == TransmissionType.SITE:
+      raise ValueError(
+        "BuiltinDcMotorActuatorCfg does not support SITE transmission. "
+        "Use BuiltinMotorActuatorCfg for site transmission."
+      )
+
+    if self.mode in (DcMotorInputMode.POSITION, DcMotorInputMode.VELOCITY):
+      if self.stiffness <= 0.0:
+        raise ValueError(f"{self.mode.name} mode requires stiffness > 0.")
+      if self.voltage_limit <= 0.0:
+        raise ValueError(f"{self.mode.name} mode requires voltage_limit > 0.")
+    else:
+      if self.stiffness != 0.0 or self.damping != 0.0 or self.integral_gain != 0.0:
+        raise ValueError(
+          "stiffness, damping, and integral_gain are unused in VOLTAGE mode."
+        )
+
+    for name in (
+      "integral_gain",
+      "integral_limit",
+      "slew_rate",
+      "inductance",
+      "electrical_time_constant",
+    ):
+      if getattr(self, name) < 0.0:
+        raise ValueError(f"{name} must be non-negative.")
+
+  def build(
+    self, entity: Entity, target_ids: list[int], target_names: list[str]
+  ) -> BuiltinDcMotorActuator:
+    return BuiltinDcMotorActuator(self, entity, target_ids, target_names)
+
+
+class BuiltinDcMotorActuator(Actuator[BuiltinDcMotorActuatorCfg]):
+  """MuJoCo native ``<dcmotor>``: one actuator per target."""
+
+  def edit_spec(self, spec: mujoco.MjSpec, target_names: list[str]) -> None:
+    cfg = self.cfg
+    motorconst, resistance, nominal = cfg.motor_params._pack()
+    saturation = (
+      [cfg.effort_limit, 0.0, 0.0] if cfg.effort_limit is not None else [0.0] * 3
+    )
+    controller = [
+      cfg.stiffness,  # kp
+      cfg.integral_gain,  # ki
+      cfg.damping,  # kd
+      cfg.slew_rate,  # slewmax
+      cfg.integral_limit,  # Imax (anti-windup)
+      cfg.voltage_limit,  # v_max
+    ]
+    # SITE is rejected in __post_init__, so only JOINT and TENDON remain.
+    trntype = (
+      mujoco.mjtTrn.mjTRN_JOINT
+      if cfg.transmission_type == TransmissionType.JOINT
+      else mujoco.mjtTrn.mjTRN_TENDON
+    )
+
+    for target_name in target_names:
+      actuator = spec.add_actuator(name=target_name, target=target_name)
+      actuator.trntype = trntype
+      actuator.gear[0] = cfg.gear
+      actuator.set_to_dcmotor(
+        motorconst=motorconst,
+        resistance=resistance,
+        nominal=nominal,
+        saturation=saturation,
+        controller=controller,
+        cogging=_or_zeros(cfg.cogging, 3),
+        inductance=[cfg.inductance, cfg.electrical_time_constant],
+        thermal=_or_zeros(cfg.thermal, 6),
+        lugre=_or_zeros(cfg.lugre, 5),
+        input_mode=cfg.mode,
+      )
+
+      apply_target_overrides(
+        spec,
+        target_name,
+        cfg.transmission_type,
+        armature=cfg.armature,
+        frictionloss=cfg.frictionloss,
+        viscous_damping=cfg.viscous_damping,
+      )
+
+      self._mjs_actuators.append(actuator)
+
+  def compute(self, cmd: ActuatorCmd) -> torch.Tensor:
+    if self.cfg.mode == DcMotorInputMode.POSITION:
+      return cmd.position_target
+    if self.cfg.mode == DcMotorInputMode.VELOCITY:
+      return cmd.velocity_target
+    # voltage mode: ctrl is the drive voltage carried in effort_target.
+    return cmd.effort_target
+
+
@dataclass(kw_only=True)
 class BuiltinVelocityActuatorCfg(ActuatorCfg):
  """Configuration for MuJoCo built-in velocity actuator.
@@ -182,10 +514,6 @@ class BuiltinVelocityActuatorCfg(ActuatorCfg):
 class BuiltinVelocityActuator(Actuator[BuiltinVelocityActuatorCfg]):
  """MuJoCo built-in velocity actuator."""

-  @property
-  def command_field(self) -> CommandField:
-    return "velocity"
-
  def __init__(
    self,
    cfg: BuiltinVelocityActuatorCfg,
@@ -260,10 +588,6 @@ class BuiltinMuscleActuatorCfg(ActuatorCfg):
 class BuiltinMuscleActuator(Actuator[BuiltinMuscleActuatorCfg]):
  """MuJoCo built-in muscle actuator."""

-  @property
-  def command_field(self) -> CommandField:
-    return "effort"
-
  def __init__(
    self,
    cfg: BuiltinMuscleActuatorCfg,
@@ -33,6 +33,9 @@ class DcMotorActuatorCfg(IdealPdActuatorCfg):
  Note: effort_limit should be explicitly set to a realistic value for proper
  motor modeling. Using the default (inf) will trigger a warning. Use
  IdealPdActuator if unlimited torque is desired.
+
+  For a native MuJoCo ``<dcmotor>`` with back-EMF, voltage saturation, and
+  configurable ``Kt`` / ``Ke`` / ``R``, see ``BuiltinDcMotorActuator``.
  """

  saturation_effort: float
@@ -9,7 +9,7 @@ import mujoco
 import mujoco_warp as mjwarp
 import torch

-from mjlab.actuator.actuator import Actuator, ActuatorCfg, ActuatorCmd, CommandField
+from mjlab.actuator.actuator import Actuator, ActuatorCfg, ActuatorCmd
 from mjlab.utils.spec import create_motor_actuator

 if TYPE_CHECKING:
@@ -38,10 +38,6 @@ class IdealPdActuatorCfg(ActuatorCfg):
 class IdealPdActuator(Actuator, Generic[IdealPdCfgT]):
  """Ideal PD control actuator."""

-  @property
-  def command_field(self) -> CommandField:
-    return "position"
-
  def __init__(
    self,
    cfg: IdealPdCfgT,
@@ -102,11 +102,6 @@
                    diaginertia="0.00167218 0.0016161 0.000217621"/>
                  <joint name="left_ankle_roll_joint" axis="1 0 0" range="-0.2618 0.2618"/>
                  <geom class="visual" material="black" mesh="left_ankle_roll_link"/>
-                  <!-- <geom name="left_foot1_collision" class="foot_capsule" fromto="0.1 -0.026 -0.025 0.05 -0.027
-                  -0.025"/>
-                  <geom name="left_foot2_collision" class="foot_capsule" fromto="-0.045 0 -0.015 0.12 0 -0.015"
-                    size="0.02"/>
-                  <geom name="left_foot3_collision" class="foot_capsule" fromto="0.1 0.026 -0.025 0.05 0.026 -0.025"/> -->
                  <geom name="left_foot1_collision" class="foot_capsule" fromto="0.1 -0.026 -0.025 0.05 -0.027 -0.025"/>
                  <geom name="left_foot2_collision" class="foot_capsule"
                    fromto="-0.044 -0.018 -0.025 0.123 -0.018 -0.025"/>
@@ -156,11 +151,6 @@
                    diaginertia="0.00167218 0.0016161 0.000217621"/>
                  <joint name="right_ankle_roll_joint" axis="1 0 0" range="-0.2618 0.2618"/>
                  <geom class="visual" material="black" mesh="right_ankle_roll_link"/>
-                  <!-- <geom name="right_foot1_collision" class="foot_capsule" fromto="0.1 -0.026 -0.025 0.05 -0.026
-                  -0.025"/>
-                  <geom name="right_foot2_collision" class="foot_capsule" fromto="-0.045 0 -0.015 0.12 0 -0.015"
-                    size="0.02"/>
-                  <geom name="right_foot3_collision" class="foot_capsule" fromto="0.1 0.026 -0.025 0.05 0.026 -0.025"/> -->
                  <geom name="right_foot1_collision" class="foot_capsule" fromto="0.1 -0.026 -0.025 0.05 -0.026 -0.025"/>
                  <geom name="right_foot2_collision" class="foot_capsule"
                    fromto="-0.044 -0.018 -0.025 0.123 -0.018 -0.025"/>
@@ -312,6 +302,7 @@
    <gyro name="imu_ang_vel" site="imu_in_pelvis"/>
    <velocimeter name="imu_lin_vel" site="imu_in_pelvis"/>
    <accelerometer name="imu_lin_acc" site="imu_in_pelvis"/>
+    <framezaxis name="imu_upvector" objtype="body" objname="world" reftype="site" refname="imu_in_pelvis"/>
    <subtreeangmom name="root_angmom" body="pelvis"/>
  </sensor>
 </mujoco>
@@ -165,6 +165,7 @@
    <gyro name="imu_ang_vel" site="imu"/>
    <velocimeter name="imu_lin_vel" site="imu"/>
    <accelerometer name="imu_lin_acc" site="imu"/>
+    <framezaxis name="imu_upvector" objtype="body" objname="world" reftype="site" refname="imu"/>
    <subtreeangmom name="root_angmom" body="trunk"/>
  </sensor>
 </mujoco>
@@ -3,6 +3,5 @@ from mjlab.entity.entity import Entity as Entity
 from mjlab.entity.entity import EntityArticulationInfoCfg as EntityArticulationInfoCfg
 from mjlab.entity.entity import EntityCfg as EntityCfg
 from mjlab.entity.entity import EntityIndexing as EntityIndexing
-from mjlab.entity.entity import VariantCfg as VariantCfg
-from mjlab.entity.entity import VariantEntityCfg as VariantEntityCfg
-from mjlab.entity.entity import VariantMetadata as VariantMetadata
+from mjlab.entity.variants import VariantEntityCfg as VariantEntityCfg
+from mjlab.entity.variants import VariantMetadata as VariantMetadata
@@ -3,7 +3,7 @@ from __future__ import annotations
 import warnings
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Callable, Sequence
+from typing import TYPE_CHECKING, Callable, Sequence

 import mujoco
 import mujoco_warp as mjwarp
@@ -18,14 +18,13 @@ from mjlab.entity.data import EntityData
 from mjlab.utils import spec_config as spec_cfg
 from mjlab.utils.lab_api.string import resolve_matching_names
 from mjlab.utils.mujoco import dof_width, qpos_width
-from mjlab.utils.spec import (
-  auto_wrap_fixed_base_mocap,
-  copy_mesh_data,
-  validate_variant_structure,
-)
+from mjlab.utils.spec import auto_wrap_fixed_base_mocap
 from mjlab.utils.string import resolve_expr
 from mjlab.utils.xml import fix_spec_xml, strip_buffer_textures

+if TYPE_CHECKING:
+  from mjlab.entity.variants import VariantMetadata
+

@dataclass(frozen=False)
 class EntityIndexing:
@@ -67,79 +66,6 @@ class EntityIndexing:
    return self.bodies[0].id


-@dataclass
-class VariantCfg:
-  """One object variant for per-world mesh randomization.
-
-  Each variant provides a ``spec_fn`` that returns an MjSpec for one object.
-  The ``weight`` controls what fraction of worlds use this variant.
-  """
-
-  spec_fn: Callable[[], mujoco.MjSpec]
-  weight: float = 1.0
-
-
-@dataclass(frozen=True)
-class BodyInertialMetadata:
-  """Explicit inertial properties for one body in a mesh variant."""
-
-  body_name: str
-  mass: float
-  ipos: tuple[float, float, float]
-  inertia: tuple[float, float, float]
-  iquat: tuple[float, float, float, float]
-
-
-@dataclass
-class VariantMetadata:
-  """Bookkeeping produced by Entity when merging variant specs."""
-
-  variant_names: tuple[str, ...]
-  variant_weights: tuple[float, ...]
-  # Per-variant ordered mesh names for each geom slot. Shorter variants
-  # have None for padding slots that should be disabled (dataid = -1).
-  variant_mesh_names: tuple[tuple[str | None, ...], ...]
-  num_mesh_geoms: int  # Max mesh geom count after padding.
-  # Per-variant explicit body inertials. Names are local to the variant spec;
-  # build_mesh_variant_model prefixes them with the scene entity name when
-  # applying them.
-  variant_body_inertials: tuple[tuple[BodyInertialMetadata, ...], ...] = ()
-
-
-def _iter_body_tree(body: mujoco.MjsBody):
-  yield body
-  for child in body.bodies:
-    yield from _iter_body_tree(child)
-
-
-def _collect_explicit_body_inertials(
-  root_body: mujoco.MjsBody,
-) -> tuple[BodyInertialMetadata, ...]:
-  inertials: list[BodyInertialMetadata] = []
-  for body in _iter_body_tree(root_body):
-    if not body.name or not body.explicitinertial:
-      continue
-    inertials.append(
-      BodyInertialMetadata(
-        body_name=body.name,
-        mass=float(body.mass),
-        ipos=(float(body.ipos[0]), float(body.ipos[1]), float(body.ipos[2])),
-        inertia=(
-          float(body.inertia[0]),
-          float(body.inertia[1]),
-          float(body.inertia[2]),
-        ),
-        iquat=(
-          float(body.iquat[0]),
-          float(body.iquat[1]),
-          float(body.iquat[2]),
-          float(body.iquat[3]),
-        ),
-      )
-    )
-  return tuple(inertials)
-
-
@dataclass
 class EntityCfg:
  @dataclass
@@ -187,50 +113,6 @@ class EntityArticulationInfoCfg:
  soft_joint_pos_limit_factor: float = 1.0


-def _variant_spec_fn_unset() -> mujoco.MjSpec:
-  """Sentinel default for ``VariantEntityCfg.spec_fn``.
-
-  ``VariantEntityCfg`` builds its spec from ``variants`` via
-  ``Entity._build_merged_spec``; the inherited ``spec_fn`` field is unused.
-  Identity comparison against this sentinel detects accidental user overrides.
-  """
-  raise AssertionError(
-    "VariantEntityCfg.spec_fn should never be called; the merged spec is "
-    "built from `variants`."
-  )
-
-
-@dataclass
-class VariantEntityCfg(EntityCfg):
-  """Entity config for per-world mesh variants.
-
-  Instead of a single ``spec_fn``, provide a dict of named variants.
-  Each world gets a variant assigned proportionally by weight. The
-  merged spec (with all variant meshes and padded geoms) is built
-  automatically.
-
-  All variants must share the same kinematic structure (same bodies,
-  joints, joint types). Only mesh geoms can differ.
-
-  Variant assignment is fixed at ``Simulation`` initialization; it does
-  not resample on episode reset. Pass the per-variant spec via
-  :class:`VariantCfg` rather than setting ``spec_fn`` directly.
-  """
-
-  variants: dict[str, VariantCfg] = field(default_factory=dict)
-  """Named mesh variants with weights."""
-
-  spec_fn: Callable[[], mujoco.MjSpec] = field(default=_variant_spec_fn_unset)
-  """Unused on ``VariantEntityCfg``; the merged spec is built from ``variants``."""
-
-  def __post_init__(self) -> None:
-    if self.spec_fn is not _variant_spec_fn_unset:
-      raise ValueError(
-        "VariantEntityCfg.spec_fn cannot be set; pass per-variant specs via "
-        "VariantCfg(spec_fn=...) inside `variants` instead."
-      )
-
-
 class Entity:
  """An entity represents a physical object in the simulation.

@@ -272,130 +154,13 @@ class Entity:
    self._add_initial_state_keyframe()

  def _build_spec(self) -> None:
+    from mjlab.entity.variants import VariantEntityCfg, build_merged_variant_spec
+
    if isinstance(self.cfg, VariantEntityCfg):
-      self._build_merged_spec()
+      self._spec, self._variant_metadata = build_merged_variant_spec(self.cfg)
    else:
      self._spec = auto_wrap_fixed_base_mocap(self.cfg.spec_fn)()

-  def _build_merged_spec(self) -> None:
-    """Build a merged spec from multiple variant specs.
-
-    Validates that all variants share the same kinematic structure,
-    merges all mesh assets into a single spec, and pads the body to
-    the max mesh geom count across variants.
-    """
-    assert isinstance(self.cfg, VariantEntityCfg)
-    variants = self.cfg.variants
-    if not variants:
-      raise ValueError("VariantEntityCfg.variants must contain at least one entry.")
-
-    variant_names: list[str] = []
-    variant_weights: list[float] = []
-    variant_specs: list[mujoco.MjSpec] = []
-    for name, vcfg in variants.items():
-      variant_names.append(name)
-      variant_weights.append(vcfg.weight)
-      variant_specs.append(vcfg.spec_fn())
-
-    # Find root body in each variant.
-    variant_bodies: list[mujoco.MjsBody] = []
-    for i, spec in enumerate(variant_specs):
-      children = list(spec.worldbody.bodies)
-      if len(children) != 1:
-        raise ValueError(
-          f"Variant '{variant_names[i]}' must have exactly one "
-          f"root body under worldbody, got {len(children)}."
-        )
-      variant_bodies.append(children[0])
-
-    validate_variant_structure(variant_names, variant_bodies)
-
-    # Variant entities must be floating-base. Mocap auto-wrap is not applied
-    # for variant entities, so fixed-base variants would silently stack at
-    # the world origin. Variants share joint structure (validated above), so
-    # checking the first is sufficient.
-    ref_joints = list(variant_bodies[0].joints)
-    if not ref_joints or ref_joints[0].type != mujoco.mjtJoint.mjJNT_FREE:
-      raise ValueError(
-        "VariantEntityCfg requires floating-base variants. Each variant's "
-        "root body must declare a free joint via body.add_freejoint(); "
-        "fixed-base variants are not supported."
-      )
-
-    variant_body_inertials = tuple(
-      _collect_explicit_body_inertials(body) for body in variant_bodies
-    )
-
-    # Collect original mesh names per variant BEFORE any renaming.
-    variant_orig_mesh_names: list[list[str]] = []
-    variant_mesh_geom_counts: list[int] = []
-    for body in variant_bodies:
-      orig_names = [
-        g.meshname for g in body.geoms if g.type == mujoco.mjtGeom.mjGEOM_MESH
-      ]
-      variant_orig_mesh_names.append(orig_names)
-      variant_mesh_geom_counts.append(len(orig_names))
-
-    max_mesh_geoms = max(variant_mesh_geom_counts)
-
-    # Use first variant as template. Prefix ALL mesh names with
-    # variant name to avoid collisions across variants.
-    template_spec = variant_specs[0]
-    template_body = variant_bodies[0]
-
-    # Rename template meshes first.
-    template_prefix = f"{variant_names[0]}/"
-    old_to_new: dict[str, str] = {}
-    for mesh in template_spec.meshes:
-      new_name = f"{template_prefix}{mesh.name}"
-      old_to_new[mesh.name] = new_name
-      mesh.name = new_name
-    for g in template_body.geoms:
-      if g.meshname in old_to_new:
-        g.meshname = old_to_new[g.meshname]
-
-    # Copy mesh assets from other variants.
-    for i in range(1, len(variant_specs)):
-      prefix = f"{variant_names[i]}/"
-      for mesh in variant_specs[i].meshes:
-        new_mesh = template_spec.add_mesh()
-        new_mesh.name = f"{prefix}{mesh.name}"
-        copy_mesh_data(mesh, new_mesh)
-
-    # Pad body to max mesh geom count.
-    current_count = variant_mesh_geom_counts[0]
-    if max_mesh_geoms > current_count:
-      longest_idx = max(
-        range(len(variant_mesh_geom_counts)),
-        key=lambda j: variant_mesh_geom_counts[j],
-      )
-      longest_prefix = f"{variant_names[longest_idx]}/"
-      longest_names = variant_orig_mesh_names[longest_idx]
-      for k in range(current_count, max_mesh_geoms):
-        geom = template_body.add_geom()
-        geom.type = mujoco.mjtGeom.mjGEOM_MESH
-        geom.meshname = f"{longest_prefix}{longest_names[k]}"
-        geom.contype = 1
-        geom.conaffinity = 1
-
-    # Build variant_mesh_names: use original names with variant prefix.
-    variant_mesh_name_lists: list[tuple[str | None, ...]] = []
-    for i, orig_names in enumerate(variant_orig_mesh_names):
-      prefix = f"{variant_names[i]}/"
-      names: list[str | None] = [f"{prefix}{n}" for n in orig_names]
-      while len(names) < max_mesh_geoms:
-        names.append(None)
-      variant_mesh_name_lists.append(tuple(names))
-
-    self._variant_metadata = VariantMetadata(
-      variant_names=tuple(variant_names),
-      variant_weights=tuple(variant_weights),
-      variant_mesh_names=tuple(variant_mesh_name_lists),
-      num_mesh_geoms=max_mesh_geoms,
-      variant_body_inertials=variant_body_inertials,
-    )
-    self._spec = template_spec
-
  @property
  def variant_metadata(self) -> VariantMetadata | None:
    return self._variant_metadata
@@ -404,6 +169,16 @@ class Entity:
    self._all_joints = self._spec.joints
    self._free_joint = None
    self._non_free_joints = tuple(self._all_joints)
+
+    free_joints = [j for j in self._all_joints if j.type == mujoco.mjtJoint.mjJNT_FREE]
+    if len(free_joints) > 1:
+      raise ValueError(
+        f"Entity spec has {len(free_joints)} freejoints. An Entity models a "
+        "single rigid- or articulated-body system with at most one freejoint, "
+        "which serves as its root. Model each detached floating body as its own "
+        "entry in SceneCfg.entities instead."
+      )
+
    if self._all_joints and self._all_joints[0].type == mujoco.mjtJoint.mjJNT_FREE:
      self._free_joint = self._all_joints[0]
      if not self._free_joint.name:
@@ -184,7 +184,7 @@ class ManagerBasedRlEnv:
    # Initialize base environment state.
    self.cfg = cfg
    if self.cfg.seed is not None:
-      self.cfg.seed = self.seed(self.cfg.seed, device=device)
+      self.cfg.seed = self.seed(self.cfg.seed)
    self._sim_step_counter = 0
    self.extras = {}
    self.obs_buf = {}
@@ -194,21 +194,13 @@ class ManagerBasedRlEnv:

    # Initialize scene and simulation.
    self.scene = Scene(self.cfg.scene, device=device)
-    if self.scene.has_mesh_variants:
-      self.sim = Simulation(
-        num_envs=self.scene.num_envs,
-        cfg=self.cfg.sim,
-        spec=self.scene.spec,
-        variant_info=self.scene.collect_variant_info(),
-        device=device,
-      )
-    else:
-      self.sim = Simulation(
-        num_envs=self.scene.num_envs,
-        cfg=self.cfg.sim,
-        model=self.scene.compile(),
-        device=device,
-      )
+    self.sim = Simulation(
+      num_envs=self.scene.num_envs,
+      cfg=self.cfg.sim,
+      spec=self.scene.spec,
+      variant_info=self.scene.collect_variant_info(),
+      device=device,
+    )

    self.scene.initialize(
      mj_model=self.sim.mj_model,
@@ -373,6 +365,7 @@ class ManagerBasedRlEnv:
      env_ids = torch.arange(self.num_envs, dtype=torch.int64, device=self.device)
    if seed is not None:
      self.seed(seed)
+    self.extras["log"] = dict()
    self._reset_idx(env_ids)
    self.scene.write_data_to_sim()
    self.sim.forward()
@@ -422,6 +415,7 @@ class ManagerBasedRlEnv:
        "reset(env_ids=...) before calling step() again when auto_reset=False."
      )

+    self.extras["log"] = dict()
    self.action_manager.process_action(action.to(self.device))

    for _ in range(self.cfg.decimation):
@@ -484,6 +478,9 @@ class ManagerBasedRlEnv:
      self.extras,
    )

+  def get_observations(self) -> dict:
+    return self.observation_manager.compute()
+
  def render(self) -> np.ndarray | None:
    if self.render_mode == "human" or self.render_mode is None:
      return None
@@ -506,11 +503,12 @@ class ManagerBasedRlEnv:
      self._offline_renderer.close()
    self.recorder_manager.close()

-  def seed(self, seed: int = -1, device: str | torch.device | None = None) -> int:
+  @staticmethod
+  def seed(seed: int = -1) -> int:
    if seed == -1:
      seed = np.random.randint(0, 10_000)
    print_info(f"Setting seed: {seed}")
-    random_utils.seed_rng(seed, device=device if device is not None else self.device)
+    random_utils.seed_rng(seed)
    return seed

  def update_visualizers(self, visualizer: DebugVisualizer) -> None:
@@ -564,7 +562,6 @@ class ManagerBasedRlEnv:
      )

    # NOTE: This is order sensitive.
-    self.extras["log"] = dict()
    # observation manager.
    info = self.observation_manager.reset(env_ids)
    self.extras["log"].update(info)
@@ -7,23 +7,65 @@ from typing import TYPE_CHECKING, Literal
 import torch

 from mjlab.actuator import (
-    BuiltinPositionActuator,
-    BuiltinVelocityActuator,
-    BuiltinMotorActuator,
-    IdealPdActuator,
+  BuiltinDcMotorActuator,
+  BuiltinPdActuator,
+  BuiltinPositionActuator,
+  IdealPdActuator,
 )
+from mjlab.actuator.actuator import TransmissionType
+from mjlab.actuator.builtin_actuator import DcMotorInputMode
 from mjlab.actuator.xml_actuator import XmlActuator
 from mjlab.entity import Entity
 from mjlab.managers.event_manager import requires_model_fields
 from mjlab.managers.scene_entity_config import SceneEntityCfg

 from ._core import _DEFAULT_ASSET_CFG
-from ._types import resolve_distribution
+from ._types import Operation, resolve_distribution, resolve_operation

 if TYPE_CHECKING:
  from mjlab.envs import ManagerBasedRlEnv


+def _resolve_actuators(asset: Entity, asset_cfg: SceneEntityCfg) -> list:
+  """Resolve actuator objects from SceneEntityCfg.
+
+  SceneEntityCfg actuator_ids/name resolution is based on spec actuators, while
+  runtime ``asset.actuators`` may contain grouped actuator objects. For grouped
+  actuators, map matched actuator target names back to the owning actuator object.
+  """
+  if asset_cfg.actuator_names is not None:
+    matched_names = asset_cfg.actuator_names
+    if isinstance(matched_names, str):
+      matched_names = [matched_names]
+    resolved = []
+    for actuator in asset.actuators:
+      if any(name in actuator.target_names for name in matched_names):
+        resolved.append(actuator)
+    return resolved
+
+  if isinstance(asset_cfg.actuator_ids, list):
+    if all(0 <= i < len(asset.actuators) for i in asset_cfg.actuator_ids):
+      return [asset.actuators[i] for i in asset_cfg.actuator_ids]
+
+    resolved = []
+    seen = set()
+    actuator_names = asset.actuator_names
+    for i in asset_cfg.actuator_ids:
+      if not (0 <= i < len(actuator_names)):
+        continue
+      target_name = actuator_names[i]
+      for actuator_idx, actuator in enumerate(asset.actuators):
+        if target_name in actuator.target_names and actuator_idx not in seen:
+          resolved.append(actuator)
+          seen.add(actuator_idx)
+          break
+    return resolved
+  elif isinstance(asset_cfg.actuator_ids, slice):
+    return asset.actuators[asset_cfg.actuator_ids]
+  else:
+    return [asset.actuators[asset_cfg.actuator_ids]]
+
+
@requires_model_fields("actuator_gainprm", "actuator_biasprm")
 def pd_gains(
  env: ManagerBasedRlEnv,
@@ -32,7 +74,7 @@ def pd_gains(
  kd_range: tuple[float, float],
  asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
  distribution: Literal["uniform", "log_uniform"] = "uniform",
-  operation: Literal["scale", "abs"] = "scale",
+  operation: Operation | str = "scale",
 ) -> None:
  """Randomize PD stiffness and damping gains.

@@ -46,6 +88,11 @@ def pd_gains(
    operation: "scale" multiplies default gains by sampled values, "abs" sets
      absolute values.
  """
+  op = resolve_operation(operation)
+  if op.name not in ("scale", "abs"):
+    raise ValueError(
+      f"pd_gains only supports 'scale' and 'abs' operations, got {op.name!r}"
+    )
  asset: Entity = env.scene[asset_cfg.name]

  if env_ids is None:
@@ -53,34 +100,36 @@ def pd_gains(
  else:
    env_ids = env_ids.to(env.device, dtype=torch.int)

-  if isinstance(asset_cfg.actuator_ids, list):
-    actuators = [asset.actuators[i] for i in asset_cfg.actuator_ids]
-  elif isinstance(asset_cfg.actuator_ids, slice):
-    actuators = asset.actuators[asset_cfg.actuator_ids]
-  else:
-    actuators = [asset.actuators[asset_cfg.actuator_ids]]
+  actuators = _resolve_actuators(asset, asset_cfg)

  for actuator in actuators:
    ctrl_ids = actuator.global_ctrl_ids
+    # Each target needs one kp draw and one kd draw. For single-element
+    # actuators that's len(ctrl_ids) of each; for BuiltinPd the ctrl tensor
+    # has 2*N entries but only N independent kp/kd values, so we sample
+    # num_targets to avoid throwing the other half away.
+    n_gains = (
+      actuator.num_targets if isinstance(actuator, BuiltinPdActuator) else len(ctrl_ids)
+    )

    dist = resolve_distribution(distribution)
    kp_samples = dist.sample(
      torch.tensor(kp_range[0], device=env.device),
      torch.tensor(kp_range[1], device=env.device),
-      (len(env_ids), len(ctrl_ids)),
+      (len(env_ids), n_gains),
      env.device,
    )
    kd_samples = dist.sample(
      torch.tensor(kd_range[0], device=env.device),
      torch.tensor(kd_range[1], device=env.device),
-      (len(env_ids), len(ctrl_ids)),
+      (len(env_ids), n_gains),
      env.device,
    )

    if isinstance(actuator, BuiltinPositionActuator) or (
      isinstance(actuator, XmlActuator) and actuator.command_field == "position"
    ):
-      if operation == "scale":
+      if op.name == "scale":
        default_gainprm = env.sim.get_default_field("actuator_gainprm")
        default_biasprm = env.sim.get_default_field("actuator_biasprm")
        env.sim.model.actuator_gainprm[env_ids[:, None], ctrl_ids, 0] = (
@@ -92,15 +141,69 @@ def pd_gains(
        env.sim.model.actuator_biasprm[env_ids[:, None], ctrl_ids, 2] = (
          default_biasprm[ctrl_ids, 2] * kd_samples
        )
-      elif operation == "abs":
+      else:
+        assert op.name == "abs"
        env.sim.model.actuator_gainprm[env_ids[:, None], ctrl_ids, 0] = kp_samples
        env.sim.model.actuator_biasprm[env_ids[:, None], ctrl_ids, 1] = -kp_samples
        env.sim.model.actuator_biasprm[env_ids[:, None], ctrl_ids, 2] = -kd_samples

+    elif isinstance(actuator, BuiltinDcMotorActuator):
+      if actuator.cfg.mode == DcMotorInputMode.VOLTAGE:
+        raise ValueError(
+          "dr.pd_gains does not apply to BuiltinDcMotorActuator in VOLTAGE "
+          "mode (no internal PID gains to scale)."
+        )
+      # DC motor stores kp at gainprm[4] and kd at gainprm[6] (set via
+      # set_to_dcmotor). The bias slots carry back-EMF / cogging, not the PD,
+      # so we only touch gainprm.
+      if op.name == "scale":
+        default_gainprm = env.sim.get_default_field("actuator_gainprm")
+        env.sim.model.actuator_gainprm[env_ids[:, None], ctrl_ids, 4] = (
+          default_gainprm[ctrl_ids, 4] * kp_samples
+        )
+        env.sim.model.actuator_gainprm[env_ids[:, None], ctrl_ids, 6] = (
+          default_gainprm[ctrl_ids, 6] * kd_samples
+        )
+      else:
+        assert op.name == "abs"
+        env.sim.model.actuator_gainprm[env_ids[:, None], ctrl_ids, 4] = kp_samples
+        env.sim.model.actuator_gainprm[env_ids[:, None], ctrl_ids, 6] = kd_samples
+
+    elif isinstance(actuator, BuiltinPdActuator):
+      # ctrl_ids is laid out as [pos_0..pos_{N-1}, vel_0..vel_{N-1}], so the
+      # first N rows carry kp and the next N carry kd.
+      n = actuator.num_targets
+      pos_ids = ctrl_ids[:n]
+      vel_ids = ctrl_ids[n:]
+      if op.name == "scale":
+        default_gainprm = env.sim.get_default_field("actuator_gainprm")
+        default_biasprm = env.sim.get_default_field("actuator_biasprm")
+        env.sim.model.actuator_gainprm[env_ids[:, None], pos_ids, 0] = (
+          default_gainprm[pos_ids, 0] * kp_samples
+        )
+        env.sim.model.actuator_biasprm[env_ids[:, None], pos_ids, 1] = (
+          default_biasprm[pos_ids, 1] * kp_samples
+        )
+        env.sim.model.actuator_gainprm[env_ids[:, None], vel_ids, 0] = (
+          default_gainprm[vel_ids, 0] * kd_samples
+        )
+        env.sim.model.actuator_biasprm[env_ids[:, None], vel_ids, 2] = (
+          default_biasprm[vel_ids, 2] * kd_samples
+        )
+      else:
+        assert op.name == "abs"
+        env.sim.model.actuator_gainprm[env_ids[:, None], pos_ids, 0] = kp_samples
+        env.sim.model.actuator_biasprm[env_ids[:, None], pos_ids, 1] = -kp_samples
+        env.sim.model.actuator_gainprm[env_ids[:, None], vel_ids, 0] = kd_samples
+        env.sim.model.actuator_biasprm[env_ids[:, None], vel_ids, 2] = -kd_samples
+      # biasprm[2] on the position half stays zero by construction. Writing
+      # anything else here would inject damping into the position element on
+      # top of the velocity element, silently double-counting kd.
+
    elif isinstance(actuator, IdealPdActuator):
      assert actuator.stiffness is not None
      assert actuator.damping is not None
-      if operation == "scale":
+      if op.name == "scale":
        assert actuator.default_stiffness is not None
        assert actuator.default_damping is not None
        actuator.set_gains(
@@ -108,25 +211,26 @@ def pd_gains(
          kp=actuator.default_stiffness[env_ids] * kp_samples,
          kd=actuator.default_damping[env_ids] * kd_samples,
        )
-      elif operation == "abs":
+      else:
+        assert op.name == "abs"
        actuator.set_gains(env_ids, kp=kp_samples, kd=kd_samples)

    else:
      raise TypeError(
-        f"pd_gains only supports BuiltinPositionActuator, "
-        f"XmlActuator (position), and IdealPdActuator, "
-        f"got {type(actuator).__name__}"
+        f"pd_gains only supports BuiltinPositionActuator, BuiltinPdActuator, "
+        f"BuiltinDcMotorActuator (position/velocity mode), XmlActuator (position), "
+        f"and IdealPdActuator, got {type(actuator).__name__}"
      )


-@requires_model_fields("actuator_forcerange")
+@requires_model_fields("actuator_forcerange", "jnt_actfrcrange", "tendon_actfrcrange")
 def effort_limits(
  env: ManagerBasedRlEnv,
  env_ids: torch.Tensor | None,
  effort_limit_range: tuple[float, float],
  asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
  distribution: Literal["uniform", "log_uniform"] = "uniform",
-  operation: Literal["scale", "abs"] = "scale",
+  operation: Operation | str = "scale",
 ) -> None:
  """Randomize actuator effort limits.

@@ -138,6 +242,11 @@ def effort_limits(
    distribution: Distribution type ("uniform" or "log_uniform").
    operation: "scale" multiplies default limits, "abs" sets absolute values.
  """
+  op = resolve_operation(operation)
+  if op.name not in ("scale", "abs"):
+    raise ValueError(
+      f"effort_limits only supports 'scale' and 'abs' operations, got {op.name!r}"
+    )
  asset: Entity = env.scene[asset_cfg.name]

  if env_ids is None:
@@ -145,36 +254,32 @@ def effort_limits(
  else:
    env_ids = env_ids.to(env.device, dtype=torch.int)

-  if isinstance(asset_cfg.actuator_ids, list):
-    actuators = [asset.actuators[i] for i in asset_cfg.actuator_ids]
-  else:
-    actuators = asset.actuators[asset_cfg.actuator_ids]
+  actuators = _resolve_actuators(asset, asset_cfg)

  if not isinstance(actuators, list):
    actuators = [actuators]

  for actuator in actuators:
    ctrl_ids = actuator.global_ctrl_ids
-    num_actuators = len(ctrl_ids)
+    # One effort sample per target. For single-element actuators this matches
+    # ctrl_ids; for BuiltinPd the limit lives on the joint/tendon, so one
+    # sample per target is sufficient regardless of the two-element ctrl.
+    n_samples = (
+      actuator.num_targets if isinstance(actuator, BuiltinPdActuator) else len(ctrl_ids)
+    )

    dist = resolve_distribution(distribution)
    effort_samples = dist.sample(
      torch.tensor(effort_limit_range[0], device=env.device),
      torch.tensor(effort_limit_range[1], device=env.device),
-      (len(env_ids), num_actuators),
+      (len(env_ids), n_samples),
      env.device,
    )

-    if isinstance(
-        actuator,
-        (
-            BuiltinPositionActuator,
-            BuiltinVelocityActuator,
-            BuiltinMotorActuator,
-            XmlActuator,
-        ),
+    if isinstance(actuator, (BuiltinPositionActuator, BuiltinDcMotorActuator)) or (
+      isinstance(actuator, XmlActuator) and actuator.command_field == "position"
    ):
-      if operation == "scale":
+      if op.name == "scale":
        default_forcerange = env.sim.get_default_field("actuator_forcerange")
        env.sim.model.actuator_forcerange[env_ids[:, None], ctrl_ids, 0] = (
          default_forcerange[ctrl_ids, 0] * effort_samples
@@ -182,7 +287,8 @@ def effort_limits(
        env.sim.model.actuator_forcerange[env_ids[:, None], ctrl_ids, 1] = (
          default_forcerange[ctrl_ids, 1] * effort_samples
        )
-      elif operation == "abs":
+      else:
+        assert op.name == "abs"
        env.sim.model.actuator_forcerange[
          env_ids[:, None], ctrl_ids, 0
        ] = -effort_samples
@@ -192,18 +298,42 @@ def effort_limits(

    elif isinstance(actuator, IdealPdActuator):
      assert actuator.force_limit is not None
-      if operation == "scale":
+      if op.name == "scale":
        assert actuator.default_force_limit is not None
        actuator.set_effort_limit(
          env_ids,
          effort_limit=actuator.default_force_limit[env_ids] * effort_samples,
        )
-      elif operation == "abs":
+      else:
+        assert op.name == "abs"
        actuator.set_effort_limit(env_ids, effort_limit=effort_samples)

+    elif isinstance(actuator, BuiltinPdActuator):
+      # BuiltinPd's effort_limit lives on the joint/tendon as a sum-clamp
+      # (jnt_actfrcrange / tendon_actfrcrange), not on per-element forcerange.
+      if actuator.transmission_type == TransmissionType.JOINT:
+        field = "jnt_actfrcrange"
+        target_global_ids = asset.indexing.joint_ids[actuator.target_ids]
+      else:
+        field = "tendon_actfrcrange"
+        target_global_ids = asset.indexing.tendon_ids[actuator.target_ids]
+      arr = getattr(env.sim.model, field)
+      if op.name == "scale":
+        default = env.sim.get_default_field(field)
+        arr[env_ids[:, None], target_global_ids, 0] = (
+          default[target_global_ids, 0] * effort_samples
+        )
+        arr[env_ids[:, None], target_global_ids, 1] = (
+          default[target_global_ids, 1] * effort_samples
+        )
+      else:
+        assert op.name == "abs"
+        arr[env_ids[:, None], target_global_ids, 0] = -effort_samples
+        arr[env_ids[:, None], target_global_ids, 1] = effort_samples
+
    else:
      raise TypeError(
-        f"effort_limits only supports BuiltinPositionActuator, BuiltinVelocityActuator, "
-        f"BuiltinMotorActuator, XmlActuator, and IdealPdActuator, "
+        f"effort_limits only supports BuiltinPositionActuator, BuiltinPdActuator, "
+        f"BuiltinDcMotorActuator, XmlActuator (position), and IdealPdActuator, "
        f"got {type(actuator).__name__}"
      )
@@ -21,6 +21,38 @@ if TYPE_CHECKING:
  from mjlab.viewer.debug_visualizer import DebugVisualizer

 _DEFAULT_ASSET_CFG = SceneEntityCfg("robot")
+_SE3_KEYS = ("x", "y", "z", "roll", "pitch", "yaw")
+
+
+def _sample_se3_range(
+  range_dict: dict[str, tuple[float, float]] | None,
+  shape: tuple[int, ...],
+  device: str,
+) -> torch.Tensor:
+  """Sample uniform ``[x, y, z, roll, pitch, yaw]`` offsets.
+
+  ``range_dict`` maps any subset of those keys to ``(min, max)`` ranges; missing
+  keys default to ``(0.0, 0.0)`` (no offset). ``None`` is treated as empty. The
+  returned tensor has the requested ``shape`` whose last dimension must be 6.
+  """
+  range_dict = range_dict or {}
+  range_list = [range_dict.get(key, (0.0, 0.0)) for key in _SE3_KEYS]
+  ranges = torch.tensor(range_list, device=device)
+  return sample_uniform(ranges[:, 0], ranges[:, 1], shape, device=device)
+
+
+def resolve_env_ids(
+  env: ManagerBasedRlEnv, env_ids: torch.Tensor | None
+) -> torch.Tensor:
+  """Return ``env_ids`` unchanged, or all environment indices if ``None``.
+
+  Event functions receive ``env_ids=None`` to mean "all environments" (a full
+  reset, or a global-time interval term). This normalizes that sentinel to a
+  concrete index tensor so the function body can assume a real ``torch.Tensor``.
+  """
+  if env_ids is None:
+    return torch.arange(env.num_envs, device=env.device, dtype=torch.int)
+  return env_ids


 def randomize_terrain(env: ManagerBasedRlEnv, env_ids: torch.Tensor | None) -> None:
@@ -29,8 +61,7 @@ def randomize_terrain(env: ManagerBasedRlEnv, env_ids: torch.Tensor | None) -> N
  This picks a random terrain type (column) and difficulty level (row) for each
  environment. Useful for play/evaluation mode to test on varied terrains.
  """
-  if env_ids is None:
-    env_ids = torch.arange(env.num_envs, device=env.device, dtype=torch.int)
+  env_ids = resolve_env_ids(env, env_ids)

  terrain = env.scene.terrain
  if terrain is not None:
@@ -48,8 +79,7 @@ def reset_scene_to_default(

  Automatically applies env_origins offset to position all entities correctly.
  """
-  if env_ids is None:
-    env_ids = torch.arange(env.num_envs, device=env.device, dtype=torch.int)
+  env_ids = resolve_env_ids(env, env_ids)

  for entity in env.scene.entities.values():
    if not isinstance(entity, Entity):
@@ -105,19 +135,12 @@ def reset_root_state_uniform(
    velocity_range: Velocity range (only used for floating-base entities).
    asset_cfg: Asset configuration.
  """
-  if env_ids is None:
-    env_ids = torch.arange(env.num_envs, device=env.device, dtype=torch.int)
+  env_ids = resolve_env_ids(env, env_ids)

  asset: Entity = env.scene[asset_cfg.name]

  # Pose.
-  range_list = [
-    pose_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]
-  ]
-  ranges = torch.tensor(range_list, device=env.device)
-  pose_samples = sample_uniform(
-    ranges[:, 0], ranges[:, 1], (len(env_ids), 6), device=env.device
-  )
+  pose_samples = _sample_se3_range(pose_range, (len(env_ids), 6), env.device)

  # Fixed-based entities with mocap=True.
  if asset.is_fixed_base:
@@ -157,16 +180,7 @@ def reset_root_state_uniform(
  orientations = quat_mul(root_states[:, 3:7], orientations_delta)

  # Velocities.
-  if velocity_range is None:
-    velocity_range = {}
-  range_list = [
-    velocity_range.get(key, (0.0, 0.0))
-    for key in ["x", "y", "z", "roll", "pitch", "yaw"]
-  ]
-  ranges = torch.tensor(range_list, device=env.device)
-  vel_samples = sample_uniform(
-    ranges[:, 0], ranges[:, 1], (len(env_ids), 6), device=env.device
-  )
+  vel_samples = _sample_se3_range(velocity_range, (len(env_ids), 6), env.device)
  velocities = root_states[:, 7:13] + vel_samples

  asset.write_root_link_pose_to_sim(
@@ -199,8 +213,7 @@ def reset_root_state_from_flat_patches(
    velocity_range: Optional velocity range (floating-base only).
    asset_cfg: Asset configuration.
  """
-  if env_ids is None:
-    env_ids = torch.arange(env.num_envs, device=env.device, dtype=torch.int)
+  env_ids = resolve_env_ids(env, env_ids)

  terrain = env.scene.terrain
  if terrain is None or patch_name not in terrain.flat_patches:
@@ -230,15 +243,7 @@ def reset_root_state_from_flat_patches(
  root_states = default_root_state[env_ids].clone()

  # Apply optional pose range offset.
-  if pose_range is None:
-    pose_range = {}
-  range_list = [
-    pose_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]
-  ]
-  ranges = torch.tensor(range_list, device=env.device)
-  pose_samples = sample_uniform(
-    ranges[:, 0], ranges[:, 1], (len(env_ids), 6), device=env.device
-  )
+  pose_samples = _sample_se3_range(pose_range, (len(env_ids), 6), env.device)

  # Position: flat patch position + optional offset. Use patch z instead of default.
  final_positions = positions.clone()
@@ -262,16 +267,7 @@ def reset_root_state_from_flat_patches(
    return

  # Velocities.
-  if velocity_range is None:
-    velocity_range = {}
-  vel_range_list = [
-    velocity_range.get(key, (0.0, 0.0))
-    for key in ["x", "y", "z", "roll", "pitch", "yaw"]
-  ]
-  vel_ranges = torch.tensor(vel_range_list, device=env.device)
-  vel_samples = sample_uniform(
-    vel_ranges[:, 0], vel_ranges[:, 1], (len(env_ids), 6), device=env.device
-  )
+  vel_samples = _sample_se3_range(velocity_range, (len(env_ids), 6), env.device)
  velocities = root_states[:, 7:13] + vel_samples

  asset.write_root_link_pose_to_sim(
@@ -287,8 +283,7 @@ def reset_joints_by_offset(
  velocity_range: tuple[float, float],
  asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
 ) -> None:
-  if env_ids is None:
-    env_ids = torch.arange(env.num_envs, device=env.device, dtype=torch.int)
+  env_ids = resolve_env_ids(env, env_ids)

  asset: Entity = env.scene[asset_cfg.name]
  default_joint_pos = asset.data.default_joint_pos
@@ -320,28 +315,60 @@ def reset_joints_by_offset(

 def push_by_setting_velocity(
  env: ManagerBasedRlEnv,
-  env_ids: torch.Tensor,
+  env_ids: torch.Tensor | None,
  velocity_range: dict[str, tuple[float, float]],
  asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
 ) -> None:
+  """Push an entity by overwriting its root velocity with a sampled offset.
+
+  This is an *instantaneous, mass-independent* kick: it adds a uniformly sampled
+  delta directly to the root velocity, ignoring inertia and contact dynamics. It
+  is the cheapest disturbance and the standard locomotion "push the robot" term.
+  Use with ``mode="interval"``.
+
+  For force-based disturbances that respect the entity's dynamics, see
+  :func:`apply_external_force_torque` (a constant wrench you manage yourself) or
+  :class:`apply_body_impulse` (transient, self-managing impulses).
+  """
+  env_ids = resolve_env_ids(env, env_ids)
  asset: Entity = env.scene[asset_cfg.name]
  vel_w = asset.data.root_link_vel_w[env_ids]
-  range_list = [
-    velocity_range.get(key, (0.0, 0.0))
-    for key in ["x", "y", "z", "roll", "pitch", "yaw"]
-  ]
-  ranges = torch.tensor(range_list, device=env.device)
-  vel_w += sample_uniform(ranges[:, 0], ranges[:, 1], vel_w.shape, device=env.device)
+  vel_w += _sample_se3_range(velocity_range, vel_w.shape, env.device)
  asset.write_root_link_velocity_to_sim(vel_w, env_ids=env_ids)


 def apply_external_force_torque(
  env: ManagerBasedRlEnv,
-  env_ids: torch.Tensor,
+  env_ids: torch.Tensor | None,
  force_range: tuple[float, float],
  torque_range: tuple[float, float],
  asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
 ) -> None:
+  """Apply a single *constant* external wrench to bodies.
+
+  Samples a force and torque once and writes them to ``xfrc_applied``. The wrench
+  is **stateless and never expires**: MuJoCo holds it constant on every physics
+  step until something overwrites or zeroes it. There is no duration, cooldown,
+  or auto-clear.
+
+  **When to use this vs.** :class:`apply_body_impulse`:
+
+  - Use ``apply_external_force_torque`` for a *steady, episode-long* disturbance
+    such as a fixed payload, a constant wind, or a sustained load. The intended
+    pattern is ``mode="reset"``: re-randomize the wrench each episode so it holds
+    for that episode's duration. Because it never turns itself off, **you are
+    responsible for clearing or overwriting it** (e.g. via the next reset). It is
+    *not* suited to transient bumps on its own.
+
+  - Use :class:`apply_body_impulse` for *transient, repeated, randomized*
+    disturbances during an episode (bumps, gusts, collisions). It runs a full
+    cooldown -> trigger -> sustain -> expire lifecycle per environment, zeroing
+    the wrench automatically when each impulse ends, and ticks on ``mode="step"``.
+
+  For an instantaneous, mass-independent kick instead of a force, see
+  :func:`push_by_setting_velocity`.
+  """
+  env_ids = resolve_env_ids(env, env_ids)
  asset: Entity = env.scene[asset_cfg.name]
  num_bodies = (
    len(asset_cfg.body_ids)
@@ -385,6 +412,10 @@ class apply_body_impulse:
  applied.

  Use with ``mode="step"``.
+
+  For a *constant* episode-long wrench instead of transient impulses, see
+  :func:`apply_external_force_torque`. For an instantaneous, mass-independent
+  velocity kick, see :func:`push_by_setting_velocity`.
  """

  @dataclass
@@ -422,9 +453,16 @@ class apply_body_impulse:
      else self._asset.num_bodies
    )

+    self._cooldown_s: tuple[float, float] = cfg.params["cooldown_s"]
    self._time_remaining = torch.zeros(self._num_envs, device=self._device)
-    self._interval_time_left = torch.zeros(self._num_envs, device=self._device)
    self._active = torch.zeros(self._num_envs, device=self._device, dtype=torch.bool)
+    # Pre-sample the initial cooldown so the first impulse is preceded by a cooldown
+    # rather than firing immediately at t=0.
+    self._interval_time_left = self._sample_cooldown(self._num_envs)
+
+  def _sample_cooldown(self, n: int) -> torch.Tensor:
+    low, high = self._cooldown_s
+    return sample_uniform(low, high, n, self._device)

  def __call__(
    self,
@@ -446,13 +484,14 @@ class apply_body_impulse:
      torque_range: ``(min, max)`` uniform range for each torque component (Nm).
      duration_s: ``(min, max)`` uniform range for impulse duration in seconds.
      cooldown_s: ``(min, max)`` uniform range for the cooldown between consecutive
-        impulses in seconds.
+        impulses in seconds. Captured at init so the first impulse can be
+        preceded by a sampled cooldown; the kwarg passed here is unused.
      asset_cfg: Entity and body selection. ``body_ids`` on the config selects which
        bodies receive forces.
      body_point_offset: Optional ``(x, y, z)`` offset in the body frame where the
        force is applied. Generates additional torque via ``cross(offset, force)``.
    """
-    del env, env_ids, asset_cfg  # Unused.
+    del env, env_ids, asset_cfg, cooldown_s  # Unused at call time.
    dt = self._step_dt

    # Decrement timers for active envs.
@@ -468,11 +507,7 @@ class apply_body_impulse:
      )
      self._active[expired_ids] = False
      self._time_remaining[expired_ids] = 0.0
-      int_low, int_high = cooldown_s
-      self._interval_time_left[expired_ids] = (
-        torch.rand(len(expired_ids), device=self._device) * (int_high - int_low)
-        + int_low
-      )
+      self._interval_time_left[expired_ids] = self._sample_cooldown(len(expired_ids))

    # Decrement interval timers.
    self._interval_time_left -= dt
@@ -514,10 +549,7 @@ class apply_body_impulse:
    self._active[trigger_ids] = True

    # Resample interval timers.
-    int_low, int_high = cooldown_s
-    self._interval_time_left[trigger_ids] = (
-      torch.rand(n, device=self._device) * (int_high - int_low) + int_low
-    )
+    self._interval_time_left[trigger_ids] = self._sample_cooldown(n)

  def debug_vis(self, visualizer: DebugVisualizer) -> None:
    """Draw arrows for active impulse forces."""
@@ -553,13 +585,7 @@ class apply_body_impulse:
    if env_ids is None:
      env_ids = slice(None)

-    # Clear forces for reset envs.
-    if isinstance(env_ids, slice):
-      reset_ids = env_ids
-    else:
-      reset_ids = env_ids
-
-    if self._active[reset_ids].any():
+    if self._active[env_ids].any():
      if isinstance(env_ids, slice):
        active_ids = self._active.nonzero(as_tuple=False).squeeze(-1)
      else:
@@ -573,6 +599,7 @@ class apply_body_impulse:
          zeros, zeros, env_ids=active_ids, body_ids=self._body_ids
        )

-    self._time_remaining[reset_ids] = 0.0
-    self._interval_time_left[reset_ids] = 0.0
-    self._active[reset_ids] = False
+    n = self._num_envs if isinstance(env_ids, slice) else len(env_ids)
+    self._time_remaining[env_ids] = 0.0
+    self._interval_time_left[env_ids] = self._sample_cooldown(n)
+    self._active[env_ids] = False
@@ -106,6 +106,23 @@ def builtin_sensor(env: ManagerBasedRlEnv, sensor_name: str) -> torch.Tensor:
  return sensor.data


+def projected_gravity_from_sensor(
+  env: ManagerBasedRlEnv, sensor_name: str
+) -> torch.Tensor:
+  """Projected gravity from a ``framezaxis`` up-vector sensor.
+
+  The sensor is expected to output the world Z-axis expressed in the sensor's frame
+  (e.g. ``framezaxis`` with ``objtype=body objname=world`` and ``reftype=site``). That
+  is the body-frame "up" vector, so it is negated to point along gravity.
+
+  Unlike :func:`projected_gravity`, which uses the root body orientation, this reads
+  the sensor's site frame and therefore reflects IMU site pose randomization.
+  """
+  sensor = env.scene[sensor_name]
+  assert isinstance(sensor, BuiltinSensor)
+  return -sensor.data
+
+
 def height_scan(
  env: ManagerBasedRlEnv,
  sensor_name: str,
@@ -290,9 +290,12 @@ class EventManager(ManagerBase):
        fired = True
      elif mode == "reset":
        assert global_env_step_count is not None
+        # Reset events require concrete indices: callers (e.g. ManagerBasedRlEnv)
+        # resolve None to all environments upstream. Enforce that here so a future
+        # caller passing None fails loudly instead of leaking a slice into event
+        # functions, which only understand None or a tensor.
+        assert env_ids is not None, "reset events require concrete env_ids, got None"
        min_step_count = term_cfg.min_step_count_between_reset
-        if env_ids is None:
-          env_ids = slice(None)
        if min_step_count == 0:
          self._reset_term_last_triggered_step_id[index][env_ids] = (
            global_env_step_count
@@ -11,7 +11,7 @@ import numpy as np
 import torch

 from mjlab.entity import Entity, EntityCfg
-from mjlab.entity.entity import VariantMetadata
+from mjlab.entity.variants import VariantMetadata
 from mjlab.sensor import BuiltinSensor, RayCastSensor, Sensor, SensorCfg
 from mjlab.sensor.camera_sensor import CameraSensor
 from mjlab.sensor.sensor_context import SensorContext
@@ -59,7 +59,7 @@ class Scene:
    self._default_env_origins: torch.Tensor | None = None
    self._sensor_context: SensorContext | None = None

-    self._spec = mujoco.MjSpec.from_string(_SCENE_XML.read_text())
+    self._spec = mujoco.MjSpec.from_file(str(_SCENE_XML))
    if self._cfg.extent is not None:
      self._spec.stat.extent = self._cfg.extent
    self._add_terrain()
@@ -132,11 +132,6 @@ class Scene:
  def device(self) -> str:
    return self._device

-  @property
-  def has_mesh_variants(self) -> bool:
-    """True if any entity declares per-world mesh variants."""
-    return any(ent.variant_metadata is not None for ent in self._entities.values())
-
  def collect_variant_info(
    self,
  ) -> list[tuple[str, VariantMetadata]]:
@@ -59,7 +59,7 @@ def run_train(task_id: str, cfg: TrainConfig, log_dir: Path) -> None:
    os.environ["MUJOCO_EGL_DEVICE_ID"] = str(local_rank)
    device = f"cuda:{local_rank}"
    # Set seed to have diversity in different processes.
-    seed = cfg.agent.seed + local_rank
+    seed = cfg.agent.seed + rank

  configure_torch_backends()

@@ -197,7 +197,6 @@ def launch_training(task_id: str, args: TrainConfig | None = None):
    os.environ["CUDA_VISIBLE_DEVICES"] = ""
  else:
    os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, selected_gpus))
-  os.environ["MUJOCO_GL"] = "egl"

  if num_gpus <= 1:
    # CPU or single GPU: run directly without torchrunx.
@@ -428,7 +428,7 @@ class ContactSensor(Sensor[ContactData]):
    normal = data.normal
    tangent = data.tangent
    tangent2 = torch.cross(normal, tangent, dim=-1)
-    R = torch.stack([tangent, tangent2, normal], dim=-1)
+    R = torch.stack([normal, tangent, tangent2], dim=-1)

    has_contact = torch.norm(normal, dim=-1, keepdim=True) > 1e-8

@@ -435,7 +435,7 @@ class RayCastSensor(Sensor[RayCastData]):
    self._model: mjwarp.Model | None = None
    self._mj_model: mujoco.MjModel | None = None
    self._device: str | None = None
-    self._wp_device: wp.context.Device | None = None
+    self._wp_device: wp.Device | None = None

    # Per-frame info: list of (frame_type, obj_id, body_id).
    self._frame_infos: list[tuple[Literal["body", "site", "geom"], int, int]] = []
@@ -1,335 +0,0 @@
-"""Per-world mesh variant support.
-
-Sibling of :mod:`mjlab.sim.randomization`: that module expands singleton
-model fields into per-world arrays for DR; this one writes per-world
-arrays whose rows differ by mesh variant.
-"""
-
-from __future__ import annotations
-
-from dataclasses import dataclass
-from typing import Callable
-
-import mujoco
-import mujoco_warp as mjwarp
-import numpy as np
-import warp as wp
-
-from mjlab.entity.entity import BodyInertialMetadata, VariantMetadata
-
-# Fields that depend on mesh geometry and must be compiled per-variant.
-VARIANT_DEPENDENT_FIELDS = (
-  "geom_size",
-  "geom_rbound",
-  "geom_aabb",
-  "geom_pos",
-  "geom_quat",
-  "body_mass",
-  "body_subtreemass",
-  "body_inertia",
-  "body_invweight0",
-  "body_ipos",
-  "body_iquat",
-)
-
-
-@dataclass
-class MeshVariantResult:
-  """Output of :func:`build_mesh_variant_model`."""
-
-  wp_model: mjwarp.Model
-  mj_model: mujoco.MjModel
-  # Maps entity prefix -> array of variant indices per world.
-  world_to_variant: dict[str, np.ndarray]
-
-
-def _find_entity_mesh_geom_ids(
-  model: mujoco.MjModel,
-  entity_prefix: str,
-) -> list[int]:
-  """Find all mesh geom IDs belonging to an entity, including padding."""
-  named_ids: list[int] = []
-  for gid in range(model.ngeom):
-    gname = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_GEOM, gid)
-    if (
-      gname
-      and gname.startswith(entity_prefix)
-      and model.geom_type[gid] == mujoco.mjtGeom.mjGEOM_MESH
-    ):
-      named_ids.append(gid)
-  if not named_ids:
-    return []
-  # Include unnamed padding geoms on the same body.
-  body_id = model.geom_bodyid[named_ids[0]]
-  all_ids = set(named_ids)
-  for gid in range(model.ngeom):
-    if (
-      model.geom_bodyid[gid] == body_id
-      and model.geom_type[gid] == mujoco.mjtGeom.mjGEOM_MESH
-    ):
-      all_ids.add(gid)
-  return sorted(all_ids)
-
-
-def allocate_worlds(
-  weights: tuple[float, ...],
-  nworld: int,
-) -> list[int]:
-  """Assign worlds proportionally by weight (largest-remainder method).
-
-  Returns a list of length *nworld* containing variant indices. Weights
-  must be non-negative with at least one positive entry.
-  """
-  if any(w < 0 for w in weights):
-    raise ValueError(f"weights must be non-negative, got {weights}.")
-  total = sum(weights)
-  if total <= 0:
-    raise ValueError(f"weights must have a positive sum, got {weights}.")
-  quotas = [(w / total) * nworld for w in weights]
-  floors = [int(q) for q in quotas]
-  remainders = sorted(
-    ((quotas[i] - floors[i], i) for i in range(len(weights))),
-    key=lambda x: -x[0],
-  )
-  allocated = sum(floors)
-  for j in range(nworld - allocated):
-    floors[remainders[j][1]] += 1
-  assignment: list[int] = []
-  for idx, count in enumerate(floors):
-    assignment.extend([idx] * count)
-  return assignment
-
-
-def build_mesh_variant_model(
-  spec: mujoco.MjSpec,
-  nworld: int,
-  variant_info: list[tuple[str, VariantMetadata]],
-  configure_model: Callable[[mujoco.MjModel], None] | None = None,
-) -> MeshVariantResult:
-  """Build a warp Model with per-world mesh assignments.
-
-  Args:
-    spec: Scene spec (already merged with padded variant geoms).
-    nworld: Number of simulation worlds.
-    variant_info: List of ``(entity_prefix, metadata)`` pairs for
-      entities that have mesh variants.
-    configure_model: Optional callback to configure the compiled
-      MjModel before ``put_model`` (e.g., setting solver options).
-
-  Returns:
-    A :class:`MeshVariantResult` containing the warp model, host
-    model, and per-entity world-to-variant mappings.
-  """
-  spec = spec.copy()
-  model = spec.compile()
-  if configure_model is not None:
-    configure_model(model)
-
-  # Start from base dataid tiled for all worlds.
-  base_dataid = model.geom_dataid.copy()
-  dataid_table = np.tile(base_dataid, (nworld, 1))
-
-  world_to_variant: dict[str, np.ndarray] = {}
-
-  for entity_prefix, metadata in variant_info:
-    # Allocate worlds by weight.
-    assignment = allocate_worlds(metadata.variant_weights, nworld)
-    w2v = np.array(assignment, dtype=np.int32)
-    world_to_variant[entity_prefix] = w2v
-
-    mesh_geom_ids = _find_entity_mesh_geom_ids(model, entity_prefix)
-    nslots = len(mesh_geom_ids)
-
-    # Resolve every (variant, slot) -> mesh_id once. Mesh names in the merged
-    # spec are variant-prefixed ("mug/visual_mesh"); after attaching to the
-    # scene they also carry the entity prefix ("object/mug/visual_mesh").
-    # Padding slots are -1.
-    nvariants = len(metadata.variant_mesh_names)
-    variant_slot_ids = np.full((nvariants, nslots), -1, dtype=np.int64)
-    for v_idx, mesh_names in enumerate(metadata.variant_mesh_names):
-      for slot in range(min(nslots, len(mesh_names))):
-        name = mesh_names[slot]
-        if name is None:
-          continue
-        full = f"{entity_prefix}{name}"
-        mid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_MESH, full)
-        if mid < 0:
-          variant_label = metadata.variant_names[v_idx]
-          raise ValueError(
-            f"Mesh '{full}' (variant '{variant_label}', slot {slot}) "
-            f"not found in compiled model."
-          )
-        variant_slot_ids[v_idx, slot] = mid
-
-    # Vectorized scatter: row-select by variant assignment, write into the
-    # mesh-geom columns of the per-world dataid table.
-    dataid_table[:, mesh_geom_ids] = variant_slot_ids[w2v]
-
-  # Build warp model.
-  m = mjwarp.put_model(model)
-  m.geom_dataid = wp.array(dataid_table, dtype=int)
-
-  # Populate dependent per-world fields.
-  _populate_dependent_fields(
-    m, spec, model, dataid_table, nworld, variant_info, world_to_variant
-  )
-
-  return MeshVariantResult(
-    wp_model=m,
-    mj_model=model,
-    world_to_variant=world_to_variant,
-  )
-
-
-def _populate_dependent_fields(
-  m: mjwarp.Model,
-  spec: mujoco.MjSpec,
-  padded_model: mujoco.MjModel,
-  dataid_table: np.ndarray,
-  nworld: int,
-  variant_info: list[tuple[str, VariantMetadata]],
-  world_to_variant: dict[str, np.ndarray],
-) -> None:
-  """Compile each unique variant and write per-world dependent fields.
-
-  Each unique variant is compiled from a fresh ``spec.copy()``; the input
-  ``spec`` is not mutated.
-  """
-  # Find unique dataid rows.
-  unique_rows: dict[tuple[int, ...], int] = {}
-  for w in range(nworld):
-    key = tuple(dataid_table[w])
-    if key not in unique_rows:
-      unique_rows[key] = w
-
-  # Map padded_model geom IDs to geom names (stable across spec copies).
-  geom_id_to_name: dict[int, str] = {}
-  for g in spec.geoms:
-    if not g.name:
-      continue
-    gid = mujoco.mj_name2id(padded_model, mujoco.mjtObj.mjOBJ_GEOM, g.name)
-    if gid >= 0:
-      geom_id_to_name[gid] = g.name
-
-  # Collect all variant geom IDs in padded_model.
-  all_variant_geom_ids: set[int] = set()
-  for entity_prefix, _ in variant_info:
-    all_variant_geom_ids.update(_find_entity_mesh_geom_ids(padded_model, entity_prefix))
-
-  # Bodies any variant marks as explicit-inertial: must be reset on the
-  # fresh spec copy before applying this variant's inertials. Variants
-  # without an explicit inertial fall back to MuJoCo's mesh-derived path
-  # during compile, so we clear the diagonal inertial fields. Do NOT
-  # assign ``body.fullinertia``: any assignment (even zeros) flags the
-  # field as user-specified and ``spec.compile()`` then rejects it as
-  # conflicting with ``body.inertia``.
-  variant_inertial_body_names: set[str] = set()
-  for entity_prefix, metadata in variant_info:
-    for variant_inertials in metadata.variant_body_inertials:
-      for inertial in variant_inertials:
-        variant_inertial_body_names.add(f"{entity_prefix}{inertial.body_name}")
-
-  # Compile each unique variant from a fresh spec copy.
-  compiled_variants: dict[tuple[int, ...], mujoco.MjModel] = {}
-  for key, first_world in unique_rows.items():
-    variant_spec = spec.copy()
-    geoms_by_name = {g.name: g for g in variant_spec.geoms if g.name}
-    bodies_by_name = {b.name: b for b in variant_spec.bodies if b.name}
-
-    # Apply this variant's mesh selection per geom slot.
-    for gid in all_variant_geom_ids:
-      name = geom_id_to_name.get(gid)
-      if name is None:
-        continue
-      geom = geoms_by_name[name]
-      mesh_id = int(dataid_table[first_world, gid])
-      if mesh_id >= 0:
-        mesh_name = mujoco.mj_id2name(padded_model, mujoco.mjtObj.mjOBJ_MESH, mesh_id)
-        geom.meshname = mesh_name
-        geom.contype = 1
-        geom.conaffinity = 1
-      else:
-        geom.contype = 0
-        geom.conaffinity = 0
-        geom.mass = 0.0
-
-    for body_name in variant_inertial_body_names:
-      body = bodies_by_name.get(body_name)
-      if body is None:
-        continue
-      body.explicitinertial = 0
-      body.mass = 0.0
-      body.inertia = np.zeros(3, dtype=np.float64)
-      body.ipos = np.zeros(3, dtype=np.float64)
-      body.iquat = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64)
-
-    for entity_prefix, metadata in variant_info:
-      variant_idx = int(world_to_variant[entity_prefix][first_world])
-      if variant_idx >= len(metadata.variant_body_inertials):
-        continue
-      for inertial in metadata.variant_body_inertials[variant_idx]:
-        _apply_body_inertial(
-          bodies_by_name,
-          f"{entity_prefix}{inertial.body_name}",
-          inertial,
-        )
-
-    compiled_variants[key] = variant_spec.compile()
-
-  # Build per-world numpy arrays.
-  ngeom = padded_model.ngeom
-  nbody = padded_model.nbody
-
-  geom_size = np.zeros((nworld, ngeom, 3), dtype=np.float32)
-  geom_rbound = np.zeros((nworld, ngeom), dtype=np.float32)
-  geom_aabb = np.zeros((nworld, ngeom, 2, 3), dtype=np.float32)
-  geom_pos = np.zeros((nworld, ngeom, 3), dtype=np.float32)
-  geom_quat = np.zeros((nworld, ngeom, 4), dtype=np.float32)
-  body_mass = np.zeros((nworld, nbody), dtype=np.float32)
-  body_subtreemass = np.zeros((nworld, nbody), dtype=np.float32)
-  body_inertia = np.zeros((nworld, nbody, 3), dtype=np.float32)
-  body_invweight0 = np.zeros((nworld, nbody, 2), dtype=np.float32)
-  body_ipos = np.zeros((nworld, nbody, 3), dtype=np.float32)
-  body_iquat = np.zeros((nworld, nbody, 4), dtype=np.float32)
-
-  for w in range(nworld):
-    key = tuple(dataid_table[w])
-    ref = compiled_variants[key]
-    geom_size[w] = ref.geom_size
-    geom_rbound[w] = ref.geom_rbound
-    geom_aabb[w] = ref.geom_aabb.reshape(ngeom, 2, 3)
-    geom_pos[w] = ref.geom_pos
-    geom_quat[w] = ref.geom_quat
-    body_mass[w] = ref.body_mass
-    body_subtreemass[w] = ref.body_subtreemass
-    body_inertia[w] = ref.body_inertia
-    body_invweight0[w] = ref.body_invweight0
-    body_ipos[w] = ref.body_ipos
-    body_iquat[w] = ref.body_iquat
-
-  m.geom_size = wp.array(geom_size, dtype=wp.vec3)
-  m.geom_rbound = wp.array(geom_rbound, dtype=float)
-  m.geom_aabb = wp.array(geom_aabb, dtype=wp.vec3)
-  m.geom_pos = wp.array(geom_pos, dtype=wp.vec3)
-  m.geom_quat = wp.array(geom_quat, dtype=wp.quat)
-  m.body_mass = wp.array(body_mass, dtype=float)
-  m.body_subtreemass = wp.array(body_subtreemass, dtype=float)
-  m.body_inertia = wp.array(body_inertia, dtype=wp.vec3)
-  m.body_invweight0 = wp.array(body_invweight0, dtype=wp.vec2)
-  m.body_ipos = wp.array(body_ipos, dtype=wp.vec3)
-  m.body_iquat = wp.array(body_iquat, dtype=wp.quat)
-
-
-def _apply_body_inertial(
-  bodies_by_name: dict[str, mujoco.MjsBody],
-  body_name: str,
-  inertial: BodyInertialMetadata,
-) -> None:
-  body = bodies_by_name.get(body_name)
-  if body is None:
-    raise ValueError(f"Body '{body_name}' not found in compiled variant spec.")
-  body.explicitinertial = 1
-  body.mass = inertial.mass
-  body.ipos = np.asarray(inertial.ipos, dtype=np.float64)
-  body.inertia = np.asarray(inertial.inertia, dtype=np.float64)
-  body.iquat = np.asarray(inertial.iquat, dtype=np.float64)
@@ -10,14 +10,14 @@ import mujoco_warp as mjwarp
 import torch
 import warp as wp

+from mjlab.entity.variants import VARIANT_DEPENDENT_FIELDS, build_variant_model
 from mjlab.managers.event_manager import RecomputeLevel
-from mjlab.sim.mesh_variants import VARIANT_DEPENDENT_FIELDS, build_mesh_variant_model
 from mjlab.sim.randomization import expand_model_fields
 from mjlab.sim.sim_data import TorchArray, WarpBridge
 from mjlab.utils.nan_guard import NanGuard, NanGuardCfg

 if TYPE_CHECKING:
-  from mjlab.entity.entity import VariantMetadata
+  from mjlab.entity.variants import VariantMetadata
  from mjlab.sensor.sensor_context import SensorContext

 # Type aliases for better IDE support while maintaining runtime compatibility
@@ -246,7 +246,7 @@ class Simulation:
    they are rendering or inspecting.
    """
    with wp.ScopedDevice(self.wp_device):
-      result = build_mesh_variant_model(
+      result = build_variant_model(
        spec,
        self.num_envs,
        variant_info,
@@ -275,9 +275,10 @@ class Simulation:
    # viewer syncs them per-world.
    self._expanded_fields.update(VARIANT_DEPENDENT_FIELDS)
    self._expanded_fields.add("geom_dataid")
+    self._expanded_fields.add("geom_matid")

    # Stash variant assignments as torch tensors keyed by bare entity name
-    # (mesh_variants emits "<name>/" prefixes; strip the trailing slash for
+    # (build_variant_model emits "<name>/" prefixes; strip the trailing slash for
    # the public API).
    for prefix, arr in result.world_to_variant.items():
      key = prefix.rstrip("/")
@@ -525,7 +526,7 @@ class Simulation:
    if not self.wp_device.is_cuda:
      return False

-    driver_ver = wp.context.runtime.driver_version
+    driver_ver = wp.get_cuda_driver_version()
    has_mempool = wp.is_mempool_enabled(self.wp_device)

    if driver_ver is None:
@@ -14,9 +14,10 @@ def compute_mpkpe(command: MotionCommand) -> torch.Tensor:
  """Compute Mean Per-Keybody Position Error (MPKPE).

  MPKPE measures the average Euclidean distance between the reference and
-  actual positions of all key bodies in world frame.
+  actual key body positions in the global world frame. It captures all
+  tracking error, including global translation and heading drift.
  """
-  pos_error = command.body_pos_relative_w - command.robot_body_pos_w
+  pos_error = command.body_pos_w - command.robot_body_pos_w
  per_body_error = torch.norm(pos_error, dim=-1)  # (num_envs, num_bodies)
  return per_body_error.mean(dim=-1)  # (num_envs,)

@@ -24,29 +25,25 @@ def compute_mpkpe(command: MotionCommand) -> torch.Tensor:
 def compute_root_relative_mpkpe(command: MotionCommand) -> torch.Tensor:
  """Compute Root-relative Mean Per-Keybody Position Error (R-MPKPE).

-  R-MPKPE measures pose error independent of global drift by computing
-  positions relative to the root/anchor body.
+  R-MPKPE measures intrinsic pose error independent of global drift. It
+  uses ``body_pos_relative_w``, the reference re-anchored to the robot's
+  current root position and heading each step (the same quantity the
+  tracking reward optimizes), so both global translation and yaw drift are
+  removed and only the local body pose error remains.
  """
-  # Compute reference positions relative to reference anchor.
-  ref_anchor_pos = command.anchor_pos_w.unsqueeze(1)  # (num_envs, 1, 3)
-  ref_rel_pos = command.body_pos_w - ref_anchor_pos  # (num_envs, num_bodies, 3)
-
-  # Compute robot positions relative to robot anchor.
-  robot_anchor_pos = command.robot_anchor_pos_w.unsqueeze(1)  # (num_envs, 1, 3)
-  robot_rel_pos = (
-    command.robot_body_pos_w - robot_anchor_pos
-  )  # (num_envs, num_bodies, 3)
-
-  # Compute error between relative positions.
-  pos_error = ref_rel_pos - robot_rel_pos
+  pos_error = command.body_pos_relative_w - command.robot_body_pos_w
  per_body_error = torch.norm(pos_error, dim=-1)  # (num_envs, num_bodies)
  return per_body_error.mean(dim=-1)  # (num_envs,)


 def compute_joint_velocity_error(command: MotionCommand) -> torch.Tensor:
-  """Compute average joint velocity error."""
+  """Compute root-mean-square joint velocity error.
+
+  Uses an RMS over joints (rather than a raw L2 norm) so the value is a
+  per-joint quantity, comparable across robots with different DOF counts.
+  """
  vel_error = command.joint_vel - command.robot_joint_vel
-  return torch.norm(vel_error, dim=-1)  # (num_envs,)
+  return torch.sqrt(torch.mean(vel_error**2, dim=-1))  # (num_envs,)


 def compute_ee_position_error(
@@ -93,6 +90,18 @@ def _get_body_indices(
    body_names: Names of bodies to find.

  Returns:
-    List of indices into command.cfg.body_names.
+    List of indices into command.cfg.body_names, in the order requested.
+
+  Raises:
+    ValueError: If any requested body name is not tracked by the command.
+      Silently dropping unknown names would otherwise report a spurious
+      zero error for misconfigured end-effector lists.
  """
-  return [i for i, name in enumerate(command.cfg.body_names) if name in body_names]
+  name_to_index = {name: i for i, name in enumerate(command.cfg.body_names)}
+  missing = [name for name in body_names if name not in name_to_index]
+  if missing:
+    raise ValueError(
+      f"Body names {missing} are not tracked by the command. "
+      f"Available bodies: {tuple(command.cfg.body_names)}."
+    )
+  return [name_to_index[name] for name in body_names]
@@ -95,7 +95,9 @@ class MotionTrackingOnPolicyRunner(MjlabOnPolicyRunner):
    try:
      self.export_policy_to_onnx(str(policy_dir), filename)
      run_name: str = (
-        wandb.run.name if self.logger.logger_type == "wandb" and wandb.run else "local"
+        wandb.run.name
+        if self.logger.logger_type in ("wandb", "WandbLogWriter") and wandb.run
+        else "local"
      )  # type: ignore[assignment]
      metadata = get_base_metadata(self.env.unwrapped, run_name)
      motion_term = cast(
@@ -108,7 +110,10 @@ class MotionTrackingOnPolicyRunner(MjlabOnPolicyRunner):
        }
      )
      attach_metadata_to_onnx(str(onnx_path), metadata)
-      if self.logger.logger_type in ["wandb"] and self.cfg["upload_model"]:
+      if (
+        self.logger.logger_type in ("wandb", "WandbLogWriter")
+        and self.cfg["upload_model"]
+      ):
        wandb.save(str(onnx_path), base_path=str(policy_dir))
        if self.registry_name is not None:
          wandb.run.use_artifact(self.registry_name)  # type: ignore
@@ -6,6 +6,7 @@ import json
 import sys
 from dataclasses import asdict, dataclass
 from pathlib import Path
+from types import SimpleNamespace
 from typing import cast

 import torch
@@ -97,35 +98,61 @@ def run_evaluate(task_id: str, cfg: EvaluateConfig) -> dict[str, float]:
  all_joint_vel_error: list[torch.Tensor] = []
  all_ee_pos_error: list[torch.Tensor] = []
  all_ee_ori_error: list[torch.Tensor] = []
+  all_active: list[torch.Tensor] = []

  done_envs = torch.zeros(cfg.num_envs, dtype=torch.bool, device=device)
  success = torch.zeros(cfg.num_envs, dtype=torch.bool, device=device)

  obs = env.get_observations()
-  env.unwrapped.command_manager.compute(dt=env.unwrapped.step_dt)

  print(f"[INFO] Running {cfg.num_envs} evaluation episodes...")

  step = 0
  while not done_envs.all():
+    # Snapshot the reference frame the upcoming step will be scored against.
+    # env.step computes the reward (against the current reference) and only
+    # afterwards advances the command's motion frame, so reading the
+    # reference after stepping would pair the robot with the *next* frame.
+    # We snapshot here and pair it with the post-step robot state below,
+    # matching how the reward is computed.
+    ref = SimpleNamespace(
+      num_envs=command.num_envs,
+      device=command.device,
+      cfg=command.cfg,
+      body_pos_w=command.body_pos_w.clone(),
+      body_pos_relative_w=command.body_pos_relative_w.clone(),
+      body_quat_relative_w=command.body_quat_relative_w.clone(),
+      joint_vel=command.joint_vel.clone(),
+    )
+
    with torch.no_grad():
      actions = policy(obs)
    obs, _, dones, _ = env.step(actions)

-    # Compute metrics for active envs.
+    # Pair the snapshotted reference with the post-step robot state.
+    ref.robot_body_pos_w = command.robot_body_pos_w
+    ref.robot_body_quat_w = command.robot_body_quat_w
+    ref.robot_joint_vel = command.robot_joint_vel
+    ref_command = cast(MotionCommand, ref)
+
+    # Accumulate metrics for envs still running this step. active.any() is
+    # always true here: the loop runs only while some env is not done, and
+    # done_envs is updated below after this point.
    active = ~done_envs
-    if active.any():
-      all_mpkpe.append(torch.where(active, compute_mpkpe(command), 0.0))
-      all_r_mpkpe.append(torch.where(active, compute_root_relative_mpkpe(command), 0.0))
-      all_joint_vel_error.append(
-        torch.where(active, compute_joint_velocity_error(command), 0.0)
-      )
-      all_ee_pos_error.append(
-        torch.where(active, compute_ee_position_error(command, ee_body_names), 0.0)
-      )
-      all_ee_ori_error.append(
-        torch.where(active, compute_ee_orientation_error(command, ee_body_names), 0.0)
-      )
+    all_active.append(active.float())
+    all_mpkpe.append(torch.where(active, compute_mpkpe(ref_command), 0.0))
+    all_r_mpkpe.append(
+      torch.where(active, compute_root_relative_mpkpe(ref_command), 0.0)
+    )
+    all_joint_vel_error.append(
+      torch.where(active, compute_joint_velocity_error(ref_command), 0.0)
+    )
+    all_ee_pos_error.append(
+      torch.where(active, compute_ee_position_error(ref_command, ee_body_names), 0.0)
+    )
+    all_ee_ori_error.append(
+      torch.where(active, compute_ee_orientation_error(ref_command, ee_body_names), 0.0)
+    )

    # Track completions.
    terminated = env.unwrapped.termination_manager.terminated
@@ -142,7 +169,7 @@ def run_evaluate(task_id: str, cfg: EvaluateConfig) -> dict[str, float]:
      )
    step += 1

-  # Compute mean metrics.
+  # Compute mean metrics over the steps each env was active.
  stacks = [
    all_mpkpe,
    all_r_mpkpe,
@@ -151,7 +178,7 @@ def run_evaluate(task_id: str, cfg: EvaluateConfig) -> dict[str, float]:
    all_ee_ori_error,
  ]
  stacks = [torch.stack(s, dim=0) for s in stacks]
-  active_steps = (stacks[0] != 0).sum(dim=0).float().clamp(min=1)
+  active_steps = torch.stack(all_active, dim=0).sum(dim=0).clamp(min=1)
  means = [s.sum(dim=0) / active_steps for s in stacks]

  metrics = {
@@ -24,14 +24,32 @@ from mjlab.terrains.terrain_generator import (
 )
 from mjlab.terrains.utils import find_flat_patches_from_heightfield

+# Smallest positive hfield elevation/base size, in meters. MuJoCo rejects
+# non-positive hfield sizes, so flat heightfields (difficulty 0) are clamped to
+# this instead of zero.
+_MIN_HFIELD_HEIGHT = 1e-3
+
+# Physical height (meters) that maps to full color saturation. Heights are
+# colored on this fixed absolute scale rather than normalized per patch, so a
+# given height reads the same color across every terrain and small-amplitude
+# terrain stays gently tinted instead of stretching into rainbow noise.
+_COLOR_SCALE = 0.75
+

 def color_by_height(
  spec: mujoco.MjSpec,
  noise: np.ndarray,
  unique_id: str,
-  normalized_elevation: np.ndarray,
+  physical_heights: np.ndarray,
  texture_size: int = 128,
 ) -> str:
+  """Build a height-colored texture for a heightfield.
+
+  Diverging colormap anchored at the ground plane (z=0): cool blue below ground,
+  green at z=0, warm red above. ``physical_heights`` is the surface height of
+  each cell in meters relative to z=0; it is colored on the fixed ``_COLOR_SCALE``
+  so color encodes absolute height consistently across all terrains.
+  """
  texture_name = f"hf_texture_{unique_id}"
  texture = spec.add_texture(
    name=texture_name,
@@ -40,16 +58,20 @@ def color_by_height(
    height=texture_size,
  )

-  texture_elevation = ndimage.zoom(
-    normalized_elevation,
+  texture_height = ndimage.zoom(
+    physical_heights,
    (texture_size / noise.shape[0], texture_size / noise.shape[1]),
    order=1,
  )
-  texture_elevation = np.asarray(texture_elevation)
+  texture_height = np.asarray(texture_height)

-  hue = 0.5 - texture_elevation * 0.45
-  saturation = 0.6 - texture_elevation * 0.2
-  value = 0.4 + texture_elevation * 0.3
+  # Signed deviation from the ground plane in [-1, 1] on a fixed absolute scale.
+  signed = np.clip(texture_height / _COLOR_SCALE, -1.0, 1.0)
+
+  # signed=+1 -> hue 0.0 (red, high), 0 -> 0.33 (green, ground), -1 -> 0.66 (blue, low).
+  hue = 0.33 - 0.33 * signed
+  saturation = 0.45 + 0.25 * np.abs(signed)
+  value = 0.45 + 0.25 * np.abs(signed)

  c = value * saturation
  x = c * (1 - np.abs((hue * 6) % 2 - 1))
@@ -326,7 +348,8 @@ class HfPyramidSlopedTerrainCfg(SubTerrainCfg):
    else:
      hfield_z_offset = 0

-    material_name = color_by_height(spec, noise, unique_id, normalized_elevation)
+    physical_heights = hfield_z_offset + normalized_elevation * max_physical_height
+    material_name = color_by_height(spec, noise, unique_id, physical_heights)

    hfield_geom = body.add_geom(
      type=mujoco.mjtGeom.mjGEOM_HFIELD,
@@ -378,14 +401,24 @@ class HfRandomUniformTerrainCfg(SubTerrainCfg):
  border_width: float = 0.0
  """Width of the flat border around the terrain edges, in meters. Must be >=
  horizontal_scale if non-zero."""
+  scale_with_difficulty: bool = False
+  """If False (default), the roughness is fixed and ``difficulty`` is ignored,
+  matching upstream behavior. If True, the noise amplitude scales linearly with
+  difficulty (flat at 0, full ``noise_range`` at 1) so the terrain progresses in
+  a curriculum."""

  def function(
    self, difficulty: float, spec: mujoco.MjSpec, rng: np.random.Generator
  ) -> TerrainOutput:
-    del difficulty  # Unused.
-
    body = spec.body("terrain")

+    # When difficulty scaling is enabled, ramp the noise amplitude from flat (0)
+    # to the full configured range (1). Otherwise use the full range regardless
+    # of difficulty (difficulty is ignored).
+    scale = difficulty if self.scale_with_difficulty else 1.0
+    noise_lo = self.noise_range[0] * scale
+    noise_hi = self.noise_range[1] * scale
+
    if self.border_width > 0 and self.border_width < self.horizontal_scale:
      raise ValueError(
        f"Border width ({self.border_width}) must be >= horizontal scale "
@@ -419,8 +452,8 @@ class HfRandomUniformTerrainCfg(SubTerrainCfg):
      width_downsampled = int(inner_size[0] / downsampled_scale)
      length_downsampled = int(inner_size[1] / downsampled_scale)

-      height_min = int(self.noise_range[0] / self.vertical_scale)
-      height_max = int(self.noise_range[1] / self.vertical_scale)
+      height_min = int(noise_lo / self.vertical_scale)
+      height_max = int(noise_hi / self.vertical_scale)
      height_step = int(self.noise_step / self.vertical_scale)

      height_range = np.arange(height_min, height_max + height_step, height_step)
@@ -443,8 +476,8 @@ class HfRandomUniformTerrainCfg(SubTerrainCfg):
    else:
      width_downsampled = int(self.size[0] / downsampled_scale)
      length_downsampled = int(self.size[1] / downsampled_scale)
-      height_min = int(self.noise_range[0] / self.vertical_scale)
-      height_max = int(self.noise_range[1] / self.vertical_scale)
+      height_min = int(noise_lo / self.vertical_scale)
+      height_max = int(noise_hi / self.vertical_scale)
      height_step = int(self.noise_step / self.vertical_scale)

      height_range = np.arange(height_min, height_max + height_step, height_step)
@@ -489,7 +522,8 @@ class HfRandomUniformTerrainCfg(SubTerrainCfg):
      userdata=normalized_elevation.flatten().astype(np.float32).tolist(),
    )

-    material_name = color_by_height(spec, noise, unique_id, normalized_elevation)
+    physical_heights = normalized_elevation * max_physical_height
+    material_name = color_by_height(spec, noise, unique_id, physical_heights)

    hfield_geom = body.add_geom(
      type=mujoco.mjtGeom.mjGEOM_HFIELD,
@@ -498,7 +532,7 @@ class HfRandomUniformTerrainCfg(SubTerrainCfg):
      material=material_name,
    )

-    spawn_height = (self.noise_range[0] + self.noise_range[1]) / 2
+    spawn_height = (noise_lo + noise_hi) / 2
    origin = np.array([self.size[0] / 2, self.size[1] / 2, spawn_height])

    flat_patches = _compute_flat_patches(
@@ -616,7 +650,11 @@ class HfWaveTerrainCfg(SubTerrainCfg):
      userdata=normalized_elevation.flatten().astype(np.float32).tolist(),
    )

-    material_name = color_by_height(spec, noise, unique_id, normalized_elevation)
+    # The wave oscillates around z=0 (geom is offset down by half the range).
+    physical_heights = (
+      normalized_elevation * max_physical_height - max_physical_height / 2
+    )
+    material_name = color_by_height(spec, noise, unique_id, physical_heights)

    hfield_geom = body.add_geom(
      type=mujoco.mjtGeom.mjGEOM_HFIELD,
@@ -783,7 +821,9 @@ class HfDiscreteObstaclesTerrainCfg(SubTerrainCfg):
    else:
      hfield_z_offset = 0

-    material_name = color_by_height(spec, noise, unique_id, normalized_elevation)
+    # Physical surface height per cell (pits negative, bumps positive about z=0).
+    physical_heights = hfield_z_offset + normalized_elevation * max_physical_height
+    material_name = color_by_height(spec, noise, unique_id, physical_heights)

    hfield_geom = body.add_geom(
      type=mujoco.mjtGeom.mjGEOM_HFIELD,
@@ -887,8 +927,13 @@ class HfPerlinNoiseTerrainCfg(SubTerrainCfg):
      noise_range = noise_max - noise_min if noise_max > noise_min else 1.0
      normalized_elevation = ((noise_raw - noise_min) / noise_range).astype(np.float32)

-    max_physical_height = target_height
-    base_thickness = max_physical_height * self.base_thickness_ratio
+    # MuJoCo requires positive hfield elevation and base sizes. At difficulty 0
+    # (target_height == 0) the surface is flat; clamp to a small positive height
+    # so compilation does not fail with "size parameter is not positive".
+    max_physical_height = max(target_height, _MIN_HFIELD_HEIGHT)
+    base_thickness = max(
+      max_physical_height * self.base_thickness_ratio, _MIN_HFIELD_HEIGHT
+    )

    unique_id = uuid.uuid4().hex
    field = spec.add_hfield(
@@ -904,8 +949,9 @@ class HfPerlinNoiseTerrainCfg(SubTerrainCfg):
      userdata=normalized_elevation.flatten().tolist(),
    )

+    physical_heights = normalized_elevation * max_physical_height
    material_name = color_by_height(
-      spec, normalized_elevation, unique_id, normalized_elevation
+      spec, normalized_elevation, unique_id, physical_heights
    )

    hfield_geom = body.add_geom(
@@ -11,7 +11,6 @@ References:
 from __future__ import annotations

 from dataclasses import dataclass
-from typing import Tuple

 import mujoco
 import numpy as np
@@ -23,30 +22,19 @@ from mjlab.terrains.terrain_generator import (
 )
 from mjlab.terrains.utils import make_border, make_plane
 from mjlab.utils.color import (
-  HSV,
  brand_ramp,
-  clamp,
  darken_rgba,
-  hsv_to_rgb,
-  rgb_to_hsv,
 )

 _MUJOCO_BLUE = (0.20, 0.45, 0.95)
 _MUJOCO_RED = (0.90, 0.30, 0.30)
 _MUJOCO_GREEN = (0.25, 0.80, 0.45)

-
-def _get_platform_color(
-  base_rgb: Tuple[float, float, float],
-  desaturation_factor: float = 0.4,
-  lightening_factor: float = 0.25,
-) -> Tuple[float, float, float, float]:
-  hsv = rgb_to_hsv(base_rgb)
-  new_s = hsv.s * desaturation_factor
-  new_v = clamp(hsv.v + lightening_factor)
-  new_hsv = HSV(hsv.h, new_s, new_v)
-  r, g, b = hsv_to_rgb(new_hsv)
-  return (r, g, b, 1.0)
+# Minimum vertical extent of a flat border frame, in meters. The border top sits
+# flush at z=0 and extends downward, so this depth is not visible from above; it
+# only guarantees the frame is solid (never a degenerate zero-height geom) when
+# the step height collapses to zero at difficulty 0.
+_MIN_BORDER_HEIGHT = 0.05


@dataclass(kw_only=True)
@@ -107,13 +95,18 @@ class BoxPyramidStairsTerrainCfg(SubTerrainCfg):
    border_rgba = darken_rgba(first_step_rgba, 0.85)

    if self.border_width > 0.0 and not self.holes:
-      border_center = (0.5 * self.size[0], 0.5 * self.size[1], -step_height / 2)
+      # Decouple the border's vertical extent from step_height so difficulty 0
+      # (step_height == 0) still produces a solid, gap-free frame instead of
+      # being skipped or generating degenerate zero-height geoms. The top stays
+      # flush with the ground at z=0.
+      border_height = max(step_height, _MIN_BORDER_HEIGHT)
+      border_center = (0.5 * self.size[0], 0.5 * self.size[1], -border_height / 2)
      border_inner_size = (
        self.size[0] - 2 * self.border_width,
        self.size[1] - 2 * self.border_width,
      )
      border_boxes = make_border(
-        body, self.size, border_inner_size, step_height, border_center
+        body, self.size, border_inner_size, border_height, border_center
      )
      boxes.extend(border_boxes)
      for _ in range(len(border_boxes)):
@@ -280,13 +273,16 @@ class BoxInvertedPyramidStairsTerrainCfg(BoxPyramidStairsTerrainCfg):
    border_rgba = darken_rgba(first_step_rgba, 0.85)

    if self.border_width > 0.0 and not self.holes:
-      border_center = (0.5 * self.size[0], 0.5 * self.size[1], -0.5 * step_height)
+      # See BoxPyramidStairsTerrainCfg: keep the border solid and flush at z=0
+      # even when step_height collapses to 0 at difficulty 0.
+      border_height = max(step_height, _MIN_BORDER_HEIGHT)
+      border_center = (0.5 * self.size[0], 0.5 * self.size[1], -0.5 * border_height)
      border_inner_size = (
        self.size[0] - 2 * self.border_width,
        self.size[1] - 2 * self.border_width,
      )
      border_boxes = make_border(
-        body, self.size, border_inner_size, step_height, border_center
+        body, self.size, border_inner_size, border_height, border_center
      )
      boxes.extend(border_boxes)
      for _ in range(len(border_boxes)):
@@ -546,8 +542,7 @@ class BoxRandomGridTerrainCfg(SubTerrainCfg):
      pos=(self.size[0] / 2, self.size[1] / 2, platform_center_z),
    )
    boxes_list.append(box)
-    platform_rgba = _get_platform_color(_MUJOCO_GREEN)
-    box_colors.append(platform_rgba)
+    box_colors.append(brand_ramp(_MUJOCO_GREEN, 0.5))

    origin = np.array([self.size[0] / 2, self.size[1] / 2, grid_height])

@@ -575,6 +570,22 @@ class BoxRandomGridTerrainCfg(SubTerrainCfg):
    half_border_width = border_width / 2
    neg_half_terrain = -terrain_height / 2

+    # Mark cells under the center platform as visited so they are never emitted
+    # or merged; the platform box covers that region and would otherwise z-fight
+    # with the cells beneath it.
+    platform_half = self.platform_width / 2
+    terrain_center = self.size[0] / 2
+    platform_min = terrain_center - platform_half
+    platform_max = terrain_center + platform_half
+    for i in range(num_boxes_x):
+      cx = half_border_width + (i + 0.5) * self.grid_width
+      if not (platform_min <= cx <= platform_max):
+        continue
+      for j in range(num_boxes_y):
+        cy = half_border_width + (j + 0.5) * self.grid_width
+        if platform_min <= cy <= platform_max:
+          visited[i, j] = True
+
    # Quantize heights to create more merging opportunities
    quantized_heights = (
      np.round(height_map / self.height_merge_threshold) * self.height_merge_threshold
@@ -588,7 +599,12 @@ class BoxRandomGridTerrainCfg(SubTerrainCfg):
        # Find rectangular region with similar height
        height = quantized_heights[i, j]

-        normalized_height = (height + grid_height) / (2 * grid_height)
+        # grid_height == 0 (difficulty 0) means a flat grid; use the midpoint
+        # color and avoid dividing by zero.
+        if grid_height > 0:
+          normalized_height = (height + grid_height) / (2 * grid_height)
+        else:
+          normalized_height = 0.5
        t = float(np.clip(normalized_height, 0.0, 1.0))
        rgba = brand_ramp(_MUJOCO_GREEN, t)

@@ -653,14 +669,10 @@ class BoxRandomGridTerrainCfg(SubTerrainCfg):
    half_border_width = border_width / 2
    neg_half_terrain = -terrain_height / 2

-    if self.holes:
-      platform_half = self.platform_width / 2
-      terrain_center = self.size[0] / 2
-      platform_min = terrain_center - platform_half
-      platform_max = terrain_center + platform_half
-    else:
-      platform_min = None
-      platform_max = None
+    platform_half = self.platform_width / 2
+    terrain_center = self.size[0] / 2
+    platform_min = terrain_center - platform_half
+    platform_max = terrain_center + platform_half

    for i in range(num_boxes_x):
      box_center_x = half_border_width + (i + 0.5) * self.grid_width
@@ -678,11 +690,24 @@ class BoxRandomGridTerrainCfg(SubTerrainCfg):
          if not (in_x_strip or in_y_strip):
            continue

+        # Skip cells under the center platform so the platform is the only
+        # geometry there. Otherwise the platform box sits on top of these cells
+        # and the coplanar faces z-fight.
+        if (platform_min <= box_center_x <= platform_max) and (
+          platform_min <= box_center_y <= platform_max
+        ):
+          continue
+
        height_noise = height_map[i, j]
        box_height = terrain_height + height_noise
        box_center_z = neg_half_terrain + height_noise / 2

-        normalized_height = (height_noise + grid_height) / (2 * grid_height)
+        # grid_height == 0 (difficulty 0) means a flat grid; use the midpoint
+        # color and avoid dividing by zero.
+        if grid_height > 0:
+          normalized_height = (height_noise + grid_height) / (2 * grid_height)
+        else:
+          normalized_height = 0.5
        t = float(np.clip(normalized_height, 0.0, 1.0))
        rgba = brand_ramp(_MUJOCO_GREEN, t)
        box_colors.append(rgba)
@@ -744,13 +769,17 @@ class BoxRandomSpreadTerrainCfg(SubTerrainCfg):
      )
      geometries.append(TerrainGeometry(geom=floor_geom, color=(0.4, 0.4, 0.4, 1.0)))

-    # Platform
-    platform_geom = body.add_geom(
-      type=mujoco.mjtGeom.mjGEOM_BOX,
-      size=(self.platform_width / 2, self.platform_width / 2, terrain_height / 2),
-      pos=(self.size[0] / 2, self.size[1] / 2, -terrain_height / 2),
-    )
-    geometries.append(TerrainGeometry(geom=platform_geom, color=(0.4, 0.4, 0.4, 1.0)))
+    # Center platform. When a floor is present it already provides flat ground at
+    # z=0 across the (box-free) center, so an extra platform box would only
+    # duplicate that surface and z-fight with the floor. Add the platform only
+    # when there is no floor, where it is the sole ground at the spawn point.
+    if not self.add_floor:
+      platform_geom = body.add_geom(
+        type=mujoco.mjtGeom.mjGEOM_BOX,
+        size=(self.platform_width / 2, self.platform_width / 2, terrain_height / 2),
+        pos=(self.size[0] / 2, self.size[1] / 2, -terrain_height / 2),
+      )
+      geometries.append(TerrainGeometry(geom=platform_geom, color=(0.4, 0.4, 0.4, 1.0)))

    platform_half = self.platform_width / 2
    terrain_center = self.size[0] / 2
@@ -840,13 +869,15 @@ class BoxOpenStairsTerrainCfg(SubTerrainCfg):
    border_rgba = darken_rgba(first_step_rgba, 0.85)

    if self.border_width > 0.0:
-      border_center = (0.5 * self.size[0], 0.5 * self.size[1], -step_height / 2)
+      # Keep the border solid and flush at z=0 even if step_height is 0.
+      border_height = max(step_height, _MIN_BORDER_HEIGHT)
+      border_center = (0.5 * self.size[0], 0.5 * self.size[1], -border_height / 2)
      border_inner_size = (
        self.size[0] - 2 * self.border_width,
        self.size[1] - 2 * self.border_width,
      )
      border_boxes = make_border(
-        body, self.size, border_inner_size, step_height, border_center
+        body, self.size, border_inner_size, border_height, border_center
      )
      for box in border_boxes:
        geometries.append(TerrainGeometry(geom=box, color=border_rgba))
@@ -1132,7 +1163,12 @@ class BoxRandomStairsTerrainCfg(SubTerrainCfg):
@dataclass(kw_only=True)
 class BoxSteppingStonesTerrainCfg(SubTerrainCfg):
  stone_size_range: tuple[float, float] = (0.4, 0.8)
+  """Max and min stone side length, in meters. Stones shrink from the max toward
+  the min as difficulty increases, which widens the gaps between them."""
  stone_distance_range: tuple[float, float] = (0.2, 0.5)
+  """Gap between stones, in meters. The lower bound seeds the (fixed) grid
+  density; the gap itself grows with difficulty as the stones shrink, so the
+  upper bound is no longer used directly."""
  stone_height: float = 0.2
  stone_height_variation: float = 0.1
  stone_size_variation: float = 0.1
@@ -1152,23 +1188,56 @@ class BoxSteppingStonesTerrainCfg(SubTerrainCfg):
    displacement_range = self.displacement_range * difficulty
    stone_height_variation = self.stone_height_variation * difficulty

-    # Increase distance between stones with difficulty.
-    d_low, d_high = self.stone_distance_range
-    avg_distance = d_low + difficulty * (d_high - d_low)
-
-    # Decrease stone size with difficulty (larger stones are easier).
+    # Decrease stone size with difficulty (larger stones are easier). With the
+    # grid pitch held fixed (below), shrinking stones means the gaps between them
+    # grow, which is the actual difficulty curriculum.
    s_min, s_max = self.stone_size_range
    avg_stone_size = s_max - difficulty * (s_max - s_min)
-    spacing = avg_stone_size + avg_distance

-    # Aggressive grid density to reach borders.
+    # Difficulty-INDEPENDENT grid. The count and pitch are fixed across difficulty
+    # so the layout never re-tiles (previously, num = floor(inner / spacing) + 1
+    # jumped by one as the difficulty-varying spacing crossed an integer boundary,
+    # shifting every stone at once). The pitch exactly spans the inner region so
+    # edge stones always reach the borders. Density is seeded by the tightest
+    # nominal spacing (largest stones + smallest gap).
    inner_w = self.size[0] - 2 * self.border_width
    inner_h = self.size[1] - 2 * self.border_width
-    num_x = int(np.floor(inner_w / spacing)) + 1
-    num_y = int(np.floor(inner_h / spacing)) + 1
+    nominal_spacing = s_max + self.stone_distance_range[0]
+    num_x = max(2, int(np.floor(inner_w / nominal_spacing)) + 1)
+    num_y = max(2, int(np.floor(inner_h / nominal_spacing)) + 1)
+    pitch_x = inner_w / (num_x - 1)
+    pitch_y = inner_h / (num_y - 1)

-    offset_x = self.border_width + (inner_w - (num_x - 1) * spacing) / 2
-    offset_y = self.border_width + (inner_h - (num_y - 1) * spacing) / 2
+    # Inter-stone gap (grows with difficulty as stones shrink).
+    gap_x = max(0.0, pitch_x - avg_stone_size)
+    gap_y = max(0.0, pitch_y - avg_stone_size)
+
+    # Snap the central platform out to the grid. It is at least the configured
+    # width and reaches to exactly one gap before the nearest *full* stone, so the
+    # ring of stones around it are whole (no clipped slivers that pop in and out
+    # with difficulty) and sit one consistent gap away. The platform simply
+    # absorbs the stones that would otherwise be partially under it.
+    center_x, center_y = self.size[0] / 2, self.size[1] / 2
+    half_stone = avg_stone_size / 2
+    a0 = self.platform_width / 2
+
+    def _snapped_half(center: float, pitch: float, gap: float, num: int) -> float:
+      # Nearest grid stone that can stay full while the platform is >= a0 wide.
+      threshold = center + a0 + half_stone + gap
+      i_keep = min(num - 1, int(np.ceil((threshold - self.border_width) / pitch)))
+      c_keep = self.border_width + i_keep * pitch
+      return max(a0, c_keep - half_stone - gap - center)
+
+    platform_half_x = _snapped_half(center_x, pitch_x, gap_x, num_x)
+    platform_half_y = _snapped_half(center_y, pitch_y, gap_y, num_y)
+    platform_min_x, platform_max_x = (
+      center_x - platform_half_x,
+      center_x + platform_half_x,
+    )
+    platform_min_y, platform_max_y = (
+      center_y - platform_half_y,
+      center_y + platform_half_y,
+    )

    border_rgba = darken_rgba(brand_ramp(_MUJOCO_GREEN, 0.0), 0.85)
    z_center = (self.stone_height - self.floor_depth) / 2
@@ -1195,25 +1264,20 @@ class BoxSteppingStonesTerrainCfg(SubTerrainCfg):
    )
    geometries.append(TerrainGeometry(geom=floor_geom, color=(0.1, 0.1, 0.1, 1.0)))

-    # Platform Column.
+    # Platform Column (grid-snapped, see above).
    platform_geom = body.add_geom(
      type=mujoco.mjtGeom.mjGEOM_BOX,
      size=(
-        np.maximum(1e-6, self.platform_width / 2),
-        np.maximum(1e-6, self.platform_width / 2),
+        np.maximum(1e-6, platform_half_x),
+        np.maximum(1e-6, platform_half_y),
        np.maximum(1e-6, half_height),
      ),
-      pos=(self.size[0] / 2, self.size[1] / 2, z_center),
+      pos=(center_x, center_y, z_center),
    )
    geometries.append(
      TerrainGeometry(geom=platform_geom, color=brand_ramp(_MUJOCO_GREEN, 0.5))
    )

-    platform_half = self.platform_width / 2
-    terrain_center = self.size[0] / 2
-    platform_min = terrain_center - platform_half
-    platform_max = terrain_center + platform_half
-
    inner_min_x, inner_max_x = self.border_width, self.size[0] - self.border_width
    inner_min_y, inner_max_y = self.border_width, self.size[1] - self.border_width

@@ -1221,12 +1285,17 @@ class BoxSteppingStonesTerrainCfg(SubTerrainCfg):
      for j in range(num_y):
        base_size = avg_stone_size

-        # Proposed position with displacement.
+        # Proposed position on the fixed grid with random displacement. Centers
+        # span border to (size - border), so edge stones reach the borders.
        px = (
-          offset_x + i * spacing + rng.uniform(-displacement_range, displacement_range)
+          self.border_width
+          + i * pitch_x
+          + rng.uniform(-displacement_range, displacement_range)
        )
        py = (
-          offset_y + j * spacing + rng.uniform(-displacement_range, displacement_range)
+          self.border_width
+          + j * pitch_y
+          + rng.uniform(-displacement_range, displacement_range)
        )

        # Randomized size.
@@ -1237,10 +1306,11 @@ class BoxSteppingStonesTerrainCfg(SubTerrainCfg):
        x_min, x_max = px - size_x / 2, px + size_x / 2
        y_min, y_max = py - size_y / 2, py + size_y / 2

-        # Skip stones centered inside the platform. Stones whose edges
-        # extend under the platform are kept; the platform covers the overlap.
-        if (platform_min <= px <= platform_max) and (
-          platform_min <= py <= platform_max
+        # Drop stones whose center lies under the (grid-snapped) platform; the
+        # platform absorbs them. Every remaining stone stays full size and sits
+        # one gap from the platform, so there are no clipped slivers.
+        if (platform_min_x <= px <= platform_max_x) and (
+          platform_min_y <= py <= platform_max_y
        ):
          continue

@@ -1296,6 +1366,7 @@ class BoxNarrowBeamsTerrainCfg(SubTerrainCfg):
  def function(
    self, difficulty: float, spec: mujoco.MjSpec, rng: np.random.Generator
  ) -> TerrainOutput:
+    del rng  # Beam layout is deterministic.
    body = spec.body("terrain")
    geometries = []

@@ -1306,6 +1377,19 @@ class BoxNarrowBeamsTerrainCfg(SubTerrainCfg):
    w_min, w_max = self.beam_width_range
    beam_width = w_max - difficulty * (w_max - w_min)

+    # Shrink the square platform so its corners stay within the beams' angular
+    # coverage rather than protruding into the pit between beams. A corner sits at
+    # radius r*sqrt(2) and, in the worst case, pi/num_beams away from the nearest
+    # beam, so it is covered while r*sqrt(2)*sin(pi/num_beams) <= beam_width/2.
+    # Beams thin with difficulty, so the safe radius shrinks with it. The beams
+    # attach at this same radius (below), so shrinking never opens a fall gap.
+    spacing_sin = float(np.sin(np.pi / num_beams)) if num_beams > 1 else 0.0
+    if spacing_sin > 1e-9:
+      max_no_protrude = beam_width / (2.0 * np.sqrt(2.0) * spacing_sin)
+      platform_radius = float(min(self.platform_width / 2.0, max_no_protrude))
+    else:
+      platform_radius = self.platform_width / 2.0
+
    border_rgba = darken_rgba(brand_ramp(_MUJOCO_BLUE, 0.0), 0.85)
    z_center = (self.beam_height - self.floor_depth) / 2
    half_height = (self.beam_height + self.floor_depth) / 2
@@ -1335,8 +1419,8 @@ class BoxNarrowBeamsTerrainCfg(SubTerrainCfg):
    platform_geom = body.add_geom(
      type=mujoco.mjtGeom.mjGEOM_BOX,
      size=(
-        np.maximum(1e-6, self.platform_width / 2),
-        np.maximum(1e-6, self.platform_width / 2),
+        np.maximum(1e-6, platform_radius),
+        np.maximum(1e-6, platform_radius),
        np.maximum(1e-6, half_height),
      ),
      pos=(self.size[0] / 2, self.size[1] / 2, z_center),
@@ -1347,7 +1431,6 @@ class BoxNarrowBeamsTerrainCfg(SubTerrainCfg):

    inner_size = self.size[0] - 2 * self.border_width
    center_x, center_y = self.size[0] / 2, self.size[1] / 2
-    platform_radius = self.platform_width / 2

    # Radial beams as columns.
    angles = np.linspace(0, 2 * np.pi, num_beams, endpoint=False)
@@ -1526,6 +1609,8 @@ class BoxNestedRingsTerrainCfg(SubTerrainCfg):
  ring_width_range: tuple[float, float] = (0.3, 0.6)
  gap_range: tuple[float, float] = (0.0, 0.2)
  height_range: tuple[float, float] = (0.1, 0.4)
+  """Min and max ring height, in meters. All rings share a single fixed height
+  taken as the midpoint of this range; difficulty does not scale it."""
  platform_width: float = 1.0
  border_width: float = 0.25
  floor_depth: float = 2.0
@@ -1533,20 +1618,25 @@ class BoxNestedRingsTerrainCfg(SubTerrainCfg):
  def function(
    self, difficulty: float, spec: mujoco.MjSpec, rng: np.random.Generator
  ) -> TerrainOutput:
+    del rng  # Ring layout is deterministic.
    body = spec.body("terrain")
    geometries = []

-    # Difficulty scaling: wider width range and higher average height.
-    h_scale = 1.0 + difficulty * 0.5
+    # Concentric ridges of a single fixed height. Difficulty controls
+    # gap-crossing only: gaps widen and rings narrow, so the terrain reads
+    # consistently across difficulty instead of weakly scaling height.
    w_min, w_max = self.ring_width_range
    ring_width = w_max - difficulty * (w_max - w_min)

+    ring_height = 0.5 * (self.height_range[0] + self.height_range[1])
+    ring_rgba = brand_ramp(_MUJOCO_BLUE, 0.6)
+
    border_rgba = darken_rgba(brand_ramp(_MUJOCO_BLUE, 0.0), 0.85)
-    # Use ground level z=0 as top of border/beams for consistency with NarrowBeams.
-    # In beam terrain, border top was at beam_height.

    if self.border_width > 0.0:
-      border_h = 0.5
+      # Outer border wall matches the ring height so there is no arbitrary
+      # crossover between the two as difficulty changes.
+      border_h = ring_height
      border_center = (
        0.5 * self.size[0],
        0.5 * self.size[1],
@@ -1582,12 +1672,9 @@ class BoxNestedRingsTerrainCfg(SubTerrainCfg):
    gap_min, gap_max = self.gap_range
    gap = gap_min + difficulty * (gap_max - gap_min)

-    for k in range(self.num_rings):
-      # Ring k: randomized height.
-      h = rng.uniform(self.height_range[0], self.height_range[1]) * h_scale
-
-      t = k / max(self.num_rings - 1, 1)
-      rgba = brand_ramp(_MUJOCO_BLUE, t)
+    for _ in range(self.num_rings):
+      h = ring_height
+      rgba = ring_rgba

      # Outer dimensions of this ring.
      ring_outer_size = (
@@ -1655,7 +1742,8 @@ class BoxNestedRingsTerrainCfg(SubTerrainCfg):
      ),  # Fill the ring hole + gap area.
      np.maximum(1e-2, current_outer_size[1] + 2 * gap),
    )
-    platform_h = 0.2
+    # Center pad sits flush with the ring height.
+    platform_h = ring_height

    platform_half_h = (platform_h + self.floor_depth) / 2
    platform_z = (platform_h - self.floor_depth) / 2
@@ -1,7 +1,6 @@
 from __future__ import annotations

 import abc
-import time
 from dataclasses import dataclass, field
 from typing import Literal

@@ -194,16 +193,9 @@ class TerrainGenerator:
    body = spec.worldbody.add_body(name="terrain")

    if self.cfg.curriculum:
-      tic = time.perf_counter()
      self._generate_curriculum_terrains(spec)
-      toc = time.perf_counter()
-      print(f"Curriculum terrain generation took {toc - tic:.4f} seconds.")
-
    else:
-      tic = time.perf_counter()
      self._generate_random_terrains(spec)
-      toc = time.perf_counter()
-      print(f"Terrain generation took {toc - tic:.4f} seconds.")

    self._add_terrain_border(spec)
    self._add_grid_lights(spec)
@@ -257,11 +249,11 @@ class TerrainGenerator:
    # One column per terrain type — proportion is only for spawning.
    sub_terrains_cfgs = list(self.cfg.sub_terrains.values())

+    lower, upper = self.cfg.difficulty_range
    for sub_col in range(self._num_cols):
      for sub_row in range(self.cfg.num_rows):
-        lower, upper = self.cfg.difficulty_range
-        difficulty = (sub_row + self.np_rng.uniform()) / self.cfg.num_rows
-        difficulty = lower + (upper - lower) * difficulty
+        t = sub_row / max(self.cfg.num_rows - 1, 1)
+        difficulty = lower + (upper - lower) * t
        world_position = self._get_sub_terrain_position(sub_row, sub_col)
        spawn_origin = self._create_terrain_geom(
          spec,
@@ -209,8 +209,9 @@ class CircularBuffer:

    # Backfill entire history with first frame for newly initialized batches.
    is_first_push = self._num_pushes == 0
-    if torch.any(is_first_push):
-      self._buffer[:, is_first_push] = data[is_first_push]
+    torch.where(
+      is_first_push[None, :, None], data[None, :, :], self._buffer, out=self._buffer
+    )

    self._num_pushes += 1

@@ -236,8 +237,5 @@ class CircularBuffer:
    pushes = self._num_pushes.clamp_min(1)
    valid = torch.minimum(key, pushes - 1).clamp_min(0)

-    if torch.all(valid == 0):
-      return self._buffer[self._pointer]
-
    idx = torch.remainder(self._pointer - valid, self._max_len)
    return self._buffer[idx, self._all_indices]
@@ -3,10 +3,12 @@
 import os
 from typing import Literal

+GpuId = int | str
+

 def select_gpus(
  gpu_ids: list[int] | Literal["all"] | None,
-) -> tuple[list[int] | None, int]:
+) -> tuple[list[GpuId] | None, int]:
  """Select GPUs based on CUDA_VISIBLE_DEVICES and user specification.

  This function treats the `gpu_ids` parameter as indices into the existing
@@ -19,7 +21,8 @@ def select_gpus(

  Returns:
    A tuple of (selected_gpu_ids, num_gpus) where:
-    - selected_gpu_ids: List of physical GPU IDs to use, or None for CPU mode
+    - selected_gpu_ids: List of physical GPU IDs (int for numeric, str for MIG
+      UUIDs), or None for CPU mode
    - num_gpus: Number of GPUs selected (0 for CPU mode)

  Examples:
@@ -50,8 +53,11 @@ def select_gpus(

  if existing_visible_devices is not None:
    # Parse existing CUDA_VISIBLE_DEVICES.
-    available_gpus = [
-      int(x.strip()) for x in existing_visible_devices.split(",") if x.strip()
+    # Use int for numeric IDs, keep as string for MIG UUIDs.
+    available_gpus: list[GpuId] = [
+      int(x.strip()) if x.strip().isdigit() else x.strip()
+      for x in existing_visible_devices.split(",")
+      if x.strip()
    ]
    # Empty CUDA_VISIBLE_DEVICES means CPU mode.
    if not available_gpus:
@@ -60,15 +66,16 @@ def select_gpus(
    # If not set, default to all available GPUs.
    import torch.cuda

-    available_gpus = list(range(torch.cuda.device_count()))
+    available_gpus: list[GpuId] = list(range(torch.cuda.device_count()))

  # Map gpu_ids indices to actual GPU IDs.
+  selected: list[GpuId]
  if gpu_ids == "all":
-    selected_gpus = available_gpus
+    selected = available_gpus
  else:
    # gpu_ids are indices into available_gpus.
-    selected_gpus = [available_gpus[i] for i in gpu_ids]
+    selected = [available_gpus[i] for i in gpu_ids]

-  num_gpus = len(selected_gpus)
+  num_gpus = len(selected)

-  return selected_gpus, num_gpus
+  return selected, num_gpus
@@ -6,18 +6,9 @@ import torch
 import warp as wp


-def seed_rng(
-  seed: int,
-  torch_deterministic: bool = False,
-  device: str | torch.device | None = None,
-) -> None:
+def seed_rng(seed: int, torch_deterministic: bool = False) -> None:
  """Seed all random number generators for reproducibility.

-  When ``device`` is a CPU device, ``wp.rand_init`` is skipped so that Warp's
-  CUDA runtime is not initialized on machines where a GPU is visible but the
-  caller has explicitly opted into CPU-only execution. When ``device`` is
-  ``None``, behavior is unchanged (Warp is seeded).
-
  Note: MuJoCo Warp is not fully deterministic yet.
  See: https://github.com/google-deepmind/mujoco_warp/issues/562
  """
@@ -26,8 +17,7 @@ def seed_rng(
  random.seed(seed)
  np.random.seed(seed)

-  if device is None or torch.device(device).type != "cpu":
-    wp.rand_init(wp.int32(seed))
+  wp.rand_init(wp.int32(seed))

  # Ref: https://docs.pytorch.org/docs/stable/notes/randomness.html
  torch.manual_seed(seed)  # Seed RNG for all devices.
@@ -120,6 +120,34 @@ _TRANSMISSION_TYPE_MAP = {
 }


+def apply_target_overrides(
+  spec: mujoco.MjSpec,
+  target_name: str,
+  transmission_type: TransmissionType,
+  *,
+  armature: float | None,
+  frictionloss: float | None,
+  viscous_damping: float | None,
+) -> None:
+  """Apply joint- or tendon-level overrides. ``None`` preserves the XML value.
+
+  SITE transmission is a no-op (sites have no armature / frictionloss / damping);
+  callers using SITE should not pass non-None overrides.
+  """
+  if transmission_type == TransmissionType.JOINT:
+    target = spec.joint(target_name)
+  elif transmission_type == TransmissionType.TENDON:
+    target = spec.tendon(target_name)
+  else:
+    return
+  if armature is not None:
+    target.armature = armature
+  if frictionloss is not None:
+    target.frictionloss = frictionloss
+  if viscous_damping is not None:
+    target.damping[0] = viscous_damping
+
+
 def auto_wrap_fixed_base_mocap(
  spec_fn: Callable[[], mujoco.MjSpec],
 ) -> Callable[[], mujoco.MjSpec]:
@@ -235,21 +263,14 @@ def create_motor_actuator(
  actuator.ctrllimited = True
  actuator.ctrlrange[:] = np.array([-effort_limit, effort_limit])

-  # Set armature, frictionloss, and viscous_damping (None = preserve XML value).
-  if transmission_type == TransmissionType.JOINT:
-    if armature is not None:
-      spec.joint(joint_name).armature = armature
-    if frictionloss is not None:
-      spec.joint(joint_name).frictionloss = frictionloss
-    if viscous_damping is not None:
-      spec.joint(joint_name).damping[0] = viscous_damping
-  elif transmission_type == TransmissionType.TENDON:
-    if armature is not None:
-      spec.tendon(joint_name).armature = armature
-    if frictionloss is not None:
-      spec.tendon(joint_name).frictionloss = frictionloss
-    if viscous_damping is not None:
-      spec.tendon(joint_name).damping[0] = viscous_damping
+  apply_target_overrides(
+    spec,
+    joint_name,
+    transmission_type,
+    armature=armature,
+    frictionloss=frictionloss,
+    viscous_damping=viscous_damping,
+  )

  return actuator

@@ -265,14 +286,21 @@ def create_position_actuator(
  frictionloss: float | None = None,
  viscous_damping: float | None = None,
  transmission_type: TransmissionType = TransmissionType.JOINT,
+  actuator_name: str | None = None,
 ) -> mujoco.MjsActuator:
  """Creates a <position> actuator.

  An important note about this actuator is that we set `ctrllimited` to False. This is
  because we want to allow the policy to output setpoints that are outside the kinematic
  limits of the joint.
+
+  ``actuator_name`` defaults to ``joint_name``; pass a distinct value when multiple
+  actuators target the same joint (e.g. paired position+velocity elements).
  """
-  actuator = spec.add_actuator(name=joint_name, target=joint_name)
+  actuator = spec.add_actuator(
+    name=actuator_name if actuator_name is not None else joint_name,
+    target=joint_name,
+  )

  actuator.trntype = _TRANSMISSION_TYPE_MAP[transmission_type]
  actuator.dyntype = mujoco.mjtDyn.mjDYN_NONE
@@ -314,21 +342,14 @@ def create_position_actuator(
    actuator.forcelimited = False
    # No forcerange needed.

-  # Set armature, frictionloss, and viscous_damping (None = preserve XML value).
-  if transmission_type == TransmissionType.JOINT:
-    if armature is not None:
-      spec.joint(joint_name).armature = armature
-    if frictionloss is not None:
-      spec.joint(joint_name).frictionloss = frictionloss
-    if viscous_damping is not None:
-      spec.joint(joint_name).damping[0] = viscous_damping
-  elif transmission_type == TransmissionType.TENDON:
-    if armature is not None:
-      spec.tendon(joint_name).armature = armature
-    if frictionloss is not None:
-      spec.tendon(joint_name).frictionloss = frictionloss
-    if viscous_damping is not None:
-      spec.tendon(joint_name).damping[0] = viscous_damping
+  apply_target_overrides(
+    spec,
+    joint_name,
+    transmission_type,
+    armature=armature,
+    frictionloss=frictionloss,
+    viscous_damping=viscous_damping,
+  )

  return actuator

@@ -343,14 +364,21 @@ def create_velocity_actuator(
  frictionloss: float | None = None,
  viscous_damping: float | None = None,
  transmission_type: TransmissionType = TransmissionType.JOINT,
+  actuator_name: str | None = None,
 ) -> mujoco.MjsActuator:
  """Creates a <velocity> actuator.

  Control inputs are not clamped so that velocity commands work for any joint,
  including continuous joints that have no range defined. Force output is still
  bounded when effort_limit is set.
+
+  ``actuator_name`` defaults to ``joint_name``; pass a distinct value when multiple
+  actuators target the same joint (e.g. paired position+velocity elements).
  """
-  actuator = spec.add_actuator(name=joint_name, target=joint_name)
+  actuator = spec.add_actuator(
+    name=actuator_name if actuator_name is not None else joint_name,
+    target=joint_name,
+  )

  actuator.trntype = _TRANSMISSION_TYPE_MAP[transmission_type]
  actuator.dyntype = mujoco.mjtDyn.mjDYN_NONE
@@ -369,21 +397,14 @@ def create_velocity_actuator(
  else:
    actuator.forcelimited = False

-  # Set armature, frictionloss, and viscous_damping (None = preserve XML value).
-  if transmission_type == TransmissionType.JOINT:
-    if armature is not None:
-      spec.joint(joint_name).armature = armature
-    if frictionloss is not None:
-      spec.joint(joint_name).frictionloss = frictionloss
-    if viscous_damping is not None:
-      spec.joint(joint_name).damping[0] = viscous_damping
-  elif transmission_type == TransmissionType.TENDON:
-    if armature is not None:
-      spec.tendon(joint_name).armature = armature
-    if frictionloss is not None:
-      spec.tendon(joint_name).frictionloss = frictionloss
-    if viscous_damping is not None:
-      spec.tendon(joint_name).damping[0] = viscous_damping
+  apply_target_overrides(
+    spec,
+    joint_name,
+    transmission_type,
+    armature=armature,
+    frictionloss=frictionloss,
+    viscous_damping=viscous_damping,
+  )

  return actuator

@@ -467,54 +488,52 @@ def copy_mesh_data(src: mujoco.MjsMesh, dst: mujoco.MjsMesh) -> None:
  dst.smoothnormal = src.smoothnormal


-def validate_variant_structure(
-  names: list[str],
-  bodies: list[mujoco.MjsBody],
-) -> None:
-  """Validate that variant specs share the same kinematic structure.
+def copy_texture_data(src: mujoco.MjsTexture, dst: mujoco.MjsTexture) -> None:
+  """Copy texture data from *src* to *dst*.

-  Checks that all variants have the same number of child bodies, the same number of
-  joints, the same joint types, and the same joint names. Raises ``ValueError`` with a
-  descriptive message if any differ.
+  Copies the file path or builtin/data fields, format, dimensions, and color
+  settings. The ``name`` field is NOT copied; set it on *dst* before calling.
  """
-  ref_name = names[0]
-  ref_body = bodies[0]
-  ref_joints = list(ref_body.joints)
-  ref_joint_types = [j.type for j in ref_joints]
-  ref_joint_names = [j.name for j in ref_joints]
-  ref_sub_bodies = list(ref_body.bodies)
+  assert dst.name, "dst.name must be set before copy_texture_data."
+  dst.type = src.type
+  dst.colorspace = src.colorspace
+  dst.builtin = src.builtin
+  dst.mark = src.mark
+  dst.rgb1[:] = src.rgb1
+  dst.rgb2[:] = src.rgb2
+  dst.markrgb[:] = src.markrgb
+  dst.random = src.random
+  dst.gridsize[:] = src.gridsize
+  dst.gridlayout = src.gridlayout
+  dst.width = src.width
+  dst.height = src.height
+  dst.nchannel = src.nchannel
+  dst.hflip = src.hflip
+  dst.vflip = src.vflip
+  if src.file:
+    dst.file = src.file
+  if len(src.cubefiles) > 0:
+    dst.cubefiles = src.cubefiles
+  if len(src.data) > 0:
+    dst.data = src.data
+  if src.content_type:
+    dst.content_type = src.content_type

-  for i in range(1, len(names)):
-    other_name = names[i]
-    other_body = bodies[i]

-    other_sub_bodies = list(other_body.bodies)
-    if len(other_sub_bodies) != len(ref_sub_bodies):
-      raise ValueError(
-        f"Variant '{other_name}' has {len(other_sub_bodies)} "
-        f"child bodies, but '{ref_name}' has "
-        f"{len(ref_sub_bodies)}."
-      )
+def copy_material_data(src: mujoco.MjsMaterial, dst: mujoco.MjsMaterial) -> None:
+  """Copy material data from *src* to *dst*.

-    other_joints = list(other_body.joints)
-    if len(other_joints) != len(ref_joints):
-      raise ValueError(
-        f"Variant '{other_name}' has {len(other_joints)} "
-        f"joints, but '{ref_name}' has {len(ref_joints)}."
-      )
-
-    other_joint_types = [j.type for j in other_joints]
-    if other_joint_types != ref_joint_types:
-      raise ValueError(
-        f"Variant '{other_name}' has joint types "
-        f"{other_joint_types}, but '{ref_name}' has "
-        f"{ref_joint_types}."
-      )
-
-    other_joint_names = [j.name for j in other_joints]
-    if other_joint_names != ref_joint_names:
-      raise ValueError(
-        f"Variant '{other_name}' has joint names "
-        f"{other_joint_names}, but '{ref_name}' has "
-        f"{ref_joint_names}."
-      )
+  Copies appearance settings (rgba, specular, shininess, ...) and texture
+  bindings. The ``name`` field is NOT copied; set it on *dst* before calling.
+  """
+  assert dst.name, "dst.name must be set before copy_material_data."
+  dst.rgba[:] = src.rgba
+  dst.emission = src.emission
+  dst.specular = src.specular
+  dst.shininess = src.shininess
+  dst.reflectance = src.reflectance
+  dst.roughness = src.roughness
+  dst.metallic = src.metallic
+  dst.texuniform = src.texuniform
+  dst.texrepeat[:] = src.texrepeat
+  dst.textures = list(src.textures)
@@ -179,6 +179,26 @@ class DebugVisualizer(ABC):
    """
    ...

+  @abstractmethod
+  def add_box(
+    self,
+    center: np.ndarray,
+    size: np.ndarray,
+    mat: np.ndarray,
+    color: tuple[float, float, float, float],
+    label: str | None = None,
+  ) -> None:
+    """Add an axis-oriented box visualization.
+
+    Args:
+      center: Center position (3D vector).
+      size: Half-extents along each local axis (3D vector: a, b, c).
+      mat: 3x3 rotation matrix (or flattened 9-element array).
+      color: RGBA color (values 0-1).
+      label: Optional label for this box.
+    """
+    ...
+
  @abstractmethod
  def clear(self) -> None:
    """Clear all debug visualizations."""
@@ -242,5 +262,8 @@ class NullDebugVisualizer:
  def add_ellipsoid(self, center, size, mat, color, label=None) -> None:
    pass

+  def add_box(self, center, size, mat, color, label=None) -> None:
+    pass
+
  def clear(self) -> None:
    pass
@@ -25,6 +25,7 @@ VIEWER_MODEL_FIELDS = frozenset(
  {
    "qpos0",  # Needed for correct mj_forward kinematics (qpos - qpos0).
    "geom_dataid",  # Per-world mesh variants.
+    "geom_matid",  # Per-world material variants.
    "geom_rgba",
    "geom_size",
    "geom_pos",
@@ -238,6 +238,31 @@ class MujocoNativeDebugVisualizer(DebugVisualizer):
      rgba=np.asarray(color, dtype=np.float32),
    )

+  @override
+  def add_box(
+    self,
+    center: np.ndarray,
+    size: np.ndarray,
+    mat: np.ndarray,
+    color: tuple[float, float, float, float],
+    label: str | None = None,
+  ) -> None:
+    """Add a box visualization using MuJoCo's box geometry."""
+    del label  # Unused.
+
+    self.scn.ngeom += 1
+    geom = self.scn.geoms[self.scn.ngeom - 1]
+    geom.category = mujoco.mjtCatBit.mjCAT_DECOR
+
+    mujoco.mjv_initGeom(
+      geom=geom,
+      type=mujoco.mjtGeom.mjGEOM_BOX.value,
+      size=np.asarray(size, dtype=np.float64),
+      pos=np.asarray(center, dtype=np.float64),
+      mat=np.asarray(mat, dtype=np.float64).flatten(),
+      rgba=np.asarray(color, dtype=np.float32),
+    )
+
  @override
  def clear(self) -> None:
    """Clear debug visualizations by resetting geom count."""
@@ -57,6 +57,7 @@ class OffscreenRenderer:
    self._opt = mujoco.MjvOption()
    self._pert = mujoco.MjvPerturb()
    self._catmask = mujoco.mjtCatBit.mjCAT_DYNAMIC
+    self._extra_env_ids: list[int] | None = None

  @property
  def renderer(self) -> mujoco.Renderer:
@@ -134,9 +135,17 @@ class OffscreenRenderer:

    We render a small local neighborhood around ``env_idx`` instead of the first
    N environments, so videos stay focused on the tracked robot and nearby peers.
+
+    The neighbor set is computed once and cached. ``env_origins`` can mutate during
+    training (e.g. the terrain curriculum reassigns origins on reset), so recomputing
+    every frame would make the context robots pop in and out, causing video flicker.
    """
+    if self._extra_env_ids is not None:
+      return self._extra_env_ids
+
    if self._cfg.max_extra_envs <= 0 or nworld <= 1:
-      return []
+      self._extra_env_ids = []
+      return self._extra_env_ids

    k = min(self._cfg.max_extra_envs, nworld - 1)
    origins = self._scene.env_origins[:nworld].cpu().numpy()
@@ -146,7 +155,8 @@ class OffscreenRenderer:

    nearest = np.argpartition(dist2, kth=k - 1)[:k]
    nearest = nearest[np.argsort(dist2[nearest])]
-    return [int(i) for i in nearest]
+    self._extra_env_ids = [int(i) for i in nearest]
+    return self._extra_env_ids

  def _sync_model_fields(self, env_idx: int) -> None:
    """Sync visually relevant per-world model fields into the host MjModel."""
@@ -263,6 +263,7 @@ class MjlabViserScene(ViserMujocoScene, DebugVisualizer):
    self._queued_spheres: list = []
    self._queued_cylinders: list = []
    self._queued_ellipsoids: list = []
+    self._queued_boxes: list = []

    # Batched mesh handles for simple primitives.
    def _shaft_mesh() -> trimesh.Trimesh:
@@ -287,12 +288,19 @@ class MjlabViserScene(ViserMujocoScene, DebugVisualizer):
      "ellipsoids",
      lambda: trimesh.creation.icosphere(subdivisions=2, radius=1.0),
    )
+    # Unit half-extents so that scaling by the box size yields the requested
+    # half-extents (extents=2 spans -1 to 1 along each axis).
+    self._boxes = _BatchedPrimitive(
+      "boxes",
+      lambda: trimesh.creation.box(extents=(2.0, 2.0, 2.0)),
+    )
    self._all_primitives = [
      self._arrow_shafts,
      self._arrow_heads,
      self._spheres,
      self._cylinders,
      self._ellipsoids,
+      self._boxes,
    ]

    # Ghost mesh state.
@@ -955,6 +963,27 @@ class MjlabViserScene(ViserMujocoScene, DebugVisualizer):
      )
    )

+  @override
+  def add_box(
+    self,
+    center: np.ndarray | torch.Tensor,
+    size: np.ndarray | torch.Tensor,
+    mat: np.ndarray | torch.Tensor,
+    color: tuple[float, float, float, float],
+    label: str | None = None,
+  ) -> None:
+    if not self.debug_visualization_enabled:
+      return
+    del label
+    self._queued_boxes.append(
+      (
+        np.asarray(_to_numpy(center), dtype=np.float32).copy(),
+        np.asarray(_to_numpy(size), dtype=np.float32).copy(),
+        np.asarray(_to_numpy(mat), dtype=np.float32).reshape(3, 3).copy(),
+        color,
+      )
+    )
+
  @override
  def clear(self) -> None:
    """Clear all debug visualization queues."""
@@ -962,6 +991,7 @@ class MjlabViserScene(ViserMujocoScene, DebugVisualizer):
    self._queued_spheres.clear()
    self._queued_cylinders.clear()
    self._queued_ellipsoids.clear()
+    self._queued_boxes.clear()
    self._queued_ghosts.clear()

  def clear_debug_all(self) -> None:
@@ -1039,6 +1069,7 @@ class MjlabViserScene(ViserMujocoScene, DebugVisualizer):
    self._sync_spheres()
    self._sync_cylinders()
    self._sync_ellipsoids()
+    self._sync_boxes()

  def _sync_spheres(self) -> None:
    if not self._queued_spheres:
@@ -1125,6 +1156,32 @@ class MjlabViserScene(ViserMujocoScene, DebugVisualizer):
      opacity,
    )

+  def _sync_boxes(self) -> None:
+    if not self._queued_boxes:
+      self._boxes.remove()
+      return
+    n = len(self._queued_boxes)
+    positions = np.zeros((n, 3), dtype=np.float32)
+    wxyzs = np.zeros((n, 4), dtype=np.float32)
+    scales = np.zeros((n, 3), dtype=np.float32)
+    colors = np.zeros((n, 3), dtype=np.uint8)
+    opacity = 1.0
+    for i, (center, size, mat, color) in enumerate(self._queued_boxes):
+      positions[i] = center + self._scene_offset
+      wxyzs[i] = vtf.SO3.from_matrix(mat).wxyz
+      scales[i] = size
+      colors[i] = _color_uint8(color)
+      opacity = color[3]
+    self._boxes.sync(
+      self.server,
+      self.env_idx,
+      positions,
+      wxyzs,
+      scales,
+      colors,
+      opacity,
+    )
+
  def _sync_ghosts(self) -> None:
    """Render queued ghosts as one batched handle per (model, body)."""
    if not self._queued_ghosts:
@@ -106,17 +106,25 @@ def initialize_entity(entity: Entity, device: str, num_envs: int = 1):

 def make_scene_and_sim(
  device: str,
-  xml: str,
+  xml: str | dict[str, str],
  sensors: tuple,
  num_envs: int = 1,
  sim_cfg: SimulationCfg | None = None,
 ) -> tuple[Scene, Simulation]:
-  """Create a scene and simulation from inline XML with sensors wired up."""
-  entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(xml))
+  """Create a scene and simulation from inline XML with sensors wired up.
+
+  ``xml`` may be a single XML string (registered as the ``robot`` entity) or a
+  mapping of entity name to XML string for multi-entity scenes.
+  """
+  xml_by_entity = {"robot": xml} if isinstance(xml, str) else xml
+  entities = {
+    name: EntityCfg(spec_fn=lambda s=s: mujoco.MjSpec.from_string(s))
+    for name, s in xml_by_entity.items()
+  }
  scene_cfg = SceneCfg(
    num_envs=num_envs,
    env_spacing=5.0,
-    entities={"robot": entity_cfg},
+    entities=entities,
    sensors=sensors,
  )
  scene = Scene(scene_cfg, device)
@@ -0,0 +1,653 @@
+"""Tests for BuiltinDcMotorActuator.
+
+Covers wiring of MuJoCo's native ``<dcmotor>`` element through mjlab: the
+three input modes (voltage / position / velocity), torque saturation,
+config validation, and DR integration.
+"""
+
+import math
+from unittest.mock import Mock
+
+import mujoco
+import pytest
+import torch
+from conftest import (
+  create_entity_with_actuator,
+  get_test_device,
+  initialize_entity,
+  load_fixture_xml,
+)
+
+from mjlab.actuator import (
+  BuiltinDcMotorActuator,
+  BuiltinDcMotorActuatorCfg,
+  DcMotorDatasheetParams,
+  DcMotorInputMode,
+  DcMotorPhysicalParams,
+)
+from mjlab.actuator.actuator import TransmissionType
+from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
+from mjlab.envs.mdp import dr
+from mjlab.managers.scene_entity_config import SceneEntityCfg
+from mjlab.scene import Scene, SceneCfg
+from mjlab.sim.sim import Simulation, SimulationCfg
+
+ROBOT_XML = load_fixture_xml("floating_base_articulated")
+
+# Motor characterization used throughout (resolves to K=0.24, R=2.88).
+V_NOM, TAU_STALL, OMEGA_NL = 24.0, 2.0, 100.0
+K = V_NOM / OMEGA_NL
+R = K * V_NOM / TAU_STALL
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+DATASHEET = DcMotorDatasheetParams(
+  nominal_voltage=V_NOM, stall_torque=TAU_STALL, no_load_speed=OMEGA_NL
+)
+
+
+def _make_cfg(
+  *,
+  mode: DcMotorInputMode = DcMotorInputMode.POSITION,
+  stiffness=5.0,
+  damping=0.5,
+  voltage_limit=24.0,
+  **extra,
+) -> BuiltinDcMotorActuatorCfg:
+  """Build a cfg with sensible PID defaults. ``extra`` forwards any other
+  BuiltinDcMotorActuatorCfg kwarg (effort_limit, integral_gain, thermal,
+  delay_*, etc.)."""
+  return BuiltinDcMotorActuatorCfg(
+    target_names_expr=("joint.*",),
+    mode=mode,
+    motor_params=DATASHEET,
+    stiffness=stiffness,
+    damping=damping,
+    voltage_limit=voltage_limit,
+    **extra,
+  )
+
+
+def _make_entity(**kwargs) -> Entity:
+  return create_entity_with_actuator(ROBOT_XML, _make_cfg(**kwargs))
+
+
+def _make_initialized(device, **kwargs):
+  """Build entity from cfg kwargs and initialize it through the sim."""
+  return initialize_entity(_make_entity(**kwargs), device)
+
+
+def _drive(
+  entity: Entity,
+  sim,
+  device: str,
+  *,
+  pos_target=None,
+  vel_target=None,
+  effort_target=None,
+  q0=None,
+  qd0=None,
+) -> None:
+  zero = torch.zeros(1, 2, device=device)
+  entity.write_joint_state_to_sim(
+    position=q0 if q0 is not None else zero,
+    velocity=qd0 if qd0 is not None else zero,
+  )
+  entity.set_joint_position_target(pos_target if pos_target is not None else zero)
+  entity.set_joint_velocity_target(vel_target if vel_target is not None else zero)
+  entity.set_joint_effort_target(effort_target if effort_target is not None else zero)
+  entity.write_data_to_sim()
+  sim.forward()
+
+
+# Wiring sanity.
+
+
+def test_kr_packed_into_gainprm(device):
+  """The XML compiler derives K and R from the nominal triplet."""
+  _, sim = initialize_entity(_make_entity(effort_limit=1.5), device)
+  m = sim.mj_model
+  for i in range(2):
+    assert m.actuator_gainprm[i, 0] == pytest.approx(R, abs=1e-6)
+    assert m.actuator_gainprm[i, 1] == pytest.approx(K, abs=1e-6)
+    assert m.actuator_gainprm[i, 4] == pytest.approx(5.0)  # kp
+    assert m.actuator_gainprm[i, 6] == pytest.approx(0.5)  # kd
+    assert m.actuator_gainprm[i, 7] == pytest.approx(24.0)  # Vmax
+    assert m.actuator_gainprm[i, 8] == pytest.approx(1.0)  # input_mode=position
+    assert m.actuator_gaintype[i] == mujoco.mjtGain.mjGAIN_DCMOTOR
+    assert m.actuator_biastype[i] == mujoco.mjtBias.mjBIAS_DCMOTOR
+    # No activation state: ki=0, no inductance, no thermal/lugre/slew.
+    assert m.actuator_actnum[i] == 0
+
+
+def test_motor_const_path(device):
+  """Physical params pack K = sqrt(Kt*Ke) and R verbatim."""
+  cfg = BuiltinDcMotorActuatorCfg(
+    target_names_expr=("joint.*",),
+    mode=DcMotorInputMode.VOLTAGE,
+    motor_params=DcMotorPhysicalParams(kt=0.1, ke=0.05, resistance=2.0),
+  )
+  _, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
+  m = sim.mj_model
+  for i in range(2):
+    assert m.actuator_gainprm[i, 0] == pytest.approx(2.0, abs=1e-6)
+    assert m.actuator_gainprm[i, 1] == pytest.approx((0.1 * 0.05) ** 0.5, abs=1e-6)
+
+
+# Stateless motor physics.
+
+
+def test_voltage_mode_steady_state(device):
+  """At rest, ctrl = V -> tau = K * V / R."""
+  cfg = BuiltinDcMotorActuatorCfg(
+    target_names_expr=("joint.*",),
+    mode=DcMotorInputMode.VOLTAGE,
+    motor_params=DATASHEET,
+  )
+  entity, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
+  V = torch.tensor([[10.0, -5.0]], device=device)
+  _drive(entity, sim, device, effort_target=V)
+  v_adr = entity.indexing.joint_v_adr
+  expected = K * V[0] / R
+  assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
+
+
+def test_voltage_mode_voltage_limit_zero_is_noop(device):
+  """Docstring promises ``voltage_limit=0`` disables clamping. Verify against
+  MuJoCo's ``dcmotor_voltage`` (which only clamps when ``Vmax > 0``)."""
+  cfg = BuiltinDcMotorActuatorCfg(
+    target_names_expr=("joint.*",),
+    mode=DcMotorInputMode.VOLTAGE,
+    motor_params=DATASHEET,
+    voltage_limit=0.0,
+  )
+  entity, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
+  V = torch.tensor([[1000.0, 0.0]], device=device)  # absurdly high voltage.
+  _drive(entity, sim, device, effort_target=V)
+  v_adr = entity.indexing.joint_v_adr
+  expected = K * V[0] / R
+  assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-2)
+
+
+def test_back_emf_reduces_torque_at_velocity(device):
+  """Same V, joint moving at omega: tau = K * (V - K * omega) / R."""
+  cfg = BuiltinDcMotorActuatorCfg(
+    target_names_expr=("joint.*",),
+    mode=DcMotorInputMode.VOLTAGE,
+    motor_params=DATASHEET,
+  )
+  entity, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
+  V = torch.tensor([[10.0, 0.0]], device=device)
+  omega0 = torch.tensor([[2.0, 0.0]], device=device)
+  _drive(entity, sim, device, effort_target=V, qd0=omega0)
+  v_adr = entity.indexing.joint_v_adr
+  expected = K * (V[0] - K * omega0[0]) / R
+  assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
+
+
+def test_position_mode_pid_at_rest(device):
+  """kd=0, no Vmax clamp: tau = K * kp * (target - q) / R."""
+  # voltage_limit must be >0 (cfg invariant), pick it big enough not to clamp.
+  entity, sim = initialize_entity(
+    _make_entity(damping=0.0, voltage_limit=1000.0), device
+  )
+  pos = torch.tensor([[0.1, -0.05]], device=device)
+  _drive(entity, sim, device, pos_target=pos)
+  v_adr = entity.indexing.joint_v_adr
+  expected = K * 5.0 * pos[0] / R
+  assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
+
+
+def test_position_mode_voltage_clamp(device):
+  """Huge position error -> PID voltage saturates at Vmax."""
+  entity, sim = initialize_entity(
+    _make_entity(stiffness=100.0, damping=0.0, voltage_limit=2.0),
+    device,
+  )
+  # kp * err = 100 * 0.5 = 50 V, well above Vmax=2.
+  pos = torch.tensor([[0.5, 0.0]], device=device)
+  _drive(entity, sim, device, pos_target=pos)
+  v_adr = entity.indexing.joint_v_adr
+  qfrc = sim.data.qfrc_actuator[0, v_adr]
+  expected_first = K * 2.0 / R  # tau at clamped V.
+  assert qfrc[0].item() == pytest.approx(expected_first, abs=1e-4)
+  assert qfrc[1].item() == pytest.approx(0.0, abs=1e-4)
+
+
+def test_velocity_mode_pid(device):
+  """P-only velocity tracking: tau = K * kp * (target - qdot) / R."""
+  entity, sim = initialize_entity(
+    _make_entity(mode=DcMotorInputMode.VELOCITY, damping=0.0, voltage_limit=1000.0),
+    device,
+  )
+  qd0 = torch.tensor([[1.0, 0.0]], device=device)
+  vel_target = torch.tensor([[3.0, 0.0]], device=device)
+  _drive(entity, sim, device, vel_target=vel_target, qd0=qd0)
+  v_adr = entity.indexing.joint_v_adr
+  # back-EMF subtracts K*omega; this is folded into the dcmotor bias.
+  # voltage = kp*(target - qdot); tau = K*(voltage - K*omega)/R.
+  voltage = 5.0 * (vel_target[0] - qd0[0])
+  expected = K * (voltage - K * qd0[0]) / R
+  assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
+
+
+def test_effort_limit_clamps_torque(device):
+  """forcerange clamps the algebraic torque output."""
+  entity, sim = initialize_entity(
+    _make_entity(stiffness=100.0, damping=0.0, voltage_limit=1000.0, effort_limit=0.1),
+    device,
+  )
+  m = sim.mj_model
+  for i in range(2):
+    assert m.actuator_forcelimited[i] == 1
+    assert m.actuator_forcerange[i, 0] == pytest.approx(-0.1)
+    assert m.actuator_forcerange[i, 1] == pytest.approx(0.1)
+
+  # Unclamped tau would be K * 100 * 0.5 / R ~= K*50/R, well above 0.1.
+  pos = torch.tensor([[0.5, 0.0]], device=device)
+  _drive(entity, sim, device, pos_target=pos)
+  v_adr = entity.indexing.joint_v_adr
+  qfrc = sim.data.qfrc_actuator[0, v_adr]
+  assert qfrc[0].item() == pytest.approx(0.1, abs=1e-4)
+  assert qfrc[1].item() == pytest.approx(0.0, abs=1e-4)
+
+
+# Cogging.
+
+
+def test_cogging_packed_into_biasprm(device):
+  """``cogging=(A, Np, phi)`` packs into ``biasprm[0:3]``."""
+  cfg = BuiltinDcMotorActuatorCfg(
+    target_names_expr=("joint.*",),
+    mode=DcMotorInputMode.VOLTAGE,
+    motor_params=DATASHEET,
+    cogging=(0.5, 4.0, 0.1),
+  )
+  _, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
+  m = sim.mj_model
+  for i in range(2):
+    assert m.actuator_biasprm[i, 0] == pytest.approx(0.5)
+    assert m.actuator_biasprm[i, 1] == pytest.approx(4.0)
+    assert m.actuator_biasprm[i, 2] == pytest.approx(0.1)
+
+
+def test_cogging_contributes_torque(device):
+  """At ctrl=0 (no electromagnetic torque), qfrc_actuator equals the cogging
+  term ``A * sin(Np * q + phi)`` evaluated at the joint angle."""
+  A, Np, phi = 0.5, 4.0, 0.1
+  cfg = BuiltinDcMotorActuatorCfg(
+    target_names_expr=("joint.*",),
+    mode=DcMotorInputMode.VOLTAGE,
+    motor_params=DATASHEET,
+    cogging=(A, Np, phi),
+  )
+  entity, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
+  q0, q1 = 0.3, -0.2
+  _drive(
+    entity,
+    sim,
+    device,
+    q0=torch.tensor([[q0, q1]], device=device),
+    effort_target=torch.zeros(1, 2, device=device),
+  )
+  v_adr = entity.indexing.joint_v_adr
+  qfrc = sim.data.qfrc_actuator[0, v_adr]
+  assert qfrc[0].item() == pytest.approx(A * math.sin(Np * q0 + phi), abs=1e-5)
+  assert qfrc[1].item() == pytest.approx(A * math.sin(Np * q1 + phi), abs=1e-5)
+
+
+def test_cogging_bypasses_effort_limit(device):
+  """Cogging is added *after* the forcerange clamp (MuJoCo's intentional
+  model: ``effort_limit`` bounds electromagnetic torque, cogging is
+  mechanical). Total torque can exceed ``effort_limit`` by up to the
+  cogging amplitude."""
+  A, Np, phi = 0.5, 0.0, math.pi / 2  # sin(pi/2)=1, so cogging = A at any q.
+  cfg = BuiltinDcMotorActuatorCfg(
+    target_names_expr=("joint.*",),
+    mode=DcMotorInputMode.VOLTAGE,
+    motor_params=DATASHEET,
+    cogging=(A, Np, phi),
+    effort_limit=0.05,  # An order of magnitude below A.
+  )
+  entity, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
+  # Pick a voltage large enough that the electromagnetic torque alone
+  # would saturate forcerange at +/- 0.05.
+  V = torch.tensor([[100.0, 0.0]], device=device)
+  _drive(entity, sim, device, effort_target=V)
+  v_adr = entity.indexing.joint_v_adr
+  qfrc = sim.data.qfrc_actuator[0, v_adr]
+  # joint1: electromagnetic clamped to +0.05, plus cogging A=0.5.
+  assert qfrc[0].item() == pytest.approx(0.05 + A, abs=1e-5)
+  # joint2: zero voltage, electromagnetic=0, only cogging.
+  assert qfrc[1].item() == pytest.approx(A, abs=1e-5)
+
+
+# Optional stateful extensions (integral, slew, inductance, thermal, LuGre).
+# Each behavior check compares against a baseline with the feature disabled
+# so that removing the wiring in edit_spec causes the comparison to fail.
+
+
+def _step_n(entity, sim, device, n: int, *, pos_target=None, eff_target=None):
+  zero = torch.zeros(1, 2, device=device)
+  entity.write_joint_state_to_sim(position=zero, velocity=zero)
+  for _ in range(n):
+    entity.set_joint_position_target(pos_target if pos_target is not None else zero)
+    entity.set_joint_velocity_target(zero)
+    entity.set_joint_effort_target(eff_target if eff_target is not None else zero)
+    entity.write_data_to_sim()
+    sim.step()
+
+
+def _qfrc(entity, sim) -> torch.Tensor:
+  return sim.data.qfrc_actuator[0, entity.indexing.joint_v_adr].clone()
+
+
+def test_integral_gain_ramps_torque(device):
+  """Integrator in position mode ramps torque over time even with ``kp``
+  and ``kd`` near zero."""
+  # stiffness must be > 0 (validation); choose tiny so ki dominates.
+  base = dict(
+    mode=DcMotorInputMode.POSITION, stiffness=1e-4, damping=0.0, voltage_limit=24.0
+  )
+  ent_off, sim_off = _make_initialized(device, **base, integral_gain=0.0)
+  ent_on, sim_on = _make_initialized(device, **base, integral_gain=10.0)
+
+  target = torch.tensor([[0.5, 0.0]], device=device)
+  for sim, ent in ((sim_off, ent_off), (sim_on, ent_on)):
+    _step_n(ent, sim, device, n=20, pos_target=target)
+  assert _qfrc(ent_on, sim_on)[0].abs() > 100.0 * _qfrc(ent_off, sim_off)[0].abs()
+
+
+def test_slew_rate_limits_voltage(device):
+  """``slew_rate`` rate-limits ``ctrl``: after one step, effective voltage
+  is far below the requested input."""
+  base = dict(
+    mode=DcMotorInputMode.VOLTAGE, stiffness=0.0, damping=0.0, voltage_limit=0.0
+  )
+  ent_off, sim_off = _make_initialized(device, **base, slew_rate=0.0)
+  ent_on, sim_on = _make_initialized(device, **base, slew_rate=10.0)
+
+  V = torch.tensor([[100.0, 0.0]], device=device)
+  for sim, ent in ((sim_off, ent_off), (sim_on, ent_on)):
+    _step_n(ent, sim, device, n=1, eff_target=V)
+  assert _qfrc(ent_off, sim_off)[0] > 100.0 * _qfrc(ent_on, sim_on)[0]
+
+
+def test_inductance_lags_current(device):
+  """Large ``inductance`` (te >> dt) suppresses early-step torque."""
+  base = dict(
+    mode=DcMotorInputMode.VOLTAGE, stiffness=0.0, damping=0.0, voltage_limit=0.0
+  )
+  ent_off, sim_off = _make_initialized(device, **base, inductance=0.0)
+  ent_on, sim_on = _make_initialized(device, **base, inductance=1.0)
+
+  V = torch.tensor([[10.0, 0.0]], device=device)
+  for sim, ent in ((sim_off, ent_off), (sim_on, ent_on)):
+    _step_n(ent, sim, device, n=2, eff_target=V)
+  assert _qfrc(ent_off, sim_off)[0].abs() > 10.0 * _qfrc(ent_on, sim_on)[0].abs()
+
+
+def test_thermal_decays_torque(device):
+  """I^2R heating raises T, which raises effective resistance and decays
+  torque over time."""
+  # Params chosen for visible effect in a handful of steps without going
+  # numerically unstable: small C (fast heating) and modest alpha.
+  base = dict(
+    mode=DcMotorInputMode.VOLTAGE, stiffness=0.0, damping=0.0, voltage_limit=0.0
+  )
+  ent_off, sim_off = _make_initialized(device, **base)
+  ent_on, sim_on = _make_initialized(
+    device, **base, thermal=(1.0, 0.1, 0.0, 0.01, 0.0, 0.0)
+  )
+
+  V = torch.tensor([[100.0, 0.0]], device=device)
+  for sim, ent in ((sim_off, ent_off), (sim_on, ent_on)):
+    _step_n(ent, sim, device, n=5, eff_target=V)
+  assert _qfrc(ent_on, sim_on)[0].abs() < 0.5 * _qfrc(ent_off, sim_off)[0].abs()
+
+
+def test_lugre_subtracts_friction(device):
+  """LuGre friction subtracts a velocity-dependent force after the
+  ``effort_limit`` clamp (mechanical, like cogging)."""
+  # Static comparison at v>0, ctrl=0; avoids feedback between LuGre slowing
+  # the joint and back-EMF easing off under sim.step().
+  #   no LuGre:  qfrc = -K^2 * v / R              (back-EMF only)
+  #   w/ LuGre:  qfrc = -K^2 * v / R - sigma1*v - ...
+  base = dict(
+    mode=DcMotorInputMode.VOLTAGE, stiffness=0.0, damping=0.0, voltage_limit=0.0
+  )
+  ent_off, sim_off = _make_initialized(device, **base)
+  ent_on, sim_on = _make_initialized(
+    device, **base, lugre=(1e4, 100.0, 0.1, 0.15, 0.01)
+  )
+
+  zero = torch.zeros(1, 2, device=device)
+  v0 = torch.tensor([[1.0, 0.0]], device=device)
+  for sim, ent in ((sim_off, ent_off), (sim_on, ent_on)):
+    ent.write_joint_state_to_sim(position=zero, velocity=v0)
+    ent.set_joint_position_target(zero)
+    ent.set_joint_velocity_target(zero)
+    ent.set_joint_effort_target(zero)
+    ent.write_data_to_sim()
+    sim.forward()
+  assert abs(_qfrc(ent_on, sim_on)[0]) > 100.0 * abs(_qfrc(ent_off, sim_off)[0])
+
+
+# Config validation.
+
+
+def test_pid_mode_requires_gains():
+  with pytest.raises(ValueError, match="stiffness"):
+    BuiltinDcMotorActuatorCfg(
+      target_names_expr=("j",),
+      mode=DcMotorInputMode.POSITION,
+      motor_params=DATASHEET,
+      voltage_limit=1.0,
+    )
+  with pytest.raises(ValueError, match="voltage_limit"):
+    BuiltinDcMotorActuatorCfg(
+      target_names_expr=("j",),
+      mode=DcMotorInputMode.POSITION,
+      motor_params=DATASHEET,
+      stiffness=1.0,
+    )
+
+
+def test_voltage_mode_rejects_pid_gains():
+  with pytest.raises(ValueError, match="VOLTAGE"):
+    BuiltinDcMotorActuatorCfg(
+      target_names_expr=("j",),
+      mode=DcMotorInputMode.VOLTAGE,
+      motor_params=DATASHEET,
+      stiffness=1.0,
+    )
+
+
+def test_site_rejected():
+  with pytest.raises(ValueError, match="SITE"):
+    BuiltinDcMotorActuatorCfg(
+      target_names_expr=("j",),
+      motor_params=DATASHEET,
+      stiffness=1.0,
+      voltage_limit=1.0,
+      transmission_type=TransmissionType.SITE,
+    )
+
+
+# Joint-level passthrough.
+
+
+def test_armature_applied(device):
+  _, sim = initialize_entity(_make_entity(armature=0.7), device)
+  m = sim.mj_model
+  for jname in ("joint1", "joint2"):
+    dof_id = m.jnt_dofadr[m.joint(jname).id]
+    assert m.dof_armature[dof_id] == pytest.approx(0.7)
+
+
+# Domain randomization.
+
+
+def _scene_env(
+  device,
+  num_envs=2,
+  mode: DcMotorInputMode = DcMotorInputMode.POSITION,
+):
+  def spec_fn():
+    spec = mujoco.MjSpec.from_string(ROBOT_XML)
+    for a in list(spec.actuators):
+      spec.delete(a)
+    return spec
+
+  entity_cfg = EntityCfg(
+    spec_fn=spec_fn,
+    articulation=EntityArticulationInfoCfg(
+      actuators=(
+        BuiltinDcMotorActuatorCfg(
+          target_names_expr=("joint.*",),
+          mode=mode,
+          motor_params=DATASHEET,
+          stiffness=5.0 if mode != DcMotorInputMode.VOLTAGE else 0.0,
+          damping=0.5 if mode != DcMotorInputMode.VOLTAGE else 0.0,
+          voltage_limit=24.0 if mode != DcMotorInputMode.VOLTAGE else 0.0,
+          effort_limit=50.0,
+        ),
+      )
+    ),
+  )
+  scene_cfg = SceneCfg(num_envs=num_envs, entities={"robot": entity_cfg})
+  scene = Scene(scene_cfg, device)
+  model = scene.compile()
+  sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
+  scene.initialize(model, sim.model, sim.data)
+
+  env = Mock()
+  env.num_envs = num_envs
+  env.device = device
+  env.scene = {"robot": scene["robot"]}
+  env.sim = sim
+  return env
+
+
+@pytest.mark.parametrize(
+  "operation, kp_in, kd_in, kp_expected, kd_expected",
+  [
+    # scale: multiplies the configured defaults (kp=5.0, kd=0.5).
+    ("scale", 2.0, 3.0, 2.0 * 5.0, 3.0 * 0.5),
+    # abs: writes the value directly.
+    ("abs", 10.0, 2.0, 10.0, 2.0),
+  ],
+)
+def test_dr_pd_gains_position_mode(
+  device, operation, kp_in, kd_in, kp_expected, kd_expected
+):
+  env = _scene_env(device)
+  robot = env.scene["robot"]
+  act = robot.actuators[0]
+  assert isinstance(act, BuiltinDcMotorActuator)
+  ctrl_ids = act.global_ctrl_ids
+  env.sim.expand_model_fields(("actuator_gainprm", "actuator_biasprm"))
+
+  dr.pd_gains(
+    env,
+    env_ids=torch.tensor([0], device=device),
+    kp_range=(kp_in, kp_in),
+    kd_range=(kd_in, kd_in),
+    asset_cfg=SceneEntityCfg("robot"),
+    operation=operation,
+  )
+
+  m = env.sim.model
+  n = len(ctrl_ids)
+  assert torch.allclose(
+    m.actuator_gainprm[0, ctrl_ids, 4], torch.full((n,), kp_expected, device=device)
+  )
+  assert torch.allclose(
+    m.actuator_gainprm[0, ctrl_ids, 6], torch.full((n,), kd_expected, device=device)
+  )
+  # Other env untouched (cfg defaults).
+  assert torch.allclose(m.actuator_gainprm[1, ctrl_ids, 4], torch.tensor(5.0))
+  assert torch.allclose(m.actuator_gainprm[1, ctrl_ids, 6], torch.tensor(0.5))
+
+
+def test_dr_pd_gains_voltage_mode_rejected(device):
+  env = _scene_env(device, mode=DcMotorInputMode.VOLTAGE)
+  env.sim.expand_model_fields(("actuator_gainprm", "actuator_biasprm"))
+  with pytest.raises(ValueError, match="VOLTAGE"):
+    dr.pd_gains(
+      env,
+      env_ids=torch.tensor([0], device=device),
+      kp_range=(1.0, 1.0),
+      kd_range=(1.0, 1.0),
+      asset_cfg=SceneEntityCfg("robot"),
+    )
+
+
+def test_dr_effort_limits_writes_forcerange(device):
+  env = _scene_env(device)
+  robot = env.scene["robot"]
+  act = robot.actuators[0]
+  ctrl_ids = act.global_ctrl_ids
+  env.sim.expand_model_fields(
+    ("actuator_forcerange", "jnt_actfrcrange", "tendon_actfrcrange")
+  )
+
+  dr.effort_limits(
+    env,
+    env_ids=torch.tensor([0], device=device),
+    effort_limit_range=(123.0, 123.0),
+    asset_cfg=SceneEntityCfg("robot"),
+    operation="abs",
+  )
+
+  m = env.sim.model
+  n = len(ctrl_ids)
+  assert torch.allclose(
+    m.actuator_forcerange[0, ctrl_ids, 0],
+    torch.full((n,), -123.0, device=device),
+  )
+  assert torch.allclose(
+    m.actuator_forcerange[0, ctrl_ids, 1],
+    torch.full((n,), 123.0, device=device),
+  )
+  # Env 1 keeps the configured default of 50.
+  assert torch.allclose(m.actuator_forcerange[1, ctrl_ids, 1], torch.tensor(50.0))
+
+
+# Delay.
+
+
+def test_delay_position_mode(device):
+  """A 2-step lag should make position-mode torque reference step-0 target."""
+  entity, sim = initialize_entity(
+    _make_entity(
+      stiffness=10.0,
+      damping=0.0,
+      voltage_limit=1000.0,
+      delay_min_lag=2,
+      delay_max_lag=2,
+    ),
+    device,
+  )
+  zero = torch.zeros(1, 2, device=device)
+  entity.write_joint_state_to_sim(position=zero, velocity=zero)
+  targets = [
+    torch.tensor([[0.1, 0.0]], device=device),
+    torch.tensor([[0.3, 0.0]], device=device),
+    torch.tensor([[0.5, 0.0]], device=device),
+  ]
+  for p in targets:
+    entity.set_joint_position_target(p)
+    entity.set_joint_velocity_target(zero)
+    entity.set_joint_effort_target(zero)
+    entity.write_data_to_sim()
+    sim.forward()
+
+  v_adr = entity.indexing.joint_v_adr
+  # With lag=2 and three writes, the effective target is targets[0].
+  expected = K * 10.0 * targets[0][0] / R
+  assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
@@ -0,0 +1,468 @@
+"""Tests for BuiltinPdActuator.
+
+Covers the unique surface of the actuator: paired <position>/<velocity>
+elements per target, joint/tendon-level actfrcrange sum-clamp, DR for both
+gains and effort limits, delay synchronization, and the ordering invariant
+that DR depends on.
+"""
+
+from unittest.mock import Mock
+
+import mujoco
+import pytest
+import torch
+from conftest import (
+  create_entity_with_actuator,
+  get_test_device,
+  initialize_entity,
+  load_fixture_xml,
+)
+
+from mjlab.actuator import BuiltinPdActuator, BuiltinPdActuatorCfg
+from mjlab.actuator.actuator import TransmissionType
+from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
+from mjlab.envs.mdp import dr
+from mjlab.managers.scene_entity_config import SceneEntityCfg
+from mjlab.scene import Scene, SceneCfg
+from mjlab.sim.sim import Simulation, SimulationCfg
+
+ROBOT_XML = load_fixture_xml("floating_base_articulated")
+KP = 100.0
+KD = 10.0
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+def _make_entity(
+  *,
+  effort_limit: float | None = 50.0,
+  armature: float | None = None,
+  delay_max_lag: int = 0,
+  delay_min_lag: int = 0,
+  delay_hold_prob: float = 0.0,
+) -> Entity:
+  cfg = BuiltinPdActuatorCfg(
+    target_names_expr=("joint.*",),
+    stiffness=KP,
+    damping=KD,
+    effort_limit=effort_limit,
+    armature=armature,
+    delay_min_lag=delay_min_lag,
+    delay_max_lag=delay_max_lag,
+    delay_hold_prob=delay_hold_prob,
+  )
+  return create_entity_with_actuator(ROBOT_XML, cfg)
+
+
+def _at_rest_with_targets(
+  entity: Entity,
+  sim,
+  device: str,
+  pos_target: torch.Tensor,
+  vel_target: torch.Tensor,
+) -> None:
+  entity.write_joint_state_to_sim(
+    position=torch.zeros(1, 2, device=device),
+    velocity=torch.zeros(1, 2, device=device),
+  )
+  entity.set_joint_position_target(pos_target)
+  entity.set_joint_velocity_target(vel_target)
+  entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
+  entity.write_data_to_sim()
+  sim.forward()
+
+
+# ---------------------------------------------------------------------------
+# Structural invariants
+# ---------------------------------------------------------------------------
+
+
+def test_two_ctrls_per_target_with_pos_then_vel_layout(device):
+  """Each target gets one <position> + one <velocity>, in halves."""
+  entity, sim = initialize_entity(_make_entity(), device)
+  act = entity.actuators[0]
+  assert isinstance(act, BuiltinPdActuator)
+
+  n = act.num_targets
+  assert n == len(act.target_names) == 2
+  assert len(act.ctrl_ids) == 2 * n
+  assert len(act.global_ctrl_ids) == 2 * n
+
+  names = [sim.mj_model.actuator(i).name for i in act.global_ctrl_ids.tolist()]
+  assert names[:n] == [f"{name}_pd_pos" for name in act.target_names]
+  assert names[n:] == [f"{name}_pd_vel" for name in act.target_names]
+
+
+def test_site_transmission_rejected():
+  with pytest.raises(ValueError, match="SITE"):
+    BuiltinPdActuatorCfg(
+      target_names_expr=("x",),
+      stiffness=1.0,
+      damping=1.0,
+      transmission_type=TransmissionType.SITE,
+    )
+
+
+def test_armature_applied_once(device):
+  """Joint armature must come from the position element only; double-applying
+  would silently double dof_armature."""
+  _, sim = initialize_entity(_make_entity(armature=0.7), device)
+  m = sim.mj_model
+  for jname in ("joint1", "joint2"):
+    dof_id = m.jnt_dofadr[m.joint(jname).id]
+    assert m.dof_armature[dof_id] == pytest.approx(0.7)
+
+
+# ---------------------------------------------------------------------------
+# Force computation
+# ---------------------------------------------------------------------------
+
+
+def test_position_only(device):
+  """Zero vel target: qfrc = kp * pos_target."""
+  entity, sim = initialize_entity(_make_entity(effort_limit=None), device)
+  pos = torch.tensor([[0.1, -0.05]], device=device)
+  _at_rest_with_targets(entity, sim, device, pos, torch.zeros(1, 2, device=device))
+  v_adr = entity.indexing.joint_v_adr
+  assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], KP * pos[0], atol=1e-4)
+
+
+def test_velocity_only(device):
+  """Zero pos target, joint at rest: qfrc = kd * vel_target."""
+  entity, sim = initialize_entity(_make_entity(effort_limit=None), device)
+  vel = torch.tensor([[0.3, -0.2]], device=device)
+  _at_rest_with_targets(entity, sim, device, torch.zeros(1, 2, device=device), vel)
+  v_adr = entity.indexing.joint_v_adr
+  assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], KD * vel[0], atol=1e-4)
+
+
+def test_pd_superposition(device):
+  """Both targets nonzero: qfrc = kp * pos_target + kd * vel_target."""
+  entity, sim = initialize_entity(_make_entity(effort_limit=None), device)
+  pos = torch.tensor([[0.1, -0.05]], device=device)
+  vel = torch.tensor([[0.2, -0.1]], device=device)
+  _at_rest_with_targets(entity, sim, device, pos, vel)
+  v_adr = entity.indexing.joint_v_adr
+  expected = KP * pos[0] + KD * vel[0]
+  assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
+
+
+def test_actfrcrange_sum_clamp(device):
+  """A pos error big enough to make kp*err exceed effort_limit must be
+  clamped at the joint, not allowed to ride through the unbounded element."""
+  entity, sim = initialize_entity(_make_entity(effort_limit=5.0), device)
+  # kp * 10.0 = 1000, well over the 5.0 clamp.
+  pos = torch.tensor([[10.0, 0.0]], device=device)
+  _at_rest_with_targets(entity, sim, device, pos, torch.zeros(1, 2, device=device))
+  v_adr = entity.indexing.joint_v_adr
+  qfrc = sim.data.qfrc_actuator[0, v_adr]
+  assert qfrc[0].item() == pytest.approx(5.0, abs=1e-4)
+  assert qfrc[1].item() == pytest.approx(0.0, abs=1e-4)
+
+
+def test_effort_limit_none_leaves_joint_unlimited(device):
+  """effort_limit=None: jnt_actfrclimited stays 0 on the targeted joints."""
+  _, sim = initialize_entity(_make_entity(effort_limit=None), device)
+  m = sim.mj_model
+  for jname in ("joint1", "joint2"):
+    jid = m.joint(jname).id
+    assert m.jnt_actfrclimited[jid] == 0
+
+
+def test_actuator_forcerange_not_set(device):
+  """We deliberately leave per-element forcerange unset; the limit lives on
+  the joint. Inspection of actuator_force[i] thus shows the unclamped value."""
+  entity, sim = initialize_entity(_make_entity(effort_limit=5.0), device)
+  m = sim.mj_model
+  for ctrl_id in entity.actuators[0].global_ctrl_ids.tolist():
+    assert m.actuator_forcelimited[ctrl_id] == 0
+
+
+# ---------------------------------------------------------------------------
+# Delay synchronization
+# ---------------------------------------------------------------------------
+
+
+def test_delay_syncs_pos_and_vel(device):
+  """The shared delay buffer must lag pos and vel together."""
+  entity, sim = initialize_entity(
+    _make_entity(effort_limit=None, delay_min_lag=2, delay_max_lag=2),
+    device,
+  )
+  pos_targets = [
+    torch.tensor([[0.1, 0.0]], device=device),
+    torch.tensor([[0.3, 0.0]], device=device),
+    torch.tensor([[0.5, 0.0]], device=device),
+  ]
+  vel_targets = [
+    torch.tensor([[1.0, 0.0]], device=device),
+    torch.tensor([[2.0, 0.0]], device=device),
+    torch.tensor([[3.0, 0.0]], device=device),
+  ]
+  entity.write_joint_state_to_sim(
+    position=torch.zeros(1, 2, device=device),
+    velocity=torch.zeros(1, 2, device=device),
+  )
+  for p, v in zip(pos_targets, vel_targets, strict=True):
+    entity.set_joint_position_target(p)
+    entity.set_joint_velocity_target(v)
+    entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
+    entity.write_data_to_sim()
+    sim.forward()
+
+  v_adr = entity.indexing.joint_v_adr
+  # With lag=2, both halves should reference step-0 values.
+  expected = KP * pos_targets[0][0] + KD * vel_targets[0][0]
+  assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
+
+
+def test_reset_clears_delay_buffer(device):
+  entity, _ = initialize_entity(_make_entity(delay_min_lag=1, delay_max_lag=3), device)
+  act = entity.actuators[0]
+  assert act._delay_buffer is not None
+  entity.set_joint_position_target(torch.full((1, 2), 0.5, device=device))
+  entity.set_joint_velocity_target(torch.zeros(1, 2, device=device))
+  entity.write_data_to_sim()
+
+  entity.reset(torch.tensor([0], device=device))
+  assert act._delay_buffer.current_lags[0] == 0
+
+
+# ---------------------------------------------------------------------------
+# Domain randomization
+# ---------------------------------------------------------------------------
+
+
+def _scene_env(device, transmission=TransmissionType.JOINT, num_envs=2):
+  """Build a real scene/sim with one BuiltinPd-driven entity for DR tests."""
+  if transmission == TransmissionType.JOINT:
+    xml = ROBOT_XML
+    targets = ("joint.*",)
+  else:
+    xml = load_fixture_xml("tendon_finger")
+    # tendon_finger ships with motor/position/velocity actuators; we need a
+    # bare spec so BuiltinPd can attach to the tendon without name clashes.
+    targets = ("finger_tendon",)
+
+  def spec_fn():
+    spec = mujoco.MjSpec.from_string(xml)
+    # Strip any pre-existing actuators so BuiltinPd's added elements own ctrl.
+    for a in list(spec.actuators):
+      spec.delete(a)
+    return spec
+
+  entity_cfg = EntityCfg(
+    spec_fn=spec_fn,
+    articulation=EntityArticulationInfoCfg(
+      actuators=(
+        BuiltinPdActuatorCfg(
+          target_names_expr=targets,
+          stiffness=KP,
+          damping=KD,
+          effort_limit=50.0,
+          transmission_type=transmission,
+        ),
+      )
+    ),
+  )
+  scene_cfg = SceneCfg(num_envs=num_envs, entities={"robot": entity_cfg})
+  scene = Scene(scene_cfg, device)
+  model = scene.compile()
+  sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
+  scene.initialize(model, sim.model, sim.data)
+
+  env = Mock()
+  env.num_envs = num_envs
+  env.device = device
+  env.scene = {"robot": scene["robot"]}
+  env.sim = sim
+  return env
+
+
+def test_dr_pd_gains_scales_halves_independently(device):
+  env = _scene_env(device)
+  robot = env.scene["robot"]
+  act = robot.actuators[0]
+  assert isinstance(act, BuiltinPdActuator)
+  n = act.num_targets
+  pos_ids = act.global_ctrl_ids[:n]
+  vel_ids = act.global_ctrl_ids[n:]
+
+  # Expand fields so DR can write per-env.
+  env.sim.expand_model_fields(("actuator_gainprm", "actuator_biasprm"))
+
+  dr.pd_gains(
+    env,
+    env_ids=torch.tensor([0], device=device),
+    kp_range=(2.0, 2.0),
+    kd_range=(3.0, 3.0),
+    asset_cfg=SceneEntityCfg("robot"),
+    operation="scale",
+  )
+
+  m = env.sim.model
+  # Position half: gainprm[0] and biasprm[1] both scaled by kp=2, biasprm[2]
+  # must stay zero (no kd injection).
+  assert torch.allclose(
+    m.actuator_gainprm[0, pos_ids, 0],
+    torch.full((n,), 2.0 * KP, device=device),
+  )
+  assert torch.allclose(
+    m.actuator_biasprm[0, pos_ids, 1],
+    torch.full((n,), -2.0 * KP, device=device),
+  )
+  assert torch.allclose(
+    m.actuator_biasprm[0, pos_ids, 2], torch.zeros(n, device=device)
+  )
+  # Velocity half: gainprm[0] and biasprm[2] both scaled by kd=3, biasprm[1]
+  # stays zero (no kp injection).
+  assert torch.allclose(
+    m.actuator_gainprm[0, vel_ids, 0],
+    torch.full((n,), 3.0 * KD, device=device),
+  )
+  assert torch.allclose(
+    m.actuator_biasprm[0, vel_ids, 2],
+    torch.full((n,), -3.0 * KD, device=device),
+  )
+  assert torch.allclose(
+    m.actuator_biasprm[0, vel_ids, 1], torch.zeros(n, device=device)
+  )
+  # The other env must be untouched.
+  assert torch.allclose(m.actuator_gainprm[1, pos_ids, 0], torch.tensor(KP))
+  assert torch.allclose(m.actuator_gainprm[1, vel_ids, 0], torch.tensor(KD))
+
+
+def test_dr_pd_gains_abs_writes_correct_columns(device):
+  env = _scene_env(device)
+  robot = env.scene["robot"]
+  act = robot.actuators[0]
+  assert isinstance(act, BuiltinPdActuator)
+  n = act.num_targets
+  pos_ids = act.global_ctrl_ids[:n]
+  vel_ids = act.global_ctrl_ids[n:]
+  env.sim.expand_model_fields(("actuator_gainprm", "actuator_biasprm"))
+
+  dr.pd_gains(
+    env,
+    env_ids=torch.tensor([0, 1], device=device),
+    kp_range=(200.0, 200.0),
+    kd_range=(25.0, 25.0),
+    asset_cfg=SceneEntityCfg("robot"),
+    operation="abs",
+  )
+
+  m = env.sim.model
+  assert torch.allclose(
+    m.actuator_gainprm[:, pos_ids, 0],
+    torch.full((env.num_envs, n), 200.0, device=device),
+  )
+  assert torch.allclose(
+    m.actuator_biasprm[:, pos_ids, 1],
+    torch.full((env.num_envs, n), -200.0, device=device),
+  )
+  assert torch.allclose(
+    m.actuator_biasprm[:, pos_ids, 2],
+    torch.zeros(env.num_envs, n, device=device),
+  )
+  assert torch.allclose(
+    m.actuator_gainprm[:, vel_ids, 0],
+    torch.full((env.num_envs, n), 25.0, device=device),
+  )
+  assert torch.allclose(
+    m.actuator_biasprm[:, vel_ids, 2],
+    torch.full((env.num_envs, n), -25.0, device=device),
+  )
+
+
+def test_dr_effort_limits_writes_jnt_actfrcrange(device):
+  env = _scene_env(device, transmission=TransmissionType.JOINT)
+  robot = env.scene["robot"]
+  act = robot.actuators[0]
+  assert isinstance(act, BuiltinPdActuator)
+  env.sim.expand_model_fields(
+    ("actuator_forcerange", "jnt_actfrcrange", "tendon_actfrcrange")
+  )
+
+  joint_ids = robot.indexing.joint_ids[act.target_ids]
+  pre_forcerange = env.sim.model.actuator_forcerange.clone()
+
+  dr.effort_limits(
+    env,
+    env_ids=torch.tensor([0], device=device),
+    effort_limit_range=(123.0, 123.0),
+    asset_cfg=SceneEntityCfg("robot"),
+    operation="abs",
+  )
+
+  m = env.sim.model
+  # The joint sum-clamp was rewritten on env 0 only.
+  assert torch.allclose(
+    m.jnt_actfrcrange[0, joint_ids],
+    torch.tensor([[-123.0, 123.0]] * len(joint_ids), device=device),
+  )
+  assert torch.allclose(
+    m.jnt_actfrcrange[1, joint_ids],
+    torch.tensor([[-50.0, 50.0]] * len(joint_ids), device=device),
+  )
+  # Per-element actuator_forcerange must be untouched for BuiltinPd: that
+  # field belongs to the existing single-element actuator semantic.
+  assert torch.allclose(m.actuator_forcerange, pre_forcerange)
+
+
+def test_dr_effort_limits_scale_multiplies_default(device):
+  """``scale`` multiplies the configured ``effort_limit`` (50.0) by the sample."""
+  env = _scene_env(device, transmission=TransmissionType.JOINT)
+  robot = env.scene["robot"]
+  act = robot.actuators[0]
+  assert isinstance(act, BuiltinPdActuator)
+  env.sim.expand_model_fields(
+    ("actuator_forcerange", "jnt_actfrcrange", "tendon_actfrcrange")
+  )
+  joint_ids = robot.indexing.joint_ids[act.target_ids]
+
+  dr.effort_limits(
+    env,
+    env_ids=torch.tensor([0], device=device),
+    effort_limit_range=(2.0, 2.0),
+    asset_cfg=SceneEntityCfg("robot"),
+    operation="scale",
+  )
+
+  m = env.sim.model
+  # Default is [-50, 50], scaled by 2 -> [-100, 100].
+  assert torch.allclose(
+    m.jnt_actfrcrange[0, joint_ids],
+    torch.tensor([[-100.0, 100.0]] * len(joint_ids), device=device),
+  )
+
+
+def test_dr_effort_limits_writes_tendon_actfrcrange(device):
+  env = _scene_env(device, transmission=TransmissionType.TENDON)
+  robot = env.scene["robot"]
+  act = robot.actuators[0]
+  assert isinstance(act, BuiltinPdActuator)
+  env.sim.expand_model_fields(
+    ("actuator_forcerange", "jnt_actfrcrange", "tendon_actfrcrange")
+  )
+  tendon_ids = robot.indexing.tendon_ids[act.target_ids]
+
+  dr.effort_limits(
+    env,
+    env_ids=torch.tensor([0], device=device),
+    effort_limit_range=(77.0, 77.0),
+    asset_cfg=SceneEntityCfg("robot"),
+    operation="abs",
+  )
+
+  m = env.sim.model
+  assert torch.allclose(
+    m.tendon_actfrcrange[0, tendon_ids],
+    torch.tensor([[-77.0, 77.0]] * len(tendon_ids), device=device),
+  )
+  assert torch.allclose(
+    m.tendon_actfrcrange[1, tendon_ids],
+    torch.tensor([[-50.0, 50.0]] * len(tendon_ids), device=device),
+  )
@@ -1072,3 +1072,37 @@ def test_history_captures_impact_forces(device):
  assert torch.all(max_force_seen > steady_state_force * 1.5), (
    f"Peak force {max_force_seen} should be significantly above mg={steady_state_force}"
  )
+
+
+def test_global_frame_maxforce_rotation(device):
+  """A box at rest on a plane has its contact normals all vertical."""
+  cfg = ContactSensorCfg(
+    name="box_contact",
+    primary=ContactMatch(mode="geom", pattern="box_geom", entity="box"),
+    fields=("found", "force", "normal", "tangent"),
+    reduce="maxforce",
+    global_frame=True,
+  )
+  scene, sim = create_scene_with_sensor(FALLING_BOX_XML, "box", cfg, device)
+
+  root_state = torch.zeros((2, 13), device=sim.device)
+  root_state[:, 2] = 0.11
+  root_state[:, 3] = 1.0
+  scene["box"].write_root_state_to_sim(root_state)
+  for _ in range(150):
+    sim.step()
+    scene.update(dt=sim.cfg.mujoco.timestep)
+
+  sensor_force = scene["box_contact"].data.force[:, 0, :]
+
+  # On a flat plane the contact normal is vertical, so a correctly rotated
+  # global-frame force should have its magnitude entirely on the z axis.
+  assert torch.all(sensor_force[:, 0].abs() < 0.05), (
+    f"sensor_force x-component should be ~0, got {sensor_force[:, 0].tolist()}"
+  )
+  assert torch.all(sensor_force[:, 1].abs() < 0.05), (
+    f"sensor_force y-component should be ~0, got {sensor_force[:, 1].tolist()}"
+  )
+  assert torch.all(sensor_force[:, 2].abs() > 1.0), (
+    f"sensor_force z-component should be non-trivial, got {sensor_force[:, 2].tolist()}"
+  )
@@ -125,6 +125,72 @@ def test_delayed_ideal_applies_delay(device):
  assert torch.allclose(qfrc, expected_torque, atol=1e-4)


+def test_delayed_ideal_delays_velocity(device):
+  """Velocity targets share the same delay as position targets.
+
+  Regression test: the velocity reference used to bypass the delay buffer, so
+  the damping term consumed the latest target instead of the delayed one.
+  """
+  entity = create_entity_with_delayed_ideal(delay_min_lag=2, delay_max_lag=2)
+  entity, sim = initialize_entity(entity, device)
+
+  joint_pos = torch.zeros(1, 2, device=device)
+  joint_vel = torch.zeros(1, 2, device=device)
+  entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  # Only the velocity target varies; position and effort stay zero.
+  vel_targets = [
+    torch.tensor([[0.1, 0.2]], device=device),
+    torch.tensor([[0.3, 0.4]], device=device),
+    torch.tensor([[0.5, 0.6]], device=device),
+  ]
+
+  for vel_target in vel_targets:
+    entity.set_joint_position_target(joint_pos)
+    entity.set_joint_velocity_target(vel_target)
+    entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
+    entity.write_data_to_sim()
+    sim.forward()
+
+  joint_v_adr = entity.indexing.joint_v_adr
+  qfrc = sim.data.qfrc_actuator[0, joint_v_adr]
+
+  # With lag=2, the damping term uses the velocity target from step 0:
+  # kd * (delayed_vel_target - 0) = 10.0 * [0.1, 0.2].
+  expected_torque = 10.0 * vel_targets[0][0]
+  assert torch.allclose(qfrc, expected_torque, atol=1e-4)
+
+
+def test_delayed_ideal_delays_effort(device):
+  """Feedforward effort targets share the same delay as position targets."""
+  entity = create_entity_with_delayed_ideal(delay_min_lag=2, delay_max_lag=2)
+  entity, sim = initialize_entity(entity, device)
+
+  joint_pos = torch.zeros(1, 2, device=device)
+  joint_vel = torch.zeros(1, 2, device=device)
+  entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  effort_targets = [
+    torch.tensor([[1.0, 2.0]], device=device),
+    torch.tensor([[3.0, 4.0]], device=device),
+    torch.tensor([[5.0, 6.0]], device=device),
+  ]
+
+  for effort_target in effort_targets:
+    entity.set_joint_position_target(joint_pos)
+    entity.set_joint_velocity_target(joint_vel)
+    entity.set_joint_effort_target(effort_target)
+    entity.write_data_to_sim()
+    sim.forward()
+
+  joint_v_adr = entity.indexing.joint_v_adr
+  qfrc = sim.data.qfrc_actuator[0, joint_v_adr]
+
+  # With lag=2, the feedforward term uses the effort target from step 0.
+  expected_torque = effort_targets[0][0]
+  assert torch.allclose(qfrc, expected_torque, atol=1e-4)
+
+
 def test_delayed_actuator_reset(device):
  """Test that reset clears the delay buffer."""
  entity = create_entity_with_delayed_builtin(delay_min_lag=1, delay_max_lag=3)
@@ -243,6 +243,27 @@ def test_unnamed_freejoint_gets_default_name():
  assert "floating_base_joint" in entity.all_joint_names


+def test_multiple_freejoints_raises():
+  """An entity with more than one freejoint is rejected at construction."""
+  xml = """
+  <mujoco>
+    <worldbody>
+      <body name="object_a" pos="0 0 1">
+        <freejoint/>
+        <geom type="box" size="0.1 0.1 0.1" mass="0.1"/>
+      </body>
+      <body name="object_b" pos="1 0 1">
+        <freejoint/>
+        <geom type="box" size="0.1 0.1 0.1" mass="0.1"/>
+      </body>
+    </worldbody>
+  </mujoco>
+  """
+  cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(xml))
+  with pytest.raises(ValueError, match="2 freejoints"):
+    Entity(cfg)
+
+
 def test_find_methods():
  """Test find methods with exact and regex matches."""
  entity = create_floating_articulated_entity()
@@ -125,7 +125,9 @@ def test_dr_fields_registered_in_event_manager(device):
  assert "actuator_gainprm" in manager.domain_randomization_fields
  assert "actuator_biasprm" in manager.domain_randomization_fields
  assert "actuator_forcerange" in manager.domain_randomization_fields
-  assert len(manager.domain_randomization_fields) == 5
+  assert "jnt_actfrcrange" in manager.domain_randomization_fields
+  assert "tendon_actfrcrange" in manager.domain_randomization_fields
+  assert len(manager.domain_randomization_fields) == 7


 def test_recompute_level_ordering():
@@ -418,6 +420,62 @@ def test_effort_limits_scale_no_accumulation(device):
  assert abs(actual_upper - 200.0) < 1e-5


+def test_pd_gains_accepts_operation_object(device):
+  """dr.scale / dr.abs Operation objects produce the same result as strings."""
+  env_str, ideal_str = _make_pd_env(device)
+  env_obj, ideal_obj = _make_pd_env(device)
+
+  ids = torch.tensor([0], device=device)
+  kwargs = dict(
+    kp_range=(1.5, 1.5), kd_range=(2.0, 2.0), asset_cfg=SceneEntityCfg("robot")
+  )
+
+  torch.manual_seed(0)
+  dr.pd_gains(env_str, ids, operation="scale", **kwargs)
+  torch.manual_seed(0)
+  dr.pd_gains(env_obj, ids, operation=dr.scale, **kwargs)
+
+  assert torch.allclose(
+    env_str.sim.model.actuator_gainprm[0], env_obj.sim.model.actuator_gainprm[0]
+  )
+  assert torch.allclose(ideal_str.stiffness, ideal_obj.stiffness)
+
+
+def test_effort_limits_accepts_operation_object(device):
+  """dr.abs Operation object produces the same result as the string."""
+  env_str, ideal_str = _make_effort_env(device)
+  env_obj, ideal_obj = _make_effort_env(device)
+
+  ids = torch.tensor([0], device=device)
+  kwargs = dict(effort_limit_range=(150.0, 150.0), asset_cfg=SceneEntityCfg("robot"))
+
+  dr.effort_limits(env_str, ids, operation="abs", **kwargs)
+  dr.effort_limits(env_obj, ids, operation=dr.abs, **kwargs)
+
+  assert torch.allclose(
+    env_str.sim.model.actuator_forcerange[0], env_obj.sim.model.actuator_forcerange[0]
+  )
+  assert torch.allclose(ideal_str.force_limit, ideal_obj.force_limit)
+
+
+def test_pd_gains_rejects_unsupported_operation(device):
+  """Operations other than scale/abs raise ValueError."""
+  env, _ = _make_pd_env(device)
+  ids = torch.tensor([0], device=device)
+
+  with pytest.raises(ValueError, match="only supports 'scale' and 'abs'"):
+    dr.pd_gains(env, ids, kp_range=(1.0, 1.0), kd_range=(1.0, 1.0), operation=dr.add)
+
+
+def test_effort_limits_rejects_unsupported_operation(device):
+  """Operations other than scale/abs raise ValueError."""
+  env, _ = _make_effort_env(device)
+  ids = torch.tensor([0], device=device)
+
+  with pytest.raises(ValueError, match="only supports 'scale' and 'abs'"):
+    dr.effort_limits(env, ids, effort_limit_range=(1.0, 1.0), operation=dr.add)
+
+
 # ===========================================================================
 # Section 3: Other events
 # ===========================================================================
@@ -503,7 +561,9 @@ def test_step_mode_fires_every_call(device):
  assert call_count[0] == 5


-def _make_impulse_env(device, num_envs=2, num_bodies=1, body_ids=None):
+def _make_impulse_env(
+  device, num_envs=2, num_bodies=1, body_ids=None, cooldown_s=(0.0, 0.0)
+):
  """Create a mock env for apply_body_impulse tests."""
  if body_ids is None:
    body_ids = [0]
@@ -523,7 +583,7 @@ def _make_impulse_env(device, num_envs=2, num_bodies=1, body_ids=None):

  asset_cfg = SceneEntityCfg("robot", body_ids=body_ids)
  term_cfg = Mock()
-  term_cfg.params = {"asset_cfg": asset_cfg}
+  term_cfg.params = {"asset_cfg": asset_cfg, "cooldown_s": cooldown_s}
  impulse = events.apply_body_impulse(cfg=term_cfg, env=env)
  return env, mock_entity, asset_cfg, impulse

@@ -531,11 +591,13 @@ def _make_impulse_env(device, num_envs=2, num_bodies=1, body_ids=None):
 def test_apply_body_impulse_basic(device):
  """Impulse is applied and cleared after duration expires."""
  env, mock_entity, asset_cfg, impulse = _make_impulse_env(
-    device, num_envs=2, num_bodies=3, body_ids=[1]
+    device, num_envs=2, num_bodies=3, body_ids=[1], cooldown_s=(10.0, 10.0)
  )

-  # First call: cooldown_s starts at 0 and gets decremented by dt,
-  # so it becomes <= 0 and triggers.
+  # Skip the initial cooldown so the first call triggers immediately;
+  # the trigger/sustain/expire cycle is what's under test here.
+  impulse._interval_time_left[:] = 0.0
+
  impulse(
    env,
    None,
@@ -643,6 +705,43 @@ def test_apply_body_impulse_reset_clears(device):
  assert env_ids_arg[0].item() == 0


+def test_apply_body_impulse_initial_cooldown(device):
+  """The first call after init/reset enters cooldown, not an immediate impulse.
+
+  Regression test for #973.
+  """
+  env, mock_entity, asset_cfg, impulse = _make_impulse_env(
+    device, num_envs=1, num_bodies=1, body_ids=[0], cooldown_s=(0.05, 0.05)
+  )
+
+  def step():
+    impulse(
+      env,
+      None,
+      force_range=(10.0, 10.0),
+      torque_range=(0.0, 0.0),
+      duration_s=(1.0, 1.0),
+      cooldown_s=(0.05, 0.05),  # ~2.5 steps at dt=0.02
+      asset_cfg=asset_cfg,
+    )
+
+  # First two steps consume the sampled cooldown; impulse must not fire yet.
+  step()
+  assert not impulse._active.any()
+  step()
+  assert not impulse._active.any()
+
+  # Third step crosses the cooldown boundary and triggers.
+  step()
+  assert impulse._active.all()
+
+  # Reset re-enters cooldown: next step should not immediately re-trigger.
+  impulse.reset(env_ids=torch.tensor([0], device=device))
+  assert not impulse._active.any()
+  step()
+  assert not impulse._active.any()
+
+
 # ===========================================================================
 # Section 5: Recomputation integration
 # ===========================================================================
@@ -113,3 +113,16 @@ def test_select_gpus_cpu_mode_empty_cuda_visible_devices():
  selected, num = select_gpus([0])
  assert selected is None
  assert num == 0
+
+
+def test_select_gpus_mig_uuids():
+  """Handles MIG GPU UUIDs in CUDA_VISIBLE_DEVICES."""
+  os.environ["CUDA_VISIBLE_DEVICES"] = "MIG-GPU-abc-123,MIG-GPU-def-456"
+
+  selected, num = select_gpus("all")
+  assert selected == ["MIG-GPU-abc-123", "MIG-GPU-def-456"]
+  assert num == 2
+
+  selected, num = select_gpus([0])
+  assert selected == ["MIG-GPU-abc-123"]
+  assert num == 1
@@ -1,961 +0,0 @@
-"""Tests for per-world mesh variant support."""
-
-from __future__ import annotations
-
-from typing import Any, cast
-
-import mujoco
-import numpy as np
-import pytest
-import torch
-
-from mjlab.entity import EntityCfg, VariantCfg, VariantEntityCfg
-from mjlab.sim.mesh_variants import allocate_worlds, build_mesh_variant_model
-from mjlab.viewer.model_sync import (
-  disable_model_sameframe_shortcuts,
-  sync_model_fields,
-)
-
-# Helpers: variant specs with visual + collision mesh geoms.
-
-
-def _sphere_2col_spec() -> mujoco.MjSpec:
-  """Sphere: 1 visual + 2 collision geoms."""
-  spec = mujoco.MjSpec()
-  mv = spec.add_mesh()
-  mv.name = "visual"
-  mv.make_sphere(subdivision=3)
-  for i in range(2):
-    mc = spec.add_mesh()
-    mc.name = f"col_{i}"
-    mc.make_sphere(subdivision=1)
-  body = spec.worldbody.add_body()
-  body.name = "prop"
-  body.add_freejoint()
-  gv = body.add_geom()
-  gv.name = "visual"
-  gv.type = mujoco.mjtGeom.mjGEOM_MESH
-  gv.meshname = "visual"
-  gv.contype = 0
-  gv.conaffinity = 0
-  for i in range(2):
-    gc = body.add_geom()
-    gc.name = f"col_{i}"
-    gc.type = mujoco.mjtGeom.mjGEOM_MESH
-    gc.meshname = f"col_{i}"
-  return spec
-
-
-def _cone_4col_spec() -> mujoco.MjSpec:
-  """Cone: 1 visual + 4 collision geoms (more than sphere)."""
-  spec = mujoco.MjSpec()
-  mv = spec.add_mesh()
-  mv.name = "visual"
-  mv.make_cone(nedge=8, radius=0.05)
-  for i in range(4):
-    mc = spec.add_mesh()
-    mc.name = f"col_{i}"
-    mc.make_sphere(subdivision=1)
-  body = spec.worldbody.add_body()
-  body.name = "prop"
-  body.add_freejoint()
-  gv = body.add_geom()
-  gv.name = "visual"
-  gv.type = mujoco.mjtGeom.mjGEOM_MESH
-  gv.meshname = "visual"
-  gv.contype = 0
-  gv.conaffinity = 0
-  for i in range(4):
-    gc = body.add_geom()
-    gc.name = f"col_{i}"
-    gc.type = mujoco.mjtGeom.mjGEOM_MESH
-    gc.meshname = f"col_{i}"
-  return spec
-
-
-def _simple_sphere_spec() -> mujoco.MjSpec:
-  """Single-geom sphere for simple tests."""
-  spec = mujoco.MjSpec()
-  m = spec.add_mesh()
-  m.name = "sphere"
-  m.make_sphere(subdivision=2)
-  body = spec.worldbody.add_body()
-  body.name = "prop"
-  body.add_freejoint()
-  g = body.add_geom()
-  g.name = "visual"
-  g.type = mujoco.mjtGeom.mjGEOM_MESH
-  g.meshname = "sphere"
-  return spec
-
-
-def _simple_cone_spec() -> mujoco.MjSpec:
-  """Single-geom cone for simple tests."""
-  spec = mujoco.MjSpec()
-  m = spec.add_mesh()
-  m.name = "cone"
-  m.make_cone(nedge=8, radius=0.05)
-  body = spec.worldbody.add_body()
-  body.name = "prop"
-  body.add_freejoint()
-  g = body.add_geom()
-  g.name = "visual"
-  g.type = mujoco.mjtGeom.mjGEOM_MESH
-  g.meshname = "cone"
-  return spec
-
-
-def _hinge_spec() -> mujoco.MjSpec:
-  """Object with a hinge joint (incompatible with freejoint variants)."""
-  spec = mujoco.MjSpec()
-  m = spec.add_mesh()
-  m.name = "box"
-  m.make_sphere(subdivision=1)
-  body = spec.worldbody.add_body()
-  body.name = "prop"
-  j = body.add_joint()
-  j.name = "hinge"
-  j.type = mujoco.mjtJoint.mjJNT_HINGE
-  g = body.add_geom()
-  g.name = "visual"
-  g.type = mujoco.mjtGeom.mjGEOM_MESH
-  g.meshname = "box"
-  return spec
-
-
-def _build_scene_with_variants(
-  variant_a_fn, variant_b_fn, *, weight_a=0.5, weight_b=0.5
-):
-  """Build a scene spec + variant_info from two variant spec_fns."""
-  cfg = VariantEntityCfg(
-    variants={
-      "a": VariantCfg(spec_fn=variant_a_fn, weight=weight_a),
-      "b": VariantCfg(spec_fn=variant_b_fn, weight=weight_b),
-    },
-  )
-  entity = cfg.build()
-  assert entity.variant_metadata is not None
-  scene_spec = mujoco.MjSpec()
-  frame = scene_spec.worldbody.add_frame()
-  scene_spec.attach(entity.spec, prefix="object/", frame=frame)
-  return scene_spec, [("object/", entity.variant_metadata)]
-
-
-# allocate_worlds.
-
-
-def test_allocate_worlds_proportional():
-  result = allocate_worlds((0.6, 0.4), 10)
-  assert len(result) == 10
-  assert result.count(0) == 6
-  assert result.count(1) == 4
-
-
-def test_allocate_worlds_uniform():
-  result = allocate_worlds((1.0, 1.0), 8)
-  assert result.count(0) == 4
-  assert result.count(1) == 4
-
-
-def test_allocate_worlds_single_variant():
-  result = allocate_worlds((1.0,), 5)
-  assert result == [0, 0, 0, 0, 0]
-
-
-def test_allocate_worlds_zero_weight_skips_variant():
-  """A zero-weight variant gets zero worlds; the rest split nworld."""
-  result = allocate_worlds((1.0, 0.0, 1.0), 10)
-  assert len(result) == 10
-  assert result.count(1) == 0
-  assert result.count(0) == 5
-  assert result.count(2) == 5
-
-
-def test_allocate_worlds_rejects_negative_weight():
-  with pytest.raises(ValueError, match="non-negative"):
-    allocate_worlds((1.0, -0.1), 10)
-
-
-def test_allocate_worlds_rejects_all_zero():
-  with pytest.raises(ValueError, match="positive sum"):
-    allocate_worlds((0.0, 0.0), 10)
-
-
-def test_allocate_worlds_largest_remainder_sums_to_nworld():
-  """Largest-remainder rounding must always allocate exactly nworld worlds."""
-  for nworld in (3, 7, 100, 1000):
-    result = allocate_worlds((1.0, 1.0, 1.0), nworld)
-    assert len(result) == nworld
-    # Difference between any two variant counts is at most 1 (uniform).
-    counts = [result.count(i) for i in range(3)]
-    assert max(counts) - min(counts) <= 1
-
-
-# Entity merging.
-
-
-def test_entity_builds_with_variants():
-  cfg = VariantEntityCfg(
-    variants={
-      "sphere": VariantCfg(spec_fn=_simple_sphere_spec, weight=0.5),
-      "cone": VariantCfg(spec_fn=_simple_cone_spec, weight=0.5),
-    },
-  )
-  entity = cfg.build()
-  meta = entity.variant_metadata
-  assert meta is not None
-  assert meta.variant_names == ("sphere", "cone")
-  assert meta.num_mesh_geoms == 1
-  mesh_names = [m.name for m in entity.spec.meshes]
-  assert any("sphere" in n for n in mesh_names)
-  assert any("cone" in n for n in mesh_names)
-
-
-def test_multi_geom_body_padding():
-  """Sphere (3 geoms) + cone (5 geoms) -> body padded to 5 mesh geoms."""
-  cfg = VariantEntityCfg(
-    variants={
-      "sphere": VariantCfg(spec_fn=_sphere_2col_spec, weight=0.5),
-      "cone": VariantCfg(spec_fn=_cone_4col_spec, weight=0.5),
-    },
-  )
-  entity = cfg.build()
-  meta = entity.variant_metadata
-  assert meta is not None
-  assert meta.num_mesh_geoms == 5  # max(3, 5)
-  # Sphere: 3 real + 2 padding (None).
-  assert sum(1 for n in meta.variant_mesh_names[0] if n is None) == 2
-  # Cone: 5 real, no padding.
-  assert all(n is not None for n in meta.variant_mesh_names[1])
-
-
-# Validation.
-
-
-def test_mismatched_joint_structure_raises():
-  cfg = VariantEntityCfg(
-    variants={
-      "sphere": VariantCfg(spec_fn=_simple_sphere_spec, weight=0.5),
-      "hinge": VariantCfg(spec_fn=_hinge_spec, weight=0.5),
-    },
-  )
-  with pytest.raises(ValueError, match="joint"):
-    cfg.build()
-
-
-def test_single_variant_builds():
-  """A single variant degenerates cleanly; useful for templated variant sets."""
-  cfg = VariantEntityCfg(
-    variants={"only": VariantCfg(spec_fn=_simple_sphere_spec)},
-  )
-  entity = cfg.build()
-  assert entity.variant_metadata is not None
-  assert entity.variant_metadata.variant_names == ("only",)
-
-
-def test_empty_variants_raises():
-  cfg = VariantEntityCfg(variants={})
-  with pytest.raises(ValueError, match="at least one"):
-    cfg.build()
-
-
-def _fixed_base_sphere_spec() -> mujoco.MjSpec:
-  """Fixed-base sphere variant (no free joint): currently unsupported."""
-  spec = mujoco.MjSpec()
-  m = spec.add_mesh(name="sphere")
-  m.make_sphere(subdivision=2)
-  body = spec.worldbody.add_body(name="prop")
-  body.add_geom(type=mujoco.mjtGeom.mjGEOM_MESH, meshname="sphere")
-  return spec
-
-
-def test_fixed_base_variants_rejected():
-  """Variants must be floating-base; fixed-base raises with a clear message."""
-  cfg = VariantEntityCfg(
-    variants={
-      "a": VariantCfg(spec_fn=_fixed_base_sphere_spec, weight=0.5),
-      "b": VariantCfg(spec_fn=_fixed_base_sphere_spec, weight=0.5),
-    },
-  )
-  with pytest.raises(ValueError, match="floating-base"):
-    cfg.build()
-
-
-def test_setting_spec_fn_on_variant_cfg_raises():
-  """VariantEntityCfg.spec_fn is unused; setting it should fail loudly."""
-  with pytest.raises(ValueError, match="spec_fn cannot be set"):
-    VariantEntityCfg(
-      variants={"only": VariantCfg(spec_fn=_simple_sphere_spec)},
-      spec_fn=_simple_sphere_spec,
-    )
-
-
-def test_no_variants_unchanged():
-  cfg = EntityCfg(spec_fn=_simple_sphere_spec)
-  entity = cfg.build()
-  assert entity.variant_metadata is None
-
-
-# build_mesh_variant_model: dataid and dependent fields.
-
-
-def test_dataid_assigned_per_world():
-  """Each world's geom_dataid points to its variant's meshes."""
-  scene_spec, vi = _build_scene_with_variants(_simple_sphere_spec, _simple_cone_spec)
-  result = build_mesh_variant_model(scene_spec, 4, vi)
-
-  dataid = result.wp_model.geom_dataid.numpy()
-  assert dataid.shape == (4, result.mj_model.ngeom)
-  assert dataid.ndim == 2
-
-  w2v = result.world_to_variant["object/"]
-  assert w2v[0] == 0  # variant a (sphere)
-  assert w2v[2] == 1  # variant b (cone)
-
-  # Sphere and cone worlds must have different dataid values.
-  assert not np.array_equal(dataid[0], dataid[2])
-
-
-def test_padding_slots_get_disabled():
-  """Shorter variant's padding geom slots have dataid == -1."""
-  scene_spec, vi = _build_scene_with_variants(_sphere_2col_spec, _cone_4col_spec)
-  result = build_mesh_variant_model(scene_spec, 4, vi)
-
-  dataid = result.wp_model.geom_dataid.numpy()
-  w2v = result.world_to_variant["object/"]
-
-  # Find a sphere world (variant 0, 3 mesh geoms -> 2 padding slots).
-  sphere_world = int(np.where(w2v == 0)[0][0])
-  # Find mesh geom columns (skip non-mesh geoms like worldbody).
-  mesh_geom_ids = [
-    gid
-    for gid in range(result.mj_model.ngeom)
-    if result.mj_model.geom_type[gid] == mujoco.mjtGeom.mjGEOM_MESH
-  ]
-  sphere_dataid = dataid[sphere_world, mesh_geom_ids]
-  # Last 2 mesh geom slots should be -1 (disabled padding).
-  assert sphere_dataid[-1] == -1
-  assert sphere_dataid[-2] == -1
-  # Padding slots must still be collision-enabled in the template/warp model.
-  # Short variants are disabled by per-world dataid=-1; long variants need the
-  # same slots enabled so their extra hulls can collide.
-  assert np.all(result.mj_model.geom_contype[mesh_geom_ids[-2:]] == 1)
-  assert np.all(result.mj_model.geom_conaffinity[mesh_geom_ids[-2:]] == 1)
-  assert np.all(result.wp_model.geom_contype.numpy()[mesh_geom_ids[-2:]] == 1)
-  assert np.all(result.wp_model.geom_conaffinity.numpy()[mesh_geom_ids[-2:]] == 1)
-  # First 3 should be valid (>= 0).
-  assert all(d >= 0 for d in sphere_dataid[:3])
-
-
-def test_dependent_fields_match_individual_compilation():
-  """Per-world body_mass matches independently compiled variant models."""
-  scene_spec, vi = _build_scene_with_variants(_simple_sphere_spec, _simple_cone_spec)
-  result = build_mesh_variant_model(scene_spec, 4, vi)
-
-  # Compile each variant independently for reference values.
-  sphere_model = _simple_sphere_spec().compile()
-  cone_model = _simple_cone_spec().compile()
-
-  body_mass = result.wp_model.body_mass.numpy()
-  w2v = result.world_to_variant["object/"]
-
-  sphere_w = int(np.where(w2v == 0)[0][0])
-  cone_w = int(np.where(w2v == 1)[0][0])
-
-  # The object body is the last body in the scene.
-  obj_body = result.mj_model.nbody - 1
-
-  # Mass should match individually compiled models.
-  np.testing.assert_allclose(
-    body_mass[sphere_w, obj_body],
-    sphere_model.body_mass[-1],
-    atol=1e-4,
-  )
-  np.testing.assert_allclose(
-    body_mass[cone_w, obj_body],
-    cone_model.body_mass[-1],
-    atol=1e-4,
-  )
-
-  # Sphere and cone should have different masses.
-  assert not np.isclose(body_mass[sphere_w, obj_body], body_mass[cone_w, obj_body])
-
-
-def test_select_default_values_uses_per_world_variant_defaults():
-  """Per-world defaults are indexed by env first, then by entity."""
-  from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
-  from mjlab.envs.mdp.dr._core import _select_default_values
-  from mjlab.scene import SceneCfg
-  from mjlab.terrains import TerrainEntityCfg
-
-  def _explicit_variant(
-    mesh_name: str,
-    mass: float,
-    inertia: tuple[float, float, float],
-    *,
-    cone: bool = False,
-  ) -> mujoco.MjSpec:
-    spec = mujoco.MjSpec()
-    mesh = spec.add_mesh()
-    mesh.name = mesh_name
-    if cone:
-      mesh.make_cone(nedge=8, radius=0.05)
-    else:
-      mesh.make_sphere(subdivision=1)
-    body = spec.worldbody.add_body(name="prop")
-    body.add_freejoint()
-    body.explicitinertial = 1
-    body.mass = mass
-    body.ipos[:] = (0.0, 0.0, 0.0)
-    body.inertia[:] = inertia
-    body.iquat[:] = (1.0, 0.0, 0.0, 0.0)
-    body.add_geom(
-      name="visual",
-      type=mujoco.mjtGeom.mjGEOM_MESH,
-      meshname=mesh_name,
-      contype=0,
-      conaffinity=0,
-      mass=0.0,
-    )
-    return spec
-
-  object_cfg = VariantEntityCfg(
-    variants={
-      "sphere": VariantCfg(
-        lambda: _explicit_variant("sphere", 0.2, (1e-4, 2e-4, 3e-4)),
-        weight=0.5,
-      ),
-      "cone": VariantCfg(
-        lambda: _explicit_variant("cone", 0.7, (4e-4, 5e-4, 6e-4), cone=True),
-        weight=0.5,
-      ),
-    },
-    init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
-  )
-  env_cfg = ManagerBasedRlEnvCfg(
-    decimation=1,
-    scene=SceneCfg(
-      terrain=TerrainEntityCfg(terrain_type="plane"),
-      num_envs=4,
-      env_spacing=1.0,
-      entities={"object": object_cfg},
-    ),
-  )
-
-  env = ManagerBasedRlEnv(cfg=env_cfg, device="cpu")
-  try:
-    obj_body = int(env.scene["object"].indexing.root_body_id)
-    env_ids = torch.arange(env.num_envs, device=env.device)
-    body_ids = torch.tensor([obj_body], device=env.device)
-
-    for field in ("body_mass", "body_ipos", "body_inertia", "body_iquat"):
-      selected = _select_default_values(env, field, env_ids, body_ids)
-      torch.testing.assert_close(
-        selected[:, 0],
-        getattr(env.sim.model, field)[:, obj_body],
-      )
-  finally:
-    env.close()
-
-
-def test_viser_builds_per_world_mesh_handles_for_variants():
-  """Viser dynamic meshes must not collapse all worlds onto env0's mesh."""
-  from contextlib import nullcontext
-
-  from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
-  from mjlab.scene import SceneCfg
-  from mjlab.terrains import TerrainEntityCfg
-  from mjlab.viewer.viser.scene import MjlabViserScene, _PerWorldMeshGroup
-
-  class _Handle:
-    def __init__(self, **kwargs):
-      self.visible = kwargs.get("visible", True)
-      self.batched_positions = kwargs.get("batched_positions", np.zeros((0, 3)))
-      self.batched_wxyzs = kwargs.get("batched_wxyzs", np.zeros((0, 4)))
-      self.batched_scales = kwargs.get("batched_scales")
-      self.batched_colors = kwargs.get("batched_colors")
-      self.batched_opacities = kwargs.get("batched_opacities")
-      self.position = kwargs.get("position", np.zeros(3))
-      self.wxyz = kwargs.get("wxyz", np.array([1.0, 0.0, 0.0, 0.0]))
-
-    def remove(self) -> None:
-      pass
-
-  class _Scene:
-    def __init__(self):
-      self.batched: list[tuple[tuple, dict, _Handle]] = []
-
-    def configure_environment_map(self, **_kwargs) -> None:
-      pass
-
-    def add_frame(self, *_args, **kwargs) -> _Handle:
-      return _Handle(**kwargs)
-
-    def add_grid(self, *_args, **kwargs) -> _Handle:
-      return _Handle(**kwargs)
-
-    def add_mesh_trimesh(self, *_args, **kwargs) -> _Handle:
-      return _Handle(**kwargs)
-
-    def add_batched_meshes_trimesh(self, *args, **kwargs) -> _Handle:
-      handle = _Handle(**kwargs)
-      self.batched.append((args, kwargs, handle))
-      return handle
-
-    def add_batched_meshes_simple(self, *args, **kwargs) -> _Handle:
-      handle = _Handle(**kwargs)
-      self.batched.append((args, kwargs, handle))
-      return handle
-
-  class _Server:
-    def __init__(self):
-      self.scene = _Scene()
-
-    def atomic(self):
-      return nullcontext()
-
-    def flush(self) -> None:
-      pass
-
-  env_cfg = ManagerBasedRlEnvCfg(
-    decimation=1,
-    scene=SceneCfg(
-      terrain=TerrainEntityCfg(terrain_type="plane"),
-      num_envs=4,
-      env_spacing=1.0,
-      entities={
-        "object": VariantEntityCfg(
-          variants={
-            "sphere": VariantCfg(_simple_sphere_spec, weight=0.5),
-            "cone": VariantCfg(_simple_cone_spec, weight=0.5),
-          },
-          init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
-        )
-      },
-    ),
-  )
-
-  env = ManagerBasedRlEnv(cfg=env_cfg, device="cpu")
-  try:
-    env.sim.expand_model_fields(("geom_rgba",))
-    env.sim.model.geom_rgba[:, :, :3] = torch.linspace(
-      0.2,
-      0.9,
-      env.num_envs,
-      device=env.device,
-    )[:, None, None]
-    server = _Server()
-    scene = MjlabViserScene(
-      cast(Any, server),
-      env.sim.mj_model,
-      env.num_envs,
-      sim_model=env.sim.model,
-      expanded_fields=env.sim.expanded_fields,
-    )
-    groups = [mg for mg in scene._mesh_groups if isinstance(mg, _PerWorldMeshGroup)]
-
-    assert groups
-    assert sum(len(mg.env_ids) for mg in groups) >= env.num_envs
-
-    body_xpos = env.sim.data.xpos.cpu().numpy()
-    body_xmat = env.sim.data.xmat.cpu().numpy()
-    mocap_pos = (
-      env.sim.data.mocap_pos.cpu().numpy() if env.sim.mj_model.nmocap > 0 else None
-    )
-    mocap_quat = (
-      env.sim.data.mocap_quat.cpu().numpy() if env.sim.mj_model.nmocap > 0 else None
-    )
-    scene.show_only_selected = True
-    scene.update_from_arrays(body_xpos, body_xmat, mocap_pos, mocap_quat, env_idx=0)
-    scene.update_from_arrays(body_xpos, body_xmat, mocap_pos, mocap_quat, env_idx=1)
-
-    assert any(mg.handle.visible for mg in groups)
-
-    handle_count = len(server.scene.batched)
-    env.sim.model.geom_rgba[:, :, :3] = torch.linspace(
-      0.9,
-      0.2,
-      env.num_envs,
-      device=env.device,
-    )[:, None, None]
-    scene.update_from_arrays(body_xpos, body_xmat, mocap_pos, mocap_quat, env_idx=0)
-    assert len(server.scene.batched) > handle_count
-  finally:
-    env.close()
-
-
-def test_viser_convex_hulls_are_per_variant():
-  """Convex-hull handles must differ across variants, not all show env0's hull."""
-  from contextlib import nullcontext
-
-  from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
-  from mjlab.scene import SceneCfg
-  from mjlab.terrains import TerrainEntityCfg
-  from mjlab.viewer.viser.scene import MjlabViserScene, _PerWorldHullGroup
-
-  class _Handle:
-    def __init__(self, **kwargs):
-      self.visible = kwargs.get("visible", True)
-      self.batched_positions = kwargs.get("batched_positions", np.zeros((0, 3)))
-      self.batched_wxyzs = kwargs.get("batched_wxyzs", np.zeros((0, 4)))
-      self.batched_scales = kwargs.get("batched_scales")
-      self.batched_colors = kwargs.get("batched_colors")
-      self.batched_opacities = kwargs.get("batched_opacities")
-      self.position = kwargs.get("position", np.zeros(3))
-      self.wxyz = kwargs.get("wxyz", np.array([1.0, 0.0, 0.0, 0.0]))
-      self.vertices = kwargs.get("vertices")
-      self.faces = kwargs.get("faces")
-
-    def remove(self) -> None:
-      pass
-
-  class _Scene:
-    def __init__(self):
-      self.batched: list[tuple[tuple, dict, _Handle]] = []
-
-    def configure_environment_map(self, **_kwargs) -> None:
-      pass
-
-    def add_frame(self, *_args, **kwargs) -> _Handle:
-      return _Handle(**kwargs)
-
-    def add_grid(self, *_args, **kwargs) -> _Handle:
-      return _Handle(**kwargs)
-
-    def add_mesh_trimesh(self, *_args, **kwargs) -> _Handle:
-      return _Handle(**kwargs)
-
-    def add_batched_meshes_trimesh(self, *args, **kwargs) -> _Handle:
-      handle = _Handle(**kwargs)
-      self.batched.append((args, kwargs, handle))
-      return handle
-
-    def add_batched_meshes_simple(self, path, vertices, faces, **kwargs) -> _Handle:
-      # Capture the mesh identity so the test can compare hull shapes.
-      kwargs = dict(kwargs)
-      kwargs["vertices"] = np.asarray(vertices)
-      kwargs["faces"] = np.asarray(faces)
-      handle = _Handle(**kwargs)
-      self.batched.append(((path,), kwargs, handle))
-      return handle
-
-  class _Server:
-    def __init__(self):
-      self.scene = _Scene()
-
-    def atomic(self):
-      return nullcontext()
-
-    def flush(self) -> None:
-      pass
-
-  # Sphere and cone produce visibly different convex hulls.
-  env_cfg = ManagerBasedRlEnvCfg(
-    decimation=1,
-    scene=SceneCfg(
-      terrain=TerrainEntityCfg(terrain_type="plane"),
-      num_envs=4,
-      env_spacing=1.0,
-      entities={
-        "object": VariantEntityCfg(
-          variants={
-            "sphere": VariantCfg(_simple_sphere_spec, weight=0.5),
-            "cone": VariantCfg(_simple_cone_spec, weight=0.5),
-          },
-          init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
-        )
-      },
-    ),
-  )
-
-  env = ManagerBasedRlEnv(cfg=env_cfg, device="cpu")
-  try:
-    server = _Server()
-    scene = MjlabViserScene(
-      cast(Any, server),
-      env.sim.mj_model,
-      env.num_envs,
-      sim_model=env.sim.model,
-      expanded_fields=env.sim.expanded_fields,
-    )
-    groups: list[_PerWorldHullGroup] = list(scene._hull_per_world_groups)
-    # Two distinct variants -> at least two hull handles on the same body.
-    assert len(groups) >= 2, f"expected >=2 hull variants, got {len(groups)}"
-    all_envs = np.concatenate([g.env_ids for g in groups])
-    assert sorted(all_envs.tolist()) == list(range(env.num_envs))
-    # Hulls must be shape-distinct, not all copies of env0's hull.
-    shapes = {(g.handle.vertices.shape, g.handle.faces.shape) for g in groups}
-    assert len(shapes) >= 2, (
-      f"hull variants collapsed to one shape: {shapes} "
-      "(all envs would share env0's hull)"
-    )
-
-    body_xpos = env.sim.data.xpos.cpu().numpy()
-    body_xmat = env.sim.data.xmat.cpu().numpy()
-    scene.show_convex_hull = True
-    scene.show_only_selected = True
-    for target_env in range(env.num_envs):
-      scene.update_from_arrays(body_xpos, body_xmat, env_idx=target_env)
-      visible_groups = [g for g in groups if g.handle.visible]
-      assert len(visible_groups) == 1
-      assert target_env in visible_groups[0].env_ids
-      assert visible_groups[0].handle.batched_positions.shape[0] == 1
-
-    scene.show_only_selected = False
-    scene.update_from_arrays(body_xpos, body_xmat, env_idx=0)
-    assert all(g.handle.visible for g in groups)
-  finally:
-    env.close()
-
-
-# DR consistency on variant scenes.
-
-
-def _explicit_mass_variant(
-  mesh_name: str,
-  mass: float,
-  *,
-  cone: bool = False,
-) -> mujoco.MjSpec:
-  """Build a single-geom freejoint variant with an explicit body mass."""
-  spec = mujoco.MjSpec()
-  mesh = spec.add_mesh()
-  mesh.name = mesh_name
-  if cone:
-    mesh.make_cone(nedge=8, radius=0.05)
-  else:
-    mesh.make_sphere(subdivision=1)
-  body = spec.worldbody.add_body(name="prop")
-  body.add_freejoint()
-  body.explicitinertial = 1
-  body.mass = mass
-  body.ipos[:] = (0.0, 0.0, 0.0)
-  body.inertia[:] = (1e-4, 1e-4, 1e-4)
-  body.iquat[:] = (1.0, 0.0, 0.0, 0.0)
-  body.add_geom(
-    name="visual",
-    type=mujoco.mjtGeom.mjGEOM_MESH,
-    meshname=mesh_name,
-    contype=0,
-    conaffinity=0,
-    mass=0.0,
-  )
-  return spec
-
-
-def test_dr_body_mass_scale_preserves_variant_baseline():
-  """``dr.body_mass`` scale must use each variant's own baseline.
-
-  This is the load-bearing claim of ``_per_world_default_fields``: scaling
-  body_mass on a variant scene by a per-env factor must produce
-  ``variant_default[env] * scale[env]``, not ``template_default * scale[env]``.
-  """
-  from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
-  from mjlab.envs.mdp import dr
-  from mjlab.managers.event_manager import EventTermCfg
-  from mjlab.managers.scene_entity_config import SceneEntityCfg
-  from mjlab.scene import SceneCfg
-  from mjlab.terrains import TerrainEntityCfg
-
-  light_mass = 0.1
-  heavy_mass = 1.0
-  scale = 2.0
-
-  object_cfg = VariantEntityCfg(
-    variants={
-      "light": VariantCfg(
-        lambda: _explicit_mass_variant("light", light_mass), weight=0.5
-      ),
-      "heavy": VariantCfg(
-        lambda: _explicit_mass_variant("heavy", heavy_mass, cone=True), weight=0.5
-      ),
-    },
-    init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
-  )
-  env_cfg = ManagerBasedRlEnvCfg(
-    decimation=1,
-    scene=SceneCfg(
-      terrain=TerrainEntityCfg(terrain_type="plane"),
-      num_envs=4,
-      env_spacing=1.0,
-      entities={"object": object_cfg},
-    ),
-    events={
-      "scale_mass": EventTermCfg(
-        func=dr.body_mass,
-        mode="startup",
-        params={
-          "asset_cfg": SceneEntityCfg("object", body_names=("prop",)),
-          "operation": "scale",
-          "ranges": (scale, scale),  # deterministic factor
-        },
-      ),
-    },
-  )
-
-  with pytest.warns(UserWarning, match="dr.body_mass only randomizes mass"):
-    env = ManagerBasedRlEnv(cfg=env_cfg, device="cpu")
-  try:
-    obj_body = int(env.scene["object"].indexing.root_body_id)
-    w2v = env.sim.world_to_variant["object"]
-    actual = env.sim.model.body_mass[:, obj_body].cpu()
-
-    variant_baseline = torch.tensor([light_mass, heavy_mass], dtype=actual.dtype)
-    expected = variant_baseline[w2v.cpu()] * scale
-    torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5)
-
-    # Sanity: at least one env per variant, otherwise the test is vacuous.
-    assert (w2v == 0).any() and (w2v == 1).any()
-  finally:
-    env.close()
-
-
-# Full env lifecycle.
-
-
-def test_env_step_with_variants():
-  """Build a full ManagerBasedRlEnv with variants; step without crashing."""
-  from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
-  from mjlab.envs.mdp.events import reset_root_state_uniform
-  from mjlab.managers.event_manager import EventTermCfg
-  from mjlab.managers.scene_entity_config import SceneEntityCfg
-  from mjlab.scene import SceneCfg
-  from mjlab.terrains import TerrainEntityCfg
-
-  object_cfg = VariantEntityCfg(
-    variants={
-      "sphere": VariantCfg(_simple_sphere_spec, weight=0.5),
-      "cone": VariantCfg(_simple_cone_spec, weight=0.5),
-    },
-    init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
-  )
-
-  env_cfg = ManagerBasedRlEnvCfg(
-    decimation=2,
-    scene=SceneCfg(
-      terrain=TerrainEntityCfg(terrain_type="plane"),
-      num_envs=4,
-      env_spacing=1.0,
-      entities={"object": object_cfg},
-    ),
-    events={
-      "reset": EventTermCfg(
-        func=reset_root_state_uniform,
-        mode="reset",
-        params={
-          "pose_range": {},
-          "velocity_range": {},
-          "asset_cfg": SceneEntityCfg("object"),
-        },
-      ),
-    },
-  )
-
-  env = ManagerBasedRlEnv(cfg=env_cfg, device="cpu")
-  obs, _ = env.reset()
-  actions = torch.zeros(env.num_envs, 0)
-  for _ in range(10):
-    obs, rew, term, trunc, info = env.step(actions)
-  # No NaN in positions.
-  qpos = env.sim.data.qpos[:].cpu().numpy()
-  assert np.all(np.isfinite(qpos))
-  env.close()
-
-
-# Viewer: sameframe shortcut fix.
-
-
-def _viewer_regression_sphere_spec() -> mujoco.MjSpec:
-  spec = mujoco.MjSpec()
-  m = spec.add_mesh()
-  m.name = "sphere"
-  m.make_sphere(subdivision=3)
-  m.scale[:] = (0.05, 0.05, 0.05)
-  body = spec.worldbody.add_body()
-  body.name = "prop"
-  body.add_freejoint()
-  g = body.add_geom()
-  g.name = "visual"
-  g.type = mujoco.mjtGeom.mjGEOM_MESH
-  g.meshname = "sphere"
-  return spec
-
-
-def _viewer_regression_cone_spec() -> mujoco.MjSpec:
-  spec = mujoco.MjSpec()
-  m = spec.add_mesh()
-  m.name = "cone"
-  m.make_cone(nedge=16, radius=0.04)
-  m.scale[:] = (0.05, 0.05, 0.05)
-  body = spec.worldbody.add_body()
-  body.name = "prop"
-  body.add_freejoint()
-  g = body.add_geom()
-  g.name = "visual"
-  g.type = mujoco.mjtGeom.mjGEOM_MESH
-  g.meshname = "cone"
-  return spec
-
-
-def test_sameframe_fix_makes_host_forward_match_variant():
-  """Clearing sameframe shortcuts aligns host mj_forward with variant."""
-  base_model = _viewer_regression_sphere_spec().compile()
-  cone_model = _viewer_regression_cone_spec().compile()
-
-  # Sync cone's kinematic fields onto sphere's model (like viewer does).
-  for field in (
-    "geom_size",
-    "geom_pos",
-    "geom_quat",
-    "body_mass",
-    "body_inertia",
-    "body_ipos",
-    "body_iquat",
-  ):
-    getattr(base_model, field)[:] = getattr(cone_model, field)
-
-  base_data = mujoco.MjData(base_model)
-  base_data.qpos[:] = cone_model.qpos0
-  base_data.qpos[2] = 0.05
-  mujoco.mj_forward(base_model, base_data)
-
-  cone_data = mujoco.MjData(cone_model)
-  cone_data.qpos[:] = cone_model.qpos0
-  cone_data.qpos[2] = 0.05
-  mujoco.mj_forward(cone_model, cone_data)
-
-  # Before fix: positions differ due to stale sameframe flags.
-  assert not np.allclose(base_data.geom_xpos, cone_data.geom_xpos)
-
-  # After fix: clearing sameframe makes them match.
-  disable_model_sameframe_shortcuts(base_model)
-  mujoco.mj_forward(base_model, base_data)
-  np.testing.assert_allclose(base_data.geom_xpos, cone_data.geom_xpos, atol=1e-6)
-
-
-def test_sync_model_fields_copies_only_requested_env_fields():
-  """Viewer model sync copies explicit fields and leaves others unchanged."""
-  model = _simple_sphere_spec().compile()
-
-  class _SimModel:
-    geom_rgba = torch.tensor(
-      [
-        [[0.1, 0.2, 0.3, 0.4]],
-        [[0.5, 0.6, 0.7, 0.8]],
-      ],
-      dtype=torch.float32,
-    )
-    geom_pos = torch.tensor(
-      [
-        [[1.0, 2.0, 3.0]],
-        [[4.0, 5.0, 6.0]],
-      ],
-      dtype=torch.float32,
-    )
-
-  original_geom_pos = model.geom_pos.copy()
-
-  sync_model_fields(model, _SimModel(), {"geom_rgba"}, env_idx=1)
-
-  np.testing.assert_allclose(model.geom_rgba, [[0.5, 0.6, 0.7, 0.8]])
-  np.testing.assert_allclose(model.geom_pos, original_geom_pos)
@@ -0,0 +1,213 @@
+"""Tests for sensor-based projected gravity (framezaxis up-vector sensor).
+
+The shipped robots expose a ``framezaxis`` sensor that outputs the world Z-axis in the
+IMU site frame; negating it gives projected gravity. These tests check the sensor (and
+the ``projected_gravity_from_sensor`` observation that wraps it) against an independent
+ground-truth computation, and verify that -- unlike the entity-data
+``projected_gravity_b`` -- it tracks the IMU site orientation, which is what makes IMU
+mounting domain randomization observable.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING, cast
+
+import mujoco
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.entity import EntityCfg
+from mjlab.envs.mdp import dr
+from mjlab.envs.mdp.observations import projected_gravity_from_sensor
+from mjlab.managers.scene_entity_config import SceneEntityCfg
+from mjlab.scene import Scene, SceneCfg
+from mjlab.sim.sim import Simulation, SimulationCfg
+
+if TYPE_CHECKING:
+  from mjlab.envs import ManagerBasedRlEnv
+
+# Gravity points along world -Z; projected gravity is this expressed in a body frame.
+_GRAVITY_DIR_W = (0.0, 0.0, -1.0)
+
+
+def _quat_to_mat(q: tuple[float, float, float, float]) -> torch.Tensor:
+  """Rotation matrix from a (w, x, y, z) quaternion. Independent of MuJoCo/mjlab."""
+  w, x, y, z = q
+  return torch.tensor(
+    [
+      [1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y)],
+      [2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x)],
+      [2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y)],
+    ],
+    dtype=torch.float64,
+  )
+
+
+def _expected_projected_gravity(q: tuple[float, float, float, float]) -> torch.Tensor:
+  """Ground-truth projected gravity for a body with world orientation ``q``.
+
+  proj = R(q)^T @ g_world, computed from an explicit rotation matrix so it does not
+  share a code path with the sensor or with ``projected_gravity_b``.
+  """
+  g_w = torch.tensor(_GRAVITY_DIR_W, dtype=torch.float64)
+  return _quat_to_mat(q).T @ g_w
+
+
+class Env:
+  """Minimal env stub for driving observation and dr functions in tests."""
+
+  def __init__(self, scene, sim, device):
+    self.scene = scene
+    self.sim = sim
+    self.num_envs = scene.num_envs
+    self.device = device
+
+
+def _make_env(scene, sim, device) -> ManagerBasedRlEnv:
+  """Build the env stub, typed as the real env for the functions under test."""
+  return cast("ManagerBasedRlEnv", Env(scene, sim, device))
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+def _robot_xml(site_euler: str = "0 0 0") -> str:
+  """Free-floating box with an IMU site and the framezaxis up-vector sensor."""
+  return f"""
+    <mujoco>
+      <worldbody>
+        <body name="base" pos="0 0 1">
+          <freejoint name="free_joint"/>
+          <geom name="base_geom" type="box" size="0.2 0.2 0.1" mass="5.0"/>
+          <site name="imu" pos="0.05 0 0" euler="{site_euler}"/>
+        </body>
+      </worldbody>
+      <sensor>
+        <framezaxis name="imu_upvector" objtype="body" objname="world"
+                    reftype="site" refname="imu"/>
+      </sensor>
+    </mujoco>
+  """
+
+
+def _build(xml: str, device: str, num_envs: int = 2):
+  entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(xml))
+  scene = Scene(
+    SceneCfg(num_envs=num_envs, env_spacing=3.0, entities={"robot": entity_cfg}),
+    device,
+  )
+  model = scene.compile()
+  sim = Simulation(
+    num_envs=num_envs, cfg=SimulationCfg(njmax=20), model=model, device=device
+  )
+  scene.initialize(sim.mj_model, sim.model, sim.data)
+  return scene, sim
+
+
+def _set_root_quat(robot, q: tuple[float, float, float, float], device: str) -> None:
+  root_state = robot.data.default_root_state.clone()
+  root_state[:, 3:7] = torch.tensor(q, device=device, dtype=root_state.dtype)
+  robot.write_root_state_to_sim(root_state)
+
+
+def test_sensor_matches_ground_truth_when_site_aligned(device):
+  """Sensor and entity both equal hand-computed projected gravity for a tilted base."""
+  scene, sim = _build(_robot_xml(), device)
+  robot = scene["robot"]
+
+  # Compose a 0.6 rad roll with a 0.3 rad pitch into a single root quaternion.
+  ax = (math.cos(0.3), math.sin(0.3), 0.0, 0.0)
+  ay = (math.cos(0.15), 0.0, math.sin(0.15), 0.0)
+  q = (
+    ax[0] * ay[0] - ax[1] * ay[1] - ax[2] * ay[2] - ax[3] * ay[3],
+    ax[0] * ay[1] + ax[1] * ay[0] + ax[2] * ay[3] - ax[3] * ay[2],
+    ax[0] * ay[2] - ax[1] * ay[3] + ax[2] * ay[0] + ax[3] * ay[1],
+    ax[0] * ay[3] + ax[1] * ay[2] - ax[2] * ay[1] + ax[3] * ay[0],
+  )
+  _set_root_quat(robot, q, device)
+  sim.forward()
+
+  expected = _expected_projected_gravity(q).to(device=device, dtype=torch.float32)
+  # Guard against a vacuous pass: the tilt must actually move gravity off straight-down.
+  straight_down = torch.tensor(_GRAVITY_DIR_W, device=device)
+  assert (expected - straight_down).abs().max() > 0.3
+
+  sensor_grav = -scene["robot/imu_upvector"].data
+  entity_grav = robot.data.projected_gravity_b
+  torch.testing.assert_close(sensor_grav[0], expected, atol=1e-5, rtol=0)
+  torch.testing.assert_close(entity_grav[0], expected, atol=1e-5, rtol=0)
+
+
+def test_observation_fn_tracks_site_orientation(device):
+  """The observation fn reflects IMU site tilt; the entity-data version does not.
+
+  With the base upright but the IMU site rolled 30 deg about x, projected gravity in the
+  site frame is (0, -sin30, -cos30). The entity-data version stays straight-down because
+  it uses the root body orientation and is blind to the site.
+  """
+  scene_rot, sim_rot = _build(_robot_xml(site_euler="30 0 0"), device)
+  scene_flat, sim_flat = _build(_robot_xml(site_euler="0 0 0"), device)
+  sim_rot.forward()
+  sim_flat.forward()
+
+  # Drive through the actual shipped observation function, not the raw sensor.
+  env_rot = _make_env(scene_rot, sim_rot, device)
+  env_flat = _make_env(scene_flat, sim_flat, device)
+  grav_rot = projected_gravity_from_sensor(env_rot, "robot/imu_upvector")
+  grav_flat = projected_gravity_from_sensor(env_flat, "robot/imu_upvector")
+
+  expected_rot = torch.tensor(
+    [0.0, -math.sin(math.radians(30)), -math.cos(math.radians(30))], device=device
+  )
+  straight_down = torch.tensor(_GRAVITY_DIR_W, device=device)
+  torch.testing.assert_close(grav_rot[0], expected_rot, atol=1e-5, rtol=0)
+  torch.testing.assert_close(grav_flat[0], straight_down, atol=1e-5, rtol=0)
+
+  # The entity-data version is unchanged by the site rotation (so it cannot be used to
+  # observe IMU mounting randomization), confirming why the sensor path is needed.
+  entity_rot = scene_rot["robot"].data.projected_gravity_b
+  torch.testing.assert_close(entity_rot[0], straight_down, atol=1e-5, rtol=0)
+
+
+@pytest.mark.filterwarnings(
+  "ignore:Use of index_put_ on expanded tensors is deprecated:UserWarning"
+)
+def test_site_quat_randomization_changes_sensor(device):
+  """The full DR path: running ``dr.site_quat`` perturbs the gravity observation.
+
+  This is what the G1 example configs rely on -- randomizing the IMU site orientation
+  must show up in the sensor-based projected gravity, per-environment.
+  """
+  scene, sim = _build(_robot_xml(), device, num_envs=4)
+  sim.expand_model_fields(("site_quat",))
+  env = _make_env(scene, sim, device)
+
+  sim.forward()
+  straight_down = torch.tensor(_GRAVITY_DIR_W, device=device)
+  before = projected_gravity_from_sensor(env, "robot/imu_upvector").clone()
+  # Upright base + identity site quat => straight-down gravity in every env.
+  torch.testing.assert_close(before, straight_down.expand_as(before), atol=1e-5, rtol=0)
+
+  torch.manual_seed(0)
+  dr.site_quat(
+    env,
+    env_ids=None,
+    roll_range=(-0.3, 0.3),
+    pitch_range=(-0.3, 0.3),
+    yaw_range=(-0.3, 0.3),
+    asset_cfg=SceneEntityCfg("robot", site_names=("imu",)),
+  )
+  sim.forward()
+  after = projected_gravity_from_sensor(env, "robot/imu_upvector")
+
+  # Randomization moved the reading off straight-down and made it env-dependent.
+  assert (after - straight_down).abs().max() > 0.05
+  assert not torch.allclose(after, before, atol=1e-3)
+  assert torch.unique(after, dim=0).shape[0] >= 2
+  # The perturbation is a rotation, so gravity stays a unit vector.
+  norms = torch.linalg.norm(after, dim=-1)
+  torch.testing.assert_close(norms, torch.ones_like(norms), atol=1e-5, rtol=0)
@@ -1,30 +0,0 @@
-"""Tests for mjlab.utils.random."""
-
-import subprocess
-import sys
-import textwrap
-
-
-def test_seed_rng_cpu_device_does_not_initialize_warp_cuda() -> None:
-  """seed_rng(device="cpu") must not initialize Warp's CUDA runtime.
-
-  Runs in a subprocess so that Warp is guaranteed uninitialized before the
-  call.
-  """
-  script = textwrap.dedent("""
-    import warp as wp
-    from mjlab.utils.random import seed_rng
-
-    assert wp._src.context.runtime is None, "Warp must not be initialized yet"
-    seed_rng(42, device="cpu")
-    rt = wp._src.context.runtime
-    if rt is not None:
-        cuda = [d for d in wp.get_devices() if "cuda" in str(d)]
-        assert not cuda, f"seed_rng(device='cpu') initialized CUDA devices {cuda}"
-  """)
-  result = subprocess.run(
-    [sys.executable, "-c", script], capture_output=True, text=True
-  )
-  assert result.returncode == 0, (
-    f"subprocess failed:\nstdout={result.stdout}\nstderr={result.stderr}"
-  )
@@ -936,7 +936,7 @@ def test_multi_frame_body_exclusion(device):
  should skip body_b's own geom but HIT body_a's platform. Frame A's
  rays should skip body_a and hit the floor.
  """
-  xml = """
+  body_a_xml = """
    <mujoco>
      <worldbody>
        <geom name="floor" type="plane" size="10 10 0.1" pos="0 0 0"/>
@@ -945,6 +945,12 @@ def test_multi_frame_body_exclusion(device):
          <geom name="geom_a" type="box" size="2 2 0.1" mass="5.0"/>
          <site name="site_a" pos="0 0 0"/>
        </body>
+      </worldbody>
+    </mujoco>
+  """
+  body_b_xml = """
+    <mujoco>
+      <worldbody>
        <body name="body_b" pos="0 0 3">
          <freejoint name="free_b"/>
          <geom name="geom_b" type="box" size="0.5 0.5 0.5" mass="5.0"/>
@@ -957,15 +963,17 @@ def test_multi_frame_body_exclusion(device):
  cfg = RayCastSensorCfg(
    name="multi",
    frame=(
-      ObjRef(type="site", name="site_a", entity="robot"),
-      ObjRef(type="site", name="site_b", entity="robot"),
+      ObjRef(type="site", name="site_a", entity="body_a"),
+      ObjRef(type="site", name="site_b", entity="body_b"),
    ),
    pattern=GridPatternCfg(size=(0.0, 0.0), resolution=0.1),
    max_distance=10.0,
    exclude_parent_body=True,
  )

-  scene, sim = make_scene_and_sim(device, xml, (cfg,))
+  scene, sim = make_scene_and_sim(
+    device, {"body_a": body_a_xml, "body_b": body_b_xml}, (cfg,)
+  )
  sim.step()
  sim.sense()

@@ -4,6 +4,7 @@ import ast
 import tempfile
 from dataclasses import asdict
 from pathlib import Path
+from unittest.mock import MagicMock, patch

 import mujoco
 import onnx
@@ -521,3 +522,85 @@ def test_onnx_motion_model_clamps_out_of_bounds_time_step():
    _, joint_pos, *_ = model(x, time_step)

  torch.testing.assert_close(joint_pos, motion.joint_pos[num_steps - 1 : num_steps])
+
+
+def _make_tracking_runner_shell(registry_name, logger_type, upload_model=True):
+  """Build a MotionTrackingOnPolicyRunner with all heavy parts mocked out."""
+  from mjlab.tasks.tracking.rl.runner import MotionTrackingOnPolicyRunner
+
+  runner = MotionTrackingOnPolicyRunner.__new__(MotionTrackingOnPolicyRunner)
+  runner.registry_name = registry_name
+  runner.cfg = {"upload_model": upload_model}
+  runner.logger = MagicMock()
+  runner.logger.logger_type = logger_type
+
+  mock_motion_term = MagicMock()
+  mock_motion_term.cfg.anchor_body_name = "pelvis"
+  mock_motion_term.cfg.body_names = ["body1"]
+  runner.env = MagicMock()
+  runner.env.unwrapped.command_manager.get_term.return_value = mock_motion_term
+  return runner
+
+
+@pytest.mark.parametrize("logger_type", ["wandb", "WandbLogWriter"])
+def test_tracking_runner_registers_artifact_for_wandb_logger_types(
+  logger_type, monkeypatch, tmp_path
+):
+  """use_artifact is called for both legacy 'wandb' and current 'WandbLogWriter' logger types.
+
+  Regression test: rsl-rl-lib 5.4 renamed the WandB logger type from 'wandb'
+  to 'WandbLogWriter'. If only 'wandb' is checked, use_artifact is silently
+  skipped and the nightly report fails with 'No motion artifact found in the run.'
+  """
+  from mjlab.rl.runner import MjlabOnPolicyRunner
+  from mjlab.tasks.tracking.rl import runner as runner_mod
+
+  runner = _make_tracking_runner_shell("org/motions/motion:latest", logger_type)
+
+  monkeypatch.setattr(MjlabOnPolicyRunner, "save", lambda *a, **kw: None)
+  monkeypatch.setattr(runner_mod, "get_base_metadata", lambda *a: {})
+  monkeypatch.setattr(runner_mod, "attach_metadata_to_onnx", lambda *a: None)
+  monkeypatch.setattr(
+    runner.env.unwrapped.__class__,
+    "export_policy_to_onnx",
+    lambda *a, **kw: None,
+    raising=False,
+  )
+
+  checkpoint = tmp_path / "run-dir" / "model_100.pt"
+  checkpoint.parent.mkdir()
+  checkpoint.touch()
+
+  mock_run = MagicMock()
+  mock_run.name = "test-run"
+
+  with patch.object(runner_mod, "wandb") as mock_wandb:
+    mock_wandb.run = mock_run
+    runner.export_policy_to_onnx = MagicMock()
+    runner.save(str(checkpoint))
+
+  mock_run.use_artifact.assert_called_once_with("org/motions/motion:latest")
+
+
+def test_tracking_runner_does_not_register_artifact_for_tensorboard(
+  monkeypatch, tmp_path
+):
+  """use_artifact is NOT called when using the tensorboard logger."""
+  from mjlab.rl.runner import MjlabOnPolicyRunner
+  from mjlab.tasks.tracking.rl import runner as runner_mod
+
+  runner = _make_tracking_runner_shell("org/motions/motion:latest", "tensorboard")
+
+  monkeypatch.setattr(MjlabOnPolicyRunner, "save", lambda *a, **kw: None)
+  monkeypatch.setattr(runner_mod, "get_base_metadata", lambda *a: {})
+  monkeypatch.setattr(runner_mod, "attach_metadata_to_onnx", lambda *a: None)
+
+  checkpoint = tmp_path / "run-dir" / "model_100.pt"
+  checkpoint.parent.mkdir()
+  checkpoint.touch()
+
+  with patch.object(runner_mod, "wandb") as mock_wandb:
+    runner.export_policy_to_onnx = MagicMock()
+    runner.save(str(checkpoint))
+
+  mock_wandb.run.use_artifact.assert_not_called()
@@ -2,8 +2,15 @@

 import mujoco
 import numpy as np
+import pytest

-from mjlab.terrains.primitive_terrains import BoxSteppingStonesTerrainCfg
+from mjlab.terrains.config import ALL_TERRAIN_PRESETS
+from mjlab.terrains.primitive_terrains import (
+  _MIN_BORDER_HEIGHT,
+  BoxInvertedPyramidStairsTerrainCfg,
+  BoxPyramidStairsTerrainCfg,
+  BoxSteppingStonesTerrainCfg,
+)

 _CFG = BoxSteppingStonesTerrainCfg(
  proportion=1.0,
@@ -37,12 +44,10 @@ def _generate_stones(
    if geom is None:
      continue
    pos, size = geom.pos, geom.size
-    # Skip platform, floor, and border geoms.
-    is_platform = (
-      np.isclose(pos[0], center)
-      and np.isclose(pos[1], center)
-      and np.isclose(size[0], cfg.platform_width / 2, atol=1e-4)
-    )
+    # Skip platform, floor, and border geoms. The platform is the geom centered
+    # exactly at the patch center (its size is grid-snapped, not the configured
+    # width, so it is identified by position alone).
+    is_platform = np.isclose(pos[0], center) and np.isclose(pos[1], center)
    is_full_span = np.isclose(size[0], cfg.size[0] / 2) or np.isclose(
      size[1], cfg.size[1] / 2
    )
@@ -74,3 +79,50 @@ def test_stone_size_decreases_with_difficulty():
    sizes[difficulty] = np.mean([hx + hy for _, _, hx, hy in stones])

  assert sizes[0.0] > sizes[1.0]
+
+
+@pytest.mark.parametrize(
+  "cfg_cls", [BoxPyramidStairsTerrainCfg, BoxInvertedPyramidStairsTerrainCfg]
+)
+def test_pyramid_stairs_border_present_at_zero_difficulty(cfg_cls):
+  """At difficulty 0 the step height collapses to 0, but the flat border frame
+  must still be generated as solid, non-degenerate geometry (regression for the
+  empty-boundary bug, issue #1033)."""
+  cfg = cfg_cls(
+    size=(8.0, 8.0),
+    step_height_range=(0.0, 0.2),
+    step_width=0.3,
+    platform_width=3.0,
+    border_width=1.0,
+  )
+  spec = mujoco.MjSpec()
+  spec.worldbody.add_body(name="terrain")
+  output = cfg.function(difficulty=0.0, spec=spec, rng=np.random.default_rng(0))
+
+  # The border frame sits below z=0 (top flush at ground level); inner step
+  # boxes are centered at z=0. Identify the frame by its downward offset.
+  border_geoms = [
+    g.geom for g in output.geometries if g.geom is not None and g.geom.pos[2] < -1e-4
+  ]
+  assert len(border_geoms) == 4, "Expected four border frame boxes."
+  for geom in border_geoms:
+    # Each frame box must be solid, not a degenerate zero-height geom, and its
+    # top must be flush with the ground plane at z=0.
+    assert geom.size[2] >= _MIN_BORDER_HEIGHT / 2 - 1e-9
+    assert np.isclose(geom.pos[2] + geom.size[2], 0.0, atol=1e-6)
+
+
+@pytest.mark.parametrize("preset_name", sorted(ALL_TERRAIN_PRESETS))
+@pytest.mark.parametrize("difficulty", [0.0, 1.0])
+def test_preset_compiles_across_difficulty(preset_name, difficulty):
+  """Every terrain preset must generate compilable MuJoCo geometry across the
+  full difficulty range. Difficulty 0 is exercised explicitly because curriculum
+  row 0 lands there deterministically, which previously produced degenerate
+  geometry (zero-height hfields, NaN colors, missing borders)."""
+  cfg = ALL_TERRAIN_PRESETS[preset_name](size=(8.0, 8.0))
+  spec = mujoco.MjSpec()
+  spec.worldbody.add_body(name="terrain")
+  cfg.function(difficulty=difficulty, spec=spec, rng=np.random.default_rng(0))
+  # Compiling validates geom/hfield sizes and rgba values (catches NaNs and
+  # non-positive sizes that MuJoCo rejects).
+  spec.compile()
@@ -1,5 +1,6 @@
 """Tests for motion tracking evaluation metrics."""

+import math
 from unittest.mock import Mock

 import pytest
@@ -34,11 +35,11 @@ def mock_command():


 def test_mpkpe_zero_when_positions_match(mock_command):
-  """Test MPKPE is zero when positions are identical."""
+  """Test MPKPE is zero when global positions are identical."""
  num_bodies = len(mock_command.cfg.body_names)
  positions = torch.rand(mock_command.num_envs, num_bodies, 3)

-  mock_command.body_pos_relative_w = positions.clone()
+  mock_command.body_pos_w = positions.clone()
  mock_command.robot_body_pos_w = positions.clone()

  mpkpe = compute_mpkpe(mock_command)
@@ -48,10 +49,10 @@ def test_mpkpe_zero_when_positions_match(mock_command):


 def test_mpkpe_correct_error(mock_command):
-  """Test MPKPE computes correct mean error."""
+  """Test MPKPE computes the correct mean global error."""
  num_bodies = len(mock_command.cfg.body_names)

-  mock_command.body_pos_relative_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
+  mock_command.body_pos_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
  mock_command.robot_body_pos_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
  mock_command.robot_body_pos_w[:, :, 0] = 1.0  # 1 unit offset in x

@@ -60,58 +61,71 @@ def test_mpkpe_correct_error(mock_command):
  assert torch.allclose(mpkpe, torch.ones(mock_command.num_envs), atol=1e-6)


-def test_r_mpkpe_invariant_to_global_translation(mock_command):
-  """Test R-MPKPE is invariant to global translation."""
+def test_mpkpe_uses_global_reference(mock_command):
+  """MPKPE must read the global reference, not the drift-cancelled one.
+
+  Pins issue #1006: setting body_pos_relative_w to match the robot exactly
+  would yield zero error if it were (incorrectly) used; the metric must
+  instead follow body_pos_w.
+  """
  num_bodies = len(mock_command.cfg.body_names)
+  robot_pos = torch.rand(mock_command.num_envs, num_bodies, 3)
+  mock_command.robot_body_pos_w = robot_pos.clone()
+  mock_command.body_pos_relative_w = robot_pos.clone()  # zero error if misused
+  mock_command.body_pos_w = robot_pos.clone()
+  mock_command.body_pos_w[:, :, 0] += 1.0  # 1 unit of global drift

-  mock_command.anchor_pos_w = torch.zeros(mock_command.num_envs, 3)
-  mock_command.body_pos_w = torch.rand(mock_command.num_envs, num_bodies, 3)
-  mock_command.robot_anchor_pos_w = torch.zeros(mock_command.num_envs, 3)
-  mock_command.robot_body_pos_w = mock_command.body_pos_w.clone()
+  mpkpe = compute_mpkpe(mock_command)

-  r_mpkpe_1 = compute_root_relative_mpkpe(mock_command)
-
-  # Translate everything by large offset.
-  offset = torch.tensor([100.0, 200.0, 300.0])
-  mock_command.anchor_pos_w = offset.expand(mock_command.num_envs, 3).clone()
-  mock_command.body_pos_w = mock_command.body_pos_w + offset
-  mock_command.robot_anchor_pos_w = offset.expand(mock_command.num_envs, 3).clone()
-  mock_command.robot_body_pos_w = mock_command.robot_body_pos_w + offset
-
-  r_mpkpe_2 = compute_root_relative_mpkpe(mock_command)
-
-  assert torch.allclose(r_mpkpe_1, r_mpkpe_2, atol=1e-5)
+  assert torch.allclose(mpkpe, torch.ones(mock_command.num_envs), atol=1e-6)


-def test_r_mpkpe_detects_relative_error(mock_command):
-  """Test R-MPKPE detects errors in relative positions."""
+def test_r_mpkpe_zero_when_relative_positions_match(mock_command):
+  """R-MPKPE is zero when re-anchored positions are identical."""
  num_bodies = len(mock_command.cfg.body_names)
+  positions = torch.rand(mock_command.num_envs, num_bodies, 3)

-  mock_command.anchor_pos_w = torch.zeros(mock_command.num_envs, 3)
-  mock_command.body_pos_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
-  mock_command.body_pos_w[:, :, 0] = 1.0  # Bodies 1 unit from anchor
+  mock_command.body_pos_relative_w = positions.clone()
+  mock_command.robot_body_pos_w = positions.clone()

-  mock_command.robot_anchor_pos_w = torch.zeros(mock_command.num_envs, 3)
-  mock_command.robot_body_pos_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
-  mock_command.robot_body_pos_w[:, :, 0] = 2.0  # Bodies 2 units from anchor
+  r_mpkpe = compute_root_relative_mpkpe(mock_command)
+
+  assert r_mpkpe.shape == (mock_command.num_envs,)
+  assert torch.allclose(r_mpkpe, torch.zeros(mock_command.num_envs), atol=1e-6)
+
+
+def test_r_mpkpe_uses_relative_reference(mock_command):
+  """R-MPKPE reads the re-anchored reference, not the global one.
+
+  Setting body_pos_w to match the robot exactly would yield zero error if
+  it were (incorrectly) used; the metric must instead follow
+  body_pos_relative_w.
+  """
+  num_bodies = len(mock_command.cfg.body_names)
+  robot_pos = torch.rand(mock_command.num_envs, num_bodies, 3)
+  mock_command.robot_body_pos_w = robot_pos.clone()
+  mock_command.body_pos_w = robot_pos.clone()  # zero error if misused
+  mock_command.body_pos_relative_w = robot_pos.clone()
+  mock_command.body_pos_relative_w[:, :, 0] += 1.0  # 1 unit of local pose error

  r_mpkpe = compute_root_relative_mpkpe(mock_command)

  assert torch.allclose(r_mpkpe, torch.ones(mock_command.num_envs), atol=1e-6)


-def test_joint_velocity_error(mock_command):
-  """Test joint velocity error computes correct L2 norm."""
+def test_joint_velocity_error_rms(mock_command):
+  """Joint velocity error is the per-joint RMS of the velocity error."""
  num_joints = 3

  mock_command.joint_vel = torch.zeros(mock_command.num_envs, num_joints)
  mock_command.robot_joint_vel = torch.zeros(mock_command.num_envs, num_joints)
  mock_command.robot_joint_vel[:, 0] = 3.0
-  mock_command.robot_joint_vel[:, 1] = 4.0  # Error [3, 4, 0] has norm 5
+  mock_command.robot_joint_vel[:, 1] = 4.0  # Error [3, 4, 0]

  error = compute_joint_velocity_error(mock_command)

-  assert torch.allclose(error, torch.ones(mock_command.num_envs) * 5.0, atol=1e-6)
+  expected = math.sqrt((3.0**2 + 4.0**2 + 0.0**2) / num_joints)
+  assert torch.allclose(error, torch.ones(mock_command.num_envs) * expected, atol=1e-6)


 def test_ee_position_error_only_uses_specified_bodies(mock_command):
@@ -153,3 +167,13 @@ def test_ee_orientation_error_detects_rotation(mock_command):
  # Error should be approximately pi/2 radians.
  expected = torch.ones(mock_command.num_envs) * (3.14159 / 2)
  assert torch.allclose(error, expected, atol=0.01)
+
+
+def test_ee_metrics_raise_on_unknown_body(mock_command):
+  """Unknown end-effector names raise instead of silently scoring zero."""
+  num_bodies = len(mock_command.cfg.body_names)
+  mock_command.body_pos_relative_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
+  mock_command.robot_body_pos_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
+
+  with pytest.raises(ValueError, match="not tracked"):
+    compute_ee_position_error(mock_command, ("nonexistent_body",))
@@ -4,7 +4,7 @@ import mujoco
 import pytest
 from conftest import get_test_device

-from mjlab.actuator import XmlActuatorCfg
+from mjlab.actuator import XmlActuator, XmlActuatorCfg
 from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
 from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg, mdp
 from mjlab.managers.observation_manager import ObservationGroupCfg, ObservationTermCfg
@@ -160,6 +160,7 @@ def test_xml_actuator_explicit_command_field_bypasses_detection():
  entity.compile()

  actuator = entity._actuators[0]
+  assert isinstance(actuator, XmlActuator)
  assert actuator.command_field == "effort"
  assert actuator._target_names == ["joint1"]

@@ -42,7 +42,7 @@ conflicts = [[

 [manifest]
 constraints = [
-    { name = "gitpython", specifier = ">=3.1.47" },
+    { name = "gitpython", specifier = ">=3.1.49" },
    { name = "lxml", specifier = ">=6.1.0" },
 ]
 overrides = [{ name = "mujoco", specifier = ">=3.8.0.dev0", index = "https://py.mujoco.org/" }]
@@ -823,14 +823,14 @@ wheels = [

 [[package]]
 name = "gitpython"
-version = "3.1.47"
+version = "3.1.50"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
    { name = "gitdb" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/c1/bd/50db468e9b1310529a19fce651b3b0e753b5c07954d486cba31bbee9a5d5/gitpython-3.1.47.tar.gz", hash = "sha256:dba27f922bd2b42cb54c87a8ab3cb6beb6bf07f3d564e21ac848913a05a8a3cd", size = 216978, upload-time = "2026-04-22T02:44:44.059Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/33/f6/354ae6491228b5eb40e10d89c4d13c651fe1cf7556e35ebdded50cff57ce/gitpython-3.1.50.tar.gz", hash = "sha256:80da2d12504d52e1f998772dc5baf6e553f8d2fcfe1fcc226c9d9a2ee3372dcc", size = 219798, upload-time = "2026-05-06T04:01:26.571Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/f2/c5/a1bc0996af85757903cf2bf444a7824e68e0035ce63fb41d6f76f9def68b/gitpython-3.1.47-py3-none-any.whl", hash = "sha256:489f590edfd6d20571b2c0e72c6a6ac6915ee8b8cd04572330e3842207a78905", size = 209547, upload-time = "2026-04-22T02:44:41.271Z" },
+    { url = "https://files.pythonhosted.org/packages/20/7a/1c6e3562dfd8950adbb11ffbc65d21e7c89d01a6e4f137fa981056de25c5/gitpython-3.1.50-py3-none-any.whl", hash = "sha256:d352abe2908d07355014abdd21ddf798c2a961469239afec4962e9da884858f9", size = 212507, upload-time = "2026-05-06T04:01:23.799Z" },
 ]

 [[package]]
@@ -1647,7 +1647,7 @@ wheels = [

 [[package]]
 name = "mjlab"
-version = "1.3.0"
+version = "1.4.0"
 source = { editable = "." }
 dependencies = [
    { name = "imageio-ffmpeg" },
@@ -1658,6 +1658,8 @@ dependencies = [
    { name = "onnxscript" },
    { name = "prettytable" },
    { name = "rsl-rl-lib" },
+    { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-5-mjlab-cpu' and extra == 'extra-5-mjlab-cu128')" },
+    { name = "scipy", version = "1.16.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-5-mjlab-cpu' and extra == 'extra-5-mjlab-cu128')" },
    { name = "tensorboard" },
    { name = "tensordict" },
    { name = "torch", version = "2.9.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or (extra == 'extra-5-mjlab-cpu' and extra == 'extra-5-mjlab-cu128') or (extra != 'extra-5-mjlab-cpu' and extra != 'extra-5-mjlab-cu128')" },
@@ -1715,12 +1717,13 @@ docs = [
 requires-dist = [
    { name = "imageio-ffmpeg" },
    { name = "mediapy", specifier = ">=1.2.6" },
-    { name = "mjviser", git = "https://github.com/mujocolab/mjviser?rev=1bdfd6fe79066b847a5f430000fcfbb53ec31a6f" },
-    { name = "mujoco", specifier = ">=3.8.0", index = "https://py.mujoco.org/" },
-    { name = "mujoco-warp", git = "https://github.com/google-deepmind/mujoco_warp?rev=6f235d4" },
+    { name = "mjviser", specifier = ">=0.0.14" },
+    { name = "mujoco", specifier = "~=3.8.0", index = "https://py.mujoco.org/" },
+    { name = "mujoco-warp", git = "https://github.com/google-deepmind/mujoco_warp?rev=88b55fc2696960b927bc12584994bb8412b36558" },
    { name = "onnxscript", specifier = ">=0.5.4" },
    { name = "prettytable" },
-    { name = "rsl-rl-lib", specifier = "==5.2.0" },
+    { name = "rsl-rl-lib", specifier = "==5.4.0" },
+    { name = "scipy", specifier = ">=1.15" },
    { name = "tensorboard", specifier = ">=2.20.0" },
    { name = "tensordict" },
    { name = "torch", specifier = ">=2.7.0" },
@@ -1732,7 +1735,7 @@ requires-dist = [
    { name = "tqdm" },
    { name = "trimesh", specifier = ">=4.8.3" },
    { name = "tyro", specifier = ">=1.0.1" },
-    { name = "viser", specifier = ">=1.0.26" },
+    { name = "viser", specifier = ">=1.0.27" },
    { name = "wandb", specifier = ">=0.22.3" },
    { name = "warp-lang", marker = "sys_platform != 'darwin'", specifier = ">=1.12.0", index = "https://pypi.nvidia.com/" },
    { name = "warp-lang", marker = "sys_platform == 'darwin'", specifier = ">=1.12.0" },
@@ -1767,8 +1770,8 @@ docs = [

 [[package]]
 name = "mjviser"
-version = "0.0.13"
-source = { git = "https://github.com/mujocolab/mjviser?rev=1bdfd6fe79066b847a5f430000fcfbb53ec31a6f#1bdfd6fe79066b847a5f430000fcfbb53ec31a6f" }
+version = "0.0.14"
+source = { registry = "https://pypi.org/simple" }
 dependencies = [
    { name = "mujoco" },
    { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-5-mjlab-cpu' and extra == 'extra-5-mjlab-cu128')" },
@@ -1777,6 +1780,10 @@ dependencies = [
    { name = "trimesh" },
    { name = "viser" },
 ]
+sdist = { url = "https://files.pythonhosted.org/packages/1d/e4/eef89b279fb1811b5f120a99ac9284c32ff2ca4fad5e6f5c93035f72ba9a/mjviser-0.0.14.tar.gz", hash = "sha256:ebde2203dab89959a13ae549b4d3e5e5cf9eb69de11a1a2fd759cbe8f8c641f3", size = 29576, upload-time = "2026-05-07T03:35:13.128Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/9c/c2/4534d678ad1b3f7dee6fd83110112800287be21ba007d988abb9ddc8e0ac/mjviser-0.0.14-py3-none-any.whl", hash = "sha256:4b09f8e90506fc4a71d76fc628872147157947e9292428213ab78cfe137c9a26", size = 32296, upload-time = "2026-05-07T03:35:14.252Z" },
+]

 [[package]]
 name = "ml-dtypes"
@@ -1958,8 +1965,8 @@ wheels = [

 [[package]]
 name = "mujoco-warp"
-version = "3.8.0"
-source = { git = "https://github.com/google-deepmind/mujoco_warp?rev=6f235d4#6f235d46cb2ecf8f37c8f967f8dd9c87d0ca5807" }
+version = "3.8.0.2"
+source = { git = "https://github.com/google-deepmind/mujoco_warp?rev=88b55fc2696960b927bc12584994bb8412b36558#88b55fc2696960b927bc12584994bb8412b36558" }
 dependencies = [
    { name = "absl-py" },
    { name = "etils", extra = ["epath"] },
@@ -2542,7 +2549,7 @@ wheels = [

 [[package]]
 name = "paramiko"
-version = "4.0.0"
+version = "5.0.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
    { name = "bcrypt" },
@@ -2550,9 +2557,9 @@ dependencies = [
    { name = "invoke" },
    { name = "pynacl" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/1f/e7/81fdcbc7f190cdb058cffc9431587eb289833bdd633e2002455ca9bb13d4/paramiko-4.0.0.tar.gz", hash = "sha256:6a25f07b380cc9c9a88d2b920ad37167ac4667f8d9886ccebd8f90f654b5d69f", size = 1630743, upload-time = "2025-08-04T01:02:03.711Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/62/93/dcc25d52f49022ae6175d15e6bd751f1acc99b98bc61fc55e5155a7be2e7/paramiko-5.0.0.tar.gz", hash = "sha256:36763b5b95c2a0dcfdf1abc48e48156ee425b21efe2f0e787c2dd5a95c0e5e79", size = 1548586, upload-time = "2026-05-09T18:28:52.256Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/a9/90/a744336f5af32c433bd09af7854599682a383b37cfd78f7de263de6ad6cb/paramiko-4.0.0-py3-none-any.whl", hash = "sha256:0e20e00ac666503bf0b4eda3b6d833465a2b7aff2e2b3d79a8bba5ef144ee3b9", size = 223932, upload-time = "2025-08-04T01:02:02.029Z" },
+    { url = "https://files.pythonhosted.org/packages/82/5b/eadf6d45de38d30ab603f49393b6cd2cbe7e233af8cf90197e32782b68a9/paramiko-5.0.0-py3-none-any.whl", hash = "sha256:b7044611c30140d9a75261653210e2002977b71a0497ff3ba0d98d7edbf62f7c", size = 208919, upload-time = "2026-05-09T18:28:50.295Z" },
 ]

 [[package]]
@@ -3253,7 +3260,7 @@ wheels = [

 [[package]]
 name = "rsl-rl-lib"
-version = "5.2.0"
+version = "5.4.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
    { name = "gitpython" },
@@ -3261,6 +3268,7 @@ dependencies = [
    { name = "numpy", version = "2.3.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-5-mjlab-cpu' and extra == 'extra-5-mjlab-cu128')" },
    { name = "onnx" },
    { name = "onnxscript" },
+    { name = "tensorboard" },
    { name = "tensordict" },
    { name = "torch", version = "2.9.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or (extra == 'extra-5-mjlab-cpu' and extra == 'extra-5-mjlab-cu128') or (extra != 'extra-5-mjlab-cpu' and extra != 'extra-5-mjlab-cu128')" },
    { name = "torch", version = "2.9.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "(sys_platform != 'darwin' and extra == 'extra-5-mjlab-cu128') or (extra == 'extra-5-mjlab-cpu' and extra == 'extra-5-mjlab-cu128')" },
@@ -3268,9 +3276,9 @@ dependencies = [
    { name = "torchvision", version = "0.24.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or extra != 'extra-5-mjlab-cpu' or (extra == 'extra-5-mjlab-cpu' and extra == 'extra-5-mjlab-cu128')" },
    { name = "torchvision", version = "0.25.0", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform != 'darwin' and extra == 'extra-5-mjlab-cpu') or (extra == 'extra-5-mjlab-cpu' and extra == 'extra-5-mjlab-cu128')" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/d7/a8/fb9aae0573a83dd510e228085f5e6ff9076a91f2bff6699270f07d31854c/rsl_rl_lib-5.2.0.tar.gz", hash = "sha256:cbb4eee96af9574495208381115d45d68ad1c0403710a2bc6512a2ee5bf57124", size = 60558, upload-time = "2026-04-23T12:40:54.259Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/0f/51/2d4c95b3642c0f659fbcddcd17fa0401903733250fa382f51f9b913bb85f/rsl_rl_lib-5.4.0.tar.gz", hash = "sha256:e1aa5cd5771f2d9a9e7a7ba5456b942ab588410cfd397b3c63e32d00a0744f0d", size = 65902, upload-time = "2026-05-27T10:42:39.656Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/7e/3a/3e8f39049cc5a8994bfdb2dd09d727f90a2781b9755adf59e49fa509500e/rsl_rl_lib-5.2.0-py3-none-any.whl", hash = "sha256:fc767059f329a184527dd10766c3382354f6deca4b049f23febc8140c3dc6029", size = 86451, upload-time = "2026-04-23T12:40:52.762Z" },
+    { url = "https://files.pythonhosted.org/packages/f8/dd/6797be77ce0cc12881f29cd9363b791f15276fde6a136a73443b6b6c2ce3/rsl_rl_lib-5.4.0-py3-none-any.whl", hash = "sha256:b30a0e59dac0ef7236f8f793c9bedfa8f2b8f8d29c2965140feeb1439153ebab", size = 92871, upload-time = "2026-05-27T10:42:38.415Z" },
 ]

 [[package]]
@@ -4593,11 +4601,11 @@ wheels = [

 [[package]]
 name = "urllib3"
-version = "2.6.3"
+version = "2.7.0"
 source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/c7/24/5f1b3bdffd70275f6661c76461e25f024d5a38a46f04aaca912426a2b1d3/urllib3-2.6.3.tar.gz", hash = "sha256:1b62b6884944a57dbe321509ab94fd4d3b307075e0c2eae991ac71ee15ad38ed", size = 435556, upload-time = "2026-01-07T16:24:43.925Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl", hash = "sha256:bf272323e553dfb2e87d9bfd225ca7b0f467b919d7bbd355436d3fd37cb0acd4", size = 131584, upload-time = "2026-01-07T16:24:42.685Z" },
+    { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" },
 ]

 [[package]]
@@ -4672,7 +4680,7 @@ wheels = [

 [[package]]
 name = "viser"
-version = "1.0.26"
+version = "1.0.27"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
    { name = "imageio" },
@@ -4688,9 +4696,9 @@ dependencies = [
    { name = "yourdfpy" },
    { name = "zstandard" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/12/ce/82a0e50fae21f5e02fcc5d9aff2ab59dccb9c319b6c4cf528f2228049b05/viser-1.0.26.tar.gz", hash = "sha256:dc08c6f505e70324b0603bdddf9714c00ac828c259ee49abd8ad094bfc90c91c", size = 4828261, upload-time = "2026-03-30T11:43:19.513Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/fd/f5/48adb4e5e4234f48e96a1e7fc50cca6731280df0c279833e333963f9ea5c/viser-1.0.27.tar.gz", hash = "sha256:87e3239d6c1c2c003db93ac4072430ec790e336ffe7214781f035e54faebc0af", size = 4897986, upload-time = "2026-05-06T10:30:47.556Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/91/f7/762a2d5fab509d0c632b271e21e634462397cc02cca649771c3e9d2e0bcc/viser-1.0.26-py3-none-any.whl", hash = "sha256:03b177b4ef584f58f7b74fdf44cccb165b8a220ffd90728ef5c1e3d1b1fcf258", size = 4922888, upload-time = "2026-03-30T11:43:21.355Z" },
+    { url = "https://files.pythonhosted.org/packages/7c/ad/8ae712579e294b4395fb39f7d65524b51fc7b731eacce26af096b7e59b61/viser-1.0.27-py3-none-any.whl", hash = "sha256:8da5b7934416e6e2d3a7ebcf39fc840f21030b51eb63231e8cfef457bfb49031", size = 4998748, upload-time = "2026-05-06T10:30:49.965Z" },
 ]

 [[package]]
@@ -0,0 +1,57 @@
+# MuJoCo 独立工具集
+
+本目录保存不依赖训练循环的 MuJoCo、姿态、IK、动力学和 MPC 分析工具。脚本通过父目录读取 `../mjcf/scene.xml` 与 `../mjcf/wheelleg.xml`，因此应从 `rc_mjlab` 工程根目录运行。
+
+## 工具分类
+
+| 入口 | 用途 |
+| --- | --- |
+| `posture_tool.py` | 基于解析运动学快速生成站立和低姿态参数表 |
+| `rl_friendly_opt.py` | 按轮心位置、关节力矩和雅可比条件数筛选适合 RL 的姿态 |
+| `posture_optimizer.py` | 解析计算与 MuJoCo 扫描结合的姿态优化 |
+| `static_posture_optimizer.py` | 在重力和地面接触下评估静态站姿、支撑域和离地间隙 |
+| `ik_diff_sweep.py` | 扫描 IK 姿态和差速轮跟踪参数，可导出 JSON |
+| `run.py` | 启动完整 MuJoCo 控制、GUI 和 MPC 调试链路 |
+| `robot.py`、`controller.py` | 仿真机器人接口和控制器 |
+| `dynamics.py`、`mpc.py`、`mpc_controller.py` | Pinocchio 动力学与 OSQP MPC |
+
+## 依赖
+
+执行工程根目录的 `uv sync` 后，训练环境已经提供 NumPy、SciPy 和 MuJoCo。不同工具还需要：
+
+- 纯解析工具：Python、NumPy。
+- MuJoCo 扫描：`mujoco`、NumPy。
+- 完整 MPC：`pinocchio`、`osqp`、SciPy。
+- GUI：系统可用的 Tk/Tkinter。
+
+Pinocchio 和 OSQP 没有加入训练环境锁文件，因为它们只服务于可选 MPC 工具，且 Pinocchio 的安装方式与操作系统、Conda/Python 环境有关。
+
+## 常用命令
+
+在 `05_software/train/rc_mjlab` 下执行：
+
+```bash
+# 不启动 MuJoCo 的快速姿态表
+uv run python mujoco_sim/posture_tool.py
+uv run python mujoco_sim/rl_friendly_opt.py
+
+# 姿态扫描
+uv run python mujoco_sim/posture_optimizer.py --analyze
+uv run python mujoco_sim/static_posture_optimizer.py --quick
+
+# IK 与差速轮参数快速扫描
+uv run python mujoco_sim/ik_diff_sweep.py --quick
+
+# 完整 GUI/MPC 仿真，需要可选依赖
+uv run python mujoco_sim/run.py
+```
+
+## 参数边界
+
+这是一份历史工具快照，保留当时用于分析和调参的常量：
+
+- `config.py`、`posture_optimizer.py` 和 `static_posture_optimizer.py` 中的解析质量常量为 `12.3 kg`。
+- 当前新版 MJCF 的惯性质量合计约为 `18.0377 kg`。
+- `config.py` 的姿态表默认值为髋俯仰 `0.666`、膝关节 `-1.546`；比赛训练架构的 Rough 默认姿态为 `0.550/-1.125`。
+
+MuJoCo 直接加载模型的工具会使用 MJCF 内的质量和惯性；显式读取 `ROBOT_MASS` 的解析计算和 MPC 工具仍使用历史常量。使用输出作为新版本控制参数前，应先根据目标机械状态完成质量、惯性和默认姿态复核。
@@ -1,6 +1,5 @@
 """Robot constants and control parameters."""

-import numpy as np
 from pathlib import Path

 # Paths
@@ -10,48 +9,73 @@ MJCF_PATH = REPO_ROOT / "mjcf" / "wheelleg.xml"

 # Robot geometry
 WHEEL_RADIUS = 0.10  # m
-WHEEL_TRACK = 0.32   # m (left-right distance)
+WHEEL_TRACK = 0.32   # m, left-right distance
 ROBOT_MASS = 12.3    # kg
 MAX_TORQUE = 17.0    # Nm per joint
 MAX_JOINT_VEL = 13.0 # rad/s

-# Leg link lengths (from MJCF)
-L_THIGH = 0.25  # m
-L_CALF = 0.20   # m (to wheel center)
+# Leg link lengths from MJCF, measured to wheel center.
+L_THIGH = 0.25
+L_CALF = 0.20

 # Leg names and joint ordering
 LEG_NAMES = ("fl", "fr", "rl", "rr")
 LEG_JOINTS = ("hip_abduction_joint", "hip_pitch_joint", "knee_joint")
 WHEEL_JOINT = "wheel_joint"

-# Default standing pose (from go2w_sim2sim: [0, 0.8, -1.5])
+# Default standing pose aligned with the soft wheel-X height table.
+# height ~= 0.37m, wheel x-offset ~= 0, peak/RMS leg torque balanced.
 DEFAULT_JOINT_ANGLES = {
    "hip_abduction": 0.0,
-    "hip_pitch": 0.93,
-    "knee": -1.65,
+    "hip_pitch": 0.666,
+    "knee": -1.546,
 }

-# Actuator modes (MJCF native):
-#   Leg joints: position PD (kp=120, kd=8), ctrl = target angle
-#   Wheel joints: velocity (gain=0.5), ctrl = target velocity (rad/s)
+# Actuator modes, configured at runtime:
+#   Leg joints: position PD, ctrl = target angle
+#   Wheel joints: velocity, ctrl = target velocity in rad/s

 # Control rates
-SIM_DT = 0.002       # 500 Hz (from scene.xml)
-CTRL_DT = 0.004      # 250 Hz control loop
+SIM_DT = 0.002
+CTRL_DT = 0.02
 CTRL_DECIMATION = int(CTRL_DT / SIM_DT)

 # Wheel drive
-WHEEL_VEL_MAX = 10.0  # rad/s max wheel command
+WHEEL_VEL_MAX = 10.0

-# Body pose control gains (for height/roll/pitch compensation)
-KP_HEIGHT = 3.0    # rad/m error → joint angle correction
-KP_ROLL = 0.5      # compensation gain
-KP_PITCH = 0.5     # compensation gain
+# Body pose control gains for height/roll/pitch compensation.
+KP_HEIGHT = 3.0
+KP_ROLL = 0.5
+KP_PITCH = 0.5
+
+# Calibrated height-to-joint-angle table.
+# Constraint: avoid large wheel-center X offset from the hip/leg. This is a
+# soft support-geometry guardrail, not a strict x=0 requirement.
+# The optimizer also considers peak motor torque and RMS torque, so one hot
+# motor is not hidden by a low average across all motors.
+# Format: (height_m, hip_pitch_rad, knee_rad)
+HEIGHT_TABLE = [
+    (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),
+]

 # Gait parameters
-GAIT_FREQ = 2.5     # Hz
-GAIT_DUTY = 0.6     # stance fraction
-SWING_HEIGHT = 0.06  # m
+GAIT_FREQ = 2.5
+GAIT_DUTY = 0.6
+SWING_HEIGHT = 0.06

 # Trot phase offsets: FL/RR in phase, FR/RL in phase
 PHASE_OFFSETS = {"fl": 0.0, "fr": 0.5, "rl": 0.5, "rr": 0.0}
@@ -14,10 +14,13 @@ from dynamics import Dynamics
 from mpc_controller import MPCController
 from config import (
    LEG_NAMES, DEFAULT_JOINT_ANGLES, WHEEL_RADIUS, WHEEL_TRACK,
-    WHEEL_VEL_MAX, KP_ROLL, KP_PITCH,
+    WHEEL_VEL_MAX, KP_ROLL, KP_PITCH, HEIGHT_TABLE,
    GAIT_FREQ, GAIT_DUTY, SWING_HEIGHT, PHASE_OFFSETS,
 )

+RL_ROUGH_Q = np.array([0.0, 0.550, -1.125], dtype=float)
+LEG_STATE_IDX = np.array([0, 1, 2, 4, 5, 6, 8, 9, 10, 12, 13, 14], dtype=int)
+

 class Controller:
    """Wheeled-legged robot controller."""
@@ -30,7 +33,8 @@ class Controller:
        self.vel_x = 0.0       # m/s forward
        self.vel_y = 0.0       # m/s lateral
        self.yaw_rate = 0.0    # rad/s
-        self.height = 0.33     # m desired body height
+        self.height = 0.37     # m desired body height (wheel center under hip)
+        self.wheel_posture = "table"  # "table" follows height slider; "rl" matches src/robot default

        # Mode: "wheel", "trot", or "mpc"
        self.mode = "wheel"
@@ -49,6 +53,30 @@ class Controller:
        self._smooth_vx = 0.0
        self._smooth_vy = 0.0
        self._smooth_yaw = 0.0
+        self._last_leg_targets = np.tile(
+            np.array(
+                [
+                    DEFAULT_JOINT_ANGLES["hip_abduction"],
+                    DEFAULT_JOINT_ANGLES["hip_pitch"],
+                    DEFAULT_JOINT_ANGLES["knee"],
+                ],
+                dtype=float,
+            ),
+            4,
+        )
+
+        # Wheel-mode sensor feedback.
+        self.yaw_rate_kp = 0.45
+        self.roll_comp_gain = KP_ROLL
+        self.pitch_comp_gain = KP_PITCH
+        self.encoder_posture_kp = 0.12
+        self.encoder_posture_max = 0.025
+        self.encoder_guard_start = 0.28
+        self.encoder_guard_stop = 0.65
+        self.imu_guard_start = np.deg2rad(12.0)
+        self.imu_guard_stop = np.deg2rad(28.0)
+        self.yaw_wheel_gain = 1.0
+        self.max_yaw_wheel_speed = 4.0

        # Default leg angles
        self._default_q = np.array([
@@ -111,7 +139,7 @@ class Controller:
            side = 1.0 if leg[1] == "l" else -1.0
            leg_targets[i*3] = side * 0.3     # fl/rl: +0.3, fr/rr: -0.3
            leg_targets[i*3+1] = 1.5          # hip pitch
-            leg_targets[i*3+2] = -2.65        # knee fully folded
+            leg_targets[i*3+2] = -2.65        # knee hard limit from MJCF
        return leg_targets, np.zeros(4)

    # ─────────────────────────────────────────────────────────────────────
@@ -125,35 +153,84 @@ class Controller:
        pure lateral motion). Uses hip_abduction lean for small lateral force.
        For significant lateral motion, use trot mode.
        """
-        wheel_targets = self._differential_drive(self._smooth_vx, self._smooth_yaw)
        leg_targets = self._posture_control(state)
+        safe_vx, safe_yaw = self._wheel_velocity_envelope(self._smooth_vx, self._smooth_yaw)
+        yaw_feedback = safe_yaw + self.yaw_rate_kp * (safe_yaw - float(state.ang_vel[2]))
+        wheel_targets = self._differential_drive(safe_vx, yaw_feedback)
+        wheel_targets *= self._sensor_command_scale(state, leg_targets)
+        self._last_leg_targets = leg_targets.copy()
        return leg_targets, wheel_targets

    def _posture_control(self, state: RobotState) -> np.ndarray:
-        """Leg joint targets: table-interpolated height control."""
+        """Leg joint targets from the soft wheel-X height table."""
        leg_targets = np.zeros(12)

-        # Calibrated height→angle lookup (measured from simulation)
-        _H = [0.157, 0.248, 0.311, 0.366, 0.411, 0.448]
-        _HIP = [1.5, 1.2, 1.0, 0.8, 0.6, 0.4]
-        _KNEE = [-2.5, -2.1, -1.8, -1.5, -1.2, -0.9]
+        # Calibrated height→angle lookup (minimizes motor torque at each height)
+        _H  = [r[0] for r in HEIGHT_TABLE]
+        _HIP = [r[1] for r in HEIGHT_TABLE]
+        _KNEE = [r[2] for r in HEIGHT_TABLE]

        h_clamp = np.clip(self.height, _H[0], _H[-1])
        q_hip_base = float(np.interp(h_clamp, _H, _HIP))
        q_knee_base = float(np.interp(h_clamp, _H, _KNEE))
+        q_ab_base = 0.0
+        if self.wheel_posture == "rl":
+            q_ab_base, q_hip_base, q_knee_base = RL_ROUGH_Q

-        roll_corr = -KP_ROLL * state.rpy[0]
-        pitch_corr = -KP_PITCH * state.rpy[1]
+        roll_corr = -self.roll_comp_gain * state.rpy[0]
+        pitch_corr = -self.pitch_comp_gain * state.rpy[1]
        lateral_lean = 0.3 * self.vel_y

        for i, leg in enumerate(LEG_NAMES):
            side = 1.0 if leg[1] == "l" else -1.0
-            leg_targets[i*3] = np.clip(side * roll_corr + lateral_lean, -0.5, 0.5)
+            leg_targets[i*3] = np.clip(q_ab_base + side * roll_corr + lateral_lean, -0.5, 0.5)
            leg_targets[i*3+1] = np.clip(q_hip_base + pitch_corr, -1.0, 2.5)
-            leg_targets[i*3+2] = np.clip(q_knee_base, -2.6, -0.3)
+            leg_targets[i*3+2] = np.clip(q_knee_base, -2.65, -0.3)

+        encoder_err = self._last_leg_targets - state.joint_pos[LEG_STATE_IDX]
+        leg_targets += np.clip(
+            self.encoder_posture_kp * encoder_err,
+            -self.encoder_posture_max,
+            self.encoder_posture_max,
+        )
+        leg_targets[0::3] = np.clip(leg_targets[0::3], -0.5, 0.5)
+        leg_targets[1::3] = np.clip(leg_targets[1::3], -1.0, 2.5)
+        leg_targets[2::3] = np.clip(leg_targets[2::3], -2.65, -0.3)
        return leg_targets

+    def _wheel_velocity_envelope(self, vel_x: float, yaw_rate: float) -> tuple[float, float]:
+        """Limit x/yaw combinations that are unsafe for the RL posture."""
+        ax = abs(vel_x)
+        if ax >= 0.8:
+            yaw_lim = 0.35
+        elif ax >= 0.5:
+            yaw_lim = 0.55
+        elif ax >= 0.25:
+            yaw_lim = 0.75
+        else:
+            yaw_lim = 1.0
+        return float(vel_x), float(np.clip(yaw_rate, -yaw_lim, yaw_lim))
+
+    def _sensor_command_scale(self, state: RobotState, leg_targets: np.ndarray) -> float:
+        """Back off wheels when IMU or encoder feedback says posture is degrading."""
+        leg_error = float(np.max(np.abs(state.joint_pos[LEG_STATE_IDX] - leg_targets)))
+        tilt = float(np.hypot(state.rpy[0], state.rpy[1]))
+        scale = 1.0
+
+        if leg_error >= self.encoder_guard_stop:
+            scale = 0.0
+        elif leg_error > self.encoder_guard_start:
+            span = max(1e-6, self.encoder_guard_stop - self.encoder_guard_start)
+            scale *= 1.0 - (leg_error - self.encoder_guard_start) / span
+
+        if tilt >= self.imu_guard_stop:
+            scale = 0.0
+        elif tilt > self.imu_guard_start:
+            span = max(1e-6, self.imu_guard_stop - self.imu_guard_start)
+            scale *= 1.0 - (tilt - self.imu_guard_start) / span
+
+        return float(np.clip(scale, 0.0, 1.0))
+
    # ─────────────────────────────────────────────────────────────────────
    # TROT MODE
    # ─────────────────────────────────────────────────────────────────────
@@ -199,9 +276,9 @@ class Controller:

    def _stance_leg_target(self, state: RobotState, leg: str) -> np.ndarray:
        """Stance leg: table-interpolated height + attitude compensation."""
-        _H = [0.157, 0.248, 0.311, 0.366, 0.411, 0.448]
-        _HIP = [1.5, 1.2, 1.0, 0.8, 0.6, 0.4]
-        _KNEE = [-2.5, -2.1, -1.8, -1.5, -1.2, -0.9]
+        _H  = [r[0] for r in HEIGHT_TABLE]
+        _HIP = [r[1] for r in HEIGHT_TABLE]
+        _KNEE = [r[2] for r in HEIGHT_TABLE]

        h_clamp = np.clip(self.height, _H[0], _H[-1])
        q_hip = float(np.interp(h_clamp, _H, _HIP))
@@ -215,7 +292,7 @@ class Controller:
        return np.array([
            np.clip(side * roll_corr + lateral_lean, -0.5, 0.5),
            np.clip(q_hip + pitch_corr, -1.0, 2.5),
-            np.clip(q_knee, -2.6, -0.3),
+            np.clip(q_knee, -2.65, -0.3),
        ])

    # ─────────────────────────────────────────────────────────────────────
@@ -224,8 +301,11 @@ class Controller:

    def _differential_drive(self, vel_x: float, yaw_rate: float) -> np.ndarray:
        """4 wheel velocities from body commands."""
-        vel_left = (vel_x - 0.5 * WHEEL_TRACK * yaw_rate) / WHEEL_RADIUS
-        vel_right = (vel_x + 0.5 * WHEEL_TRACK * yaw_rate) / WHEEL_RADIUS
+        linear_wheel = vel_x / WHEEL_RADIUS
+        yaw_wheel = self.yaw_wheel_gain * 0.5 * WHEEL_TRACK * yaw_rate / WHEEL_RADIUS
+        yaw_wheel = float(np.clip(yaw_wheel, -self.max_yaw_wheel_speed, self.max_yaw_wheel_speed))
+        vel_left = linear_wheel - yaw_wheel
+        vel_right = linear_wheel + yaw_wheel
        targets = np.zeros(4)
        for i, leg in enumerate(LEG_NAMES):
            targets[i] = vel_left if leg[1] == "l" else vel_right
@@ -40,7 +40,7 @@ class GUI:
        self._slider(cf, "Vel X", self.vel_x_var, -1.5, 1.5)
        self._slider(cf, "Vel Y*", self.vel_y_var, -0.5, 0.5)
        self._slider(cf, "Yaw", self.yaw_var, -2.0, 2.0)
-        self._slider(cf, "Height", self.height_var, 0.16, 0.45)
+        self._slider(cf, "Height", self.height_var, 0.17, 0.45)

        ttk.Label(cf, text="* Vel Y: trot mode only (diff-drive can't sidestep)",
                  font=("", 8)).pack(anchor="w", padx=8)
@@ -0,0 +1,389 @@
+"""Sweep wheel-mode IK postures for differential-drive tracking.
+
+The sweep is intentionally small and reproducible:
+  1. Generate ab=0 leg postures in the requested height range.
+  2. Keep candidates with good static geometry from rl_friendly_opt.
+  3. Simulate forward, yaw, and arc commands in MuJoCo.
+  4. Rank by attitude, x-speed tracking, yaw-rate tracking, and wheel contact.
+
+Usage:
+    uv run python mujoco_sim/ik_diff_sweep.py --quick
+    uv run python mujoco_sim/ik_diff_sweep.py --height-min 0.15 --height-max 0.42
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+
+import mujoco
+import numpy as np
+
+THIS_DIR = Path(__file__).resolve().parent
+if str(THIS_DIR) not in sys.path:
+    sys.path.insert(0, str(THIS_DIR))
+
+from config import LEG_JOINTS, LEG_NAMES, SCENE_XML, WHEEL_JOINT, WHEEL_RADIUS  # noqa: E402
+from rl_friendly_opt import get_all, rl_cost  # noqa: E402
+from robot import Robot  # noqa: E402
+
+
+@dataclass(frozen=True)
+class Candidate:
+    height: float
+    ab: float
+    hip: float
+    knee: float
+    static_cost: float
+    r_hip_x: float
+    cond: float
+    max_tau: float
+
+
+@dataclass(frozen=True)
+class Trial:
+    name: str
+    vx: float
+    yaw_rate: float
+    duration: float
+
+
+def wrap_pi(x: float) -> float:
+    return (x + math.pi) % (2.0 * math.pi) - math.pi
+
+
+def body_track(model: mujoco.MjModel, data: mujoco.MjData) -> float:
+    wheel_bids = []
+    for leg in LEG_NAMES:
+        bid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, f"{leg}_wheel_Link")
+        wheel_bids.append(bid)
+    y = data.xipos[wheel_bids, 1]
+    return float(np.mean(y[[0, 2]]) - np.mean(y[[1, 3]]))
+
+
+def build_maps(robot: Robot) -> tuple[dict[str, int], dict[str, int], dict[str, int]]:
+    act: dict[str, int] = {}
+    qadr: dict[str, int] = {}
+    vadr: dict[str, int] = {}
+    for leg in LEG_NAMES:
+        for jt in (*LEG_JOINTS, WHEEL_JOINT):
+            name = f"{leg}_{jt}"
+            act[name] = mujoco.mj_name2id(robot.model, mujoco.mjtObj.mjOBJ_ACTUATOR, name)
+            jid = mujoco.mj_name2id(robot.model, mujoco.mjtObj.mjOBJ_JOINT, name)
+            qadr[name] = robot.model.jnt_qposadr[jid]
+            vadr[name] = robot.model.jnt_dofadr[jid]
+    return act, qadr, vadr
+
+
+def set_posture(robot: Robot, cand: Candidate, act: dict[str, int], qadr: dict[str, int]) -> None:
+    mujoco.mj_resetData(robot.model, robot.data)
+    for leg in LEG_NAMES:
+        side_ab = cand.ab if leg[1] == "l" else -cand.ab
+        for jt, val in zip(LEG_JOINTS, (side_ab, cand.hip, cand.knee)):
+            name = f"{leg}_{jt}"
+            robot.data.qpos[qadr[name]] = val
+            robot.data.ctrl[act[name]] = val
+        robot.data.ctrl[act[f"{leg}_{WHEEL_JOINT}"]] = 0.0
+    robot.data.qpos[:3] = [0.0, 0.0, max(0.25, cand.height + 0.08)]
+    robot.data.qpos[3:7] = [1.0, 0.0, 0.0, 0.0]
+    robot.data.qvel[:] = 0.0
+    mujoco.mj_forward(robot.model, robot.data)
+
+
+def wheel_targets(
+    vx: float,
+    yaw_rate: float,
+    track: float,
+    max_wheel: float,
+    yaw_gain: float,
+    wheel_model: str,
+    linear_gain: float,
+    direct_yaw_gain: float,
+    wheel_signs: np.ndarray,
+) -> np.ndarray:
+    if wheel_model == "direct":
+        left = linear_gain * vx - direct_yaw_gain * yaw_rate
+        right = linear_gain * vx + direct_yaw_gain * yaw_rate
+    else:
+        left = (vx - yaw_gain * 0.5 * track * yaw_rate) / WHEEL_RADIUS
+        right = (vx + yaw_gain * 0.5 * track * yaw_rate) / WHEEL_RADIUS
+    raw = np.array([left, right, left, right], dtype=float)
+    return np.clip(raw * wheel_signs, -max_wheel, max_wheel)
+
+
+def run_trial(robot: Robot, cand: Candidate, trial: Trial, args: argparse.Namespace) -> dict:
+    act, qadr, vadr = build_maps(robot)
+    set_posture(robot, cand, act, qadr)
+
+    ctrl_dt = args.control_dt
+    sim_dt = robot.model.opt.timestep
+    steps_per_ctrl = max(1, int(round(ctrl_dt / sim_dt)))
+    track = body_track(robot.model, robot.data) if args.track_source == "model" else args.track_width
+    wheel_signs = np.array(args.wheel_signs, dtype=float)
+    wheel_cmd = wheel_targets(
+        trial.vx,
+        trial.yaw_rate,
+        track,
+        args.max_wheel_speed,
+        args.yaw_gain,
+        args.wheel_model,
+        args.linear_gain,
+        args.direct_yaw_gain,
+        wheel_signs,
+    )
+
+    for _ in range(int(round(args.settle / sim_dt))):
+        for leg in LEG_NAMES:
+            vals = (cand.ab if leg[1] == "l" else -cand.ab, cand.hip, cand.knee)
+            for jt, val in zip(LEG_JOINTS, vals):
+                robot.data.ctrl[act[f"{leg}_{jt}"]] = val
+            robot.data.ctrl[act[f"{leg}_{WHEEL_JOINT}"]] = 0.0
+        robot.step()
+
+    state0 = robot.get_state()
+    yaw0 = float(state0.rpy[2])
+    x0 = float(state0.pos[0])
+
+    max_roll = 0.0
+    max_pitch = 0.0
+    max_tilt = 0.0
+    max_wheel_air = -1e9
+    wheel_err_sum = 0.0
+    samples = 0
+    max_leg_err = 0.0
+    body_vx_sum = 0.0
+    yaw_unwrapped = 0.0
+    last_yaw = yaw0
+    leg_target = np.array([cand.ab, cand.hip, cand.knee] * 4, dtype=float)
+
+    total_steps = int(round(trial.duration / sim_dt))
+    cmd = np.zeros(4, dtype=float)
+    max_delta = args.wheel_accel_limit * ctrl_dt
+    wheel_body_ids = [
+        mujoco.mj_name2id(robot.model, mujoco.mjtObj.mjOBJ_BODY, f"{leg}_wheel_Link")
+        for leg in LEG_NAMES
+    ]
+    wheel_vadr = np.array([vadr[f"{leg}_{WHEEL_JOINT}"] for leg in LEG_NAMES], dtype=int)
+    leg_qadr = np.array([qadr[f"{leg}_{jt}"] for leg in LEG_NAMES for jt in LEG_JOINTS], dtype=int)
+
+    for step in range(total_steps):
+        if step % steps_per_ctrl == 0:
+            cmd = cmd + np.clip(wheel_cmd - cmd, -max_delta, max_delta)
+            for leg in LEG_NAMES:
+                vals = (cand.ab if leg[1] == "l" else -cand.ab, cand.hip, cand.knee)
+                for jt, val in zip(LEG_JOINTS, vals):
+                    robot.data.ctrl[act[f"{leg}_{jt}"]] = val
+            for i, leg in enumerate(LEG_NAMES):
+                robot.data.ctrl[act[f"{leg}_{WHEEL_JOINT}"]] = cmd[i]
+
+        robot.step()
+
+        if step % steps_per_ctrl == 0:
+            st = robot.get_state()
+            roll, pitch = float(st.rpy[0]), float(st.rpy[1])
+            yaw_now = float(st.rpy[2])
+            yaw_unwrapped += wrap_pi(yaw_now - last_yaw)
+            last_yaw = yaw_now
+            body_vx_sum += float(st.rot[:, 0].dot(st.lin_vel))
+            max_roll = max(max_roll, abs(roll))
+            max_pitch = max(max_pitch, abs(pitch))
+            max_tilt = max(max_tilt, math.hypot(roll, pitch))
+            wheel_air = robot.data.xipos[wheel_body_ids, 2] - WHEEL_RADIUS
+            max_wheel_air = max(max_wheel_air, float(np.max(wheel_air)))
+            wheel_err_sum += float(np.mean(np.abs(robot.data.qvel[wheel_vadr] - cmd)))
+            max_leg_err = max(max_leg_err, float(np.max(np.abs(robot.data.qpos[leg_qadr] - leg_target))))
+            samples += 1
+
+    st = robot.get_state()
+    elapsed = max(1e-6, float(st.time - state0.time))
+    world_x_rate = (float(st.pos[0]) - x0) / elapsed
+    x_rate = body_vx_sum / max(1, samples)
+    yaw_rate = yaw_unwrapped / max(1e-6, samples * steps_per_ctrl * sim_dt)
+    x_err = abs(x_rate - trial.vx)
+    yaw_err = abs(yaw_rate - trial.yaw_rate)
+    return {
+        "trial": trial.name,
+        "x_rate": x_rate,
+        "world_x_rate": world_x_rate,
+        "yaw_rate": yaw_rate,
+        "x_err": x_err,
+        "yaw_err": yaw_err,
+        "max_roll_deg": math.degrees(max_roll),
+        "max_pitch_deg": math.degrees(max_pitch),
+        "max_tilt_deg": math.degrees(max_tilt),
+        "max_wheel_air_m": max_wheel_air,
+        "mean_wheel_err": wheel_err_sum / max(1, samples),
+        "max_leg_err": max_leg_err,
+        "track": track,
+        "wheel_cmd": [float(x) for x in wheel_cmd],
+    }
+
+
+def generate_candidates(args: argparse.Namespace) -> list[Candidate]:
+    if args.fixed_hip is not None or args.fixed_knee is not None:
+        if args.fixed_hip is None or args.fixed_knee is None:
+            raise SystemExit("--fixed-hip and --fixed-knee must be provided together")
+        ab = float(args.fixed_ab)
+        hip = float(args.fixed_hip)
+        knee = float(args.fixed_knee)
+        r = get_all(ab, hip, knee)
+        return [
+            Candidate(
+                height=float(r["z"]),
+                ab=ab,
+                hip=hip,
+                knee=knee,
+                static_cost=float(rl_cost(r)),
+                r_hip_x=float(r["r_hip_x_mag"]),
+                cond=float(r["cond"]),
+                max_tau=float(r["max_tau"]),
+            )
+        ]
+
+    cands: list[Candidate] = []
+    h_targets = np.arange(args.height_min, args.height_max + 0.5 * args.height_step, args.height_step)
+    ab_values = [0.0] if args.ab_max <= 1e-9 else np.arange(0.0, args.ab_max + 1e-9, args.ab_step)
+    hip_values = np.arange(args.hip_min, args.hip_max + 0.5 * args.hip_step, args.hip_step)
+    knee_values = np.arange(args.knee_min, args.knee_max + 0.5 * args.knee_step, args.knee_step)
+    for ht in h_targets:
+        bucket: list[Candidate] = []
+        for ab in ab_values:
+            for hip in hip_values:
+                for knee in knee_values:
+                    r = get_all(float(ab), float(hip), float(knee))
+                    if float(r["z"]) < args.height_min or float(r["z"]) > args.height_max:
+                        continue
+                    if abs(float(r["z"]) - float(ht)) > args.height_tol:
+                        continue
+                    if r["wz"] >= r["kz"]:
+                        continue
+                    if abs(r["r_hip_x_mag"]) > args.max_wheel_x:
+                        continue
+                    cost = float(rl_cost(r))
+                    bucket.append(
+                        Candidate(
+                            height=float(r["z"]),
+                            ab=float(ab),
+                            hip=float(hip),
+                            knee=float(knee),
+                            static_cost=cost,
+                            r_hip_x=float(r["r_hip_x_mag"]),
+                            cond=float(r["cond"]),
+                            max_tau=float(r["max_tau"]),
+                        )
+                    )
+        bucket.sort(key=lambda c: c.static_cost)
+        cands.extend(bucket[: args.top_per_height])
+    return cands
+
+
+def score_result(cand: Candidate, trials: list[dict]) -> float:
+    score = 0.08 * cand.static_cost
+    for t in trials:
+        score += 8.0 * t["x_err"]
+        score += 10.0 * t["yaw_err"]
+        score += 0.08 * t["max_tilt_deg"]
+        score += 0.03 * max(0.0, t["max_pitch_deg"] - 8.0) ** 2
+        score += 20.0 * max(0.0, t["max_wheel_air_m"] - 0.015)
+        score += 1.5 * t["mean_wheel_err"]
+        score += 2.0 * max(0.0, t["max_leg_err"] - 0.35)
+    return float(score)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--height-min", type=float, default=0.15)
+    parser.add_argument("--height-max", type=float, default=0.42)
+    parser.add_argument("--height-step", type=float, default=0.01)
+    parser.add_argument("--height-tol", type=float, default=0.004)
+    parser.add_argument("--top-per-height", type=int, default=1)
+    parser.add_argument("--ab-max", type=float, default=0.0)
+    parser.add_argument("--ab-step", type=float, default=0.04)
+    parser.add_argument("--hip-min", type=float, default=0.25)
+    parser.add_argument("--hip-max", type=float, default=1.05)
+    parser.add_argument("--hip-step", type=float, default=0.025)
+    parser.add_argument("--knee-min", type=float, default=-2.65)
+    parser.add_argument("--knee-max", type=float, default=-0.85)
+    parser.add_argument("--knee-step", type=float, default=0.025)
+    parser.add_argument("--max-wheel-x", type=float, default=0.09)
+    parser.add_argument("--fixed-ab", type=float, default=0.0)
+    parser.add_argument("--fixed-hip", type=float, default=None)
+    parser.add_argument("--fixed-knee", type=float, default=None)
+    parser.add_argument("--duration", type=float, default=4.0)
+    parser.add_argument("--settle", type=float, default=1.5)
+    parser.add_argument("--control-dt", type=float, default=0.02)
+    parser.add_argument("--vx", type=float, default=0.6)
+    parser.add_argument("--yaw", type=float, default=0.3)
+    parser.add_argument("--arc-yaw", type=float, default=0.15)
+    parser.add_argument("--yaw-gain", type=float, default=1.0)
+    parser.add_argument("--wheel-model", choices=("diff", "direct"), default="diff")
+    parser.add_argument("--linear-gain", type=float, default=12.5)
+    parser.add_argument("--direct-yaw-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("--track-source", choices=("model", "fixed"), default="model")
+    parser.add_argument("--track-width", type=float, default=0.394)
+    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.",
+    )
+    parser.add_argument("--quick", action="store_true")
+    parser.add_argument("--json", type=Path, default=None)
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    if args.quick:
+        args.height_step = 0.02
+        args.hip_step = 0.05
+        args.knee_step = 0.05
+
+    candidates = generate_candidates(args)
+    if not candidates:
+        raise SystemExit("No candidates found")
+
+    trials = [
+        Trial("forward", args.vx, 0.0, args.duration),
+        Trial("yaw", 0.0, args.yaw, args.duration),
+        Trial("arc", args.vx, args.arc_yaw, args.duration),
+    ]
+    robot = Robot(SCENE_XML)
+    rows = []
+    for i, cand in enumerate(candidates, 1):
+        trial_rows = [run_trial(robot, cand, t, args) for t in trials]
+        rows.append({"candidate": cand.__dict__, "trials": trial_rows, "score": score_result(cand, trial_rows)})
+        if i % 10 == 0:
+            print(f"tested {i}/{len(candidates)}")
+
+    rows.sort(key=lambda r: r["score"])
+    if args.json:
+        args.json.write_text(json.dumps(rows, indent=2), encoding="utf-8")
+
+    print("Top IK postures for differential drive tracking")
+    print("rank score height ab hip knee static xhip cond tau | fwd_x yaw_wz arc_x arc_wz max_tilt max_pitch")
+    for rank, row in enumerate(rows[:10], 1):
+        c = row["candidate"]
+        by = {t["trial"]: t for t in row["trials"]}
+        max_tilt = max(t["max_tilt_deg"] for t in row["trials"])
+        max_pitch = max(t["max_pitch_deg"] for t in row["trials"])
+        print(
+            f"{rank:>2} {row['score']:>7.2f} {c['height']:.3f} {c['ab']:.2f} {c['hip']:.3f} {c['knee']:.3f} "
+            f"{c['static_cost']:.1f} {c['r_hip_x']:.3f} {c['cond']:.2f} {c['max_tau']:.2f} | "
+            f"{by['forward']['x_rate']:.3f} {by['yaw']['yaw_rate']:.3f} {by['arc']['x_rate']:.3f} {by['arc']['yaw_rate']:.3f} "
+            f"{max_tilt:.1f} {max_pitch:.1f}"
+        )
+
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
@@ -5,7 +5,7 @@ Integrates: gait scheduler + reference trajectory + ConvexMPC solver +

 Architecture (following go2-convex-mpc):
  - MPC runs at ~50 Hz (every MPC_DECIMATION control steps)
-  - Swing/stance leg controller runs at control rate (250 Hz)
+  - Swing/stance leg controller runs at control rate (50 Hz)
  - Wheel drive: stance legs use differential drive, swing legs coast
 """

@@ -20,7 +20,7 @@ from config import (
 )

 # MPC update decimation (relative to control loop)
-MPC_DECIMATION = max(1, int(MPC_DT / CTRL_DT))  # ~5 steps at 250Hz
+MPC_DECIMATION = max(1, int(MPC_DT / CTRL_DT))  # 1 step at 50Hz


 class MPCController:
@@ -0,0 +1,367 @@
+"""Posture optimizer for wheeled-leg standing/crawl height table.
+
+The table is not a pure "minimum average torque" table. For crawl and low-bar
+traversal, the wheel center should not be far from the hip/leg in sagittal X,
+otherwise the robot is no longer really using the wheel as the support/drive
+point. This is a soft guardrail, not a strict x=0 constraint. The score combines:
+
+1. wheel center X offset from hip
+2. peak single-motor holding torque
+3. RMS torque, used as a proxy for I^2R heating
+
+Usage:
+    python posture_optimizer.py              # MuJoCo sweep
+    python posture_optimizer.py --quick      # coarse MuJoCo sweep
+    python posture_optimizer.py --analyze    # analytical-only sweep
+"""
+
+import argparse
+import math
+import time
+from pathlib import Path
+
+import numpy as np
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+
+try:
+    import mujoco
+except ImportError:
+    mujoco = None
+
+SCENE_XML = REPO_ROOT / "mjcf" / "scene.xml"
+WHEEL_RADIUS = 0.10
+HIP_Z_OFFSET = 0.054
+L1, L2 = 0.25, 0.20
+ROBOT_MASS = 12.3
+G = 9.81
+F_PER_LEG = ROBOT_MASS * G / 4.0
+MAX_TORQUE = 17.0
+
+LEG_NAMES = ("fl", "fr", "rl", "rr")
+LEG_JOINTS = ("hip_abduction_joint", "hip_pitch_joint", "knee_joint")
+
+# 0.15m is not a good default table target for ab=0.0. 0.17m is the practical
+# default crawl height, while lower crawl can be evaluated with abduction.
+KNEE_MIN = -2.65
+HEIGHT_MIN = 0.17
+HEIGHT_MAX = 0.46
+SOFT_WHEEL_X_OFFSET = 0.05
+HARD_WHEEL_X_OFFSET = 0.08
+
+_ACTUATOR_NAMES = [f"{leg}_{jt}" for leg in LEG_NAMES for jt in LEG_JOINTS]
+
+
+def compute_fk(hip, knee):
+    """Return wheel-center x offset and base height from (hip_pitch, knee)."""
+    x = L1 * math.sin(hip) + L2 * math.sin(hip + knee)
+    z = L1 * math.cos(hip) + L2 * math.cos(hip + knee)
+    base_height = WHEEL_RADIUS + z - HIP_Z_OFFSET
+    return x, base_height
+
+
+def posture_cost(x_foot, torques):
+    """Score one posture by support geometry, peak torque, and RMS torque."""
+    tau = np.asarray(torques, dtype=float)
+    peak_torque = float(np.max(np.abs(tau)))
+    rms_torque = float(np.sqrt(np.mean(np.square(tau))))
+    mean_i2r = float(np.mean(np.square(tau)))
+    x_penalty = max(0.0, abs(x_foot) - SOFT_WHEEL_X_OFFSET)
+    cost = (
+        0.5 * (abs(x_foot) / SOFT_WHEEL_X_OFFSET) ** 2
+        + 8.0 * (x_penalty / max(1e-6, HARD_WHEEL_X_OFFSET - SOFT_WHEEL_X_OFFSET)) ** 2
+        + 3.0 * (peak_torque / MAX_TORQUE) ** 2
+        + (rms_torque / MAX_TORQUE) ** 2
+    )
+    return cost, peak_torque, rms_torque, mean_i2r
+
+
+def analyze_analytical():
+    """Analytical sweep using static GRF moments."""
+    hip_range = np.arange(0.0, 1.6, 0.002)
+    knee_range = np.arange(KNEE_MIN, -0.4, 0.002)
+
+    results = []
+    for hip in hip_range:
+        for knee in knee_range:
+            x_foot, height = compute_fk(hip, knee)
+            if height < HEIGHT_MIN or height > HEIGHT_MAX:
+                continue
+
+            tau_hip = F_PER_LEG * x_foot
+            x_knee_to_foot = L2 * math.sin(hip + knee)
+            tau_knee = F_PER_LEG * x_knee_to_foot
+            tau_abduction = 0.0
+            cost, peak, rms, mean_i2r = posture_cost(
+                x_foot, (tau_abduction, tau_hip, tau_knee))
+
+            results.append({
+                "hip": float(hip),
+                "knee": float(knee),
+                "height": float(height),
+                "x_foot": float(x_foot),
+                "tau_abd": tau_abduction,
+                "tau_hip": float(tau_hip),
+                "tau_knee": float(tau_knee),
+                "tau_peak": peak,
+                "tau_rms": rms,
+                "mean_i2r": mean_i2r,
+                "cost": cost,
+            })
+
+    return results
+
+
+def run_mujoco_sweep(quick=False):
+    """MuJoCo sweep measuring actual actuator forces at steady state."""
+    if mujoco is None:
+        raise RuntimeError("mujoco is not installed; use --analyze for analytical mode")
+
+    model = mujoco.MjModel.from_xml_path(str(SCENE_XML))
+    data = mujoco.MjData(model)
+
+    act_ids = {}
+    for name in _ACTUATOR_NAMES:
+        act_ids[name] = mujoco.mj_name2id(
+            model, mujoco.mjtObj.mjOBJ_ACTUATOR, name)
+
+    for i in range(model.nu):
+        name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_ACTUATOR, i)
+        model.actuator_biastype[i] = 1
+        model.actuator_gaintype[i] = 0
+        model.actuator_forcelimited[i] = 0
+        if "wheel" not in name:
+            model.actuator_gainprm[i, 0] = 60.0
+            model.actuator_biasprm[i, 0] = 0.0
+            model.actuator_biasprm[i, 1] = -60.0
+            model.actuator_biasprm[i, 2] = -3.0
+            model.actuator_ctrlrange[i] = [-3.14, 3.14]
+        else:
+            model.actuator_gainprm[i, 0] = 2.0
+            model.actuator_biasprm[i, 0] = 0.0
+            model.actuator_biasprm[i, 1] = 0.0
+            model.actuator_biasprm[i, 2] = -2.0
+            model.actuator_ctrlrange[i] = [-20.0, 20.0]
+
+    jnt_ids = {}
+    for name in _ACTUATOR_NAMES:
+        jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, name)
+        jnt_ids[name] = model.jnt_qposadr[jid]
+
+    if quick:
+        hip_range = np.arange(0.3, 1.5, 0.10)
+        knee_range = np.arange(KNEE_MIN, -0.6, 0.10)
+    else:
+        hip_range = np.arange(0.0, 1.6, 0.04)
+        knee_range = np.arange(KNEE_MIN, -0.4, 0.04)
+
+    results = []
+    total = 0
+    valid = 0
+
+    for hip in hip_range:
+        for knee in knee_range:
+            total += 1
+            x_foot, height = compute_fk(hip, knee)
+            if height < HEIGHT_MIN or height > HEIGHT_MAX:
+                continue
+
+            mujoco.mj_resetData(model, data)
+
+            for leg in LEG_NAMES:
+                for jt, val in zip(LEG_JOINTS, (0.0, hip, knee)):
+                    name = f"{leg}_{jt}"
+                    data.qpos[jnt_ids[name]] = val
+
+            data.qpos[2] = height
+            data.qpos[3] = 1.0
+            data.qpos[4:7] = 0.0
+
+            for leg in LEG_NAMES:
+                for jt, val in zip(LEG_JOINTS, (0.0, hip, knee)):
+                    name = f"{leg}_{jt}"
+                    data.ctrl[act_ids[name]] = val
+                name_w = f"{leg}_wheel_joint"
+                if name_w in act_ids:
+                    data.ctrl[act_ids[name_w]] = 0.0
+
+            mujoco.mj_forward(model, data)
+
+            for _ in range(500):
+                mujoco.mj_step(model, data)
+
+            roll, pitch, _ = _get_rpy(data, model)
+            if abs(roll) > 0.8 or abs(pitch) > 0.8:
+                continue
+
+            torque_buf = []
+            for _ in range(100):
+                mujoco.mj_step(model, data)
+                torque_buf.append([data.actuator_force[act_ids[name]]
+                                   for name in _ACTUATOR_NAMES])
+            tau_avg = np.array(torque_buf).mean(axis=0)
+
+            tau_hip_val = tau_avg[1]
+            tau_knee_val = tau_avg[2]
+            tau_abd_val = tau_avg[0]
+            cost, peak, rms, mean_i2r = posture_cost(x_foot, tau_avg)
+
+            q_hip_actual = float(data.qpos[jnt_ids["fl_hip_pitch_joint"]])
+            q_knee_actual = float(data.qpos[jnt_ids["fl_knee_joint"]])
+
+            valid += 1
+            results.append({
+                "hip": float(hip),
+                "knee": float(knee),
+                "height": float(f"{height:.4f}"),
+                "x_foot": float(f"{x_foot:.4f}"),
+                "tau_abd": float(f"{tau_abd_val:.4f}"),
+                "tau_hip": float(f"{tau_hip_val:.4f}"),
+                "tau_knee": float(f"{tau_knee_val:.4f}"),
+                "tau_peak": float(f"{peak:.4f}"),
+                "tau_rms": float(f"{rms:.4f}"),
+                "mean_i2r": float(f"{mean_i2r:.4f}"),
+                "cost": float(f"{cost:.4f}"),
+                "q_hip_actual": float(f"{q_hip_actual:.4f}"),
+                "q_knee_actual": float(f"{q_knee_actual:.4f}"),
+            })
+
+            if valid % 20 == 0:
+                print(f"  [{valid}/{total}] hip={hip:.2f} knee={knee:.2f} "
+                      f"h={height:.3f} x={x_foot:+.4f} "
+                      f"peak={peak:.2f} rms={rms:.2f} cost={cost:.4f}")
+
+    return results
+
+
+def _get_rpy(data, model):
+    """Extract roll and pitch from MuJoCo data."""
+    base_bid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "base_link")
+    rot = data.xmat[base_bid].reshape(3, 3)
+    roll = math.atan2(rot[2, 1], rot[2, 2])
+    pitch = math.atan2(-rot[2, 0], math.sqrt(rot[2, 1] ** 2 + rot[2, 2] ** 2))
+    return roll, pitch, 0.0
+
+
+def print_top_results(results, n=10):
+    """Print the top-N results with lowest cost."""
+    sorted_r = sorted(results, key=lambda r: r["cost"])
+
+    print(f"\n{'=' * 104}")
+    print(f"TOP {n} CONFIGURATIONS (soft wheel-X + peak/RMS torque score)")
+    print(f"{'=' * 104}")
+    print(f"{'Rank':>4} {'hip':>6} {'knee':>7} {'height':>7} {'x_foot':>8} "
+          f"{'tau_abd':>8} {'tau_hip':>8} {'tau_knee':>8} "
+          f"{'peak':>8} {'rms':>8} {'cost':>9}")
+    print(f"{'-' * 104}")
+
+    for i, r in enumerate(sorted_r[:n]):
+        print(f"{i + 1:>4} {r['hip']:>6.3f} {r['knee']:>7.3f} "
+              f"{r['height']:>7.3f} {r.get('x_foot', 0):>8.4f} "
+              f"{r.get('tau_abd', 0):>8.3f} {r['tau_hip']:>8.3f} "
+              f"{r['tau_knee']:>8.3f} {r.get('tau_peak', 0):>8.3f} "
+              f"{r.get('tau_rms', 0):>8.3f} {r['cost']:>9.4f}")
+
+    best = sorted_r[0]
+    print(f"\nBEST: hip={best['hip']:.3f} knee={best['knee']:.3f} "
+          f"z={best['height']:.3f}m x={best.get('x_foot', 0):+.4f}m "
+          f"peak={best.get('tau_peak', 0):.3f}Nm "
+          f"rms={best.get('tau_rms', 0):.3f}Nm cost={best['cost']:.4f}\n")
+
+    return sorted_r
+
+
+def compute_height_table(results):
+    """Build height-to-angle lookup with soft wheel-X support guardrail."""
+    sorted_r = sorted(
+        (r for r in results if abs(r.get("x_foot", 999.0)) <= HARD_WHEEL_X_OFFSET),
+        key=lambda r: r["height"],
+    )
+    if not sorted_r:
+        raise RuntimeError("No candidates satisfy HARD_WHEEL_X_OFFSET")
+
+    h_range = np.arange(0.17, 0.46, 0.02)
+    table_h, table_hip, table_knee = [], [], []
+
+    for h_target in h_range:
+        candidates = [(r, abs(r["height"] - h_target)) for r in sorted_r]
+        candidates.sort(key=lambda x: (x[1], x[0]["cost"]))
+        best = candidates[0][0]
+        table_h.append(best["height"])
+        table_hip.append(best["hip"])
+        table_knee.append(best["knee"])
+
+    return {
+        "height": [round(h, 3) for h in table_h],
+        "hip": [round(h, 3) for h in table_hip],
+        "knee": [round(k, 3) for k in table_knee],
+    }
+
+
+def export_calibrated_table(table):
+    """Print the new height table in copy-paste format."""
+    print(f"\n{'=' * 80}")
+    print("CALIBRATED HEIGHT TABLE (soft wheel-X support guardrail)")
+    print(f"{'=' * 80}")
+    print(f"_H = {table['height']}")
+    print(f"_HIP = {table['hip']}")
+    print(f"_KNEE = {table['knee']}")
+    print(f"{'=' * 80}\n")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Find wheeled-leg posture table")
+    parser.add_argument("--quick", action="store_true", help="Coarse MuJoCo sweep")
+    parser.add_argument("--analyze", action="store_true", help="Analytical only")
+    parser.add_argument("--mujoco", action="store_true", default=True,
+                        help="Run MuJoCo simulation when available")
+    args = parser.parse_args()
+
+    print("=" * 72)
+    print("WHEELED-LEG POSTURE OPTIMIZER")
+    print("=" * 72)
+    print(f"Robot mass: {ROBOT_MASS} kg, F_per_leg: {F_PER_LEG:.1f} N")
+    print(f"L1={L1}m, L2={L2}m, wheel_r={WHEEL_RADIUS}m")
+    print(f"Height range: [{HEIGHT_MIN}, {HEIGHT_MAX}] m")
+    print(f"Knee min hard limit: {KNEE_MIN} rad")
+    print(f"Soft wheel X offset: {SOFT_WHEEL_X_OFFSET} m")
+    print(f"Hard wheel X offset: {HARD_WHEEL_X_OFFSET} m\n")
+
+    t0 = time.time()
+    if args.analyze or mujoco is None:
+        print("[Analytical mode]")
+        results = analyze_analytical()
+    else:
+        print("[MuJoCo simulation mode]")
+        results = run_mujoco_sweep(quick=args.quick)
+
+    elapsed = time.time() - t0
+    print(f"Evaluated {len(results)} valid configurations in {elapsed:.1f}s")
+    if not results:
+        print("No valid configurations found")
+        return
+
+    best_results = print_top_results(results, n=15)
+
+    x_def, h_def = compute_fk(0.666, -1.546)
+    print(f"Current config default: hip=0.666, knee=-1.546 "
+          f"=> z={h_def:.3f}m, x={x_def:+.4f}m")
+
+    table = compute_height_table(best_results)
+    export_calibrated_table(table)
+
+    best = best_results[0]
+    print("=" * 72)
+    print("RECOMMENDED DEFAULT")
+    print("=" * 72)
+    print(f"hip_abduction: 0.0")
+    print(f"hip_pitch:     {best['hip']:.3f}")
+    print(f"knee:          {best['knee']:.3f}")
+    print(f"height:        {best['height']:.3f} m")
+    print(f"x_foot:        {best.get('x_foot', 0):+.4f} m")
+    print(f"tau_peak:      {best.get('tau_peak', 0):.3f} Nm")
+    print(f"tau_rms:       {best.get('tau_rms', 0):.3f} Nm")
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,189 @@
+"""Fast posture table helper based on wheelleg.xml link offsets.
+
+This tool is useful because it uses the real FL leg offsets from MJCF instead
+of the simplified two-link geometry used by posture_optimizer.py. It is still a
+static single-leg approximation, so use it to choose candidate crawl/standing
+poses, then verify in MuJoCo and on the robot at low speed.
+"""
+
+import math
+
+F_PER_LEG = 12.3 * 9.81 / 4.0
+KNEE_MIN = -2.65
+HIP_MIN = -2.58
+HIP_MAX = 2.58
+AB_MIN = -0.436
+AB_MAX = 0.611
+HARD_WHEEL_X_OFFSET = 0.08
+
+
+def rx(a):
+    c, s = math.cos(a), math.sin(a)
+    return ((1, 0, 0), (0, c, -s), (0, s, c))
+
+
+def ry(a):
+    c, s = math.cos(a), math.sin(a)
+    return ((c, 0, s), (0, 1, 0), (-s, 0, c))
+
+
+def mv(m, v):
+    return [
+        m[0][0] * v[0] + m[0][1] * v[1] + m[0][2] * v[2],
+        m[1][0] * v[0] + m[1][1] * v[1] + m[1][2] * v[2],
+        m[2][0] * v[0] + m[2][1] * v[1] + m[2][2] * v[2],
+    ]
+
+
+def add(a, b):
+    return [a[0] + b[0], a[1] + b[1], a[2] + b[2]]
+
+
+def get_posture(q_ab, q_hip, q_knee):
+    """Return static FK/torque metrics for one FL leg.
+
+    Returns:
+        base_z, tau_ab, tau_hip, tau_knee, i2r_total, knee_z, wheel_z, wheel_x_from_hip
+    """
+    # FL offsets from mjcf/wheelleg.xml.
+    t_ab = [0.32826, 0.066172, 0.053981]
+    t_hip = [0.06389, -0.027344, 0.00010727]
+    t_knee = [0.0, 0.1035, -0.25]
+    t_wheel = [0.0, 0.014699, -0.20011]
+
+    p = mv(ry(q_knee), t_wheel)
+    p = add(p, t_knee)
+    p = mv(ry(q_hip), p)
+    p = add(p, t_hip)
+    p = mv(rx(q_ab), p)
+    p = add(p, t_ab)
+
+    base_z = 0.10 - p[2]
+    wheel_z = p[2]
+
+    knee_pos = mv(ry(q_hip), t_knee)
+    knee_pos = mv(rx(q_ab), knee_pos)
+    knee_z = knee_pos[2] + t_ab[2]
+
+    joint_knee = mv(ry(q_hip), add(t_knee, t_hip))
+    joint_knee = mv(rx(q_ab), joint_knee)
+    joint_knee = add(t_ab, joint_knee)
+
+    joint_hip = mv(rx(q_ab), t_hip)
+    joint_hip = add(t_ab, joint_hip)
+
+    r_knee = [p[0] - joint_knee[0], p[1] - joint_knee[1], p[2] - joint_knee[2]]
+    r_hip = [p[0] - joint_hip[0], p[1] - joint_hip[1], p[2] - joint_hip[2]]
+    r_ab = [p[0] - t_ab[0], p[1] - t_ab[1], p[2] - t_ab[2]]
+
+    tau_ab = r_ab[1] * F_PER_LEG
+    tau_hip = -r_hip[0] * F_PER_LEG
+    tau_knee = -r_knee[0] * F_PER_LEG
+    i2r_total = tau_ab * tau_ab + tau_hip * tau_hip + tau_knee * tau_knee
+    wheel_x_from_hip = p[0] - joint_hip[0]
+    return base_z, tau_ab, tau_hip, tau_knee, i2r_total, knee_z, wheel_z, wheel_x_from_hip
+
+
+def score_candidate(z, z_target, tau_ab, tau_hip, tau_knee, wheel_x, x_target=0.0):
+    peak = max(abs(tau_ab), abs(tau_hip), abs(tau_knee))
+    rms = math.sqrt((tau_ab * tau_ab + tau_hip * tau_hip + tau_knee * tau_knee) / 3.0)
+    x_err = wheel_x - x_target
+    x_over = max(0.0, abs(wheel_x) - 0.05)
+    return (
+        3000.0 * (z - z_target) ** 2
+        + 2.5 * (peak / 17.0) ** 2
+        + (rms / 17.0) ** 2
+        + 0.4 * (x_err / 0.05) ** 2
+        + 6.0 * (x_over / 0.03) ** 2
+    ), peak, rms
+
+
+def find_best(z_target, ab_range=(0.0, 0.0), step=0.002, x_target=0.0, hard_wheel_x_offset=HARD_WHEEL_X_OFFSET):
+    """Find one static posture near target height without violating hard limits."""
+    best = None
+    ab0, ab1 = ab_range
+    n_ab = max(1, int(round((ab1 - ab0) / step)) + 1)
+    n_hip = int(round((1.6 - 0.0) / step)) + 1
+    n_knee = int(round((-0.4 - KNEE_MIN) / step)) + 1
+
+    for ia in range(n_ab):
+        ab = ab0 + ia * step
+        if ab < AB_MIN or ab > AB_MAX:
+            continue
+        for ih in range(n_hip):
+            hip = ih * step
+            if hip < HIP_MIN or hip > HIP_MAX:
+                continue
+            for ik in range(n_knee):
+                knee = KNEE_MIN + ik * step
+                z, ta, th, tk, i2r, kz, wz, wx = get_posture(ab, hip, knee)
+                if abs(z - z_target) > 0.0015:
+                    continue
+                if abs(wx) > hard_wheel_x_offset:
+                    continue
+                if wz >= kz:
+                    continue
+                cost, peak, rms = score_candidate(z, z_target, ta, th, tk, wx, x_target=x_target)
+                cand = (cost, ab, hip, knee, z, ta, th, tk, peak, rms, i2r, wx)
+                if best is None or cand[0] < best[0]:
+                    best = cand
+    return best
+
+
+def print_table(
+    z_targets,
+    name,
+    ab_range=(0.0, 0.0),
+    step=0.002,
+    x_target=0.0,
+    hard_wheel_x_offset=HARD_WHEEL_X_OFFSET,
+):
+    print(f"\n{'=' * 96}")
+    print(name)
+    print(f"{'=' * 96}")
+    print(f"{'z':>6} {'ab':>6} {'hip':>7} {'knee':>7} "
+          f"{'tau_ab':>8} {'tau_hip':>8} {'tau_knee':>9} "
+          f"{'peak':>8} {'rms':>8} {'x_hip':>8}")
+    print("-" * 96)
+    for zt in z_targets:
+        best = find_best(
+            zt,
+            ab_range=ab_range,
+            step=step,
+            x_target=x_target,
+            hard_wheel_x_offset=hard_wheel_x_offset,
+        )
+        if best is None:
+            print(f"{zt:>6.3f} no valid config")
+            continue
+        _, ab, hip, knee, z, ta, th, tk, peak, rms, _, wx = best
+        print(f"{z:>6.3f} {ab:>6.3f} {hip:>7.3f} {knee:>7.3f} "
+              f"{ta:>8.3f} {th:>8.3f} {tk:>9.3f} "
+              f"{peak:>8.3f} {rms:>8.3f} {wx:>8.4f}")
+
+
+def print_crawl_default():
+    best = find_best(0.17, ab_range=(0.0, 0.0), step=0.002)
+    if best is None:
+        return
+    _, ab, hip, knee, z, ta, th, tk, peak, rms, _, wx = best
+    print("\nSuggested runtime crawl_default_dof_pos:")
+    print(
+        f"[{ab:.3f}, {hip:.3f}, {knee:.3f}, "
+        f"{ab:.3f}, {hip:.3f}, {knee:.3f}, "
+        f"{ab:.3f}, {hip:.3f}, {knee:.3f}, "
+        f"{ab:.3f}, {hip:.3f}, {knee:.3f}, "
+        "0.0, 0.0, 0.0, 0.0]"
+    )
+    print(f"# z={z:.3f}, peak={peak:.3f}Nm, rms={rms:.3f}Nm, wheel_x_from_hip={wx:+.4f}m")
+
+
+if __name__ == "__main__":
+    print_table([round(0.36 + 0.01 * i, 2) for i in range(10)],
+                "STANDING candidates from MJCF FL geometry", step=0.004)
+    print_table([round(0.17 + 0.01 * i, 2) for i in range(9)],
+                "CRAWL candidates from MJCF FL geometry", step=0.002)
+    print_table([round(0.10 + 0.01 * i, 2) for i in range(6)],
+                "LOW CRAWL candidates, fixed FL abduction = +0.2",
+                ab_range=(0.2, 0.2), step=0.002, hard_wheel_x_offset=0.45)
+    print_crawl_default()
@@ -0,0 +1,200 @@
+"""RlFriendlyPostureOpt — 结合电机发热 + RL友好度约束的站/爬姿优化
+
+RL友好度约束（从实机经验总结，静力学可计算）：
+  1. 髋关节力臂 ≥ 0.08m   — 不让髋闲置（动态响应差）
+  2. 三电机不均衡 ≤ 1.3x   — 不让单电机先超载
+  3. 运动学条件数 κ ≤ 3.5 — 不让有效传动比过高（放大控制噪声）
+
+使用方法：
+  uv run python mujoco_sim/rl_friendly_opt.py          # 打印推荐
+  uv run python -c "from mujoco_sim.rl_friendly_opt import get_all; print(get_all(0, 0.8, -1.22))"
+
+基于 mjcf/wheelleg.xml 的 FL 腿运动学。
+"""
+import math
+import numpy as np
+
+# ---------------------------------------------------------------------------
+# 运动学常数（来自 wheelleg.xml FL 腿）
+# ---------------------------------------------------------------------------
+T_AB = [0.32826, 0.066172, 0.053981]
+T_HIP = [0.06389, -0.027344, 0.00010727]
+T_KNEE = [0.0, 0.1035, -0.25]
+T_WHEEL = [0.0, 0.014699, -0.20011]
+F_PER_LEG = 12.3 * 9.81 / 4.0
+KNEE_MIN, KNEE_MAX = -2.65, 2.65
+HIP_MIN, HIP_MAX = -2.58, 2.58
+
+# ---------------------------------------------------------------------------
+# 工具函数
+# ---------------------------------------------------------------------------
+def rx(a):
+    c = math.cos(a); s = math.sin(a)
+    return ((1, 0, 0), (0, c, -s), (0, s, c))
+
+def ry(a):
+    c = math.cos(a); s = math.sin(a)
+    return ((c, 0, s), (0, 1, 0), (-s, 0, c))
+
+def mv(m, v):
+    return [m[0][0]*v[0] + m[0][1]*v[1] + m[0][2]*v[2],
+            m[1][0]*v[0] + m[1][1]*v[1] + m[1][2]*v[2],
+            m[2][0]*v[0] + m[2][1]*v[1] + m[2][2]*v[2]]
+
+def add(a, b): return [a[0]+b[0], a[1]+b[1], a[2]+b[2]]
+
+# ---------------------------------------------------------------------------
+# 主函数
+# ---------------------------------------------------------------------------
+def get_all(q_ab, q_hip, q_knee):
+    """FK + 力矩 + 运动学指标。
+
+    返回 dict:
+      z, ab, hip, knee, tau_ab, tau_hip, tau_knee,
+      max_tau, i2r, imbal, cond, r_hip_x_mag, calf_deg
+    """
+    # --- FK ---
+    p = mv(ry(q_knee), T_WHEEL)
+    p = add(p, T_KNEE)
+    p = mv(ry(q_hip), p)
+    p = add(p, T_HIP)
+    p = mv(rx(q_ab), p)
+    wh = add(p, T_AB)
+    base_z = 0.10 - wh[2]
+    wheel_z = wh[2]
+
+    # 膝位置（用于 wheel-below-knee）
+    pk = mv(ry(q_hip), T_KNEE)
+    pk = mv(rx(q_ab), pk)
+    kz = pk[2] + T_AB[2]
+
+    # 关节位置
+    jk = mv(ry(q_hip), add(T_KNEE, T_HIP))
+    jk = mv(rx(q_ab), jk)
+    jk = add(T_AB, jk)
+    jh = mv(rx(q_ab), T_HIP)
+    jh = add(T_AB, jh)
+
+    # 力矩
+    rk = [wh[0]-jk[0], wh[1]-jk[1], wh[2]-jk[2]]
+    rh = [wh[0]-jh[0], wh[1]-jh[1], wh[2]-jh[2]]
+    ra = [wh[0]-T_AB[0], wh[1]-T_AB[1], wh[2]-T_AB[2]]
+    tau_ab = ra[1] * F_PER_LEG
+    tau_hip = -rh[0] * F_PER_LEG
+    tau_knee = -rk[0] * F_PER_LEG
+
+    # --- Jacobian（有限差分） ---
+    eps = 1e-6
+    def foot_rel_hip(h, k):
+        fp = mv(ry(k), T_WHEEL)
+        fp = add(fp, T_KNEE)
+        fp = mv(ry(h), fp)
+        return [fp[0] + T_AB[0] - jh[0], fp[2] + T_AB[2] - jh[2]]
+    fp0 = foot_rel_hip(q_hip, q_knee)
+    fph = foot_rel_hip(q_hip + eps, q_knee)
+    fpk = foot_rel_hip(q_hip, q_knee + eps)
+    J = np.array([
+        [(fph[0]-fp0[0])/eps, (fpk[0]-fp0[0])/eps],
+        [(fph[1]-fp0[1])/eps, (fpk[1]-fp0[1])/eps],
+    ])
+    s = np.linalg.svd(J, compute_uv=False)
+    cond = s[0] / s[-1] if s[-1] > 1e-10 else 999.0
+    min_sv = s[-1]
+
+    # --- 小腿角度（相对铅垂线） ---
+    calf_x = (-0.20011) * math.sin(q_knee)
+    calf_z = (-0.20011) * math.cos(q_knee)
+    cv_x = calf_x * math.cos(q_hip) + calf_z * math.sin(q_hip)
+    cv_z = -calf_x * math.sin(q_hip) + calf_z * math.cos(q_hip)
+    calf_deg = math.degrees(math.atan2(cv_x, -cv_z))
+
+    abs_taus = [abs(tau_ab), abs(tau_hip), abs(tau_knee)]
+    return {
+        'z': base_z, 'ab': q_ab, 'hip': q_hip, 'knee': q_knee,
+        'tau_ab': tau_ab, 'tau_hip': tau_hip, 'tau_knee': tau_knee,
+        'max_tau': max(abs_taus),
+        'i2r': tau_ab**2 + tau_hip**2 + tau_knee**2,
+        'imbal': max(abs_taus) / max(1e-10, min(abs_taus)),
+        'cond': cond, 'min_sv': min_sv,
+        'r_hip_x_mag': abs(rh[0]),
+        'calf_deg': calf_deg,
+        'kz': kz, 'wz': wheel_z,
+    }
+
+
+def rl_cost(r):
+    """RL友好度综合成本（越小越好）。
+
+    约束来源：
+      c1 — 瓶颈电机发热           τ²/τ_max²      主目标
+      c2 — 电机不均衡 > 1.3x      (imbal-1.3)²    单电机先超载
+      c3 — 髋力臂 < 8cm           (0.08 - r_hip)  髋闲置→动态响应差
+      c4 — 有效传动比 κ > 3.5     (κ - 3.5)       高刚度→冲击传递大
+      c5 — 腿的被动刚度 > 1.3x    (stiff-1.3)     刚度比→冲击吸收（新！）
+    """
+    c1 = (r['max_tau'] / 17.0) ** 2
+    c2 = max(0.0, (r['imbal'] - 1.3) / 1.0) ** 2
+    c3 = max(0.0, (0.08 - r['r_hip_x_mag'])) / 0.08
+    c4 = max(0.0, (r['cond'] - 3.5)) / 5.0
+    # stiffness ratio normalized to z=0.40 (σ_min≈0.115)
+    stiff = (0.115 / r['min_sv']) ** 2
+    c5 = max(0.0, (stiff - 1.3)) / 3.0
+    return 100.0*c1 + 50.0*c2 + 80.0*c3 + 30.0*c4 + 40.0*c5
+
+
+def sweep_z(z_targets, name, ab_max=0.44, tol=0.004):
+    """遍历 z 扫描最优姿态。"""
+    print(f"\n{'='*100}")
+    print(f" {name}")
+    print(f"{'='*100}")
+    print(f"{'z_tgt':>6} {'z':>6} {'ab':>5} {'hip':>6} {'knee':>6}  | "
+          f"{'maxτ':>6} {'imbal':>6} {'r_hip':>6} {'κ':>5}  | "
+          f"{'c1热':>6} {'c2均':>6} {'c3髋':>6} {'c4奇':>6} {'cost':>7}")
+    print("-"*100)
+    results = []
+    for zt in z_targets:
+        best_cost = float('inf')
+        best_r = None
+        for ab in [round(i*0.02, 2) for i in range(int(ab_max/0.02)+1)]:
+            for h in [round(i*0.01, 2) for i in range(260)]:
+                for kn in [round(-2.65+i*0.01, 2) for i in range(256)]:
+                    if h + kn > -0.1: continue
+                    r = get_all(ab, h, kn)
+                    if abs(r['z'] - zt) > tol: continue
+                    if r['wz'] >= r['kz']: continue
+                    if r['tau_ab'] < 0 or r['tau_hip'] < 0: continue
+                    cost = rl_cost(r)
+                    if cost < best_cost:
+                        best_cost = cost; best_r = r
+        if best_r:
+            r = best_r
+            c1s = 100.0*(r['max_tau']/17.0)**2
+            c2s = 50.0*max(0.0,(r['imbal']-1.3))**2
+            c3s = 80.0*max(0.0,(0.08-r['r_hip_x_mag']))/0.08
+            c4s = 30.0*max(0.0,(r['cond']-3.5))/5.0
+            print(f"{zt:>6.2f} {r['z']:>6.3f} {r['ab']:>5.2f} {r['hip']:>6.2f} {r['knee']:>6.2f}  | "
+                  f"{r['max_tau']:>6.3f} {r['imbal']:>6.1f}x {r['r_hip_x_mag']:>6.3f} {r['cond']:>5.1f}  | "
+                  f"{c1s:>6.1f} {c2s:>6.1f} {c3s:>6.1f} {c4s:>6.1f} {best_cost:>7.1f}")
+            results.append((zt, r))
+        else:
+            print(f"{zt:>6.2f}  —  no valid")
+    return results
+
+
+if __name__ == "__main__":
+    stand = sweep_z([round(0.36+0.01*i, 2) for i in range(10)],
+                    "STANDING — heat + RL constraints (ab free)", 0.44)
+    crawl = sweep_z([round(0.10+0.01*i, 2) for i in range(9)],
+                    "CRAWL — heat + RL constraints (ab_max=0.25)", 0.25)
+
+    print(f"\n{'='*100}")
+    print(" RECOMMENDATION")
+    print(f"{'='*100}")
+    if stand:
+        zt, r = min(stand, key=lambda x: rl_cost(x[1]))
+        print(f"  Standing: z={r['z']:.3f}  ab={r['ab']:.2f}  hip={r['hip']:.2f}  knee={r['knee']:.2f}")
+        print(f"    maxτ={r['max_tau']:.3f}  imbal={r['imbal']:.1f}x  κ={r['cond']:.1f}  r_hip={r['r_hip_x_mag']:.3f}m")
+    if crawl:
+        zt, r = min(crawl, key=lambda x: rl_cost(x[1]))
+        print(f"  Crawl:    z={r['z']:.3f}  ab={r['ab']:.2f}  hip={r['hip']:.2f}  knee={r['knee']:.2f}")
+        print(f"    maxτ={r['max_tau']:.3f}  imbal={r['imbal']:.1f}x  κ={r['cond']:.1f}  r_hip={r['r_hip_x_mag']:.3f}m")
@@ -5,11 +5,13 @@ Configures actuators as proper PD controllers at runtime:
  - Wheel joints: force = gain*(ctrl - qvel), ctrl = target velocity (rad/s)
 """

+import math
 import numpy as np
 import mujoco
 from dataclasses import dataclass
 from config import (SCENE_XML, LEG_NAMES, LEG_JOINTS, WHEEL_JOINT,
-                    DEFAULT_JOINT_ANGLES, WHEEL_RADIUS, WHEEL_TRACK)
+                    DEFAULT_JOINT_ANGLES, WHEEL_RADIUS, WHEEL_TRACK,
+                    L_THIGH, L_CALF)


@dataclass
@@ -100,14 +102,14 @@ class Robot:
                adr = self._joint_qpos_adr[name]
                self.data.qpos[adr] = DEFAULT_JOINT_ANGLES[key]

-        # Compute correct base height from default angles using 2R FK
-        # leg_length = sqrt(L1^2 + L2^2 - 2*L1*L2*cos(pi + knee))
-        import math
-        L1, L2 = 0.25, 0.20
+        # Compute correct base height from default angles using exact FK
+        # z_base = wheel_radius + L1*cos(hip) + L2*cos(hip+knee) - hip_z_offset
+        hip = DEFAULT_JOINT_ANGLES["hip_pitch"]
        knee = DEFAULT_JOINT_ANGLES["knee"]
-        leg_length = math.sqrt(L1**2 + L2**2 - 2*L1*L2*math.cos(math.pi + knee))
-        # base_z = wheel_radius + leg_length - hip_z_offset
-        base_z = 0.10 + leg_length - 0.054
+        base_z = (WHEEL_RADIUS
+                  + L_THIGH * math.cos(hip)
+                  + L_CALF * math.cos(hip + knee)
+                  - 0.054)
        self.data.qpos[2] = base_z
        self.data.qpos[3] = 1.0  # quat w

@@ -52,7 +52,7 @@ def main():
        elif c == 'q':
            ctrl.height = min(ctrl.height + 0.02, 0.45)
        elif c == 'e':
-            ctrl.height = max(ctrl.height - 0.02, 0.15)
+            ctrl.height = max(ctrl.height - 0.02, 0.17)
        elif c == '1':
            ctrl.mode = "wheel"; ctrl.prone = False
        elif c == '2':
@@ -0,0 +1,528 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+"""MuJoCo static posture optimizer for wheeled-leg standing defaults.
+
+The old posture tools are mostly analytical. This script keeps the fast MJCF
+kinematics as a candidate generator, then evaluates the best candidates in
+MuJoCo with gravity and floor contact enabled.
+
+The score is meant for RL default pose / real deployment:
+  - low peak and RMS standing torque, so one hot motor is not hidden by average
+  - wheel contact point close to the hip in X for wheel speed tracking
+  - COM projection margin inside the four-wheel support rectangle
+  - non-singular leg Jacobian for posture control authority
+  - underbody and knee clearance for obstacle tolerance
+
+Usage:
+    uv run python mujoco_sim/static_posture_optimizer.py
+    uv run python mujoco_sim/static_posture_optimizer.py --quick
+    uv run python mujoco_sim/static_posture_optimizer.py --ab-max 0.08
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+
+import mujoco
+import numpy as np
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from mujoco_sim.rl_friendly_opt import get_all  # noqa: E402
+
+SCENE_XML = REPO_ROOT / "mjcf" / "scene.xml"
+
+LEG_NAMES = ("fl", "fr", "rl", "rr")
+LEG_JOINTS = ("hip_abduction_joint", "hip_pitch_joint", "knee_joint")
+WHEEL_JOINT = "wheel_joint"
+LEG_ACTUATORS = [f"{leg}_{jt}" for leg in LEG_NAMES for jt in LEG_JOINTS]
+
+ROBOT_MASS = 12.3
+G = 9.81
+MAX_TORQUE = 17.0
+WHEEL_RADIUS = 0.10
+
+HIP_MIN, HIP_MAX = -2.58, 2.58
+KNEE_MIN, KNEE_MAX = -2.65, 2.65
+HIP_SCAN = (0.20, 1.15)
+KNEE_SCAN = (-1.90, -0.65)
+
+SOFT_WHEEL_X = 0.045
+HARD_WHEEL_X = 0.085
+MIN_COM_MARGIN = 0.045
+MIN_KNEE_CLEARANCE = 0.105
+MIN_UNDERBODY_CLEARANCE = 0.33
+MAX_COND = 3.7
+
+
+@dataclass(frozen=True)
+class Candidate:
+    ab: float
+    hip: float
+    knee: float
+    z_fk: float
+    analytic_cost: float
+
+
+@dataclass
+class StaticResult:
+    z_target: float
+    z: float
+    ab: float
+    hip: float
+    knee: float
+    cost: float
+    peak_tau: float
+    rms_tau: float
+    mean_i2r: float
+    imbal: float
+    wheel_x: float
+    cond: float
+    min_sv: float
+    com_margin_x: float
+    com_margin_y: float
+    support_margin: float
+    normal_cv: float
+    body_clearance: float
+    knee_clearance: float
+    roll: float
+    pitch: float
+    height_err: float
+
+
+def _configure_actuators(model: mujoco.MjModel) -> None:
+    """Configure leg joints as position PD and wheels as zero-velocity motors."""
+    for i in range(model.nu):
+        name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_ACTUATOR, i) or ""
+        model.actuator_biastype[i] = 1
+        model.actuator_gaintype[i] = 0
+        model.actuator_forcelimited[i] = 0
+        if "wheel" in name:
+            model.actuator_gainprm[i, 0] = 2.0
+            model.actuator_biasprm[i, 0] = 0.0
+            model.actuator_biasprm[i, 1] = 0.0
+            model.actuator_biasprm[i, 2] = -2.0
+            model.actuator_ctrlrange[i] = [-20.0, 20.0]
+        else:
+            model.actuator_gainprm[i, 0] = 60.0
+            model.actuator_biasprm[i, 0] = 0.0
+            model.actuator_biasprm[i, 1] = -60.0
+            model.actuator_biasprm[i, 2] = -3.0
+            model.actuator_ctrlrange[i] = [-3.14, 3.14]
+
+
+def _ids(model: mujoco.MjModel):
+    act = {}
+    qadr = {}
+    for leg in LEG_NAMES:
+        for jt in (*LEG_JOINTS, WHEEL_JOINT):
+            name = f"{leg}_{jt}"
+            act[name] = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, name)
+            jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, name)
+            qadr[name] = model.jnt_qposadr[jid]
+    bodies = {
+        "base": mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "base_link"),
+        **{
+            f"{leg}_wheel": mujoco.mj_name2id(
+                model, mujoco.mjtObj.mjOBJ_BODY, f"{leg}_wheel_Link"
+            )
+            for leg in LEG_NAMES
+        },
+        **{
+            f"{leg}_knee": mujoco.mj_name2id(
+                model, mujoco.mjtObj.mjOBJ_BODY, f"{leg}_knee_Link"
+            )
+            for leg in LEG_NAMES
+        },
+    }
+    return act, qadr, bodies
+
+
+def _mirrored_ab(leg: str, ab: float) -> float:
+    return ab if leg[1] == "l" else -ab
+
+
+def _set_pose(
+    model: mujoco.MjModel,
+    data: mujoco.MjData,
+    act: dict[str, int],
+    qadr: dict[str, int],
+    cand: Candidate,
+) -> None:
+    mujoco.mj_resetData(model, data)
+    for leg in LEG_NAMES:
+        vals = (_mirrored_ab(leg, cand.ab), cand.hip, cand.knee)
+        for jt, val in zip(LEG_JOINTS, vals):
+            name = f"{leg}_{jt}"
+            data.qpos[qadr[name]] = val
+            data.ctrl[act[name]] = val
+        data.ctrl[act[f"{leg}_{WHEEL_JOINT}"]] = 0.0
+    data.qpos[0:3] = [0.0, 0.0, cand.z_fk]
+    data.qpos[3] = 1.0
+    data.qpos[4:7] = 0.0
+    mujoco.mj_forward(model, data)
+
+
+def _rpy(data: mujoco.MjData, bodies: dict[str, int]) -> tuple[float, float]:
+    rot = data.xmat[bodies["base"]].reshape(3, 3)
+    roll = math.atan2(rot[2, 1], rot[2, 2])
+    pitch = math.atan2(-rot[2, 0], math.sqrt(rot[2, 1] ** 2 + rot[2, 2] ** 2))
+    return roll, pitch
+
+
+def _robot_com(model: mujoco.MjModel, data: mujoco.MjData) -> np.ndarray:
+    masses = model.body_mass[1:]
+    return (data.xipos[1:] * masses[:, None]).sum(axis=0) / masses.sum()
+
+
+def _support_metrics(
+    model: mujoco.MjModel,
+    data: mujoco.MjData,
+    bodies: dict[str, int],
+) -> tuple[float, float, float, float]:
+    wheel_xy = np.array([data.xpos[bodies[f"{leg}_wheel"]][:2] for leg in LEG_NAMES])
+    com_xy = _robot_com(model, data)[:2]
+    min_xy = wheel_xy.min(axis=0)
+    max_xy = wheel_xy.max(axis=0)
+    margin_low = com_xy - min_xy
+    margin_high = max_xy - com_xy
+    margin_x = float(min(margin_low[0], margin_high[0]))
+    margin_y = float(min(margin_low[1], margin_high[1]))
+    support_margin = float(min(margin_x, margin_y))
+
+    normal = []
+    for leg in LEG_NAMES:
+        bid = bodies[f"{leg}_wheel"]
+        fz = 0.0
+        for i in range(data.ncon):
+            con = data.contact[i]
+            b1 = model.geom_bodyid[con.geom1]
+            b2 = model.geom_bodyid[con.geom2]
+            if b1 == bid or b2 == bid:
+                wrench = np.zeros(6)
+                mujoco.mj_contactForce(model, data, i, wrench)
+                fz += abs(float(wrench[0]))
+        normal.append(fz)
+    normal = np.asarray(normal, dtype=float)
+    if normal.sum() < 1e-6:
+        normal_cv = 9.99
+    else:
+        normal_cv = float(normal.std() / max(1e-6, normal.mean()))
+    return margin_x, margin_y, support_margin, normal_cv
+
+
+def _clearance_metrics(data: mujoco.MjData, bodies: dict[str, int]) -> tuple[float, float]:
+    base_z = float(data.xpos[bodies["base"]][2])
+    # The collision box in wheelleg.xml is centered at z=0.054 with half-height 0.073.
+    body_clearance = base_z + 0.054 - 0.073
+    knee_z = min(float(data.xpos[bodies[f"{leg}_knee"]][2]) for leg in LEG_NAMES)
+    wheel_z = min(float(data.xpos[bodies[f"{leg}_wheel"]][2]) for leg in LEG_NAMES)
+    return body_clearance, knee_z - wheel_z
+
+
+def _static_cost(r: StaticResult) -> float:
+    wheel_over = max(0.0, abs(r.wheel_x) - SOFT_WHEEL_X)
+    support_short = max(0.0, MIN_COM_MARGIN - r.support_margin)
+    body_short = max(0.0, MIN_UNDERBODY_CLEARANCE - r.body_clearance)
+    knee_short = max(0.0, MIN_KNEE_CLEARANCE - r.knee_clearance)
+    cond_over = max(0.0, r.cond - MAX_COND)
+    tilt = math.hypot(r.roll, r.pitch)
+
+    return (
+        6500.0 * r.height_err**2
+        + 2.6 * (r.peak_tau / MAX_TORQUE) ** 2
+        + 1.1 * (r.rms_tau / MAX_TORQUE) ** 2
+        + 1.8 * max(0.0, r.imbal - 1.6) ** 2
+        + 0.65 * (abs(r.wheel_x) / SOFT_WHEEL_X) ** 2
+        + 9.0 * (wheel_over / max(1e-6, HARD_WHEEL_X - SOFT_WHEEL_X)) ** 2
+        + 10.0 * (support_short / MIN_COM_MARGIN) ** 2
+        + 2.0 * (r.normal_cv / 0.35) ** 2
+        + 30.0 * (cond_over / 1.0) ** 2
+        + 4.0 * (body_short / 0.06) ** 2
+        + 2.0 * (knee_short / 0.04) ** 2
+        + 1.0 * (tilt / 0.05) ** 2
+    )
+
+
+def _analytic_cost(r: dict, z_target: float) -> float:
+    wheel_over = max(0.0, abs(r["r_hip_x_mag"]) - SOFT_WHEEL_X)
+    cond_over = max(0.0, r["cond"] - MAX_COND)
+    return (
+        1400.0 * (r["z"] - z_target) ** 2
+        + 2.2 * (r["max_tau"] / MAX_TORQUE) ** 2
+        + 0.7 * (math.sqrt(r["i2r"] / 3.0) / MAX_TORQUE) ** 2
+        + 0.55 * (abs(r["r_hip_x_mag"]) / SOFT_WHEEL_X) ** 2
+        + 8.0 * (wheel_over / max(1e-6, HARD_WHEEL_X - SOFT_WHEEL_X)) ** 2
+        + 2.0 * (cond_over / 2.0) ** 2
+    )
+
+
+def generate_candidates(
+    z_target: float,
+    step: float,
+    ab_max: float,
+    keep: int,
+    z_tol: float,
+) -> list[Candidate]:
+    candidates: list[Candidate] = []
+    ab_values = np.arange(0.0, ab_max + 0.5 * step, step)
+    hip_values = np.arange(HIP_SCAN[0], HIP_SCAN[1] + 0.5 * step, step)
+    knee_values = np.arange(KNEE_SCAN[0], KNEE_SCAN[1] + 0.5 * step, step)
+    for ab in ab_values:
+        for hip in hip_values:
+            if hip < HIP_MIN or hip > HIP_MAX:
+                continue
+            for knee in knee_values:
+                if knee < KNEE_MIN or knee > KNEE_MAX:
+                    continue
+                r = get_all(float(ab), float(hip), float(knee))
+                if abs(r["z"] - z_target) > z_tol:
+                    continue
+                if r["wz"] >= r["kz"]:
+                    continue
+                if abs(r["r_hip_x_mag"]) > HARD_WHEEL_X:
+                    continue
+                if r["max_tau"] > MAX_TORQUE * 1.15:
+                    continue
+                candidates.append(
+                    Candidate(
+                        ab=float(ab),
+                        hip=float(hip),
+                        knee=float(knee),
+                        z_fk=float(r["z"]),
+                        analytic_cost=_analytic_cost(r, z_target),
+                    )
+                )
+    candidates.sort(key=lambda c: c.analytic_cost)
+    return candidates[:keep]
+
+
+def evaluate_candidate(
+    model: mujoco.MjModel,
+    data: mujoco.MjData,
+    act: dict[str, int],
+    qadr: dict[str, int],
+    bodies: dict[str, int],
+    cand: Candidate,
+    z_target: float,
+    settle_steps: int,
+    avg_steps: int,
+) -> StaticResult | None:
+    _set_pose(model, data, act, qadr, cand)
+    for _ in range(settle_steps):
+        mujoco.mj_step(model, data)
+
+    roll, pitch = _rpy(data, bodies)
+    if abs(roll) > 0.35 or abs(pitch) > 0.35:
+        return None
+
+    tau_buf = []
+    for _ in range(avg_steps):
+        mujoco.mj_step(model, data)
+        tau_buf.append([data.actuator_force[act[name]] for name in LEG_ACTUATORS])
+    tau = np.asarray(tau_buf, dtype=float).mean(axis=0)
+    tau_abs = np.abs(tau)
+    peak = float(tau_abs.max())
+    rms = float(np.sqrt(np.mean(tau * tau)))
+    mean_i2r = float(np.mean(tau * tau))
+
+    sagittal_abs = []
+    for name, value in zip(LEG_ACTUATORS, tau):
+        if "hip_pitch" in name or "knee" in name:
+            sagittal_abs.append(abs(float(value)))
+    sagittal_abs = np.asarray(sagittal_abs, dtype=float)
+    imbal = float(sagittal_abs.max() / max(1e-6, sagittal_abs.mean()))
+
+    fk = get_all(cand.ab, cand.hip, cand.knee)
+    margin_x, margin_y, support_margin, normal_cv = _support_metrics(model, data, bodies)
+    body_clearance, knee_clearance = _clearance_metrics(data, bodies)
+    z = float(data.xpos[bodies["base"]][2])
+
+    result = StaticResult(
+        z_target=z_target,
+        z=z,
+        ab=cand.ab,
+        hip=cand.hip,
+        knee=cand.knee,
+        cost=0.0,
+        peak_tau=peak,
+        rms_tau=rms,
+        mean_i2r=mean_i2r,
+        imbal=imbal,
+        wheel_x=float(fk["r_hip_x_mag"]),
+        cond=float(fk["cond"]),
+        min_sv=float(fk["min_sv"]),
+        com_margin_x=margin_x,
+        com_margin_y=margin_y,
+        support_margin=support_margin,
+        normal_cv=normal_cv,
+        body_clearance=body_clearance,
+        knee_clearance=knee_clearance,
+        roll=roll,
+        pitch=pitch,
+        height_err=z - z_target,
+    )
+    result.cost = _static_cost(result)
+    return result
+
+
+def optimize(
+    z_targets: list[float],
+    step: float,
+    ab_max: float,
+    keep: int,
+    z_tol: float,
+    settle_steps: int,
+    avg_steps: int,
+) -> list[StaticResult]:
+    model = mujoco.MjModel.from_xml_path(str(SCENE_XML))
+    data = mujoco.MjData(model)
+    _configure_actuators(model)
+    act, qadr, bodies = _ids(model)
+
+    results: list[StaticResult] = []
+    for zt in z_targets:
+        candidates = generate_candidates(zt, step=step, ab_max=ab_max, keep=keep, z_tol=z_tol)
+        best: StaticResult | None = None
+        for cand in candidates:
+            result = evaluate_candidate(
+                model,
+                data,
+                act,
+                qadr,
+                bodies,
+                cand,
+                zt,
+                settle_steps=settle_steps,
+                avg_steps=avg_steps,
+            )
+            if result is None:
+                continue
+            if best is None or result.cost < best.cost:
+                best = result
+        if best is None:
+            print(f"{zt:.2f}: no valid MuJoCo-static candidate from {len(candidates)} seeds")
+            continue
+        results.append(best)
+        print_result(best)
+    return results
+
+
+def print_header() -> None:
+    print("\nMuJoCo static posture optimization")
+    print("z_tgt      z    ab    hip   knee | peak   rms  imbal wheelX  cond | comX  comY  clrB  clrK | cost")
+    print("-" * 111)
+
+
+def print_result(r: StaticResult) -> None:
+    print(
+        f"{r.z_target:5.2f} {r.z:6.3f} {r.ab:5.2f} {r.hip:6.3f} {r.knee:6.3f} | "
+        f"{r.peak_tau:5.2f} {r.rms_tau:5.2f} {r.imbal:6.2f} "
+        f"{r.wheel_x:6.3f} {r.cond:5.2f} | "
+        f"{r.com_margin_x:5.3f} {r.com_margin_y:5.3f} "
+        f"{r.body_clearance:5.3f} {r.knee_clearance:5.3f} | "
+        f"{r.cost:6.2f}"
+    )
+
+
+def print_tables(results: list[StaticResult]) -> None:
+    if not results:
+        return
+    print("\nCopy-paste tables:")
+    print("_H_TARGET = " + repr([round(r.z_target, 3) for r in results]))
+    print("_Z_STATIC = " + repr([round(r.z, 3) for r in results]))
+    print("_HIP = " + repr([round(r.hip, 3) for r in results]))
+    print("_KNEE = " + repr([round(r.knee, 3) for r in results]))
+    print("_ABD_LEFT = " + repr([round(r.ab, 3) for r in results]))
+    print("_ABD_RIGHT = " + repr([round(-r.ab, 3) for r in results]))
+
+    best = min(results, key=lambda r: r.cost)
+    print("\nRecommended default:")
+    print(
+        f"z={best.z_target:.2f}, ab={best.ab:.3f}, hip={best.hip:.3f}, "
+        f"knee={best.knee:.3f}, peak={best.peak_tau:.2f}Nm, "
+        f"rms={best.rms_tau:.2f}Nm, support_margin={best.support_margin:.3f}m"
+    )
+    pose = []
+    for leg in LEG_NAMES:
+        pose.extend([_mirrored_ab(leg, best.ab), best.hip, best.knee])
+    pose.extend([0.0, 0.0, 0.0, 0.0])
+    print("default_dof_pos = " + repr([round(v, 3) for v in pose]))
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--z-min", type=float, default=0.36)
+    parser.add_argument("--z-max", type=float, default=0.45)
+    parser.add_argument("--z-step", type=float, default=0.01)
+    parser.add_argument("--grid-step", type=float, default=0.01)
+    parser.add_argument("--ab-max", type=float, default=0.0)
+    parser.add_argument("--keep", type=int, default=45)
+    parser.add_argument("--z-tol", type=float, default=0.035)
+    parser.add_argument("--settle-steps", type=int, default=350)
+    parser.add_argument("--avg-steps", type=int, default=80)
+    parser.add_argument("--quick", action="store_true", help="Coarser and faster scan")
+    args = parser.parse_args()
+
+    if args.quick:
+        args.grid_step = max(args.grid_step, 0.02)
+        args.keep = min(args.keep, 20)
+        args.z_tol = max(args.z_tol, 0.040)
+        args.settle_steps = min(args.settle_steps, 220)
+        args.avg_steps = min(args.avg_steps, 40)
+
+    n = int(round((args.z_max - args.z_min) / args.z_step)) + 1
+    z_targets = [round(args.z_min + i * args.z_step, 3) for i in range(n)]
+
+    print_header()
+    results = optimize(
+        z_targets,
+        step=args.grid_step,
+        ab_max=args.ab_max,
+        keep=args.keep,
+        z_tol=args.z_tol,
+        settle_steps=args.settle_steps,
+        avg_steps=args.avg_steps,
+    )
+    print_tables(results)
+
+
+if __name__ == "__main__":
+    main()
@@ -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` 的航点和避障几何定义；默认使用 `points_20260715_120154.json` 与 `1hao.xml`。
+
+运行时生成的日志、临时 XML、`route_check_runs/` 和批量实验输出不纳入版本库。人工打点形成的路线快照保存在 `../tools/nav_tools/points/`，大量重复仿真轨迹仍不复制。
@@ -0,0 +1,110 @@
+import argparse
+import torch
+import torch.nn as nn
+import numpy as np
+from pathlib import Path
+import onnxruntime as ort
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+
+class PolicyMLP(nn.Module):
+    def __init__(self, obs_dim=53, action_dim=16):
+        super().__init__()
+        self.register_buffer("obs_mean", torch.zeros(obs_dim))
+        self.register_buffer("obs_std", torch.ones(obs_dim))
+        self.net = nn.Sequential(
+            nn.Linear(obs_dim, 512), nn.ELU(),
+            nn.Linear(512, 256), nn.ELU(),
+            nn.Linear(256, 128), nn.ELU(),
+            nn.Linear(128, action_dim),
+        )
+
+    def forward(self, x):
+        x = (x - self.obs_mean) / torch.clamp(self.obs_std, min=1e-6)
+        return self.net(x)
+
+def load_policy(model_path, device):
+    ckpt = torch.load(model_path, map_location=device, weights_only=False)
+    state_dict = ckpt["actor_state_dict"]
+    weight_key = "mlp.0.weight" if "mlp.0.weight" in state_dict else "net.0.weight"
+    obs_dim = state_dict[weight_key].shape[1]
+
+    output_key = "mlp.6.weight" if "mlp.6.weight" in state_dict else "net.6.weight"
+    action_dim = state_dict[output_key].shape[0]
+
+    model = PolicyMLP(obs_dim=obs_dim, action_dim=action_dim)
+    my_sd = {}
+    for k, v in state_dict.items():
+        if k.startswith("mlp."):
+            my_sd[k.replace("mlp.", "net.")] = v
+        elif k.startswith("net."):
+            my_sd[k] = v
+        elif k == "obs_normalizer._mean":
+            my_sd["obs_mean"] = v.squeeze()
+        elif k == "obs_normalizer._var":
+            my_sd["obs_std"] = torch.sqrt(v.squeeze() + 1e-5)
+
+    model.load_state_dict(my_sd, strict=False)
+    model.eval()
+    model.to(device)
+    return model, obs_dim
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--pt-path",
+        "--pt_path",
+        dest="pt_path",
+        type=Path,
+        default=PROJECT_ROOT / "model_rough.pt",
+        help="PyTorch checkpoint to export (default: ../model_rough.pt).",
+    )
+    args = parser.parse_args()
+
+    pt_path = args.pt_path.expanduser().resolve()
+    if not pt_path.exists():
+        print(f"File not found: {pt_path}")
+        return
+
+    device = torch.device("cpu")
+    print(f"Loading {pt_path}...")
+    model, obs_dim = load_policy(pt_path, device)
+
+    onnx_path = pt_path.with_suffix(".onnx")
+
+    dummy_input = torch.randn(1, obs_dim, device=device)
+
+    print(f"Exporting to {onnx_path}...")
+    torch.onnx.export(
+        model,
+        dummy_input,
+        str(onnx_path),
+        export_params=True,
+        opset_version=14,
+        do_constant_folding=True,
+        input_names=["obs"],
+        output_names=["action"],
+        dynamic_axes={"obs": {0: "batch_size"}, "action": {0: "batch_size"}}
+    )
+
+    print("Verifying ONNX export...")
+    try:
+        session = ort.InferenceSession(str(onnx_path))
+        with torch.no_grad():
+            pt_out = model(dummy_input).numpy()
+        onnx_out = session.run(["action"], {"obs": dummy_input.numpy()})[0]
+
+        max_diff = np.max(np.abs(pt_out - onnx_out))
+        mean_diff = np.mean(np.abs(pt_out - onnx_out))
+        print(f"ONNX vs PyTorch - max_diff: {max_diff:.6f}, mean_diff: {mean_diff:.6f}")
+
+        if max_diff < 1e-4:
+            print("ONNX export verified OK.")
+        else:
+            print("WARNING: ONNX export has significant divergence from PyTorch model.")
+    except ImportError:
+        print("onnxruntime not installed. Skipping verification. Install with: pip install onnxruntime")
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,238 @@
+#!/usr/bin/env python3
+"""Sweep IK compensation parameters in the standalone sim2sim scene."""
+
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import json
+import sys
+from pathlib import Path
+from typing import Any
+
+THIS_DIR = Path(__file__).resolve().parent
+PROJECT_ROOT = THIS_DIR.parent
+if str(THIS_DIR) not in sys.path:
+    sys.path.insert(0, str(THIS_DIR))
+
+SIM_PATH = THIS_DIR / "ik_slalom_sim2sim.py"
+spec = importlib.util.spec_from_file_location("ik_slalom_sim2sim", SIM_PATH)
+if spec is None or spec.loader is None:
+    raise RuntimeError(f"Cannot load {SIM_PATH}")
+sim = importlib.util.module_from_spec(spec)
+sys.modules["ik_slalom_sim2sim"] = sim
+spec.loader.exec_module(sim)
+
+
+TRIALS = [
+    {"name": "forward", "speed": 1.0, "yaw": 0.0, "target_vx": 1.0, "target_yaw": 0.0},
+    {"name": "yaw", "speed": 0.0, "yaw": 1.0, "target_vx": 0.0, "target_yaw": 1.0},
+    {"name": "arc", "speed": 1.0, "yaw": 1.0, "target_vx": 1.0, "target_yaw": 1.0},
+]
+
+
+def parse_float_list(text: str) -> list[float]:
+    return [float(x.strip()) for x in text.split(",") if x.strip()]
+
+
+def parse_bool_list(text: str) -> list[bool]:
+    out: list[bool] = []
+    for item in text.split(","):
+        key = item.strip().lower()
+        if not key:
+            continue
+        if key in {"1", "true", "on", "yes"}:
+            out.append(True)
+        elif key in {"0", "false", "off", "no"}:
+            out.append(False)
+        else:
+            raise argparse.ArgumentTypeError(f"Invalid bool item: {item}")
+    return out
+
+
+def make_sim_args(args: argparse.Namespace, trial: dict[str, float | str], cfg: dict[str, Any]) -> argparse.Namespace:
+    argv = [
+        "ik_slalom_sim2sim.py",
+        "--test",
+        str(trial["name"]),
+        "--duration",
+        str(args.duration),
+        "--settle",
+        str(args.settle),
+        "--speed",
+        str(trial["speed"]),
+        "--yaw-rate",
+        str(trial["yaw"]),
+        "--posture",
+        "custom",
+        "--custom-abduction",
+        str(args.custom_abduction),
+        "--custom-hip",
+        str(args.custom_hip),
+        "--custom-knee",
+        str(args.custom_knee),
+        "--wheel-model",
+        "direct",
+        "--linear-wheel-gain",
+        str(args.linear_wheel_gain),
+        "--direct-yaw-wheel-gain",
+        str(args.direct_yaw_wheel_gain),
+        "--max-wheel-speed",
+        str(args.max_wheel_speed),
+        "--wheel-accel-limit",
+        str(args.wheel_accel_limit),
+        "--yaw-rate-kp",
+        str(cfg["yaw_rate_kp"]),
+        "--encoder-posture-kp",
+        str(cfg["encoder_posture_kp"]),
+        "--encoder-posture-max",
+        str(cfg["encoder_posture_max"]),
+        "--roll-comp-gain",
+        str(cfg["roll_comp_gain"]),
+        "--pitch-comp-gain",
+        str(cfg["pitch_comp_gain"]),
+        "--no-realtime",
+    ]
+    argv.append("--imu-posture" if cfg["imu_posture"] else "--no-imu-posture")
+    argv.append("--encoder-guard" if cfg["encoder_guard"] else "--no-encoder-guard")
+    argv.append("--imu-guard" if cfg["imu_guard"] else "--no-imu-guard")
+    old_argv = sys.argv
+    try:
+        sys.argv = argv
+        return sim.parse_args()
+    finally:
+        sys.argv = old_argv
+
+
+def score_trial(out: dict[str, Any], trial: dict[str, float | str]) -> dict[str, float]:
+    vx = float(out["mean_body_vx_mps"])
+    yaw = float(out["mean_yaw_rate_rad_s"])
+    vx_err = abs(vx - float(trial["target_vx"]))
+    yaw_err = abs(yaw - float(trial["target_yaw"]))
+    return {
+        "vx": vx,
+        "yaw": yaw,
+        "imu_gyro_z": float(out["mean_imu_gyro_z_rad_s"]),
+        "vx_err": vx_err,
+        "yaw_err": yaw_err,
+        "err": vx_err + yaw_err,
+    }
+
+
+def run_sweep(args: argparse.Namespace) -> list[dict[str, Any]]:
+    rows: list[dict[str, Any]] = []
+    for imu_posture in args.imu_posture_values:
+        for encoder_guard in args.encoder_guard_values:
+            for imu_guard in args.imu_guard_values:
+                for encoder_posture_kp in args.encoder_posture_kps:
+                    for encoder_posture_max in args.encoder_posture_maxs:
+                        for yaw_rate_kp in args.yaw_rate_kps:
+                            for roll_comp_gain in args.roll_comp_gains:
+                                for pitch_comp_gain in args.pitch_comp_gains:
+                                    cfg = {
+                                        "imu_posture": imu_posture,
+                                        "encoder_guard": encoder_guard,
+                                        "imu_guard": imu_guard,
+                                        "encoder_posture_kp": encoder_posture_kp,
+                                        "encoder_posture_max": encoder_posture_max,
+                                        "yaw_rate_kp": yaw_rate_kp,
+                                        "roll_comp_gain": roll_comp_gain,
+                                        "pitch_comp_gain": pitch_comp_gain,
+                                    }
+                                    detail: list[dict[str, Any]] = []
+                                    speed_error = 0.0
+                                    max_tilt = 0.0
+                                    max_leg = 0.0
+                                    mean_wheel_err = 0.0
+                                    stable_all = True
+                                    for trial in TRIALS:
+                                        sim_args = make_sim_args(args, trial, cfg)
+                                        out = sim.run_one(str(trial["name"]), sim_args)
+                                        trial_score = score_trial(out, trial)
+                                        trial_score["test"] = str(trial["name"])
+                                        detail.append(trial_score)
+                                        speed_error += trial_score["err"]
+                                        max_tilt = max(max_tilt, float(out["max_tilt_deg"]))
+                                        max_leg = max(max_leg, float(out["max_leg_encoder_error_rad"]))
+                                        mean_wheel_err += float(out["mean_wheel_speed_error_rad_s"])
+                                        stable_all = stable_all and bool(out["stable"])
+
+                                    score = (
+                                        speed_error
+                                        + args.tilt_weight * max_tilt
+                                        + args.leg_error_weight * max_leg
+                                        + args.wheel_error_weight * (mean_wheel_err / len(TRIALS))
+                                    )
+                                    row = {
+                                        **cfg,
+                                        "score": round(score, 6),
+                                        "speed_error_sum": round(speed_error, 6),
+                                        "max_tilt_deg": round(max_tilt, 5),
+                                        "max_leg_encoder_error_rad": round(max_leg, 6),
+                                        "mean_wheel_speed_error_rad_s": round(mean_wheel_err / len(TRIALS), 6),
+                                        "stable_all": stable_all,
+                                        "detail": detail,
+                                    }
+                                    rows.append(row)
+                                    print(
+                                        "DONE "
+                                        + json.dumps(
+                                            {k: v for k, v in row.items() if k != "detail"},
+                                            ensure_ascii=False,
+                                        ),
+                                        flush=True,
+                                    )
+    rows.sort(key=lambda r: float(r["score"]))
+    return rows
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--duration", type=float, default=3.0)
+    parser.add_argument("--settle", type=float, default=1.5)
+    parser.add_argument("--custom-abduction", type=float, default=0.2)
+    parser.add_argument("--custom-hip", type=float, default=1.697)
+    parser.add_argument("--custom-knee", type=float, default=-2.650)
+    parser.add_argument("--linear-wheel-gain", type=float, default=12.5)
+    parser.add_argument("--direct-yaw-wheel-gain", type=float, default=8.0)
+    parser.add_argument("--max-wheel-speed", type=float, default=12.0)
+    parser.add_argument("--wheel-accel-limit", type=float, default=35.0)
+    parser.add_argument("--imu-posture-values", type=parse_bool_list, default=[True, False])
+    parser.add_argument("--encoder-guard-values", type=parse_bool_list, default=[True])
+    parser.add_argument("--imu-guard-values", type=parse_bool_list, default=[True])
+    parser.add_argument("--encoder-posture-kps", type=parse_float_list, default=[0.0, 0.05, 0.15, 0.30])
+    parser.add_argument("--encoder-posture-maxs", type=parse_float_list, default=[0.03])
+    parser.add_argument("--yaw-rate-kps", type=parse_float_list, default=[0.0, 0.4, 0.8])
+    parser.add_argument("--roll-comp-gains", type=parse_float_list, default=[0.35])
+    parser.add_argument("--pitch-comp-gains", type=parse_float_list, default=[0.35])
+    parser.add_argument("--tilt-weight", type=float, default=0.02)
+    parser.add_argument("--leg-error-weight", type=float, default=0.5)
+    parser.add_argument("--wheel-error-weight", type=float, default=0.0)
+    parser.add_argument("--top", type=int, default=12)
+    parser.add_argument("--json", type=Path, default=None)
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    rows = run_sweep(args)
+    if args.json:
+        args.json.write_text(json.dumps(rows, indent=2), encoding="utf-8")
+
+    print("\nTop compensation parameter sets")
+    print("rank score speed_err tilt leg_err wheel_err imu enc_kp yaw_kp enc_guard imu_guard")
+    for i, row in enumerate(rows[: args.top], 1):
+        print(
+            f"{i:2d} {row['score']:7.4f} {row['speed_error_sum']:7.4f} "
+            f"{row['max_tilt_deg']:5.2f} {row['max_leg_encoder_error_rad']:7.4f} "
+            f"{row['mean_wheel_speed_error_rad_s']:7.4f} "
+            f"{int(row['imu_posture'])} {row['encoder_posture_kp']:6.3f} "
+            f"{row['yaw_rate_kp']:6.3f} {int(row['encoder_guard'])} {int(row['imu_guard'])}"
+        )
+        for d in row["detail"]:
+            print(f"   {d['test']:<7} vx={d['vx']:+.3f} yaw={d['yaw']:+.3f} err={d['err']:.3f}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
zeitvex	55094bf682	[tools] 整理导航打点工具与抽样PCD	2026-07-27 13:39:52 +08:00
zeitvex	4ee4af028c	[sim] 整理后期Sim2Sim与比赛Rough策略	2026-07-27 13:20:41 +08:00
zeitvex	c05c1cb162	[sim] 添加后期MuJoCo姿态与MPC工具集	2026-07-27 12:57:03 +08:00
zeitvex	d8c5d34091	[train] 整理比赛最终训练架构	2026-07-27 12:48:12 +08:00
zeitvex	60f7a08e91	[train] 增强Sim2Real训练随机化配置	2026-07-27 12:39:27 +08:00
zeitvex	b08956aec7	[train] 更新新MJCF与第一版完整训练框架	2026-07-27 12:31:21 +08:00