Compare commits
4 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 zeitvex
|
c05c1cb162 | ||
|
zeitvex
|
d8c5d34091 | ||
|
zeitvex
|
60f7a08e91 | ||
|
zeitvex
|
b08956aec7 |
@@ -35,6 +35,7 @@ logs/
|
||||
checkpoints/
|
||||
wandb/
|
||||
sim2sim_log_*.txt
|
||||
**/sim2sim_temp.xml
|
||||
|
||||
# IDE and operating system files
|
||||
.idea/
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
# 训练代码演进
|
||||
|
||||
本项目将“训练代码架构”和“训练产生的模型 checkpoint”分别管理。代码、奖励、课程或观测契约发生变化时形成新的 Git 版本；同一架构下继续训练产生的模型编号作为实验和部署工件记录。
|
||||
|
||||
| 版本 | 来源快照 | 主要目的 |
|
||||
| --- | --- | --- |
|
||||
| `v0.4.0` | `uni_mjlab(1)` | 第一份完整的新 MJCF 与新 mjlab 训练工程 |
|
||||
| `v0.5.0` | `uni_mjlab_new` | 扩大观测、延迟和动力学随机化，探索更强 Sim2Real 鲁棒性 |
|
||||
| `v0.6.0` | `best` | 面向比赛越障重新设计奖励、课程、站姿和诊断体系 |
|
||||
|
||||
## 从随机化增强到比赛训练
|
||||
|
||||
`v0.6.0` 不是简单地继续增大 `v0.5.0` 的随机范围。实践中重新降低了部分噪声、延迟和动力学随机化强度，并把训练重点转向可控的比赛任务课程：
|
||||
|
||||
- 投影重力噪声由 `±0.08` 调回 `±0.05`。
|
||||
- 腿和轮动作最大随机延迟由 4 步调回 2 步。
|
||||
- 摩擦、刚度和阻尼随机化调回较窄范围。
|
||||
- 移除该阶段的连续机身外力扰动与腿部质量随机化。
|
||||
- 将 x、y 和 yaw 跟踪拆成独立奖励和独立指令课程。
|
||||
- 障碍由课程逐步释放，高墙训练地形改为五道重复横墙。
|
||||
- 增加楼梯侧向漂移和偏航漂移惩罚。
|
||||
- 增加大量只用于训练评估的误差、轮速、姿态和接触指标。
|
||||
|
||||
这种调整反映的是从“广泛鲁棒性探索”转向“比赛场景定向优化”，不表示 `v0.5.0` 被删除；它仍由对应 Tag 完整保留。
|
||||
|
||||
## 模型与比赛部署
|
||||
|
||||
比赛训练可能先获得基模，再修改参数继续训练和筛选 checkpoint。最终关系为：
|
||||
|
||||
```text
|
||||
训练代码架构：v0.6.0 / best
|
||||
比赛 Rough 策略：model_6800.onnx
|
||||
比赛真机工程：last_not_slalom_1050
|
||||
比赛得分：1050
|
||||
```
|
||||
|
||||
`model_6800.onnx` 是比赛最终部署工件，不用模型编号替代训练代码版本号。它将在最终比赛部署版本中与运行配置一起归档。
|
||||
@@ -0,0 +1,53 @@
|
||||
# 版本演进
|
||||
|
||||
本项目使用同一条 `16dof` 主线和里程碑 Tag 保存线性演进，不在源码目录中复制历史版本。
|
||||
|
||||
| Tag | 阶段 | 核心内容 |
|
||||
| --- | --- | --- |
|
||||
| `v0.1.0` | 8DOF 中期检查 | 8DOF 串联足机械与大疆 A 板实机版本 |
|
||||
| `v0.2.0` | 16DOF 机械 | 16DOF 串联轮足机械 CAD 与 STEP |
|
||||
| `v0.3.0` | 第一代软件闭环 | 早期训练、MJCF、MuJoCo、Sim2Sim、IK 与 Python Sim2Real |
|
||||
| `v0.3.1` | 实机记录 | 补充第一代 Sim2Real 实机视频 |
|
||||
| `v0.4.0` | 新训练基线 | 第一份完整的新 MJCF、新 mjlab 框架和 Rough 策略工程 |
|
||||
| `v0.5.0` | 随机化增强 | 扩大观测、延迟和动力学随机化，加入持续外力扰动 |
|
||||
| `v0.6.0` | 比赛训练架构 | 分轴奖励、自适应指令课程、障碍释放课程和比赛站姿 |
|
||||
| `v0.7.0` | MuJoCo 工具 | 姿态优化、IK 扫描、动力学、MPC 和 GUI 调试工具 |
|
||||
|
||||
## `v0.4.0` 的模型变化
|
||||
|
||||
- 机械 CAD 不变。
|
||||
- MJCF 更新整机质量和惯性参数，旧、新 `wheelleg.xml` 的 SHA-256 不同。
|
||||
- mjlab 上游基准从 `00409797` 更新到 `40f8d93e`。
|
||||
- 保留轮腿分组执行器随机化所需的本地补丁。
|
||||
- 本阶段归档 `model_rough.pt`，不将生成日志、缓存和临时 XML 纳入版本库。
|
||||
|
||||
## `v0.5.0` 的训练变化
|
||||
|
||||
- MJCF、mjlab 基准和已有模型文件保持不变。
|
||||
- 投影重力噪声由 `±0.05` 扩大到 `±0.08`。
|
||||
- 腿与轮动作的最大随机延迟由 2 步增加到 4 步。
|
||||
- 地面摩擦随机范围由 `0.3–1.0` 扩大到 `0.15–1.25`。
|
||||
- 执行器刚度和阻尼缩放由 `0.9–1.1` 扩大到 `0.5–1.5`。
|
||||
- 增加膝部和轮部质量的 `0.7–1.3` 随机缩放。
|
||||
- 增加作用于机身的连续随机外力和力矩扰动。
|
||||
|
||||
## `v0.6.0` 的比赛训练架构
|
||||
|
||||
- 保持 `v0.4.0` 引入的新 MJCF 和 mjlab 框架不变。
|
||||
- 将线速度奖励拆分为 x/y 两轴，并独立配置偏航角速度奖励。
|
||||
- 增加自适应 x/y/yaw 指令范围课程。
|
||||
- 增加障碍地形逐步释放与更严格的地形晋级逻辑。
|
||||
- 增加楼梯横向速度和偏航漂移约束。
|
||||
- 默认站姿调整为髋俯仰 `0.550`、膝关节 `-1.125`，初始机身高度为 `0.42 m`。
|
||||
- 增加速度误差、轮速跟踪、动作和姿态等训练诊断指标。
|
||||
- 该 Tag 保存比赛训练代码架构，不把每次继续训练产生的 checkpoint 误记为新的软件版本。
|
||||
|
||||
训练阶段的详细关系见 [`training_evolution.md`](training_evolution.md)。
|
||||
|
||||
## `v0.7.0` 的独立 MuJoCo 工具
|
||||
|
||||
- 比赛训练架构、MJCF、模型和依赖锁文件保持 `v0.6.0` 状态不变。
|
||||
- 增加解析姿态表、RL 友好姿态筛选和 MuJoCo 静态姿态优化。
|
||||
- 增加 IK/差速轮参数扫描，可导出 JSON 结果。
|
||||
- 增加 Robot、Controller、Dynamics、MPCController 和 GUI 调试链路。
|
||||
- 记录历史工具常量与新版 MJCF 质量、比赛默认站姿之间的参数边界，避免将分析结果直接当作已校准真机参数。
|
||||
@@ -1,6 +1,6 @@
|
||||
# 软件
|
||||
|
||||
本目录当前保存 16DOF 轮足机器人的第一代完整软件闭环。
|
||||
本目录保存 16DOF 轮足机器人的训练、仿真和真机软件演进。
|
||||
|
||||
```text
|
||||
05_software/
|
||||
@@ -19,7 +19,7 @@ MJCF + mjlab task
|
||||
v
|
||||
PPO 训练策略
|
||||
|
|
||||
+----> MuJoCo 独立模型调试
|
||||
+----> MuJoCo 姿态 / IK / MPC 调试
|
||||
|
|
||||
+----> Sim2Sim 策略验证
|
||||
|
|
||||
@@ -28,7 +28,7 @@ MJCF + mjlab task
|
||||
IK real --------------------------------> 电机
|
||||
```
|
||||
|
||||
`rc_mjlab` 在早期版本中是自包含工程。训练、MJCF、独立 MuJoCo、Sim2Sim 和策略权重通过相对路径绑定，因此本次保留其原始内部布局，没有为了目录外观拆散。
|
||||
`rc_mjlab` 是自包含工程。训练、MJCF、MuJoCo、Sim2Sim 和策略权重通过相对路径绑定，因此保留其内部布局，没有为了目录外观拆散。第一代完整闭环见 `v0.3.0`，第一份新版 MJCF 与训练框架见 `v0.4.0`，随机化增强版见 `v0.5.0`，比赛最终训练架构见 `v0.6.0`，后期 MuJoCo 工具集见 `v0.7.0`。
|
||||
|
||||
详细说明见：
|
||||
|
||||
|
||||
@@ -1,15 +1,25 @@
|
||||
# 第一代强化学习与仿真工程
|
||||
# 强化学习与仿真工程
|
||||
|
||||
`rc_mjlab/` 是 16DOF 轮足机器人的第一代自包含训练与仿真工程。
|
||||
`rc_mjlab/` 保存 16DOF 轮足机器人的当前训练与 Sim2Sim 工程。历史快照由 Git Tag 保留，不在目录中复制 `old`、`new` 或 `final` 版本。
|
||||
|
||||
当前内容对应 `v0.7.0`：训练代码保持 `v0.6.0` 的比赛架构，新增后期 MuJoCo 姿态、IK、动力学和 MPC 工具。训练过程可能先获得基模，再调整奖励、课程和环境参数继续训练；模型 checkpoint 的变化不等同于软件架构变化。
|
||||
|
||||
## 内容
|
||||
|
||||
- `src/robot`：Flat、Rough、Crawl 训练任务和自定义 MDP
|
||||
- `mjcf`：轮足机器人 MuJoCo 模型和网格
|
||||
- `mujoco_sim`：不依赖策略的独立 MuJoCo/MPC 调试工具
|
||||
- `sim2sim`：策略加载、交互控制和比赛地形验证
|
||||
- `mujoco_sim`：不依赖训练循环的姿态、IK、动力学和 MPC 分析
|
||||
- `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 步，扩大摩擦、刚度和阻尼随机化，增加腿部质量随机化与连续外力/力矩扰动。
|
||||
|
||||
`v0.6.0` 在 `v0.5.0` 之后转向比赛任务优化：降低部分过强随机化，加入分轴速度跟踪奖励、自适应指令课程、障碍地形释放课程、楼梯横向/偏航约束以及更完整的训练诊断。详细对比见 [`../../01_doc/training_evolution.md`](../../01_doc/training_evolution.md)。
|
||||
|
||||
`v0.7.0` 不修改比赛训练架构，增加独立 MuJoCo 工具；入口和参数边界见 [`rc_mjlab/mujoco_sim/README.md`](rc_mjlab/mujoco_sim/README.md)。
|
||||
|
||||
工程命令和任务说明见 [`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.7.0`：保留 `v0.6.0` 的比赛训练架构，并加入后期 MuJoCo 独立工具集。当前目录中的 `model_rough.pt` 是早期参考权重；比赛最终使用的 `model_6800.onnx` 将随最终部署版本归档。
|
||||
|
||||
---
|
||||
|
||||
## 🛠️ 项目简介
|
||||
@@ -51,8 +53,8 @@ rc_mjlab/
|
||||
│ ├── wheelleg.xml # 机器人 MuJoCo 模型（含网格引用）
|
||||
│ ├── scene.xml # mjlab 场景入口文件
|
||||
│ └── meshes/ # STL/OBJ 碰撞与外观网格
|
||||
├── mujoco_sim/ # 独立 MPC 仿真调试工具（不依赖 RL 训练）
|
||||
├── logs/ # 训练日志（rsl_rl 格式，按任务名/日期/checkpoint 归档）
|
||||
├── mujoco_sim/ # 姿态、IK、动力学和 MPC 独立工具
|
||||
├── 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:
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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",))
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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"]
|
||||
|
||||
|
||||
Generated
+34
-26
@@ -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]]
|
||||
|
||||
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,57 @@
|
||||
# MuJoCo 独立工具集
|
||||
|
||||
本目录保存不依赖训练循环的 MuJoCo、姿态、IK、动力学和 MPC 分析工具。脚本通过父目录读取 `../mjcf/scene.xml` 与 `../mjcf/wheelleg.xml`，因此应从 `rc_mjlab` 工程根目录运行。
|
||||
|
||||
## 工具分类
|
||||
|
||||
| 入口 | 用途 |
|
||||
| --- | --- |
|
||||
| `posture_tool.py` | 基于解析运动学快速生成站立和低姿态参数表 |
|
||||
| `rl_friendly_opt.py` | 按轮心位置、关节力矩和雅可比条件数筛选适合 RL 的姿态 |
|
||||
| `posture_optimizer.py` | 解析计算与 MuJoCo 扫描结合的姿态优化 |
|
||||
| `static_posture_optimizer.py` | 在重力和地面接触下评估静态站姿、支撑域和离地间隙 |
|
||||
| `ik_diff_sweep.py` | 扫描 IK 姿态和差速轮跟踪参数，可导出 JSON |
|
||||
| `run.py` | 启动完整 MuJoCo 控制、GUI 和 MPC 调试链路 |
|
||||
| `robot.py`、`controller.py` | 仿真机器人接口和控制器 |
|
||||
| `dynamics.py`、`mpc.py`、`mpc_controller.py` | Pinocchio 动力学与 OSQP MPC |
|
||||
|
||||
## 依赖
|
||||
|
||||
执行工程根目录的 `uv sync` 后，训练环境已经提供 NumPy、SciPy 和 MuJoCo。不同工具还需要：
|
||||
|
||||
- 纯解析工具：Python、NumPy。
|
||||
- MuJoCo 扫描：`mujoco`、NumPy。
|
||||
- 完整 MPC：`pinocchio`、`osqp`、SciPy。
|
||||
- GUI：系统可用的 Tk/Tkinter。
|
||||
|
||||
Pinocchio 和 OSQP 没有加入训练环境锁文件，因为它们只服务于可选 MPC 工具，且 Pinocchio 的安装方式与操作系统、Conda/Python 环境有关。
|
||||
|
||||
## 常用命令
|
||||
|
||||
在 `05_software/train/rc_mjlab` 下执行：
|
||||
|
||||
```bash
|
||||
# 不启动 MuJoCo 的快速姿态表
|
||||
uv run python mujoco_sim/posture_tool.py
|
||||
uv run python mujoco_sim/rl_friendly_opt.py
|
||||
|
||||
# 姿态扫描
|
||||
uv run python mujoco_sim/posture_optimizer.py --analyze
|
||||
uv run python mujoco_sim/static_posture_optimizer.py --quick
|
||||
|
||||
# IK 与差速轮参数快速扫描
|
||||
uv run python mujoco_sim/ik_diff_sweep.py --quick
|
||||
|
||||
# 完整 GUI/MPC 仿真，需要可选依赖
|
||||
uv run python mujoco_sim/run.py
|
||||
```
|
||||
|
||||
## 参数边界
|
||||
|
||||
这是一份历史工具快照，保留当时用于分析和调参的常量：
|
||||
|
||||
- `config.py`、`posture_optimizer.py` 和 `static_posture_optimizer.py` 中的解析质量常量为 `12.3 kg`。
|
||||
- 当前新版 MJCF 的惯性质量合计约为 `18.0377 kg`。
|
||||
- `config.py` 的姿态表默认值为髋俯仰 `0.666`、膝关节 `-1.546`；比赛训练架构的 Rough 默认姿态为 `0.550/-1.125`。
|
||||
|
||||
MuJoCo 直接加载模型的工具会使用 MJCF 内的质量和惯性；显式读取 `ROBOT_MASS` 的解析计算和 MPC 工具仍使用历史常量。使用输出作为新版本控制参数前，应先根据目标机械状态完成质量、惯性和默认姿态复核。
|
||||
@@ -1,6 +1,5 @@
|
||||
"""Robot constants and control parameters."""
|
||||
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
|
||||
# Paths
|
||||
@@ -10,48 +9,73 @@ MJCF_PATH = REPO_ROOT / "mjcf" / "wheelleg.xml"
|
||||
|
||||
# Robot geometry
|
||||
WHEEL_RADIUS = 0.10 # m
|
||||
WHEEL_TRACK = 0.32 # m (left-right distance)
|
||||
WHEEL_TRACK = 0.32 # m, left-right distance
|
||||
ROBOT_MASS = 12.3 # kg
|
||||
MAX_TORQUE = 17.0 # Nm per joint
|
||||
MAX_JOINT_VEL = 13.0 # rad/s
|
||||
|
||||
# Leg link lengths (from MJCF)
|
||||
L_THIGH = 0.25 # m
|
||||
L_CALF = 0.20 # m (to wheel center)
|
||||
# Leg link lengths from MJCF, measured to wheel center.
|
||||
L_THIGH = 0.25
|
||||
L_CALF = 0.20
|
||||
|
||||
# Leg names and joint ordering
|
||||
LEG_NAMES = ("fl", "fr", "rl", "rr")
|
||||
LEG_JOINTS = ("hip_abduction_joint", "hip_pitch_joint", "knee_joint")
|
||||
WHEEL_JOINT = "wheel_joint"
|
||||
|
||||
# Default standing pose (from go2w_sim2sim: [0, 0.8, -1.5])
|
||||
# Default standing pose aligned with the soft wheel-X height table.
|
||||
# height ~= 0.37m, wheel x-offset ~= 0, peak/RMS leg torque balanced.
|
||||
DEFAULT_JOINT_ANGLES = {
|
||||
"hip_abduction": 0.0,
|
||||
"hip_pitch": 0.93,
|
||||
"knee": -1.65,
|
||||
"hip_pitch": 0.666,
|
||||
"knee": -1.546,
|
||||
}
|
||||
|
||||
# Actuator modes (MJCF native):
|
||||
# Leg joints: position PD (kp=120, kd=8), ctrl = target angle
|
||||
# Wheel joints: velocity (gain=0.5), ctrl = target velocity (rad/s)
|
||||
# Actuator modes, configured at runtime:
|
||||
# Leg joints: position PD, ctrl = target angle
|
||||
# Wheel joints: velocity, ctrl = target velocity in rad/s
|
||||
|
||||
# Control rates
|
||||
SIM_DT = 0.002 # 500 Hz (from scene.xml)
|
||||
CTRL_DT = 0.004 # 250 Hz control loop
|
||||
SIM_DT = 0.002
|
||||
CTRL_DT = 0.02
|
||||
CTRL_DECIMATION = int(CTRL_DT / SIM_DT)
|
||||
|
||||
# Wheel drive
|
||||
WHEEL_VEL_MAX = 10.0 # rad/s max wheel command
|
||||
WHEEL_VEL_MAX = 10.0
|
||||
|
||||
# Body pose control gains (for height/roll/pitch compensation)
|
||||
KP_HEIGHT = 3.0 # rad/m error → joint angle correction
|
||||
KP_ROLL = 0.5 # compensation gain
|
||||
KP_PITCH = 0.5 # compensation gain
|
||||
# Body pose control gains for height/roll/pitch compensation.
|
||||
KP_HEIGHT = 3.0
|
||||
KP_ROLL = 0.5
|
||||
KP_PITCH = 0.5
|
||||
|
||||
# Calibrated height-to-joint-angle table.
|
||||
# Constraint: avoid large wheel-center X offset from the hip/leg. This is a
|
||||
# soft support-geometry guardrail, not a strict x=0 requirement.
|
||||
# The optimizer also considers peak motor torque and RMS torque, so one hot
|
||||
# motor is not hidden by a low average across all motors.
|
||||
# Format: (height_m, hip_pitch_rad, knee_rad)
|
||||
HEIGHT_TABLE = [
|
||||
(0.17, 0.914, -2.628),
|
||||
(0.19, 0.926, -2.528),
|
||||
(0.21, 0.924, -2.428),
|
||||
(0.23, 0.912, -2.328),
|
||||
(0.25, 0.892, -2.226),
|
||||
(0.27, 0.864, -2.122),
|
||||
(0.29, 0.834, -2.014),
|
||||
(0.31, 0.798, -1.906),
|
||||
(0.33, 0.758, -1.792),
|
||||
(0.35, 0.714, -1.672),
|
||||
(0.37, 0.666, -1.546),
|
||||
(0.39, 0.612, -1.412),
|
||||
(0.41, 0.552, -1.266),
|
||||
(0.43, 0.484, -1.104),
|
||||
(0.45, 0.404, -0.918),
|
||||
]
|
||||
|
||||
# Gait parameters
|
||||
GAIT_FREQ = 2.5 # Hz
|
||||
GAIT_DUTY = 0.6 # stance fraction
|
||||
SWING_HEIGHT = 0.06 # m
|
||||
GAIT_FREQ = 2.5
|
||||
GAIT_DUTY = 0.6
|
||||
SWING_HEIGHT = 0.06
|
||||
|
||||
# Trot phase offsets: FL/RR in phase, FR/RL in phase
|
||||
PHASE_OFFSETS = {"fl": 0.0, "fr": 0.5, "rl": 0.5, "rr": 0.0}
|
||||
|
||||
@@ -14,10 +14,13 @@ from dynamics import Dynamics
|
||||
from mpc_controller import MPCController
|
||||
from config import (
|
||||
LEG_NAMES, DEFAULT_JOINT_ANGLES, WHEEL_RADIUS, WHEEL_TRACK,
|
||||
WHEEL_VEL_MAX, KP_ROLL, KP_PITCH,
|
||||
WHEEL_VEL_MAX, KP_ROLL, KP_PITCH, HEIGHT_TABLE,
|
||||
GAIT_FREQ, GAIT_DUTY, SWING_HEIGHT, PHASE_OFFSETS,
|
||||
)
|
||||
|
||||
RL_ROUGH_Q = np.array([0.0, 0.550, -1.125], dtype=float)
|
||||
LEG_STATE_IDX = np.array([0, 1, 2, 4, 5, 6, 8, 9, 10, 12, 13, 14], dtype=int)
|
||||
|
||||
|
||||
class Controller:
|
||||
"""Wheeled-legged robot controller."""
|
||||
@@ -30,7 +33,8 @@ class Controller:
|
||||
self.vel_x = 0.0 # m/s forward
|
||||
self.vel_y = 0.0 # m/s lateral
|
||||
self.yaw_rate = 0.0 # rad/s
|
||||
self.height = 0.33 # m desired body height
|
||||
self.height = 0.37 # m desired body height (wheel center under hip)
|
||||
self.wheel_posture = "table" # "table" follows height slider; "rl" matches src/robot default
|
||||
|
||||
# Mode: "wheel", "trot", or "mpc"
|
||||
self.mode = "wheel"
|
||||
@@ -49,6 +53,30 @@ class Controller:
|
||||
self._smooth_vx = 0.0
|
||||
self._smooth_vy = 0.0
|
||||
self._smooth_yaw = 0.0
|
||||
self._last_leg_targets = np.tile(
|
||||
np.array(
|
||||
[
|
||||
DEFAULT_JOINT_ANGLES["hip_abduction"],
|
||||
DEFAULT_JOINT_ANGLES["hip_pitch"],
|
||||
DEFAULT_JOINT_ANGLES["knee"],
|
||||
],
|
||||
dtype=float,
|
||||
),
|
||||
4,
|
||||
)
|
||||
|
||||
# Wheel-mode sensor feedback.
|
||||
self.yaw_rate_kp = 0.45
|
||||
self.roll_comp_gain = KP_ROLL
|
||||
self.pitch_comp_gain = KP_PITCH
|
||||
self.encoder_posture_kp = 0.12
|
||||
self.encoder_posture_max = 0.025
|
||||
self.encoder_guard_start = 0.28
|
||||
self.encoder_guard_stop = 0.65
|
||||
self.imu_guard_start = np.deg2rad(12.0)
|
||||
self.imu_guard_stop = np.deg2rad(28.0)
|
||||
self.yaw_wheel_gain = 1.0
|
||||
self.max_yaw_wheel_speed = 4.0
|
||||
|
||||
# Default leg angles
|
||||
self._default_q = np.array([
|
||||
@@ -111,7 +139,7 @@ class Controller:
|
||||
side = 1.0 if leg[1] == "l" else -1.0
|
||||
leg_targets[i*3] = side * 0.3 # fl/rl: +0.3, fr/rr: -0.3
|
||||
leg_targets[i*3+1] = 1.5 # hip pitch
|
||||
leg_targets[i*3+2] = -2.65 # knee fully folded
|
||||
leg_targets[i*3+2] = -2.65 # knee hard limit from MJCF
|
||||
return leg_targets, np.zeros(4)
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
@@ -125,35 +153,84 @@ class Controller:
|
||||
pure lateral motion). Uses hip_abduction lean for small lateral force.
|
||||
For significant lateral motion, use trot mode.
|
||||
"""
|
||||
wheel_targets = self._differential_drive(self._smooth_vx, self._smooth_yaw)
|
||||
leg_targets = self._posture_control(state)
|
||||
safe_vx, safe_yaw = self._wheel_velocity_envelope(self._smooth_vx, self._smooth_yaw)
|
||||
yaw_feedback = safe_yaw + self.yaw_rate_kp * (safe_yaw - float(state.ang_vel[2]))
|
||||
wheel_targets = self._differential_drive(safe_vx, yaw_feedback)
|
||||
wheel_targets *= self._sensor_command_scale(state, leg_targets)
|
||||
self._last_leg_targets = leg_targets.copy()
|
||||
return leg_targets, wheel_targets
|
||||
|
||||
def _posture_control(self, state: RobotState) -> np.ndarray:
|
||||
"""Leg joint targets: table-interpolated height control."""
|
||||
"""Leg joint targets from the soft wheel-X height table."""
|
||||
leg_targets = np.zeros(12)
|
||||
|
||||
# Calibrated height→angle lookup (measured from simulation)
|
||||
_H = [0.157, 0.248, 0.311, 0.366, 0.411, 0.448]
|
||||
_HIP = [1.5, 1.2, 1.0, 0.8, 0.6, 0.4]
|
||||
_KNEE = [-2.5, -2.1, -1.8, -1.5, -1.2, -0.9]
|
||||
# Calibrated height→angle lookup (minimizes motor torque at each height)
|
||||
_H = [r[0] for r in HEIGHT_TABLE]
|
||||
_HIP = [r[1] for r in HEIGHT_TABLE]
|
||||
_KNEE = [r[2] for r in HEIGHT_TABLE]
|
||||
|
||||
h_clamp = np.clip(self.height, _H[0], _H[-1])
|
||||
q_hip_base = float(np.interp(h_clamp, _H, _HIP))
|
||||
q_knee_base = float(np.interp(h_clamp, _H, _KNEE))
|
||||
q_ab_base = 0.0
|
||||
if self.wheel_posture == "rl":
|
||||
q_ab_base, q_hip_base, q_knee_base = RL_ROUGH_Q
|
||||
|
||||
roll_corr = -KP_ROLL * state.rpy[0]
|
||||
pitch_corr = -KP_PITCH * state.rpy[1]
|
||||
roll_corr = -self.roll_comp_gain * state.rpy[0]
|
||||
pitch_corr = -self.pitch_comp_gain * state.rpy[1]
|
||||
lateral_lean = 0.3 * self.vel_y
|
||||
|
||||
for i, leg in enumerate(LEG_NAMES):
|
||||
side = 1.0 if leg[1] == "l" else -1.0
|
||||
leg_targets[i*3] = np.clip(side * roll_corr + lateral_lean, -0.5, 0.5)
|
||||
leg_targets[i*3] = np.clip(q_ab_base + side * roll_corr + lateral_lean, -0.5, 0.5)
|
||||
leg_targets[i*3+1] = np.clip(q_hip_base + pitch_corr, -1.0, 2.5)
|
||||
leg_targets[i*3+2] = np.clip(q_knee_base, -2.6, -0.3)
|
||||
leg_targets[i*3+2] = np.clip(q_knee_base, -2.65, -0.3)
|
||||
|
||||
encoder_err = self._last_leg_targets - state.joint_pos[LEG_STATE_IDX]
|
||||
leg_targets += np.clip(
|
||||
self.encoder_posture_kp * encoder_err,
|
||||
-self.encoder_posture_max,
|
||||
self.encoder_posture_max,
|
||||
)
|
||||
leg_targets[0::3] = np.clip(leg_targets[0::3], -0.5, 0.5)
|
||||
leg_targets[1::3] = np.clip(leg_targets[1::3], -1.0, 2.5)
|
||||
leg_targets[2::3] = np.clip(leg_targets[2::3], -2.65, -0.3)
|
||||
return leg_targets
|
||||
|
||||
def _wheel_velocity_envelope(self, vel_x: float, yaw_rate: float) -> tuple[float, float]:
|
||||
"""Limit x/yaw combinations that are unsafe for the RL posture."""
|
||||
ax = abs(vel_x)
|
||||
if ax >= 0.8:
|
||||
yaw_lim = 0.35
|
||||
elif ax >= 0.5:
|
||||
yaw_lim = 0.55
|
||||
elif ax >= 0.25:
|
||||
yaw_lim = 0.75
|
||||
else:
|
||||
yaw_lim = 1.0
|
||||
return float(vel_x), float(np.clip(yaw_rate, -yaw_lim, yaw_lim))
|
||||
|
||||
def _sensor_command_scale(self, state: RobotState, leg_targets: np.ndarray) -> float:
|
||||
"""Back off wheels when IMU or encoder feedback says posture is degrading."""
|
||||
leg_error = float(np.max(np.abs(state.joint_pos[LEG_STATE_IDX] - leg_targets)))
|
||||
tilt = float(np.hypot(state.rpy[0], state.rpy[1]))
|
||||
scale = 1.0
|
||||
|
||||
if leg_error >= self.encoder_guard_stop:
|
||||
scale = 0.0
|
||||
elif leg_error > self.encoder_guard_start:
|
||||
span = max(1e-6, self.encoder_guard_stop - self.encoder_guard_start)
|
||||
scale *= 1.0 - (leg_error - self.encoder_guard_start) / span
|
||||
|
||||
if tilt >= self.imu_guard_stop:
|
||||
scale = 0.0
|
||||
elif tilt > self.imu_guard_start:
|
||||
span = max(1e-6, self.imu_guard_stop - self.imu_guard_start)
|
||||
scale *= 1.0 - (tilt - self.imu_guard_start) / span
|
||||
|
||||
return float(np.clip(scale, 0.0, 1.0))
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# TROT MODE
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
@@ -199,9 +276,9 @@ class Controller:
|
||||
|
||||
def _stance_leg_target(self, state: RobotState, leg: str) -> np.ndarray:
|
||||
"""Stance leg: table-interpolated height + attitude compensation."""
|
||||
_H = [0.157, 0.248, 0.311, 0.366, 0.411, 0.448]
|
||||
_HIP = [1.5, 1.2, 1.0, 0.8, 0.6, 0.4]
|
||||
_KNEE = [-2.5, -2.1, -1.8, -1.5, -1.2, -0.9]
|
||||
_H = [r[0] for r in HEIGHT_TABLE]
|
||||
_HIP = [r[1] for r in HEIGHT_TABLE]
|
||||
_KNEE = [r[2] for r in HEIGHT_TABLE]
|
||||
|
||||
h_clamp = np.clip(self.height, _H[0], _H[-1])
|
||||
q_hip = float(np.interp(h_clamp, _H, _HIP))
|
||||
@@ -215,7 +292,7 @@ class Controller:
|
||||
return np.array([
|
||||
np.clip(side * roll_corr + lateral_lean, -0.5, 0.5),
|
||||
np.clip(q_hip + pitch_corr, -1.0, 2.5),
|
||||
np.clip(q_knee, -2.6, -0.3),
|
||||
np.clip(q_knee, -2.65, -0.3),
|
||||
])
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
@@ -224,8 +301,11 @@ class Controller:
|
||||
|
||||
def _differential_drive(self, vel_x: float, yaw_rate: float) -> np.ndarray:
|
||||
"""4 wheel velocities from body commands."""
|
||||
vel_left = (vel_x - 0.5 * WHEEL_TRACK * yaw_rate) / WHEEL_RADIUS
|
||||
vel_right = (vel_x + 0.5 * WHEEL_TRACK * yaw_rate) / WHEEL_RADIUS
|
||||
linear_wheel = vel_x / WHEEL_RADIUS
|
||||
yaw_wheel = self.yaw_wheel_gain * 0.5 * WHEEL_TRACK * yaw_rate / WHEEL_RADIUS
|
||||
yaw_wheel = float(np.clip(yaw_wheel, -self.max_yaw_wheel_speed, self.max_yaw_wheel_speed))
|
||||
vel_left = linear_wheel - yaw_wheel
|
||||
vel_right = linear_wheel + yaw_wheel
|
||||
targets = np.zeros(4)
|
||||
for i, leg in enumerate(LEG_NAMES):
|
||||
targets[i] = vel_left if leg[1] == "l" else vel_right
|
||||
|
||||
@@ -40,7 +40,7 @@ class GUI:
|
||||
self._slider(cf, "Vel X", self.vel_x_var, -1.5, 1.5)
|
||||
self._slider(cf, "Vel Y*", self.vel_y_var, -0.5, 0.5)
|
||||
self._slider(cf, "Yaw", self.yaw_var, -2.0, 2.0)
|
||||
self._slider(cf, "Height", self.height_var, 0.16, 0.45)
|
||||
self._slider(cf, "Height", self.height_var, 0.17, 0.45)
|
||||
|
||||
ttk.Label(cf, text="* Vel Y: trot mode only (diff-drive can't sidestep)",
|
||||
font=("", 8)).pack(anchor="w", padx=8)
|
||||
|
||||
@@ -0,0 +1,389 @@
|
||||
"""Sweep wheel-mode IK postures for differential-drive tracking.
|
||||
|
||||
The sweep is intentionally small and reproducible:
|
||||
1. Generate ab=0 leg postures in the requested height range.
|
||||
2. Keep candidates with good static geometry from rl_friendly_opt.
|
||||
3. Simulate forward, yaw, and arc commands in MuJoCo.
|
||||
4. Rank by attitude, x-speed tracking, yaw-rate tracking, and wheel contact.
|
||||
|
||||
Usage:
|
||||
uv run python mujoco_sim/ik_diff_sweep.py --quick
|
||||
uv run python mujoco_sim/ik_diff_sweep.py --height-min 0.15 --height-max 0.42
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
|
||||
THIS_DIR = Path(__file__).resolve().parent
|
||||
if str(THIS_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(THIS_DIR))
|
||||
|
||||
from config import LEG_JOINTS, LEG_NAMES, SCENE_XML, WHEEL_JOINT, WHEEL_RADIUS # noqa: E402
|
||||
from rl_friendly_opt import get_all, rl_cost # noqa: E402
|
||||
from robot import Robot # noqa: E402
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Candidate:
|
||||
height: float
|
||||
ab: float
|
||||
hip: float
|
||||
knee: float
|
||||
static_cost: float
|
||||
r_hip_x: float
|
||||
cond: float
|
||||
max_tau: float
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Trial:
|
||||
name: str
|
||||
vx: float
|
||||
yaw_rate: float
|
||||
duration: float
|
||||
|
||||
|
||||
def wrap_pi(x: float) -> float:
|
||||
return (x + math.pi) % (2.0 * math.pi) - math.pi
|
||||
|
||||
|
||||
def body_track(model: mujoco.MjModel, data: mujoco.MjData) -> float:
|
||||
wheel_bids = []
|
||||
for leg in LEG_NAMES:
|
||||
bid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, f"{leg}_wheel_Link")
|
||||
wheel_bids.append(bid)
|
||||
y = data.xipos[wheel_bids, 1]
|
||||
return float(np.mean(y[[0, 2]]) - np.mean(y[[1, 3]]))
|
||||
|
||||
|
||||
def build_maps(robot: Robot) -> tuple[dict[str, int], dict[str, int], dict[str, int]]:
|
||||
act: dict[str, int] = {}
|
||||
qadr: dict[str, int] = {}
|
||||
vadr: dict[str, int] = {}
|
||||
for leg in LEG_NAMES:
|
||||
for jt in (*LEG_JOINTS, WHEEL_JOINT):
|
||||
name = f"{leg}_{jt}"
|
||||
act[name] = mujoco.mj_name2id(robot.model, mujoco.mjtObj.mjOBJ_ACTUATOR, name)
|
||||
jid = mujoco.mj_name2id(robot.model, mujoco.mjtObj.mjOBJ_JOINT, name)
|
||||
qadr[name] = robot.model.jnt_qposadr[jid]
|
||||
vadr[name] = robot.model.jnt_dofadr[jid]
|
||||
return act, qadr, vadr
|
||||
|
||||
|
||||
def set_posture(robot: Robot, cand: Candidate, act: dict[str, int], qadr: dict[str, int]) -> None:
|
||||
mujoco.mj_resetData(robot.model, robot.data)
|
||||
for leg in LEG_NAMES:
|
||||
side_ab = cand.ab if leg[1] == "l" else -cand.ab
|
||||
for jt, val in zip(LEG_JOINTS, (side_ab, cand.hip, cand.knee)):
|
||||
name = f"{leg}_{jt}"
|
||||
robot.data.qpos[qadr[name]] = val
|
||||
robot.data.ctrl[act[name]] = val
|
||||
robot.data.ctrl[act[f"{leg}_{WHEEL_JOINT}"]] = 0.0
|
||||
robot.data.qpos[:3] = [0.0, 0.0, max(0.25, cand.height + 0.08)]
|
||||
robot.data.qpos[3:7] = [1.0, 0.0, 0.0, 0.0]
|
||||
robot.data.qvel[:] = 0.0
|
||||
mujoco.mj_forward(robot.model, robot.data)
|
||||
|
||||
|
||||
def wheel_targets(
|
||||
vx: float,
|
||||
yaw_rate: float,
|
||||
track: float,
|
||||
max_wheel: float,
|
||||
yaw_gain: float,
|
||||
wheel_model: str,
|
||||
linear_gain: float,
|
||||
direct_yaw_gain: float,
|
||||
wheel_signs: np.ndarray,
|
||||
) -> np.ndarray:
|
||||
if wheel_model == "direct":
|
||||
left = linear_gain * vx - direct_yaw_gain * yaw_rate
|
||||
right = linear_gain * vx + direct_yaw_gain * yaw_rate
|
||||
else:
|
||||
left = (vx - yaw_gain * 0.5 * track * yaw_rate) / WHEEL_RADIUS
|
||||
right = (vx + yaw_gain * 0.5 * track * yaw_rate) / WHEEL_RADIUS
|
||||
raw = np.array([left, right, left, right], dtype=float)
|
||||
return np.clip(raw * wheel_signs, -max_wheel, max_wheel)
|
||||
|
||||
|
||||
def run_trial(robot: Robot, cand: Candidate, trial: Trial, args: argparse.Namespace) -> dict:
|
||||
act, qadr, vadr = build_maps(robot)
|
||||
set_posture(robot, cand, act, qadr)
|
||||
|
||||
ctrl_dt = args.control_dt
|
||||
sim_dt = robot.model.opt.timestep
|
||||
steps_per_ctrl = max(1, int(round(ctrl_dt / sim_dt)))
|
||||
track = body_track(robot.model, robot.data) if args.track_source == "model" else args.track_width
|
||||
wheel_signs = np.array(args.wheel_signs, dtype=float)
|
||||
wheel_cmd = wheel_targets(
|
||||
trial.vx,
|
||||
trial.yaw_rate,
|
||||
track,
|
||||
args.max_wheel_speed,
|
||||
args.yaw_gain,
|
||||
args.wheel_model,
|
||||
args.linear_gain,
|
||||
args.direct_yaw_gain,
|
||||
wheel_signs,
|
||||
)
|
||||
|
||||
for _ in range(int(round(args.settle / sim_dt))):
|
||||
for leg in LEG_NAMES:
|
||||
vals = (cand.ab if leg[1] == "l" else -cand.ab, cand.hip, cand.knee)
|
||||
for jt, val in zip(LEG_JOINTS, vals):
|
||||
robot.data.ctrl[act[f"{leg}_{jt}"]] = val
|
||||
robot.data.ctrl[act[f"{leg}_{WHEEL_JOINT}"]] = 0.0
|
||||
robot.step()
|
||||
|
||||
state0 = robot.get_state()
|
||||
yaw0 = float(state0.rpy[2])
|
||||
x0 = float(state0.pos[0])
|
||||
|
||||
max_roll = 0.0
|
||||
max_pitch = 0.0
|
||||
max_tilt = 0.0
|
||||
max_wheel_air = -1e9
|
||||
wheel_err_sum = 0.0
|
||||
samples = 0
|
||||
max_leg_err = 0.0
|
||||
body_vx_sum = 0.0
|
||||
yaw_unwrapped = 0.0
|
||||
last_yaw = yaw0
|
||||
leg_target = np.array([cand.ab, cand.hip, cand.knee] * 4, dtype=float)
|
||||
|
||||
total_steps = int(round(trial.duration / sim_dt))
|
||||
cmd = np.zeros(4, dtype=float)
|
||||
max_delta = args.wheel_accel_limit * ctrl_dt
|
||||
wheel_body_ids = [
|
||||
mujoco.mj_name2id(robot.model, mujoco.mjtObj.mjOBJ_BODY, f"{leg}_wheel_Link")
|
||||
for leg in LEG_NAMES
|
||||
]
|
||||
wheel_vadr = np.array([vadr[f"{leg}_{WHEEL_JOINT}"] for leg in LEG_NAMES], dtype=int)
|
||||
leg_qadr = np.array([qadr[f"{leg}_{jt}"] for leg in LEG_NAMES for jt in LEG_JOINTS], dtype=int)
|
||||
|
||||
for step in range(total_steps):
|
||||
if step % steps_per_ctrl == 0:
|
||||
cmd = cmd + np.clip(wheel_cmd - cmd, -max_delta, max_delta)
|
||||
for leg in LEG_NAMES:
|
||||
vals = (cand.ab if leg[1] == "l" else -cand.ab, cand.hip, cand.knee)
|
||||
for jt, val in zip(LEG_JOINTS, vals):
|
||||
robot.data.ctrl[act[f"{leg}_{jt}"]] = val
|
||||
for i, leg in enumerate(LEG_NAMES):
|
||||
robot.data.ctrl[act[f"{leg}_{WHEEL_JOINT}"]] = cmd[i]
|
||||
|
||||
robot.step()
|
||||
|
||||
if step % steps_per_ctrl == 0:
|
||||
st = robot.get_state()
|
||||
roll, pitch = float(st.rpy[0]), float(st.rpy[1])
|
||||
yaw_now = float(st.rpy[2])
|
||||
yaw_unwrapped += wrap_pi(yaw_now - last_yaw)
|
||||
last_yaw = yaw_now
|
||||
body_vx_sum += float(st.rot[:, 0].dot(st.lin_vel))
|
||||
max_roll = max(max_roll, abs(roll))
|
||||
max_pitch = max(max_pitch, abs(pitch))
|
||||
max_tilt = max(max_tilt, math.hypot(roll, pitch))
|
||||
wheel_air = robot.data.xipos[wheel_body_ids, 2] - WHEEL_RADIUS
|
||||
max_wheel_air = max(max_wheel_air, float(np.max(wheel_air)))
|
||||
wheel_err_sum += float(np.mean(np.abs(robot.data.qvel[wheel_vadr] - cmd)))
|
||||
max_leg_err = max(max_leg_err, float(np.max(np.abs(robot.data.qpos[leg_qadr] - leg_target))))
|
||||
samples += 1
|
||||
|
||||
st = robot.get_state()
|
||||
elapsed = max(1e-6, float(st.time - state0.time))
|
||||
world_x_rate = (float(st.pos[0]) - x0) / elapsed
|
||||
x_rate = body_vx_sum / max(1, samples)
|
||||
yaw_rate = yaw_unwrapped / max(1e-6, samples * steps_per_ctrl * sim_dt)
|
||||
x_err = abs(x_rate - trial.vx)
|
||||
yaw_err = abs(yaw_rate - trial.yaw_rate)
|
||||
return {
|
||||
"trial": trial.name,
|
||||
"x_rate": x_rate,
|
||||
"world_x_rate": world_x_rate,
|
||||
"yaw_rate": yaw_rate,
|
||||
"x_err": x_err,
|
||||
"yaw_err": yaw_err,
|
||||
"max_roll_deg": math.degrees(max_roll),
|
||||
"max_pitch_deg": math.degrees(max_pitch),
|
||||
"max_tilt_deg": math.degrees(max_tilt),
|
||||
"max_wheel_air_m": max_wheel_air,
|
||||
"mean_wheel_err": wheel_err_sum / max(1, samples),
|
||||
"max_leg_err": max_leg_err,
|
||||
"track": track,
|
||||
"wheel_cmd": [float(x) for x in wheel_cmd],
|
||||
}
|
||||
|
||||
|
||||
def generate_candidates(args: argparse.Namespace) -> list[Candidate]:
|
||||
if args.fixed_hip is not None or args.fixed_knee is not None:
|
||||
if args.fixed_hip is None or args.fixed_knee is None:
|
||||
raise SystemExit("--fixed-hip and --fixed-knee must be provided together")
|
||||
ab = float(args.fixed_ab)
|
||||
hip = float(args.fixed_hip)
|
||||
knee = float(args.fixed_knee)
|
||||
r = get_all(ab, hip, knee)
|
||||
return [
|
||||
Candidate(
|
||||
height=float(r["z"]),
|
||||
ab=ab,
|
||||
hip=hip,
|
||||
knee=knee,
|
||||
static_cost=float(rl_cost(r)),
|
||||
r_hip_x=float(r["r_hip_x_mag"]),
|
||||
cond=float(r["cond"]),
|
||||
max_tau=float(r["max_tau"]),
|
||||
)
|
||||
]
|
||||
|
||||
cands: list[Candidate] = []
|
||||
h_targets = np.arange(args.height_min, args.height_max + 0.5 * args.height_step, args.height_step)
|
||||
ab_values = [0.0] if args.ab_max <= 1e-9 else np.arange(0.0, args.ab_max + 1e-9, args.ab_step)
|
||||
hip_values = np.arange(args.hip_min, args.hip_max + 0.5 * args.hip_step, args.hip_step)
|
||||
knee_values = np.arange(args.knee_min, args.knee_max + 0.5 * args.knee_step, args.knee_step)
|
||||
for ht in h_targets:
|
||||
bucket: list[Candidate] = []
|
||||
for ab in ab_values:
|
||||
for hip in hip_values:
|
||||
for knee in knee_values:
|
||||
r = get_all(float(ab), float(hip), float(knee))
|
||||
if float(r["z"]) < args.height_min or float(r["z"]) > args.height_max:
|
||||
continue
|
||||
if abs(float(r["z"]) - float(ht)) > args.height_tol:
|
||||
continue
|
||||
if r["wz"] >= r["kz"]:
|
||||
continue
|
||||
if abs(r["r_hip_x_mag"]) > args.max_wheel_x:
|
||||
continue
|
||||
cost = float(rl_cost(r))
|
||||
bucket.append(
|
||||
Candidate(
|
||||
height=float(r["z"]),
|
||||
ab=float(ab),
|
||||
hip=float(hip),
|
||||
knee=float(knee),
|
||||
static_cost=cost,
|
||||
r_hip_x=float(r["r_hip_x_mag"]),
|
||||
cond=float(r["cond"]),
|
||||
max_tau=float(r["max_tau"]),
|
||||
)
|
||||
)
|
||||
bucket.sort(key=lambda c: c.static_cost)
|
||||
cands.extend(bucket[: args.top_per_height])
|
||||
return cands
|
||||
|
||||
|
||||
def score_result(cand: Candidate, trials: list[dict]) -> float:
|
||||
score = 0.08 * cand.static_cost
|
||||
for t in trials:
|
||||
score += 8.0 * t["x_err"]
|
||||
score += 10.0 * t["yaw_err"]
|
||||
score += 0.08 * t["max_tilt_deg"]
|
||||
score += 0.03 * max(0.0, t["max_pitch_deg"] - 8.0) ** 2
|
||||
score += 20.0 * max(0.0, t["max_wheel_air_m"] - 0.015)
|
||||
score += 1.5 * t["mean_wheel_err"]
|
||||
score += 2.0 * max(0.0, t["max_leg_err"] - 0.35)
|
||||
return float(score)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--height-min", type=float, default=0.15)
|
||||
parser.add_argument("--height-max", type=float, default=0.42)
|
||||
parser.add_argument("--height-step", type=float, default=0.01)
|
||||
parser.add_argument("--height-tol", type=float, default=0.004)
|
||||
parser.add_argument("--top-per-height", type=int, default=1)
|
||||
parser.add_argument("--ab-max", type=float, default=0.0)
|
||||
parser.add_argument("--ab-step", type=float, default=0.04)
|
||||
parser.add_argument("--hip-min", type=float, default=0.25)
|
||||
parser.add_argument("--hip-max", type=float, default=1.05)
|
||||
parser.add_argument("--hip-step", type=float, default=0.025)
|
||||
parser.add_argument("--knee-min", type=float, default=-2.65)
|
||||
parser.add_argument("--knee-max", type=float, default=-0.85)
|
||||
parser.add_argument("--knee-step", type=float, default=0.025)
|
||||
parser.add_argument("--max-wheel-x", type=float, default=0.09)
|
||||
parser.add_argument("--fixed-ab", type=float, default=0.0)
|
||||
parser.add_argument("--fixed-hip", type=float, default=None)
|
||||
parser.add_argument("--fixed-knee", type=float, default=None)
|
||||
parser.add_argument("--duration", type=float, default=4.0)
|
||||
parser.add_argument("--settle", type=float, default=1.5)
|
||||
parser.add_argument("--control-dt", type=float, default=0.02)
|
||||
parser.add_argument("--vx", type=float, default=0.6)
|
||||
parser.add_argument("--yaw", type=float, default=0.3)
|
||||
parser.add_argument("--arc-yaw", type=float, default=0.15)
|
||||
parser.add_argument("--yaw-gain", type=float, default=1.0)
|
||||
parser.add_argument("--wheel-model", choices=("diff", "direct"), default="diff")
|
||||
parser.add_argument("--linear-gain", type=float, default=12.5)
|
||||
parser.add_argument("--direct-yaw-gain", type=float, default=8.0)
|
||||
parser.add_argument("--max-wheel-speed", type=float, default=12.0)
|
||||
parser.add_argument("--wheel-accel-limit", type=float, default=35.0)
|
||||
parser.add_argument("--track-source", choices=("model", "fixed"), default="model")
|
||||
parser.add_argument("--track-width", type=float, default=0.394)
|
||||
parser.add_argument(
|
||||
"--wheel-signs",
|
||||
type=float,
|
||||
nargs=4,
|
||||
default=[1.0, 1.0, 1.0, 1.0],
|
||||
metavar=("FL", "FR", "RL", "RR"),
|
||||
help="Per-wheel velocity sign multipliers in joint order.",
|
||||
)
|
||||
parser.add_argument("--quick", action="store_true")
|
||||
parser.add_argument("--json", type=Path, default=None)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
args = parse_args()
|
||||
if args.quick:
|
||||
args.height_step = 0.02
|
||||
args.hip_step = 0.05
|
||||
args.knee_step = 0.05
|
||||
|
||||
candidates = generate_candidates(args)
|
||||
if not candidates:
|
||||
raise SystemExit("No candidates found")
|
||||
|
||||
trials = [
|
||||
Trial("forward", args.vx, 0.0, args.duration),
|
||||
Trial("yaw", 0.0, args.yaw, args.duration),
|
||||
Trial("arc", args.vx, args.arc_yaw, args.duration),
|
||||
]
|
||||
robot = Robot(SCENE_XML)
|
||||
rows = []
|
||||
for i, cand in enumerate(candidates, 1):
|
||||
trial_rows = [run_trial(robot, cand, t, args) for t in trials]
|
||||
rows.append({"candidate": cand.__dict__, "trials": trial_rows, "score": score_result(cand, trial_rows)})
|
||||
if i % 10 == 0:
|
||||
print(f"tested {i}/{len(candidates)}")
|
||||
|
||||
rows.sort(key=lambda r: r["score"])
|
||||
if args.json:
|
||||
args.json.write_text(json.dumps(rows, indent=2), encoding="utf-8")
|
||||
|
||||
print("Top IK postures for differential drive tracking")
|
||||
print("rank score height ab hip knee static xhip cond tau | fwd_x yaw_wz arc_x arc_wz max_tilt max_pitch")
|
||||
for rank, row in enumerate(rows[:10], 1):
|
||||
c = row["candidate"]
|
||||
by = {t["trial"]: t for t in row["trials"]}
|
||||
max_tilt = max(t["max_tilt_deg"] for t in row["trials"])
|
||||
max_pitch = max(t["max_pitch_deg"] for t in row["trials"])
|
||||
print(
|
||||
f"{rank:>2} {row['score']:>7.2f} {c['height']:.3f} {c['ab']:.2f} {c['hip']:.3f} {c['knee']:.3f} "
|
||||
f"{c['static_cost']:.1f} {c['r_hip_x']:.3f} {c['cond']:.2f} {c['max_tau']:.2f} | "
|
||||
f"{by['forward']['x_rate']:.3f} {by['yaw']['yaw_rate']:.3f} {by['arc']['x_rate']:.3f} {by['arc']['yaw_rate']:.3f} "
|
||||
f"{max_tilt:.1f} {max_pitch:.1f}"
|
||||
)
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -5,7 +5,7 @@ Integrates: gait scheduler + reference trajectory + ConvexMPC solver +
|
||||
|
||||
Architecture (following go2-convex-mpc):
|
||||
- MPC runs at ~50 Hz (every MPC_DECIMATION control steps)
|
||||
- Swing/stance leg controller runs at control rate (250 Hz)
|
||||
- Swing/stance leg controller runs at control rate (50 Hz)
|
||||
- Wheel drive: stance legs use differential drive, swing legs coast
|
||||
"""
|
||||
|
||||
@@ -20,7 +20,7 @@ from config import (
|
||||
)
|
||||
|
||||
# MPC update decimation (relative to control loop)
|
||||
MPC_DECIMATION = max(1, int(MPC_DT / CTRL_DT)) # ~5 steps at 250Hz
|
||||
MPC_DECIMATION = max(1, int(MPC_DT / CTRL_DT)) # 1 step at 50Hz
|
||||
|
||||
|
||||
class MPCController:
|
||||
|
||||
@@ -0,0 +1,367 @@
|
||||
"""Posture optimizer for wheeled-leg standing/crawl height table.
|
||||
|
||||
The table is not a pure "minimum average torque" table. For crawl and low-bar
|
||||
traversal, the wheel center should not be far from the hip/leg in sagittal X,
|
||||
otherwise the robot is no longer really using the wheel as the support/drive
|
||||
point. This is a soft guardrail, not a strict x=0 constraint. The score combines:
|
||||
|
||||
1. wheel center X offset from hip
|
||||
2. peak single-motor holding torque
|
||||
3. RMS torque, used as a proxy for I^2R heating
|
||||
|
||||
Usage:
|
||||
python posture_optimizer.py # MuJoCo sweep
|
||||
python posture_optimizer.py --quick # coarse MuJoCo sweep
|
||||
python posture_optimizer.py --analyze # analytical-only sweep
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import math
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
try:
|
||||
import mujoco
|
||||
except ImportError:
|
||||
mujoco = None
|
||||
|
||||
SCENE_XML = REPO_ROOT / "mjcf" / "scene.xml"
|
||||
WHEEL_RADIUS = 0.10
|
||||
HIP_Z_OFFSET = 0.054
|
||||
L1, L2 = 0.25, 0.20
|
||||
ROBOT_MASS = 12.3
|
||||
G = 9.81
|
||||
F_PER_LEG = ROBOT_MASS * G / 4.0
|
||||
MAX_TORQUE = 17.0
|
||||
|
||||
LEG_NAMES = ("fl", "fr", "rl", "rr")
|
||||
LEG_JOINTS = ("hip_abduction_joint", "hip_pitch_joint", "knee_joint")
|
||||
|
||||
# 0.15m is not a good default table target for ab=0.0. 0.17m is the practical
|
||||
# default crawl height, while lower crawl can be evaluated with abduction.
|
||||
KNEE_MIN = -2.65
|
||||
HEIGHT_MIN = 0.17
|
||||
HEIGHT_MAX = 0.46
|
||||
SOFT_WHEEL_X_OFFSET = 0.05
|
||||
HARD_WHEEL_X_OFFSET = 0.08
|
||||
|
||||
_ACTUATOR_NAMES = [f"{leg}_{jt}" for leg in LEG_NAMES for jt in LEG_JOINTS]
|
||||
|
||||
|
||||
def compute_fk(hip, knee):
|
||||
"""Return wheel-center x offset and base height from (hip_pitch, knee)."""
|
||||
x = L1 * math.sin(hip) + L2 * math.sin(hip + knee)
|
||||
z = L1 * math.cos(hip) + L2 * math.cos(hip + knee)
|
||||
base_height = WHEEL_RADIUS + z - HIP_Z_OFFSET
|
||||
return x, base_height
|
||||
|
||||
|
||||
def posture_cost(x_foot, torques):
|
||||
"""Score one posture by support geometry, peak torque, and RMS torque."""
|
||||
tau = np.asarray(torques, dtype=float)
|
||||
peak_torque = float(np.max(np.abs(tau)))
|
||||
rms_torque = float(np.sqrt(np.mean(np.square(tau))))
|
||||
mean_i2r = float(np.mean(np.square(tau)))
|
||||
x_penalty = max(0.0, abs(x_foot) - SOFT_WHEEL_X_OFFSET)
|
||||
cost = (
|
||||
0.5 * (abs(x_foot) / SOFT_WHEEL_X_OFFSET) ** 2
|
||||
+ 8.0 * (x_penalty / max(1e-6, HARD_WHEEL_X_OFFSET - SOFT_WHEEL_X_OFFSET)) ** 2
|
||||
+ 3.0 * (peak_torque / MAX_TORQUE) ** 2
|
||||
+ (rms_torque / MAX_TORQUE) ** 2
|
||||
)
|
||||
return cost, peak_torque, rms_torque, mean_i2r
|
||||
|
||||
|
||||
def analyze_analytical():
|
||||
"""Analytical sweep using static GRF moments."""
|
||||
hip_range = np.arange(0.0, 1.6, 0.002)
|
||||
knee_range = np.arange(KNEE_MIN, -0.4, 0.002)
|
||||
|
||||
results = []
|
||||
for hip in hip_range:
|
||||
for knee in knee_range:
|
||||
x_foot, height = compute_fk(hip, knee)
|
||||
if height < HEIGHT_MIN or height > HEIGHT_MAX:
|
||||
continue
|
||||
|
||||
tau_hip = F_PER_LEG * x_foot
|
||||
x_knee_to_foot = L2 * math.sin(hip + knee)
|
||||
tau_knee = F_PER_LEG * x_knee_to_foot
|
||||
tau_abduction = 0.0
|
||||
cost, peak, rms, mean_i2r = posture_cost(
|
||||
x_foot, (tau_abduction, tau_hip, tau_knee))
|
||||
|
||||
results.append({
|
||||
"hip": float(hip),
|
||||
"knee": float(knee),
|
||||
"height": float(height),
|
||||
"x_foot": float(x_foot),
|
||||
"tau_abd": tau_abduction,
|
||||
"tau_hip": float(tau_hip),
|
||||
"tau_knee": float(tau_knee),
|
||||
"tau_peak": peak,
|
||||
"tau_rms": rms,
|
||||
"mean_i2r": mean_i2r,
|
||||
"cost": cost,
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def run_mujoco_sweep(quick=False):
|
||||
"""MuJoCo sweep measuring actual actuator forces at steady state."""
|
||||
if mujoco is None:
|
||||
raise RuntimeError("mujoco is not installed; use --analyze for analytical mode")
|
||||
|
||||
model = mujoco.MjModel.from_xml_path(str(SCENE_XML))
|
||||
data = mujoco.MjData(model)
|
||||
|
||||
act_ids = {}
|
||||
for name in _ACTUATOR_NAMES:
|
||||
act_ids[name] = mujoco.mj_name2id(
|
||||
model, mujoco.mjtObj.mjOBJ_ACTUATOR, name)
|
||||
|
||||
for i in range(model.nu):
|
||||
name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_ACTUATOR, i)
|
||||
model.actuator_biastype[i] = 1
|
||||
model.actuator_gaintype[i] = 0
|
||||
model.actuator_forcelimited[i] = 0
|
||||
if "wheel" not in name:
|
||||
model.actuator_gainprm[i, 0] = 60.0
|
||||
model.actuator_biasprm[i, 0] = 0.0
|
||||
model.actuator_biasprm[i, 1] = -60.0
|
||||
model.actuator_biasprm[i, 2] = -3.0
|
||||
model.actuator_ctrlrange[i] = [-3.14, 3.14]
|
||||
else:
|
||||
model.actuator_gainprm[i, 0] = 2.0
|
||||
model.actuator_biasprm[i, 0] = 0.0
|
||||
model.actuator_biasprm[i, 1] = 0.0
|
||||
model.actuator_biasprm[i, 2] = -2.0
|
||||
model.actuator_ctrlrange[i] = [-20.0, 20.0]
|
||||
|
||||
jnt_ids = {}
|
||||
for name in _ACTUATOR_NAMES:
|
||||
jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, name)
|
||||
jnt_ids[name] = model.jnt_qposadr[jid]
|
||||
|
||||
if quick:
|
||||
hip_range = np.arange(0.3, 1.5, 0.10)
|
||||
knee_range = np.arange(KNEE_MIN, -0.6, 0.10)
|
||||
else:
|
||||
hip_range = np.arange(0.0, 1.6, 0.04)
|
||||
knee_range = np.arange(KNEE_MIN, -0.4, 0.04)
|
||||
|
||||
results = []
|
||||
total = 0
|
||||
valid = 0
|
||||
|
||||
for hip in hip_range:
|
||||
for knee in knee_range:
|
||||
total += 1
|
||||
x_foot, height = compute_fk(hip, knee)
|
||||
if height < HEIGHT_MIN or height > HEIGHT_MAX:
|
||||
continue
|
||||
|
||||
mujoco.mj_resetData(model, data)
|
||||
|
||||
for leg in LEG_NAMES:
|
||||
for jt, val in zip(LEG_JOINTS, (0.0, hip, knee)):
|
||||
name = f"{leg}_{jt}"
|
||||
data.qpos[jnt_ids[name]] = val
|
||||
|
||||
data.qpos[2] = height
|
||||
data.qpos[3] = 1.0
|
||||
data.qpos[4:7] = 0.0
|
||||
|
||||
for leg in LEG_NAMES:
|
||||
for jt, val in zip(LEG_JOINTS, (0.0, hip, knee)):
|
||||
name = f"{leg}_{jt}"
|
||||
data.ctrl[act_ids[name]] = val
|
||||
name_w = f"{leg}_wheel_joint"
|
||||
if name_w in act_ids:
|
||||
data.ctrl[act_ids[name_w]] = 0.0
|
||||
|
||||
mujoco.mj_forward(model, data)
|
||||
|
||||
for _ in range(500):
|
||||
mujoco.mj_step(model, data)
|
||||
|
||||
roll, pitch, _ = _get_rpy(data, model)
|
||||
if abs(roll) > 0.8 or abs(pitch) > 0.8:
|
||||
continue
|
||||
|
||||
torque_buf = []
|
||||
for _ in range(100):
|
||||
mujoco.mj_step(model, data)
|
||||
torque_buf.append([data.actuator_force[act_ids[name]]
|
||||
for name in _ACTUATOR_NAMES])
|
||||
tau_avg = np.array(torque_buf).mean(axis=0)
|
||||
|
||||
tau_hip_val = tau_avg[1]
|
||||
tau_knee_val = tau_avg[2]
|
||||
tau_abd_val = tau_avg[0]
|
||||
cost, peak, rms, mean_i2r = posture_cost(x_foot, tau_avg)
|
||||
|
||||
q_hip_actual = float(data.qpos[jnt_ids["fl_hip_pitch_joint"]])
|
||||
q_knee_actual = float(data.qpos[jnt_ids["fl_knee_joint"]])
|
||||
|
||||
valid += 1
|
||||
results.append({
|
||||
"hip": float(hip),
|
||||
"knee": float(knee),
|
||||
"height": float(f"{height:.4f}"),
|
||||
"x_foot": float(f"{x_foot:.4f}"),
|
||||
"tau_abd": float(f"{tau_abd_val:.4f}"),
|
||||
"tau_hip": float(f"{tau_hip_val:.4f}"),
|
||||
"tau_knee": float(f"{tau_knee_val:.4f}"),
|
||||
"tau_peak": float(f"{peak:.4f}"),
|
||||
"tau_rms": float(f"{rms:.4f}"),
|
||||
"mean_i2r": float(f"{mean_i2r:.4f}"),
|
||||
"cost": float(f"{cost:.4f}"),
|
||||
"q_hip_actual": float(f"{q_hip_actual:.4f}"),
|
||||
"q_knee_actual": float(f"{q_knee_actual:.4f}"),
|
||||
})
|
||||
|
||||
if valid % 20 == 0:
|
||||
print(f" [{valid}/{total}] hip={hip:.2f} knee={knee:.2f} "
|
||||
f"h={height:.3f} x={x_foot:+.4f} "
|
||||
f"peak={peak:.2f} rms={rms:.2f} cost={cost:.4f}")
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def _get_rpy(data, model):
|
||||
"""Extract roll and pitch from MuJoCo data."""
|
||||
base_bid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "base_link")
|
||||
rot = data.xmat[base_bid].reshape(3, 3)
|
||||
roll = math.atan2(rot[2, 1], rot[2, 2])
|
||||
pitch = math.atan2(-rot[2, 0], math.sqrt(rot[2, 1] ** 2 + rot[2, 2] ** 2))
|
||||
return roll, pitch, 0.0
|
||||
|
||||
|
||||
def print_top_results(results, n=10):
|
||||
"""Print the top-N results with lowest cost."""
|
||||
sorted_r = sorted(results, key=lambda r: r["cost"])
|
||||
|
||||
print(f"\n{'=' * 104}")
|
||||
print(f"TOP {n} CONFIGURATIONS (soft wheel-X + peak/RMS torque score)")
|
||||
print(f"{'=' * 104}")
|
||||
print(f"{'Rank':>4} {'hip':>6} {'knee':>7} {'height':>7} {'x_foot':>8} "
|
||||
f"{'tau_abd':>8} {'tau_hip':>8} {'tau_knee':>8} "
|
||||
f"{'peak':>8} {'rms':>8} {'cost':>9}")
|
||||
print(f"{'-' * 104}")
|
||||
|
||||
for i, r in enumerate(sorted_r[:n]):
|
||||
print(f"{i + 1:>4} {r['hip']:>6.3f} {r['knee']:>7.3f} "
|
||||
f"{r['height']:>7.3f} {r.get('x_foot', 0):>8.4f} "
|
||||
f"{r.get('tau_abd', 0):>8.3f} {r['tau_hip']:>8.3f} "
|
||||
f"{r['tau_knee']:>8.3f} {r.get('tau_peak', 0):>8.3f} "
|
||||
f"{r.get('tau_rms', 0):>8.3f} {r['cost']:>9.4f}")
|
||||
|
||||
best = sorted_r[0]
|
||||
print(f"\nBEST: hip={best['hip']:.3f} knee={best['knee']:.3f} "
|
||||
f"z={best['height']:.3f}m x={best.get('x_foot', 0):+.4f}m "
|
||||
f"peak={best.get('tau_peak', 0):.3f}Nm "
|
||||
f"rms={best.get('tau_rms', 0):.3f}Nm cost={best['cost']:.4f}\n")
|
||||
|
||||
return sorted_r
|
||||
|
||||
|
||||
def compute_height_table(results):
|
||||
"""Build height-to-angle lookup with soft wheel-X support guardrail."""
|
||||
sorted_r = sorted(
|
||||
(r for r in results if abs(r.get("x_foot", 999.0)) <= HARD_WHEEL_X_OFFSET),
|
||||
key=lambda r: r["height"],
|
||||
)
|
||||
if not sorted_r:
|
||||
raise RuntimeError("No candidates satisfy HARD_WHEEL_X_OFFSET")
|
||||
|
||||
h_range = np.arange(0.17, 0.46, 0.02)
|
||||
table_h, table_hip, table_knee = [], [], []
|
||||
|
||||
for h_target in h_range:
|
||||
candidates = [(r, abs(r["height"] - h_target)) for r in sorted_r]
|
||||
candidates.sort(key=lambda x: (x[1], x[0]["cost"]))
|
||||
best = candidates[0][0]
|
||||
table_h.append(best["height"])
|
||||
table_hip.append(best["hip"])
|
||||
table_knee.append(best["knee"])
|
||||
|
||||
return {
|
||||
"height": [round(h, 3) for h in table_h],
|
||||
"hip": [round(h, 3) for h in table_hip],
|
||||
"knee": [round(k, 3) for k in table_knee],
|
||||
}
|
||||
|
||||
|
||||
def export_calibrated_table(table):
|
||||
"""Print the new height table in copy-paste format."""
|
||||
print(f"\n{'=' * 80}")
|
||||
print("CALIBRATED HEIGHT TABLE (soft wheel-X support guardrail)")
|
||||
print(f"{'=' * 80}")
|
||||
print(f"_H = {table['height']}")
|
||||
print(f"_HIP = {table['hip']}")
|
||||
print(f"_KNEE = {table['knee']}")
|
||||
print(f"{'=' * 80}\n")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Find wheeled-leg posture table")
|
||||
parser.add_argument("--quick", action="store_true", help="Coarse MuJoCo sweep")
|
||||
parser.add_argument("--analyze", action="store_true", help="Analytical only")
|
||||
parser.add_argument("--mujoco", action="store_true", default=True,
|
||||
help="Run MuJoCo simulation when available")
|
||||
args = parser.parse_args()
|
||||
|
||||
print("=" * 72)
|
||||
print("WHEELED-LEG POSTURE OPTIMIZER")
|
||||
print("=" * 72)
|
||||
print(f"Robot mass: {ROBOT_MASS} kg, F_per_leg: {F_PER_LEG:.1f} N")
|
||||
print(f"L1={L1}m, L2={L2}m, wheel_r={WHEEL_RADIUS}m")
|
||||
print(f"Height range: [{HEIGHT_MIN}, {HEIGHT_MAX}] m")
|
||||
print(f"Knee min hard limit: {KNEE_MIN} rad")
|
||||
print(f"Soft wheel X offset: {SOFT_WHEEL_X_OFFSET} m")
|
||||
print(f"Hard wheel X offset: {HARD_WHEEL_X_OFFSET} m\n")
|
||||
|
||||
t0 = time.time()
|
||||
if args.analyze or mujoco is None:
|
||||
print("[Analytical mode]")
|
||||
results = analyze_analytical()
|
||||
else:
|
||||
print("[MuJoCo simulation mode]")
|
||||
results = run_mujoco_sweep(quick=args.quick)
|
||||
|
||||
elapsed = time.time() - t0
|
||||
print(f"Evaluated {len(results)} valid configurations in {elapsed:.1f}s")
|
||||
if not results:
|
||||
print("No valid configurations found")
|
||||
return
|
||||
|
||||
best_results = print_top_results(results, n=15)
|
||||
|
||||
x_def, h_def = compute_fk(0.666, -1.546)
|
||||
print(f"Current config default: hip=0.666, knee=-1.546 "
|
||||
f"=> z={h_def:.3f}m, x={x_def:+.4f}m")
|
||||
|
||||
table = compute_height_table(best_results)
|
||||
export_calibrated_table(table)
|
||||
|
||||
best = best_results[0]
|
||||
print("=" * 72)
|
||||
print("RECOMMENDED DEFAULT")
|
||||
print("=" * 72)
|
||||
print(f"hip_abduction: 0.0")
|
||||
print(f"hip_pitch: {best['hip']:.3f}")
|
||||
print(f"knee: {best['knee']:.3f}")
|
||||
print(f"height: {best['height']:.3f} m")
|
||||
print(f"x_foot: {best.get('x_foot', 0):+.4f} m")
|
||||
print(f"tau_peak: {best.get('tau_peak', 0):.3f} Nm")
|
||||
print(f"tau_rms: {best.get('tau_rms', 0):.3f} Nm")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,189 @@
|
||||
"""Fast posture table helper based on wheelleg.xml link offsets.
|
||||
|
||||
This tool is useful because it uses the real FL leg offsets from MJCF instead
|
||||
of the simplified two-link geometry used by posture_optimizer.py. It is still a
|
||||
static single-leg approximation, so use it to choose candidate crawl/standing
|
||||
poses, then verify in MuJoCo and on the robot at low speed.
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
F_PER_LEG = 12.3 * 9.81 / 4.0
|
||||
KNEE_MIN = -2.65
|
||||
HIP_MIN = -2.58
|
||||
HIP_MAX = 2.58
|
||||
AB_MIN = -0.436
|
||||
AB_MAX = 0.611
|
||||
HARD_WHEEL_X_OFFSET = 0.08
|
||||
|
||||
|
||||
def rx(a):
|
||||
c, s = math.cos(a), math.sin(a)
|
||||
return ((1, 0, 0), (0, c, -s), (0, s, c))
|
||||
|
||||
|
||||
def ry(a):
|
||||
c, s = math.cos(a), math.sin(a)
|
||||
return ((c, 0, s), (0, 1, 0), (-s, 0, c))
|
||||
|
||||
|
||||
def mv(m, v):
|
||||
return [
|
||||
m[0][0] * v[0] + m[0][1] * v[1] + m[0][2] * v[2],
|
||||
m[1][0] * v[0] + m[1][1] * v[1] + m[1][2] * v[2],
|
||||
m[2][0] * v[0] + m[2][1] * v[1] + m[2][2] * v[2],
|
||||
]
|
||||
|
||||
|
||||
def add(a, b):
|
||||
return [a[0] + b[0], a[1] + b[1], a[2] + b[2]]
|
||||
|
||||
|
||||
def get_posture(q_ab, q_hip, q_knee):
|
||||
"""Return static FK/torque metrics for one FL leg.
|
||||
|
||||
Returns:
|
||||
base_z, tau_ab, tau_hip, tau_knee, i2r_total, knee_z, wheel_z, wheel_x_from_hip
|
||||
"""
|
||||
# FL offsets from mjcf/wheelleg.xml.
|
||||
t_ab = [0.32826, 0.066172, 0.053981]
|
||||
t_hip = [0.06389, -0.027344, 0.00010727]
|
||||
t_knee = [0.0, 0.1035, -0.25]
|
||||
t_wheel = [0.0, 0.014699, -0.20011]
|
||||
|
||||
p = mv(ry(q_knee), t_wheel)
|
||||
p = add(p, t_knee)
|
||||
p = mv(ry(q_hip), p)
|
||||
p = add(p, t_hip)
|
||||
p = mv(rx(q_ab), p)
|
||||
p = add(p, t_ab)
|
||||
|
||||
base_z = 0.10 - p[2]
|
||||
wheel_z = p[2]
|
||||
|
||||
knee_pos = mv(ry(q_hip), t_knee)
|
||||
knee_pos = mv(rx(q_ab), knee_pos)
|
||||
knee_z = knee_pos[2] + t_ab[2]
|
||||
|
||||
joint_knee = mv(ry(q_hip), add(t_knee, t_hip))
|
||||
joint_knee = mv(rx(q_ab), joint_knee)
|
||||
joint_knee = add(t_ab, joint_knee)
|
||||
|
||||
joint_hip = mv(rx(q_ab), t_hip)
|
||||
joint_hip = add(t_ab, joint_hip)
|
||||
|
||||
r_knee = [p[0] - joint_knee[0], p[1] - joint_knee[1], p[2] - joint_knee[2]]
|
||||
r_hip = [p[0] - joint_hip[0], p[1] - joint_hip[1], p[2] - joint_hip[2]]
|
||||
r_ab = [p[0] - t_ab[0], p[1] - t_ab[1], p[2] - t_ab[2]]
|
||||
|
||||
tau_ab = r_ab[1] * F_PER_LEG
|
||||
tau_hip = -r_hip[0] * F_PER_LEG
|
||||
tau_knee = -r_knee[0] * F_PER_LEG
|
||||
i2r_total = tau_ab * tau_ab + tau_hip * tau_hip + tau_knee * tau_knee
|
||||
wheel_x_from_hip = p[0] - joint_hip[0]
|
||||
return base_z, tau_ab, tau_hip, tau_knee, i2r_total, knee_z, wheel_z, wheel_x_from_hip
|
||||
|
||||
|
||||
def score_candidate(z, z_target, tau_ab, tau_hip, tau_knee, wheel_x, x_target=0.0):
|
||||
peak = max(abs(tau_ab), abs(tau_hip), abs(tau_knee))
|
||||
rms = math.sqrt((tau_ab * tau_ab + tau_hip * tau_hip + tau_knee * tau_knee) / 3.0)
|
||||
x_err = wheel_x - x_target
|
||||
x_over = max(0.0, abs(wheel_x) - 0.05)
|
||||
return (
|
||||
3000.0 * (z - z_target) ** 2
|
||||
+ 2.5 * (peak / 17.0) ** 2
|
||||
+ (rms / 17.0) ** 2
|
||||
+ 0.4 * (x_err / 0.05) ** 2
|
||||
+ 6.0 * (x_over / 0.03) ** 2
|
||||
), peak, rms
|
||||
|
||||
|
||||
def find_best(z_target, ab_range=(0.0, 0.0), step=0.002, x_target=0.0, hard_wheel_x_offset=HARD_WHEEL_X_OFFSET):
|
||||
"""Find one static posture near target height without violating hard limits."""
|
||||
best = None
|
||||
ab0, ab1 = ab_range
|
||||
n_ab = max(1, int(round((ab1 - ab0) / step)) + 1)
|
||||
n_hip = int(round((1.6 - 0.0) / step)) + 1
|
||||
n_knee = int(round((-0.4 - KNEE_MIN) / step)) + 1
|
||||
|
||||
for ia in range(n_ab):
|
||||
ab = ab0 + ia * step
|
||||
if ab < AB_MIN or ab > AB_MAX:
|
||||
continue
|
||||
for ih in range(n_hip):
|
||||
hip = ih * step
|
||||
if hip < HIP_MIN or hip > HIP_MAX:
|
||||
continue
|
||||
for ik in range(n_knee):
|
||||
knee = KNEE_MIN + ik * step
|
||||
z, ta, th, tk, i2r, kz, wz, wx = get_posture(ab, hip, knee)
|
||||
if abs(z - z_target) > 0.0015:
|
||||
continue
|
||||
if abs(wx) > hard_wheel_x_offset:
|
||||
continue
|
||||
if wz >= kz:
|
||||
continue
|
||||
cost, peak, rms = score_candidate(z, z_target, ta, th, tk, wx, x_target=x_target)
|
||||
cand = (cost, ab, hip, knee, z, ta, th, tk, peak, rms, i2r, wx)
|
||||
if best is None or cand[0] < best[0]:
|
||||
best = cand
|
||||
return best
|
||||
|
||||
|
||||
def print_table(
|
||||
z_targets,
|
||||
name,
|
||||
ab_range=(0.0, 0.0),
|
||||
step=0.002,
|
||||
x_target=0.0,
|
||||
hard_wheel_x_offset=HARD_WHEEL_X_OFFSET,
|
||||
):
|
||||
print(f"\n{'=' * 96}")
|
||||
print(name)
|
||||
print(f"{'=' * 96}")
|
||||
print(f"{'z':>6} {'ab':>6} {'hip':>7} {'knee':>7} "
|
||||
f"{'tau_ab':>8} {'tau_hip':>8} {'tau_knee':>9} "
|
||||
f"{'peak':>8} {'rms':>8} {'x_hip':>8}")
|
||||
print("-" * 96)
|
||||
for zt in z_targets:
|
||||
best = find_best(
|
||||
zt,
|
||||
ab_range=ab_range,
|
||||
step=step,
|
||||
x_target=x_target,
|
||||
hard_wheel_x_offset=hard_wheel_x_offset,
|
||||
)
|
||||
if best is None:
|
||||
print(f"{zt:>6.3f} no valid config")
|
||||
continue
|
||||
_, ab, hip, knee, z, ta, th, tk, peak, rms, _, wx = best
|
||||
print(f"{z:>6.3f} {ab:>6.3f} {hip:>7.3f} {knee:>7.3f} "
|
||||
f"{ta:>8.3f} {th:>8.3f} {tk:>9.3f} "
|
||||
f"{peak:>8.3f} {rms:>8.3f} {wx:>8.4f}")
|
||||
|
||||
|
||||
def print_crawl_default():
|
||||
best = find_best(0.17, ab_range=(0.0, 0.0), step=0.002)
|
||||
if best is None:
|
||||
return
|
||||
_, ab, hip, knee, z, ta, th, tk, peak, rms, _, wx = best
|
||||
print("\nSuggested runtime crawl_default_dof_pos:")
|
||||
print(
|
||||
f"[{ab:.3f}, {hip:.3f}, {knee:.3f}, "
|
||||
f"{ab:.3f}, {hip:.3f}, {knee:.3f}, "
|
||||
f"{ab:.3f}, {hip:.3f}, {knee:.3f}, "
|
||||
f"{ab:.3f}, {hip:.3f}, {knee:.3f}, "
|
||||
"0.0, 0.0, 0.0, 0.0]"
|
||||
)
|
||||
print(f"# z={z:.3f}, peak={peak:.3f}Nm, rms={rms:.3f}Nm, wheel_x_from_hip={wx:+.4f}m")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print_table([round(0.36 + 0.01 * i, 2) for i in range(10)],
|
||||
"STANDING candidates from MJCF FL geometry", step=0.004)
|
||||
print_table([round(0.17 + 0.01 * i, 2) for i in range(9)],
|
||||
"CRAWL candidates from MJCF FL geometry", step=0.002)
|
||||
print_table([round(0.10 + 0.01 * i, 2) for i in range(6)],
|
||||
"LOW CRAWL candidates, fixed FL abduction = +0.2",
|
||||
ab_range=(0.2, 0.2), step=0.002, hard_wheel_x_offset=0.45)
|
||||
print_crawl_default()
|
||||
@@ -0,0 +1,200 @@
|
||||
"""RlFriendlyPostureOpt — 结合电机发热 + RL友好度约束的站/爬姿优化
|
||||
|
||||
RL友好度约束（从实机经验总结，静力学可计算）：
|
||||
1. 髋关节力臂 ≥ 0.08m — 不让髋闲置（动态响应差）
|
||||
2. 三电机不均衡 ≤ 1.3x — 不让单电机先超载
|
||||
3. 运动学条件数 κ ≤ 3.5 — 不让有效传动比过高（放大控制噪声）
|
||||
|
||||
使用方法：
|
||||
uv run python mujoco_sim/rl_friendly_opt.py # 打印推荐
|
||||
uv run python -c "from mujoco_sim.rl_friendly_opt import get_all; print(get_all(0, 0.8, -1.22))"
|
||||
|
||||
基于 mjcf/wheelleg.xml 的 FL 腿运动学。
|
||||
"""
|
||||
import math
|
||||
import numpy as np
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 运动学常数（来自 wheelleg.xml FL 腿）
|
||||
# ---------------------------------------------------------------------------
|
||||
T_AB = [0.32826, 0.066172, 0.053981]
|
||||
T_HIP = [0.06389, -0.027344, 0.00010727]
|
||||
T_KNEE = [0.0, 0.1035, -0.25]
|
||||
T_WHEEL = [0.0, 0.014699, -0.20011]
|
||||
F_PER_LEG = 12.3 * 9.81 / 4.0
|
||||
KNEE_MIN, KNEE_MAX = -2.65, 2.65
|
||||
HIP_MIN, HIP_MAX = -2.58, 2.58
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 工具函数
|
||||
# ---------------------------------------------------------------------------
|
||||
def rx(a):
|
||||
c = math.cos(a); s = math.sin(a)
|
||||
return ((1, 0, 0), (0, c, -s), (0, s, c))
|
||||
|
||||
def ry(a):
|
||||
c = math.cos(a); s = math.sin(a)
|
||||
return ((c, 0, s), (0, 1, 0), (-s, 0, c))
|
||||
|
||||
def mv(m, v):
|
||||
return [m[0][0]*v[0] + m[0][1]*v[1] + m[0][2]*v[2],
|
||||
m[1][0]*v[0] + m[1][1]*v[1] + m[1][2]*v[2],
|
||||
m[2][0]*v[0] + m[2][1]*v[1] + m[2][2]*v[2]]
|
||||
|
||||
def add(a, b): return [a[0]+b[0], a[1]+b[1], a[2]+b[2]]
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 主函数
|
||||
# ---------------------------------------------------------------------------
|
||||
def get_all(q_ab, q_hip, q_knee):
|
||||
"""FK + 力矩 + 运动学指标。
|
||||
|
||||
返回 dict:
|
||||
z, ab, hip, knee, tau_ab, tau_hip, tau_knee,
|
||||
max_tau, i2r, imbal, cond, r_hip_x_mag, calf_deg
|
||||
"""
|
||||
# --- FK ---
|
||||
p = mv(ry(q_knee), T_WHEEL)
|
||||
p = add(p, T_KNEE)
|
||||
p = mv(ry(q_hip), p)
|
||||
p = add(p, T_HIP)
|
||||
p = mv(rx(q_ab), p)
|
||||
wh = add(p, T_AB)
|
||||
base_z = 0.10 - wh[2]
|
||||
wheel_z = wh[2]
|
||||
|
||||
# 膝位置（用于 wheel-below-knee）
|
||||
pk = mv(ry(q_hip), T_KNEE)
|
||||
pk = mv(rx(q_ab), pk)
|
||||
kz = pk[2] + T_AB[2]
|
||||
|
||||
# 关节位置
|
||||
jk = mv(ry(q_hip), add(T_KNEE, T_HIP))
|
||||
jk = mv(rx(q_ab), jk)
|
||||
jk = add(T_AB, jk)
|
||||
jh = mv(rx(q_ab), T_HIP)
|
||||
jh = add(T_AB, jh)
|
||||
|
||||
# 力矩
|
||||
rk = [wh[0]-jk[0], wh[1]-jk[1], wh[2]-jk[2]]
|
||||
rh = [wh[0]-jh[0], wh[1]-jh[1], wh[2]-jh[2]]
|
||||
ra = [wh[0]-T_AB[0], wh[1]-T_AB[1], wh[2]-T_AB[2]]
|
||||
tau_ab = ra[1] * F_PER_LEG
|
||||
tau_hip = -rh[0] * F_PER_LEG
|
||||
tau_knee = -rk[0] * F_PER_LEG
|
||||
|
||||
# --- Jacobian（有限差分） ---
|
||||
eps = 1e-6
|
||||
def foot_rel_hip(h, k):
|
||||
fp = mv(ry(k), T_WHEEL)
|
||||
fp = add(fp, T_KNEE)
|
||||
fp = mv(ry(h), fp)
|
||||
return [fp[0] + T_AB[0] - jh[0], fp[2] + T_AB[2] - jh[2]]
|
||||
fp0 = foot_rel_hip(q_hip, q_knee)
|
||||
fph = foot_rel_hip(q_hip + eps, q_knee)
|
||||
fpk = foot_rel_hip(q_hip, q_knee + eps)
|
||||
J = np.array([
|
||||
[(fph[0]-fp0[0])/eps, (fpk[0]-fp0[0])/eps],
|
||||
[(fph[1]-fp0[1])/eps, (fpk[1]-fp0[1])/eps],
|
||||
])
|
||||
s = np.linalg.svd(J, compute_uv=False)
|
||||
cond = s[0] / s[-1] if s[-1] > 1e-10 else 999.0
|
||||
min_sv = s[-1]
|
||||
|
||||
# --- 小腿角度（相对铅垂线） ---
|
||||
calf_x = (-0.20011) * math.sin(q_knee)
|
||||
calf_z = (-0.20011) * math.cos(q_knee)
|
||||
cv_x = calf_x * math.cos(q_hip) + calf_z * math.sin(q_hip)
|
||||
cv_z = -calf_x * math.sin(q_hip) + calf_z * math.cos(q_hip)
|
||||
calf_deg = math.degrees(math.atan2(cv_x, -cv_z))
|
||||
|
||||
abs_taus = [abs(tau_ab), abs(tau_hip), abs(tau_knee)]
|
||||
return {
|
||||
'z': base_z, 'ab': q_ab, 'hip': q_hip, 'knee': q_knee,
|
||||
'tau_ab': tau_ab, 'tau_hip': tau_hip, 'tau_knee': tau_knee,
|
||||
'max_tau': max(abs_taus),
|
||||
'i2r': tau_ab**2 + tau_hip**2 + tau_knee**2,
|
||||
'imbal': max(abs_taus) / max(1e-10, min(abs_taus)),
|
||||
'cond': cond, 'min_sv': min_sv,
|
||||
'r_hip_x_mag': abs(rh[0]),
|
||||
'calf_deg': calf_deg,
|
||||
'kz': kz, 'wz': wheel_z,
|
||||
}
|
||||
|
||||
|
||||
def rl_cost(r):
|
||||
"""RL友好度综合成本（越小越好）。
|
||||
|
||||
约束来源：
|
||||
c1 — 瓶颈电机发热 τ²/τ_max² 主目标
|
||||
c2 — 电机不均衡 > 1.3x (imbal-1.3)² 单电机先超载
|
||||
c3 — 髋力臂 < 8cm (0.08 - r_hip) 髋闲置→动态响应差
|
||||
c4 — 有效传动比 κ > 3.5 (κ - 3.5) 高刚度→冲击传递大
|
||||
c5 — 腿的被动刚度 > 1.3x (stiff-1.3) 刚度比→冲击吸收（新！）
|
||||
"""
|
||||
c1 = (r['max_tau'] / 17.0) ** 2
|
||||
c2 = max(0.0, (r['imbal'] - 1.3) / 1.0) ** 2
|
||||
c3 = max(0.0, (0.08 - r['r_hip_x_mag'])) / 0.08
|
||||
c4 = max(0.0, (r['cond'] - 3.5)) / 5.0
|
||||
# stiffness ratio normalized to z=0.40 (σ_min≈0.115)
|
||||
stiff = (0.115 / r['min_sv']) ** 2
|
||||
c5 = max(0.0, (stiff - 1.3)) / 3.0
|
||||
return 100.0*c1 + 50.0*c2 + 80.0*c3 + 30.0*c4 + 40.0*c5
|
||||
|
||||
|
||||
def sweep_z(z_targets, name, ab_max=0.44, tol=0.004):
|
||||
"""遍历 z 扫描最优姿态。"""
|
||||
print(f"\n{'='*100}")
|
||||
print(f" {name}")
|
||||
print(f"{'='*100}")
|
||||
print(f"{'z_tgt':>6} {'z':>6} {'ab':>5} {'hip':>6} {'knee':>6} | "
|
||||
f"{'maxτ':>6} {'imbal':>6} {'r_hip':>6} {'κ':>5} | "
|
||||
f"{'c1热':>6} {'c2均':>6} {'c3髋':>6} {'c4奇':>6} {'cost':>7}")
|
||||
print("-"*100)
|
||||
results = []
|
||||
for zt in z_targets:
|
||||
best_cost = float('inf')
|
||||
best_r = None
|
||||
for ab in [round(i*0.02, 2) for i in range(int(ab_max/0.02)+1)]:
|
||||
for h in [round(i*0.01, 2) for i in range(260)]:
|
||||
for kn in [round(-2.65+i*0.01, 2) for i in range(256)]:
|
||||
if h + kn > -0.1: continue
|
||||
r = get_all(ab, h, kn)
|
||||
if abs(r['z'] - zt) > tol: continue
|
||||
if r['wz'] >= r['kz']: continue
|
||||
if r['tau_ab'] < 0 or r['tau_hip'] < 0: continue
|
||||
cost = rl_cost(r)
|
||||
if cost < best_cost:
|
||||
best_cost = cost; best_r = r
|
||||
if best_r:
|
||||
r = best_r
|
||||
c1s = 100.0*(r['max_tau']/17.0)**2
|
||||
c2s = 50.0*max(0.0,(r['imbal']-1.3))**2
|
||||
c3s = 80.0*max(0.0,(0.08-r['r_hip_x_mag']))/0.08
|
||||
c4s = 30.0*max(0.0,(r['cond']-3.5))/5.0
|
||||
print(f"{zt:>6.2f} {r['z']:>6.3f} {r['ab']:>5.2f} {r['hip']:>6.2f} {r['knee']:>6.2f} | "
|
||||
f"{r['max_tau']:>6.3f} {r['imbal']:>6.1f}x {r['r_hip_x_mag']:>6.3f} {r['cond']:>5.1f} | "
|
||||
f"{c1s:>6.1f} {c2s:>6.1f} {c3s:>6.1f} {c4s:>6.1f} {best_cost:>7.1f}")
|
||||
results.append((zt, r))
|
||||
else:
|
||||
print(f"{zt:>6.2f} — no valid")
|
||||
return results
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
stand = sweep_z([round(0.36+0.01*i, 2) for i in range(10)],
|
||||
"STANDING — heat + RL constraints (ab free)", 0.44)
|
||||
crawl = sweep_z([round(0.10+0.01*i, 2) for i in range(9)],
|
||||
"CRAWL — heat + RL constraints (ab_max=0.25)", 0.25)
|
||||
|
||||
print(f"\n{'='*100}")
|
||||
print(" RECOMMENDATION")
|
||||
print(f"{'='*100}")
|
||||
if stand:
|
||||
zt, r = min(stand, key=lambda x: rl_cost(x[1]))
|
||||
print(f" Standing: z={r['z']:.3f} ab={r['ab']:.2f} hip={r['hip']:.2f} knee={r['knee']:.2f}")
|
||||
print(f" maxτ={r['max_tau']:.3f} imbal={r['imbal']:.1f}x κ={r['cond']:.1f} r_hip={r['r_hip_x_mag']:.3f}m")
|
||||
if crawl:
|
||||
zt, r = min(crawl, key=lambda x: rl_cost(x[1]))
|
||||
print(f" Crawl: z={r['z']:.3f} ab={r['ab']:.2f} hip={r['hip']:.2f} knee={r['knee']:.2f}")
|
||||
print(f" maxτ={r['max_tau']:.3f} imbal={r['imbal']:.1f}x κ={r['cond']:.1f} r_hip={r['r_hip_x_mag']:.3f}m")
|
||||
@@ -5,11 +5,13 @@ Configures actuators as proper PD controllers at runtime:
|
||||
- Wheel joints: force = gain*(ctrl - qvel), ctrl = target velocity (rad/s)
|
||||
"""
|
||||
|
||||
import math
|
||||
import numpy as np
|
||||
import mujoco
|
||||
from dataclasses import dataclass
|
||||
from config import (SCENE_XML, LEG_NAMES, LEG_JOINTS, WHEEL_JOINT,
|
||||
DEFAULT_JOINT_ANGLES, WHEEL_RADIUS, WHEEL_TRACK)
|
||||
DEFAULT_JOINT_ANGLES, WHEEL_RADIUS, WHEEL_TRACK,
|
||||
L_THIGH, L_CALF)
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -100,14 +102,14 @@ class Robot:
|
||||
adr = self._joint_qpos_adr[name]
|
||||
self.data.qpos[adr] = DEFAULT_JOINT_ANGLES[key]
|
||||
|
||||
# Compute correct base height from default angles using 2R FK
|
||||
# leg_length = sqrt(L1^2 + L2^2 - 2*L1*L2*cos(pi + knee))
|
||||
import math
|
||||
L1, L2 = 0.25, 0.20
|
||||
# Compute correct base height from default angles using exact FK
|
||||
# z_base = wheel_radius + L1*cos(hip) + L2*cos(hip+knee) - hip_z_offset
|
||||
hip = DEFAULT_JOINT_ANGLES["hip_pitch"]
|
||||
knee = DEFAULT_JOINT_ANGLES["knee"]
|
||||
leg_length = math.sqrt(L1**2 + L2**2 - 2*L1*L2*math.cos(math.pi + knee))
|
||||
# base_z = wheel_radius + leg_length - hip_z_offset
|
||||
base_z = 0.10 + leg_length - 0.054
|
||||
base_z = (WHEEL_RADIUS
|
||||
+ L_THIGH * math.cos(hip)
|
||||
+ L_CALF * math.cos(hip + knee)
|
||||
- 0.054)
|
||||
self.data.qpos[2] = base_z
|
||||
self.data.qpos[3] = 1.0 # quat w
|
||||
|
||||
|
||||
@@ -52,7 +52,7 @@ def main():
|
||||
elif c == 'q':
|
||||
ctrl.height = min(ctrl.height + 0.02, 0.45)
|
||||
elif c == 'e':
|
||||
ctrl.height = max(ctrl.height - 0.02, 0.15)
|
||||
ctrl.height = max(ctrl.height - 0.02, 0.17)
|
||||
elif c == '1':
|
||||
ctrl.mode = "wheel"; ctrl.prone = False
|
||||
elif c == '2':
|
||||
|
||||
@@ -0,0 +1,528 @@
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
"""MuJoCo static posture optimizer for wheeled-leg standing defaults.
|
||||
|
||||
The old posture tools are mostly analytical. This script keeps the fast MJCF
|
||||
kinematics as a candidate generator, then evaluates the best candidates in
|
||||
MuJoCo with gravity and floor contact enabled.
|
||||
|
||||
The score is meant for RL default pose / real deployment:
|
||||
- low peak and RMS standing torque, so one hot motor is not hidden by average
|
||||
- wheel contact point close to the hip in X for wheel speed tracking
|
||||
- COM projection margin inside the four-wheel support rectangle
|
||||
- non-singular leg Jacobian for posture control authority
|
||||
- underbody and knee clearance for obstacle tolerance
|
||||
|
||||
Usage:
|
||||
uv run python mujoco_sim/static_posture_optimizer.py
|
||||
uv run python mujoco_sim/static_posture_optimizer.py --quick
|
||||
uv run python mujoco_sim/static_posture_optimizer.py --ab-max 0.08
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import math
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[1]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from mujoco_sim.rl_friendly_opt import get_all # noqa: E402
|
||||
|
||||
SCENE_XML = REPO_ROOT / "mjcf" / "scene.xml"
|
||||
|
||||
LEG_NAMES = ("fl", "fr", "rl", "rr")
|
||||
LEG_JOINTS = ("hip_abduction_joint", "hip_pitch_joint", "knee_joint")
|
||||
WHEEL_JOINT = "wheel_joint"
|
||||
LEG_ACTUATORS = [f"{leg}_{jt}" for leg in LEG_NAMES for jt in LEG_JOINTS]
|
||||
|
||||
ROBOT_MASS = 12.3
|
||||
G = 9.81
|
||||
MAX_TORQUE = 17.0
|
||||
WHEEL_RADIUS = 0.10
|
||||
|
||||
HIP_MIN, HIP_MAX = -2.58, 2.58
|
||||
KNEE_MIN, KNEE_MAX = -2.65, 2.65
|
||||
HIP_SCAN = (0.20, 1.15)
|
||||
KNEE_SCAN = (-1.90, -0.65)
|
||||
|
||||
SOFT_WHEEL_X = 0.045
|
||||
HARD_WHEEL_X = 0.085
|
||||
MIN_COM_MARGIN = 0.045
|
||||
MIN_KNEE_CLEARANCE = 0.105
|
||||
MIN_UNDERBODY_CLEARANCE = 0.33
|
||||
MAX_COND = 3.7
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Candidate:
|
||||
ab: float
|
||||
hip: float
|
||||
knee: float
|
||||
z_fk: float
|
||||
analytic_cost: float
|
||||
|
||||
|
||||
@dataclass
|
||||
class StaticResult:
|
||||
z_target: float
|
||||
z: float
|
||||
ab: float
|
||||
hip: float
|
||||
knee: float
|
||||
cost: float
|
||||
peak_tau: float
|
||||
rms_tau: float
|
||||
mean_i2r: float
|
||||
imbal: float
|
||||
wheel_x: float
|
||||
cond: float
|
||||
min_sv: float
|
||||
com_margin_x: float
|
||||
com_margin_y: float
|
||||
support_margin: float
|
||||
normal_cv: float
|
||||
body_clearance: float
|
||||
knee_clearance: float
|
||||
roll: float
|
||||
pitch: float
|
||||
height_err: float
|
||||
|
||||
|
||||
def _configure_actuators(model: mujoco.MjModel) -> None:
|
||||
"""Configure leg joints as position PD and wheels as zero-velocity motors."""
|
||||
for i in range(model.nu):
|
||||
name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_ACTUATOR, i) or ""
|
||||
model.actuator_biastype[i] = 1
|
||||
model.actuator_gaintype[i] = 0
|
||||
model.actuator_forcelimited[i] = 0
|
||||
if "wheel" in name:
|
||||
model.actuator_gainprm[i, 0] = 2.0
|
||||
model.actuator_biasprm[i, 0] = 0.0
|
||||
model.actuator_biasprm[i, 1] = 0.0
|
||||
model.actuator_biasprm[i, 2] = -2.0
|
||||
model.actuator_ctrlrange[i] = [-20.0, 20.0]
|
||||
else:
|
||||
model.actuator_gainprm[i, 0] = 60.0
|
||||
model.actuator_biasprm[i, 0] = 0.0
|
||||
model.actuator_biasprm[i, 1] = -60.0
|
||||
model.actuator_biasprm[i, 2] = -3.0
|
||||
model.actuator_ctrlrange[i] = [-3.14, 3.14]
|
||||
|
||||
|
||||
def _ids(model: mujoco.MjModel):
|
||||
act = {}
|
||||
qadr = {}
|
||||
for leg in LEG_NAMES:
|
||||
for jt in (*LEG_JOINTS, WHEEL_JOINT):
|
||||
name = f"{leg}_{jt}"
|
||||
act[name] = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, name)
|
||||
jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, name)
|
||||
qadr[name] = model.jnt_qposadr[jid]
|
||||
bodies = {
|
||||
"base": mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "base_link"),
|
||||
**{
|
||||
f"{leg}_wheel": mujoco.mj_name2id(
|
||||
model, mujoco.mjtObj.mjOBJ_BODY, f"{leg}_wheel_Link"
|
||||
)
|
||||
for leg in LEG_NAMES
|
||||
},
|
||||
**{
|
||||
f"{leg}_knee": mujoco.mj_name2id(
|
||||
model, mujoco.mjtObj.mjOBJ_BODY, f"{leg}_knee_Link"
|
||||
)
|
||||
for leg in LEG_NAMES
|
||||
},
|
||||
}
|
||||
return act, qadr, bodies
|
||||
|
||||
|
||||
def _mirrored_ab(leg: str, ab: float) -> float:
|
||||
return ab if leg[1] == "l" else -ab
|
||||
|
||||
|
||||
def _set_pose(
|
||||
model: mujoco.MjModel,
|
||||
data: mujoco.MjData,
|
||||
act: dict[str, int],
|
||||
qadr: dict[str, int],
|
||||
cand: Candidate,
|
||||
) -> None:
|
||||
mujoco.mj_resetData(model, data)
|
||||
for leg in LEG_NAMES:
|
||||
vals = (_mirrored_ab(leg, cand.ab), cand.hip, cand.knee)
|
||||
for jt, val in zip(LEG_JOINTS, vals):
|
||||
name = f"{leg}_{jt}"
|
||||
data.qpos[qadr[name]] = val
|
||||
data.ctrl[act[name]] = val
|
||||
data.ctrl[act[f"{leg}_{WHEEL_JOINT}"]] = 0.0
|
||||
data.qpos[0:3] = [0.0, 0.0, cand.z_fk]
|
||||
data.qpos[3] = 1.0
|
||||
data.qpos[4:7] = 0.0
|
||||
mujoco.mj_forward(model, data)
|
||||
|
||||
|
||||
def _rpy(data: mujoco.MjData, bodies: dict[str, int]) -> tuple[float, float]:
|
||||
rot = data.xmat[bodies["base"]].reshape(3, 3)
|
||||
roll = math.atan2(rot[2, 1], rot[2, 2])
|
||||
pitch = math.atan2(-rot[2, 0], math.sqrt(rot[2, 1] ** 2 + rot[2, 2] ** 2))
|
||||
return roll, pitch
|
||||
|
||||
|
||||
def _robot_com(model: mujoco.MjModel, data: mujoco.MjData) -> np.ndarray:
|
||||
masses = model.body_mass[1:]
|
||||
return (data.xipos[1:] * masses[:, None]).sum(axis=0) / masses.sum()
|
||||
|
||||
|
||||
def _support_metrics(
|
||||
model: mujoco.MjModel,
|
||||
data: mujoco.MjData,
|
||||
bodies: dict[str, int],
|
||||
) -> tuple[float, float, float, float]:
|
||||
wheel_xy = np.array([data.xpos[bodies[f"{leg}_wheel"]][:2] for leg in LEG_NAMES])
|
||||
com_xy = _robot_com(model, data)[:2]
|
||||
min_xy = wheel_xy.min(axis=0)
|
||||
max_xy = wheel_xy.max(axis=0)
|
||||
margin_low = com_xy - min_xy
|
||||
margin_high = max_xy - com_xy
|
||||
margin_x = float(min(margin_low[0], margin_high[0]))
|
||||
margin_y = float(min(margin_low[1], margin_high[1]))
|
||||
support_margin = float(min(margin_x, margin_y))
|
||||
|
||||
normal = []
|
||||
for leg in LEG_NAMES:
|
||||
bid = bodies[f"{leg}_wheel"]
|
||||
fz = 0.0
|
||||
for i in range(data.ncon):
|
||||
con = data.contact[i]
|
||||
b1 = model.geom_bodyid[con.geom1]
|
||||
b2 = model.geom_bodyid[con.geom2]
|
||||
if b1 == bid or b2 == bid:
|
||||
wrench = np.zeros(6)
|
||||
mujoco.mj_contactForce(model, data, i, wrench)
|
||||
fz += abs(float(wrench[0]))
|
||||
normal.append(fz)
|
||||
normal = np.asarray(normal, dtype=float)
|
||||
if normal.sum() < 1e-6:
|
||||
normal_cv = 9.99
|
||||
else:
|
||||
normal_cv = float(normal.std() / max(1e-6, normal.mean()))
|
||||
return margin_x, margin_y, support_margin, normal_cv
|
||||
|
||||
|
||||
def _clearance_metrics(data: mujoco.MjData, bodies: dict[str, int]) -> tuple[float, float]:
|
||||
base_z = float(data.xpos[bodies["base"]][2])
|
||||
# The collision box in wheelleg.xml is centered at z=0.054 with half-height 0.073.
|
||||
body_clearance = base_z + 0.054 - 0.073
|
||||
knee_z = min(float(data.xpos[bodies[f"{leg}_knee"]][2]) for leg in LEG_NAMES)
|
||||
wheel_z = min(float(data.xpos[bodies[f"{leg}_wheel"]][2]) for leg in LEG_NAMES)
|
||||
return body_clearance, knee_z - wheel_z
|
||||
|
||||
|
||||
def _static_cost(r: StaticResult) -> float:
|
||||
wheel_over = max(0.0, abs(r.wheel_x) - SOFT_WHEEL_X)
|
||||
support_short = max(0.0, MIN_COM_MARGIN - r.support_margin)
|
||||
body_short = max(0.0, MIN_UNDERBODY_CLEARANCE - r.body_clearance)
|
||||
knee_short = max(0.0, MIN_KNEE_CLEARANCE - r.knee_clearance)
|
||||
cond_over = max(0.0, r.cond - MAX_COND)
|
||||
tilt = math.hypot(r.roll, r.pitch)
|
||||
|
||||
return (
|
||||
6500.0 * r.height_err**2
|
||||
+ 2.6 * (r.peak_tau / MAX_TORQUE) ** 2
|
||||
+ 1.1 * (r.rms_tau / MAX_TORQUE) ** 2
|
||||
+ 1.8 * max(0.0, r.imbal - 1.6) ** 2
|
||||
+ 0.65 * (abs(r.wheel_x) / SOFT_WHEEL_X) ** 2
|
||||
+ 9.0 * (wheel_over / max(1e-6, HARD_WHEEL_X - SOFT_WHEEL_X)) ** 2
|
||||
+ 10.0 * (support_short / MIN_COM_MARGIN) ** 2
|
||||
+ 2.0 * (r.normal_cv / 0.35) ** 2
|
||||
+ 30.0 * (cond_over / 1.0) ** 2
|
||||
+ 4.0 * (body_short / 0.06) ** 2
|
||||
+ 2.0 * (knee_short / 0.04) ** 2
|
||||
+ 1.0 * (tilt / 0.05) ** 2
|
||||
)
|
||||
|
||||
|
||||
def _analytic_cost(r: dict, z_target: float) -> float:
|
||||
wheel_over = max(0.0, abs(r["r_hip_x_mag"]) - SOFT_WHEEL_X)
|
||||
cond_over = max(0.0, r["cond"] - MAX_COND)
|
||||
return (
|
||||
1400.0 * (r["z"] - z_target) ** 2
|
||||
+ 2.2 * (r["max_tau"] / MAX_TORQUE) ** 2
|
||||
+ 0.7 * (math.sqrt(r["i2r"] / 3.0) / MAX_TORQUE) ** 2
|
||||
+ 0.55 * (abs(r["r_hip_x_mag"]) / SOFT_WHEEL_X) ** 2
|
||||
+ 8.0 * (wheel_over / max(1e-6, HARD_WHEEL_X - SOFT_WHEEL_X)) ** 2
|
||||
+ 2.0 * (cond_over / 2.0) ** 2
|
||||
)
|
||||
|
||||
|
||||
def generate_candidates(
|
||||
z_target: float,
|
||||
step: float,
|
||||
ab_max: float,
|
||||
keep: int,
|
||||
z_tol: float,
|
||||
) -> list[Candidate]:
|
||||
candidates: list[Candidate] = []
|
||||
ab_values = np.arange(0.0, ab_max + 0.5 * step, step)
|
||||
hip_values = np.arange(HIP_SCAN[0], HIP_SCAN[1] + 0.5 * step, step)
|
||||
knee_values = np.arange(KNEE_SCAN[0], KNEE_SCAN[1] + 0.5 * step, step)
|
||||
for ab in ab_values:
|
||||
for hip in hip_values:
|
||||
if hip < HIP_MIN or hip > HIP_MAX:
|
||||
continue
|
||||
for knee in knee_values:
|
||||
if knee < KNEE_MIN or knee > KNEE_MAX:
|
||||
continue
|
||||
r = get_all(float(ab), float(hip), float(knee))
|
||||
if abs(r["z"] - z_target) > z_tol:
|
||||
continue
|
||||
if r["wz"] >= r["kz"]:
|
||||
continue
|
||||
if abs(r["r_hip_x_mag"]) > HARD_WHEEL_X:
|
||||
continue
|
||||
if r["max_tau"] > MAX_TORQUE * 1.15:
|
||||
continue
|
||||
candidates.append(
|
||||
Candidate(
|
||||
ab=float(ab),
|
||||
hip=float(hip),
|
||||
knee=float(knee),
|
||||
z_fk=float(r["z"]),
|
||||
analytic_cost=_analytic_cost(r, z_target),
|
||||
)
|
||||
)
|
||||
candidates.sort(key=lambda c: c.analytic_cost)
|
||||
return candidates[:keep]
|
||||
|
||||
|
||||
def evaluate_candidate(
|
||||
model: mujoco.MjModel,
|
||||
data: mujoco.MjData,
|
||||
act: dict[str, int],
|
||||
qadr: dict[str, int],
|
||||
bodies: dict[str, int],
|
||||
cand: Candidate,
|
||||
z_target: float,
|
||||
settle_steps: int,
|
||||
avg_steps: int,
|
||||
) -> StaticResult | None:
|
||||
_set_pose(model, data, act, qadr, cand)
|
||||
for _ in range(settle_steps):
|
||||
mujoco.mj_step(model, data)
|
||||
|
||||
roll, pitch = _rpy(data, bodies)
|
||||
if abs(roll) > 0.35 or abs(pitch) > 0.35:
|
||||
return None
|
||||
|
||||
tau_buf = []
|
||||
for _ in range(avg_steps):
|
||||
mujoco.mj_step(model, data)
|
||||
tau_buf.append([data.actuator_force[act[name]] for name in LEG_ACTUATORS])
|
||||
tau = np.asarray(tau_buf, dtype=float).mean(axis=0)
|
||||
tau_abs = np.abs(tau)
|
||||
peak = float(tau_abs.max())
|
||||
rms = float(np.sqrt(np.mean(tau * tau)))
|
||||
mean_i2r = float(np.mean(tau * tau))
|
||||
|
||||
sagittal_abs = []
|
||||
for name, value in zip(LEG_ACTUATORS, tau):
|
||||
if "hip_pitch" in name or "knee" in name:
|
||||
sagittal_abs.append(abs(float(value)))
|
||||
sagittal_abs = np.asarray(sagittal_abs, dtype=float)
|
||||
imbal = float(sagittal_abs.max() / max(1e-6, sagittal_abs.mean()))
|
||||
|
||||
fk = get_all(cand.ab, cand.hip, cand.knee)
|
||||
margin_x, margin_y, support_margin, normal_cv = _support_metrics(model, data, bodies)
|
||||
body_clearance, knee_clearance = _clearance_metrics(data, bodies)
|
||||
z = float(data.xpos[bodies["base"]][2])
|
||||
|
||||
result = StaticResult(
|
||||
z_target=z_target,
|
||||
z=z,
|
||||
ab=cand.ab,
|
||||
hip=cand.hip,
|
||||
knee=cand.knee,
|
||||
cost=0.0,
|
||||
peak_tau=peak,
|
||||
rms_tau=rms,
|
||||
mean_i2r=mean_i2r,
|
||||
imbal=imbal,
|
||||
wheel_x=float(fk["r_hip_x_mag"]),
|
||||
cond=float(fk["cond"]),
|
||||
min_sv=float(fk["min_sv"]),
|
||||
com_margin_x=margin_x,
|
||||
com_margin_y=margin_y,
|
||||
support_margin=support_margin,
|
||||
normal_cv=normal_cv,
|
||||
body_clearance=body_clearance,
|
||||
knee_clearance=knee_clearance,
|
||||
roll=roll,
|
||||
pitch=pitch,
|
||||
height_err=z - z_target,
|
||||
)
|
||||
result.cost = _static_cost(result)
|
||||
return result
|
||||
|
||||
|
||||
def optimize(
|
||||
z_targets: list[float],
|
||||
step: float,
|
||||
ab_max: float,
|
||||
keep: int,
|
||||
z_tol: float,
|
||||
settle_steps: int,
|
||||
avg_steps: int,
|
||||
) -> list[StaticResult]:
|
||||
model = mujoco.MjModel.from_xml_path(str(SCENE_XML))
|
||||
data = mujoco.MjData(model)
|
||||
_configure_actuators(model)
|
||||
act, qadr, bodies = _ids(model)
|
||||
|
||||
results: list[StaticResult] = []
|
||||
for zt in z_targets:
|
||||
candidates = generate_candidates(zt, step=step, ab_max=ab_max, keep=keep, z_tol=z_tol)
|
||||
best: StaticResult | None = None
|
||||
for cand in candidates:
|
||||
result = evaluate_candidate(
|
||||
model,
|
||||
data,
|
||||
act,
|
||||
qadr,
|
||||
bodies,
|
||||
cand,
|
||||
zt,
|
||||
settle_steps=settle_steps,
|
||||
avg_steps=avg_steps,
|
||||
)
|
||||
if result is None:
|
||||
continue
|
||||
if best is None or result.cost < best.cost:
|
||||
best = result
|
||||
if best is None:
|
||||
print(f"{zt:.2f}: no valid MuJoCo-static candidate from {len(candidates)} seeds")
|
||||
continue
|
||||
results.append(best)
|
||||
print_result(best)
|
||||
return results
|
||||
|
||||
|
||||
def print_header() -> None:
|
||||
print("\nMuJoCo static posture optimization")
|
||||
print("z_tgt z ab hip knee | peak rms imbal wheelX cond | comX comY clrB clrK | cost")
|
||||
print("-" * 111)
|
||||
|
||||
|
||||
def print_result(r: StaticResult) -> None:
|
||||
print(
|
||||
f"{r.z_target:5.2f} {r.z:6.3f} {r.ab:5.2f} {r.hip:6.3f} {r.knee:6.3f} | "
|
||||
f"{r.peak_tau:5.2f} {r.rms_tau:5.2f} {r.imbal:6.2f} "
|
||||
f"{r.wheel_x:6.3f} {r.cond:5.2f} | "
|
||||
f"{r.com_margin_x:5.3f} {r.com_margin_y:5.3f} "
|
||||
f"{r.body_clearance:5.3f} {r.knee_clearance:5.3f} | "
|
||||
f"{r.cost:6.2f}"
|
||||
)
|
||||
|
||||
|
||||
def print_tables(results: list[StaticResult]) -> None:
|
||||
if not results:
|
||||
return
|
||||
print("\nCopy-paste tables:")
|
||||
print("_H_TARGET = " + repr([round(r.z_target, 3) for r in results]))
|
||||
print("_Z_STATIC = " + repr([round(r.z, 3) for r in results]))
|
||||
print("_HIP = " + repr([round(r.hip, 3) for r in results]))
|
||||
print("_KNEE = " + repr([round(r.knee, 3) for r in results]))
|
||||
print("_ABD_LEFT = " + repr([round(r.ab, 3) for r in results]))
|
||||
print("_ABD_RIGHT = " + repr([round(-r.ab, 3) for r in results]))
|
||||
|
||||
best = min(results, key=lambda r: r.cost)
|
||||
print("\nRecommended default:")
|
||||
print(
|
||||
f"z={best.z_target:.2f}, ab={best.ab:.3f}, hip={best.hip:.3f}, "
|
||||
f"knee={best.knee:.3f}, peak={best.peak_tau:.2f}Nm, "
|
||||
f"rms={best.rms_tau:.2f}Nm, support_margin={best.support_margin:.3f}m"
|
||||
)
|
||||
pose = []
|
||||
for leg in LEG_NAMES:
|
||||
pose.extend([_mirrored_ab(leg, best.ab), best.hip, best.knee])
|
||||
pose.extend([0.0, 0.0, 0.0, 0.0])
|
||||
print("default_dof_pos = " + repr([round(v, 3) for v in pose]))
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--z-min", type=float, default=0.36)
|
||||
parser.add_argument("--z-max", type=float, default=0.45)
|
||||
parser.add_argument("--z-step", type=float, default=0.01)
|
||||
parser.add_argument("--grid-step", type=float, default=0.01)
|
||||
parser.add_argument("--ab-max", type=float, default=0.0)
|
||||
parser.add_argument("--keep", type=int, default=45)
|
||||
parser.add_argument("--z-tol", type=float, default=0.035)
|
||||
parser.add_argument("--settle-steps", type=int, default=350)
|
||||
parser.add_argument("--avg-steps", type=int, default=80)
|
||||
parser.add_argument("--quick", action="store_true", help="Coarser and faster scan")
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.quick:
|
||||
args.grid_step = max(args.grid_step, 0.02)
|
||||
args.keep = min(args.keep, 20)
|
||||
args.z_tol = max(args.z_tol, 0.040)
|
||||
args.settle_steps = min(args.settle_steps, 220)
|
||||
args.avg_steps = min(args.avg_steps, 40)
|
||||
|
||||
n = int(round((args.z_max - args.z_min) / args.z_step)) + 1
|
||||
z_targets = [round(args.z_min + i * args.z_step, 3) for i in range(n)]
|
||||
|
||||
print_header()
|
||||
results = optimize(
|
||||
z_targets,
|
||||
step=args.grid_step,
|
||||
ab_max=args.ab_max,
|
||||
keep=args.keep,
|
||||
z_tol=args.z_tol,
|
||||
settle_steps=args.settle_steps,
|
||||
avg_steps=args.avg_steps,
|
||||
)
|
||||
print_tables(results)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -52,7 +52,11 @@ from ..mdp.only_positive_rewards import enable_only_positive_rewards
|
||||
from ..mdp.rewards import (
|
||||
track_linear_velocity,
|
||||
track_linear_velocity_l1,
|
||||
track_linear_velocity_x,
|
||||
track_linear_velocity_y,
|
||||
track_angular_velocity,
|
||||
track_angular_velocity_z,
|
||||
stair_lateral_yaw_drift_l2,
|
||||
base_height_l2,
|
||||
safe_base_lin_vel,
|
||||
safe_foot_contact,
|
||||
@@ -74,9 +78,39 @@ 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,
|
||||
tracking_lin_vel_error,
|
||||
tracking_yaw_vel_error,
|
||||
tracking_lin_vel_x_error,
|
||||
tracking_lin_vel_y_error,
|
||||
tracking_lin_vel_along_command_error,
|
||||
actual_lin_vel_orthogonal_command_mean,
|
||||
command_lin_vel_mean,
|
||||
command_yaw_vel_abs_mean,
|
||||
actual_lin_vel_mean,
|
||||
tracking_lin_vel_error_band_mean,
|
||||
tracking_lin_vel_axis_error_band_mean,
|
||||
command_band_active,
|
||||
wheel_raw_action_abs_mean,
|
||||
wheel_target_vel_abs_mean,
|
||||
wheel_actual_vel_abs_mean,
|
||||
wheel_target_actual_vel_error_mean,
|
||||
wheel_actual_to_target_vel_ratio_mean,
|
||||
wheel_target_actual_sign_agreement,
|
||||
upright_metric,
|
||||
base_ground_contact_metric,
|
||||
)
|
||||
from ..mdp.curriculums import (
|
||||
command_axis_levels_vel,
|
||||
command_levels_adaptive,
|
||||
terrain_levels_obstacle_release,
|
||||
terrain_levels_ramp_strict,
|
||||
terrain_levels_vel_strict,
|
||||
)
|
||||
from ..mdp.curriculums import terrain_levels_vel_strict
|
||||
from ..mdp.commands import UniformThresholdVelocityCommandCfg
|
||||
|
||||
# Constant Definitions
|
||||
@@ -174,7 +208,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 +226,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,
|
||||
@@ -239,8 +273,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"),
|
||||
},
|
||||
),
|
||||
@@ -254,13 +295,9 @@ def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
|
||||
params={"asset_cfg": SceneEntityCfg("robot", body_names=("base_link",)),
|
||||
"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.3, 1.0)},
|
||||
),
|
||||
"actuator_stiffness": EventTermCfg(
|
||||
func=envs_dr.joint_stiffness, mode="startup",
|
||||
@@ -270,16 +307,7 @@ def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
|
||||
func=envs_dr.joint_damping, mode="startup",
|
||||
params={"asset_cfg": SceneEntityCfg("robot"), "ranges": (0.9, 1.1), "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,16 +315,6 @@ def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
|
||||
"ranges": (-1.0, 3.0),
|
||||
},
|
||||
),
|
||||
"continuous_disturbance": EventTermCfg(
|
||||
func=apply_continuous_disturbance, mode="step",
|
||||
params={
|
||||
"asset_cfg": SceneEntityCfg("robot", body_names=("base_link",)),
|
||||
"force_range": (-15.0, 15.0),
|
||||
"torque_range": (-10.0, 10.0),
|
||||
"resample_time_range": (0.5, 2.0),
|
||||
"time_constant": 0.5,
|
||||
},
|
||||
),
|
||||
}
|
||||
|
||||
# ------------------
|
||||
@@ -346,8 +364,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),
|
||||
)
|
||||
|
||||
@@ -379,39 +397,97 @@ def rough_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
|
||||
terrain_generator=TerrainGeneratorCfg(
|
||||
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)),
|
||||
"flat": BoxFlatTerrainCfg(proportion=0.15, size=(8.0, 8.0)),
|
||||
"pyramid_stairs": BoxPyramidStairsTerrainCfg(proportion=0.05, step_height_range=(0.0, 0.20), step_width=0.30, size=(8.0, 8.0)),
|
||||
"pyramid_stairs_inv": BoxInvertedPyramidStairsTerrainCfg(proportion=0.35, step_height_range=(0.0, 0.20), 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.20), 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.10, 0.35),
|
||||
wall_centers_x=(2.1, 3.2, 4.3, 5.4, 6.5),
|
||||
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=terrain_levels_obstacle_release,
|
||||
params={
|
||||
"command_name": "twist",
|
||||
"initial_terrain_names": ("flat", "random_rough", "perlin_noise", "sloped_terrain", "pyramid_stairs"),
|
||||
"release_schedule": (
|
||||
(200 * 24, ("random_grid",)),
|
||||
(500 * 24, ("pyramid_stairs_inv",)),
|
||||
(700 * 24, ("rc_wall",)),
|
||||
),
|
||||
},
|
||||
)
|
||||
cfg.curriculum["command_x_levels"] = CurriculumTermCfg(
|
||||
func=command_levels_adaptive,
|
||||
params={
|
||||
"command_name": "twist",
|
||||
"reward_term_name": "track_lin_vel_x_exp",
|
||||
"axis": "x",
|
||||
"initial_range": (-0.5, 0.5),
|
||||
"delta_command": 0.05,
|
||||
"target_ratio": 0.8,
|
||||
"ema_alpha": 0.5,
|
||||
},
|
||||
)
|
||||
cfg.curriculum["command_y_levels"] = CurriculumTermCfg(
|
||||
func=command_levels_adaptive,
|
||||
params={
|
||||
"command_name": "twist",
|
||||
"reward_term_name": "track_lin_vel_y_exp",
|
||||
"axis": "y",
|
||||
"initial_range": (-0.5, 0.5),
|
||||
"delta_command": 0.05,
|
||||
"target_ratio": 0.8,
|
||||
"ema_alpha": 0.5,
|
||||
},
|
||||
)
|
||||
cfg.curriculum["command_yaw_levels"] = CurriculumTermCfg(
|
||||
func=command_levels_adaptive,
|
||||
params={
|
||||
"command_name": "twist",
|
||||
"reward_term_name": "track_ang_vel_z_exp",
|
||||
"axis": "yaw",
|
||||
"initial_range": (-0.5, 0.5),
|
||||
"delta_command": 0.05,
|
||||
"target_ratio": 0.8,
|
||||
"ema_alpha": 0.5,
|
||||
},
|
||||
)
|
||||
|
||||
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"].rel_forward_envs = 0.30
|
||||
cfg.commands["twist"].rel_lateral_envs = 0.20
|
||||
cfg.commands["twist"].rel_yaw_envs = 0.20
|
||||
cfg.commands["twist"].ranges.lin_vel_x = (-1.0, 1.0)
|
||||
cfg.commands["twist"].ranges.lin_vel_y = (-1.0, 1.0)
|
||||
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",
|
||||
params={
|
||||
"pose_range": {"z": (0.40, 0.45), "yaw": (-math.pi, math.pi)},
|
||||
"pose_range": {"z": (0.42, 0.42), "yaw": (-math.pi, math.pi)},
|
||||
"velocity_range": {"x": (-0.2, 0.2), "y": (-0.1, 0.1), "yaw": (-0.2, 0.2)},
|
||||
"asset_cfg": SceneEntityCfg("robot"),
|
||||
},
|
||||
@@ -425,43 +501,54 @@ def rough_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
|
||||
# ------------------
|
||||
# Rewards Integration
|
||||
# ------------------
|
||||
cfg.rewards["track_lin_vel"] = RewardTermCfg(
|
||||
func=track_linear_velocity_l1,
|
||||
weight=4.5,
|
||||
params={"std": 0.5, "command_name": "twist"}
|
||||
cfg.rewards.pop("track_lin_vel", None)
|
||||
cfg.rewards.pop("track_ang_vel", None)
|
||||
cfg.rewards["track_lin_vel_x_exp"] = RewardTermCfg(
|
||||
func=track_linear_velocity_x,
|
||||
weight=1.0,
|
||||
params={"std": 0.25, "command_name": "twist"},
|
||||
)
|
||||
cfg.rewards["track_lin_vel_y_exp"] = RewardTermCfg(
|
||||
func=track_linear_velocity_y,
|
||||
weight=1.0,
|
||||
params={"std": 0.25, "command_name": "twist"},
|
||||
)
|
||||
cfg.rewards["track_ang_vel_z_exp"] = RewardTermCfg(
|
||||
func=track_angular_velocity_z,
|
||||
weight=1.0,
|
||||
params={"std": 0.25, "command_name": "twist"},
|
||||
)
|
||||
cfg.rewards["stair_lateral_yaw_drift"] = RewardTermCfg(
|
||||
func=stair_lateral_yaw_drift_l2,
|
||||
weight=-1.0,
|
||||
params={
|
||||
"terrain_names": ("pyramid_stairs", "pyramid_stairs_inv", "random_grid"),
|
||||
"y_scale": 1.0,
|
||||
"yaw_scale": 1.0,
|
||||
"asset_cfg": SceneEntityCfg("robot"),
|
||||
},
|
||||
)
|
||||
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)
|
||||
@@ -470,8 +557,8 @@ def rough_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
|
||||
weight=-0.05,
|
||||
params={
|
||||
"mirror_joints": [
|
||||
["fl_(hip_abduction|hip_pitch|knee)_joint", "rr_(hip_abduction|hip_pitch|knee)_joint"],
|
||||
["fr_(hip_abduction|hip_pitch|knee)_joint", "rl_(hip_abduction|hip_pitch|knee)_joint"]
|
||||
["fl_(hip_pitch|knee)_joint", "rr_(hip_pitch|knee)_joint"],
|
||||
["fr_(hip_pitch|knee)_joint", "rl_(hip_pitch|knee)_joint"]
|
||||
],
|
||||
"asset_cfg": SceneEntityCfg("robot")
|
||||
}
|
||||
@@ -480,64 +567,157 @@ def rough_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
|
||||
# 移除 variable_posture 及其产生的静止奖励陷阱，换用极轻微的偏离惩罚
|
||||
cfg.rewards.pop("stand_still", None)
|
||||
cfg.rewards["stand_still"] = RewardTermCfg(func=stand_still, weight=-2.0, params={"command_name": "twist", "command_threshold": 0.1})
|
||||
|
||||
|
||||
cfg.rewards.pop("hip_deviation", None)
|
||||
cfg.rewards.pop("variable_posture", None)
|
||||
|
||||
|
||||
cfg.rewards.pop("joint_deviation_l2", None)
|
||||
cfg.rewards["joint_pos_penalty"] = RewardTermCfg(
|
||||
func=joint_pos_penalty,
|
||||
weight=-0.8,
|
||||
|
||||
# 针对 ab 关节施加较严厉的惩罚，防止在 yaw 时乱撇腿
|
||||
cfg.rewards["joint_pos_penalty_ab"] = RewardTermCfg(
|
||||
func=joint_pos_penalty,
|
||||
weight=-1.0,
|
||||
params={
|
||||
"stand_still_scale": 5.0,
|
||||
"velocity_threshold": 0.5,
|
||||
"command_threshold": 0.1,
|
||||
"asset_cfg": SceneEntityCfg("robot", joint_names=(".*_hip_abduction_joint", ".*_hip_pitch_joint", ".*_knee_joint")),
|
||||
"asset_cfg": SceneEntityCfg("robot", joint_names=(".*_hip_abduction_joint",)),
|
||||
"command_name": "twist"
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
# 针对 pitch 和 knee 关节施加较宽松的惩罚，保留跨越障碍的抬腿自由度
|
||||
cfg.rewards["joint_pos_penalty_sagittal"] = RewardTermCfg(
|
||||
func=joint_pos_penalty,
|
||||
weight=-0.3,
|
||||
params={
|
||||
"stand_still_scale": 5.0,
|
||||
"velocity_threshold": 0.5,
|
||||
"command_threshold": 0.1,
|
||||
"asset_cfg": SceneEntityCfg("robot", joint_names=(".*_hip_pitch_joint", ".*_knee_joint")),
|
||||
"command_name": "twist"
|
||||
}
|
||||
)
|
||||
|
||||
# 🌟 强力约束同侧外展关节平行对称，消除转向时的前后剪刀式摆动
|
||||
cfg.rewards["abduction_mirror"] = RewardTermCfg(
|
||||
func=joint_mirror,
|
||||
weight=-0.5, # 施加合理惩罚，限制前后腿同侧外展关节反向运动
|
||||
params={
|
||||
"mirror_joints": [
|
||||
["fl_hip_abduction_joint", "rl_hip_abduction_joint"],
|
||||
["fr_hip_abduction_joint", "rr_hip_abduction_joint"]
|
||||
],
|
||||
"asset_cfg": SceneEntityCfg("robot")
|
||||
}
|
||||
)
|
||||
|
||||
cfg.rewards["feet_contact_without_cmd"] = RewardTermCfg(
|
||||
func=feet_contact_without_cmd,
|
||||
weight=0.1,
|
||||
func=feet_contact_without_cmd,
|
||||
weight=0.1,
|
||||
params={"command_name": "twist", "sensor_name": "feet_ground_contact"}
|
||||
)
|
||||
cfg.rewards["feet_air_time"].weight = 0.15
|
||||
cfg.rewards["upward"] = RewardTermCfg(func=upward, weight=0.5)
|
||||
|
||||
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.42
|
||||
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
|
||||
cfg.scene.terrain.env_spacing = 2.5
|
||||
|
||||
cfg.metrics.update(
|
||||
{
|
||||
"tracking_lin_vel_error": MetricsTermCfg(func=tracking_lin_vel_error, params={"command_name": "twist"}),
|
||||
"tracking_lin_vel_x_error": MetricsTermCfg(func=tracking_lin_vel_x_error, params={"command_name": "twist"}),
|
||||
"tracking_lin_vel_y_error": MetricsTermCfg(func=tracking_lin_vel_y_error, params={"command_name": "twist"}),
|
||||
"tracking_lin_vel_along_cmd_error": MetricsTermCfg(
|
||||
func=tracking_lin_vel_along_command_error, params={"command_name": "twist"}
|
||||
),
|
||||
"actual_lin_vel_orthogonal_cmd": MetricsTermCfg(
|
||||
func=actual_lin_vel_orthogonal_command_mean, params={"command_name": "twist"}
|
||||
),
|
||||
"tracking_yaw_vel_error": MetricsTermCfg(func=tracking_yaw_vel_error, params={"command_name": "twist"}),
|
||||
"cmd_lin_vel": MetricsTermCfg(func=command_lin_vel_mean, params={"command_name": "twist"}),
|
||||
"cmd_yaw_vel": MetricsTermCfg(func=command_yaw_vel_abs_mean, params={"command_name": "twist"}),
|
||||
"actual_lin_vel": MetricsTermCfg(func=actual_lin_vel_mean),
|
||||
"tracking_lin_vel_error_cmd_0_03": MetricsTermCfg(
|
||||
func=tracking_lin_vel_error_band_mean,
|
||||
params={"command_name": "twist", "min_speed": 0.0, "max_speed": 0.3},
|
||||
),
|
||||
"tracking_lin_vel_error_cmd_03_07": MetricsTermCfg(
|
||||
func=tracking_lin_vel_error_band_mean,
|
||||
params={"command_name": "twist", "min_speed": 0.3, "max_speed": 0.7},
|
||||
),
|
||||
"tracking_lin_vel_error_cmd_07_up": MetricsTermCfg(
|
||||
func=tracking_lin_vel_error_band_mean,
|
||||
params={"command_name": "twist", "min_speed": 0.7, "max_speed": 10.0},
|
||||
),
|
||||
"tracking_lin_vel_x_error_cmd_0_03": MetricsTermCfg(
|
||||
func=tracking_lin_vel_axis_error_band_mean,
|
||||
params={"axis": 0, "command_name": "twist", "min_speed": 0.0, "max_speed": 0.3},
|
||||
),
|
||||
"tracking_lin_vel_x_error_cmd_03_07": MetricsTermCfg(
|
||||
func=tracking_lin_vel_axis_error_band_mean,
|
||||
params={"axis": 0, "command_name": "twist", "min_speed": 0.3, "max_speed": 0.7},
|
||||
),
|
||||
"tracking_lin_vel_x_error_cmd_07_up": MetricsTermCfg(
|
||||
func=tracking_lin_vel_axis_error_band_mean,
|
||||
params={"axis": 0, "command_name": "twist", "min_speed": 0.7, "max_speed": 10.0},
|
||||
),
|
||||
"tracking_lin_vel_y_error_cmd_0_03": MetricsTermCfg(
|
||||
func=tracking_lin_vel_axis_error_band_mean,
|
||||
params={"axis": 1, "command_name": "twist", "min_speed": 0.0, "max_speed": 0.3},
|
||||
),
|
||||
"tracking_lin_vel_y_error_cmd_03_07": MetricsTermCfg(
|
||||
func=tracking_lin_vel_axis_error_band_mean,
|
||||
params={"axis": 1, "command_name": "twist", "min_speed": 0.3, "max_speed": 0.7},
|
||||
),
|
||||
"tracking_lin_vel_y_error_cmd_07_up": MetricsTermCfg(
|
||||
func=tracking_lin_vel_axis_error_band_mean,
|
||||
params={"axis": 1, "command_name": "twist", "min_speed": 0.7, "max_speed": 10.0},
|
||||
),
|
||||
"cmd_band_0_03": MetricsTermCfg(
|
||||
func=command_band_active, params={"command_name": "twist", "min_speed": 0.0, "max_speed": 0.3}
|
||||
),
|
||||
"cmd_band_03_07": MetricsTermCfg(
|
||||
func=command_band_active, params={"command_name": "twist", "min_speed": 0.3, "max_speed": 0.7}
|
||||
),
|
||||
"cmd_band_07_up": MetricsTermCfg(
|
||||
func=command_band_active, params={"command_name": "twist", "min_speed": 0.7, "max_speed": 10.0}
|
||||
),
|
||||
"wheel_raw_action_abs": MetricsTermCfg(func=wheel_raw_action_abs_mean),
|
||||
"wheel_target_vel_abs": MetricsTermCfg(func=wheel_target_vel_abs_mean),
|
||||
"wheel_actual_vel_abs": MetricsTermCfg(func=wheel_actual_vel_abs_mean),
|
||||
"wheel_target_actual_vel_error": MetricsTermCfg(func=wheel_target_actual_vel_error_mean),
|
||||
"wheel_actual_to_target_vel_ratio": MetricsTermCfg(func=wheel_actual_to_target_vel_ratio_mean),
|
||||
"wheel_target_actual_sign_agreement": MetricsTermCfg(func=wheel_target_actual_sign_agreement),
|
||||
"upright": MetricsTermCfg(func=upright_metric),
|
||||
"base_ground_contact_rate": MetricsTermCfg(
|
||||
func=base_ground_contact_metric, params={"sensor_name": "base_ground_contact"}
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
if play:
|
||||
cfg.episode_length_s = int(1e9)
|
||||
cfg.observations["actor"].enable_corruption = False
|
||||
@@ -663,7 +843,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,
|
||||
num_learning_epochs=5,
|
||||
num_mini_batches=4,
|
||||
learning_rate=2.0e-4,
|
||||
entropy_coef=0.003,#第一轮为0.003，第二轮0.001，第三轮0.0008
|
||||
num_learning_epochs=5,#2048为5 4096为3
|
||||
num_mini_batches=4,#2048为4 4096为8
|
||||
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
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -21,14 +21,20 @@ class UniformThresholdVelocityCommand(UniformVelocityCommand):
|
||||
|
||||
def __init__(self, cfg: UniformThresholdVelocityCommandCfg, env: ManagerBasedRlEnv):
|
||||
super().__init__(cfg, env)
|
||||
self.is_lateral_env = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
|
||||
self.is_yaw_env = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
|
||||
self.was_climbing = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
|
||||
# 缓存地形类型的索引（台阶、反向台阶、垂直短墙），实现高容错动态查找
|
||||
self._climbing_indices = []
|
||||
self._flat_index = -1
|
||||
terrain = getattr(self._env.scene, "terrain", None)
|
||||
if terrain is not None and getattr(terrain.cfg, "terrain_generator", None) is not None:
|
||||
sub_terrain_names = list(terrain.cfg.terrain_generator.sub_terrains.keys())
|
||||
for name in ["pyramid_stairs", "pyramid_stairs_inv", "rc_wall"]:
|
||||
if name in sub_terrain_names:
|
||||
self._climbing_indices.append(sub_terrain_names.index(name))
|
||||
if "flat" in sub_terrain_names:
|
||||
self._flat_index = sub_terrain_names.index("flat")
|
||||
|
||||
def _resample_command(self, env_ids: torch.Tensor) -> None:
|
||||
# 1. 调用基类的标准采样
|
||||
@@ -42,6 +48,41 @@ class UniformThresholdVelocityCommand(UniformVelocityCommand):
|
||||
self.vel_command_b[small_cmd_ids, :] = 0.0
|
||||
self.vel_command_w[small_cmd_ids, :] = 0.0
|
||||
|
||||
# 重置并采样单 Y（只横移）与单 Z（只原地转向）指令分布
|
||||
self.is_lateral_env[env_ids] = False
|
||||
self.is_yaw_env[env_ids] = False
|
||||
|
||||
# 对没有设为前进且没有静止的激活环境进行单轴采样
|
||||
active_non_fwd_mask = (~self.is_forward_env[env_ids]) & (~self.is_standing_env[env_ids])
|
||||
active_non_fwd_ids = env_ids[active_non_fwd_mask]
|
||||
|
||||
if len(active_non_fwd_ids) > 0:
|
||||
r = torch.empty(len(active_non_fwd_ids), device=self.device)
|
||||
# 采样单 Y 占比
|
||||
self.is_lateral_env[active_non_fwd_ids] = r.uniform_(0.0, 1.0) <= self.cfg.rel_lateral_envs
|
||||
lat_ids = active_non_fwd_ids[self.is_lateral_env[active_non_fwd_ids]]
|
||||
if len(lat_ids) > 0:
|
||||
self.vel_command_b[lat_ids, 0] = 0.0
|
||||
y_signs = self.vel_command_b[lat_ids, 1].sign()
|
||||
y_signs[y_signs == 0] = 1.0
|
||||
self.vel_command_b[lat_ids, 1] = y_signs * self.vel_command_b[lat_ids, 1].abs().clamp(min=0.3)
|
||||
self.vel_command_b[lat_ids, 2] = 0.0
|
||||
|
||||
# 对不是 forward 也不是 lateral 的环境采样单 Z 占比
|
||||
non_lat_mask = ~self.is_lateral_env[active_non_fwd_ids]
|
||||
non_lat_ids = active_non_fwd_ids[non_lat_mask]
|
||||
|
||||
if len(non_lat_ids) > 0:
|
||||
r_yaw = torch.empty(len(non_lat_ids), device=self.device)
|
||||
self.is_yaw_env[non_lat_ids] = r_yaw.uniform_(0.0, 1.0) <= self.cfg.rel_yaw_envs
|
||||
yaw_ids = non_lat_ids[self.is_yaw_env[non_lat_ids]]
|
||||
if len(yaw_ids) > 0:
|
||||
self.vel_command_b[yaw_ids, 0] = 0.0
|
||||
self.vel_command_b[yaw_ids, 1] = 0.0
|
||||
z_signs = self.vel_command_b[yaw_ids, 2].sign()
|
||||
z_signs[z_signs == 0] = 1.0
|
||||
self.vel_command_b[yaw_ids, 2] = z_signs * self.vel_command_b[yaw_ids, 2].abs().clamp(min=0.3)
|
||||
|
||||
# 3. 地形自适应重采样限制：若在爬行地形，强制纯前进方向且速度 >= 0.3 m/s
|
||||
terrain = getattr(self._env.scene, "terrain", None)
|
||||
terrain_types = getattr(terrain, "terrain_types", None)
|
||||
@@ -58,6 +99,37 @@ class UniformThresholdVelocityCommand(UniformVelocityCommand):
|
||||
self.vel_command_b[climbing_env_ids, 2] = 0.0
|
||||
self.is_heading_env[climbing_env_ids] = True
|
||||
|
||||
# 4. 在平地（flat）地形上生成 35% 纯 Y 和 35% 纯 Z 指令，排除静止环境
|
||||
if terrain_types is not None and self._flat_index != -1:
|
||||
is_flat = terrain_types == self._flat_index
|
||||
flat_env_ids = env_ids[is_flat[env_ids]]
|
||||
if len(flat_env_ids) > 0:
|
||||
active_flat_mask = ~self.is_standing_env[flat_env_ids]
|
||||
active_flat_ids = flat_env_ids[active_flat_mask]
|
||||
if len(active_flat_ids) > 0:
|
||||
r = torch.empty(len(active_flat_ids), device=self.device).uniform_(0.0, 1.0)
|
||||
y_only_mask = r < 0.35
|
||||
z_only_mask = (r >= 0.35) & (r < 0.70)
|
||||
|
||||
y_only_ids = active_flat_ids[y_only_mask]
|
||||
if len(y_only_ids) > 0:
|
||||
self.vel_command_b[y_only_ids, 0] = 0.0 # x = 0
|
||||
self.vel_command_b[y_only_ids, 2] = 0.0 # yaw = 0
|
||||
if self.cfg.heading_command:
|
||||
self.heading_target[y_only_ids] = self.robot.data.heading_w[y_only_ids]
|
||||
self.is_heading_env[y_only_ids] = True
|
||||
|
||||
z_only_ids = active_flat_ids[z_only_mask]
|
||||
if len(z_only_ids) > 0:
|
||||
self.vel_command_b[z_only_ids, 0] = 0.0 # x = 0
|
||||
self.vel_command_b[z_only_ids, 1] = 0.0 # y = 0
|
||||
if self.cfg.heading_command:
|
||||
yaw_delta = torch.empty(len(z_only_ids), device=self.device).uniform_(
|
||||
-self.cfg.yaw_only_heading_range, self.cfg.yaw_only_heading_range
|
||||
)
|
||||
self.heading_target[z_only_ids] = self.robot.data.heading_w[z_only_ids] + yaw_delta
|
||||
self.is_heading_env[z_only_ids] = True
|
||||
|
||||
def _update_command(self) -> None:
|
||||
# 调用基类的每步更新
|
||||
super()._update_command()
|
||||
@@ -69,7 +141,14 @@ class UniformThresholdVelocityCommand(UniformVelocityCommand):
|
||||
is_climbing = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
|
||||
for idx in self._climbing_indices:
|
||||
is_climbing |= (terrain_types == idx)
|
||||
|
||||
|
||||
# 【出台阶实时重采样】：检查从爬坡状态刚刚离开环境的机器人，进行即时指令重采样
|
||||
left_climbing_mask = self.was_climbing & ~is_climbing
|
||||
if left_climbing_mask.any():
|
||||
left_climbing_ids = torch.where(left_climbing_mask)[0]
|
||||
self._resample_command(left_climbing_ids)
|
||||
|
||||
# 【入台阶实时截断】：对于正在爬楼梯/翻越短墙的环境，强行截断其横向横移指令，并进行温和偏航对齐
|
||||
climbing_env_ids = is_climbing.nonzero(as_tuple=False).flatten()
|
||||
if len(climbing_env_ids) > 0:
|
||||
self.vel_command_b[climbing_env_ids, 1] = 0.0
|
||||
@@ -79,6 +158,7 @@ class UniformThresholdVelocityCommand(UniformVelocityCommand):
|
||||
min=-0.3,
|
||||
max=0.3
|
||||
)
|
||||
self.was_climbing = is_climbing
|
||||
|
||||
# 5. 高速侧向解耦（适用于平地/斜坡等混合路面）：当前进速度 >= 0.8 m/s 时，清空侧向指令，防止高速甩尾甩飞
|
||||
high_speed_mask = self.vel_command_b[:, 0].abs() >= 0.8
|
||||
@@ -90,3 +170,6 @@ class UniformThresholdVelocityCommand(UniformVelocityCommand):
|
||||
@dataclass(kw_only=True)
|
||||
class UniformThresholdVelocityCommandCfg(UniformVelocityCommandCfg):
|
||||
class_type: type = UniformThresholdVelocityCommand
|
||||
yaw_only_heading_range: float = 1.0
|
||||
rel_lateral_envs: float = 0.0
|
||||
rel_yaw_envs: float = 0.0
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user