[train] 增强Sim2Real训练随机化配置

[train] 更新新MJCF与第一版完整训练框架
2026-07-27 12:39:27 +08:00 · 2026-07-27 12:31:21 +08:00
116 changed files with 7681 additions and 6760 deletions
@@ -35,6 +35,7 @@ logs/
 checkpoints/
 wandb/
 sim2sim_log_*.txt
+**/sim2sim_temp.xml

 # IDE and operating system files
 .idea/
@@ -0,0 +1,30 @@
+# 版本演进
+
+本项目使用同一条 `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.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` 随机缩放。
+- 增加作用于机身的连续随机外力和力矩扰动。
@@ -1,11 +1,11 @@
 # 软件

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

 ```text
 05_software/
 ├─ train/
-│  └─ rc_mjlab/       # 训练、MJCF、MuJoCo、Sim2Sim 和本地 mjlab 依赖
+│  └─ rc_mjlab/       # 训练、MJCF、Sim2Sim 和本地 mjlab 依赖
 └─ real/
   ├─ ik_real/        # IK 轨迹与早期真机控制
   └─ sim2real/       # 第一代 Python 策略真机部署
@@ -18,8 +18,6 @@ MJCF + mjlab task
        |
        v
   PPO 训练策略
-        |
-        +----> MuJoCo 独立模型调试
        |
        +----> Sim2Sim 策略验证
        |
@@ -28,7 +26,7 @@ MJCF + mjlab task
 IK real --------------------------------> 电机
 ```

-`rc_mjlab` 在早期版本中是自包含工程。训练、MJCF、独立 MuJoCo、Sim2Sim 和策略权重通过相对路径绑定，因此本次保留其原始内部布局，没有为了目录外观拆散。
+`rc_mjlab` 是自包含工程。训练、MJCF、Sim2Sim 和策略权重通过相对路径绑定，因此保留其内部布局，没有为了目录外观拆散。第一代完整闭环见 `v0.3.0`，第一份新版 MJCF 与训练框架见 `v0.4.0`，增强 Sim2Real 随机化的第二版训练配置见 `v0.5.0`。

 详细说明见：

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

-`rc_mjlab/` 是 16DOF 轮足机器人的第一代自包含训练与仿真工程。
+`rc_mjlab/` 保存 16DOF 轮足机器人的当前训练与 Sim2Sim 工程。历史快照由 Git Tag 保留，不在目录中复制 `old`、`new` 或 `final` 版本。
+
+当前内容对应 `v0.5.0`，在 `v0.4.0` 的完整新版训练工程上增强了面向 Sim2Real 的随机化配置。

 ## 内容

 - `src/robot`：Flat、Rough、Crawl 训练任务和自定义 MDP
 - `mjcf`：轮足机器人 MuJoCo 模型和网格
- `mujoco_sim`：不依赖策略的独立 MuJoCo/MPC 调试工具
 - `sim2sim`：策略加载、交互控制和比赛地形验证
 - `mjlab`：固定版本的本地训练框架依赖
- `model_rough.pt`、`model_crawl.pt`：对应的早期策略权重
+- `model_rough.pt`：本阶段 Rough 策略权重
 - `pyproject.toml`、`uv.lock`：Python 环境与依赖锁定

+与 `v0.3.0` 相比，本版本更新了 MJCF 质量和惯性参数，并将 mjlab 上游基准从 `00409797` 更新到 `40f8d93e`。机械 CAD 未发生变化。
+
+与 `v0.4.0` 相比，本版本没有再次修改 MJCF 和训练框架，只调整训练环境配置：投影重力噪声从 `±0.05` 扩大到 `±0.08`，最大动作延迟从 2 步增加到 4 步，扩大摩擦、刚度和阻尼随机化，增加腿部质量随机化与连续外力/力矩扰动。
+
 工程命令和任务说明见 [`rc_mjlab/README.md`](rc_mjlab/README.md)，本地依赖来源见 [`rc_mjlab/DEPENDENCIES.md`](rc_mjlab/DEPENDENCIES.md)。
@@ -2,14 +2,14 @@

 ## Python 环境

- Python `>=3.10`
+- Python `>=3.10,<3.14`
 - `uv` 依赖管理
- MuJoCo development wheel
+- MuJoCo `3.8` 系列
 - `mjlab[cu128]`
 - PyTorch CUDA 12.8 环境
 - `pynput`

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

 ```toml
 [tool.uv.sources]
@@ -19,17 +19,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`、本地缓存、生成日志和运行时临时文件。

 ## 基本入口

@@ -2,6 +2,8 @@

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

+> 本目录对应 `v0.5.0`：在第一份完整的新 MJCF 与新框架训练工程上，扩大观测噪声、动作延迟和动力学随机化范围，并加入持续外力扰动。该快照包含 `model_rough.pt`；未包含独立 `mujoco_sim` 工具和单独的 Crawl 策略权重，相关早期内容仍可通过 `v0.3.0` 查看。
+
 ---

 ## 🛠️ 项目简介
@@ -51,8 +53,7 @@ rc_mjlab/
 │   ├── wheelleg.xml              # 机器人 MuJoCo 模型（含网格引用）
 │   ├── scene.xml                 # mjlab 场景入口文件
 │   └── meshes/                   # STL/OBJ 碰撞与外观网格
-├── mujoco_sim/                   # 独立 MPC 仿真调试工具（不依赖 RL 训练）
-├── logs/                         # 训练日志（rsl_rl 格式，按任务名/日期/checkpoint 归档）
+├── model_rough.pt                # 本阶段用于回放和 Sim2Sim 的 Rough 策略
 ├── pyproject.toml                # 项目依赖（uv 管理，含清华镜像源加速）
 └── uv.lock                       # 精确依赖锁定文件
 ```
@@ -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]]
@@ -1,57 +0,0 @@
-"""Robot constants and control parameters."""
-
-import numpy as np
-from pathlib import Path
-
-# Paths
-REPO_ROOT = Path(__file__).resolve().parents[1]
-SCENE_XML = REPO_ROOT / "mjcf" / "scene.xml"
-MJCF_PATH = REPO_ROOT / "mjcf" / "wheelleg.xml"
-
-# Robot geometry
-WHEEL_RADIUS = 0.10  # m
-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 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_JOINT_ANGLES = {
-    "hip_abduction": 0.0,
-    "hip_pitch": 0.93,
-    "knee": -1.65,
-}
-
-# 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)
-
-# Control rates
-SIM_DT = 0.002       # 500 Hz (from scene.xml)
-CTRL_DT = 0.004      # 250 Hz control loop
-CTRL_DECIMATION = int(CTRL_DT / SIM_DT)
-
-# Wheel drive
-WHEEL_VEL_MAX = 10.0  # rad/s max wheel command
-
-# 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
-
-# Gait parameters
-GAIT_FREQ = 2.5     # Hz
-GAIT_DUTY = 0.6     # stance fraction
-SWING_HEIGHT = 0.06  # m
-
-# 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}
@@ -1,287 +0,0 @@
-"""Main controller: wheel mode + trot mode for wheeled-legged robot.
-
-Wheel mode: differential drive + leg posture hold (height/roll/pitch compensation)
-Trot mode: quadruped gait with wheel-assisted propulsion
-
-Actuator interface:
-  - Leg joints: ctrl = target angle (PD: kp=60, kd=3)
-  - Wheel joints: ctrl = target velocity in rad/s (gain=2.0)
-"""
-
-import numpy as np
-from robot import Robot, RobotState
-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,
-    GAIT_FREQ, GAIT_DUTY, SWING_HEIGHT, PHASE_OFFSETS,
-)
-
-
-class Controller:
-    """Wheeled-legged robot controller."""
-
-    def __init__(self, robot: Robot):
-        self.robot = robot
-        self.dynamics = Dynamics()
-
-        # User commands
-        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
-
-        # Mode: "wheel", "trot", or "mpc"
-        self.mode = "wheel"
-
-        # Prone (lie down) state
-        self.prone = False
-
-        # MPC controller
-        self._mpc_ctrl = MPCController(robot)
-        self._mpc_active = False  # track torque mode state
-
-        # Gait state
-        self._gait_phase = 0.0
-
-        # Smoothed commands for trot mode (avoid sudden jumps)
-        self._smooth_vx = 0.0
-        self._smooth_vy = 0.0
-        self._smooth_yaw = 0.0
-
-        # Default leg angles
-        self._default_q = np.array([
-            DEFAULT_JOINT_ANGLES["hip_abduction"],
-            DEFAULT_JOINT_ANGLES["hip_pitch"],
-            DEFAULT_JOINT_ANGLES["knee"],
-        ])
-
-        # Swing leg memory
-        self._swing_start_foot = {leg: np.zeros(3) for leg in LEG_NAMES}
-        self._last_contact = {leg: True for leg in LEG_NAMES}
-
-    def compute(self, state: RobotState, dt: float) -> tuple[np.ndarray, np.ndarray]:
-        # Smooth all velocity commands (both modes)
-        alpha = min(dt * 3.0, 1.0)  # ~0.33s time constant
-        self._smooth_vx += alpha * (self.vel_x - self._smooth_vx)
-        self._smooth_vy += alpha * (self.vel_y - self._smooth_vy)
-        self._smooth_yaw += alpha * (self.yaw_rate - self._smooth_yaw)
-
-        if self.prone:
-            self._ensure_position_mode()
-            return self._prone_mode()
-        if self.mode == "mpc":
-            return self._mpc_mode(state, dt)
-        if self.mode == "wheel":
-            self._ensure_position_mode()
-            return self._wheel_mode(state, dt)
-        else:
-            self._ensure_position_mode()
-            return self._trot_mode(state, dt)
-
-    def _mpc_mode(self, state: RobotState, dt: float):
-        """MPC locomotion: MIT motor protocol (PD + MPC feedforward torque)."""
-        # Switch to torque mode if not already
-        if not self._mpc_active:
-            self.robot.enable_torque_mode()
-            self._mpc_active = True
-
-        # Sync commands to MPC controller
-        self._mpc_ctrl.vel_x = self.vel_x
-        self._mpc_ctrl.vel_y = self.vel_y
-        self._mpc_ctrl.yaw_rate = self.yaw_rate
-        self._mpc_ctrl.height = self.height
-
-        # Compute and apply (sets ctrl directly via set_ctrl_mit)
-        self._mpc_ctrl.compute(state, dt)
-        # Return dummy - ctrl already set
-        return np.zeros(12), np.zeros(4)
-
-    def _ensure_position_mode(self):
-        """Switch back to position PD mode if coming from MPC."""
-        if self._mpc_active:
-            self.robot.enable_position_mode()
-            self._mpc_active = False
-
-    def _prone_mode(self):
-        """Lie down: actual prone pose from real robot."""
-        leg_targets = np.zeros(12)
-        for i, leg in enumerate(LEG_NAMES):
-            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
-        return leg_targets, np.zeros(4)
-
-    # ─────────────────────────────────────────────────────────────────────
-    # WHEEL MODE
-    # ─────────────────────────────────────────────────────────────────────
-
-    def _wheel_mode(self, state: RobotState, dt: float):
-        """Wheel drive + leg posture hold.
-
-        vel_y: limited effect in wheel mode (differential drive cannot produce
-        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)
-        return leg_targets, wheel_targets
-
-    def _posture_control(self, state: RobotState) -> np.ndarray:
-        """Leg joint targets: table-interpolated height control."""
-        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]
-
-        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))
-
-        roll_corr = -KP_ROLL * state.rpy[0]
-        pitch_corr = -KP_PITCH * 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+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)
-
-        return leg_targets
-
-    # ─────────────────────────────────────────────────────────────────────
-    # TROT MODE
-    # ─────────────────────────────────────────────────────────────────────
-
-    def _trot_mode(self, state: RobotState, dt: float):
-        """Trot gait with wheel assist."""
-        # Advance gait phase
-        self._gait_phase = (self._gait_phase + dt * GAIT_FREQ) % 1.0
-
-        # Contact state
-        contacts = {}
-        for leg in LEG_NAMES:
-            phase = (self._gait_phase + PHASE_OFFSETS[leg]) % 1.0
-            contacts[leg] = phase < GAIT_DUTY
-
-        # Pinocchio update
-        q_pin, dq_pin = self.robot.get_qpos_qvel_for_pinocchio()
-        self.dynamics.update(q_pin, dq_pin)
-
-        leg_targets = np.zeros(12)
-        wheel_targets = np.zeros(4)
-
-        for i, leg in enumerate(LEG_NAMES):
-            if contacts[leg]:
-                # Stance: posture hold
-                leg_targets[i*3:(i+1)*3] = self._stance_leg_target(state, leg)
-                self._swing_start_foot[leg] = self.dynamics.get_foot_pos(leg)
-                self._last_contact[leg] = True
-                # Wheel: drive with smoothed velocity
-                wheel_targets[i] = self._differential_drive_single(
-                    self._smooth_vx, self._smooth_yaw, leg)
-            else:
-                # Swing: IK trajectory
-                swing_phase = self._get_swing_phase(leg)
-                target_foot = self._compute_swing_target(leg, state, swing_phase)
-                q_ik = self.dynamics.inverse_kinematics(leg, target_foot, q_pin)
-                leg_targets[i*3:(i+1)*3] = q_ik
-                self._last_contact[leg] = False
-                # Wheel: zero (free during swing)
-                wheel_targets[i] = 0.0
-
-        return leg_targets, wheel_targets
-
-    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_clamp = np.clip(self.height, _H[0], _H[-1])
-        q_hip = float(np.interp(h_clamp, _H, _HIP))
-        q_knee = float(np.interp(h_clamp, _H, _KNEE))
-
-        roll_corr = -KP_ROLL * state.rpy[0]
-        pitch_corr = -KP_PITCH * state.rpy[1]
-        side = 1.0 if leg[1] == "l" else -1.0
-        lateral_lean = 0.3 * self.vel_y
-
-        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),
-        ])
-
-    # ─────────────────────────────────────────────────────────────────────
-    # DIFFERENTIAL DRIVE
-    # ─────────────────────────────────────────────────────────────────────
-
-    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
-        targets = np.zeros(4)
-        for i, leg in enumerate(LEG_NAMES):
-            targets[i] = vel_left if leg[1] == "l" else vel_right
-        return np.clip(targets, -WHEEL_VEL_MAX, WHEEL_VEL_MAX)
-
-    def _differential_drive_single(self, vel_x: float, yaw_rate: float, leg: str) -> float:
-        if leg[1] == "l":
-            v = (vel_x - 0.5 * WHEEL_TRACK * yaw_rate) / WHEEL_RADIUS
-        else:
-            v = (vel_x + 0.5 * WHEEL_TRACK * yaw_rate) / WHEEL_RADIUS
-        return np.clip(v, -WHEEL_VEL_MAX, WHEEL_VEL_MAX)
-
-    # ─────────────────────────────────────────────────────────────────────
-    # SWING TRAJECTORY
-    # ─────────────────────────────────────────────────────────────────────
-
-    def _get_swing_phase(self, leg: str) -> float:
-        phase = (self._gait_phase + PHASE_OFFSETS[leg]) % 1.0
-        if phase < GAIT_DUTY:
-            return 0.0
-        return (phase - GAIT_DUTY) / (1.0 - GAIT_DUTY)
-
-    def _compute_swing_target(self, leg: str, state: RobotState,
-                              swing_phase: float) -> np.ndarray:
-        """Swing foot target with Raibert heuristic using COMMANDED velocity."""
-        p_start = self._swing_start_foot[leg]
-        p_end = self._compute_touchdown(leg, state)
-
-        s = swing_phase
-        s_mj = 10*s**3 - 15*s**4 + 6*s**5
-
-        pos = p_start + (p_end - p_start) * s_mj
-
-        # Z lift
-        z_lift = 64.0 * s**3 * (1.0 - s)**3
-        pos[2] = p_start[2] + SWING_HEIGHT * z_lift
-
-        return pos
-
-    def _compute_touchdown(self, leg: str, state: RobotState) -> np.ndarray:
-        """Raibert heuristic using COMMANDED velocity.
-
-        When commands are zero, foot lands at its takeoff position (no net motion).
-        When commands are nonzero, foot placement is offset by commanded velocity.
-        """
-        # Base: land where the foot took off (zero net displacement)
-        td = self._swing_start_foot[leg].copy()
-
-        # Add commanded velocity offset (Raibert-style)
-        t_stance = (1.0 / GAIT_FREQ) * GAIT_DUTY
-        yaw = state.rpy[2]
-        c, s = np.cos(yaw), np.sin(yaw)
-        R_z = np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]])
-        cmd_vel_world = R_z @ np.array([self._smooth_vx, self._smooth_vy, 0.0])
-
-        td[0] += cmd_vel_world[0] * t_stance * 0.5
-        td[1] += cmd_vel_world[1] * t_stance * 0.5
-        td[2] = WHEEL_RADIUS  # ground level
-        return td
@@ -1,97 +0,0 @@
-"""Pinocchio dynamics: FK, Jacobian, IK for the wheeled-legged robot."""
-
-import numpy as np
-import pinocchio as pin
-from config import MJCF_PATH, LEG_NAMES
-
-# Foot frame names in Pinocchio model (wheel link centers)
-FOOT_FRAMES = {leg: f"{leg}_wheel_Link" for leg in LEG_NAMES}
-
-# Leg joint names for each leg
-_LEG_JOINT_NAMES = {
-    leg: [f"{leg}_{jt}" for jt in ("hip_abduction_joint", "hip_pitch_joint", "knee_joint")]
-    for leg in LEG_NAMES
-}
-
-
-class Dynamics:
-    """Pinocchio-based kinematics/dynamics. Deployable on real hardware."""
-
-    def __init__(self):
-        self.model = pin.buildModelFromMJCF(str(MJCF_PATH))
-        self.data = self.model.createData()
-
-        # Cache frame IDs
-        self._foot_fids = {}
-        for leg, fname in FOOT_FRAMES.items():
-            self._foot_fids[leg] = self.model.getFrameId(fname)
-
-        # Cache joint velocity indices for each leg (3 joints)
-        self._leg_v_indices = {}
-        for leg, jnames in _LEG_JOINT_NAMES.items():
-            indices = []
-            for jn in jnames:
-                jid = self.model.getJointId(jn)
-                indices.append(self.model.joints[jid].idx_v)
-            self._leg_v_indices[leg] = indices
-
-        # Cache joint config indices for each leg
-        self._leg_q_indices = {}
-        for leg, jnames in _LEG_JOINT_NAMES.items():
-            indices = []
-            for jn in jnames:
-                jid = self.model.getJointId(jn)
-                indices.append(self.model.joints[jid].idx_q)
-            self._leg_q_indices[leg] = indices
-
-    def update(self, q: np.ndarray, dq: np.ndarray):
-        """Forward kinematics + Jacobians.
-
-        Args:
-            q: Pinocchio config (nq=23: pos3, quat_xyzw4, joints16)
-            dq: Pinocchio velocity (nv=22: v_body3, w_body3, joints16)
-        """
-        pin.forwardKinematics(self.model, self.data, q, dq)
-        pin.updateFramePlacements(self.model, self.data)
-        pin.computeJointJacobians(self.model, self.data, q)
-
-    def get_foot_pos(self, leg: str) -> np.ndarray:
-        """Foot (wheel center) position in world frame (3,)."""
-        return self.data.oMf[self._foot_fids[leg]].translation.copy()
-
-    def get_foot_jacobian_leg(self, leg: str) -> np.ndarray:
-        """3x3 linear Jacobian of foot w.r.t. 3 leg joints (world frame)."""
-        fid = self._foot_fids[leg]
-        J_full = pin.getFrameJacobian(
-            self.model, self.data, fid, pin.LOCAL_WORLD_ALIGNED)[:3, :]
-        cols = self._leg_v_indices[leg]
-        return J_full[:, cols]
-
-    def inverse_kinematics(self, leg: str, target_pos: np.ndarray,
-                           q_current: np.ndarray, max_iter=30, eps=1e-4) -> np.ndarray:
-        """Numerical IK for one leg. Returns (3,) joint angles.
-
-        Args:
-            leg: Leg name
-            target_pos: Desired foot position in world frame (3,)
-            q_current: Current full Pinocchio config (nq=23)
-        """
-        q = q_current.copy()
-        fid = self._foot_fids[leg]
-        q_indices = self._leg_q_indices[leg]
-
-        for _ in range(max_iter):
-            pin.forwardKinematics(self.model, self.data, q)
-            pin.updateFramePlacements(self.model, self.data)
-            err = target_pos - self.data.oMf[fid].translation
-            if np.linalg.norm(err) < eps:
-                break
-            pin.computeJointJacobians(self.model, self.data, q)
-            J = pin.getFrameJacobian(
-                self.model, self.data, fid, pin.LOCAL_WORLD_ALIGNED)[:3, :]
-            J_leg = J[:, self._leg_v_indices[leg]]
-            dq = np.linalg.solve(J_leg.T @ J_leg + 1e-6 * np.eye(3), J_leg.T @ err)
-            for i, idx in enumerate(q_indices):
-                q[idx] += dq[i]
-
-        return np.array([q[idx] for idx in q_indices])
@@ -1,127 +0,0 @@
-"""GUI control panel for the wheeled-legged robot."""
-
-import tkinter as tk
-from tkinter import ttk
-
-
-class GUI:
-    """Tkinter control panel: sliders + gait buttons + status display."""
-
-    def __init__(self, controller):
-        self.ctrl = controller
-        self.root = tk.Tk()
-        self.root.title("WheelLeg Control")
-        self.root.geometry("400x500")
-        self.root.protocol("WM_DELETE_WINDOW", self._on_close)
-        self._closed = False
-
-        self._build()
-
-    def _build(self):
-        # Mode buttons
-        mf = ttk.LabelFrame(self.root, text="Mode")
-        mf.pack(fill="x", padx=8, pady=4)
-        for mode in ("wheel", "trot", "mpc"):
-            ttk.Button(mf, text=mode.upper(),
-                       command=lambda m=mode: self._set_mode(m)
-                       ).pack(side="left", padx=4, expand=True)
-        ttk.Button(mf, text="PRONE/STAND",
-                   command=self._toggle_prone).pack(side="left", padx=4, expand=True)
-
-        # Command sliders
-        cf = ttk.LabelFrame(self.root, text="Commands")
-        cf.pack(fill="x", padx=8, pady=4)
-
-        self.vel_x_var = tk.DoubleVar(value=0.0)
-        self.vel_y_var = tk.DoubleVar(value=0.0)
-        self.yaw_var = tk.DoubleVar(value=0.0)
-        self.height_var = tk.DoubleVar(value=self.ctrl.height)
-
-        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)
-
-        ttk.Label(cf, text="* Vel Y: trot mode only (diff-drive can't sidestep)",
-                  font=("", 8)).pack(anchor="w", padx=8)
-
-        ttk.Button(cf, text="Reset", command=self._reset).pack(pady=4)
-
-        # Status display
-        sf = ttk.LabelFrame(self.root, text="Status")
-        sf.pack(fill="both", expand=True, padx=8, pady=4)
-        self.status_text = tk.Text(sf, height=12, width=45, font=("Consolas", 9))
-        self.status_text.pack(fill="both", expand=True, padx=4, pady=4)
-
-    def _slider(self, parent, label, var, lo, hi):
-        f = ttk.Frame(parent)
-        f.pack(fill="x", padx=4, pady=2)
-        ttk.Label(f, text=label, width=7).pack(side="left")
-        ttk.Scale(f, from_=lo, to=hi, variable=var,
-                  command=lambda *_: self._sync()).pack(side="left", fill="x", expand=True)
-        lbl = ttk.Label(f, text="0.00", width=6)
-        lbl.pack(side="left")
-        var.trace_add("write", lambda *_, v=var, l=lbl: l.config(text=f"{v.get():.2f}"))
-
-    def _set_mode(self, mode):
-        self.ctrl.mode = mode
-        self.ctrl.prone = False
-
-    def _toggle_prone(self):
-        self.ctrl.prone = not self.ctrl.prone
-
-    def _sync(self):
-        self.ctrl.vel_x = self.vel_x_var.get()
-        self.ctrl.vel_y = self.vel_y_var.get()
-        self.ctrl.yaw_rate = self.yaw_var.get()
-        self.ctrl.height = self.height_var.get()
-
-    def _reset(self):
-        self.vel_x_var.set(0.0)
-        self.vel_y_var.set(0.0)
-        self.yaw_var.set(0.0)
-        self._sync()
-
-    def _on_close(self):
-        self._closed = True
-        self.root.destroy()
-
-    @property
-    def closed(self):
-        return self._closed
-
-    def update_status(self, state, step):
-        """Update status text with current robot state."""
-        txt = (
-            f"Mode: {self.ctrl.mode}  Step: {step}\n"
-            f"Pos:  x={state.pos[0]:.3f} y={state.pos[1]:.3f} z={state.pos[2]:.3f}\n"
-            f"RPY:  r={np.degrees(state.rpy[0]):.1f}° p={np.degrees(state.rpy[1]):.1f}° "
-            f"y={np.degrees(state.rpy[2]):.1f}°\n"
-            f"Vel:  vx={state.lin_vel[0]:.3f} vy={state.lin_vel[1]:.3f} vz={state.lin_vel[2]:.3f}\n"
-            f"Cmd:  vx={self.ctrl.vel_x:.2f} yaw={self.ctrl.yaw_rate:.2f} h={self.ctrl.height:.3f}\n"
-            f"─────────────────────────────────\n"
-        )
-        # Joint angles (compact)
-        for i, leg in enumerate(("FL", "FR", "RL", "RR")):
-            q = state.joint_pos[i*4:i*4+3]
-            w = state.joint_vel[i*4+3]
-            txt += f"{leg}: [{q[0]:+.2f} {q[1]:+.2f} {q[2]:+.2f}] w={w:+.1f}\n"
-
-        self.status_text.delete("1.0", tk.END)
-        self.status_text.insert(tk.END, txt)
-
-    def tick(self):
-        """Process GUI events. Returns False if window closed."""
-        if self._closed:
-            return False
-        try:
-            self.root.update_idletasks()
-            self.root.update()
-            return True
-        except tk.TclError:
-            self._closed = True
-            return False
-
-
-# Need numpy for degrees conversion in update_status
-import numpy as np
@@ -1,242 +0,0 @@
-"""Convex MPC solver for wheeled-legged robot.
-
-Centroidal dynamics: single rigid body model with 4 contact forces.
-State: x = [pos(3), rpy(3), vel(3), omega(3)] = 12
-Input: u = [f1(3), f2(3), f3(3), f4(3)] = 12
-Friction pyramid constraints on each foot.
-
-Reference: MIT Cheetah 3 Convex MPC (Di Carlo et al.)
-"""
-
-import numpy as np
-from scipy import sparse
-from scipy.linalg import block_diag
-import osqp
-
-from config import ROBOT_MASS, LEG_NAMES
-
-# MPC parameters
-MPC_HORIZON = 10       # prediction steps
-MPC_DT = 0.02         # 50 Hz MPC update
-MU = 0.6              # friction coefficient
-FZ_MAX = 200.0        # max vertical force per leg
-FZ_MIN = 10.0         # min vertical force (stance)
-NX = 12               # state dim
-NU = 12               # input dim (4 legs × 3D force)
-
-# Cost weights: [pos_x, pos_y, pos_z, roll, pitch, yaw, vx, vy, vz, wx, wy, wz]
-Q_WEIGHTS = np.array([2.0, 2.0, 50.0, 50.0, 50.0, 10.0, 2.0, 2.0, 1.0, 1.0, 1.0, 1.0])
-R_WEIGHTS = np.array([1e-6] * 12)
-
-
-def _skew(v):
-    return np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
-
-
-class ConvexMPC:
-    """Convex MPC: solves QP for optimal ground reaction forces."""
-
-    def __init__(self, mass=ROBOT_MASS, inertia=None):
-        self.mass = mass
-        # Approximate body inertia (diagonal, world-aligned)
-        if inertia is None:
-            self.I_body = np.diag([0.07, 0.26, 0.24])
-        else:
-            self.I_body = np.array(inertia).reshape(3, 3)
-
-        self.N = MPC_HORIZON
-        self.dt = MPC_DT
-        self.Q = np.diag(Q_WEIGHTS)
-        self.R = np.diag(R_WEIGHTS)
-        self.gravity = np.array([0, 0, -9.81])
-
-        self._last_forces = np.zeros(NU)
-
-    def solve(self, x0, x_ref, foot_positions, contact_schedule):
-        """Solve MPC QP.
-
-        Args:
-            x0: (12,) current state [pos, rpy, vel, omega]
-            x_ref: (12, N) reference trajectory over horizon
-            foot_positions: (4, 3) foot positions in world frame (relative to CoM)
-            contact_schedule: (4, N) binary contact table (1=stance)
-
-        Returns:
-            forces: (12,) optimal forces for current timestep [f1x,f1y,f1z,...,f4x,f4y,f4z]
-        """
-        N = self.N
-
-        # Build dynamics matrices
-        Ad, Bd_list, gd = self._discretize_dynamics(x0, foot_positions)
-
-        # Build QP: min 0.5 z'Hz + f'z  s.t. lb <= Az <= ub, lbx <= z <= ubx
-        # Decision variables: z = [x1,...,xN, u0,...,uN-1]
-        nvars = N * NX + N * NU
-
-        # --- Hessian ---
-        H_diag = np.concatenate([np.tile(2 * Q_WEIGHTS, N), np.tile(2 * R_WEIGHTS, N)])
-        H = sparse.diags(H_diag, format='csc')
-
-        # --- Gradient ---
-        g = np.zeros(nvars)
-        for k in range(N):
-            g[k*NX:(k+1)*NX] = -2 * self.Q @ x_ref[:, k]
-
-        # --- Dynamics equality constraints ---
-        # x_{k+1} = Ad @ x_k + Bd_k @ u_k + gd
-        # Rewrite: x_{k+1} - Ad @ x_k - Bd_k @ u_k = gd (for k>0)
-        #           x_1 - Bd_0 @ u_0 = Ad @ x0 + gd (for k=0)
-        n_eq = N * NX
-        A_eq = np.zeros((n_eq, nvars))
-        b_eq = np.zeros(n_eq)
-
-        # k=0: x_1 = Ad @ x0 + Bd_0 @ u_0 + gd
-        A_eq[0:NX, 0:NX] = np.eye(NX)  # x_1
-        A_eq[0:NX, N*NX:N*NX+NU] = -Bd_list[0]  # -Bd_0 @ u_0
-        b_eq[0:NX] = Ad @ x0 + gd
-
-        for k in range(1, N):
-            row = k * NX
-            # x_{k+1}
-            A_eq[row:row+NX, k*NX:(k+1)*NX] = np.eye(NX)
-            # -Ad @ x_k
-            A_eq[row:row+NX, (k-1)*NX:k*NX] = -Ad
-            # -Bd_k @ u_k
-            A_eq[row:row+NX, N*NX+k*NU:N*NX+(k+1)*NU] = -Bd_list[k]
-            b_eq[row:row+NX] = gd
-
-        # --- Friction pyramid inequality constraints ---
-        # For each stance leg at each timestep: 4 faces
-        # fx - mu*fz <= 0, -fx - mu*fz <= 0, fy - mu*fz <= 0, -fy - mu*fz <= 0
-        n_ineq_max = 4 * 4 * N
-        A_ineq = np.zeros((n_ineq_max, nvars))
-        u_ineq = np.zeros(n_ineq_max)
-
-        row = 0
-        for k in range(N):
-            u_base = N * NX + k * NU
-            for leg in range(4):
-                if contact_schedule[leg, k] == 1:
-                    fx_idx = u_base + leg * 3
-                    fy_idx = u_base + leg * 3 + 1
-                    fz_idx = u_base + leg * 3 + 2
-
-                    # Friction pyramid: stance leg
-                    A_ineq[row, fx_idx] = 1.0
-                    A_ineq[row, fz_idx] = -MU
-                    row += 1
-                    A_ineq[row, fx_idx] = -1.0
-                    A_ineq[row, fz_idx] = -MU
-                    row += 1
-                    A_ineq[row, fy_idx] = 1.0
-                    A_ineq[row, fz_idx] = -MU
-                    row += 1
-                    A_ineq[row, fy_idx] = -1.0
-                    A_ineq[row, fz_idx] = -MU
-                    row += 1
-
-        A_ineq = A_ineq[:row]
-        u_ineq = u_ineq[:row]
-
-        # Stack constraints
-        A_full = np.vstack([A_eq, A_ineq])
-        l_full = np.concatenate([b_eq, -np.inf * np.ones(row)])
-        u_full = np.concatenate([b_eq, u_ineq])
-
-        # --- Box constraints on forces (as identity rows in A) ---
-        A_box = np.zeros((N * NU, nvars))
-        l_box = -np.inf * np.ones(N * NU)
-        u_box = np.inf * np.ones(N * NU)
-
-        for k in range(N):
-            u_base = N * NX + k * NU
-            for leg in range(4):
-                idx = u_base + leg * 3
-                box_row = k * NU + leg * 3
-                # Identity rows for fx, fy, fz
-                for j in range(3):
-                    A_box[box_row + j, idx + j] = 1.0
-
-                if contact_schedule[leg, k] == 1:
-                    # Stance: fz bounded
-                    l_box[box_row + 2] = FZ_MIN
-                    u_box[box_row + 2] = FZ_MAX
-                else:
-                    # Swing: all forces = 0
-                    l_box[box_row:box_row+3] = 0.0
-                    u_box[box_row:box_row+3] = 0.0
-
-        # Final constraint matrix
-        A_full = np.vstack([A_full, A_box])
-        l_full = np.concatenate([l_full, l_box])
-        u_full = np.concatenate([u_full, u_box])
-
-        # --- Solve with OSQP ---
-        A_sparse = sparse.csc_matrix(A_full)
-        H_sparse = sparse.triu(H, format='csc')
-
-        solver = osqp.OSQP()
-        solver.setup(H_sparse, g, A_sparse, l_full, u_full,
-                     eps_abs=1e-4, eps_rel=1e-4,
-                     max_iter=500, polish=True, verbose=False,
-                     warm_start=True)
-
-        # Warm start with previous solution
-        if self._last_forces is not None:
-            x_warm = np.zeros(nvars)
-            x_warm[N*NX:N*NX+NU] = self._last_forces
-            solver.warm_start(x=x_warm)
-
-        result = solver.solve()
-
-        if result.info.status == 'solved' or result.info.status == 'solved_inaccurate':
-            # Extract first timestep forces
-            forces = result.x[N*NX:N*NX+NU]
-            self._last_forces = forces.copy()
-        else:
-            forces = self._last_forces
-
-        return forces
-
-    def _discretize_dynamics(self, x0, foot_positions):
-        """Build discrete-time centroidal dynamics.
-
-        State: [pos, rpy, vel, omega] (12)
-        Continuous: dx/dt = Ac @ x + Bc @ u + gc
-        Discrete: x_{k+1} = Ad @ x + Bd @ u + gd
-        """
-        m = self.mass
-        I_inv = np.linalg.inv(self.I_body)
-        dt = self.dt
-        yaw = x0[5]
-        cy, sy = np.cos(yaw), np.sin(yaw)
-
-        # Rotation for rpy rate ≈ R_z^T @ omega
-        R_zT = np.array([[cy, sy, 0], [-sy, cy, 0], [0, 0, 1]])
-
-        # Ac (12×12)
-        Ac = np.zeros((NX, NX))
-        Ac[0:3, 6:9] = np.eye(3)       # pos_dot = vel
-        Ac[3:6, 9:12] = R_zT           # rpy_dot ≈ R_z^T @ omega
-
-        # Ad = I + Ac*dt (first-order)
-        Ad = np.eye(NX) + Ac * dt
-
-        # Bc varies per timestep (foot positions change contact point)
-        Bd_list = []
-        for k in range(self.N):
-            Bc = np.zeros((NX, NU))
-            for leg in range(4):
-                r = foot_positions[leg]
-                # vel_dot += f/m
-                Bc[6:9, leg*3:(leg+1)*3] = np.eye(3) / m
-                # omega_dot += I^{-1} @ (r × f)
-                Bc[9:12, leg*3:(leg+1)*3] = I_inv @ _skew(r)
-            Bd = Bc * dt
-            Bd_list.append(Bd)
-
-        # Gravity contribution
-        gd = np.zeros(NX)
-        gd[6:9] = self.gravity * dt  # vel += g*dt
-
-        return Ad, Bd_list, gd
@@ -1,261 +0,0 @@
-"""MPC controller integration for wheeled-legged robot.
-
-Integrates: gait scheduler + reference trajectory + ConvexMPC solver +
-            swing leg control + stance force mapping + wheel drive.
-
-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)
-  - Wheel drive: stance legs use differential drive, swing legs coast
-"""
-
-import numpy as np
-from robot import Robot, RobotState
-from dynamics import Dynamics
-from mpc import ConvexMPC, MPC_DT
-from config import (
-    LEG_NAMES, WHEEL_RADIUS, WHEEL_TRACK, WHEEL_VEL_MAX,
-    CTRL_DT, GAIT_FREQ, GAIT_DUTY, SWING_HEIGHT, PHASE_OFFSETS,
-    DEFAULT_JOINT_ANGLES, ROBOT_MASS,
-)
-
-# MPC update decimation (relative to control loop)
-MPC_DECIMATION = max(1, int(MPC_DT / CTRL_DT))  # ~5 steps at 250Hz
-
-
-class MPCController:
-    """Convex MPC locomotion controller for wheeled-legged robot."""
-
-    def __init__(self, robot: Robot):
-        self.robot = robot
-        self.dynamics = Dynamics()
-        self.mpc = ConvexMPC(mass=ROBOT_MASS)
-
-        # User commands
-        self.vel_x = 0.0
-        self.vel_y = 0.0
-        self.yaw_rate = 0.0
-        self.height = 0.35  # actual standing height with default joint angles
-
-        # Gait state - start at phase 0 with all legs in stance (duty=0.6)
-        self._gait_phase = 0.0
-        self._step_count = 0
-        self._initialized = False
-
-        # MPC solution cache - initialize with gravity compensation
-        self._mpc_forces = np.zeros(12)
-        self._init_gravity_comp()
-
-        # Swing trajectory state
-        self._swing_start_foot = {leg: np.zeros(3) for leg in LEG_NAMES}
-        self._swing_start_time = {leg: 0.0 for leg in LEG_NAMES}
-        self._last_contact = {leg: True for leg in LEG_NAMES}
-
-        # Smoothed commands
-        self._smooth_vx = 0.0
-        self._smooth_vy = 0.0
-        self._smooth_yaw = 0.0
-
-    def _init_gravity_comp(self):
-        """Pre-fill MPC forces with static gravity compensation."""
-        fz_per_leg = ROBOT_MASS * 9.81 / 4.0
-        for i in range(4):
-            self._mpc_forces[i*3 + 2] = fz_per_leg
-
-    def compute(self, state: RobotState, dt: float):
-        """Main MPC control loop.
-
-        Uses MIT motor protocol: tau = kp*(q_des-q) + kd*(dq_des-dq) + tau_ff
-        where tau_ff comes from MPC force mapping via Jacobian transpose.
-
-        Returns:
-            tau_legs: (12,) feedforward torques for MIT mode
-            wheel_targets: (4,) wheel velocity targets
-        """
-        # Smooth commands
-        alpha = min(dt * 3.0, 1.0)
-        self._smooth_vx += alpha * (self.vel_x - self._smooth_vx)
-        self._smooth_vy += alpha * (self.vel_y - self._smooth_vy)
-        self._smooth_yaw += alpha * (self.yaw_rate - self._smooth_yaw)
-
-        # Update Pinocchio
-        q_pin, dq_pin = self.robot.get_qpos_qvel_for_pinocchio()
-        self.dynamics.update(q_pin, dq_pin)
-
-        # Initialize foot positions on first call
-        if not self._initialized:
-            for leg in LEG_NAMES:
-                self._swing_start_foot[leg] = self.dynamics.get_foot_pos(leg)
-            self._initialized = True
-
-        # Decide if we should trot or just stand
-        moving = (abs(self._smooth_vx) > 0.02 or
-                  abs(self._smooth_vy) > 0.02 or
-                  abs(self._smooth_yaw) > 0.05)
-
-        if moving:
-            self._gait_phase = (self._gait_phase + dt * GAIT_FREQ) % 1.0
-        else:
-            self._gait_phase = 0.0  # all legs in stance
-
-        # Contact schedule
-        contacts = {}
-        for leg in LEG_NAMES:
-            phase = (self._gait_phase + PHASE_OFFSETS[leg]) % 1.0
-            contacts[leg] = phase < GAIT_DUTY
-
-        # Get foot positions relative to CoM
-        foot_positions = np.zeros((4, 3))
-        for i, leg in enumerate(LEG_NAMES):
-            foot_positions[i] = self.dynamics.get_foot_pos(leg) - state.pos
-
-        # --- Run MPC at lower rate ---
-        if self._step_count % MPC_DECIMATION == 0:
-            x0 = self._build_state_vector(state)
-            x_ref = self._build_reference(state)
-            contact_table = self._build_contact_table()
-            self._mpc_forces = self.mpc.solve(x0, x_ref, foot_positions, contact_table)
-
-        self._step_count += 1
-
-        # --- Compute feedforward torques and desired joint positions ---
-        tau_ff = np.zeros(12)
-        q_des = np.zeros(12)
-        dq_des = np.zeros(12)
-        kp = np.zeros(12)
-        kd = np.zeros(12)
-        wheel_targets = np.zeros(4)
-
-        for i, leg in enumerate(LEG_NAMES):
-            if contacts[leg]:
-                # Stance: MPC force → feedforward torque, PD holds posture
-                f_leg = self._mpc_forces[i*3:(i+1)*3]
-                J = self.dynamics.get_foot_jacobian_leg(leg)
-                tau_ff[i*3:(i+1)*3] = J.T @ (-f_leg)
-
-                # PD target: default standing angles (posture hold)
-                q_des[i*3] = DEFAULT_JOINT_ANGLES["hip_abduction"]
-                q_des[i*3+1] = DEFAULT_JOINT_ANGLES["hip_pitch"]
-                q_des[i*3+2] = DEFAULT_JOINT_ANGLES["knee"]
-                kp[i*3:(i+1)*3] = [40.0, 40.0, 40.0]
-                kd[i*3:(i+1)*3] = [3.0, 3.0, 3.0]
-
-                # Record foot position
-                self._swing_start_foot[leg] = self.dynamics.get_foot_pos(leg)
-                self._last_contact[leg] = True
-
-                # Wheel drive
-                wheel_targets[i] = self._wheel_cmd(leg)
-            else:
-                # Swing: IK target position, strong PD, no feedforward
-                if self._last_contact[leg]:
-                    self._swing_start_foot[leg] = self.dynamics.get_foot_pos(leg)
-                    self._swing_start_time[leg] = state.time
-                    self._last_contact[leg] = False
-
-                q_ik = self._swing_leg_ik(leg, state, q_pin)
-                q_des[i*3:(i+1)*3] = q_ik
-                kp[i*3:(i+1)*3] = [60.0, 60.0, 60.0]  # strong PD for swing
-                kd[i*3:(i+1)*3] = [3.0, 3.0, 3.0]
-                # tau_ff stays 0 for swing
-
-                wheel_targets[i] = 0.0
-
-        # Use MIT protocol via robot interface
-        self.robot.set_ctrl_mit(q_des, dq_des, kp, kd, tau_ff, wheel_targets)
-        # Return dummy (actual ctrl is set directly above)
-        return None, None
-
-    def _build_state_vector(self, state: RobotState):
-        """Build MPC state: [pos, rpy, vel, omega]."""
-        return np.concatenate([state.pos, state.rpy, state.lin_vel, state.ang_vel])
-
-    def _build_reference(self, state: RobotState):
-        """Build reference trajectory over MPC horizon."""
-        N = self.mpc.N
-        x_ref = np.zeros((12, N))
-
-        yaw = state.rpy[2]
-        cy, sy = np.cos(yaw), np.sin(yaw)
-        R_z = np.array([[cy, -sy, 0], [sy, cy, 0], [0, 0, 1]])
-        vel_world = R_z @ np.array([self._smooth_vx, self._smooth_vy, 0.0])
-
-        for k in range(N):
-            t = (k + 1) * self.mpc.dt
-            # Position: integrate from current
-            x_ref[0, k] = state.pos[0] + vel_world[0] * t
-            x_ref[1, k] = state.pos[1] + vel_world[1] * t
-            x_ref[2, k] = self.height
-            # RPY: keep roll/pitch zero, integrate yaw
-            x_ref[3, k] = 0.0
-            x_ref[4, k] = 0.0
-            x_ref[5, k] = yaw + self._smooth_yaw * t
-            # Velocity
-            x_ref[6, k] = vel_world[0]
-            x_ref[7, k] = vel_world[1]
-            x_ref[8, k] = 0.0
-            # Angular velocity
-            x_ref[9, k] = 0.0
-            x_ref[10, k] = 0.0
-            x_ref[11, k] = self._smooth_yaw
-
-        return x_ref
-
-    def _build_contact_table(self):
-        """Build contact schedule over MPC horizon."""
-        N = self.mpc.N
-        table = np.zeros((4, N), dtype=int)
-        for k in range(N):
-            future_phase = (self._gait_phase + (k + 1) * self.mpc.dt * GAIT_FREQ) % 1.0
-            for i, leg in enumerate(LEG_NAMES):
-                leg_phase = (future_phase + PHASE_OFFSETS[leg]) % 1.0
-                table[i, k] = 1 if leg_phase < GAIT_DUTY else 0
-        return table
-
-    def _swing_leg_ik(self, leg: str, state: RobotState, q_pin: np.ndarray):
-        """Swing leg: compute IK target joint angles for trajectory."""
-        swing_phase = self._get_swing_phase(leg)
-
-        p_start = self._swing_start_foot[leg]
-        p_end = self._compute_touchdown(leg, state)
-
-        s = swing_phase
-        s_mj = 10*s**3 - 15*s**4 + 6*s**5
-
-        pos_des = p_start + (p_end - p_start) * s_mj
-        # Z lift
-        z_lift = 64.0 * s**3 * (1.0 - s)**3
-        pos_des[2] = p_start[2] + SWING_HEIGHT * z_lift
-
-        # IK to get joint angles
-        q_ik = self.dynamics.inverse_kinematics(leg, pos_des, q_pin)
-        return q_ik
-
-    def _get_swing_phase(self, leg: str) -> float:
-        phase = (self._gait_phase + PHASE_OFFSETS[leg]) % 1.0
-        if phase < GAIT_DUTY:
-            return 0.0
-        return (phase - GAIT_DUTY) / (1.0 - GAIT_DUTY)
-
-    def _compute_touchdown(self, leg: str, state: RobotState) -> np.ndarray:
-        """Raibert heuristic for touchdown position."""
-        td = self._swing_start_foot[leg].copy()
-        t_stance = GAIT_DUTY / GAIT_FREQ
-
-        yaw = state.rpy[2]
-        cy, sy = np.cos(yaw), np.sin(yaw)
-        R_z = np.array([[cy, -sy, 0], [sy, cy, 0], [0, 0, 1]])
-        cmd_vel_world = R_z @ np.array([self._smooth_vx, self._smooth_vy, 0.0])
-
-        td[0] += cmd_vel_world[0] * t_stance * 0.5
-        td[1] += cmd_vel_world[1] * t_stance * 0.5
-        td[2] = WHEEL_RADIUS
-        return td
-
-    def _wheel_cmd(self, leg: str) -> float:
-        """Differential drive for a single wheel."""
-        if leg[1] == "l":
-            v = (self._smooth_vx - 0.5 * WHEEL_TRACK * self._smooth_yaw) / WHEEL_RADIUS
-        else:
-            v = (self._smooth_vx + 0.5 * WHEEL_TRACK * self._smooth_yaw) / WHEEL_RADIUS
-        return np.clip(v, -WHEEL_VEL_MAX, WHEEL_VEL_MAX)
@@ -1,229 +0,0 @@
-"""MuJoCo interface for the wheeled-legged robot.
-
-Configures actuators as proper PD controllers at runtime:
-  - Leg joints: force = kp*(ctrl - qpos) - kd*qvel, ctrl = target angle
-  - Wheel joints: force = gain*(ctrl - qvel), ctrl = target velocity (rad/s)
-"""
-
-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)
-
-
-@dataclass
-class RobotState:
-    """Robot state from MuJoCo."""
-    pos: np.ndarray       # (3,) world position
-    quat: np.ndarray      # (4,) quaternion (w,x,y,z) MuJoCo convention
-    rot: np.ndarray       # (3,3) body→world rotation
-    rpy: np.ndarray       # (3,) roll, pitch, yaw
-    lin_vel: np.ndarray   # (3,) world frame linear velocity
-    ang_vel: np.ndarray   # (3,) body frame angular velocity
-    joint_pos: np.ndarray # (16,) all joint positions [fl3+wheel, fr3+wheel, rl3+wheel, rr3+wheel]
-    joint_vel: np.ndarray # (16,) all joint velocities
-    time: float
-
-
-class Robot:
-    """MuJoCo simulation interface with proper PD actuator configuration."""
-
-    # Leg PD gains (tuned for 12.3kg robot)
-    LEG_KP = 60.0
-    LEG_KD = 3.0
-    # Wheel velocity gain
-    WHEEL_KP = 2.0
-
-    def __init__(self, xml_path=None):
-        self.model = mujoco.MjModel.from_xml_path(str(xml_path or SCENE_XML))
-        self.data = mujoco.MjData(self.model)
-
-        # Cache IDs
-        self._base_bid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "base_link")
-        self._actuator_ids = {}  # name → actuator index
-        self._joint_qpos_adr = {}  # name → qpos address
-        self._joint_qvel_adr = {}  # name → qvel address
-
-        # Build joint/actuator maps
-        self._ctrl_order = []
-        for leg in LEG_NAMES:
-            for jt in (*LEG_JOINTS, WHEEL_JOINT):
-                name = f"{leg}_{jt}"
-                aid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, name)
-                jid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, name)
-                self._actuator_ids[name] = aid
-                self._joint_qpos_adr[name] = self.model.jnt_qposadr[jid]
-                self._joint_qvel_adr[name] = self.model.jnt_dofadr[jid]
-                self._ctrl_order.append(name)
-
-        # Configure actuators as proper PD controllers
-        self._configure_actuators()
-
-    def _configure_actuators(self):
-        """Set actuators to proper PD mode.
-
-        Leg joints: force = kp*(ctrl - qpos) - kd*qvel
-        Wheels: force = gain*(ctrl - qvel)  (velocity tracking)
-        """
-        for i in range(self.model.nu):
-            name = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, i)
-            self.model.actuator_biastype[i] = 1  # affine bias
-            self.model.actuator_gaintype[i] = 0  # fixed gain
-            self.model.actuator_forcelimited[i] = 0  # no force clamp (17Nm is in actuatorfrcrange)
-
-            if 'wheel' not in name:
-                self.model.actuator_gainprm[i, 0] = self.LEG_KP
-                self.model.actuator_biasprm[i, 0] = 0.0
-                self.model.actuator_biasprm[i, 1] = -self.LEG_KP
-                self.model.actuator_biasprm[i, 2] = -self.LEG_KD
-                self.model.actuator_ctrlrange[i] = [-3.14, 3.14]
-            else:
-                self.model.actuator_gainprm[i, 0] = self.WHEEL_KP
-                self.model.actuator_biasprm[i, 0] = 0.0
-                self.model.actuator_biasprm[i, 1] = 0.0
-                self.model.actuator_biasprm[i, 2] = -self.WHEEL_KP
-                self.model.actuator_ctrlrange[i] = [-20.0, 20.0]
-
-    @property
-    def dt(self):
-        return self.model.opt.timestep
-
-    def reset(self):
-        """Reset to standing pose at correct height for default joint angles."""
-        mujoco.mj_resetData(self.model, self.data)
-
-        # Set default leg angles
-        for leg in LEG_NAMES:
-            for jt, key in zip(LEG_JOINTS, ("hip_abduction", "hip_pitch", "knee")):
-                name = f"{leg}_{jt}"
-                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
-        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
-        self.data.qpos[2] = base_z
-        self.data.qpos[3] = 1.0  # quat w
-
-        mujoco.mj_forward(self.model, self.data)
-
-        # Set ctrl to match initial pose (so PD doesn't jerk)
-        for leg in LEG_NAMES:
-            for jt, key in zip(LEG_JOINTS, ("hip_abduction", "hip_pitch", "knee")):
-                name = f"{leg}_{jt}"
-                self.data.ctrl[self._actuator_ids[name]] = DEFAULT_JOINT_ANGLES[key]
-            # Wheels: zero velocity
-            self.data.ctrl[self._actuator_ids[f"{leg}_{WHEEL_JOINT}"]] = 0.0
-
-    def get_state(self) -> RobotState:
-        """Extract robot state."""
-        pos = self.data.xpos[self._base_bid].copy()
-        quat = self.data.xquat[self._base_bid].copy()  # (w,x,y,z)
-        rot = self.data.xmat[self._base_bid].reshape(3, 3).copy()
-
-        rpy = np.array([
-            np.arctan2(rot[2, 1], rot[2, 2]),
-            np.arctan2(-rot[2, 0], np.sqrt(rot[2, 1]**2 + rot[2, 2]**2)),
-            np.arctan2(rot[1, 0], rot[0, 0]),
-        ])
-
-        # Base velocity (world frame)
-        lin_vel = self.data.qvel[0:3].copy()
-        ang_vel = self.data.qvel[3:6].copy()
-
-        # Joint states (16 joints: 4 legs × 4 joints each)
-        joint_pos = np.zeros(16)
-        joint_vel = np.zeros(16)
-        for i, name in enumerate(self._ctrl_order):
-            joint_pos[i] = self.data.qpos[self._joint_qpos_adr[name]]
-            joint_vel[i] = self.data.qvel[self._joint_qvel_adr[name]]
-
-        return RobotState(
-            pos=pos, quat=quat, rot=rot, rpy=rpy,
-            lin_vel=lin_vel, ang_vel=ang_vel,
-            joint_pos=joint_pos, joint_vel=joint_vel,
-            time=self.data.time,
-        )
-
-    def set_ctrl(self, leg_targets: np.ndarray, wheel_targets: np.ndarray):
-        """Set actuator commands (position PD mode).
-
-        Args:
-            leg_targets: (12,) target joint angles for legs [fl3, fr3, rl3, rr3]
-            wheel_targets: (4,) target wheel velocities [fl, fr, rl, rr] in rad/s
-        """
-        for i, leg in enumerate(LEG_NAMES):
-            for j, jt in enumerate(LEG_JOINTS):
-                name = f"{leg}_{jt}"
-                self.data.ctrl[self._actuator_ids[name]] = leg_targets[i * 3 + j]
-            name = f"{leg}_{WHEEL_JOINT}"
-            self.data.ctrl[self._actuator_ids[name]] = wheel_targets[i]
-
-    def set_ctrl_mit(self, q_des: np.ndarray, dq_des: np.ndarray,
-                     kp: np.ndarray, kd: np.ndarray, tau_ff: np.ndarray,
-                     wheel_targets: np.ndarray):
-        """MIT motor protocol: tau = kp*(q_des-q) + kd*(dq_des-dq) + tau_ff.
-
-        Computes torque in software, sends to actuators in torque mode.
-        Call enable_torque_mode() first.
-
-        Args:
-            q_des: (12,) desired joint angles
-            dq_des: (12,) desired joint velocities
-            kp: (12,) position gains (0 for pure torque)
-            kd: (12,) velocity gains
-            tau_ff: (12,) feedforward torques
-            wheel_targets: (4,) wheel velocity targets
-        """
-        for i, leg in enumerate(LEG_NAMES):
-            for j, jt in enumerate(LEG_JOINTS):
-                name = f"{leg}_{jt}"
-                aid = self._actuator_ids[name]
-                idx = i * 3 + j
-                q = self.data.qpos[self._joint_qpos_adr[name]]
-                dq = self.data.qvel[self._joint_qvel_adr[name]]
-                tau = (kp[idx] * (q_des[idx] - q)
-                       + kd[idx] * (dq_des[idx] - dq)
-                       + tau_ff[idx])
-                self.data.ctrl[aid] = np.clip(tau, -17.0, 17.0)
-            name = f"{leg}_{WHEEL_JOINT}"
-            self.data.ctrl[self._actuator_ids[name]] = wheel_targets[i]
-
-    def enable_torque_mode(self):
-        """Switch leg actuators to direct torque mode (for MPC/MIT)."""
-        for i in range(self.model.nu):
-            name = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, i)
-            if 'wheel' not in name:
-                self.model.actuator_gainprm[i, 0] = 1.0
-                self.model.actuator_biasprm[i, :3] = [0, 0, 0]
-                self.model.actuator_biastype[i] = 0
-                self.model.actuator_ctrlrange[i] = [-17.0, 17.0]
-
-    def enable_position_mode(self):
-        """Switch leg actuators back to position PD mode."""
-        self._configure_actuators()
-
-    def step(self):
-        """Advance one simulation timestep."""
-        mujoco.mj_step(self.model, self.data)
-
-    def get_qpos_qvel_for_pinocchio(self):
-        """Get full qpos/qvel for Pinocchio (reorder quaternion)."""
-        qpos = self.data.qpos.copy()
-        qvel = self.data.qvel.copy()
-        # MuJoCo quat: (w,x,y,z) → Pinocchio: (x,y,z,w)
-        w, x, y, z = qpos[3], qpos[4], qpos[5], qpos[6]
-        q_pin = np.concatenate([qpos[0:3], [x, y, z, w], qpos[7:]])
-        # MuJoCo vel is already [lin_world(3), ang_body(3), joints(16)]
-        # Pinocchio wants [lin_body(3), ang_body(3), joints(16)]
-        from scipy.spatial.transform import Rotation
-        R = Rotation.from_quat([x, y, z, w]).as_matrix()
-        v_body = R.T @ qvel[0:3]
-        dq_pin = np.concatenate([v_body, qvel[3:]])
-        return q_pin, dq_pin
@@ -1,111 +0,0 @@
-"""Main entry point: wheeled-legged robot simulation.
-
-Controls:
-  Mode: wheel (default) - differential drive + posture hold
-        trot - quadruped gait with wheel assist
-        mpc  - convex MPC locomotion (torque control)
-
-  Keyboard (in MuJoCo viewer):
-    W/S: vel_x ±0.1
-    A/D: yaw_rate ±0.2
-    Q/E: height ±0.02
-    1: wheel mode
-    2: trot mode
-    3: MPC mode
-    4: prone toggle
-    Z: reset commands
-"""
-
-import time
-import numpy as np
-import mujoco.viewer as mjv
-
-from robot import Robot
-from controller import Controller
-from gui import GUI
-from config import CTRL_DECIMATION
-
-
-def main():
-    robot = Robot()
-    robot.reset()
-    ctrl = Controller(robot)
-    gui = GUI(ctrl)
-
-    step = 0
-    sim_steps_per_ctrl = CTRL_DECIMATION
-
-    def key_callback(keycode):
-        """Called from MuJoCo render thread - only modify ctrl directly, not tkinter."""
-        try:
-            c = chr(keycode).lower()
-        except (ValueError, OverflowError):
-            return
-        if c == 'w':
-            ctrl.vel_x = min(ctrl.vel_x + 0.1, 1.5)
-        elif c == 's':
-            ctrl.vel_x = max(ctrl.vel_x - 0.1, -1.5)
-        elif c == 'a':
-            ctrl.yaw_rate = min(ctrl.yaw_rate + 0.2, 2.0)
-        elif c == 'd':
-            ctrl.yaw_rate = max(ctrl.yaw_rate - 0.2, -2.0)
-        elif c == 'q':
-            ctrl.height = min(ctrl.height + 0.02, 0.45)
-        elif c == 'e':
-            ctrl.height = max(ctrl.height - 0.02, 0.15)
-        elif c == '1':
-            ctrl.mode = "wheel"; ctrl.prone = False
-        elif c == '2':
-            ctrl.mode = "trot"; ctrl.prone = False
-        elif c == '3':
-            ctrl.mode = "mpc"; ctrl.prone = False
-        elif c == '4':
-            ctrl.prone = not ctrl.prone
-        elif c == 'z':
-            ctrl.vel_x = 0.0; ctrl.vel_y = 0.0; ctrl.yaw_rate = 0.0
-
-    with mjv.launch_passive(robot.model, robot.data, key_callback=key_callback) as viewer:
-        viewer.cam.distance = 2.5
-        viewer.cam.elevation = -20
-        viewer.cam.azimuth = 135
-
-        last_time = robot.data.time
-
-        while viewer.is_running() and not gui.closed:
-            t_start = time.perf_counter()
-
-            # Detect viewer reset (Backspace) - time jumps back to 0
-            if robot.data.time < last_time:
-                robot.reset()
-            last_time = robot.data.time
-
-            # Get state and compute control
-            state = robot.get_state()
-            leg_targets, wheel_targets = ctrl.compute(state, robot.dt * sim_steps_per_ctrl)
-
-            # Apply control and step simulation
-            # MPC mode sets ctrl directly via set_ctrl_mit, skip set_ctrl
-            if ctrl.mode != "mpc":
-                robot.set_ctrl(leg_targets, wheel_targets)
-            for _ in range(sim_steps_per_ctrl):
-                robot.step()
-
-            viewer.sync()
-            step += 1
-
-            # Update GUI every 25 steps (~10 Hz)
-            if step % 25 == 0:
-                state = robot.get_state()
-                gui.update_status(state, step)
-                if not gui.tick():
-                    break
-
-            # Real-time sync
-            elapsed = time.perf_counter() - t_start
-            target_dt = robot.dt * sim_steps_per_ctrl
-            if elapsed < target_dt:
-                time.sleep(target_dt - elapsed)
-
-
-if __name__ == "__main__":
-    main()
@@ -25,7 +25,7 @@ class MuJoCoIO:
        print("[MuJoCoIO] Loading MuJoCo model...")
        spec = mujoco.MjSpec.from_file(str(temp_xml))
        
-        # Override actuators to match mjlab exactly
+        # Override actuators to match training/runtime actuator semantics.
        self._rebuild_actuators(spec)
        
        self.m = spec.compile()
@@ -91,8 +91,10 @@ class MuJoCoIO:
        for act in actuators_to_delete:
            spec.delete(act)
        
-        KP_LEG, KD_LEG = 40.0, 1.0
-        KD_WHEEL = 0.5
+        # Keep sim2sim aligned with the training robot config and sim2real runtime:
+        # leg position PD = (50.0, 1.5), wheel velocity damping = 1.0.
+        KP_LEG, KD_LEG = 50.0, 1.5
+        KD_WHEEL = 1.0
        EFFORT_LIMIT = 17.0
        
        leg_jnames = [
@@ -77,6 +77,7 @@ class PolicyRunner:

        # Load both policy networks
        self.policies = {}
+        self.policy_obs_dims = {}
        for name, path in self.policy_paths.items():
            print(f"[PolicyRunner] Loading {name} policy from: {path}")
            if Path(path).exists():
@@ -84,6 +85,7 @@ class PolicyRunner:
            else:
                print(f"[PolicyRunner] WARNING: {name} policy file not found! Falling back to rough.")
                self.policies[name] = load_policy(self.policy_paths["rough"], device)
+            self.policy_obs_dims[name] = int(self.policies[name].obs_mean.numel())

        # Default DOF positions for each policy
        self.default_dof_poses = {
@@ -226,11 +228,23 @@ class PolicyRunner:

        raw_actions_out = {}
        for name in active_policies:
-            # Flatten observation history
-            obs_history_array = np.array(self.obs_histories[name])
-            term_dims = [3, 3, 3, 12, 12, 4, 16]
-            term_histories = np.split(obs_history_array, np.cumsum(term_dims)[:-1], axis=1)
-            flat_obs = np.concatenate([h.flatten() for h in term_histories])
+            expected_obs_dim = self.policy_obs_dims[name]
+
+            if expected_obs_dim == current_obs_53d.shape[0]:
+                # Newer policies consume the current 53D observation directly.
+                flat_obs = self.obs_histories[name][-1]
+            elif expected_obs_dim == current_obs_53d.shape[0] * self.history_length:
+                # Legacy policies expect 6-step history stacking grouped by term.
+                obs_history_array = np.array(self.obs_histories[name])
+                term_dims = [3, 3, 3, 12, 12, 4, 16]
+                term_histories = np.split(obs_history_array, np.cumsum(term_dims)[:-1], axis=1)
+                flat_obs = np.concatenate([h.flatten() for h in term_histories])
+            else:
+                raise RuntimeError(
+                    f"Policy '{name}' expects obs dim {expected_obs_dim}, "
+                    f"but sim2sim can only provide {current_obs_53d.shape[0]} or "
+                    f"{current_obs_53d.shape[0] * self.history_length}."
+                )
            
            obs_tensor = torch.tensor(flat_obs, device=self.device, dtype=torch.float32).unsqueeze(0)
            
@@ -74,7 +74,11 @@ from ..mdp.rewards import (
    joint_mirror,
    feet_contact_without_cmd,
    upright_roll_only,
-    pitch_control_penalty,
+    upward,
+    joint_power,
+    ang_vel_xy_l2,
+    undesired_contacts,
+    contact_forces,
 )
 from ..mdp.curriculums import terrain_levels_vel_strict
 from ..mdp.commands import UniformThresholdVelocityCommandCfg
@@ -149,7 +153,7 @@ def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
        ),
        "projected_gravity": ObservationTermCfg(
            func=velocity_mdp.projected_gravity,
-            noise=Unoise(n_min=-0.05, n_max=0.05),
+            noise=Unoise(n_min=-0.08, n_max=0.08),
        ),
        "command": ObservationTermCfg(
            func=velocity_mdp.generated_commands,
@@ -174,7 +178,7 @@ def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
            params={"asset_cfg": SceneEntityCfg("robot", joint_names=(".*_wheel_joint",))},
            scale=0.05, noise=Unoise(n_min=-1.0, n_max=1.0),
        ),
-        "actions": ObservationTermCfg(func=velocity_mdp.last_action, history_length=1),
+        "actions": ObservationTermCfg(func=velocity_mdp.last_action),
    }

    critic_terms = {
@@ -192,7 +196,7 @@ def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
    observations = {
        "actor": ObservationGroupCfg(
            terms=actor_terms, concatenate_terms=True,
-            enable_corruption=True, history_length=6,
+            enable_corruption=True,
        ),
        "critic": ObservationGroupCfg(
            terms=critic_terms, concatenate_terms=True, enable_corruption=False,
@@ -208,13 +212,13 @@ def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
            actuator_names=(".*_hip_abduction_joint", ".*_hip_pitch_joint", ".*_knee_joint"),
            scale={".*_hip_abduction_joint": 0.125, "^(?!.*_hip_abduction_joint).*": 0.25}, use_default_offset=True,
            control_frequency=50.0, cut_off_frequency=5.0,
-            min_delay=0, max_delay=2,
+            min_delay=0, max_delay=4,
        ),
        "wheel_joint_vel": JointVelocityDelayedLowPassActionCfg(
            entity_name="robot", actuator_names=(".*_wheel_joint",),
            scale=5.0, offset=0.0, use_default_offset=False,
            control_frequency=50.0, cut_off_frequency=15.0,
-            min_delay=0, max_delay=2,
+            min_delay=0, max_delay=4,
        ),
    }

@@ -239,8 +243,15 @@ def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
        "reset_base": EventTermCfg(
            func=envs_mdp.reset_root_state_uniform, mode="reset",
            params={
-                "pose_range": {"z": (0.30, 0.50), "yaw": (-math.pi, math.pi)},
-                "velocity_range": {"x": (-0.5, 0.5), "y": (-0.15, 0.15), "yaw": (-0.35, 0.35)},
+                "pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-math.pi, math.pi)},
+                "velocity_range": {
+                    "x": (-0.5, 0.5),
+                    "y": (-0.5, 0.5),
+                    "z": (-0.5, 0.5),
+                    "roll": (-0.5, 0.5),
+                    "pitch": (-0.5, 0.5),
+                    "yaw": (-0.5, 0.5),
+                },
                "asset_cfg": SceneEntityCfg("robot"),
            },
        ),
@@ -251,35 +262,22 @@ def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
        ),
        "base_com": EventTermCfg(
            func=envs_dr.body_com_offset, mode="startup",
-            params={"asset_cfg": SceneEntityCfg("robot", body_names=("base_link",)),
+            params={"asset_cfg": SceneEntityCfg("robot", body_names=(".*",)),
                    "operation": "add", "ranges": {0: (-0.05, 0.05), 1: (-0.05, 0.05), 2: (-0.05, 0.05)}},
        ),
-        "encoder_bias": EventTermCfg(
-            func=envs_dr.encoder_bias, mode="startup",
-            params={"asset_cfg": SceneEntityCfg("robot", joint_names=(".*",)), "bias_range": (-0.015, 0.015)},
-        ),
        "body_friction": EventTermCfg(
            func=envs_dr.geom_friction, mode="startup",
-            params={"asset_cfg": SceneEntityCfg("robot", geom_names=(".*",)), "operation": "abs", "ranges": (0.3, 1.2)},
+            params={"asset_cfg": SceneEntityCfg("robot", geom_names=(".*",)), "operation": "abs", "ranges": (0.15, 1.25)},
        ),
        "actuator_stiffness": EventTermCfg(
            func=envs_dr.joint_stiffness, mode="startup",
-            params={"asset_cfg": SceneEntityCfg("robot"), "ranges": (0.9, 1.1), "operation": "scale", "distribution": "log_uniform"},
+            params={"asset_cfg": SceneEntityCfg("robot"), "ranges": (0.5, 1.5), "operation": "scale", "distribution": "log_uniform"},
        ),
        "actuator_damping": EventTermCfg(
            func=envs_dr.joint_damping, mode="startup",
-            params={"asset_cfg": SceneEntityCfg("robot"), "ranges": (0.9, 1.1), "operation": "scale", "distribution": "log_uniform"},
+            params={"asset_cfg": SceneEntityCfg("robot"), "ranges": (0.5, 1.5), "operation": "scale", "distribution": "log_uniform"},
        ),
-        "actuator_effort_limit": EventTermCfg(
-            func=envs_dr.effort_limits, mode="startup",
-            params={
-                "asset_cfg": SceneEntityCfg("robot"),
-                "effort_limit_range": (0.8, 1.0),
-                "operation": "scale",
-                "distribution": "uniform",
-            },
-        ),
-        "payload_mass": EventTermCfg(
+        "body_mass_base": EventTermCfg(
            func=envs_dr.body_mass, mode="startup",
            params={
                "asset_cfg": SceneEntityCfg("robot", body_names=("base_link",)),
@@ -287,14 +285,22 @@ def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
                "ranges": (-1.0, 3.0),
            },
        ),
-        "continuous_disturbance": EventTermCfg(
+        "body_mass_limbs": EventTermCfg(
+            func=envs_dr.body_mass, mode="startup",
+            params={
+                "asset_cfg": SceneEntityCfg("robot", body_names=(".*_knee_.*", ".*_wheel_.*")),
+                "operation": "scale",
+                "ranges": (0.7, 1.3),
+            },
+        ),
+        "apply_continuous_disturbance": EventTermCfg(
            func=apply_continuous_disturbance, mode="step",
            params={
                "asset_cfg": SceneEntityCfg("robot", body_names=("base_link",)),
-                "force_range": (-15.0, 15.0),
+                "force_range": (-10.0, 10.0),
                "torque_range": (-10.0, 10.0),
-                "resample_time_range": (0.5, 2.0),
-                "time_constant": 0.5,
+                "resample_time_range": (5.0, 10.0),
+                "time_constant": 1.0,
            },
        ),
    }
@@ -346,8 +352,8 @@ def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
        ),
        commands=commands, actions=actions, observations=observations,
        rewards=rewards, terminations=terminations, events=events,
-        metrics=metrics, curriculum=curriculum, decimation=10, episode_length_s=20.0,
-        sim=SimulationCfg(mujoco=MujocoCfg(impratio=100, cone="elliptic")),
+        metrics=metrics, curriculum=curriculum, decimation=4, episode_length_s=20.0,
+        sim=SimulationCfg(mujoco=MujocoCfg(timestep=0.005, impratio=100, cone="elliptic")),
        viewer=ViewerConfig(body_name="base_link", distance=3.0, elevation=-20.0, azimuth=45.0),
    )

@@ -380,33 +386,36 @@ def rough_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
            size=(8.0, 8.0), border_width=20.0, num_rows=10, num_cols=20, curriculum=True,
            sub_terrains={
                "flat": BoxFlatTerrainCfg(proportion=0.05, size=(8.0, 8.0)),
-                "pyramid_stairs": BoxPyramidStairsTerrainCfg(proportion=0.25, step_height_range=(0.0, 0.3), step_width=0.30, size=(8.0, 8.0)),
-                "pyramid_stairs_inv": BoxInvertedPyramidStairsTerrainCfg(proportion=0.10, step_height_range=(0.0, 0.3), step_width=0.30, size=(8.0, 8.0)),
-                "random_grid": BoxRandomGridTerrainCfg(proportion=0.1, grid_width=0.45, grid_height_range=(0.0, 0.3), size=(8.0, 8.0)),
-                "random_rough": HfRandomUniformTerrainCfg(proportion=0.05, noise_range=(0.0, 0.06), noise_step=0.01, horizontal_scale=0.20, downsampled_scale=0.20, border_width=0.25, base_thickness_ratio=100.0, size=(8.0, 8.0)),
-                "perlin_noise": HfPerlinNoiseTerrainCfg(proportion=0.05, height_range=(0.0, 0.06), octaves=2, persistence=0.4, lacunarity=2.0, horizontal_scale=0.20, resolution=0.20, border_width=0.50, base_thickness_ratio=100.0, size=(8.0, 8.0)),
-                "rc_wall": RCWallTerrainCfg(proportion=0.25, wall_height_range=(0.0, 0.45), size=(8.0, 8.0)),
-                "sloped_terrain": HfPyramidSlopedTerrainCfg(proportion=0.15, slope_range=(0.052, 0.325), platform_width=2.0, border_width=0.25, base_thickness_ratio=100.0, horizontal_scale=0.20, size=(8.0, 8.0)),
+                "pyramid_stairs": BoxPyramidStairsTerrainCfg(proportion=0.05, step_height_range=(0.0, 0.3), step_width=0.30, size=(8.0, 8.0)),
+                "pyramid_stairs_inv": BoxInvertedPyramidStairsTerrainCfg(proportion=0.45, step_height_range=(0.0, 0.3), step_width=0.30, size=(8.0, 8.0)),
+                "random_grid": BoxRandomGridTerrainCfg(proportion=0.27, grid_width=0.45, grid_height_range=(0.0, 0.3), size=(8.0, 8.0)),
+                "random_rough": HfRandomUniformTerrainCfg(proportion=0.01, noise_range=(0.0, 0.06), noise_step=0.01, horizontal_scale=0.20, downsampled_scale=0.20, border_width=0.25, base_thickness_ratio=100.0, size=(8.0, 8.0)),
+                "perlin_noise": HfPerlinNoiseTerrainCfg(proportion=0.01, height_range=(0.0, 0.06), octaves=2, persistence=0.4, lacunarity=2.0, horizontal_scale=0.20, resolution=0.20, border_width=0.50, base_thickness_ratio=100.0, size=(8.0, 8.0)),
+                "rc_wall": RCWallTerrainCfg(proportion=0.15, wall_height_range=(0.0, 0.45), size=(8.0, 8.0)),
+                "sloped_terrain": HfPyramidSlopedTerrainCfg(proportion=0.01, slope_range=(0.052, 0.325), platform_width=2.0, border_width=0.25, base_thickness_ratio=100.0, horizontal_scale=0.20, size=(8.0, 8.0)),
            },
        ),
-        max_init_terrain_level=0,
+        max_init_terrain_level=5,
    )

-    # Disable default velocity stages command and bind strict velocity terrain curriculum
+    # Keep the custom terrain set, but align command/curriculum behavior with go2w rough.
    cfg.curriculum.pop("command_vel", None)
-    cfg.curriculum["terrain_levels"] = CurriculumTermCfg(func=terrain_levels_vel_strict, params={"command_name": "twist"})
+    cfg.curriculum["terrain_levels"] = CurriculumTermCfg(func=velocity_mdp.terrain_levels_vel, params={"command_name": "twist"})

    cfg.commands["twist"].heading_command = True
-    cfg.commands["twist"].rel_heading_envs = 0.5
-    cfg.commands["twist"].heading_control_stiffness = 0.6
+    cfg.commands["twist"].rel_heading_envs = 1.0
+    cfg.commands["twist"].heading_control_stiffness = 0.5
    cfg.commands["twist"].ranges.heading = (-math.pi, math.pi)
-    cfg.commands["twist"].rel_standing_envs = 0.2
+    cfg.commands["twist"].rel_standing_envs = 0.02
+    cfg.commands["twist"].ranges.lin_vel_x = (-1.0, 1.0)
+    cfg.commands["twist"].ranges.lin_vel_y = (-0.6, 0.6)
+    cfg.commands["twist"].ranges.ang_vel_z = (-1.0, 1.0)

    # ------------------
    # Startup & Reset Randomizations
    # ------------------
-    cfg.events["joint_friction"] = EventTermCfg(func=envs_dr.joint_friction, mode="startup", params={"asset_cfg": SceneEntityCfg("robot"), "ranges": (0.7, 1.3), "operation": "scale"})
-    cfg.events["reset_joints"] = EventTermCfg(func=envs_mdp.reset_joints_by_offset, mode="reset", params={"position_range": (0.0, 0.1), "velocity_range": (0.0, 0.0), "asset_cfg": SceneEntityCfg("robot", joint_names=(".*",))})
+    cfg.events.pop("joint_friction", None)
+    cfg.events["reset_joints"] = EventTermCfg(func=envs_mdp.reset_joints_by_offset, mode="reset", params={"position_range": (0.0, 0.0), "velocity_range": (0.0, 0.0), "asset_cfg": SceneEntityCfg("robot", joint_names=(".*",))})
    
    cfg.events["reset_base"] = EventTermCfg(
        func=envs_mdp.reset_root_state_uniform, mode="reset",
@@ -426,42 +435,36 @@ def rough_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
    # Rewards Integration
    # ------------------
    cfg.rewards["track_lin_vel"] = RewardTermCfg(
-        func=track_linear_velocity_l1,
-        weight=4.5,
+        func=track_linear_velocity,
+        weight=3.0,
+        params={"std": 0.5, "command_name": "twist"}
+    )
+    cfg.rewards["track_ang_vel"] = RewardTermCfg(
+        func=track_angular_velocity,
+        weight=1.5,
        params={"std": 0.5, "command_name": "twist"}
    )
-    cfg.rewards["track_ang_vel"].weight = 2.0
-    cfg.rewards["track_ang_vel"].params["std"] = 0.5
    
-    cfg.rewards["lin_vel_z"] = RewardTermCfg(func=lin_vel_z_l2, weight=-0.5)  # 🌟 增强垂直速度惩罚，抑制越障后的惯性暴冲
-    cfg.rewards["ang_vel_xy"] = RewardTermCfg(func=velocity_mdp.body_angular_velocity_penalty, weight=-0.3, params={"asset_cfg": SceneEntityCfg("robot", body_names=("base_link",))})  # 🌟 增强角速度惩罚，防止突发性翻转/后仰
+    cfg.rewards["lin_vel_z"] = RewardTermCfg(func=lin_vel_z_l2, weight=-2.0)
+    cfg.rewards["ang_vel_xy"] = RewardTermCfg(func=ang_vel_xy_l2, weight=-0.05, params={"asset_cfg": SceneEntityCfg("robot")})

    cfg.rewards.pop("upright", None)
-    cfg.rewards["roll_penalty"] = RewardTermCfg(
-        func=upright_roll_only,
-        weight=-1.0,
-        params={"asset_cfg": SceneEntityCfg("robot")}
-    )
+    cfg.rewards.pop("roll_penalty", None)

    # 🌟 限制俯仰角死区（Pitch Dead-zone）：允许正常爬坡时有最大 29 度（0.50 rad）的仰角，但严厉惩罚超过该仰角的“前轮悬空暴冲/后翻”
-    cfg.rewards["pitch_penalty"] = RewardTermCfg(
-        func=pitch_control_penalty,
-        weight=-1.5,
-        params={"max_pitch_rad": 0.50, "asset_cfg": SceneEntityCfg("robot")}
-    )
-    
    # 动态课程奖励与动作惩罚衰减
    cfg.rewards.pop("terrain_level_bonus", None)
-    cfg.rewards.pop("action_rate", None)
-    cfg.rewards["action_rate_curriculum"] = RewardTermCfg(func=action_rate_curriculum_l2, weight=-0.005)
+    cfg.rewards.pop("action_rate_curriculum", None)
+    cfg.rewards["action_rate"].weight = -0.01

-    cfg.rewards["joint_torques"].weight = -1e-4
+    cfg.rewards["joint_torques"].weight = -2.5e-5
+    cfg.rewards["joint_power"] = RewardTermCfg(func=joint_power, weight=-2.0e-5)
    cfg.rewards.pop("joint_acc", None)
    cfg.rewards["leg_joint_acc_l2"] = RewardTermCfg(func=envs_mdp.joint_acc_l2, weight=-2.5e-7, params={"asset_cfg": SceneEntityCfg("robot", joint_names=(".*_hip_abduction_joint", ".*_hip_pitch_joint", ".*_knee_joint"))})
    cfg.rewards["wheel_joint_acc_l2"] = RewardTermCfg(func=envs_mdp.joint_acc_l2, weight=-2.5e-9, params={"asset_cfg": SceneEntityCfg("robot", joint_names=(".*_wheel_joint",))})


-    cfg.rewards["joint_pos_limits"].weight = -0.2
+    cfg.rewards["joint_pos_limits"].weight = -5.0
    cfg.rewards.pop("leg_motion_penalty", None)
    cfg.rewards["is_terminated"].weight = 0.0
    cfg.rewards.pop("leg_symmetry", None)
@@ -487,7 +490,7 @@ def rough_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
    cfg.rewards.pop("joint_deviation_l2", None)
    cfg.rewards["joint_pos_penalty"] = RewardTermCfg(
        func=joint_pos_penalty, 
-        weight=-0.8, 
+        weight=-1.0,
        params={
            "stand_still_scale": 5.0,
            "velocity_threshold": 0.5,
@@ -502,37 +505,29 @@ def rough_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
        weight=0.1, 
        params={"command_name": "twist", "sensor_name": "feet_ground_contact"}
    )
+    cfg.rewards["feet_air_time"].weight = 0.0
+    cfg.rewards["upward"] = RewardTermCfg(func=upward, weight=1.0)
    
-    cfg.rewards["base_height_l2"].weight = -0.5
+    cfg.rewards["base_height_l2"].weight = 0.0
+    cfg.rewards["base_height_l2"].params["target_height"] = 0.40
    cfg.rewards["base_height_l2"].params["sensor_cfg"] = SceneEntityCfg("height_scanner")

    # 恢复机身碰撞惩罚为-1.0，逼迫机器人高抬腿跨越障碍，防止拖地
    cfg.rewards.pop("body_collision", None)
-    cfg.rewards["body_collision"] = RewardTermCfg(func=velocity_mdp.self_collision_cost, weight=-1.0, params={"sensor_name": "body_collision"})
+    cfg.rewards["undesired_contacts"] = RewardTermCfg(func=undesired_contacts, weight=-1.0, params={"sensor_name": "body_collision", "threshold": 1.0})
+    cfg.rewards["contact_forces"] = RewardTermCfg(func=contact_forces, weight=-1.5e-4, params={"sensor_name": "feet_ground_contact", "threshold": 100.0})

    # 🌟 严厉惩罚机身/胸部碰撞（防止硬撞高墙），逼迫机器人学会用前轮触墙并主动抬腿攀爬的“触觉反射”
-    cfg.rewards["base_collision"] = RewardTermCfg(
-        func=velocity_mdp.self_collision_cost, 
-        weight=-5.0, 
-        params={"sensor_name": "base_ground_contact"}
-    )
-
    # 彻底移除机身俯仰约束，允许机器人抬头爬高?    cfg.rewards.pop("flat_orientation", None)

-    cfg.rewards.pop("feet_air_time", None)
-
    # Remove non-applicable rewards
-    for key in ("wheel_roll_tracking", "wheel_contact_bonus", "body_ang_vel"):
+    for key in ("wheel_roll_tracking", "wheel_contact_bonus", "body_ang_vel", "terrain_level_bonus", "action_rate_curriculum"):
        cfg.rewards.pop(key, None)

-    cfg.episode_length_s = 30.0
-    cfg.sim = SimulationCfg(contact_sensor_maxmatch=128, mujoco=MujocoCfg(impratio=100, cone="elliptic", ccd_iterations=80))
+    cfg.episode_length_s = 20.0
+    cfg.sim = SimulationCfg(contact_sensor_maxmatch=128, mujoco=MujocoCfg(timestep=0.005, impratio=100, cone="elliptic", ccd_iterations=80))

    # 移除 orientation 终止，允许机器人翻倒以学习回复
-    cfg.terminations.pop("bad_orientation", None)
-    # 移除 base_ground_contact 终止，越障时机身会碰到障碍物
-    cfg.terminations.pop("base_ground_contact", None)
-
    cfg.seed = 42
    if cfg.scene.terrain is not None:
        cfg.scene.terrain.num_envs = 2048
@@ -663,7 +658,7 @@ def crawl_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
        cfg.rewards.pop(key, None)

    cfg.episode_length_s = 30.0
-    cfg.sim = SimulationCfg(contact_sensor_maxmatch=128, mujoco=MujocoCfg(impratio=100, cone="elliptic", ccd_iterations=80))
+    cfg.sim = SimulationCfg(contact_sensor_maxmatch=128, mujoco=MujocoCfg(timestep=0.005, impratio=100, cone="elliptic", ccd_iterations=80))

    # Loosen orientation bad threshold to 80 degrees for steep crawling tilts
    cfg.terminations["bad_orientation"].params["limit_angle"] = math.radians(80.0)
@@ -33,10 +33,10 @@ def rough_ppo_runner_cfg() -> RslRlOnPolicyRunnerCfg:
            value_loss_coef=1.0,
            use_clipped_value_loss=True,
            clip_param=0.2,
-            entropy_coef=0.002,
+            entropy_coef=0.001,#第一轮为0.003
            num_learning_epochs=5,
            num_mini_batches=4,
-            learning_rate=2.0e-4,
+            learning_rate=8.0e-4,
            schedule="adaptive",
            gamma=0.99,
            lam=0.95,
@@ -44,9 +44,9 @@ def rough_ppo_runner_cfg() -> RslRlOnPolicyRunnerCfg:
            max_grad_norm=1.0,
        ),
        experiment_name="robot_rough",
-        save_interval=50,
+        save_interval=100,
        num_steps_per_env=24,
-        max_iterations=15_000,
+        max_iterations=20_000,
    )


@@ -55,5 +55,3 @@ def crawl_ppo_runner_cfg() -> RslRlOnPolicyRunnerCfg:
    cfg.experiment_name = "robot_crawl"
    return cfg

-
-
@@ -28,6 +28,7 @@ def track_linear_velocity(
    actual = asset.data.root_link_lin_vel_b
    xy_error = torch.sum(torch.square(command[:, :2] - actual[:, :2]), dim=1)
    reward = torch.exp(-xy_error / std**2)
+    reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
    return reward


@@ -45,6 +46,7 @@ def track_angular_velocity(
    actual = asset.data.root_link_ang_vel_b
    z_error = torch.square(command[:, 2] - actual[:, 2])
    reward = torch.exp(-z_error / std**2)
+    reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
    return reward


@@ -94,6 +96,7 @@ def base_height_l2(
        error = root_z - target_height
        
    reward = torch.square(error)
+    reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
    return reward


@@ -304,6 +307,7 @@ def stand_still(
    angular_norm = torch.abs(command[:, 2])
    inactive = (linear_norm + angular_norm < command_threshold).float()
    reward = cost * inactive
+    reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
    return reward

 def hip_deviation(
@@ -355,6 +359,20 @@ def lin_vel_z_l2(
        asset_cfg = SceneEntityCfg("robot")
    asset: Entity = env.scene[asset_cfg.name]
    reward = torch.square(asset.data.root_link_lin_vel_b[:, 2])
+    reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
+    return reward
+
+
+def ang_vel_xy_l2(
+    env: ManagerBasedRlEnv,
+    asset_cfg: SceneEntityCfg | None = None,
+) -> torch.Tensor:
+    """Penalize xy-axis base angular velocity using the go2w kernel."""
+    if asset_cfg is None:
+        asset_cfg = SceneEntityCfg("robot")
+    asset: Entity = env.scene[asset_cfg.name]
+    reward = torch.sum(torch.square(asset.data.root_link_ang_vel_b[:, :2]), dim=1)
+    reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
    return reward


@@ -650,6 +668,8 @@ def feet_contact_without_cmd(env, command_name: str, sensor_name: str) -> torch.
    linear_norm = torch.norm(cmd[:, :2], dim=1)
    angular_norm = torch.abs(cmd[:, 2])
    reward *= (linear_norm + angular_norm) < 0.1
+    asset = env.scene["robot"]
+    reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
    return reward

 def joint_pos_penalty(
@@ -674,6 +694,7 @@ def joint_pos_penalty(
        running_reward,
        stand_still_scale * running_reward,
    )
+    reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
    return reward

 def joint_mirror(env, mirror_joints: list[list[str]], asset_cfg: SceneEntityCfg | None = None) -> torch.Tensor:
@@ -694,6 +715,50 @@ def joint_mirror(env, mirror_joints: list[list[str]], asset_cfg: SceneEntityCfg
        )
        reward += diff
    reward *= 1 / len(mirror_joints) if len(mirror_joints) > 0 else 0
+    reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
+    return reward
+
+
+def undesired_contacts(
+    env: ManagerBasedRlEnv,
+    sensor_name: str,
+    threshold: float = 1.0,
+) -> torch.Tensor:
+    """Penalize non-wheel contacts above a force threshold."""
+    from mjlab.sensor import ContactSensor
+
+    sensor: ContactSensor = env.scene[sensor_name]
+    data = sensor.data
+    if data.force_history is not None:
+        force_mag = torch.norm(data.force_history, dim=-1)
+        is_contact = torch.max(force_mag, dim=2)[0] > threshold
+    else:
+        force_mag = torch.norm(data.force, dim=-1)
+        is_contact = force_mag > threshold
+    reward = torch.sum(is_contact, dim=1).float()
+    asset = env.scene["robot"]
+    reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
+    return reward
+
+
+def contact_forces(
+    env: ManagerBasedRlEnv,
+    sensor_name: str,
+    threshold: float = 100.0,
+) -> torch.Tensor:
+    """Penalize foot contact forces above threshold."""
+    from mjlab.sensor import ContactSensor
+
+    sensor: ContactSensor = env.scene[sensor_name]
+    data = sensor.data
+    if data.force_history is not None:
+        force_mag = torch.norm(data.force_history, dim=-1)
+        peak_force = torch.max(force_mag, dim=2)[0]
+    else:
+        peak_force = torch.norm(data.force, dim=-1)
+    reward = torch.sum(torch.clamp(peak_force - threshold, min=0.0), dim=1)
+    asset = env.scene["robot"]
+    reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
    return reward


@@ -705,6 +770,14 @@ def upward(env, asset_cfg=None):
    reward = torch.square(1 - asset.data.projected_gravity_b[:, 2])
    return reward

+
+def joint_power(env: ManagerBasedRlEnv, asset_cfg: SceneEntityCfg | None = None) -> torch.Tensor:
+    """Penalty for total joint mechanical power: sum(|tau * dq|)."""
+    if asset_cfg is None:
+        asset_cfg = SceneEntityCfg("robot")
+    asset: Entity = env.scene[asset_cfg.name]
+    return torch.sum(torch.abs(asset.data.qfrc_actuator * asset.data.joint_vel), dim=1)
+
 def upright_roll_only(env, asset_cfg=None):
    if asset_cfg is None:
        from mjlab.envs.manager_based_rl_env import SceneEntityCfg
@@ -1,41 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: BSD-3-Clause
-# 
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-#
-# 1. Redistributions of source code must retain the above copyright notice, this
-# list of conditions and the following disclaimer.
-#
-# 2. Redistributions in binary form must reproduce the above copyright notice,
-# this list of conditions and the following disclaimer in the documentation
-# and/or other materials provided with the distribution.
-#
-# 3. Neither the name of the copyright holder nor the names of its
-# contributors may be used to endorse or promote products derived from
-# this software without specific prior written permission.
-#
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-#
-# Copyright (c) 2021 ETH Zurich, Nikita Rudin
-
-"""
-HimLoco RSL-RL implementation with history-informed models.
-"""
-
-# Export HIM implementations
-from .algorithms.him_ppo import *
-from .modules.him_actor_critic import *
-from .modules.him_estimator import *
-from .storage.him_rollout_storage import *
-from .runners.him_on_policy_runner import *
-from .wrappers import *
@@ -1,31 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: BSD-3-Clause
-# 
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-#
-# 1. Redistributions of source code must retain the above copyright notice, this
-# list of conditions and the following disclaimer.
-#
-# 2. Redistributions in binary form must reproduce the above copyright notice,
-# this list of conditions and the following disclaimer in the documentation
-# and/or other materials provided with the distribution.
-#
-# 3. Neither the name of the copyright holder nor the names of its
-# contributors may be used to endorse or promote products derived from
-# this software without specific prior written permission.
-#
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-#
-# Copyright (c) 2021 ETH Zurich, Nikita Rudin
-
-from .him_ppo import HIMPPO
@@ -1,192 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: BSD-3-Clause
-# 
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-#
-# 1. Redistributions of source code must retain the above copyright notice, this
-# list of conditions and the following disclaimer.
-#
-# 2. Redistributions in binary form must reproduce the above copyright notice,
-# this list of conditions and the following disclaimer in the documentation
-# and/or other materials provided with the distribution.
-#
-# 3. Neither the name of the copyright holder nor the names of its
-# contributors may be used to endorse or promote products derived from
-# this software without specific prior written permission.
-#
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-#
-# Copyright (c) 2021 ETH Zurich, Nikita Rudin
-
-import torch
-import torch.nn as nn
-import torch.optim as optim
-
-from ..modules import HIMActorCritic
-from ..storage import HIMRolloutStorage
-
-class HIMPPO:
-    actor_critic: HIMActorCritic
-    def __init__(self,
-                 actor_critic,
-                 num_learning_epochs=1,
-                 num_mini_batches=1,
-                 clip_param=0.2,
-                 gamma=0.998,
-                 lam=0.95,
-                 value_loss_coef=1.0,
-                 entropy_coef=0.0,
-                 learning_rate=1e-3,
-                 max_grad_norm=1.0,
-                 use_clipped_value_loss=True,
-                 schedule="fixed",
-                 desired_kl=0.01,
-                 device='cpu',
-                 ):
-
-        self.device = device
-
-        self.desired_kl = desired_kl
-        self.schedule = schedule
-        self.learning_rate = learning_rate
-
-        # PPO components
-        self.actor_critic = actor_critic
-        self.actor_critic.to(self.device)
-        self.storage = None # initialized later
-        self.optimizer = optim.Adam(self.actor_critic.parameters(), lr=learning_rate)
-        self.transition = HIMRolloutStorage.Transition()
-
-        # PPO parameters
-        self.clip_param = clip_param
-        self.num_learning_epochs = num_learning_epochs
-        self.num_mini_batches = num_mini_batches
-        self.value_loss_coef = value_loss_coef
-        self.entropy_coef = entropy_coef
-        self.gamma = gamma
-        self.lam = lam
-        self.max_grad_norm = max_grad_norm
-        self.use_clipped_value_loss = use_clipped_value_loss
-
-    def init_storage(self, num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, action_shape):
-        self.storage = HIMRolloutStorage(num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, action_shape, self.device)
-
-    def test_mode(self):
-        self.actor_critic.test()
-    
-    def train_mode(self):
-        self.actor_critic.train()
-
-    def act(self, obs, critic_obs):
-        # Compute the actions and values
-        self.transition.actions = self.actor_critic.act(obs).detach()
-        self.transition.values = self.actor_critic.evaluate(critic_obs).detach()
-        self.transition.actions_log_prob = self.actor_critic.get_actions_log_prob(self.transition.actions).detach()
-        self.transition.action_mean = self.actor_critic.action_mean.detach()
-        self.transition.action_sigma = self.actor_critic.action_std.detach()
-        # need to record obs and critic_obs before env.step()
-        self.transition.observations = obs
-        self.transition.critic_observations = critic_obs
-        return self.transition.actions
-    
-    def process_env_step(self, rewards, dones, infos, next_critic_obs):
-        self.transition.next_critic_observations = next_critic_obs.clone()
-        self.transition.rewards = rewards.clone()
-        self.transition.dones = dones
-        # Bootstrapping on time outs
-        if 'time_outs' in infos:
-            self.transition.rewards += self.gamma * torch.squeeze(self.transition.values * infos['time_outs'].unsqueeze(1).to(self.device), 1)
-
-        # Record the transition
-        self.storage.add_transitions(self.transition)
-        self.transition.clear()
-        self.actor_critic.reset(dones)
-    
-    def compute_returns(self, last_critic_obs):
-        last_values= self.actor_critic.evaluate(last_critic_obs).detach()
-        self.storage.compute_returns(last_values, self.gamma, self.lam)
-
-    def update(self):
-        mean_value_loss = 0
-        mean_surrogate_loss = 0
-        mean_estimation_loss = 0
-        mean_swap_loss = 0
-        
-        generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs)
-
-        for obs_batch, critic_obs_batch, actions_batch, next_critic_obs_batch, target_values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, \
-            old_mu_batch, old_sigma_batch in generator:
-                
-                self.actor_critic.act(obs_batch)
-                actions_log_prob_batch = self.actor_critic.get_actions_log_prob(actions_batch)
-                value_batch = self.actor_critic.evaluate(critic_obs_batch)
-                mu_batch = self.actor_critic.action_mean
-                sigma_batch = self.actor_critic.action_std
-                entropy_batch = self.actor_critic.entropy
-
-                # KL
-                if self.desired_kl != None and self.schedule == 'adaptive':
-                    with torch.inference_mode():
-                        kl = torch.sum(
-                            torch.log(sigma_batch / old_sigma_batch + 1.e-5) + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) / (2.0 * torch.square(sigma_batch)) - 0.5, axis=-1)
-                        kl_mean = torch.mean(kl)
-
-                        if kl_mean > self.desired_kl * 2.0:
-                            self.learning_rate = max(1e-5, self.learning_rate / 1.5)
-                        elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0:
-                            self.learning_rate = min(1e-2, self.learning_rate * 1.5)
-                        
-                        for param_group in self.optimizer.param_groups:
-                            param_group['lr'] = self.learning_rate
-
-                #Estimator Update
-                estimation_loss, swap_loss = self.actor_critic.estimator.update(obs_batch, next_critic_obs_batch, lr=self.learning_rate)
-
-                # Surrogate loss
-                ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch))
-                surrogate = -torch.squeeze(advantages_batch) * ratio
-                surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp(ratio, 1.0 - self.clip_param,
-                                                                                1.0 + self.clip_param)
-                surrogate_loss = torch.max(surrogate, surrogate_clipped).mean()
-
-                # Value function loss
-                if self.use_clipped_value_loss:
-                    value_clipped = target_values_batch + (value_batch - target_values_batch).clamp(-self.clip_param,
-                                                                                                    self.clip_param)
-                    value_losses = (value_batch - returns_batch).pow(2)
-                    value_losses_clipped = (value_clipped - returns_batch).pow(2)
-                    value_loss = torch.max(value_losses, value_losses_clipped).mean()
-                else:
-                    value_loss = (returns_batch - value_batch).pow(2).mean()
-
-                loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean()
-
-                # Gradient step
-                self.optimizer.zero_grad()
-                loss.backward()
-                nn.utils.clip_grad_norm_(self.actor_critic.parameters(), self.max_grad_norm)
-                self.optimizer.step()
-
-                mean_value_loss += value_loss.item()
-                mean_surrogate_loss += surrogate_loss.item()
-                mean_estimation_loss += estimation_loss
-                mean_swap_loss += swap_loss
-
-        num_updates = self.num_learning_epochs * self.num_mini_batches
-        mean_value_loss /= num_updates
-        mean_surrogate_loss /= num_updates
-        mean_estimation_loss /= num_updates
-        mean_swap_loss /= num_updates
-        self.storage.clear()
-
-        return mean_value_loss, mean_surrogate_loss, estimation_loss, swap_loss
@@ -1,31 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: BSD-3-Clause
-# 
-# Redistribution and use in source and binary forms, with or without
-# modification, are permitted provided that the following conditions are met:
-#
-# 1. Redistributions of source code must retain the above copyright notice, this
-# list of conditions and the following disclaimer.
-#
-# 2. Redistributions in binary form must reproduce the above copyright notice,
-# this list of conditions and the following disclaimer in the documentation
-# and/or other materials provided with the distribution.
-#
-# 3. Neither the name of the copyright holder nor the names of its
-# contributors may be used to endorse or promote products derived from
-# this software without specific prior written permission.
-#
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-#
-# Copyright (c) 2021 ETH Zurich, Nikita Rudin
-
-from .rl_cfg import *
@@ -1,180 +0,0 @@
-from dataclasses import MISSING
-from isaaclab.utils import configclass
-from typing import Literal
-
-
-@configclass
-class HIMBaseRunnerCfg:
-    """Base configuration of the runner."""
-
-    seed: int = 1
-    """The seed for the experiment. Default is 1."""
-
-    device: str = "cuda:0"
-    """The device for the rl-agent. Default is cuda:0."""
-
-    num_steps_per_env: int = MISSING
-    """The number of steps per environment per update."""
-
-    max_iterations: int = MISSING
-    """The maximum number of iterations."""
-
-    empirical_normalization: bool | None = None
-    """This parameter is deprecated and will be removed in the future.
-
-    Use `actor_obs_normalization` and `critic_obs_normalization` instead.
-    """
-
-    # obs_groups: dict[str, list[str]] = MISSING
-    # """A mapping from observation groups to observation sets.
-
-    # The keys of the dictionary are predefined observation sets used by the underlying algorithm
-    # and values are lists of observation groups provided by the environment.
-
-    # For instance, if the environment provides a dictionary of observations with groups "policy", "images",
-    # and "privileged", these can be mapped to algorithmic observation sets as follows:
-
-    # .. code-block:: python
-
-    #     obs_groups = {
-    #         "policy": ["policy", "images"],
-    #         "critic": ["policy", "privileged"],
-    #     }
-
-    # This way, the policy will receive the "policy" and "images" observations, and the critic will
-    # receive the "policy" and "privileged" observations.
-
-    # For more details, please check ``vec_env.py`` in the rsl_rl library.
-    # """
-
-    # clip_actions: float | None = None
-    # """The clipping value for actions. If None, then no clipping is done. Defaults to None.
-
-    # .. note::
-    #     This clipping is performed inside the :class:`RslRlVecEnvWrapper` wrapper.
-    # """
-
-    save_interval: int = MISSING
-    """The number of iterations between saves."""
-
-    experiment_name: str = MISSING
-    """The experiment name."""
-
-    run_name: str = ""
-    """The run name. Default is empty string.
-
-    The name of the run directory is typically the time-stamp at execution. If the run name is not empty,
-    then it is appended to the run directory's name, i.e. the logging directory's name will become
-    ``{time-stamp}_{run_name}``.
-    """
-
-    logger: Literal["tensorboard", "neptune", "wandb"] = "tensorboard"
-    """The logger to use. Default is tensorboard."""
-
-    neptune_project: str = "isaaclab"
-    """The neptune project name. Default is "isaaclab"."""
-
-    wandb_project: str = "isaaclab"
-    """The wandb project name. Default is "isaaclab"."""
-
-    resume: bool = False
-    """Whether to resume a previous training. Default is False.
-
-    This flag will be ignored for distillation.
-    """
-
-    load_run: str = ".*"
-    """The run directory to load. Default is ".*" (all).
-
-    If regex expression, the latest (alphabetical order) matching run will be loaded.
-    """
-
-    load_checkpoint: str = "model_.*.pt"
-    """The checkpoint file to load. Default is ``"model_.*.pt"`` (all).
-
-    If regex expression, the latest (alphabetical order) matching file will be loaded.
-    """
-    
-    
-@configclass
-class HIMPPOActorCriticCfg:
-    """Configuration of the HIM PPO actor-critic."""
-
-    actor_hidden_dims: list[int] = [512, 256, 128]
-    """The hidden dimensions of the actor network."""
-
-    critic_hidden_dims: list[int] = [512, 256, 128]
-    """The hidden dimensions of the critic network."""
-
-    activation: str = "elu"
-    """The activation function to use. Default is 'elu'."""
-
-    init_noise_std: float = 1.0
-    """The initial noise standard deviation for the actor. Default is 1.0."""
-    
-    normalize_obs:  bool = False
-    """Whether to normalize observations. Default is False."""
-
-@configclass
-class HIMPPPOAlgorithmCfg:
-    """Configuration of the HIM PPO algorithm."""
-    num_learning_epochs: int = 1
-    """The number of learning epochs per update. Default is 1."""
-
-    num_mini_batches: int = 1
-    """The number of mini-batches per update. Default is 1."""
-
-    clip_param: float = 0.2
-    """The clipping parameter for PPO. Default is 0.2."""
-
-    gamma: float = 0.998
-    """The discount factor. Default is 0.998."""
-
-    lam: float = 0.95
-    """The GAE lambda parameter. Default is 0.95."""
-
-    value_loss_coef: float = 1.0
-    """The coefficient for the value loss. Default is 1.0."""
-
-    entropy_coef: float = 0.0
-    """The coefficient for the entropy bonus. Default is 0.0."""
-
-    learning_rate: float = 1.0e-3
-    """The learning rate. Default is 1.0e-3."""
-
-    max_grad_norm: float = 1.0
-    """The maximum gradient norm for clipping. Default is 1.0."""
-    
-    use_clipped_value_loss: bool = True
-    """Whether to use clipped value loss. Default is True."""
-
-    schedule: str = "fixed"
-    """The learning rate schedule. Default is 'fixed'."""
-
-    desired_kl: float = 0.01
-    """The desired KL divergence for adaptive learning rate. Default is 0.01."""
-
-@configclass
-class HIMOnPolicyRunnerCfg(HIMBaseRunnerCfg):
-    """Configuration of the runner for on-policy algorithms."""
-
-    class_name: str = "HIMOnPolicyRunner"
-    """The runner class name. Default is OnPolicyRunner."""
-    
-    policy_class_name: str = "HIMActorCritic"
-    """The policy class name. Default is HIMActorCritic."""
-    
-    algorithm_class_name: str = "HIMPPO"
-    """The algorithm class name. Default is HIMPPO."""
-
-    policy: HIMPPOActorCriticCfg = MISSING
-    """The policy configuration."""
-
-    algorithm: HIMPPPOAlgorithmCfg = MISSING
-    """The algorithm configuration."""
-    
-    history_length: int = 0
-    """Number of historical time steps to stack with current observation (0 means current only). Default is 0."""
-    
-    privileged_history_length: int = 0
-    """Number of historical time steps to stack with current privileged observation. Default is 0."""
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
zeitvex	60f7a08e91	[train] 增强Sim2Real训练随机化配置	2026-07-27 12:39:27 +08:00
zeitvex	b08956aec7	[train] 更新新MJCF与第一版完整训练框架	2026-07-27 12:31:21 +08:00