[train] 更新新MJCF与第一版完整训练框架
This commit is contained in:
@@ -35,6 +35,7 @@ logs/
|
||||
checkpoints/
|
||||
wandb/
|
||||
sim2sim_log_*.txt
|
||||
**/sim2sim_temp.xml
|
||||
|
||||
# IDE and operating system files
|
||||
.idea/
|
||||
|
||||
@@ -0,0 +1,19 @@
|
||||
# 版本演进
|
||||
|
||||
本项目使用同一条 `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.4.0` 的模型变化
|
||||
|
||||
- 机械 CAD 不变。
|
||||
- MJCF 更新整机质量和惯性参数，旧、新 `wheelleg.xml` 的 SHA-256 不同。
|
||||
- mjlab 上游基准从 `00409797` 更新到 `40f8d93e`。
|
||||
- 保留轮腿分组执行器随机化所需的本地补丁。
|
||||
- 本阶段归档 `model_rough.pt`，不将生成日志、缓存和临时 XML 纳入版本库。
|
||||
@@ -1,11 +1,11 @@
|
||||
# 软件
|
||||
|
||||
本目录当前保存 16DOF 轮足机器人的第一代完整软件闭环。
|
||||
本目录保存 16DOF 轮足机器人的训练、仿真和真机软件演进。
|
||||
|
||||
```text
|
||||
05_software/
|
||||
├─ train/
|
||||
│ └─ rc_mjlab/ # 训练、MJCF、MuJoCo、Sim2Sim 和本地 mjlab 依赖
|
||||
│ └─ rc_mjlab/ # 训练、MJCF、Sim2Sim 和本地 mjlab 依赖
|
||||
└─ real/
|
||||
├─ ik_real/ # IK 轨迹与早期真机控制
|
||||
└─ sim2real/ # 第一代 Python 策略真机部署
|
||||
@@ -18,8 +18,6 @@ MJCF + mjlab task
|
||||
|
|
||||
v
|
||||
PPO 训练策略
|
||||
|
|
||||
+----> MuJoCo 独立模型调试
|
||||
|
|
||||
+----> Sim2Sim 策略验证
|
||||
|
|
||||
@@ -28,7 +26,7 @@ MJCF + mjlab task
|
||||
IK real --------------------------------> 电机
|
||||
```
|
||||
|
||||
`rc_mjlab` 在早期版本中是自包含工程。训练、MJCF、独立 MuJoCo、Sim2Sim 和策略权重通过相对路径绑定，因此本次保留其原始内部布局，没有为了目录外观拆散。
|
||||
`rc_mjlab` 是自包含工程。训练、MJCF、Sim2Sim 和策略权重通过相对路径绑定，因此保留其内部布局，没有为了目录外观拆散。第一代完整闭环见 `v0.3.0`，第一份新版 MJCF 与训练框架见 `v0.4.0`。
|
||||
|
||||
详细说明见：
|
||||
|
||||
|
||||
@@ -1,15 +1,18 @@
|
||||
# 第一代强化学习与仿真工程
|
||||
# 强化学习与仿真工程
|
||||
|
||||
`rc_mjlab/` 是 16DOF 轮足机器人的第一代自包含训练与仿真工程。
|
||||
`rc_mjlab/` 保存 16DOF 轮足机器人的当前训练与 Sim2Sim 工程。历史快照由 Git Tag 保留，不在目录中复制 `old`、`new` 或 `final` 版本。
|
||||
|
||||
当前内容对应 `v0.4.0`，是第一份完整采用新版 MJCF 和新版 mjlab 框架的训练工程。
|
||||
|
||||
## 内容
|
||||
|
||||
- `src/robot`：Flat、Rough、Crawl 训练任务和自定义 MDP
|
||||
- `mjcf`：轮足机器人 MuJoCo 模型和网格
|
||||
- `mujoco_sim`：不依赖策略的独立 MuJoCo/MPC 调试工具
|
||||
- `sim2sim`：策略加载、交互控制和比赛地形验证
|
||||
- `mjlab`：固定版本的本地训练框架依赖
|
||||
- `model_rough.pt`、`model_crawl.pt`：对应的早期策略权重
|
||||
- `model_rough.pt`：本阶段 Rough 策略权重
|
||||
- `pyproject.toml`、`uv.lock`：Python 环境与依赖锁定
|
||||
|
||||
与 `v0.3.0` 相比，本版本更新了 MJCF 质量和惯性参数，并将 mjlab 上游基准从 `00409797` 更新到 `40f8d93e`。机械 CAD 未发生变化。
|
||||
|
||||
工程命令和任务说明见 [`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.4.0`：第一份完整的新 MJCF 与新框架训练工程。该快照包含 `model_rough.pt`；未包含独立 `mujoco_sim` 工具和单独的 Crawl 策略权重，相关早期内容仍可通过 `v0.3.0` 查看。
|
||||
|
||||
---
|
||||
|
||||
## 🛠️ 项目简介
|
||||
@@ -51,8 +53,7 @@ rc_mjlab/
|
||||
│ ├── wheelleg.xml # 机器人 MuJoCo 模型（含网格引用）
|
||||
│ ├── scene.xml # mjlab 场景入口文件
|
||||
│ └── meshes/ # STL/OBJ 碰撞与外观网格
|
||||
├── mujoco_sim/ # 独立 MPC 仿真调试工具（不依赖 RL 训练）
|
||||
├── logs/ # 训练日志（rsl_rl 格式，按任务名/日期/checkpoint 归档）
|
||||
├── model_rough.pt # 本阶段用于回放和 Sim2Sim 的 Rough 策略
|
||||
├── pyproject.toml # 项目依赖（uv 管理，含清华镜像源加速）
|
||||
└── uv.lock # 精确依赖锁定文件
|
||||
```
|
||||
|
||||
@@ -1,327 +0,0 @@
|
||||
<mujoco model="go2w scene">
|
||||
<include file="C:/Users/31560/Documents/00_legged/new_rl/rc_mjlab/mjcf/wheelleg.xml"/>
|
||||
<statistic center="3.7 -9.0 0.4" extent="5.0"/>
|
||||
<visual>
|
||||
<headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0"/>
|
||||
<rgba haze="0.15 0.25 0.35 1"/>
|
||||
<global azimuth="90" elevation="-20"/>
|
||||
</visual>
|
||||
<asset>
|
||||
<texture type="skybox" builtin="gradient" rgb1="0.3 0.5 0.7" rgb2="0 0 0" width="512" height="3072"/>
|
||||
<texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3" markrgb="0.8 0.8 0.8" width="300" height="300"/>
|
||||
<material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance="0.2"/>
|
||||
<hfield name="perlin_hfield" size="1.0 0.75 0.2 0.2" file="C:/Users/31560/Documents/00_legged/new_rl/rc_mjlab/sim2sim/terrain/height_field.png"/>
|
||||
<hfield name="image_hfield" size="1.0 1.0 0.02 0.1" file="C:/Users/31560/Documents/00_legged/new_rl/rc_mjlab/sim2sim/terrain/unitree_hfield.png"/>
|
||||
</asset>
|
||||
<worldbody>
|
||||
<light pos="0 0 1.5" dir="0 0 -1" directional="true" />
|
||||
<geom name="floor" size="0 0 0.05" type="plane" material="groundplane" />
|
||||
|
||||
<!-- 30cm高墙：旋转90度，沿x轴方向放置，并与T型楼梯中心线 y=-3.50 对齐 -->
|
||||
<geom pos="1.8 -7.0 0.15"
|
||||
type="box"
|
||||
size="0.025 0.5 0.15"
|
||||
quat="0.7071068 0.0 0.0 0.7071068"
|
||||
rgba="1.0 0.9 0.4 1.0"/>
|
||||
|
||||
<!-- 沙砾碎木坑：x正方向边界与10度斜坡+x边界对齐，y正边界距斜坡y负边界4m -->
|
||||
|
||||
<geom pos="4.8361 -12.5 0.075"
|
||||
type="box"
|
||||
size="0.5 0.5 0.075"
|
||||
quat="0.0 0.0 0.0 1.0"
|
||||
rgba="0.75 0.72 0.55 1.0"/>
|
||||
|
||||
<geom pos="5.8361 -12.0 0.075"
|
||||
type="box"
|
||||
size="0.5 1.0 0.075"
|
||||
quat="0.0 0.0 0.0 1.0"
|
||||
rgba="0.75 0.72 0.55 1.0"/>
|
||||
|
||||
|
||||
|
||||
|
||||
<!-- 限高杆 -->
|
||||
<geom pos="6.2 -9.0 0.155"
|
||||
type="cylinder"
|
||||
size="0.025 0.155"
|
||||
quat="1.0 0.0 0.0 0.0"
|
||||
rgba="0.8 0.1 0.1 1.0" />
|
||||
|
||||
<geom pos="5.2 -9.0 0.155"
|
||||
type="cylinder"
|
||||
size="0.025 0.155"
|
||||
quat="1.0 0.0 0.0 0.0"
|
||||
rgba="0.8 0.1 0.1 1.0" />
|
||||
<geom pos="5.7 -9.0 0.325"
|
||||
type="cylinder"
|
||||
size="0.015 0.5"
|
||||
quat="0.7071068 0.0 0.7071068 0.0"
|
||||
rgba="1.0 0.9 0.4 1.0"/>
|
||||
|
||||
|
||||
<!-- 1m × 1m 正方形颜色块，出发区-->
|
||||
<geom pos="3.7 -9.0 0.0"
|
||||
type="box"
|
||||
size="0.5 0.5 0.001"
|
||||
rgba="1.0 0.0 0.0 0.35"
|
||||
contype="0"
|
||||
conaffinity="0" />
|
||||
|
||||
|
||||
|
||||
<!-- 10cm梯形台阶：T型楼梯，最高平台与 x=5.7 y=-3.5 平台在y轴方向对齐 -->
|
||||
|
||||
<geom pos="1.80 -4.75 0.05" type="box" size="0.15 0.5 0.05" quat="0.7071068 0.0 0.0 0.7071068" rgba="0.75 0.72 0.55 1.0" />
|
||||
<geom pos="1.80 -4.45 0.15" type="box" size="0.15 0.5 0.05" quat="0.7071068 0.0 0.0 0.7071068" rgba="0.75 0.72 0.55 1.0"/>
|
||||
<geom pos="1.80 -4.15 0.25" type="box" size="0.15 0.5 0.05" quat="0.7071068 0.0 0.0 0.7071068" rgba="0.75 0.72 0.55 1.0"/>
|
||||
|
||||
<!-- 最高平台：y = -3.50，与目标平台y轴对齐 -->
|
||||
<geom pos="1.80 -3.50 0.35" type="box" size="0.5 0.5 0.05" quat="0.7071068 0.0 0.0 0.7071068" rgba="0.75 0.72 0.55 1.0"/>
|
||||
|
||||
<geom pos="1.80 -2.85 0.25" type="box" size="0.15 0.5 0.05" quat="0.7071068 0.0 0.0 0.7071068" rgba="0.75 0.72 0.55 1.0"/>
|
||||
<geom pos="1.80 -2.55 0.15" type="box" size="0.15 0.5 0.05" quat="0.7071068 0.0 0.0 0.7071068" rgba="0.75 0.72 0.55 1.0"/>
|
||||
<geom pos="1.80 -2.25 0.05" type="box" size="0.15 0.5 0.05" quat="0.7071068 0.0 0.0 0.7071068" rgba="0.75 0.72 0.55 1.0"/>
|
||||
|
||||
|
||||
<!-- 顶部平台向 +x 方向连接地面的10cm台阶，同样y轴移动到 -3.50 -->
|
||||
|
||||
<geom pos="2.45 -3.50 0.25" type="box" size="0.15 0.5 0.05" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
|
||||
<geom pos="2.75 -3.50 0.15" type="box" size="0.15 0.5 0.05" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
|
||||
<geom pos="3.05 -3.50 0.05" type="box" size="0.15 0.5 0.05" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
|
||||
|
||||
|
||||
|
||||
|
||||
<!-- 斜坡木桥A木桥B -->
|
||||
<geom pos="1.8 -0.88 0.1" type="box" size="0.40 0.5 0.005" quat="0.701836 0.086175 -0.086175 0.701836" rgba="0.75 0.72 0.55 1.0"/>
|
||||
<geom pos="1.8 0.0 0.10" type="box" size="0.5 0.5 0.10" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
|
||||
<geom pos="2.65 0.0 0.10" type="box" size="0.2 0.5 0.10" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
|
||||
<geom pos="3.2 0.0 0.10" type="box" size="0.2 0.5 0.10" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
|
||||
<geom pos="3.75 0.0 0.10" type="box" size="0.2 0.5 0.10" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
|
||||
<geom pos="4.3 0.0 0.10" type="box" size="0.2 0.5 0.10" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
|
||||
<geom pos="4.85 0.0 0.10" type="box" size="0.2 0.5 0.10" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
|
||||
<geom pos="5.7 -0.5 0.10" type="box" size="0.5 1.0 0.10" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
|
||||
<geom pos="4.8002 -1.0 0.0951"
|
||||
type="box"
|
||||
size="0.4133 0.5 0.005"
|
||||
quat="0.992546 0.0 -0.121869 0.0"
|
||||
rgba="0.75 0.72 0.55 1.0"/>
|
||||
<geom pos="5.35 -2.25 0.10" type="box" size="0.1 0.75 0.10" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
|
||||
<geom pos="5.65 -2.25 0.10" type="box" size="0.1 0.75 0.10" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
|
||||
<geom pos="5.95 -2.25 0.10" type="box" size="0.1 0.75 0.10" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
|
||||
<geom pos="5.7 -3.5 0.10" type="box" size="0.5 0.5 0.10" quat="1.0 0.0 0.0 0.0" rgba="0.75 0.72 0.55 1.0"/>
|
||||
<!-- 10度斜坡：宽4m，斜坡 y正方向边缘 与平台 y正方向边缘 对齐 -->
|
||||
<geom pos="4.6338 -5.0 0.0951"
|
||||
type="box"
|
||||
size="0.5759 2.0 0.005"
|
||||
quat="0.9961947 0.0 -0.0871557 0.0"
|
||||
rgba="0.75 0.72 0.55 1.0"/>
|
||||
<!-- 新建10度斜坡：宽3m，高端与前一个10度斜坡高端衔接，向+x方向下坡 -->
|
||||
<geom pos="5.7681 -5.5 0.0951"
|
||||
type="box"
|
||||
size="0.5759 1.5 0.005"
|
||||
quat="0.9961947 0.0 0.0871557 0.0"
|
||||
rgba="0.75 0.72 0.55 1.0"/>
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
<!--绕杆-->
|
||||
<!-- 直径1m圆形颜色块，仅显示，不碰撞 -->
|
||||
<geom pos="1.8 -10.1 0.0"
|
||||
type="cylinder"
|
||||
size="0.1 0.001"
|
||||
rgba="1.0 0.0 0.0 0.35"
|
||||
contype="0"
|
||||
conaffinity="0" />
|
||||
|
||||
<geom pos="3.2 -12.5 0.0"
|
||||
type="cylinder"
|
||||
size="0.1 0.001"
|
||||
rgba="1.0 0.0 0.0 0.35"
|
||||
contype="0"
|
||||
conaffinity="0" />
|
||||
|
||||
<geom pos="1.55 -12.75 0.0"
|
||||
type="cylinder"
|
||||
size="0.1 0.001"
|
||||
rgba="1.0 0.0 0.0 0.35"
|
||||
contype="0"
|
||||
conaffinity="0" />
|
||||
|
||||
<!-- 原杆 -->
|
||||
<geom pos="1.8 -10.5 0.02"
|
||||
type="cylinder"
|
||||
size="0.05 0.02"
|
||||
quat="1.0 0.0 0.0 0.0"
|
||||
rgba="0.75 0.72 0.55 1.0" />
|
||||
|
||||
<geom pos="1.8 -10.5 0.37"
|
||||
type="cylinder"
|
||||
size="0.015 0.33"
|
||||
quat="1.0 0.0 0.0 0.0"
|
||||
rgba="0.75 0.72 0.55 1.0" />
|
||||
|
||||
|
||||
<!-- y轴负方向第1根：间隔1m -->
|
||||
<geom pos="1.8 -11.5 0.02"
|
||||
type="cylinder"
|
||||
size="0.05 0.02"
|
||||
quat="1.0 0.0 0.0 0.0"
|
||||
rgba="0.75 0.72 0.55 1.0" />
|
||||
|
||||
<geom pos="1.8 -11.5 0.37"
|
||||
type="cylinder"
|
||||
size="0.015 0.33"
|
||||
quat="1.0 0.0 0.0 0.0"
|
||||
rgba="0.75 0.72 0.55 1.0" />
|
||||
|
||||
|
||||
<!-- y轴负方向第2根：继续间隔1m -->
|
||||
<geom pos="1.8 -12.5 0.02"
|
||||
type="cylinder"
|
||||
size="0.05 0.02"
|
||||
quat="1.0 0.0 0.0 0.0"
|
||||
rgba="0.75 0.72 0.55 1.0" />
|
||||
|
||||
<geom pos="1.8 -12.5 0.37"
|
||||
type="cylinder"
|
||||
size="0.015 0.33"
|
||||
quat="1.0 0.0 0.0 0.0"
|
||||
rgba="0.75 0.72 0.55 1.0" />
|
||||
|
||||
|
||||
<!-- x轴正方向第3根：继续间隔1m -->
|
||||
<geom pos="2.8 -12.5 0.02"
|
||||
type="cylinder"
|
||||
size="0.05 0.02"
|
||||
quat="1.0 0.0 0.0 0.0"
|
||||
rgba="0.75 0.72 0.55 1.0" />
|
||||
|
||||
<geom pos="2.8 -12.5 0.37"
|
||||
type="cylinder"
|
||||
size="0.015 0.33"
|
||||
quat="1.0 0.0 0.0 0.0"
|
||||
rgba="0.75 0.72 0.55 1.0" />
|
||||
|
||||
|
||||
<!--===================================================================其他障碍=====================================================================================-->
|
||||
|
||||
<!-- 5cm台阶 -->
|
||||
<geom pos="1.0 2.0 0.025" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
|
||||
<geom pos="1.3 2.0 0.075" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
|
||||
<geom pos="1.6 2.0 0.125" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
|
||||
<geom pos="1.9 2.0 0.175" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
|
||||
<geom pos="2.2 2.0 0.225" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
|
||||
<geom pos="2.5 2.0 0.275" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
|
||||
<geom pos="2.8 2.0 0.325" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
|
||||
<geom pos="3.45 2.0 0.375" type="box" size="0.5 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
|
||||
<geom pos="4.1 2.0 0.325" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
|
||||
<geom pos="4.4 2.0 0.275" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
|
||||
<geom pos="4.7 2.0 0.225" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
|
||||
<geom pos="5.0 2.0 0.175" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
|
||||
<geom pos="5.3 2.0 0.125" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
|
||||
<geom pos="5.6 2.0 0.075" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
|
||||
<geom pos="5.9 2.0 0.025" type="box" size="0.15 1.0 0.025" quat="1.0 0.0 0.0 0.0" />
|
||||
|
||||
<!-- 斜坡 -->
|
||||
<geom pos="2.0 4.0 0.1" type="box" size="1.5 0.75 0.005" quat="0.9950041652780258 0.0 -0.09983341664682815 0.0" />
|
||||
|
||||
|
||||
|
||||
|
||||
<geom pos="1.4 6.0 0.165" type="box" size="0.1 0.75 0.0049999999999999975" quat="1.0 0.0 0.0 0.0"/>
|
||||
<geom pos="1.6 6.0 0.275" type="box" size="0.1 0.75 0.0049999999999999975" quat="1.0 0.0 0.0 0.0"/>
|
||||
<geom pos="1.8 6.0 0.385" type="box" size="0.1 0.75 0.0049999999999999975" quat="1.0 0.0 0.0 0.0"/>
|
||||
<geom pos="2.0 6.0 0.495" type="box" size="0.1 0.75 0.0049999999999999975" quat="1.0 0.0 0.0 0.0"/>
|
||||
<geom pos="2.2 6.0 0.605" type="box" size="0.1 0.75 0.0049999999999999975" quat="1.0 0.0 0.0 0.0"/>
|
||||
<geom pos="2.4 6.0 0.715" type="box" size="0.1 0.75 0.0049999999999999975" quat="1.0 0.0 0.0 0.0"/>
|
||||
<geom pos="2.5999999999999996 6.0 0.825" type="box" size="0.1 0.75 0.0049999999999999975" quat="1.0 0.0 0.0 0.0"/>
|
||||
<geom pos="2.8 6.0 0.9349999999999999" type="box" size="0.1 0.75 0.0049999999999999975" quat="1.0 0.0 0.0 0.0"/>
|
||||
<geom pos="3.0 6.0 1.045" type="box" size="0.1 0.75 0.0049999999999999975" quat="1.0 0.0 0.0 0.0"/>
|
||||
<geom pos="-2.3179973398407565 5.173660321080885 -0.25" type="box" size="0.2568216778785459 0.2608020098770089 0.2619541072037832" quat="0.9930658271270357 -0.05910360995856133 -0.06731065544310916 0.0761334482746007"/>
|
||||
<geom pos="-2.3179973398407565 5.3620612607436735 -0.25" type="box" size="0.23778028947050467 0.26051225137569556 0.27137457425982286" quat="0.9959802827127009 0.07435995796871331 0.0004862731479656538 -0.04993632582444846"/>
|
||||
<geom pos="-2.3179973398407565 5.545897059602109 -0.25" type="box" size="0.23841994611378936 0.2717839884518381 0.22585827399286504" quat="0.9983504429181294 -0.004811901627841528 0.03471877670768104 -0.045473566737405234"/>
|
||||
<geom pos="-2.3179973398407565 5.7436240471772795 -0.25" type="box" size="0.2552179019048769 0.2548992578792955 0.22547735976326444" quat="0.9968270877924189 -0.029673908987198697 0.06777847718526858 -0.029347814214192768"/>
|
||||
<geom pos="-2.3179973398407565 5.940214647011584 -0.25" type="box" size="0.24313116620329878 0.2372064979204117 0.26079933745117434" quat="0.9952106364954509 0.05088987714749159 -0.07605843245051555 -0.03436748846516202"/>
|
||||
<geom pos="-2.3179973398407565 6.165430585901471 -0.25" type="box" size="0.24786042386990592 0.2322559052231109 0.2644037606269708" quat="0.9936075351397807 -0.05017393314139304 0.06986162224674641 -0.07311632022760194"/>
|
||||
<geom pos="-2.3179973398407565 6.315657865031069 -0.25" type="box" size="0.23704265198840277 0.24982080672772003 0.2530694373586838" quat="0.9981716459301547 0.036179123385437884 0.04523497974247521 0.017269420947002005"/>
|
||||
<geom pos="-2.3179973398407565 6.489372835072359 -0.25" type="box" size="0.2647428927965494 0.2716292502682415 0.23725049444938928" quat="0.9954041516289313 0.019099981466976248 0.07663366510923347 0.05415761257454444"/>
|
||||
<geom pos="-2.094119617957536 5.194943631570213 -0.25" type="box" size="0.23176840693038148 0.23782936054799508 0.2282032657053922" quat="0.9973366591969123 0.030692664684553082 -0.030015392527106853 0.058962910104125923"/>
|
||||
<geom pos="-2.094119617957536 5.441090326561234 -0.25" type="box" size="0.2602142310926322 0.27213502289176367 0.2574009440402366" quat="0.9948915480473169 0.023650461407535205 0.08508706956405496 -0.04890453856469255"/>
|
||||
<geom pos="-2.094119617957536 5.642958951230403 -0.25" type="box" size="0.24800055056479955 0.24050676557282252 0.23522489807277194" quat="0.9912055235117067 0.06188914817453858 -0.09720856678264424 -0.0650525790586179"/>
|
||||
<geom pos="-2.094119617957536 5.884810142659838 -0.25" type="box" size="0.24637516954806898 0.24364583504893206 0.2682443460752295" quat="0.9964495083946477 -0.07467561549070643 -0.038754439691442995 0.003165924091257204"/>
|
||||
<geom pos="-2.094119617957536 6.129893093768145 -0.25" type="box" size="0.2378958269703999 0.25045408022075055 0.24760364656411665" quat="0.9946240462872651 -0.03963038800955788 -0.08917648605658504 0.03464091840526007"/>
|
||||
<geom pos="-2.094119617957536 6.31135792935651 -0.25" type="box" size="0.2612154064649683 0.2252849660353503 0.26933214617336004" quat="0.997742809172301 0.019590451180868298 0.06391648942914147 -0.006339033565912353"/>
|
||||
<geom pos="-2.094119617957536 6.4905285510291915 -0.25" type="box" size="0.22768983323309258 0.23025468157122184 0.2748062344121069" quat="0.9940843963100783 0.05105793020576152 0.009164090402224958 -0.0954218016127878"/>
|
||||
<geom pos="-2.094119617957536 6.650993287072798 -0.25" type="box" size="0.25443787917913946 0.24575387105878618 0.22934424641420587" quat="0.9953832160779188 -0.0075142647566912866 0.08305501342257929 0.04751477371219154"/>
|
||||
<geom pos="-1.8951777534978893 5.178742263102535 -0.25" type="box" size="0.256127371544264 0.24353861032614957 0.2714578153598507" quat="0.996547165679354 -0.008558846106514282 0.009384988866828401 0.08205129318749771"/>
|
||||
<geom pos="-1.8951777534978893 5.353646654030374 -0.25" type="box" size="0.23858897994497189 0.2426691342857623 0.26961659613969635" quat="0.997689480058865 0.014145078555906384 0.009115651940537827 -0.06582190381793727"/>
|
||||
<geom pos="-1.8951777534978893 5.575640096199311 -0.25" type="box" size="0.24430646334884096 0.2676093509240798 0.23693271670520474" quat="0.9976019172675967 -0.06794745094616256 0.006187887696065785 -0.011630503849579829"/>
|
||||
<geom pos="-1.8951777534978893 5.728196634759373 -0.25" type="box" size="0.25668697148716296 0.2743986770827004 0.23696403861156407" quat="0.9946311966964959 -0.0774649076913635 0.05866511796477466 0.03558615698054331"/>
|
||||
<geom pos="-1.8951777534978893 5.954288148947323 -0.25" type="box" size="0.2748179458028994 0.25407956175122554 0.25142548243710355" quat="0.9927720663552536 0.03885622734718657 0.06180632563910765 0.09525647469886914"/>
|
||||
<geom pos="-1.8951777534978893 6.198076908897641 -0.25" type="box" size="0.23522918501552104 0.2714895927340788 0.23659922178360912" quat="0.9944408196408429 0.027722532803612636 -0.06883505448003607 -0.07470376618171105"/>
|
||||
<geom pos="-1.8951777534978893 6.356335079330304 -0.25" type="box" size="0.253243646653926 0.26648639763264487 0.22751627090926196" quat="0.9924345646356405 0.06351743627133029 0.09661714031088077 -0.041283149154730255"/>
|
||||
<geom pos="-1.8951777534978893 6.6037178705715744 -0.25" type="box" size="0.2664884535468762 0.26472237049442093 0.2545826559482188" quat="0.996208812433608 0.04792667434016892 0.03088022810200046 -0.06570728596343135"/>
|
||||
<geom pos="-1.7450143557797366 5.181798049538029 -0.25" type="box" size="0.22764974830145798 0.2314500042225232 0.26635647118774647" quat="0.996420562587967 -0.02530598413067966 -0.01712092616039989 -0.07881968983996092"/>
|
||||
<geom pos="-1.7450143557797366 5.396539066742657 -0.25" type="box" size="0.24515338305934362 0.25502436912192245 0.23509532716059323" quat="0.9974817825235419 0.033698559354477776 -0.06057332552735373 -0.01501242370995589"/>
|
||||
<geom pos="-1.7450143557797366 5.5477605493550115 -0.25" type="box" size="0.2368415768982088 0.2653984778068547 0.25193186806340717" quat="0.9933816957092091 -0.04426535864216467 0.0892264769079807 0.057201577537394625"/>
|
||||
<geom pos="-1.7450143557797366 5.764238738853998 -0.25" type="box" size="0.257475541143157 0.25587521442146555 0.2684267956125346" quat="0.9955849335110497 0.004076118120639651 -0.09299366924131539 0.012091439447074191"/>
|
||||
<geom pos="-1.7450143557797366 5.942956213887727 -0.25" type="box" size="0.23647345784635188 0.22779489919605103 0.2690566454457882" quat="0.9997372307105926 0.020286738832439623 -0.010113822112874779 0.003410038259163467"/>
|
||||
<geom pos="-1.7450143557797366 6.162981335139796 -0.25" type="box" size="0.2336152227884161 0.23785626414299832 0.26272786991330355" quat="0.9954569549676405 0.08638675499005252 -0.03150236191317218 -0.024705881136540285"/>
|
||||
<geom pos="-1.7450143557797366 6.344025407207907 -0.25" type="box" size="0.2364006101773331 0.23674709170116234 0.2660167000427503" quat="0.9941020005048972 -0.09350614404282417 -0.05493059192638792 0.0006660998579442658"/>
|
||||
<geom pos="-1.7450143557797366 6.5791872438688825 -0.25" type="box" size="0.259653062502181 0.26359758888480966 0.27170867851854713" quat="0.9955751902075182 0.06572529509454274 0.04059253213564211 -0.05350207998588654"/>
|
||||
<geom pos="-1.4950537649717406 5.196975242498044 -0.25" type="box" size="0.24902209995429173 0.24604186796843594 0.26555264385759036" quat="0.9981753741369879 0.027059282981497578 0.02078151255996191 0.049818133312398136"/>
|
||||
<geom pos="-1.4950537649717406 5.415093649392617 -0.25" type="box" size="0.22561317519118573 0.23246623591498758 0.2516053992906602" quat="0.9961431033663982 -0.019679813255757937 0.07777937635153075 0.035524515199325105"/>
|
||||
<geom pos="-1.4950537649717406 5.596072881787104 -0.25" type="box" size="0.2545347271371952 0.2527292932516396 0.272364707011277" quat="0.9987990019474599 -0.0468132050914785 0.002764373695484487 -0.01419280718860587"/>
|
||||
<geom pos="-1.4950537649717406 5.788444207457658 -0.25" type="box" size="0.257823578756761 0.22815201323013437 0.2506904868770564" quat="0.9907800223935689 0.08207700037219429 -0.06679273748567588 -0.08459931119619979"/>
|
||||
<geom pos="-1.4950537649717406 5.962479724512656 -0.25" type="box" size="0.23447002037921025 0.260091883647859 0.2613547781123637" quat="0.994185747243896 -0.045779429335849345 -0.09537527521663794 -0.020062420196245392"/>
|
||||
<geom pos="-1.4950537649717406 6.137620949716146 -0.25" type="box" size="0.25502378303486756 0.24137626830945555 0.26521755821448284" quat="0.9954511067004865 0.04827241542514884 0.023378972724717166 -0.07874193109224191"/>
|
||||
<geom pos="-1.4950537649717406 6.345742058941403 -0.25" type="box" size="0.23763376291279575 0.27259418014745557 0.24184880568666417" quat="0.9979730938550316 -0.0521706328108408 -0.03256129013590577 0.01636127740178204"/>
|
||||
<geom pos="-1.4950537649717406 6.548966807280673 -0.25" type="box" size="0.23004532593424165 0.24736965888987583 0.22917624237245732" quat="0.993938322382385 -0.041286377558295506 -0.09233464220936707 0.04308549844056845"/>
|
||||
<geom pos="-1.2872521554407157 5.194072124237239 -0.25" type="box" size="0.2725857334366114 0.23611730648841156 0.25109418723334265" quat="0.9870880287314421 -0.09538067591152324 -0.0888119213223197 -0.09312460914696315"/>
|
||||
<geom pos="-1.2872521554407157 5.418639818976418 -0.25" type="box" size="0.2326397456179607 0.2609646699674687 0.2717115772948157" quat="0.9928206111485406 -0.07383059897166898 -0.08473752136579614 0.040936892980589765"/>
|
||||
<geom pos="-1.2872521554407157 5.655974569843163 -0.25" type="box" size="0.22633109461398734 0.25911291311168594 0.23532484499883452" quat="0.9947769304653071 -0.023834507399340135 0.09176246279588879 -0.037820963666801724"/>
|
||||
<geom pos="-1.2872521554407157 5.8981628648303595 -0.25" type="box" size="0.27278934520816167 0.2559269445001904 0.26076472835929454" quat="0.9936318869346414 0.0561019165840355 0.09234219275279033 0.0319557140415977"/>
|
||||
<geom pos="-1.2872521554407157 6.064939027275534 -0.25" type="box" size="0.2481756437845515 0.2613413397088905 0.24788858471207542" quat="0.9992310163557966 0.006919915638721734 0.03652449975321466 -0.012468024618680441"/>
|
||||
<geom pos="-1.2872521554407157 6.288670079370122 -0.25" type="box" size="0.2574127585385611 0.27445220033632356 0.22507618952620437" quat="0.9984337612791753 0.04947883993440395 -0.0010074967830211658 -0.026093173185657285"/>
|
||||
<geom pos="-1.2872521554407157 6.4987760234638605 -0.25" type="box" size="0.2442462188069663 0.2639082925274208 0.24918893213917415" quat="0.9910480688126146 0.08639971284060838 -0.0903343762232188 -0.04688832899783648"/>
|
||||
<geom pos="-1.2872521554407157 6.675611985491267 -0.25" type="box" size="0.2507168814441121 0.26699708557208374 0.26588306060638556" quat="0.9968882348111167 0.009466759168649483 -0.06920917552433178 0.03652831489763918"/>
|
||||
<geom pos="-1.0678149575697586 5.238251070535694 -0.25" type="box" size="0.2331754586513582 0.22873009754409884 0.2593258638743009" quat="0.9962563091989393 0.014506022500180732 0.069646748226122 0.049114887295595266"/>
|
||||
<geom pos="-1.0678149575697586 5.471732496077581 -0.25" type="box" size="0.2653495693182789 0.26581370557074685 0.2509273010188512" quat="0.9948019811445705 -0.05898118120008214 0.04034485064897968 -0.07254330845220838"/>
|
||||
<geom pos="-1.0678149575697586 5.691153717662407 -0.25" type="box" size="0.26309077142424103 0.26536949948987093 0.26566703066149744" quat="0.998331608401392 -0.025703385975953945 -0.05139460706256987 0.005650661991725779"/>
|
||||
<geom pos="-1.0678149575697586 5.938910564562482 -0.25" type="box" size="0.23082606038737274 0.23890770539441533 0.25941695887199245" quat="0.99040745217393 -0.09494180013051996 0.05600240658814475 -0.08332384846283197"/>
|
||||
<geom pos="-1.0678149575697586 6.13746533389903 -0.25" type="box" size="0.2676431215498913 0.2569308659994288 0.24597694927356487" quat="0.9955894165044442 0.04511524606234168 -0.003741035154326279 -0.08217258042101841"/>
|
||||
<geom pos="-1.0678149575697586 6.328213254628996 -0.25" type="box" size="0.24188540919954635 0.25556145122229207 0.2605619987765001" quat="0.9929731934387934 -0.06598945456828725 -0.06809928840246085 0.07079629875087447"/>
|
||||
<geom pos="-1.0678149575697586 6.554891603524267 -0.25" type="box" size="0.23121916641107804 0.25266417867731916 0.25489063309084503" quat="0.9984565221200774 -0.02714563973423854 0.029423903667849492 0.03849573446816443"/>
|
||||
<geom pos="-1.0678149575697586 6.718482966003368 -0.25" type="box" size="0.27418971733780156 0.2623437838864593 0.23694037285314332" quat="0.9968347573500278 0.015128289848683302 0.06990516881209438 0.03471121949050888"/>
|
||||
<geom pos="-0.8851136474992356 5.233272057743225 -0.25" type="box" size="0.2587161160845902 0.2542459313914242 0.25268742624288776" quat="0.9956323623605605 0.07455213252171443 -0.05585558926184348 0.006191260373025539"/>
|
||||
<geom pos="-0.8851136474992356 5.410759300563648 -0.25" type="box" size="0.24019625793101548 0.2509280955260936 0.26698317271101046" quat="0.9960056830073374 0.025487333399660517 0.08516393482003458 -0.008377318140719903"/>
|
||||
<geom pos="-0.8851136474992356 5.622546965631826 -0.25" type="box" size="0.22527085609750916 0.22924847380626232 0.23073331588883172" quat="0.9982201065088221 0.016292210993739946 0.013620740418279952 0.05572843307423328"/>
|
||||
<geom pos="-0.8851136474992356 5.856118124892678 -0.25" type="box" size="0.26833255247166926 0.2512767990265972 0.2502231376336179" quat="0.9914910532543448 0.0657880493114159 0.05073139628229043 -0.10021850784978792"/>
|
||||
<geom pos="-0.8851136474992356 6.010914595891683 -0.25" type="box" size="0.24926441451407527 0.22868800964152894 0.26501630221174116" quat="0.990736605756697 -0.0765632176593298 -0.09876375004658954 0.05314859727296924"/>
|
||||
<geom pos="-0.8851136474992356 6.1657945570092965 -0.25" type="box" size="0.26823427325895977 0.263134568634285 0.23692064318485426" quat="0.9903439387899988 0.07219524819688558 0.09310840228148876 0.07305856872598579"/>
|
||||
<geom pos="-0.8851136474992356 6.383066094410735 -0.25" type="box" size="0.23772929143507357 0.2619329708548053 0.23258884134606547" quat="0.9928463525385446 0.07348810370848206 0.09252710448645944 -0.017156742103069993"/>
|
||||
<geom pos="-0.8851136474992356 6.551482176174038 -0.25" type="box" size="0.25127741834221506 0.25307864976328337 0.234931895271364" quat="0.994062479132312 -0.084918495271965 -0.05565932062022146 -0.03912386445844099"/>
|
||||
<geom pos="-0.6441452943469552 5.193190826541208 -0.25" type="box" size="0.27016649011158844 0.23778885128164629 0.25859862297032477" quat="0.9957774962333289 0.08893738517887846 -0.021654911548168992 0.006955883744530942"/>
|
||||
<geom pos="-0.6441452943469552 5.407266358883468 -0.25" type="box" size="0.24982888074520473 0.26972381569065607 0.2275629646713632" quat="0.9942571418988088 0.02444442288698024 0.050336521935401994 -0.09122193010664256"/>
|
||||
<geom pos="-0.6441452943469552 5.56363948710079 -0.25" type="box" size="0.24725848381417467 0.2326432801330426 0.2476341019968084" quat="0.9935769557551848 -0.05688836927692267 -0.06858315894908604 0.0697488117593134"/>
|
||||
<geom pos="-0.6441452943469552 5.716461923308758 -0.25" type="box" size="0.226746163574222 0.25188961955216527 0.24650452053954758" quat="0.9981029078160676 0.03471313574303335 -0.0493522060742233 -0.012245136651074443"/>
|
||||
<geom pos="-0.6441452943469552 5.896430073001983 -0.25" type="box" size="0.25284949298017617 0.23421620066432108 0.2621382648463894" quat="0.9996413412447853 0.009523771195346772 -0.02374945051322109 -0.007902547492127014"/>
|
||||
<geom pos="-0.6441452943469552 6.110845613879558 -0.25" type="box" size="0.23421403581041014 0.2504556848552827 0.24096555776925194" quat="0.9982371379292613 0.05906199571247059 0.004184466367765714 0.004097238396071534"/>
|
||||
<geom pos="-0.6441452943469552 6.301410327752269 -0.25" type="box" size="0.2296817145699627 0.2701766916372966 0.22803599234835198" quat="0.9993879202059675 0.00903361239651759 0.0071600522792900426 -0.03302896372607247"/>
|
||||
<geom pos="-0.6441452943469552 6.470340592771366 -0.25" type="box" size="0.2742166750011587 0.23990439594655139 0.2609404931878803" quat="0.9954980192986956 -0.06477558432537366 -0.03478863652617152 0.059812774691783824"/>
|
||||
<geom pos="-0.43033316974190594 5.176946043237525 -0.25" type="box" size="0.2511344405622137 0.26383180387822514 0.2729516287367074" quat="0.9939152287364423 -0.054017256979398534 -0.08418383902399898 -0.0461273810376057"/>
|
||||
<geom pos="-0.43033316974190594 5.409693792302493 -0.25" type="box" size="0.23878438774624536 0.22858057493504688 0.24808981791323623" quat="0.9942721168135948 -0.06322741578652484 -0.032958930763236916 -0.0796175891553822"/>
|
||||
<geom pos="-0.43033316974190594 5.573593446642387 -0.25" type="box" size="0.2647545698926724 0.25215736713350106 0.25522885339490087" quat="0.9971918839448468 -0.004732654353889047 0.07222000940536273 0.019241071144380294"/>
|
||||
<geom pos="-0.43033316974190594 5.819531161619798 -0.25" type="box" size="0.2538318010250847 0.23661725487197 0.26323696729639623" quat="0.9929867679849618 0.07588268800385974 -0.046355180482866284 -0.07791208834634974"/>
|
||||
<geom pos="-0.43033316974190594 6.047888701789067 -0.25" type="box" size="0.25684720236482583 0.25568474221031684 0.24397094296107352" quat="0.9925256546500537 -0.08214715587137877 0.06615741576234473 0.06138294537877935"/>
|
||||
<geom pos="-0.43033316974190594 6.283676056613722 -0.25" type="box" size="0.2432692787269645 0.2742601402399693 0.2689467974776609" quat="0.9989897410981501 0.0047037010179340685 0.0014590493232283363 -0.04466814919444852"/>
|
||||
<geom pos="-0.43033316974190594 6.458239009480634 -0.25" type="box" size="0.2474330649412728 0.2551435998627988 0.23002773988805483" quat="0.9982311862818315 0.0506459597703481 0.025989046975463247 -0.01714802993392417"/>
|
||||
<geom pos="-0.43033316974190594 6.703309377175586 -0.25" type="box" size="0.2301029552870502 0.2556223065028505 0.2349557527965884" quat="0.9948672058648923 -0.034182171957127715 -0.0030330068097560517 -0.09519255582537542"/>
|
||||
<geom type="hfield" hfield="perlin_hfield" pos="-1.5 4.0 0.0" quat="1.0 0.0 0.0 0.0"/>
|
||||
<geom type="hfield" hfield="image_hfield" pos="-1.5 2.0 0.0" quat="0.7073882691671998 0.0 0.0 -0.706825181105366"/>
|
||||
</worldbody>
|
||||
</mujoco>
|
||||
@@ -26,28 +26,28 @@
|
||||
|
||||
<worldbody>
|
||||
<body name="base_link">
|
||||
<inertial pos="0.1517 0.0002 0.0542" mass="3.5" diaginertia="0.0215 0.0904 0.0985"/>
|
||||
<inertial pos="0.1517 0.0002 0.0542" mass="6.5377" fullinertia="0.0402 0.1689 0.1840 -0.0001 0.0000 0.0000"/>
|
||||
<joint type="free"/>
|
||||
<geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="base_link"/>
|
||||
<geom size="0.178 0.1175 0.073" pos="0.1518 0 0.054" type="box" rgba="0.75294 0.75294 0.75294 1"/>
|
||||
<body name="fl_hip_abduction_Link" pos="0.32826 0.066172 0.053981">
|
||||
<inertial pos="0.0488 -0.0026 0.0007" mass="0.5" diaginertia="0.0003 0.0006 0.0005"/>
|
||||
<inertial pos="0.0488 -0.0026 0.0007" mass="0.63" fullinertia="0.0004 0.0007 0.0006 0.0000 0.0000 0.0000"/>
|
||||
<joint name="fl_hip_abduction_joint" pos="0 0 0" axis="1 0 0" range="-0.436 0.611" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
|
||||
<geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="fl_hip_abduction_Link"/>
|
||||
<body name="fl_hip_pitch_Link" pos="0.06389 -0.027344 0.00010727" quat="0.999997 -0.0025023 0 0">
|
||||
<inertial pos="0.0019 0.1119 -0.048" mass="0.935" diaginertia="0.0062 0.0064 0.001"/>
|
||||
<inertial pos="0.0019 0.1119 -0.0480" mass="0.998" fullinertia="0.0066 0.0068 0.0011 0.0000 0.0000 0.0004"/>
|
||||
<joint name="fl_hip_pitch_joint" pos="0 0 0" axis="0 1 0" range="-2.58 2.58" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
|
||||
<geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="fl_hip_pitch_Link"/>
|
||||
<geom size="0.046 0.048" pos="0 0.048 0" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
|
||||
<geom size="0.0435 0.0115 0.06" pos="0 0.1155 -0.06" type="box" rgba="0.75294 0.75294 0.75294 1"/>
|
||||
<body name="fl_knee_Link" pos="0 0.1035 -0.25" quat="0.999997 0.0025023 0 0">
|
||||
<inertial pos="0.0002 0.0242 -0.1539" mass="0.651" fullinertia="0.0042 0.0045 0.0005 0 0 0.0002"/>
|
||||
<inertial pos="0.0002 0.0242 -0.1539" mass="0.6965" fullinertia="0.0045 0.0048 0.0006 0.0000 0.0000 0.0002"/>
|
||||
<joint name="fl_knee_joint" pos="0 0 0" axis="0 1 0" range="-2.65 2.65" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
|
||||
<geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="fl_knee_Link"/>
|
||||
<geom size="0.0475 0.015" pos="0 0.025 -0.20011" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
|
||||
<geom size="0.015 0.0125 0.06" pos="0 0.0125 -0.09" type="box" rgba="0.75294 0.75294 0.75294 1"/>
|
||||
<body name="fl_wheel_Link" pos="0 0.014699 -0.20011">
|
||||
<inertial pos="-0.0002 0.0407 -0.0001" mass="0.53" diaginertia="0.0017 0.0032 0.0017"/>
|
||||
<inertial pos="-0.0002 0.0407 -0.0001" mass="0.5505" fullinertia="0.0017 0.0034 0.0017 0.0000 0.0000 0.0000"/>
|
||||
<joint name="fl_wheel_joint" pos="0 0 0" axis="0 1 0" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
|
||||
<geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="fl_wheel_Link"/>
|
||||
<geom size="0.1 0.015" pos="0 0.04074 0" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
|
||||
@@ -56,23 +56,23 @@
|
||||
</body>
|
||||
</body>
|
||||
<body name="fr_hip_abduction_Link" pos="0.32826 -0.065853 0.054034">
|
||||
<inertial pos="0.0488 0.0026 0.0008" mass="0.5" diaginertia="0.0003 0.0006 0.0005"/>
|
||||
<inertial pos="0.0488 0.0026 0.0008" mass="0.63" fullinertia="0.0004 0.0007 0.0006 0.0000 0.0000 0.0000"/>
|
||||
<joint name="fr_hip_abduction_joint" pos="0 0 0" axis="1 0 0" range="-0.611 0.436" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
|
||||
<geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="fr_hip_abduction_Link"/>
|
||||
<body name="fr_hip_pitch_Link" pos="0.06389 0.027311 -0.00036027" quat="0.999976 -0.00686995 0 0">
|
||||
<inertial pos="-0.0019 -0.1119 -0.048" mass="0.935" diaginertia="0.0062 0.0064 0.001"/>
|
||||
<inertial pos="-0.0019 -0.1119 -0.048" mass="0.998" fullinertia="0.0066 0.0068 0.0011 0.0000 0.0000 0.0004"/>
|
||||
<joint name="fr_hip_pitch_joint" pos="0 0 0" axis="0 1 0" range="-2.58 2.58" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
|
||||
<geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="fr_hip_pitch_Link"/>
|
||||
<geom size="0.046 0.048" pos="0 -0.048 0" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
|
||||
<geom size="0.0435 0.0115 0.06" pos="0 -0.1155 -0.06" type="box" rgba="0.75294 0.75294 0.75294 1"/>
|
||||
<body name="fr_knee_Link" pos="-0.00075079 -0.1035 -0.25" quat="0.999976 0.00686995 0 0">
|
||||
<inertial pos="-0.0002 -0.0242 -0.1539" mass="0.651" fullinertia="0.0042 0.0045 0.0005 0 0 0.0001"/>
|
||||
<inertial pos="-0.0002 -0.0242 -0.1539" mass="0.6965" fullinertia="0.0045 0.0048 0.0006 0.0000 0.0000 0.0002"/>
|
||||
<joint name="fr_knee_joint" pos="0 0 0" axis="0 1 0" range="-2.65 2.65" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
|
||||
<geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="fr_knee_Link"/>
|
||||
<geom size="0.0475 0.015" pos="0 -0.025 -0.1998" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
|
||||
<geom size="0.015 0.0125 0.06" pos="0 -0.0125 -0.09" type="box" rgba="0.75294 0.75294 0.75294 1"/>
|
||||
<body name="fr_wheel_Link" pos="0 -0.018447 -0.1998">
|
||||
<inertial pos="0.0002 -0.0407 -0.0001" mass="0.53" diaginertia="0.0017 0.0032 0.0017"/>
|
||||
<inertial pos="0.0002 -0.0407 -0.0001" mass="0.5505" fullinertia="0.0017 0.0034 0.0017 0.0000 0.0000 0.0000"/>
|
||||
<joint name="fr_wheel_joint" pos="0 0 0" axis="0 1 0" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
|
||||
<geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="fr_wheel_Link"/>
|
||||
<geom size="0.1 0.015" pos="0 -0.040735 0" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
|
||||
@@ -81,23 +81,23 @@
|
||||
</body>
|
||||
</body>
|
||||
<body name="rl_hip_abduction_Link" pos="-0.024743 0.066141 0.054034">
|
||||
<inertial pos="-0.0488 -0.0026 -0.0008" mass="0.5" diaginertia="0.0003 0.0006 0.0005"/>
|
||||
<inertial pos="-0.0488 -0.0026 -0.0007" mass="0.63" fullinertia="0.0004 0.0007 0.0006 0.0000 0.0000 0.0000"/>
|
||||
<joint name="rl_hip_abduction_joint" pos="0 0 0" axis="1 0 0" range="-0.436 0.611" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
|
||||
<geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="rl_hip_abduction_Link"/>
|
||||
<body name="rl_hip_pitch_Link" pos="-0.06389 -0.027309 0.00045509">
|
||||
<inertial pos="0.0019 0.1119 -0.048" mass="0.935" diaginertia="0.0062 0.0064 0.001"/>
|
||||
<inertial pos="0.0019 0.1119 -0.048" mass="0.998" fullinertia="0.0066 0.0068 0.0011 0.0000 -0.0001 -0.0005"/>
|
||||
<joint name="rl_hip_pitch_joint" pos="0 0 0" axis="0 1 0" range="-2.58 2.58" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
|
||||
<geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="rl_hip_pitch_Link"/>
|
||||
<geom size="0.046 0.048" pos="0 0.048 0" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
|
||||
<geom size="0.0435 0.0115 0.06" pos="0 0.1155 -0.06" type="box" rgba="0.75294 0.75294 0.75294 1"/>
|
||||
<body name="rl_knee_Link" pos="0 0.099459 -0.25163">
|
||||
<inertial pos="0.0002 0.0242 -0.1539" mass="0.651" fullinertia="0.0042 0.0045 0.0005 0 0 -0.0003"/>
|
||||
<inertial pos="0.0002 0.0242 -0.1539" mass="0.6965" fullinertia="0.0045 0.0048 0.0006 0.0000 0.0000 -0.0003"/>
|
||||
<joint name="rl_knee_joint" pos="0 0 0" axis="0 1 0" range="-2.65 2.65" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
|
||||
<geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="rl_knee_Link"/>
|
||||
<geom size="0.0475 0.015" pos="0 0.025 -0.20027" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
|
||||
<geom size="0.015 0.0125 0.06" pos="0 0.0125 -0.09" type="box" rgba="0.75294 0.75294 0.75294 1"/>
|
||||
<body name="rl_wheel_Link" pos="0 0.012475 -0.20027">
|
||||
<inertial pos="-0.0002 0.0407 -0.0001" mass="0.53" diaginertia="0.0017 0.0032 0.0017"/>
|
||||
<inertial pos="-0.0002 0.0407 -0.0001" mass="0.5505" fullinertia="0.0017 0.0034 0.0017 0.0000 0.0000 0.0000"/>
|
||||
<joint name="rl_wheel_joint" pos="0 0 0" axis="0 1 0" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
|
||||
<geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="rl_wheel_Link"/>
|
||||
<geom size="0.1 0.015" pos="0 0.040737 0" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
|
||||
@@ -106,23 +106,23 @@
|
||||
</body>
|
||||
</body>
|
||||
<body name="rr_hip_abduction_Link" pos="-0.024743 -0.065884 0.053981">
|
||||
<inertial pos="-0.0488 0.0026 0.0008" mass="0.5" diaginertia="0.0003 0.0006 0.0005"/>
|
||||
<inertial pos="-0.0488 0.0026 0.0008" mass="0.63" fullinertia="0.0004 0.0007 0.0006 0.0000 0.0000 0.0000"/>
|
||||
<joint name="rr_hip_abduction_joint" pos="0 0 0" axis="1 0 0" range="-0.611 0.436" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
|
||||
<geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="rr_hip_abduction_Link"/>
|
||||
<body name="rr_hip_pitch_Link" pos="-0.06389 0.027341 0.00041625">
|
||||
<inertial pos="-0.002 -0.1111 -0.0498" mass="0.935" diaginertia="0.0062 0.0064 0.001"/>
|
||||
<inertial pos="-0.0020 -0.1111 -0.0498" mass="0.998" fullinertia="0.0066 0.0068 0.0011 0.0000 -0.0001 -0.0003"/>
|
||||
<joint name="rr_hip_pitch_joint" pos="0 0 0" axis="0 1 0" range="-2.58 2.58" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
|
||||
<geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="rr_hip_pitch_Link"/>
|
||||
<geom size="0.046 0.048" pos="0 -0.048 0" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
|
||||
<geom size="0.0435 0.0115 0.06" pos="0 -0.1155 -0.06" type="box" rgba="0.75294 0.75294 0.75294 1"/>
|
||||
<body name="rr_knee_Link" pos="-0.00075079 -0.099408 -0.25165">
|
||||
<inertial pos="-0.0002 -0.0225 -0.1541" mass="0.651" fullinertia="0.0042 0.0045 0.0005 0 0 -0.0001"/>
|
||||
<inertial pos="-0.0002 -0.0225 -0.1541" mass="0.6965" fullinertia="0.0045 0.0048 0.0006 0.0000 0.0000 -0.0002"/>
|
||||
<joint name="rr_knee_joint" pos="0 0 0" axis="0 1 0" range="-2.65 2.65" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
|
||||
<geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="rr_knee_Link"/>
|
||||
<geom size="0.0475 0.015" pos="0 -0.025 -0.20027" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
|
||||
<geom size="0.015 0.0125 0.06" pos="0 -0.0125 -0.09" type="box" rgba="0.75294 0.75294 0.75294 1"/>
|
||||
<body name="rr_wheel_Link" pos="0 -0.012435 -0.20027">
|
||||
<inertial pos="0.0002 -0.0407 -0.0005" mass="0.53" diaginertia="0.0017 0.0032 0.0017"/>
|
||||
<inertial pos="0.0002 -0.0407 -0.0005" mass="0.5505" fullinertia="0.0017 0.0034 0.0017 0.0000 0.0000 0.0000"/>
|
||||
<joint name="rr_wheel_joint" pos="0 0 0" axis="0 1 0" actuatorfrcrange="-17 17" damping="0.01" frictionloss="0.01" armature="0.0042"/>
|
||||
<geom type="mesh" contype="0" conaffinity="0" group="1" density="0" rgba="0.75294 0.75294 0.75294 1" mesh="rr_wheel_Link"/>
|
||||
<geom size="0.1 0.015" pos="0 -0.040737 0" quat="0.707105 0.707108 0 0" type="cylinder" rgba="0.75294 0.75294 0.75294 1"/>
|
||||
|
||||
@@ -28,9 +28,9 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
- name: Setup uv
|
||||
uses: astral-sh/setup-uv@v6
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
version: "0.9.27"
|
||||
@@ -45,9 +45,9 @@ jobs:
|
||||
matrix:
|
||||
python-version: ["3.10", "3.11", "3.12", "3.13"]
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- name: Setup uv
|
||||
uses: astral-sh/setup-uv@v6
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
enable-cache: true
|
||||
@@ -68,9 +68,9 @@ jobs:
|
||||
matrix:
|
||||
python-version: ["3.10", "3.11", "3.12", "3.13"]
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- name: Setup uv
|
||||
uses: astral-sh/setup-uv@v6
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
enable-cache: true
|
||||
@@ -84,9 +84,9 @@ jobs:
|
||||
matrix:
|
||||
python-version: ["3.10", "3.11", "3.12", "3.13"]
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- name: Setup uv
|
||||
uses: astral-sh/setup-uv@v6
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
enable-cache: true
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
name: Claude Code Review
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
types: [opened, synchronize, ready_for_review, reopened]
|
||||
# Optional: Only run on specific file changes
|
||||
# paths:
|
||||
# - "src/**/*.ts"
|
||||
# - "src/**/*.tsx"
|
||||
# - "src/**/*.js"
|
||||
# - "src/**/*.jsx"
|
||||
|
||||
jobs:
|
||||
claude-review:
|
||||
# Optional: Filter by PR author
|
||||
# if: |
|
||||
# github.event.pull_request.user.login == 'external-contributor' ||
|
||||
# github.event.pull_request.user.login == 'new-developer' ||
|
||||
# github.event.pull_request.author_association == 'FIRST_TIME_CONTRIBUTOR'
|
||||
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
issues: read
|
||||
id-token: write
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 1
|
||||
|
||||
- name: Run Claude Code Review
|
||||
id: claude-review
|
||||
uses: anthropics/claude-code-action@v1
|
||||
with:
|
||||
claude_code_oauth_token: ${{ secrets.CLAUDE_CODE_OAUTH_TOKEN }}
|
||||
plugin_marketplaces: 'https://github.com/anthropics/claude-code.git'
|
||||
plugins: 'code-review@claude-code-plugins'
|
||||
prompt: '/code-review:code-review ${{ github.repository }}/pull/${{ github.event.pull_request.number }}'
|
||||
# See https://github.com/anthropics/claude-code-action/blob/main/docs/usage.md
|
||||
# or https://code.claude.com/docs/en/cli-reference for available options
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
name: Claude Code
|
||||
|
||||
on:
|
||||
issue_comment:
|
||||
types: [created]
|
||||
pull_request_review_comment:
|
||||
types: [created]
|
||||
issues:
|
||||
types: [opened, assigned]
|
||||
pull_request_review:
|
||||
types: [submitted]
|
||||
|
||||
jobs:
|
||||
claude:
|
||||
if: |
|
||||
(github.event_name == 'issue_comment' && contains(github.event.comment.body, '@claude')) ||
|
||||
(github.event_name == 'pull_request_review_comment' && contains(github.event.comment.body, '@claude')) ||
|
||||
(github.event_name == 'pull_request_review' && contains(github.event.review.body, '@claude')) ||
|
||||
(github.event_name == 'issues' && (contains(github.event.issue.body, '@claude') || contains(github.event.issue.title, '@claude')))
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
issues: read
|
||||
id-token: write
|
||||
actions: read # Required for Claude to read CI results on PRs
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 1
|
||||
|
||||
- name: Run Claude Code
|
||||
id: claude
|
||||
uses: anthropics/claude-code-action@v1
|
||||
with:
|
||||
claude_code_oauth_token: ${{ secrets.CLAUDE_CODE_OAUTH_TOKEN }}
|
||||
|
||||
# This is an optional setting that allows Claude to read CI results on PRs
|
||||
additional_permissions: |
|
||||
actions: read
|
||||
|
||||
# Optional: Give a custom prompt to Claude. If this is not specified, Claude will perform the instructions specified in the comment that tagged it.
|
||||
# prompt: 'Update the pull request description to include a summary of changes.'
|
||||
|
||||
# Optional: Add claude_args to customize behavior and configuration
|
||||
# See https://github.com/anthropics/claude-code-action/blob/main/docs/usage.md
|
||||
# or https://code.claude.com/docs/en/cli-reference for available options
|
||||
# claude_args: '--allowed-tools Bash(gh pr *)'
|
||||
|
||||
@@ -17,7 +17,7 @@ jobs:
|
||||
build:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
|
||||
@@ -26,7 +26,7 @@ jobs:
|
||||
python-version: '3.13'
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v4
|
||||
uses: astral-sh/setup-uv@v7
|
||||
|
||||
- name: Build Sphinx Documentation
|
||||
run: uv run --group docs sphinx-multiversion docs docs/_build
|
||||
|
||||
@@ -15,9 +15,9 @@ jobs:
|
||||
contents: read
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v6
|
||||
uses: astral-sh/setup-uv@v7
|
||||
- name: Install Python 3.13
|
||||
run: uv python install 3.13
|
||||
- name: Build
|
||||
|
||||
@@ -33,9 +33,9 @@ keywords:
|
||||
- reinforcement-learning
|
||||
- robotics
|
||||
license: Apache-2.0
|
||||
commit: e2f33c6fb49caa26ec11f7b2de3c0c9aba71e9fd
|
||||
version: 1.3.0
|
||||
date-released: '2026-04-14'
|
||||
commit: 3cc461cd15e7155a8998b75ad767fae6dd448072
|
||||
version: 1.4.0
|
||||
date-released: '2026-05-26'
|
||||
preferred-citation:
|
||||
type: article
|
||||
title: >-
|
||||
|
||||
@@ -1,6 +1,10 @@
|
||||
.PHONY: sync
|
||||
sync:
|
||||
uv sync --all-extras --all-packages --group dev
|
||||
uv sync --all-packages --extra cu128 --group dev
|
||||
|
||||
.PHONY: sync-cpu
|
||||
sync-cpu:
|
||||
uv sync --all-packages --extra cpu --group dev
|
||||
|
||||
.PHONY: format
|
||||
format:
|
||||
|
||||
@@ -61,6 +61,7 @@ MuJoCo's integrator handles velocity-dependent forces.
|
||||
|
||||
**Built-in actuators** (``BuiltinPositionActuator``,
|
||||
``BuiltinVelocityActuator``, ``BuiltinMotorActuator``,
|
||||
``BuiltinPdActuator``, ``BuiltinDcMotorActuator``,
|
||||
``BuiltinMuscleActuator``) create native MuJoCo actuator elements in the
|
||||
MjSpec. The physics engine computes the control law and integrates
|
||||
velocity-dependent damping forces implicitly. This provides the best
|
||||
@@ -119,6 +120,31 @@ control.
|
||||
**BuiltinMotorActuator**: Creates ``<motor>`` actuators for direct torque
|
||||
control.
|
||||
|
||||
**BuiltinPdActuator**: Native PD that closes on both a position and a
|
||||
velocity target, implemented as paired ``<position>`` + ``<velocity>``
|
||||
actuators summing to ``kp * (p_target - q) + kd * (v_target - qdot)``.
|
||||
``BuiltinPositionActuator`` puts kd on the ``<position>`` element and
|
||||
implicitly assumes a zero velocity reference; use this when the policy
|
||||
emits a non-zero velocity target. Native delivery lets
|
||||
``implicit`` / ``implicitfast`` see the kd term in their velocity update,
|
||||
unlike ``IdealPdActuator`` which forwards Python-computed torque through
|
||||
an opaque ``<motor>``.
|
||||
|
||||
**BuiltinDcMotorActuator**: Wraps MuJoCo's native
|
||||
`<dcmotor> <https://mujoco.readthedocs.io/en/stable/XMLreference.html#actuator-dcmotor>`_
|
||||
element. Torque is ``tau = K * (V - K * omega) / R``; the back-EMF runs
|
||||
through the native bias path, so ``implicit`` / ``implicitfast`` pick up
|
||||
its velocity derivative as effective damping. Three input modes pick what
|
||||
``ctrl`` carries: VOLTAGE drives the motor directly; POSITION / VELOCITY
|
||||
close an internal PID (with anti-windup and slew limiting) against a
|
||||
single setpoint, whose Vmax-clamped output becomes torque. POSITION mode
|
||||
pins v_target = 0 (the kd term acts on raw velocity). Optional physics:
|
||||
inductance,
|
||||
thermal model with I^2R heating, cogging ripple, LuGre friction.
|
||||
``DcMotorActuator`` (the explicit version) is a software PD with a
|
||||
velocity-dependent torque clamp on top of a ``<motor>``; this is the real
|
||||
electrical model.
|
||||
|
||||
**BuiltinMuscleActuator**: Creates ``<muscle>`` actuators for
|
||||
biologically-inspired muscle dynamics with force-length-velocity
|
||||
characteristics.
|
||||
|
||||
@@ -18,6 +18,13 @@ mjlab.actuator
|
||||
- :class:`BuiltinPositionActuatorCfg`
|
||||
- :class:`BuiltinVelocityActuator`
|
||||
- :class:`BuiltinVelocityActuatorCfg`
|
||||
- :class:`BuiltinPdActuator`
|
||||
- :class:`BuiltinPdActuatorCfg`
|
||||
- :class:`BuiltinDcMotorActuator`
|
||||
- :class:`BuiltinDcMotorActuatorCfg`
|
||||
- :class:`DcMotorInputMode`
|
||||
- :class:`DcMotorDatasheetParams`
|
||||
- :class:`DcMotorPhysicalParams`
|
||||
- :class:`BuiltinMuscleActuator`
|
||||
- :class:`BuiltinMuscleActuatorCfg`
|
||||
- :class:`XmlActuator`
|
||||
@@ -84,6 +91,40 @@ Builtin Actuators
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: BuiltinPdActuator
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: BuiltinPdActuatorCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: BuiltinDcMotorActuator
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: BuiltinDcMotorActuatorCfg
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
.. autoclass:: DcMotorInputMode
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
.. autoclass:: DcMotorDatasheetParams
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
.. autoclass:: DcMotorPhysicalParams
|
||||
:members:
|
||||
:exclude-members: __init__
|
||||
:undoc-members:
|
||||
|
||||
|
||||
.. autoclass:: BuiltinMuscleActuator
|
||||
:members:
|
||||
:show-inheritance:
|
||||
|
||||
@@ -8,6 +8,89 @@ Upcoming version (not yet released)
|
||||
Added
|
||||
^^^^^
|
||||
|
||||
- Added ``BuiltinDcMotorActuator``, a native MuJoCo ``<dcmotor>`` wrapper.
|
||||
Supports voltage / position / velocity input modes with back-EMF,
|
||||
configurable motor constants, and optional integral, slew, inductance,
|
||||
thermal, LuGre, and cogging extensions.
|
||||
- Added ``scale_with_difficulty`` to ``HfRandomUniformTerrainCfg``. When
|
||||
enabled, the noise amplitude scales with difficulty (flat at 0, full
|
||||
``noise_range`` at 1) so the terrain progresses in a curriculum. Defaults to
|
||||
``False``, preserving the previous difficulty-independent behavior.
|
||||
|
||||
Changed
|
||||
^^^^^^^
|
||||
|
||||
- Bumped ``rsl-rl-lib`` from 5.2.0 to 5.4.0.
|
||||
- Curriculum-mode terrain difficulty is now deterministic across rows
|
||||
and reaches the configured ``difficulty_range`` endpoints
|
||||
(:issue:`1027`).
|
||||
- Heightfield terrains now color by absolute height with a diverging palette
|
||||
(cool below the ground plane, green at ground level, warm above) on a fixed
|
||||
scale, replacing the per-patch normalization. Color is now consistent across
|
||||
terrains, and low-amplitude terrain such as ``random_rough`` reads as gently
|
||||
tinted ground instead of high-contrast noise.
|
||||
- ``BoxNestedRingsTerrainCfg`` now builds uniform-height concentric ridges
|
||||
whose separating gaps widen with difficulty, replacing the random per-ring
|
||||
heights. Rings are colored by height (like the other terrains) and the outer
|
||||
border matches the ring height.
|
||||
- Terrain generation no longer prints timing information to stdout.
|
||||
|
||||
Fixed
|
||||
^^^^^
|
||||
|
||||
- Fixed ``select_gpus`` crashing when ``CUDA_VISIBLE_DEVICES`` contains MIG UUIDs instead of numeric indices.
|
||||
- Fixed pyramid-stairs terrains (``BoxPyramidStairsTerrainCfg``,
|
||||
``BoxInvertedPyramidStairsTerrainCfg``, and ``BoxOpenStairsTerrainCfg``)
|
||||
leaving an empty, geometry-free border at difficulty 0, where the step
|
||||
height collapses to zero. The flat border frame is now always generated as
|
||||
solid geometry flush with the ground (:issue:`1033`).
|
||||
- Fixed ``HfPerlinNoiseTerrainCfg`` failing to compile at difficulty 0, where
|
||||
the target height collapses to zero and MuJoCo rejects the non-positive
|
||||
heightfield size.
|
||||
- Fixed ``BoxRandomGridTerrainCfg`` producing NaN colors (and failing to build)
|
||||
at difficulty 0, where the grid height is zero and the color normalization
|
||||
divided by zero.
|
||||
- Fixed the center platform z-fighting with surrounding geometry in
|
||||
``BoxRandomGridTerrainCfg`` (grid cells were left underneath the platform) and
|
||||
``BoxRandomSpreadTerrainCfg`` (the platform duplicated the floor surface).
|
||||
- Fixed ``BoxNarrowBeamsTerrainCfg`` square platform corners protruding between
|
||||
the beams at high difficulty; the platform now shrinks to stay within the
|
||||
beams' angular coverage.
|
||||
- Fixed ``BoxSteppingStonesTerrainCfg`` reconfiguring abruptly at a difficulty
|
||||
threshold, where the stone grid re-tiled as its spacing crossed an integer
|
||||
boundary, and leaving an oversized gap around the center platform. The grid is
|
||||
now difficulty-independent and the platform snaps to it as a clean island.
|
||||
- Fixed ``train --video``, ``play``, and ``demo`` crashing with ``OpenGL
|
||||
platform library not loaded`` on headless Linux hosts that don't pre-set
|
||||
``MUJOCO_GL``. The default is now applied in ``mjlab/__init__.py`` (Linux
|
||||
only) so it takes effect before mujoco's GL backend selection runs.
|
||||
|
||||
Version 1.4.0 (May 26, 2026)
|
||||
----------------------------
|
||||
|
||||
Added
|
||||
^^^^^
|
||||
|
||||
- Added ``BuiltinPdActuator``, the implicit-integration version of
|
||||
``IdealPdActuator``. Same interface (position + velocity targets,
|
||||
kp/kd gains), but expresses the PD as native MuJoCo ``<position>``
|
||||
and ``<velocity>`` elements so the ``implicit`` / ``implicitfast``
|
||||
integrators include the kp/kd derivatives in their velocity update.
|
||||
The actuator stays stable at gain/timestep combinations where
|
||||
explicit Python PD would diverge, which matters when you want to
|
||||
run a real motor's stiff on-board PD gains in sim. ``effort_limit``
|
||||
is enforced as a sum-clamp on the two PD terms via
|
||||
``jnt_actfrcrange`` (or ``tendon_actfrcrange``). Supported by
|
||||
``dr.pd_gains`` and ``dr.effort_limits``.
|
||||
- Added ``mdp.projected_gravity_from_sensor``, an observation that derives
|
||||
projected gravity from a ``framezaxis`` up-vector sensor (negated) rather
|
||||
than from the root body orientation. Unlike ``mdp.projected_gravity``, it
|
||||
reflects the sensor's site frame, so it can observe IMU mounting domain
|
||||
randomization (e.g. via ``dr.site_quat``). Go1 and G1 ship an
|
||||
``imu_upvector`` sensor for this.
|
||||
- Added ``DebugVisualizer.add_box`` for drawing an axis-oriented box
|
||||
primitive, mirroring ``add_ellipsoid``. Supported by both the native
|
||||
and Viser viewers. ``size`` is the box half-extents (:issue:`992`).
|
||||
- Added ``--log-root`` CLI option to ``train``, ``play``, and ``evaluate``
|
||||
scripts for choosing where training logs are stored. Defaults to
|
||||
``logs/rsl_rl`` (unchanged behavior). Useful for directing outputs to a
|
||||
@@ -21,22 +104,41 @@ Added
|
||||
primary names in the order they appear along the per-contact axis of the
|
||||
output tensors. This makes it possible to map a contact-data column back
|
||||
to the primary it belongs to (:issue:`914`).
|
||||
- Added per-world mesh variant support via ``VariantEntityCfg`` and
|
||||
``VariantCfg``. Each world in a batched simulation can now use a
|
||||
different mesh asset for the same logical entity (e.g. world 0 holds a
|
||||
cube, world 1 a sphere), with weights controlling the proportion of
|
||||
worlds assigned to each variant. Mesh-derived constants (collision
|
||||
bounds, body inertials, subtree mass, inverse weights) are compiled
|
||||
per-variant and stored as per-world arrays in the Warp model, so domain
|
||||
randomization, the native viewer, the offscreen renderer, and the Viser
|
||||
viewer all pick up the variant assignment automatically. Variants must
|
||||
share the same kinematic structure (same bodies, joints, joint types);
|
||||
only mesh geoms may differ. Assignment is fixed at simulation init.
|
||||
See :ref:`per_world_mesh` for usage. With help from @XiangruiJiang.
|
||||
- Added per-world mesh variant support via ``VariantEntityCfg``. Each
|
||||
world in a batched simulation can now use a different mesh asset for
|
||||
the same logical entity (e.g. world 0 holds a cube, world 1 a
|
||||
sphere). Variants are passed as a ``dict[str, Callable]`` of named
|
||||
spec callables; the optional ``assignment`` field controls how worlds
|
||||
map to variants and accepts ``None`` (uniform), a ``dict[str, float]``
|
||||
of per-variant weights, or a custom ``Callable[[int], Sequence[int]]``.
|
||||
Mesh-derived constants (collision bounds, body inertials, subtree
|
||||
mass, inverse weights) are compiled per-variant and stored as
|
||||
per-world arrays in the Warp model, so domain randomization, the
|
||||
native viewer, the offscreen renderer, and the Viser viewer all pick
|
||||
up the variant assignment automatically. Variants must share the
|
||||
same kinematic structure (same bodies, joints, joint types); only
|
||||
mesh geoms may differ. Assignment is fixed at simulation init. See
|
||||
:ref:`heterogeneous_worlds` for usage. With help from @XiangruiJiang.
|
||||
- Per-world mesh variants now support per-variant materials and textures.
|
||||
Each variant can reference its own named material, which is automatically
|
||||
prefixed and scattered via ``geom_matid`` alongside the existing
|
||||
``geom_dataid`` table. Variants without a material get ``matid = -1``.
|
||||
Contribution by @omarrayyann.
|
||||
|
||||
Changed
|
||||
^^^^^^^
|
||||
|
||||
- ``Entity`` now raises a clear error at construction when its spec contains
|
||||
more than one freejoint. An entity models a single system rooted at one
|
||||
body, so it has at most one freejoint; a second one was previously accepted
|
||||
silently and only surfaced later as a cryptic shape mismatch when writing
|
||||
root state. Model each detached floating body as its own entry in
|
||||
``SceneCfg.entities`` instead.
|
||||
- Changed ``compute_root_relative_mpkpe`` to re-anchor the reference to the
|
||||
robot's root each step, removing yaw drift as well as translation so it
|
||||
measures intrinsic body pose error.
|
||||
- Changed ``compute_joint_velocity_error`` from an L2 norm to a per-joint
|
||||
RMS, so it no longer scales with the number of joints.
|
||||
- Bumped ``mujoco`` to 3.8 and ``mujoco-warp`` to 3.8.0. The ``multiccd``
|
||||
enable flag was removed in mujoco 3.8 (it became default-on), so configs
|
||||
that listed ``"multiccd"`` in ``MujocoCfg.enableflags`` need to drop it.
|
||||
@@ -68,15 +170,63 @@ Changed
|
||||
air-time fields (``current_air_time``, ``last_air_time``,
|
||||
``current_contact_time``, ``last_contact_time``) have shape ``[B, P]``,
|
||||
where ``P`` is the number of resolved primaries (:issue:`914`).
|
||||
- Event functions now share a single ``resolve_env_ids`` helper to expand
|
||||
``env_ids=None`` to all environments, replacing five copies of the same
|
||||
guard. ``push_by_setting_velocity`` and ``apply_external_force_torque``
|
||||
accept ``env_ids=None`` too, so they work as global-time interval terms.
|
||||
Documented when to use ``apply_external_force_torque`` (a constant,
|
||||
self-managed wrench) versus ``apply_body_impulse`` (transient, automatic
|
||||
impulses) versus ``push_by_setting_velocity`` (an instantaneous velocity
|
||||
kick).
|
||||
|
||||
Fixed
|
||||
^^^^^
|
||||
|
||||
- Fixed ``ManagerBasedRlEnv`` initializing Warp on all visible CUDA devices
|
||||
even when constructed with ``device="cpu"``. ``seed_rng`` now accepts a
|
||||
``device`` argument and skips ``wp.rand_init`` on CPU devices, so a
|
||||
CPU-only env no longer claims a CUDA context on machines with a visible
|
||||
GPU (:issue:`949`).
|
||||
- Removed use of deprecated ``warp-lang`` symbols (``wp.context.runtime``
|
||||
and ``wp.context.Device``) that were dropped in newer ``warp-lang``
|
||||
releases, causing ``AttributeError: module 'warp' has no attribute
|
||||
'context'`` at import/runtime. mjlab now uses
|
||||
``wp.get_cuda_driver_version()`` and ``wp.Device`` instead
|
||||
(:issue:`967`). Contribution by @rdeits.
|
||||
- Fixed the tracking ``evaluate`` script scoring each metric against the
|
||||
next motion frame; the reference is now snapshotted before each step to
|
||||
match the reward.
|
||||
- Fixed the tracking end-effector metrics silently scoring zero for an
|
||||
unknown body name; they now raise ``ValueError``.
|
||||
- Fixed ``compute_mpkpe`` measuring root-relative instead of global error;
|
||||
it now uses the global reference ``body_pos_w`` (:issue:`1006`).
|
||||
- Fixed heavy flicker in offscreen training videos on rough-terrain tasks.
|
||||
The renderer recomputed its context "neighbor" robots every frame from
|
||||
``env_origins``, which the terrain curriculum mutates on reset, so the
|
||||
neighbor set kept changing and robots popped in and out. The neighbor
|
||||
set is now computed once and cached (:issue:`979`).
|
||||
- Fixed command delay only applying to an actuator's position target.
|
||||
``IdealPdActuator`` and ``DcMotorActuator`` also use velocity and effort, which
|
||||
arrived undelayed and out of sync; all command targets now share one delay.
|
||||
Zero-reference setups are unaffected.
|
||||
- Fixed duplicate random seeds across nodes in multi-node training. The
|
||||
per-process seed offset in ``scripts/train.py`` now uses the global
|
||||
``RANK`` instead of ``LOCAL_RANK``. Contribution by @bd-pdomanico.
|
||||
- Fixed ``apply_body_impulse`` firing an impulse on the very first step (and
|
||||
the first step after every reset) instead of starting with a cooldown as
|
||||
documented. The cooldown is now sampled lazily on the first call so impulse
|
||||
timing is decorrelated from episode resets (:issue:`973`).
|
||||
- Fixed ``dr.pd_gains`` and ``dr.effort_limits`` silently no-oping when
|
||||
passed an ``Operation`` object (e.g. ``dr.scale``) instead of a string.
|
||||
Both functions now accept ``Operation | str`` like every other DR event
|
||||
and raise ``ValueError`` for unsupported operations (:issue:`971`).
|
||||
- Fixed ``ContactSensor`` with ``global_frame=True`` and
|
||||
``reduce`` ∈ {``"none"``, ``"mindist"``, ``"maxforce"``} producing forces
|
||||
rotated onto the wrong axis. The contact-frame→world rotation matrix had
|
||||
its columns ordered ``[tangent, tangent2, normal]`` instead of
|
||||
``[normal, tangent, tangent2]``, projecting the normal-force component
|
||||
onto a tangent direction. Contribution by @bd-pdomanico.
|
||||
- Fixed ``extras["log"]`` entries written by reward terms (e.g. ``Metrics/*``
|
||||
values in velocity tasks) being silently discarded on any step where at
|
||||
least one environment resets. ``_reset_idx`` was clearing the dict after
|
||||
``reward_manager.compute()`` had already populated it. The clear now
|
||||
happens at the top of ``step()`` and ``reset()`` so that all entries
|
||||
survive (:issue:`957`).
|
||||
- Fixed ``ContactSensor.compute_first_contact`` and ``compute_first_air``
|
||||
occasionally missing events when a contact began or ended right at the
|
||||
last physics substep of a control step. ``current_contact_time`` /
|
||||
|
||||
@@ -197,8 +197,8 @@ example, a ``CollisionCfg`` with ``geom_names_expr=(".*_foot.*",)``
|
||||
sets contact parameters only on foot geoms. See the asset zoo
|
||||
(``mjlab.asset_zoo.robots``) for complete examples.
|
||||
|
||||
Per-world mesh variants
|
||||
^^^^^^^^^^^^^^^^^^^^^^^
|
||||
Heterogeneous worlds
|
||||
^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
For scenes that need different mesh assets in different parallel worlds
|
||||
(for example, training a manipulation policy that generalizes across
|
||||
@@ -206,7 +206,7 @@ object shapes), use ``VariantEntityCfg`` instead of ``EntityCfg``. Each
|
||||
world is assigned a variant proportional to a configurable weight, and
|
||||
mesh-dependent compiled constants (collision bounds, body inertials,
|
||||
subtree mass) are stored as per-world arrays so domain randomization and
|
||||
viewers stay consistent. See :ref:`per_world_mesh`.
|
||||
viewers stay consistent. See :ref:`heterogeneous_worlds`.
|
||||
|
||||
Subclassing Entity
|
||||
^^^^^^^^^^^^^^^^^^
|
||||
|
||||
@@ -1,50 +1,35 @@
|
||||
.. _per_world_mesh:
|
||||
.. _heterogeneous_worlds:
|
||||
|
||||
Mesh Variants
|
||||
=============
|
||||
Heterogeneous Worlds
|
||||
====================
|
||||
|
||||
Mesh variants let a single batched simulation run with different mesh
|
||||
assets in different parallel worlds. World 0 may simulate a cube, world
|
||||
1 a sphere, and world 2 a bowl, all sharing the same compiled scene
|
||||
and the same kinematic structure. The result is a heterogeneous batch
|
||||
in which the mesh and its derived constants vary across worlds while
|
||||
everything else (the body tree, the joint structure, the contact and
|
||||
solver setup) is fixed.
|
||||
|
||||
Mesh variants are configured at the entity level through
|
||||
``VariantEntityCfg`` and ``VariantCfg``. Once configured,
|
||||
domain randomization, the native viewer, the offscreen renderer, and
|
||||
the Viser viewer all pick up the variant assignment automatically.
|
||||
mjlab can run a single batched simulation in which different parallel
|
||||
worlds use different mesh assets for the same logical entity. World 0
|
||||
may simulate a cube, world 1 a sphere, world 2 a bowl. All worlds
|
||||
share the same compiled scene and the same body and joint structure;
|
||||
only the meshes and the per-geom attributes that travel with them
|
||||
(friction, contact bits, mass, density, and a few more) differ across
|
||||
worlds. Articulated props work too (you can have a hinge or slide
|
||||
below the variant's root), as long as the joint topology matches
|
||||
across variants. The feature is exposed through ``VariantEntityCfg``.
|
||||
The full breakdown of what can and cannot vary across variants is in
|
||||
the next section.
|
||||
|
||||
|
||||
How it works
|
||||
------------
|
||||
Quickstart
|
||||
----------
|
||||
|
||||
A standard ``EntityCfg`` provides a single ``spec_fn`` that returns one
|
||||
``MjSpec``. A ``VariantEntityCfg`` provides a dictionary of named
|
||||
variants, each with its own ``spec_fn`` and a weight controlling the
|
||||
proportion of worlds that use it.
|
||||
|
||||
**All variants must declare the same kinematic structure.** The batched
|
||||
simulator assumes a single topology across worlds; per-world variation
|
||||
is confined to mesh assets and the constants derived from them. mjlab
|
||||
uses the first variant's body tree as the template and copies mesh
|
||||
assets and explicit body inertials from the others. Geom-level
|
||||
properties on later variants such as ``rgba``, friction, and material
|
||||
assignments are not propagated; control per-world appearance through
|
||||
domain randomization on ``geom_rgba`` or ``mat_rgba``. The structural
|
||||
check is enforced at construction time and raises a ``ValueError``
|
||||
describing the first mismatch. Variants must also be floating-base
|
||||
(declare a free joint on the root body); fixed-base variants are
|
||||
rejected.
|
||||
|
||||
A minimal two-variant config:
|
||||
Say you want some parallel worlds to hold a sphere and others to hold
|
||||
a cone, with a single shared scene running both at once. Define each
|
||||
variant as a function that returns an ``MjSpec``, then group them
|
||||
under one ``VariantEntityCfg``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import mujoco
|
||||
|
||||
from mjlab.entity import EntityCfg, VariantCfg, VariantEntityCfg
|
||||
from mjlab.entity import EntityCfg, VariantEntityCfg
|
||||
|
||||
|
||||
def make_sphere_spec() -> mujoco.MjSpec:
|
||||
spec = mujoco.MjSpec()
|
||||
@@ -56,40 +41,267 @@ A minimal two-variant config:
|
||||
body.add_geom(type=mujoco.mjtGeom.mjGEOM_MESH, meshname="visual")
|
||||
return spec
|
||||
|
||||
# ``make_cone_spec`` follows the same shape with
|
||||
# ``mesh.make_cone(nedge=16, radius=0.04)`` in place of the sphere call.
|
||||
|
||||
def make_cone_spec() -> mujoco.MjSpec:
|
||||
spec = mujoco.MjSpec()
|
||||
mesh = spec.add_mesh(name="visual")
|
||||
mesh.make_cone(nedge=16, radius=0.04)
|
||||
body = spec.worldbody.add_body(name="prop")
|
||||
body.add_freejoint()
|
||||
body.add_geom(type=mujoco.mjtGeom.mjGEOM_MESH, meshname="visual")
|
||||
return spec
|
||||
|
||||
|
||||
object_cfg = VariantEntityCfg(
|
||||
variants={
|
||||
"sphere": VariantCfg(spec_fn=make_sphere_spec, weight=1.0),
|
||||
"cone": VariantCfg(spec_fn=make_cone_spec, weight=2.0),
|
||||
"sphere": make_sphere_spec,
|
||||
"cone": make_cone_spec,
|
||||
},
|
||||
assignment={"cone": 2.0}, # twice as many cones as spheres
|
||||
init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
|
||||
)
|
||||
|
||||
During scene construction mjlab merges the per-variant specs into a
|
||||
single ``MjSpec`` whose mesh slots are padded to the maximum count any
|
||||
variant uses, then writes a per-world ``geom_dataid`` table that
|
||||
selects the right mesh for each world. In the merged scene
|
||||
``geom_dataid`` is no longer a flat ``(ngeom,)`` vector but a
|
||||
``(num_envs, ngeom)`` table whose rows differ by variant. A value of
|
||||
``-1`` marks a disabled mesh slot, used for variants with fewer mesh
|
||||
geoms than the maximum.
|
||||
Plug the variant entity into a :ref:`scene` exactly like a regular
|
||||
``EntityCfg``:
|
||||
|
||||
Mesh choice is entangled with several other compiled-model constants:
|
||||
geom collision bounds, geom local frames, body inertials, subtree mass,
|
||||
and inverse weights. mjlab compiles each unique row of the
|
||||
``geom_dataid`` table on the host and copies the relevant compiled
|
||||
fields into per-world arrays on the GPU, so each world's compiled
|
||||
constants stay consistent with that world's mesh selection. The full
|
||||
list of fields handled this way is in
|
||||
``mjlab.sim.mesh_variants.VARIANT_DEPENDENT_FIELDS``.
|
||||
.. code-block:: python
|
||||
|
||||
from mjlab.scene import SceneCfg
|
||||
|
||||
scene_cfg = SceneCfg(
|
||||
num_envs=4096,
|
||||
entities={"object": object_cfg},
|
||||
)
|
||||
|
||||
Twice as many worlds will hold a cone as a sphere. Variants not listed
|
||||
in the ``assignment`` dict default to weight 1.0; omit ``assignment``
|
||||
entirely for uniform allocation across all variants.
|
||||
|
||||
|
||||
What variants can differ in
|
||||
---------------------------
|
||||
|
||||
**Free to vary across variants:** the mesh asset assigned to each
|
||||
slot, the number of mesh geoms per ``(body, role)`` bucket on the
|
||||
variant body (one variant can have more collision meshes than
|
||||
another), the per-mesh-geom attributes that travel with the mesh
|
||||
(friction, contact bits, mass, density, ``condim``, and a handful of
|
||||
others), and explicit body inertial values within whichever single
|
||||
inertial mode the variants agree on per body.
|
||||
|
||||
**Must match across variants:** the body tree, joint topology,
|
||||
primitive (non-mesh) geoms, and any actuators / sensors / tendons /
|
||||
equalities. Variants must also agree on the inertial representation
|
||||
per body (mesh-derived, diagonal, or fullinertia), and may not use the
|
||||
reserved ``mjlab/pad/`` name prefix on any element. Variant entities
|
||||
must also be floating-base: the root body declares a freejoint.
|
||||
|
||||
The validator runs at entity build time and raises ``ValueError``
|
||||
naming the offending variant and the exact mismatch.
|
||||
|
||||
|
||||
How variants are assembled
|
||||
--------------------------
|
||||
|
||||
mjlab merges every variant's mesh assets into a single ``MjSpec`` and
|
||||
gives the variant body enough mesh-geom *slots* to cover the maximum
|
||||
mesh count any variant uses for each ``(body, role)`` bucket. A slot
|
||||
is identified by ``(body_path, role, ordinal)``. ``role`` is "visual"
|
||||
or "collision", derived from ``contype``/``conaffinity``;
|
||||
mujoco_warp's ``geom_contype``/``geom_conaffinity`` are 1D shared
|
||||
(not per-world), so a slot's role is fixed across worlds by
|
||||
construction.
|
||||
|
||||
A worked example
|
||||
~~~~~~~~~~~~~~~~
|
||||
|
||||
Say variant ``sphere`` has 1 visual mesh geom and 2 collision mesh
|
||||
geoms on the prop body, and variant ``cone`` has 1 visual mesh geom
|
||||
and 4 collision mesh geoms on the same body.
|
||||
|
||||
.. code-block:: text
|
||||
|
||||
sphere variant body cone variant body
|
||||
------------------- -------------------
|
||||
prop body prop body
|
||||
[visual] sphere_vis [visual] cone_vis
|
||||
[coll] sphere_col_0 [coll] cone_col_0
|
||||
[coll] sphere_col_1 [coll] cone_col_1
|
||||
[coll] cone_col_2
|
||||
[coll] cone_col_3
|
||||
|
||||
mjlab walks each variant's body tree, buckets mesh geoms by
|
||||
``(body_path, role)``, and lays the union out as slots:
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 8 18 8 12 27 27
|
||||
|
||||
* - Slot
|
||||
- body_path
|
||||
- role
|
||||
- ordinal
|
||||
- sphere fills with
|
||||
- cone fills with
|
||||
* - 0
|
||||
- /prop
|
||||
- visual
|
||||
- 0
|
||||
- sphere_vis
|
||||
- cone_vis
|
||||
* - 1
|
||||
- /prop
|
||||
- collision
|
||||
- 0
|
||||
- sphere_col_0
|
||||
- cone_col_0
|
||||
* - 2
|
||||
- /prop
|
||||
- collision
|
||||
- 1
|
||||
- sphere_col_1
|
||||
- cone_col_1
|
||||
* - 3
|
||||
- /prop
|
||||
- collision
|
||||
- 2
|
||||
- *(unfilled)*
|
||||
- cone_col_2
|
||||
* - 4
|
||||
- /prop
|
||||
- collision
|
||||
- 3
|
||||
- *(unfilled)*
|
||||
- cone_col_3
|
||||
|
||||
Five slots total. The merged scene's prop body has five mesh geoms:
|
||||
slot 0 plus four collision slots (the union of sphere's two and
|
||||
cone's four). At merge time, every variant's mesh asset is added to
|
||||
the merged spec under a unique name (e.g.
|
||||
``sphere/sphere_vis``, ``cone/cone_col_2``).
|
||||
|
||||
The merged scene compiles once into a single canonical ``MjModel``
|
||||
that every world in the batch agrees on layout-wise: same nbody,
|
||||
ngeom, same body and geom IDs. mjlab's per-world overrides on top of
|
||||
that one model are what make worlds heterogeneous.
|
||||
|
||||
What each world sees at runtime
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Worlds where ``sphere`` is active see only its three meshes; the two
|
||||
extra collision slots are disabled via per-world ``geom_dataid = -1``,
|
||||
and mujoco_warp skips them. Worlds where ``cone`` is active see all
|
||||
five meshes wired up.
|
||||
|
||||
.. list-table::
|
||||
:header-rows: 1
|
||||
:widths: 14 14 12 12 12 12 12
|
||||
|
||||
* - World
|
||||
- variant
|
||||
- slot 0
|
||||
- slot 1
|
||||
- slot 2
|
||||
- slot 3
|
||||
- slot 4
|
||||
* - 0
|
||||
- sphere
|
||||
- sphere_vis
|
||||
- sphere_col_0
|
||||
- sphere_col_1
|
||||
- **off (-1)**
|
||||
- **off (-1)**
|
||||
* - 1
|
||||
- cone
|
||||
- cone_vis
|
||||
- cone_col_0
|
||||
- cone_col_1
|
||||
- cone_col_2
|
||||
- cone_col_3
|
||||
|
||||
Three categories of per-world override carry the variation:
|
||||
|
||||
* **geom_dataid** is a ``(num_envs, ngeom)`` table. Its row for
|
||||
world W picks which compiled mesh each slot points at. ``-1`` is
|
||||
the "skip me" sentinel mujoco_warp already understands.
|
||||
* **Mesh-derived fields** (``geom_size``, ``geom_rbound``,
|
||||
``geom_aabb``, ``geom_pos``, ``geom_quat``, ``body_mass``,
|
||||
``body_subtreemass``, ``body_inertia``, ``body_invweight0``,
|
||||
``body_ipos``, ``body_iquat``) are stored as ``(num_envs, ...)``
|
||||
arrays. The values for sphere worlds reflect a sphere-shaped
|
||||
inertia tensor and sphere-sized AABBs; the values for cone worlds
|
||||
reflect the cone. The full list is in
|
||||
``mjlab.entity.variants.VARIANT_DEPENDENT_FIELDS``.
|
||||
* **Per-mesh-geom attributes** (contact bits, friction, mass,
|
||||
density, condim, group, priority, rgba, solref, solimp, margin,
|
||||
gap) are captured per variant in ``VariantGeomSpec`` at merge time
|
||||
and restored verbatim on the slot geom during the per-variant
|
||||
reference compile. So if sphere's collision geoms have
|
||||
``friction=0.5`` and cone's have ``friction=1.2``, world W's
|
||||
per-step friction reflects the assigned variant's source value.
|
||||
The one exception is ``material``, which is not propagated across
|
||||
variants; if you need per-world appearance variation use DR on
|
||||
``geom_rgba`` / ``mat_rgba``.
|
||||
|
||||
If ``sphere`` adds a body that ``cone`` lacks (or vice versa), the
|
||||
validator rejects the configuration before any of the merge logic
|
||||
runs. The slot mechanism only flexes mesh geom counts within
|
||||
matching bodies; everything structural above the geom level must
|
||||
agree.
|
||||
|
||||
.. note::
|
||||
|
||||
**Doesn't compiling the merged scene ruin the prop body's
|
||||
inertia?**
|
||||
|
||||
No, but it's worth understanding why, because the naive intuition
|
||||
says it should. If you stuck every variant's mesh geoms on the
|
||||
prop body and called ``spec.compile()``, MuJoCo would sum each
|
||||
geom's inertial contribution, and you would get a body whose mass
|
||||
and inertia tensor are a meaningless mix of every variant's shape.
|
||||
|
||||
mjlab avoids this in two layers:
|
||||
|
||||
* **The merged scene does not stick every variant's geoms on the
|
||||
body.** The prop body in the merged spec carries variant 0's
|
||||
mesh geoms (with their original mass and density) plus, for any
|
||||
slot variant 0 doesn't fill, a synthesized padding geom that has
|
||||
``mass = 0`` and ``density = 0``. Padding contributes nothing to
|
||||
body inertia. Other variants' meshes are present in the merged
|
||||
spec only as **mesh assets** (in the assets section, not as geoms
|
||||
on any body). They get wired in at runtime via per-world
|
||||
``geom_dataid`` and never affect the host compile's inertial
|
||||
sums.
|
||||
* **Per-world overrides come from per-variant source compiles.**
|
||||
Even with the above, the merged-scene compile's prop body inertia
|
||||
is only correct for variant 0. For every other variant, mjlab
|
||||
compiles that variant's original source spec in isolation (one
|
||||
body, one variant's worth of meshes), reads the resulting
|
||||
``body_mass``, ``body_inertia``, ``body_ipos``, ``body_iquat``,
|
||||
``body_invweight0``, and ``body_subtreemass``, and writes them
|
||||
into the per-world arrays at the prop body's index.
|
||||
|
||||
Net result: world W's prop body inertia is byte-equal to what you
|
||||
would get by compiling variant W's source spec on its own. There
|
||||
is a regression test
|
||||
(``test_visual_collision_split_inertia_matches_independent_compile``
|
||||
in ``tests/test_variants.py``) that asserts exactly this against
|
||||
independent per-variant compiles.
|
||||
|
||||
|
||||
World assignment
|
||||
----------------
|
||||
|
||||
mjlab assigns variants to worlds proportionally by weight using the
|
||||
How worlds get mapped to variants is controlled by the ``assignment``
|
||||
field on ``VariantEntityCfg``. It accepts three shapes:
|
||||
|
||||
* ``None`` (default): uniform allocation across variants.
|
||||
* ``dict[str, float]``: per-variant weights. Variants not listed
|
||||
default to weight 1.0.
|
||||
* ``Callable[[int], Sequence[int]]``: an explicit assignment function
|
||||
called with ``num_envs`` at simulation init.
|
||||
|
||||
Both the ``None`` and dict cases use the
|
||||
`largest remainder method
|
||||
<https://en.wikipedia.org/wiki/Largest_remainder_method>`_. Each
|
||||
variant's quota is ``q_i = (w_i / sum(w)) * num_envs``; each variant
|
||||
@@ -98,14 +310,28 @@ first receives ``floor(q_i)`` worlds, and the remaining
|
||||
fractional remainders, with ties broken by declaration order. For
|
||||
``num_envs = 10`` and weights ``(1.0, 2.0, 1.0)`` this gives
|
||||
``(3, 5, 2)`` worlds per variant. Weights are normalized internally,
|
||||
so ``(1, 2, 1)`` and ``(0.25, 0.5, 0.25)`` produce identical
|
||||
assignments. A weight of zero is allowed and produces zero worlds for
|
||||
that variant; at least one variant must have a positive weight.
|
||||
so ``{"a": 1, "b": 2, "c": 1}`` and ``{"a": 0.25, "b": 0.5, "c": 0.25}``
|
||||
produce identical assignments. A weight of zero is allowed and
|
||||
produces zero worlds for that variant; at least one variant must end
|
||||
up with positive weight.
|
||||
|
||||
Variant assignment is fixed at simulation initialization and does not
|
||||
resample on episode reset. The intended use is heterogeneous training
|
||||
across the batch, not per-episode mesh randomization. To inspect the
|
||||
assignment from user code, read ``env.sim.world_to_variant``:
|
||||
The default and dict paths are purely deterministic given
|
||||
``(assignment, num_envs)``. With ``assignment={"a": 1, "b": 1}`` and
|
||||
``num_envs = 8`` you always get ``[0, 0, 0, 0, 1, 1, 1, 1]``. There is
|
||||
no seed involved; rerunning the same config produces the same
|
||||
partition every time. Note that the partition's *boundaries* depend
|
||||
on ``num_envs``, so world W's variant is not necessarily stable when
|
||||
you change ``num_envs``. If you need explicit per-world stability
|
||||
across batch sizes (e.g. "world 0 is always variant 0, world 1 is
|
||||
always variant 1, regardless of how many envs I launch"), use a
|
||||
callable assignment as below.
|
||||
|
||||
Variant assignment is fixed at ``Simulation`` initialization and does
|
||||
not resample on episode reset. The intended use is heterogeneous
|
||||
training across the batch, not per-episode mesh randomization.
|
||||
|
||||
Read the resolved assignment from user code via
|
||||
``env.sim.world_to_variant``:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@@ -118,23 +344,53 @@ variants were declared in ``VariantEntityCfg.variants``. The dict is
|
||||
empty for non-variant scenes.
|
||||
|
||||
|
||||
Custom assignment with a callable
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
When the weighted default is not what you want, pass a callable to
|
||||
``assignment``. The callable receives ``num_envs`` and must return a
|
||||
length-``num_envs`` sequence of variant indices in
|
||||
``[0, len(variants))``. The returned sequence's length and bounds are
|
||||
validated at sim init; mismatches raise a ``ValueError`` naming the
|
||||
offending entity.
|
||||
|
||||
A few patterns:
|
||||
|
||||
**Round-robin** - cycle through variants by world index.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
cfg = VariantEntityCfg(
|
||||
variants={"a": make_a, "b": make_b, "c": make_c},
|
||||
assignment=lambda n: [w % 3 for w in range(n)],
|
||||
)
|
||||
|
||||
**Stratified halves** - first half is variant 0, second half is
|
||||
variant 1.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
cfg = VariantEntityCfg(
|
||||
variants={"easy": make_easy, "hard": make_hard},
|
||||
assignment=lambda n: [0] * (n // 2) + [1] * (n - n // 2),
|
||||
)
|
||||
|
||||
Domain randomization
|
||||
--------------------
|
||||
|
||||
Domain randomization on variant scenes preserves per-variant baselines
|
||||
automatically. When the simulation initializes, mjlab snapshots the
|
||||
variant-dependent fields (``body_mass``, ``body_inertia``,
|
||||
``geom_size``, and others listed in ``VARIANT_DEPENDENT_FIELDS``) as
|
||||
``(num_envs, ...)`` tensors and registers them in
|
||||
``sim.per_world_default_fields``. Domain randomization operations that
|
||||
read defaults (scale, additive offsets) detect this registration and
|
||||
index the per-world default array by environment, so a 10% mass scale
|
||||
variant-dependent fields as ``(num_envs, ...)`` tensors and registers
|
||||
them in ``sim.per_world_default_fields``. DR operations that read
|
||||
defaults (scale, additive offsets) detect this registration and index
|
||||
the per-world default array by environment, so a 10% mass scale
|
||||
applied across a batch containing a 100 g sphere variant and a 1 kg
|
||||
cube variant produces 10% perturbations around each variant's own
|
||||
mass, not 10% of a shared template mass. Fields that are not
|
||||
variant-dependent (``geom_friction``, ``dof_armature``,
|
||||
``dof_damping``, and so on) behave identically on variant and
|
||||
non-variant scenes.
|
||||
cube variant produces 10% perturbations *around each variant's own
|
||||
mass*, not 10% of a shared template mass.
|
||||
|
||||
Fields that are not variant-dependent (``geom_friction``,
|
||||
``dof_armature``, ``dof_damping``, and so on) behave identically on
|
||||
variant and non-variant scenes.
|
||||
|
||||
For inertial randomization the recommended path is
|
||||
``dr.pseudo_inertia``, which jointly randomizes mass, COM offset,
|
||||
@@ -169,32 +425,66 @@ Convex hull visualization is computed per variant from the variant's
|
||||
mesh vertices.
|
||||
|
||||
|
||||
Performance considerations
|
||||
--------------------------
|
||||
Performance
|
||||
-----------
|
||||
|
||||
Mesh variants do not add per-step overhead in the GPU kernels.
|
||||
Variant-dependent fields are stored as per-world arrays accessed by
|
||||
world index in the existing kernels, with no branching or dispatch
|
||||
on variant.
|
||||
**Per-step cost is unaffected by variant count.** Variant-dependent
|
||||
fields are stored as per-world arrays accessed by world index in the
|
||||
existing kernels, with no branching or dispatch on variant.
|
||||
|
||||
Initialization is the main consideration. mjlab compiles each unique
|
||||
row of the ``geom_dataid`` table by taking a fresh ``MjSpec.copy()``,
|
||||
editing the mesh selection and (if applicable) the explicit body
|
||||
inertials, and calling ``spec.compile()``. This work scales with the
|
||||
number of unique variant combinations rather than with ``num_envs``.
|
||||
For a scene with one variant entity declaring k variants, this is k
|
||||
host compiles regardless of how many worlds use each variant. With
|
||||
multiple variant entities the unique-row count is bounded by the
|
||||
product of their variant counts in the worst case, so a scene with
|
||||
two variant entities of 5 variants each could trigger up to 25 host
|
||||
compiles at init.
|
||||
**Construction cost is linear in the total variant count.** mjlab
|
||||
compiles the merged scene once to produce the canonical ``MjModel``,
|
||||
then compiles each variant's original (un-merged) source spec in
|
||||
isolation to recover that variant's per-body and per-geom mesh-derived
|
||||
fields. Each per-variant compile sees only that variant's single body
|
||||
and mesh, so its cost is independent of the total number of variants
|
||||
in the scene.
|
||||
|
||||
``MjSpec.copy()`` and ``spec.compile()`` are non-trivial operations,
|
||||
and their cost grows with scene size. For a scene with many variant
|
||||
entities or many variants per entity, the cumulative initialization
|
||||
cost can be measured in seconds. This cost is paid once at startup
|
||||
and does not affect training throughput.
|
||||
For a scene with one variant entity declaring k variants, construction
|
||||
runs ``1 + k`` compiles. With multiple variant entities, compiles
|
||||
decouple across entities: two variant entities of 5 variants each cost
|
||||
``1 + 5 + 5 = 11`` compiles, not ``1 + 5 * 5 = 26``. As an order of
|
||||
magnitude on CPU with typical procedural meshes, each per-variant
|
||||
compile takes around 1-2 ms, so a scene with 100 variants pays a few
|
||||
hundred milliseconds at startup and a scene with 1000 variants pays
|
||||
roughly two seconds.
|
||||
|
||||
The merged spec contains every variant's mesh assets simultaneously.
|
||||
Memory footprint at scene-build time scales with the total number of
|
||||
mesh vertices and faces across all declared variants.
|
||||
The merged spec contains every variant's mesh assets simultaneously,
|
||||
so memory at scene-build time scales with the total mesh vertex /
|
||||
face count across all variants. This is paid once at startup and does
|
||||
not affect training throughput.
|
||||
|
||||
|
||||
Limitations
|
||||
-----------
|
||||
|
||||
**Floating-base only.** Each variant's root body must declare a free
|
||||
joint. Fixed-base variants are rejected; mocap auto-wrapping that
|
||||
applies to non-variant entities is not applied here.
|
||||
|
||||
**Material assets are not propagated.** Each variant's ``contype``,
|
||||
``conaffinity``, ``condim``, ``friction``, ``mass``, ``density``,
|
||||
``group``, ``priority``, ``rgba``, ``solref``, ``solimp``, ``margin``,
|
||||
and ``gap`` are restored per-world during compile, but the
|
||||
``material`` reference on slot geoms inherits whichever material the
|
||||
template variant set. Use DR on ``geom_rgba`` / ``mat_rgba`` for
|
||||
per-world appearance variation.
|
||||
|
||||
**Assignment is fixed at sim init.** There is no API to swap a world
|
||||
to a different variant on episode reset. World W's mesh asset is
|
||||
whatever it was assigned at init for the lifetime of the simulation.
|
||||
Per-episode mesh randomization is not supported today; DR can vary
|
||||
scalar properties (mass, friction, color, scale) on a fixed variant
|
||||
but cannot swap one mesh for another.
|
||||
|
||||
**No support for per-world differing kinematic topology.** Variants
|
||||
must share the same body tree, joints, and actuator/sensor counts,
|
||||
so you cannot configure things like:
|
||||
|
||||
* a different number of objects per world (world 0 has two props on
|
||||
the table, world 1 has three);
|
||||
* different articulation per world (world 0's prop is an articulated
|
||||
drawer with a slider joint, world 1's prop is a rigid block).
|
||||
|
||||
True heterogeneous topology requires upstream support in mujoco_warp
|
||||
that does not currently exist.
|
||||
|
||||
@@ -42,6 +42,27 @@ Not all CUDA versions are supported by MuJoCo Warp.
|
||||
- **Recommended**: CUDA **12.4+** (for conditional execution support in CUDA
|
||||
graphs).
|
||||
|
||||
How do I run on CPU without touching the GPU?
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Passing ``device="cpu"`` puts all mjlab computation on the CPU, but it does
|
||||
**not** stop Warp from initializing the GPU. The first time Warp's runtime
|
||||
comes up, it eagerly enumerates and creates a CUDA context on **every**
|
||||
visible device, regardless of which device you requested. So on a machine
|
||||
with a visible GPU, a ``device="cpu"`` run still claims VRAM.
|
||||
|
||||
This happens inside Warp and cannot be prevented from Python once the
|
||||
package is imported. To keep the process entirely off the GPU, hide the
|
||||
devices from CUDA before launching:
|
||||
|
||||
.. code-block:: bash
|
||||
|
||||
CUDA_VISIBLE_DEVICES="" uv run train.py ...
|
||||
|
||||
With no visible CUDA devices, Warp initializes CPU-only and never allocates
|
||||
on the GPU. See `issue #949
|
||||
<https://github.com/mujocolab/mjlab/issues/949>`_ for background.
|
||||
|
||||
Performance
|
||||
-----------
|
||||
|
||||
|
||||
@@ -49,7 +49,6 @@ the geometry and how it scales with difficulty.
|
||||
terrain_generator=TerrainGeneratorCfg(
|
||||
size=(8.0, 8.0),
|
||||
num_rows=10,
|
||||
num_cols=20,
|
||||
border_width=20.0,
|
||||
curriculum=True,
|
||||
sub_terrains={
|
||||
@@ -70,10 +69,12 @@ the geometry and how it scales with difficulty.
|
||||
max_init_terrain_level=5,
|
||||
)
|
||||
|
||||
The generator creates a ``num_rows x num_cols`` grid of patches. The
|
||||
``sub_terrains`` dictionary maps names to ``SubTerrainCfg`` instances,
|
||||
and each sub-terrain's ``proportion`` weight controls how many columns
|
||||
(curriculum mode) or sampling probability (random mode) it receives.
|
||||
The generator creates a grid of patches sized ``num_rows`` by either
|
||||
``num_cols`` (random mode) or ``len(sub_terrains)`` (curriculum mode,
|
||||
where ``num_cols`` is ignored). The ``sub_terrains`` dictionary maps
|
||||
names to ``SubTerrainCfg`` instances; each sub-terrain's ``proportion``
|
||||
controls robot spawning distribution across columns in curriculum mode,
|
||||
or per-patch sampling probability in random mode.
|
||||
|
||||
|
||||
Grid layout
|
||||
@@ -82,30 +83,48 @@ Grid layout
|
||||
Two generation modes control how terrain types are distributed across
|
||||
the grid:
|
||||
|
||||
**Curriculum mode** (``curriculum=True``). Columns are deterministically
|
||||
assigned to terrain types based on their ``proportion`` weights. A type
|
||||
with proportion 0.4 in a 20-column grid gets 8 columns. All patches in
|
||||
a column share the same terrain type, and difficulty increases from row 0
|
||||
(easiest) to row ``num_rows - 1`` (hardest). This structured layout is
|
||||
what enables the curriculum system to advance environments to harder rows
|
||||
as performance improves.
|
||||
**Curriculum mode** (``curriculum=True``). Each terrain type gets exactly
|
||||
one column; the generator uses ``len(sub_terrains)`` columns regardless of
|
||||
``num_cols``. All patches in a column share the same terrain type, and
|
||||
difficulty increases from row 0 (easiest) to row ``num_rows - 1``
|
||||
(hardest). The ``proportion`` field controls how robots are distributed
|
||||
across columns at spawn time, not column count. This structured layout
|
||||
is what enables the curriculum system to advance environments to harder
|
||||
rows as performance improves.
|
||||
|
||||
**Random mode** (``curriculum=False``). Every patch independently samples
|
||||
a terrain type weighted by ``proportion`` and a difficulty from
|
||||
``difficulty_range``. This provides maximum variety but no structured
|
||||
difficulty progression.
|
||||
``difficulty_range``. ``num_cols`` is honored. This provides maximum
|
||||
variety but no structured difficulty progression.
|
||||
|
||||
|
||||
The difficulty parameter
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
|
||||
Each sub-terrain's generation function receives a ``difficulty`` value
|
||||
in ``[0, 1]``. This value linearly interpolates the terrain's
|
||||
configurable ranges. For example, a ``BoxPyramidStairsTerrainCfg`` with
|
||||
that linearly interpolates the terrain's configurable ranges. For
|
||||
example, a ``BoxPyramidStairsTerrainCfg`` with
|
||||
``step_height_range=(0.0, 0.2)`` produces flat ground at difficulty 0
|
||||
and 20 cm steps at difficulty 1. In curriculum mode, difficulty is
|
||||
determined by the row: row 0 gets the minimum, row ``num_rows - 1`` gets
|
||||
the maximum.
|
||||
and 20 cm steps at difficulty 1.
|
||||
|
||||
In curriculum mode, difficulty is determined by the row:
|
||||
``difficulty = lower + (upper - lower) * row / max(num_rows - 1, 1)``,
|
||||
where ``(lower, upper) = difficulty_range``. Row 0 is exactly
|
||||
``lower``, row ``num_rows - 1`` is exactly ``upper``, and intermediate
|
||||
rows are evenly spaced between them. All columns in a given row share
|
||||
the same difficulty scalar; the visible variation across columns comes
|
||||
from each sub-terrain type generating different geometry at the same
|
||||
difficulty.
|
||||
|
||||
.. note::
|
||||
|
||||
With ``num_rows=1`` and ``curriculum=True``, every patch is generated
|
||||
at ``difficulty = lower`` (the easiest configured difficulty). Use
|
||||
``curriculum=False`` if you want a single grid of randomly sampled
|
||||
difficulties instead.
|
||||
|
||||
In random mode, difficulty is sampled uniformly from
|
||||
``difficulty_range`` independently for every patch.
|
||||
|
||||
|
||||
Sub-terrain types
|
||||
@@ -244,17 +263,23 @@ and undulating ground that box geoms cannot represent.
|
||||
Preset configurations
|
||||
---------------------
|
||||
|
||||
mjlab ships two ready-made ``TerrainGeneratorCfg`` presets in
|
||||
mjlab ships three ready-made ``TerrainGeneratorCfg`` presets in
|
||||
``mjlab.terrains.config``:
|
||||
|
||||
``ROUGH_TERRAINS_CFG``
|
||||
A 10x20 grid with seven terrain types (flat, stairs, inverted
|
||||
stairs, slopes, inverted slopes, random rough, waves). Designed for
|
||||
locomotion training with a moderate difficulty range.
|
||||
A 10x20 random-mode grid with seven terrain types (flat, stairs,
|
||||
inverted stairs, slopes, inverted slopes, random rough, waves).
|
||||
Designed for locomotion training with a moderate difficulty range.
|
||||
Set ``curriculum=True`` via ``dataclasses.replace`` to use it as a
|
||||
curriculum grid (one column per terrain type).
|
||||
|
||||
``STAIRS_TERRAINS_CFG``
|
||||
A 10-row curriculum grid focused on stair traversal: flat plus
|
||||
three pyramid-stair variants of increasing difficulty.
|
||||
|
||||
``ALL_TERRAINS_CFG``
|
||||
A 10x16 grid with all sixteen terrain types at equal proportion.
|
||||
Useful for training on maximum terrain variety.
|
||||
A 10-row random-mode grid covering all available terrain types at
|
||||
equal proportion. Useful for training on maximum terrain variety.
|
||||
|
||||
Both can be used directly or customized with ``dataclasses.replace()``:
|
||||
|
||||
@@ -285,9 +310,9 @@ The key concepts:
|
||||
- The built-in ``terrain_levels_vel`` curriculum term promotes
|
||||
environments that track commanded velocity well and demotes
|
||||
environments that fall or fail to make progress.
|
||||
- When an environment reaches the maximum row, it is randomly reassigned
|
||||
to a lower row to prevent the policy from collapsing to a single
|
||||
difficulty level.
|
||||
- When an environment is promoted past the hardest row, it is randomly
|
||||
reassigned to any row in ``[0, num_rows)`` to prevent the policy from
|
||||
collapsing to a single difficulty level.
|
||||
|
||||
|
||||
Flat patch detection
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "uv_build"
|
||||
|
||||
[project]
|
||||
name = "mjlab"
|
||||
version = "1.3.0"
|
||||
version = "1.4.0"
|
||||
license = "Apache-2.0"
|
||||
license-files = ["LICENSE"]
|
||||
readme = { file = "README.md", content-type = "text/markdown" }
|
||||
@@ -37,15 +37,16 @@ dependencies = [
|
||||
"torch>=2.7.0",
|
||||
"torchrunx>=0.3.4",
|
||||
"warp-lang>=1.12.0",
|
||||
"mujoco-warp>=3.8.0",
|
||||
"mujoco>=3.8.0",
|
||||
"mujoco-warp>=3.8.0.3,~=3.8.0",
|
||||
"mujoco~=3.8.0",
|
||||
"trimesh>=4.8.3",
|
||||
"viser>=1.0.26",
|
||||
"mjviser>=0.0.13",
|
||||
"scipy>=1.15",
|
||||
"viser>=1.0.27",
|
||||
"mjviser>=0.0.14",
|
||||
"mediapy>=1.2.6",
|
||||
"imageio-ffmpeg",
|
||||
"tensordict",
|
||||
"rsl-rl-lib==5.2.0",
|
||||
"rsl-rl-lib==5.4.0",
|
||||
"tensorboard>=2.20.0",
|
||||
"onnxscript>=0.5.4",
|
||||
"wandb>=0.22.3",
|
||||
@@ -98,11 +99,11 @@ conflicts = [
|
||||
[{extra = "cu128"}, {extra = "cpu"}],
|
||||
]
|
||||
# The nightly index (py.mujoco.org) only has dev builds, and PEP 440 ranks
|
||||
# 3.7.0.devN < 3.7.0, so the >=3.7.0 floor in [project.dependencies] would
|
||||
# 3.8.0.devN < 3.8.0, so the ~=3.8.0 floor in [project.dependencies] would
|
||||
# reject them. This override loosens the constraint for uv resolution only.
|
||||
override-dependencies = ["mujoco>=3.8.0.dev0"]
|
||||
constraint-dependencies = [
|
||||
"GitPython>=3.1.47",
|
||||
"GitPython>=3.1.49",
|
||||
"lxml>=6.1.0",
|
||||
]
|
||||
required-environments = [
|
||||
@@ -111,12 +112,6 @@ required-environments = [
|
||||
]
|
||||
|
||||
[[tool.uv.index]]
|
||||
name = "tsinghua"
|
||||
url = "https://pypi.tuna.tsinghua.edu.cn/simple"
|
||||
default = true
|
||||
|
||||
[[tool.uv.index]]
|
||||
name = "pypi"
|
||||
url = "https://pypi.org/simple"
|
||||
|
||||
[[tool.uv.index]]
|
||||
@@ -146,8 +141,7 @@ torch = [
|
||||
{ index = "pytorch-cpu", extra = "cpu", marker = "sys_platform != 'darwin'" },
|
||||
]
|
||||
mujoco = { index = "mujoco" }
|
||||
mujoco-warp = { git = "https://github.com/google-deepmind/mujoco_warp", rev = "6f235d4" }
|
||||
mjviser = { git = "https://github.com/mujocolab/mjviser", rev = "1bdfd6fe79066b847a5f430000fcfbb53ec31a6f" }
|
||||
mujoco-warp = { git = "https://github.com/google-deepmind/mujoco_warp", rev = "88b55fc2696960b927bc12584994bb8412b36558" }
|
||||
|
||||
[tool.ruff]
|
||||
src = ["src"] # Helpful for recognizing first-party imports.
|
||||
|
||||
@@ -215,6 +215,27 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
|
||||
border-color: var(--accent);
|
||||
color: white;
|
||||
}}
|
||||
.range-selector {{
|
||||
display: flex;
|
||||
gap: 0.4rem;
|
||||
margin-bottom: 1rem;
|
||||
}}
|
||||
.range-btn {{
|
||||
background: var(--bg-card);
|
||||
border: 1px solid var(--border);
|
||||
border-radius: 6px;
|
||||
padding: 0.3rem 0.75rem;
|
||||
cursor: pointer;
|
||||
color: var(--text);
|
||||
font-size: 0.8rem;
|
||||
font-weight: 500;
|
||||
}}
|
||||
.range-btn:hover {{ border-color: var(--accent); }}
|
||||
.range-btn.active {{
|
||||
background: var(--accent);
|
||||
border-color: var(--accent);
|
||||
color: white;
|
||||
}}
|
||||
.tab-content {{ display: none; }}
|
||||
.tab-content.active {{ display: block; }}
|
||||
.tab-description {{
|
||||
@@ -303,12 +324,24 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
|
||||
|
||||
<div id="tracking" class="tab-content active">
|
||||
<p class="tab-description">Nightly motion imitation training and evaluation on Unitree G1 (1024 trials per run).</p>
|
||||
<div class="range-selector" id="range-selector">
|
||||
<button class="range-btn" data-days="30">30d</button>
|
||||
<button class="range-btn active" data-days="90">90d</button>
|
||||
<button class="range-btn" data-days="180">180d</button>
|
||||
<button class="range-btn" data-days="0">All</button>
|
||||
</div>
|
||||
<div class="charts" id="charts"></div>
|
||||
</div>
|
||||
|
||||
<div id="throughput" class="tab-content">
|
||||
<p class="tab-description">Physics simulation throughput across tasks (4096 parallel envs, NVIDIA RTX 5090).</p>
|
||||
<div class="task-grid" id="task-grid"></div>
|
||||
<div class="range-selector" id="range-selector-tp">
|
||||
<button class="range-btn" data-days="30">30d</button>
|
||||
<button class="range-btn active" data-days="90">90d</button>
|
||||
<button class="range-btn" data-days="180">180d</button>
|
||||
<button class="range-btn" data-days="0">All</button>
|
||||
</div>
|
||||
<div id="task-chart-panels"></div>
|
||||
</div>
|
||||
|
||||
@@ -393,6 +426,8 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
|
||||
}};
|
||||
|
||||
let charts = [];
|
||||
let trackingCharts = [];
|
||||
let throughputCharts = [];
|
||||
|
||||
function updateChartColors() {{
|
||||
const style = getComputedStyle(root);
|
||||
@@ -449,7 +484,7 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
|
||||
`;
|
||||
chartsContainer.appendChild(card);
|
||||
|
||||
charts.push(new Chart(card.querySelector('canvas'), {{
|
||||
const chart = new Chart(card.querySelector('canvas'), {{
|
||||
type: 'line',
|
||||
data: {{
|
||||
datasets: [
|
||||
@@ -459,7 +494,8 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
|
||||
borderColor: color,
|
||||
backgroundColor: color + '20',
|
||||
borderWidth: 2,
|
||||
pointRadius: 4,
|
||||
pointRadius: 2,
|
||||
pointHoverRadius: 5,
|
||||
tension: 0.1,
|
||||
fill: true
|
||||
}},
|
||||
@@ -531,7 +567,9 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
|
||||
}}
|
||||
}}
|
||||
}}
|
||||
}}));
|
||||
}});
|
||||
charts.push(chart);
|
||||
trackingCharts.push(chart);
|
||||
}});
|
||||
|
||||
// Tab switching
|
||||
@@ -621,7 +659,8 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
|
||||
borderColor: '#58a6ff',
|
||||
backgroundColor: '#58a6ff20',
|
||||
borderWidth: 2,
|
||||
pointRadius: 4,
|
||||
pointRadius: 2,
|
||||
pointHoverRadius: 5,
|
||||
tension: 0.1,
|
||||
fill: true
|
||||
}},
|
||||
@@ -631,7 +670,8 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
|
||||
borderColor: '#3fb950',
|
||||
backgroundColor: '#3fb95020',
|
||||
borderWidth: 2,
|
||||
pointRadius: 4,
|
||||
pointRadius: 2,
|
||||
pointHoverRadius: 5,
|
||||
tension: 0.1,
|
||||
fill: true
|
||||
}}
|
||||
@@ -697,6 +737,7 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
|
||||
}}
|
||||
}});
|
||||
charts.push(chart);
|
||||
throughputCharts.push(chart);
|
||||
throughputChartInstances[task] = {{ chart, panelId: `task-panel-${{i}}` }};
|
||||
|
||||
// Card click handler
|
||||
@@ -712,6 +753,23 @@ def generate_dashboard_html(runs: list[dict], throughput_data: list[dict]) -> st
|
||||
}} else {{
|
||||
taskGrid.innerHTML = '<p style="color: var(--text-dim)">No throughput data available. Run measure_throughput.py to generate data.</p>';
|
||||
}}
|
||||
|
||||
// Date-range windowing across both tracking and throughput charts.
|
||||
// Setting min and clearing max also resets any zoom/pan.
|
||||
function setRange(days) {{
|
||||
const min = days > 0 ? Date.now() - days * 86400000 : undefined;
|
||||
[...trackingCharts, ...throughputCharts].forEach(c => {{
|
||||
c.options.scales.x.min = min;
|
||||
c.options.scales.x.max = undefined;
|
||||
c.update();
|
||||
}});
|
||||
document.querySelectorAll('.range-btn').forEach(b =>
|
||||
b.classList.toggle('active', parseInt(b.dataset.days) === days));
|
||||
}}
|
||||
document.querySelectorAll('.range-btn').forEach(btn => {{
|
||||
btn.addEventListener('click', () => setRange(parseInt(btn.dataset.days)));
|
||||
}});
|
||||
setRange(90);
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -764,7 +822,11 @@ def main(
|
||||
if run_id in eval_results_by_id:
|
||||
print(f"Using cached result for {run_id}")
|
||||
else:
|
||||
result = evaluate_run(run_path, num_envs)
|
||||
try:
|
||||
result = evaluate_run(run_path, num_envs)
|
||||
except RuntimeError as e:
|
||||
print(f"Skipping {run_path}: {e}")
|
||||
continue
|
||||
eval_results_by_id[run_id] = result
|
||||
new_evals += 1
|
||||
else:
|
||||
@@ -783,7 +845,11 @@ def main(
|
||||
print(f"Reached eval limit ({eval_limit}), skipping remaining new runs")
|
||||
break
|
||||
run_path = f"{entity}/{project}/{run.id}"
|
||||
result = evaluate_run(run_path, num_envs)
|
||||
try:
|
||||
result = evaluate_run(run_path, num_envs)
|
||||
except RuntimeError as e:
|
||||
print(f"Skipping {run.name} ({run.id}): {e}")
|
||||
continue
|
||||
eval_results_by_id[run.id] = result
|
||||
new_evals += 1
|
||||
|
||||
|
||||
@@ -0,0 +1,115 @@
|
||||
"""Interactive single-patch terrain explorer (Viser + MuJoCo MjSpec).
|
||||
|
||||
Run with:
|
||||
uv run python scripts/tools/terrain_explorer.py
|
||||
uv run python scripts/tools/terrain_explorer.py --port 8081
|
||||
|
||||
Then open the printed URL (default http://localhost:8080).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import time
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
import viser
|
||||
from mjviser.conversions import merge_geoms
|
||||
|
||||
from mjlab.terrains.config import ALL_TERRAIN_PRESETS
|
||||
from mjlab.terrains.terrain_generator import TerrainGenerator, TerrainGeneratorCfg
|
||||
|
||||
PATCH_SIZE = (8.0, 8.0)
|
||||
|
||||
|
||||
# Per-preset overrides applied when building in the explorer (e.g. to surface
|
||||
# difficulty-driven behavior that is off by default).
|
||||
_PRESET_OVERRIDES: dict[str, dict] = {
|
||||
"random_rough": {"scale_with_difficulty": True},
|
||||
}
|
||||
|
||||
|
||||
def _build_terrain_mesh(preset_name: str, difficulty: float, seed: int):
|
||||
"""Generate a single terrain patch and return a merged trimesh (or raise)."""
|
||||
preset_fn = ALL_TERRAIN_PRESETS[preset_name]
|
||||
overrides = _PRESET_OVERRIDES.get(preset_name, {})
|
||||
generator_cfg = TerrainGeneratorCfg(
|
||||
seed=seed,
|
||||
size=PATCH_SIZE,
|
||||
num_rows=1,
|
||||
num_cols=1,
|
||||
border_width=0.0,
|
||||
curriculum=False,
|
||||
# A degenerate range pins the single patch to exactly this difficulty.
|
||||
difficulty_range=(difficulty, difficulty),
|
||||
color_scheme="height",
|
||||
sub_terrains={preset_name: preset_fn(proportion=1.0, **overrides)},
|
||||
)
|
||||
generator = TerrainGenerator(generator_cfg)
|
||||
spec = mujoco.MjSpec()
|
||||
generator.compile(spec)
|
||||
model = spec.compile()
|
||||
|
||||
terrain_body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "terrain")
|
||||
geom_ids = [i for i in range(model.ngeom) if model.geom_bodyid[i] == terrain_body_id]
|
||||
return merge_geoms(model, geom_ids)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--port", type=int, default=8080, help="Port for the viser server."
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
server = viser.ViserServer(port=args.port)
|
||||
preset_names = sorted(ALL_TERRAIN_PRESETS)
|
||||
|
||||
terrain_dropdown = server.gui.add_dropdown(
|
||||
"Terrain", options=preset_names, initial_value=preset_names[0]
|
||||
)
|
||||
difficulty_slider = server.gui.add_slider(
|
||||
"Difficulty", min=0.0, max=1.0, step=0.01, initial_value=0.0
|
||||
)
|
||||
seed_input = server.gui.add_number("Seed", initial_value=42, step=1)
|
||||
status = server.gui.add_markdown("**Status:** ready")
|
||||
|
||||
handle: viser.SceneNodeHandle | None = None
|
||||
|
||||
def update() -> None:
|
||||
nonlocal handle
|
||||
name = terrain_dropdown.value
|
||||
difficulty = float(difficulty_slider.value)
|
||||
seed = int(seed_input.value)
|
||||
status.content = f"**Status:** building `{name}` at difficulty {difficulty:.2f}..."
|
||||
try:
|
||||
mesh = _build_terrain_mesh(name, difficulty, seed)
|
||||
except Exception as e: # noqa: BLE001 - surface any generation failure in the UI.
|
||||
status.content = f"**Error:** {type(e).__name__}: {e}"
|
||||
print(f"Failed to build {name} at difficulty {difficulty}: {e}")
|
||||
return
|
||||
if handle is not None:
|
||||
handle.remove()
|
||||
handle = server.scene.add_mesh_trimesh("/terrain", mesh)
|
||||
status.content = (
|
||||
f"**Loaded** `{name}` at difficulty {difficulty:.2f} ({len(mesh.faces):,} faces)"
|
||||
)
|
||||
|
||||
terrain_dropdown.on_update(lambda _: update())
|
||||
difficulty_slider.on_update(lambda _: update())
|
||||
seed_input.on_update(lambda _: update())
|
||||
|
||||
# Top-down-ish initial camera.
|
||||
@server.on_client_connect
|
||||
def _(client: viser.ClientHandle) -> None:
|
||||
client.camera.position = np.array([10.0, 10.0, 8.0])
|
||||
client.camera.look_at = np.array([0.0, 0.0, 0.0])
|
||||
|
||||
update()
|
||||
while True:
|
||||
time.sleep(1.0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,5 +1,15 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
# Default to EGL for GPU-accelerated offscreen rendering on Linux. Must be set
|
||||
# before any mujoco import: mujoco's gl_context module captures MUJOCO_GL once
|
||||
# at load time. Override with e.g. MUJOCO_GL=osmesa on clusters without EGL.
|
||||
# Linux-only because mujoco's gl_context rejects "egl" on macOS/Windows and
|
||||
# raises at import. On those platforms we leave MUJOCO_GL alone so mujoco
|
||||
# defaults to GLFW.
|
||||
if sys.platform.startswith("linux"):
|
||||
os.environ.setdefault("MUJOCO_GL", "egl")
|
||||
|
||||
import traceback
|
||||
from importlib.metadata import entry_points
|
||||
from pathlib import Path
|
||||
|
||||
@@ -4,6 +4,12 @@ from mjlab.actuator.actuator import Actuator as Actuator
|
||||
from mjlab.actuator.actuator import ActuatorCfg as ActuatorCfg
|
||||
from mjlab.actuator.actuator import ActuatorCmd as ActuatorCmd
|
||||
from mjlab.actuator.actuator import CommandField as CommandField
|
||||
from mjlab.actuator.builtin_actuator import (
|
||||
BuiltinDcMotorActuator as BuiltinDcMotorActuator,
|
||||
)
|
||||
from mjlab.actuator.builtin_actuator import (
|
||||
BuiltinDcMotorActuatorCfg as BuiltinDcMotorActuatorCfg,
|
||||
)
|
||||
from mjlab.actuator.builtin_actuator import (
|
||||
BuiltinMotorActuator as BuiltinMotorActuator,
|
||||
)
|
||||
@@ -16,6 +22,12 @@ from mjlab.actuator.builtin_actuator import (
|
||||
from mjlab.actuator.builtin_actuator import (
|
||||
BuiltinMuscleActuatorCfg as BuiltinMuscleActuatorCfg,
|
||||
)
|
||||
from mjlab.actuator.builtin_actuator import (
|
||||
BuiltinPdActuator as BuiltinPdActuator,
|
||||
)
|
||||
from mjlab.actuator.builtin_actuator import (
|
||||
BuiltinPdActuatorCfg as BuiltinPdActuatorCfg,
|
||||
)
|
||||
from mjlab.actuator.builtin_actuator import (
|
||||
BuiltinPositionActuator as BuiltinPositionActuator,
|
||||
)
|
||||
@@ -28,6 +40,15 @@ from mjlab.actuator.builtin_actuator import (
|
||||
from mjlab.actuator.builtin_actuator import (
|
||||
BuiltinVelocityActuatorCfg as BuiltinVelocityActuatorCfg,
|
||||
)
|
||||
from mjlab.actuator.builtin_actuator import (
|
||||
DcMotorDatasheetParams as DcMotorDatasheetParams,
|
||||
)
|
||||
from mjlab.actuator.builtin_actuator import (
|
||||
DcMotorInputMode as DcMotorInputMode,
|
||||
)
|
||||
from mjlab.actuator.builtin_actuator import (
|
||||
DcMotorPhysicalParams as DcMotorPhysicalParams,
|
||||
)
|
||||
from mjlab.actuator.builtin_group import BuiltinActuatorGroup as BuiltinActuatorGroup
|
||||
from mjlab.actuator.dc_actuator import DcMotorActuator as DcMotorActuator
|
||||
from mjlab.actuator.dc_actuator import DcMotorActuatorCfg as DcMotorActuatorCfg
|
||||
|
||||
@@ -174,15 +174,6 @@ class Actuator(ABC, Generic[ActuatorCfgT]):
|
||||
"""Whether this actuator has delay configured."""
|
||||
return self.cfg.delay_max_lag > 0
|
||||
|
||||
@property
|
||||
def command_field(self) -> CommandField | None:
|
||||
"""The primary command field this actuator consumes.
|
||||
|
||||
Returns None by default. Subclasses should override to return the
|
||||
appropriate field.
|
||||
"""
|
||||
return None
|
||||
|
||||
@property
|
||||
def target_ids(self) -> torch.Tensor:
|
||||
"""Local indices of targets controlled by this actuator."""
|
||||
@@ -271,11 +262,6 @@ class Actuator(ABC, Generic[ActuatorCfgT]):
|
||||
"""Create delay buffer. Called during initialize()."""
|
||||
if not self.has_delay:
|
||||
return
|
||||
if self.command_field is None:
|
||||
raise ValueError(
|
||||
f"{self.__class__.__name__}: delay is configured (delay_max_lag="
|
||||
f"{self.cfg.delay_max_lag}) but command_field is not defined."
|
||||
)
|
||||
self._delay_buffer = DelayBuffer(
|
||||
min_lag=self.cfg.delay_min_lag,
|
||||
max_lag=self.cfg.delay_max_lag,
|
||||
@@ -287,19 +273,25 @@ class Actuator(ABC, Generic[ActuatorCfgT]):
|
||||
)
|
||||
|
||||
def apply_delay(self, cmd: ActuatorCmd) -> ActuatorCmd:
|
||||
"""Apply delay to the command_field target. No-op without delay."""
|
||||
"""Delay all command targets with one shared lag. No-op without delay.
|
||||
|
||||
Every target the policy issues (position, velocity, effort) travels the same
|
||||
command channel and experiences the same latency, so they are stacked and
|
||||
delayed together. Feedback fields (``pos``, ``vel``) are never delayed.
|
||||
"""
|
||||
if self._delay_buffer is None:
|
||||
return cmd
|
||||
cf = self.command_field
|
||||
if cf == "position":
|
||||
self._delay_buffer.append(cmd.position_target)
|
||||
return dataclasses.replace(cmd, position_target=self._delay_buffer.compute())
|
||||
elif cf == "velocity":
|
||||
self._delay_buffer.append(cmd.velocity_target)
|
||||
return dataclasses.replace(cmd, velocity_target=self._delay_buffer.compute())
|
||||
else:
|
||||
self._delay_buffer.append(cmd.effort_target)
|
||||
return dataclasses.replace(cmd, effort_target=self._delay_buffer.compute())
|
||||
targets = torch.stack(
|
||||
(cmd.position_target, cmd.velocity_target, cmd.effort_target), dim=-1
|
||||
)
|
||||
self._delay_buffer.append(targets)
|
||||
delayed = self._delay_buffer.compute()
|
||||
return dataclasses.replace(
|
||||
cmd,
|
||||
position_target=delayed[..., 0],
|
||||
velocity_target=delayed[..., 1],
|
||||
effort_target=delayed[..., 2],
|
||||
)
|
||||
|
||||
def set_lags(
|
||||
self,
|
||||
|
||||
@@ -7,19 +7,21 @@ created programmatically via the MjSpec API.
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from enum import IntEnum
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from mjlab.actuator.actuator import (
|
||||
Actuator,
|
||||
ActuatorCfg,
|
||||
ActuatorCmd,
|
||||
CommandField,
|
||||
TransmissionType,
|
||||
)
|
||||
from mjlab.utils.spec import (
|
||||
apply_target_overrides,
|
||||
create_motor_actuator,
|
||||
create_muscle_actuator,
|
||||
create_position_actuator,
|
||||
@@ -63,10 +65,6 @@ class BuiltinPositionActuatorCfg(ActuatorCfg):
|
||||
class BuiltinPositionActuator(Actuator[BuiltinPositionActuatorCfg]):
|
||||
"""MuJoCo built-in position actuator."""
|
||||
|
||||
@property
|
||||
def command_field(self) -> CommandField:
|
||||
return "position"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cfg: BuiltinPositionActuatorCfg,
|
||||
@@ -96,6 +94,102 @@ class BuiltinPositionActuator(Actuator[BuiltinPositionActuatorCfg]):
|
||||
return cmd.position_target
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class BuiltinPdActuatorCfg(ActuatorCfg):
|
||||
"""Implicit-integration version of IdealPdActuator.
|
||||
|
||||
Both consume a position target and a velocity target with kp/kd gains. The
|
||||
difference is in how the PD is delivered to MuJoCo: IdealPdActuator computes
|
||||
the PD force in Python and feeds it to a ``<motor>`` element, which MuJoCo
|
||||
sees as an opaque external force. This actuator expresses the PD as native
|
||||
MuJoCo elements (a ``<position>`` carrying kp, a ``<velocity>`` carrying kd),
|
||||
so the implicit and implicitfast integrators include the kp/kd derivatives
|
||||
in their velocity update. That makes the actuator numerically stable at
|
||||
gain/timestep combinations where explicit Python PD would diverge, which
|
||||
matters when you want to run a real motor's stiff on-board PD gains in sim.
|
||||
"""
|
||||
|
||||
stiffness: float
|
||||
"""Proportional gain (kp)."""
|
||||
damping: float
|
||||
"""Derivative gain (kd)."""
|
||||
effort_limit: float | None = None
|
||||
"""Maximum total torque applied to the joint or tendon. Enforced as a
|
||||
sum-clamp on the two PD terms via jnt_actfrcrange (JOINT) or
|
||||
tendon_actfrcrange (TENDON). None leaves the limit unset."""
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
super().__post_init__()
|
||||
if self.transmission_type == TransmissionType.SITE:
|
||||
raise ValueError(
|
||||
"BuiltinPdActuatorCfg does not support SITE transmission. "
|
||||
"Use BuiltinMotorActuatorCfg for site transmission."
|
||||
)
|
||||
|
||||
def build(
|
||||
self, entity: Entity, target_ids: list[int], target_names: list[str]
|
||||
) -> BuiltinPdActuator:
|
||||
return BuiltinPdActuator(self, entity, target_ids, target_names)
|
||||
|
||||
|
||||
class BuiltinPdActuator(Actuator[BuiltinPdActuatorCfg]):
|
||||
"""MuJoCo native PD: paired <position> + <velocity> elements per target."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cfg: BuiltinPdActuatorCfg,
|
||||
entity: Entity,
|
||||
target_ids: list[int],
|
||||
target_names: list[str],
|
||||
) -> None:
|
||||
super().__init__(cfg, entity, target_ids, target_names)
|
||||
|
||||
@property
|
||||
def num_targets(self) -> int:
|
||||
"""Number of targets. ``ctrl_ids`` is laid out as ``[pos..., vel...]``,
|
||||
each block of length ``num_targets``."""
|
||||
return len(self._target_ids_list)
|
||||
|
||||
def edit_spec(self, spec: mujoco.MjSpec, target_names: list[str]) -> None:
|
||||
# Position elements first, then velocity elements, so ctrl_ids is laid out
|
||||
# as [pos_0..pos_{N-1}, vel_0..vel_{N-1}].
|
||||
for target_name in target_names:
|
||||
pos_act = create_position_actuator(
|
||||
spec,
|
||||
target_name,
|
||||
actuator_name=f"{target_name}_pd_pos",
|
||||
stiffness=self.cfg.stiffness,
|
||||
damping=0.0, # damping lives on the <velocity> element.
|
||||
armature=self.cfg.armature,
|
||||
frictionloss=self.cfg.frictionloss,
|
||||
viscous_damping=self.cfg.viscous_damping,
|
||||
transmission_type=self.cfg.transmission_type,
|
||||
)
|
||||
self._mjs_actuators.append(pos_act)
|
||||
for target_name in target_names:
|
||||
vel_act = create_velocity_actuator(
|
||||
spec,
|
||||
target_name,
|
||||
actuator_name=f"{target_name}_pd_vel",
|
||||
damping=self.cfg.damping,
|
||||
transmission_type=self.cfg.transmission_type,
|
||||
)
|
||||
self._mjs_actuators.append(vel_act)
|
||||
# Effort limit: sum-clamp on the joint/tendon, not on each element.
|
||||
if self.cfg.effort_limit is not None:
|
||||
lim = self.cfg.effort_limit
|
||||
for target_name in target_names:
|
||||
if self.cfg.transmission_type == TransmissionType.JOINT:
|
||||
target = spec.joint(target_name)
|
||||
else:
|
||||
target = spec.tendon(target_name)
|
||||
target.actfrclimited = mujoco.mjtLimited.mjLIMITED_TRUE
|
||||
target.actfrcrange[:] = np.array([-lim, lim])
|
||||
|
||||
def compute(self, cmd: ActuatorCmd) -> torch.Tensor:
|
||||
return torch.cat((cmd.position_target, cmd.velocity_target), dim=1)
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class BuiltinMotorActuatorCfg(ActuatorCfg):
|
||||
"""Configuration for MuJoCo built-in motor actuator.
|
||||
@@ -119,10 +213,6 @@ class BuiltinMotorActuatorCfg(ActuatorCfg):
|
||||
class BuiltinMotorActuator(Actuator[BuiltinMotorActuatorCfg]):
|
||||
"""MuJoCo built-in motor actuator."""
|
||||
|
||||
@property
|
||||
def command_field(self) -> CommandField:
|
||||
return "effort"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cfg: BuiltinMotorActuatorCfg,
|
||||
@@ -151,6 +241,248 @@ class BuiltinMotorActuator(Actuator[BuiltinMotorActuatorCfg]):
|
||||
return cmd.effort_target
|
||||
|
||||
|
||||
def _or_zeros(t: tuple[float, ...] | None, n: int) -> list[float]:
|
||||
return list(t) if t is not None else [0.0] * n
|
||||
|
||||
|
||||
class DcMotorInputMode(IntEnum):
|
||||
"""What the ``ctrl`` signal of a ``<dcmotor>`` represents.
|
||||
|
||||
Values match MuJoCo's enum, consumed by mjs_setToDCMotor and read as gainprm[8].
|
||||
"""
|
||||
|
||||
VOLTAGE = 0
|
||||
POSITION = 1
|
||||
VELOCITY = 2
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DcMotorDatasheetParams:
|
||||
"""Datasheet characterization of a DC motor."""
|
||||
|
||||
nominal_voltage: float
|
||||
"""Nominal (rated) voltage V_n [V]."""
|
||||
stall_torque: float
|
||||
"""Stall torque tau_stall at V_n [N*m]."""
|
||||
no_load_speed: float
|
||||
"""No-load angular velocity omega_no_load at V_n [rad/s]."""
|
||||
|
||||
def _pack(self) -> tuple[list[float], float, list[float]]:
|
||||
"""Returns (motorconst, resistance, nominal) for set_to_dcmotor."""
|
||||
return (
|
||||
[0.0, 0.0],
|
||||
0.0,
|
||||
[self.nominal_voltage, self.stall_torque, self.no_load_speed],
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DcMotorPhysicalParams:
|
||||
"""Physical characterization of a DC motor."""
|
||||
|
||||
kt: float
|
||||
"""Torque constant [N*m/A]."""
|
||||
ke: float
|
||||
"""Back-EMF constant [V*s/rad]."""
|
||||
resistance: float
|
||||
"""Terminal resistance R [Ohm]."""
|
||||
|
||||
def _pack(self) -> tuple[list[float], float, list[float]]:
|
||||
"""Returns (motorconst, resistance, nominal) for set_to_dcmotor."""
|
||||
return [self.kt, self.ke], self.resistance, [0.0, 0.0, 0.0]
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class BuiltinDcMotorActuatorCfg(ActuatorCfg):
|
||||
"""Native MuJoCo ``<dcmotor>`` wrapper.
|
||||
|
||||
Models a DC motor: torque is derived from voltage via the motor constant K and
|
||||
back-EMF, tau = K * (V - K * omega) / R. The back-EMF term lives in biasprm, so
|
||||
MuJoCo's implicit / implicitfast integrators pick up its velocity derivative as
|
||||
effective damping.
|
||||
|
||||
Three input modes select what ctrl carries:
|
||||
|
||||
* VOLTAGE: ctrl is the drive voltage. cmd.effort_target carries volts, not torque.
|
||||
* POSITION / VELOCITY: an internal PID closes on the setpoint and the motor produces
|
||||
torque from its (Vmax-clamped) voltage output.
|
||||
|
||||
Motor characterization: pass either DcMotorDatasheetParams or DcMotorPhysicalParams
|
||||
as motor_params. mjs_setToDCMotor derives K and R (including the viscous-damping
|
||||
correction) and packs the generic gainprm / biasprm / dynprm slots.
|
||||
|
||||
Optional extensions, off by default: integral_gain / integral_limit, slew_rate,
|
||||
inductance / electrical_time_constant, thermal, lugre, cogging.
|
||||
|
||||
dr.pd_gains randomizes only kp and kd; for DR over the extensions, write directly to
|
||||
actuator_gainprm or actuator_dynprm.
|
||||
"""
|
||||
|
||||
motor_params: DcMotorDatasheetParams | DcMotorPhysicalParams
|
||||
"""Motor characterization. Datasheet form: (V_n, tau_stall, omega_no_load).
|
||||
Physical form: (Kt, Ke, R)."""
|
||||
|
||||
mode: DcMotorInputMode = DcMotorInputMode.POSITION
|
||||
"""ctrl input semantics. See class docstring."""
|
||||
|
||||
stiffness: float = 0.0
|
||||
"""PID proportional gain kp. Required in POSITION / VELOCITY mode; must be
|
||||
0 in VOLTAGE mode."""
|
||||
|
||||
damping: float = 0.0
|
||||
"""PID derivative gain kd. Used in POSITION / VELOCITY mode; must be 0 in
|
||||
VOLTAGE mode."""
|
||||
|
||||
voltage_limit: float = 0.0
|
||||
"""Max drive voltage Vmax. Required in POSITION / VELOCITY mode (clamps the
|
||||
PID output). In VOLTAGE mode it is an optional clamp on ctrl; 0 disables."""
|
||||
|
||||
integral_gain: float = 0.0
|
||||
"""PID integral gain ki. In position mode the integrator tracks
|
||||
ki * integral(target - q); in velocity mode, ki * (integral(target) - q).
|
||||
Must be 0 in VOLTAGE mode."""
|
||||
|
||||
integral_limit: float = 0.0
|
||||
"""Anti-windup clamp Imax on the integrator state. 0 disables (the
|
||||
integrator can run away)."""
|
||||
|
||||
slew_rate: float = 0.0
|
||||
"""Max rate of change of ctrl per second. 0 disables."""
|
||||
|
||||
effort_limit: float | None = None
|
||||
"""Continuous torque cap [N*m]. Sets actuator_forcerange. None leaves the
|
||||
per-element forcerange unset."""
|
||||
|
||||
gear: float = 1.0
|
||||
"""Mechanical gear ratio."""
|
||||
|
||||
inductance: float = 0.0
|
||||
"""Winding inductance L [H]. Enables first-order electrical dynamics on the
|
||||
motor current. MuJoCo internally uses te = L / R; pass
|
||||
electrical_time_constant directly to skip the divide. 0 disables."""
|
||||
|
||||
electrical_time_constant: float = 0.0
|
||||
"""Alternative to inductance: specify te [s] directly. Ignored if
|
||||
inductance > 0. 0 disables."""
|
||||
|
||||
thermal: tuple[float, float, float, float, float, float] | None = None
|
||||
"""Thermal model (R_thermal, C_thermal, tau_thermal, alpha, T0, T_ambient).
|
||||
See MuJoCo's ``<dcmotor thermal=...>`` reference for units and which of the
|
||||
first three may be underspecified. Effective resistance becomes
|
||||
R * (1 + alpha * (T + T_ambient - T0)). None disables."""
|
||||
|
||||
cogging: tuple[float, float, float] | None = None
|
||||
"""Cogging torque (amplitude, periodicity, phase) in (N*m, cycles per unit
|
||||
length, rad). Models magnetic torque ripple from rotor-stator interaction;
|
||||
at joint angle q the contribution is amplitude * sin(periodicity * q + phase).
|
||||
|
||||
Added *after* effort_limit is enforced, matching MuJoCo's physical model:
|
||||
effort_limit bounds the electromagnetic torque (the current limit), not the
|
||||
mechanical torque. Total joint torque can exceed effort_limit by up to
|
||||
amplitude. None disables."""
|
||||
|
||||
lugre: tuple[float, float, float, float, float] | None = None
|
||||
"""LuGre friction (sigma0, sigma1, F_Coulomb, F_Stribeck, v_Stribeck).
|
||||
Stick-slip friction with bristle-deflection state. Subtracted from joint
|
||||
torque after the effort_limit clamp (mechanical, like cogging). None
|
||||
disables."""
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
super().__post_init__()
|
||||
if self.transmission_type == TransmissionType.SITE:
|
||||
raise ValueError(
|
||||
"BuiltinDcMotorActuatorCfg does not support SITE transmission. "
|
||||
"Use BuiltinMotorActuatorCfg for site transmission."
|
||||
)
|
||||
|
||||
if self.mode in (DcMotorInputMode.POSITION, DcMotorInputMode.VELOCITY):
|
||||
if self.stiffness <= 0.0:
|
||||
raise ValueError(f"{self.mode.name} mode requires stiffness > 0.")
|
||||
if self.voltage_limit <= 0.0:
|
||||
raise ValueError(f"{self.mode.name} mode requires voltage_limit > 0.")
|
||||
else:
|
||||
if self.stiffness != 0.0 or self.damping != 0.0 or self.integral_gain != 0.0:
|
||||
raise ValueError(
|
||||
"stiffness, damping, and integral_gain are unused in VOLTAGE mode."
|
||||
)
|
||||
|
||||
for name in (
|
||||
"integral_gain",
|
||||
"integral_limit",
|
||||
"slew_rate",
|
||||
"inductance",
|
||||
"electrical_time_constant",
|
||||
):
|
||||
if getattr(self, name) < 0.0:
|
||||
raise ValueError(f"{name} must be non-negative.")
|
||||
|
||||
def build(
|
||||
self, entity: Entity, target_ids: list[int], target_names: list[str]
|
||||
) -> BuiltinDcMotorActuator:
|
||||
return BuiltinDcMotorActuator(self, entity, target_ids, target_names)
|
||||
|
||||
|
||||
class BuiltinDcMotorActuator(Actuator[BuiltinDcMotorActuatorCfg]):
|
||||
"""MuJoCo native ``<dcmotor>``: one actuator per target."""
|
||||
|
||||
def edit_spec(self, spec: mujoco.MjSpec, target_names: list[str]) -> None:
|
||||
cfg = self.cfg
|
||||
motorconst, resistance, nominal = cfg.motor_params._pack()
|
||||
saturation = (
|
||||
[cfg.effort_limit, 0.0, 0.0] if cfg.effort_limit is not None else [0.0] * 3
|
||||
)
|
||||
controller = [
|
||||
cfg.stiffness, # kp
|
||||
cfg.integral_gain, # ki
|
||||
cfg.damping, # kd
|
||||
cfg.slew_rate, # slewmax
|
||||
cfg.integral_limit, # Imax (anti-windup)
|
||||
cfg.voltage_limit, # v_max
|
||||
]
|
||||
# SITE is rejected in __post_init__, so only JOINT and TENDON remain.
|
||||
trntype = (
|
||||
mujoco.mjtTrn.mjTRN_JOINT
|
||||
if cfg.transmission_type == TransmissionType.JOINT
|
||||
else mujoco.mjtTrn.mjTRN_TENDON
|
||||
)
|
||||
|
||||
for target_name in target_names:
|
||||
actuator = spec.add_actuator(name=target_name, target=target_name)
|
||||
actuator.trntype = trntype
|
||||
actuator.gear[0] = cfg.gear
|
||||
actuator.set_to_dcmotor(
|
||||
motorconst=motorconst,
|
||||
resistance=resistance,
|
||||
nominal=nominal,
|
||||
saturation=saturation,
|
||||
controller=controller,
|
||||
cogging=_or_zeros(cfg.cogging, 3),
|
||||
inductance=[cfg.inductance, cfg.electrical_time_constant],
|
||||
thermal=_or_zeros(cfg.thermal, 6),
|
||||
lugre=_or_zeros(cfg.lugre, 5),
|
||||
input_mode=cfg.mode,
|
||||
)
|
||||
|
||||
apply_target_overrides(
|
||||
spec,
|
||||
target_name,
|
||||
cfg.transmission_type,
|
||||
armature=cfg.armature,
|
||||
frictionloss=cfg.frictionloss,
|
||||
viscous_damping=cfg.viscous_damping,
|
||||
)
|
||||
|
||||
self._mjs_actuators.append(actuator)
|
||||
|
||||
def compute(self, cmd: ActuatorCmd) -> torch.Tensor:
|
||||
if self.cfg.mode == DcMotorInputMode.POSITION:
|
||||
return cmd.position_target
|
||||
if self.cfg.mode == DcMotorInputMode.VELOCITY:
|
||||
return cmd.velocity_target
|
||||
# voltage mode: ctrl is the drive voltage carried in effort_target.
|
||||
return cmd.effort_target
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class BuiltinVelocityActuatorCfg(ActuatorCfg):
|
||||
"""Configuration for MuJoCo built-in velocity actuator.
|
||||
@@ -182,10 +514,6 @@ class BuiltinVelocityActuatorCfg(ActuatorCfg):
|
||||
class BuiltinVelocityActuator(Actuator[BuiltinVelocityActuatorCfg]):
|
||||
"""MuJoCo built-in velocity actuator."""
|
||||
|
||||
@property
|
||||
def command_field(self) -> CommandField:
|
||||
return "velocity"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cfg: BuiltinVelocityActuatorCfg,
|
||||
@@ -260,10 +588,6 @@ class BuiltinMuscleActuatorCfg(ActuatorCfg):
|
||||
class BuiltinMuscleActuator(Actuator[BuiltinMuscleActuatorCfg]):
|
||||
"""MuJoCo built-in muscle actuator."""
|
||||
|
||||
@property
|
||||
def command_field(self) -> CommandField:
|
||||
return "effort"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cfg: BuiltinMuscleActuatorCfg,
|
||||
|
||||
@@ -33,6 +33,9 @@ class DcMotorActuatorCfg(IdealPdActuatorCfg):
|
||||
Note: effort_limit should be explicitly set to a realistic value for proper
|
||||
motor modeling. Using the default (inf) will trigger a warning. Use
|
||||
IdealPdActuator if unlimited torque is desired.
|
||||
|
||||
For a native MuJoCo ``<dcmotor>`` with back-EMF, voltage saturation, and
|
||||
configurable ``Kt`` / ``Ke`` / ``R``, see ``BuiltinDcMotorActuator``.
|
||||
"""
|
||||
|
||||
saturation_effort: float
|
||||
|
||||
@@ -9,7 +9,7 @@ import mujoco
|
||||
import mujoco_warp as mjwarp
|
||||
import torch
|
||||
|
||||
from mjlab.actuator.actuator import Actuator, ActuatorCfg, ActuatorCmd, CommandField
|
||||
from mjlab.actuator.actuator import Actuator, ActuatorCfg, ActuatorCmd
|
||||
from mjlab.utils.spec import create_motor_actuator
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -38,10 +38,6 @@ class IdealPdActuatorCfg(ActuatorCfg):
|
||||
class IdealPdActuator(Actuator, Generic[IdealPdCfgT]):
|
||||
"""Ideal PD control actuator."""
|
||||
|
||||
@property
|
||||
def command_field(self) -> CommandField:
|
||||
return "position"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
cfg: IdealPdCfgT,
|
||||
|
||||
@@ -102,11 +102,6 @@
|
||||
diaginertia="0.00167218 0.0016161 0.000217621"/>
|
||||
<joint name="left_ankle_roll_joint" axis="1 0 0" range="-0.2618 0.2618"/>
|
||||
<geom class="visual" material="black" mesh="left_ankle_roll_link"/>
|
||||
<!-- <geom name="left_foot1_collision" class="foot_capsule" fromto="0.1 -0.026 -0.025 0.05 -0.027
|
||||
-0.025"/>
|
||||
<geom name="left_foot2_collision" class="foot_capsule" fromto="-0.045 0 -0.015 0.12 0 -0.015"
|
||||
size="0.02"/>
|
||||
<geom name="left_foot3_collision" class="foot_capsule" fromto="0.1 0.026 -0.025 0.05 0.026 -0.025"/> -->
|
||||
<geom name="left_foot1_collision" class="foot_capsule" fromto="0.1 -0.026 -0.025 0.05 -0.027 -0.025"/>
|
||||
<geom name="left_foot2_collision" class="foot_capsule"
|
||||
fromto="-0.044 -0.018 -0.025 0.123 -0.018 -0.025"/>
|
||||
@@ -156,11 +151,6 @@
|
||||
diaginertia="0.00167218 0.0016161 0.000217621"/>
|
||||
<joint name="right_ankle_roll_joint" axis="1 0 0" range="-0.2618 0.2618"/>
|
||||
<geom class="visual" material="black" mesh="right_ankle_roll_link"/>
|
||||
<!-- <geom name="right_foot1_collision" class="foot_capsule" fromto="0.1 -0.026 -0.025 0.05 -0.026
|
||||
-0.025"/>
|
||||
<geom name="right_foot2_collision" class="foot_capsule" fromto="-0.045 0 -0.015 0.12 0 -0.015"
|
||||
size="0.02"/>
|
||||
<geom name="right_foot3_collision" class="foot_capsule" fromto="0.1 0.026 -0.025 0.05 0.026 -0.025"/> -->
|
||||
<geom name="right_foot1_collision" class="foot_capsule" fromto="0.1 -0.026 -0.025 0.05 -0.026 -0.025"/>
|
||||
<geom name="right_foot2_collision" class="foot_capsule"
|
||||
fromto="-0.044 -0.018 -0.025 0.123 -0.018 -0.025"/>
|
||||
@@ -312,6 +302,7 @@
|
||||
<gyro name="imu_ang_vel" site="imu_in_pelvis"/>
|
||||
<velocimeter name="imu_lin_vel" site="imu_in_pelvis"/>
|
||||
<accelerometer name="imu_lin_acc" site="imu_in_pelvis"/>
|
||||
<framezaxis name="imu_upvector" objtype="body" objname="world" reftype="site" refname="imu_in_pelvis"/>
|
||||
<subtreeangmom name="root_angmom" body="pelvis"/>
|
||||
</sensor>
|
||||
</mujoco>
|
||||
|
||||
@@ -165,6 +165,7 @@
|
||||
<gyro name="imu_ang_vel" site="imu"/>
|
||||
<velocimeter name="imu_lin_vel" site="imu"/>
|
||||
<accelerometer name="imu_lin_acc" site="imu"/>
|
||||
<framezaxis name="imu_upvector" objtype="body" objname="world" reftype="site" refname="imu"/>
|
||||
<subtreeangmom name="root_angmom" body="trunk"/>
|
||||
</sensor>
|
||||
</mujoco>
|
||||
|
||||
@@ -3,6 +3,5 @@ from mjlab.entity.entity import Entity as Entity
|
||||
from mjlab.entity.entity import EntityArticulationInfoCfg as EntityArticulationInfoCfg
|
||||
from mjlab.entity.entity import EntityCfg as EntityCfg
|
||||
from mjlab.entity.entity import EntityIndexing as EntityIndexing
|
||||
from mjlab.entity.entity import VariantCfg as VariantCfg
|
||||
from mjlab.entity.entity import VariantEntityCfg as VariantEntityCfg
|
||||
from mjlab.entity.entity import VariantMetadata as VariantMetadata
|
||||
from mjlab.entity.variants import VariantEntityCfg as VariantEntityCfg
|
||||
from mjlab.entity.variants import VariantMetadata as VariantMetadata
|
||||
|
||||
@@ -3,7 +3,7 @@ from __future__ import annotations
|
||||
import warnings
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Callable, Sequence
|
||||
from typing import TYPE_CHECKING, Callable, Sequence
|
||||
|
||||
import mujoco
|
||||
import mujoco_warp as mjwarp
|
||||
@@ -18,14 +18,13 @@ from mjlab.entity.data import EntityData
|
||||
from mjlab.utils import spec_config as spec_cfg
|
||||
from mjlab.utils.lab_api.string import resolve_matching_names
|
||||
from mjlab.utils.mujoco import dof_width, qpos_width
|
||||
from mjlab.utils.spec import (
|
||||
auto_wrap_fixed_base_mocap,
|
||||
copy_mesh_data,
|
||||
validate_variant_structure,
|
||||
)
|
||||
from mjlab.utils.spec import auto_wrap_fixed_base_mocap
|
||||
from mjlab.utils.string import resolve_expr
|
||||
from mjlab.utils.xml import fix_spec_xml, strip_buffer_textures
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mjlab.entity.variants import VariantMetadata
|
||||
|
||||
|
||||
@dataclass(frozen=False)
|
||||
class EntityIndexing:
|
||||
@@ -67,79 +66,6 @@ class EntityIndexing:
|
||||
return self.bodies[0].id
|
||||
|
||||
|
||||
@dataclass
|
||||
class VariantCfg:
|
||||
"""One object variant for per-world mesh randomization.
|
||||
|
||||
Each variant provides a ``spec_fn`` that returns an MjSpec for one object.
|
||||
The ``weight`` controls what fraction of worlds use this variant.
|
||||
"""
|
||||
|
||||
spec_fn: Callable[[], mujoco.MjSpec]
|
||||
weight: float = 1.0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BodyInertialMetadata:
|
||||
"""Explicit inertial properties for one body in a mesh variant."""
|
||||
|
||||
body_name: str
|
||||
mass: float
|
||||
ipos: tuple[float, float, float]
|
||||
inertia: tuple[float, float, float]
|
||||
iquat: tuple[float, float, float, float]
|
||||
|
||||
|
||||
@dataclass
|
||||
class VariantMetadata:
|
||||
"""Bookkeeping produced by Entity when merging variant specs."""
|
||||
|
||||
variant_names: tuple[str, ...]
|
||||
variant_weights: tuple[float, ...]
|
||||
# Per-variant ordered mesh names for each geom slot. Shorter variants
|
||||
# have None for padding slots that should be disabled (dataid = -1).
|
||||
variant_mesh_names: tuple[tuple[str | None, ...], ...]
|
||||
num_mesh_geoms: int # Max mesh geom count after padding.
|
||||
# Per-variant explicit body inertials. Names are local to the variant spec;
|
||||
# build_mesh_variant_model prefixes them with the scene entity name when
|
||||
# applying them.
|
||||
variant_body_inertials: tuple[tuple[BodyInertialMetadata, ...], ...] = ()
|
||||
|
||||
|
||||
def _iter_body_tree(body: mujoco.MjsBody):
|
||||
yield body
|
||||
for child in body.bodies:
|
||||
yield from _iter_body_tree(child)
|
||||
|
||||
|
||||
def _collect_explicit_body_inertials(
|
||||
root_body: mujoco.MjsBody,
|
||||
) -> tuple[BodyInertialMetadata, ...]:
|
||||
inertials: list[BodyInertialMetadata] = []
|
||||
for body in _iter_body_tree(root_body):
|
||||
if not body.name or not body.explicitinertial:
|
||||
continue
|
||||
inertials.append(
|
||||
BodyInertialMetadata(
|
||||
body_name=body.name,
|
||||
mass=float(body.mass),
|
||||
ipos=(float(body.ipos[0]), float(body.ipos[1]), float(body.ipos[2])),
|
||||
inertia=(
|
||||
float(body.inertia[0]),
|
||||
float(body.inertia[1]),
|
||||
float(body.inertia[2]),
|
||||
),
|
||||
iquat=(
|
||||
float(body.iquat[0]),
|
||||
float(body.iquat[1]),
|
||||
float(body.iquat[2]),
|
||||
float(body.iquat[3]),
|
||||
),
|
||||
)
|
||||
)
|
||||
return tuple(inertials)
|
||||
|
||||
|
||||
@dataclass
|
||||
class EntityCfg:
|
||||
@dataclass
|
||||
@@ -187,50 +113,6 @@ class EntityArticulationInfoCfg:
|
||||
soft_joint_pos_limit_factor: float = 1.0
|
||||
|
||||
|
||||
def _variant_spec_fn_unset() -> mujoco.MjSpec:
|
||||
"""Sentinel default for ``VariantEntityCfg.spec_fn``.
|
||||
|
||||
``VariantEntityCfg`` builds its spec from ``variants`` via
|
||||
``Entity._build_merged_spec``; the inherited ``spec_fn`` field is unused.
|
||||
Identity comparison against this sentinel detects accidental user overrides.
|
||||
"""
|
||||
raise AssertionError(
|
||||
"VariantEntityCfg.spec_fn should never be called; the merged spec is "
|
||||
"built from `variants`."
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class VariantEntityCfg(EntityCfg):
|
||||
"""Entity config for per-world mesh variants.
|
||||
|
||||
Instead of a single ``spec_fn``, provide a dict of named variants.
|
||||
Each world gets a variant assigned proportionally by weight. The
|
||||
merged spec (with all variant meshes and padded geoms) is built
|
||||
automatically.
|
||||
|
||||
All variants must share the same kinematic structure (same bodies,
|
||||
joints, joint types). Only mesh geoms can differ.
|
||||
|
||||
Variant assignment is fixed at ``Simulation`` initialization; it does
|
||||
not resample on episode reset. Pass the per-variant spec via
|
||||
:class:`VariantCfg` rather than setting ``spec_fn`` directly.
|
||||
"""
|
||||
|
||||
variants: dict[str, VariantCfg] = field(default_factory=dict)
|
||||
"""Named mesh variants with weights."""
|
||||
|
||||
spec_fn: Callable[[], mujoco.MjSpec] = field(default=_variant_spec_fn_unset)
|
||||
"""Unused on ``VariantEntityCfg``; the merged spec is built from ``variants``."""
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.spec_fn is not _variant_spec_fn_unset:
|
||||
raise ValueError(
|
||||
"VariantEntityCfg.spec_fn cannot be set; pass per-variant specs via "
|
||||
"VariantCfg(spec_fn=...) inside `variants` instead."
|
||||
)
|
||||
|
||||
|
||||
class Entity:
|
||||
"""An entity represents a physical object in the simulation.
|
||||
|
||||
@@ -272,130 +154,13 @@ class Entity:
|
||||
self._add_initial_state_keyframe()
|
||||
|
||||
def _build_spec(self) -> None:
|
||||
from mjlab.entity.variants import VariantEntityCfg, build_merged_variant_spec
|
||||
|
||||
if isinstance(self.cfg, VariantEntityCfg):
|
||||
self._build_merged_spec()
|
||||
self._spec, self._variant_metadata = build_merged_variant_spec(self.cfg)
|
||||
else:
|
||||
self._spec = auto_wrap_fixed_base_mocap(self.cfg.spec_fn)()
|
||||
|
||||
def _build_merged_spec(self) -> None:
|
||||
"""Build a merged spec from multiple variant specs.
|
||||
|
||||
Validates that all variants share the same kinematic structure,
|
||||
merges all mesh assets into a single spec, and pads the body to
|
||||
the max mesh geom count across variants.
|
||||
"""
|
||||
assert isinstance(self.cfg, VariantEntityCfg)
|
||||
variants = self.cfg.variants
|
||||
if not variants:
|
||||
raise ValueError("VariantEntityCfg.variants must contain at least one entry.")
|
||||
|
||||
variant_names: list[str] = []
|
||||
variant_weights: list[float] = []
|
||||
variant_specs: list[mujoco.MjSpec] = []
|
||||
for name, vcfg in variants.items():
|
||||
variant_names.append(name)
|
||||
variant_weights.append(vcfg.weight)
|
||||
variant_specs.append(vcfg.spec_fn())
|
||||
|
||||
# Find root body in each variant.
|
||||
variant_bodies: list[mujoco.MjsBody] = []
|
||||
for i, spec in enumerate(variant_specs):
|
||||
children = list(spec.worldbody.bodies)
|
||||
if len(children) != 1:
|
||||
raise ValueError(
|
||||
f"Variant '{variant_names[i]}' must have exactly one "
|
||||
f"root body under worldbody, got {len(children)}."
|
||||
)
|
||||
variant_bodies.append(children[0])
|
||||
|
||||
validate_variant_structure(variant_names, variant_bodies)
|
||||
|
||||
# Variant entities must be floating-base. Mocap auto-wrap is not applied
|
||||
# for variant entities, so fixed-base variants would silently stack at
|
||||
# the world origin. Variants share joint structure (validated above), so
|
||||
# checking the first is sufficient.
|
||||
ref_joints = list(variant_bodies[0].joints)
|
||||
if not ref_joints or ref_joints[0].type != mujoco.mjtJoint.mjJNT_FREE:
|
||||
raise ValueError(
|
||||
"VariantEntityCfg requires floating-base variants. Each variant's "
|
||||
"root body must declare a free joint via body.add_freejoint(); "
|
||||
"fixed-base variants are not supported."
|
||||
)
|
||||
|
||||
variant_body_inertials = tuple(
|
||||
_collect_explicit_body_inertials(body) for body in variant_bodies
|
||||
)
|
||||
|
||||
# Collect original mesh names per variant BEFORE any renaming.
|
||||
variant_orig_mesh_names: list[list[str]] = []
|
||||
variant_mesh_geom_counts: list[int] = []
|
||||
for body in variant_bodies:
|
||||
orig_names = [
|
||||
g.meshname for g in body.geoms if g.type == mujoco.mjtGeom.mjGEOM_MESH
|
||||
]
|
||||
variant_orig_mesh_names.append(orig_names)
|
||||
variant_mesh_geom_counts.append(len(orig_names))
|
||||
|
||||
max_mesh_geoms = max(variant_mesh_geom_counts)
|
||||
|
||||
# Use first variant as template. Prefix ALL mesh names with
|
||||
# variant name to avoid collisions across variants.
|
||||
template_spec = variant_specs[0]
|
||||
template_body = variant_bodies[0]
|
||||
|
||||
# Rename template meshes first.
|
||||
template_prefix = f"{variant_names[0]}/"
|
||||
old_to_new: dict[str, str] = {}
|
||||
for mesh in template_spec.meshes:
|
||||
new_name = f"{template_prefix}{mesh.name}"
|
||||
old_to_new[mesh.name] = new_name
|
||||
mesh.name = new_name
|
||||
for g in template_body.geoms:
|
||||
if g.meshname in old_to_new:
|
||||
g.meshname = old_to_new[g.meshname]
|
||||
|
||||
# Copy mesh assets from other variants.
|
||||
for i in range(1, len(variant_specs)):
|
||||
prefix = f"{variant_names[i]}/"
|
||||
for mesh in variant_specs[i].meshes:
|
||||
new_mesh = template_spec.add_mesh()
|
||||
new_mesh.name = f"{prefix}{mesh.name}"
|
||||
copy_mesh_data(mesh, new_mesh)
|
||||
|
||||
# Pad body to max mesh geom count.
|
||||
current_count = variant_mesh_geom_counts[0]
|
||||
if max_mesh_geoms > current_count:
|
||||
longest_idx = max(
|
||||
range(len(variant_mesh_geom_counts)),
|
||||
key=lambda j: variant_mesh_geom_counts[j],
|
||||
)
|
||||
longest_prefix = f"{variant_names[longest_idx]}/"
|
||||
longest_names = variant_orig_mesh_names[longest_idx]
|
||||
for k in range(current_count, max_mesh_geoms):
|
||||
geom = template_body.add_geom()
|
||||
geom.type = mujoco.mjtGeom.mjGEOM_MESH
|
||||
geom.meshname = f"{longest_prefix}{longest_names[k]}"
|
||||
geom.contype = 1
|
||||
geom.conaffinity = 1
|
||||
|
||||
# Build variant_mesh_names: use original names with variant prefix.
|
||||
variant_mesh_name_lists: list[tuple[str | None, ...]] = []
|
||||
for i, orig_names in enumerate(variant_orig_mesh_names):
|
||||
prefix = f"{variant_names[i]}/"
|
||||
names: list[str | None] = [f"{prefix}{n}" for n in orig_names]
|
||||
while len(names) < max_mesh_geoms:
|
||||
names.append(None)
|
||||
variant_mesh_name_lists.append(tuple(names))
|
||||
|
||||
self._variant_metadata = VariantMetadata(
|
||||
variant_names=tuple(variant_names),
|
||||
variant_weights=tuple(variant_weights),
|
||||
variant_mesh_names=tuple(variant_mesh_name_lists),
|
||||
num_mesh_geoms=max_mesh_geoms,
|
||||
variant_body_inertials=variant_body_inertials,
|
||||
)
|
||||
self._spec = template_spec
|
||||
|
||||
@property
|
||||
def variant_metadata(self) -> VariantMetadata | None:
|
||||
return self._variant_metadata
|
||||
@@ -404,6 +169,16 @@ class Entity:
|
||||
self._all_joints = self._spec.joints
|
||||
self._free_joint = None
|
||||
self._non_free_joints = tuple(self._all_joints)
|
||||
|
||||
free_joints = [j for j in self._all_joints if j.type == mujoco.mjtJoint.mjJNT_FREE]
|
||||
if len(free_joints) > 1:
|
||||
raise ValueError(
|
||||
f"Entity spec has {len(free_joints)} freejoints. An Entity models a "
|
||||
"single rigid- or articulated-body system with at most one freejoint, "
|
||||
"which serves as its root. Model each detached floating body as its own "
|
||||
"entry in SceneCfg.entities instead."
|
||||
)
|
||||
|
||||
if self._all_joints and self._all_joints[0].type == mujoco.mjtJoint.mjJNT_FREE:
|
||||
self._free_joint = self._all_joints[0]
|
||||
if not self._free_joint.name:
|
||||
|
||||
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.
@@ -1,57 +0,0 @@
|
||||
"""Robot constants and control parameters."""
|
||||
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
|
||||
# Paths
|
||||
REPO_ROOT = Path(__file__).resolve().parents[1]
|
||||
SCENE_XML = REPO_ROOT / "mjcf" / "scene.xml"
|
||||
MJCF_PATH = REPO_ROOT / "mjcf" / "wheelleg.xml"
|
||||
|
||||
# Robot geometry
|
||||
WHEEL_RADIUS = 0.10 # m
|
||||
WHEEL_TRACK = 0.32 # m (left-right distance)
|
||||
ROBOT_MASS = 12.3 # kg
|
||||
MAX_TORQUE = 17.0 # Nm per joint
|
||||
MAX_JOINT_VEL = 13.0 # rad/s
|
||||
|
||||
# Leg link lengths (from MJCF)
|
||||
L_THIGH = 0.25 # m
|
||||
L_CALF = 0.20 # m (to wheel center)
|
||||
|
||||
# Leg names and joint ordering
|
||||
LEG_NAMES = ("fl", "fr", "rl", "rr")
|
||||
LEG_JOINTS = ("hip_abduction_joint", "hip_pitch_joint", "knee_joint")
|
||||
WHEEL_JOINT = "wheel_joint"
|
||||
|
||||
# Default standing pose (from go2w_sim2sim: [0, 0.8, -1.5])
|
||||
DEFAULT_JOINT_ANGLES = {
|
||||
"hip_abduction": 0.0,
|
||||
"hip_pitch": 0.93,
|
||||
"knee": -1.65,
|
||||
}
|
||||
|
||||
# Actuator modes (MJCF native):
|
||||
# Leg joints: position PD (kp=120, kd=8), ctrl = target angle
|
||||
# Wheel joints: velocity (gain=0.5), ctrl = target velocity (rad/s)
|
||||
|
||||
# Control rates
|
||||
SIM_DT = 0.002 # 500 Hz (from scene.xml)
|
||||
CTRL_DT = 0.004 # 250 Hz control loop
|
||||
CTRL_DECIMATION = int(CTRL_DT / SIM_DT)
|
||||
|
||||
# Wheel drive
|
||||
WHEEL_VEL_MAX = 10.0 # rad/s max wheel command
|
||||
|
||||
# Body pose control gains (for height/roll/pitch compensation)
|
||||
KP_HEIGHT = 3.0 # rad/m error → joint angle correction
|
||||
KP_ROLL = 0.5 # compensation gain
|
||||
KP_PITCH = 0.5 # compensation gain
|
||||
|
||||
# Gait parameters
|
||||
GAIT_FREQ = 2.5 # Hz
|
||||
GAIT_DUTY = 0.6 # stance fraction
|
||||
SWING_HEIGHT = 0.06 # m
|
||||
|
||||
# Trot phase offsets: FL/RR in phase, FR/RL in phase
|
||||
PHASE_OFFSETS = {"fl": 0.0, "fr": 0.5, "rl": 0.5, "rr": 0.0}
|
||||
@@ -1,287 +0,0 @@
|
||||
"""Main controller: wheel mode + trot mode for wheeled-legged robot.
|
||||
|
||||
Wheel mode: differential drive + leg posture hold (height/roll/pitch compensation)
|
||||
Trot mode: quadruped gait with wheel-assisted propulsion
|
||||
|
||||
Actuator interface:
|
||||
- Leg joints: ctrl = target angle (PD: kp=60, kd=3)
|
||||
- Wheel joints: ctrl = target velocity in rad/s (gain=2.0)
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
from robot import Robot, RobotState
|
||||
from dynamics import Dynamics
|
||||
from mpc_controller import MPCController
|
||||
from config import (
|
||||
LEG_NAMES, DEFAULT_JOINT_ANGLES, WHEEL_RADIUS, WHEEL_TRACK,
|
||||
WHEEL_VEL_MAX, KP_ROLL, KP_PITCH,
|
||||
GAIT_FREQ, GAIT_DUTY, SWING_HEIGHT, PHASE_OFFSETS,
|
||||
)
|
||||
|
||||
|
||||
class Controller:
|
||||
"""Wheeled-legged robot controller."""
|
||||
|
||||
def __init__(self, robot: Robot):
|
||||
self.robot = robot
|
||||
self.dynamics = Dynamics()
|
||||
|
||||
# User commands
|
||||
self.vel_x = 0.0 # m/s forward
|
||||
self.vel_y = 0.0 # m/s lateral
|
||||
self.yaw_rate = 0.0 # rad/s
|
||||
self.height = 0.33 # m desired body height
|
||||
|
||||
# Mode: "wheel", "trot", or "mpc"
|
||||
self.mode = "wheel"
|
||||
|
||||
# Prone (lie down) state
|
||||
self.prone = False
|
||||
|
||||
# MPC controller
|
||||
self._mpc_ctrl = MPCController(robot)
|
||||
self._mpc_active = False # track torque mode state
|
||||
|
||||
# Gait state
|
||||
self._gait_phase = 0.0
|
||||
|
||||
# Smoothed commands for trot mode (avoid sudden jumps)
|
||||
self._smooth_vx = 0.0
|
||||
self._smooth_vy = 0.0
|
||||
self._smooth_yaw = 0.0
|
||||
|
||||
# Default leg angles
|
||||
self._default_q = np.array([
|
||||
DEFAULT_JOINT_ANGLES["hip_abduction"],
|
||||
DEFAULT_JOINT_ANGLES["hip_pitch"],
|
||||
DEFAULT_JOINT_ANGLES["knee"],
|
||||
])
|
||||
|
||||
# Swing leg memory
|
||||
self._swing_start_foot = {leg: np.zeros(3) for leg in LEG_NAMES}
|
||||
self._last_contact = {leg: True for leg in LEG_NAMES}
|
||||
|
||||
def compute(self, state: RobotState, dt: float) -> tuple[np.ndarray, np.ndarray]:
|
||||
# Smooth all velocity commands (both modes)
|
||||
alpha = min(dt * 3.0, 1.0) # ~0.33s time constant
|
||||
self._smooth_vx += alpha * (self.vel_x - self._smooth_vx)
|
||||
self._smooth_vy += alpha * (self.vel_y - self._smooth_vy)
|
||||
self._smooth_yaw += alpha * (self.yaw_rate - self._smooth_yaw)
|
||||
|
||||
if self.prone:
|
||||
self._ensure_position_mode()
|
||||
return self._prone_mode()
|
||||
if self.mode == "mpc":
|
||||
return self._mpc_mode(state, dt)
|
||||
if self.mode == "wheel":
|
||||
self._ensure_position_mode()
|
||||
return self._wheel_mode(state, dt)
|
||||
else:
|
||||
self._ensure_position_mode()
|
||||
return self._trot_mode(state, dt)
|
||||
|
||||
def _mpc_mode(self, state: RobotState, dt: float):
|
||||
"""MPC locomotion: MIT motor protocol (PD + MPC feedforward torque)."""
|
||||
# Switch to torque mode if not already
|
||||
if not self._mpc_active:
|
||||
self.robot.enable_torque_mode()
|
||||
self._mpc_active = True
|
||||
|
||||
# Sync commands to MPC controller
|
||||
self._mpc_ctrl.vel_x = self.vel_x
|
||||
self._mpc_ctrl.vel_y = self.vel_y
|
||||
self._mpc_ctrl.yaw_rate = self.yaw_rate
|
||||
self._mpc_ctrl.height = self.height
|
||||
|
||||
# Compute and apply (sets ctrl directly via set_ctrl_mit)
|
||||
self._mpc_ctrl.compute(state, dt)
|
||||
# Return dummy - ctrl already set
|
||||
return np.zeros(12), np.zeros(4)
|
||||
|
||||
def _ensure_position_mode(self):
|
||||
"""Switch back to position PD mode if coming from MPC."""
|
||||
if self._mpc_active:
|
||||
self.robot.enable_position_mode()
|
||||
self._mpc_active = False
|
||||
|
||||
def _prone_mode(self):
|
||||
"""Lie down: actual prone pose from real robot."""
|
||||
leg_targets = np.zeros(12)
|
||||
for i, leg in enumerate(LEG_NAMES):
|
||||
side = 1.0 if leg[1] == "l" else -1.0
|
||||
leg_targets[i*3] = side * 0.3 # fl/rl: +0.3, fr/rr: -0.3
|
||||
leg_targets[i*3+1] = 1.5 # hip pitch
|
||||
leg_targets[i*3+2] = -2.65 # knee fully folded
|
||||
return leg_targets, np.zeros(4)
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# WHEEL MODE
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
def _wheel_mode(self, state: RobotState, dt: float):
|
||||
"""Wheel drive + leg posture hold.
|
||||
|
||||
vel_y: limited effect in wheel mode (differential drive cannot produce
|
||||
pure lateral motion). Uses hip_abduction lean for small lateral force.
|
||||
For significant lateral motion, use trot mode.
|
||||
"""
|
||||
wheel_targets = self._differential_drive(self._smooth_vx, self._smooth_yaw)
|
||||
leg_targets = self._posture_control(state)
|
||||
return leg_targets, wheel_targets
|
||||
|
||||
def _posture_control(self, state: RobotState) -> np.ndarray:
|
||||
"""Leg joint targets: table-interpolated height control."""
|
||||
leg_targets = np.zeros(12)
|
||||
|
||||
# Calibrated height→angle lookup (measured from simulation)
|
||||
_H = [0.157, 0.248, 0.311, 0.366, 0.411, 0.448]
|
||||
_HIP = [1.5, 1.2, 1.0, 0.8, 0.6, 0.4]
|
||||
_KNEE = [-2.5, -2.1, -1.8, -1.5, -1.2, -0.9]
|
||||
|
||||
h_clamp = np.clip(self.height, _H[0], _H[-1])
|
||||
q_hip_base = float(np.interp(h_clamp, _H, _HIP))
|
||||
q_knee_base = float(np.interp(h_clamp, _H, _KNEE))
|
||||
|
||||
roll_corr = -KP_ROLL * state.rpy[0]
|
||||
pitch_corr = -KP_PITCH * state.rpy[1]
|
||||
lateral_lean = 0.3 * self.vel_y
|
||||
|
||||
for i, leg in enumerate(LEG_NAMES):
|
||||
side = 1.0 if leg[1] == "l" else -1.0
|
||||
leg_targets[i*3] = np.clip(side * roll_corr + lateral_lean, -0.5, 0.5)
|
||||
leg_targets[i*3+1] = np.clip(q_hip_base + pitch_corr, -1.0, 2.5)
|
||||
leg_targets[i*3+2] = np.clip(q_knee_base, -2.6, -0.3)
|
||||
|
||||
return leg_targets
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# TROT MODE
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
def _trot_mode(self, state: RobotState, dt: float):
|
||||
"""Trot gait with wheel assist."""
|
||||
# Advance gait phase
|
||||
self._gait_phase = (self._gait_phase + dt * GAIT_FREQ) % 1.0
|
||||
|
||||
# Contact state
|
||||
contacts = {}
|
||||
for leg in LEG_NAMES:
|
||||
phase = (self._gait_phase + PHASE_OFFSETS[leg]) % 1.0
|
||||
contacts[leg] = phase < GAIT_DUTY
|
||||
|
||||
# Pinocchio update
|
||||
q_pin, dq_pin = self.robot.get_qpos_qvel_for_pinocchio()
|
||||
self.dynamics.update(q_pin, dq_pin)
|
||||
|
||||
leg_targets = np.zeros(12)
|
||||
wheel_targets = np.zeros(4)
|
||||
|
||||
for i, leg in enumerate(LEG_NAMES):
|
||||
if contacts[leg]:
|
||||
# Stance: posture hold
|
||||
leg_targets[i*3:(i+1)*3] = self._stance_leg_target(state, leg)
|
||||
self._swing_start_foot[leg] = self.dynamics.get_foot_pos(leg)
|
||||
self._last_contact[leg] = True
|
||||
# Wheel: drive with smoothed velocity
|
||||
wheel_targets[i] = self._differential_drive_single(
|
||||
self._smooth_vx, self._smooth_yaw, leg)
|
||||
else:
|
||||
# Swing: IK trajectory
|
||||
swing_phase = self._get_swing_phase(leg)
|
||||
target_foot = self._compute_swing_target(leg, state, swing_phase)
|
||||
q_ik = self.dynamics.inverse_kinematics(leg, target_foot, q_pin)
|
||||
leg_targets[i*3:(i+1)*3] = q_ik
|
||||
self._last_contact[leg] = False
|
||||
# Wheel: zero (free during swing)
|
||||
wheel_targets[i] = 0.0
|
||||
|
||||
return leg_targets, wheel_targets
|
||||
|
||||
def _stance_leg_target(self, state: RobotState, leg: str) -> np.ndarray:
|
||||
"""Stance leg: table-interpolated height + attitude compensation."""
|
||||
_H = [0.157, 0.248, 0.311, 0.366, 0.411, 0.448]
|
||||
_HIP = [1.5, 1.2, 1.0, 0.8, 0.6, 0.4]
|
||||
_KNEE = [-2.5, -2.1, -1.8, -1.5, -1.2, -0.9]
|
||||
|
||||
h_clamp = np.clip(self.height, _H[0], _H[-1])
|
||||
q_hip = float(np.interp(h_clamp, _H, _HIP))
|
||||
q_knee = float(np.interp(h_clamp, _H, _KNEE))
|
||||
|
||||
roll_corr = -KP_ROLL * state.rpy[0]
|
||||
pitch_corr = -KP_PITCH * state.rpy[1]
|
||||
side = 1.0 if leg[1] == "l" else -1.0
|
||||
lateral_lean = 0.3 * self.vel_y
|
||||
|
||||
return np.array([
|
||||
np.clip(side * roll_corr + lateral_lean, -0.5, 0.5),
|
||||
np.clip(q_hip + pitch_corr, -1.0, 2.5),
|
||||
np.clip(q_knee, -2.6, -0.3),
|
||||
])
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# DIFFERENTIAL DRIVE
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
def _differential_drive(self, vel_x: float, yaw_rate: float) -> np.ndarray:
|
||||
"""4 wheel velocities from body commands."""
|
||||
vel_left = (vel_x - 0.5 * WHEEL_TRACK * yaw_rate) / WHEEL_RADIUS
|
||||
vel_right = (vel_x + 0.5 * WHEEL_TRACK * yaw_rate) / WHEEL_RADIUS
|
||||
targets = np.zeros(4)
|
||||
for i, leg in enumerate(LEG_NAMES):
|
||||
targets[i] = vel_left if leg[1] == "l" else vel_right
|
||||
return np.clip(targets, -WHEEL_VEL_MAX, WHEEL_VEL_MAX)
|
||||
|
||||
def _differential_drive_single(self, vel_x: float, yaw_rate: float, leg: str) -> float:
|
||||
if leg[1] == "l":
|
||||
v = (vel_x - 0.5 * WHEEL_TRACK * yaw_rate) / WHEEL_RADIUS
|
||||
else:
|
||||
v = (vel_x + 0.5 * WHEEL_TRACK * yaw_rate) / WHEEL_RADIUS
|
||||
return np.clip(v, -WHEEL_VEL_MAX, WHEEL_VEL_MAX)
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# SWING TRAJECTORY
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
def _get_swing_phase(self, leg: str) -> float:
|
||||
phase = (self._gait_phase + PHASE_OFFSETS[leg]) % 1.0
|
||||
if phase < GAIT_DUTY:
|
||||
return 0.0
|
||||
return (phase - GAIT_DUTY) / (1.0 - GAIT_DUTY)
|
||||
|
||||
def _compute_swing_target(self, leg: str, state: RobotState,
|
||||
swing_phase: float) -> np.ndarray:
|
||||
"""Swing foot target with Raibert heuristic using COMMANDED velocity."""
|
||||
p_start = self._swing_start_foot[leg]
|
||||
p_end = self._compute_touchdown(leg, state)
|
||||
|
||||
s = swing_phase
|
||||
s_mj = 10*s**3 - 15*s**4 + 6*s**5
|
||||
|
||||
pos = p_start + (p_end - p_start) * s_mj
|
||||
|
||||
# Z lift
|
||||
z_lift = 64.0 * s**3 * (1.0 - s)**3
|
||||
pos[2] = p_start[2] + SWING_HEIGHT * z_lift
|
||||
|
||||
return pos
|
||||
|
||||
def _compute_touchdown(self, leg: str, state: RobotState) -> np.ndarray:
|
||||
"""Raibert heuristic using COMMANDED velocity.
|
||||
|
||||
When commands are zero, foot lands at its takeoff position (no net motion).
|
||||
When commands are nonzero, foot placement is offset by commanded velocity.
|
||||
"""
|
||||
# Base: land where the foot took off (zero net displacement)
|
||||
td = self._swing_start_foot[leg].copy()
|
||||
|
||||
# Add commanded velocity offset (Raibert-style)
|
||||
t_stance = (1.0 / GAIT_FREQ) * GAIT_DUTY
|
||||
yaw = state.rpy[2]
|
||||
c, s = np.cos(yaw), np.sin(yaw)
|
||||
R_z = np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]])
|
||||
cmd_vel_world = R_z @ np.array([self._smooth_vx, self._smooth_vy, 0.0])
|
||||
|
||||
td[0] += cmd_vel_world[0] * t_stance * 0.5
|
||||
td[1] += cmd_vel_world[1] * t_stance * 0.5
|
||||
td[2] = WHEEL_RADIUS # ground level
|
||||
return td
|
||||
@@ -1,97 +0,0 @@
|
||||
"""Pinocchio dynamics: FK, Jacobian, IK for the wheeled-legged robot."""
|
||||
|
||||
import numpy as np
|
||||
import pinocchio as pin
|
||||
from config import MJCF_PATH, LEG_NAMES
|
||||
|
||||
# Foot frame names in Pinocchio model (wheel link centers)
|
||||
FOOT_FRAMES = {leg: f"{leg}_wheel_Link" for leg in LEG_NAMES}
|
||||
|
||||
# Leg joint names for each leg
|
||||
_LEG_JOINT_NAMES = {
|
||||
leg: [f"{leg}_{jt}" for jt in ("hip_abduction_joint", "hip_pitch_joint", "knee_joint")]
|
||||
for leg in LEG_NAMES
|
||||
}
|
||||
|
||||
|
||||
class Dynamics:
|
||||
"""Pinocchio-based kinematics/dynamics. Deployable on real hardware."""
|
||||
|
||||
def __init__(self):
|
||||
self.model = pin.buildModelFromMJCF(str(MJCF_PATH))
|
||||
self.data = self.model.createData()
|
||||
|
||||
# Cache frame IDs
|
||||
self._foot_fids = {}
|
||||
for leg, fname in FOOT_FRAMES.items():
|
||||
self._foot_fids[leg] = self.model.getFrameId(fname)
|
||||
|
||||
# Cache joint velocity indices for each leg (3 joints)
|
||||
self._leg_v_indices = {}
|
||||
for leg, jnames in _LEG_JOINT_NAMES.items():
|
||||
indices = []
|
||||
for jn in jnames:
|
||||
jid = self.model.getJointId(jn)
|
||||
indices.append(self.model.joints[jid].idx_v)
|
||||
self._leg_v_indices[leg] = indices
|
||||
|
||||
# Cache joint config indices for each leg
|
||||
self._leg_q_indices = {}
|
||||
for leg, jnames in _LEG_JOINT_NAMES.items():
|
||||
indices = []
|
||||
for jn in jnames:
|
||||
jid = self.model.getJointId(jn)
|
||||
indices.append(self.model.joints[jid].idx_q)
|
||||
self._leg_q_indices[leg] = indices
|
||||
|
||||
def update(self, q: np.ndarray, dq: np.ndarray):
|
||||
"""Forward kinematics + Jacobians.
|
||||
|
||||
Args:
|
||||
q: Pinocchio config (nq=23: pos3, quat_xyzw4, joints16)
|
||||
dq: Pinocchio velocity (nv=22: v_body3, w_body3, joints16)
|
||||
"""
|
||||
pin.forwardKinematics(self.model, self.data, q, dq)
|
||||
pin.updateFramePlacements(self.model, self.data)
|
||||
pin.computeJointJacobians(self.model, self.data, q)
|
||||
|
||||
def get_foot_pos(self, leg: str) -> np.ndarray:
|
||||
"""Foot (wheel center) position in world frame (3,)."""
|
||||
return self.data.oMf[self._foot_fids[leg]].translation.copy()
|
||||
|
||||
def get_foot_jacobian_leg(self, leg: str) -> np.ndarray:
|
||||
"""3x3 linear Jacobian of foot w.r.t. 3 leg joints (world frame)."""
|
||||
fid = self._foot_fids[leg]
|
||||
J_full = pin.getFrameJacobian(
|
||||
self.model, self.data, fid, pin.LOCAL_WORLD_ALIGNED)[:3, :]
|
||||
cols = self._leg_v_indices[leg]
|
||||
return J_full[:, cols]
|
||||
|
||||
def inverse_kinematics(self, leg: str, target_pos: np.ndarray,
|
||||
q_current: np.ndarray, max_iter=30, eps=1e-4) -> np.ndarray:
|
||||
"""Numerical IK for one leg. Returns (3,) joint angles.
|
||||
|
||||
Args:
|
||||
leg: Leg name
|
||||
target_pos: Desired foot position in world frame (3,)
|
||||
q_current: Current full Pinocchio config (nq=23)
|
||||
"""
|
||||
q = q_current.copy()
|
||||
fid = self._foot_fids[leg]
|
||||
q_indices = self._leg_q_indices[leg]
|
||||
|
||||
for _ in range(max_iter):
|
||||
pin.forwardKinematics(self.model, self.data, q)
|
||||
pin.updateFramePlacements(self.model, self.data)
|
||||
err = target_pos - self.data.oMf[fid].translation
|
||||
if np.linalg.norm(err) < eps:
|
||||
break
|
||||
pin.computeJointJacobians(self.model, self.data, q)
|
||||
J = pin.getFrameJacobian(
|
||||
self.model, self.data, fid, pin.LOCAL_WORLD_ALIGNED)[:3, :]
|
||||
J_leg = J[:, self._leg_v_indices[leg]]
|
||||
dq = np.linalg.solve(J_leg.T @ J_leg + 1e-6 * np.eye(3), J_leg.T @ err)
|
||||
for i, idx in enumerate(q_indices):
|
||||
q[idx] += dq[i]
|
||||
|
||||
return np.array([q[idx] for idx in q_indices])
|
||||
@@ -1,127 +0,0 @@
|
||||
"""GUI control panel for the wheeled-legged robot."""
|
||||
|
||||
import tkinter as tk
|
||||
from tkinter import ttk
|
||||
|
||||
|
||||
class GUI:
|
||||
"""Tkinter control panel: sliders + gait buttons + status display."""
|
||||
|
||||
def __init__(self, controller):
|
||||
self.ctrl = controller
|
||||
self.root = tk.Tk()
|
||||
self.root.title("WheelLeg Control")
|
||||
self.root.geometry("400x500")
|
||||
self.root.protocol("WM_DELETE_WINDOW", self._on_close)
|
||||
self._closed = False
|
||||
|
||||
self._build()
|
||||
|
||||
def _build(self):
|
||||
# Mode buttons
|
||||
mf = ttk.LabelFrame(self.root, text="Mode")
|
||||
mf.pack(fill="x", padx=8, pady=4)
|
||||
for mode in ("wheel", "trot", "mpc"):
|
||||
ttk.Button(mf, text=mode.upper(),
|
||||
command=lambda m=mode: self._set_mode(m)
|
||||
).pack(side="left", padx=4, expand=True)
|
||||
ttk.Button(mf, text="PRONE/STAND",
|
||||
command=self._toggle_prone).pack(side="left", padx=4, expand=True)
|
||||
|
||||
# Command sliders
|
||||
cf = ttk.LabelFrame(self.root, text="Commands")
|
||||
cf.pack(fill="x", padx=8, pady=4)
|
||||
|
||||
self.vel_x_var = tk.DoubleVar(value=0.0)
|
||||
self.vel_y_var = tk.DoubleVar(value=0.0)
|
||||
self.yaw_var = tk.DoubleVar(value=0.0)
|
||||
self.height_var = tk.DoubleVar(value=self.ctrl.height)
|
||||
|
||||
self._slider(cf, "Vel X", self.vel_x_var, -1.5, 1.5)
|
||||
self._slider(cf, "Vel Y*", self.vel_y_var, -0.5, 0.5)
|
||||
self._slider(cf, "Yaw", self.yaw_var, -2.0, 2.0)
|
||||
self._slider(cf, "Height", self.height_var, 0.16, 0.45)
|
||||
|
||||
ttk.Label(cf, text="* Vel Y: trot mode only (diff-drive can't sidestep)",
|
||||
font=("", 8)).pack(anchor="w", padx=8)
|
||||
|
||||
ttk.Button(cf, text="Reset", command=self._reset).pack(pady=4)
|
||||
|
||||
# Status display
|
||||
sf = ttk.LabelFrame(self.root, text="Status")
|
||||
sf.pack(fill="both", expand=True, padx=8, pady=4)
|
||||
self.status_text = tk.Text(sf, height=12, width=45, font=("Consolas", 9))
|
||||
self.status_text.pack(fill="both", expand=True, padx=4, pady=4)
|
||||
|
||||
def _slider(self, parent, label, var, lo, hi):
|
||||
f = ttk.Frame(parent)
|
||||
f.pack(fill="x", padx=4, pady=2)
|
||||
ttk.Label(f, text=label, width=7).pack(side="left")
|
||||
ttk.Scale(f, from_=lo, to=hi, variable=var,
|
||||
command=lambda *_: self._sync()).pack(side="left", fill="x", expand=True)
|
||||
lbl = ttk.Label(f, text="0.00", width=6)
|
||||
lbl.pack(side="left")
|
||||
var.trace_add("write", lambda *_, v=var, l=lbl: l.config(text=f"{v.get():.2f}"))
|
||||
|
||||
def _set_mode(self, mode):
|
||||
self.ctrl.mode = mode
|
||||
self.ctrl.prone = False
|
||||
|
||||
def _toggle_prone(self):
|
||||
self.ctrl.prone = not self.ctrl.prone
|
||||
|
||||
def _sync(self):
|
||||
self.ctrl.vel_x = self.vel_x_var.get()
|
||||
self.ctrl.vel_y = self.vel_y_var.get()
|
||||
self.ctrl.yaw_rate = self.yaw_var.get()
|
||||
self.ctrl.height = self.height_var.get()
|
||||
|
||||
def _reset(self):
|
||||
self.vel_x_var.set(0.0)
|
||||
self.vel_y_var.set(0.0)
|
||||
self.yaw_var.set(0.0)
|
||||
self._sync()
|
||||
|
||||
def _on_close(self):
|
||||
self._closed = True
|
||||
self.root.destroy()
|
||||
|
||||
@property
|
||||
def closed(self):
|
||||
return self._closed
|
||||
|
||||
def update_status(self, state, step):
|
||||
"""Update status text with current robot state."""
|
||||
txt = (
|
||||
f"Mode: {self.ctrl.mode} Step: {step}\n"
|
||||
f"Pos: x={state.pos[0]:.3f} y={state.pos[1]:.3f} z={state.pos[2]:.3f}\n"
|
||||
f"RPY: r={np.degrees(state.rpy[0]):.1f}° p={np.degrees(state.rpy[1]):.1f}° "
|
||||
f"y={np.degrees(state.rpy[2]):.1f}°\n"
|
||||
f"Vel: vx={state.lin_vel[0]:.3f} vy={state.lin_vel[1]:.3f} vz={state.lin_vel[2]:.3f}\n"
|
||||
f"Cmd: vx={self.ctrl.vel_x:.2f} yaw={self.ctrl.yaw_rate:.2f} h={self.ctrl.height:.3f}\n"
|
||||
f"─────────────────────────────────\n"
|
||||
)
|
||||
# Joint angles (compact)
|
||||
for i, leg in enumerate(("FL", "FR", "RL", "RR")):
|
||||
q = state.joint_pos[i*4:i*4+3]
|
||||
w = state.joint_vel[i*4+3]
|
||||
txt += f"{leg}: [{q[0]:+.2f} {q[1]:+.2f} {q[2]:+.2f}] w={w:+.1f}\n"
|
||||
|
||||
self.status_text.delete("1.0", tk.END)
|
||||
self.status_text.insert(tk.END, txt)
|
||||
|
||||
def tick(self):
|
||||
"""Process GUI events. Returns False if window closed."""
|
||||
if self._closed:
|
||||
return False
|
||||
try:
|
||||
self.root.update_idletasks()
|
||||
self.root.update()
|
||||
return True
|
||||
except tk.TclError:
|
||||
self._closed = True
|
||||
return False
|
||||
|
||||
|
||||
# Need numpy for degrees conversion in update_status
|
||||
import numpy as np
|
||||
@@ -1,242 +0,0 @@
|
||||
"""Convex MPC solver for wheeled-legged robot.
|
||||
|
||||
Centroidal dynamics: single rigid body model with 4 contact forces.
|
||||
State: x = [pos(3), rpy(3), vel(3), omega(3)] = 12
|
||||
Input: u = [f1(3), f2(3), f3(3), f4(3)] = 12
|
||||
Friction pyramid constraints on each foot.
|
||||
|
||||
Reference: MIT Cheetah 3 Convex MPC (Di Carlo et al.)
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
from scipy import sparse
|
||||
from scipy.linalg import block_diag
|
||||
import osqp
|
||||
|
||||
from config import ROBOT_MASS, LEG_NAMES
|
||||
|
||||
# MPC parameters
|
||||
MPC_HORIZON = 10 # prediction steps
|
||||
MPC_DT = 0.02 # 50 Hz MPC update
|
||||
MU = 0.6 # friction coefficient
|
||||
FZ_MAX = 200.0 # max vertical force per leg
|
||||
FZ_MIN = 10.0 # min vertical force (stance)
|
||||
NX = 12 # state dim
|
||||
NU = 12 # input dim (4 legs × 3D force)
|
||||
|
||||
# Cost weights: [pos_x, pos_y, pos_z, roll, pitch, yaw, vx, vy, vz, wx, wy, wz]
|
||||
Q_WEIGHTS = np.array([2.0, 2.0, 50.0, 50.0, 50.0, 10.0, 2.0, 2.0, 1.0, 1.0, 1.0, 1.0])
|
||||
R_WEIGHTS = np.array([1e-6] * 12)
|
||||
|
||||
|
||||
def _skew(v):
|
||||
return np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
|
||||
|
||||
|
||||
class ConvexMPC:
|
||||
"""Convex MPC: solves QP for optimal ground reaction forces."""
|
||||
|
||||
def __init__(self, mass=ROBOT_MASS, inertia=None):
|
||||
self.mass = mass
|
||||
# Approximate body inertia (diagonal, world-aligned)
|
||||
if inertia is None:
|
||||
self.I_body = np.diag([0.07, 0.26, 0.24])
|
||||
else:
|
||||
self.I_body = np.array(inertia).reshape(3, 3)
|
||||
|
||||
self.N = MPC_HORIZON
|
||||
self.dt = MPC_DT
|
||||
self.Q = np.diag(Q_WEIGHTS)
|
||||
self.R = np.diag(R_WEIGHTS)
|
||||
self.gravity = np.array([0, 0, -9.81])
|
||||
|
||||
self._last_forces = np.zeros(NU)
|
||||
|
||||
def solve(self, x0, x_ref, foot_positions, contact_schedule):
|
||||
"""Solve MPC QP.
|
||||
|
||||
Args:
|
||||
x0: (12,) current state [pos, rpy, vel, omega]
|
||||
x_ref: (12, N) reference trajectory over horizon
|
||||
foot_positions: (4, 3) foot positions in world frame (relative to CoM)
|
||||
contact_schedule: (4, N) binary contact table (1=stance)
|
||||
|
||||
Returns:
|
||||
forces: (12,) optimal forces for current timestep [f1x,f1y,f1z,...,f4x,f4y,f4z]
|
||||
"""
|
||||
N = self.N
|
||||
|
||||
# Build dynamics matrices
|
||||
Ad, Bd_list, gd = self._discretize_dynamics(x0, foot_positions)
|
||||
|
||||
# Build QP: min 0.5 z'Hz + f'z s.t. lb <= Az <= ub, lbx <= z <= ubx
|
||||
# Decision variables: z = [x1,...,xN, u0,...,uN-1]
|
||||
nvars = N * NX + N * NU
|
||||
|
||||
# --- Hessian ---
|
||||
H_diag = np.concatenate([np.tile(2 * Q_WEIGHTS, N), np.tile(2 * R_WEIGHTS, N)])
|
||||
H = sparse.diags(H_diag, format='csc')
|
||||
|
||||
# --- Gradient ---
|
||||
g = np.zeros(nvars)
|
||||
for k in range(N):
|
||||
g[k*NX:(k+1)*NX] = -2 * self.Q @ x_ref[:, k]
|
||||
|
||||
# --- Dynamics equality constraints ---
|
||||
# x_{k+1} = Ad @ x_k + Bd_k @ u_k + gd
|
||||
# Rewrite: x_{k+1} - Ad @ x_k - Bd_k @ u_k = gd (for k>0)
|
||||
# x_1 - Bd_0 @ u_0 = Ad @ x0 + gd (for k=0)
|
||||
n_eq = N * NX
|
||||
A_eq = np.zeros((n_eq, nvars))
|
||||
b_eq = np.zeros(n_eq)
|
||||
|
||||
# k=0: x_1 = Ad @ x0 + Bd_0 @ u_0 + gd
|
||||
A_eq[0:NX, 0:NX] = np.eye(NX) # x_1
|
||||
A_eq[0:NX, N*NX:N*NX+NU] = -Bd_list[0] # -Bd_0 @ u_0
|
||||
b_eq[0:NX] = Ad @ x0 + gd
|
||||
|
||||
for k in range(1, N):
|
||||
row = k * NX
|
||||
# x_{k+1}
|
||||
A_eq[row:row+NX, k*NX:(k+1)*NX] = np.eye(NX)
|
||||
# -Ad @ x_k
|
||||
A_eq[row:row+NX, (k-1)*NX:k*NX] = -Ad
|
||||
# -Bd_k @ u_k
|
||||
A_eq[row:row+NX, N*NX+k*NU:N*NX+(k+1)*NU] = -Bd_list[k]
|
||||
b_eq[row:row+NX] = gd
|
||||
|
||||
# --- Friction pyramid inequality constraints ---
|
||||
# For each stance leg at each timestep: 4 faces
|
||||
# fx - mu*fz <= 0, -fx - mu*fz <= 0, fy - mu*fz <= 0, -fy - mu*fz <= 0
|
||||
n_ineq_max = 4 * 4 * N
|
||||
A_ineq = np.zeros((n_ineq_max, nvars))
|
||||
u_ineq = np.zeros(n_ineq_max)
|
||||
|
||||
row = 0
|
||||
for k in range(N):
|
||||
u_base = N * NX + k * NU
|
||||
for leg in range(4):
|
||||
if contact_schedule[leg, k] == 1:
|
||||
fx_idx = u_base + leg * 3
|
||||
fy_idx = u_base + leg * 3 + 1
|
||||
fz_idx = u_base + leg * 3 + 2
|
||||
|
||||
# Friction pyramid: stance leg
|
||||
A_ineq[row, fx_idx] = 1.0
|
||||
A_ineq[row, fz_idx] = -MU
|
||||
row += 1
|
||||
A_ineq[row, fx_idx] = -1.0
|
||||
A_ineq[row, fz_idx] = -MU
|
||||
row += 1
|
||||
A_ineq[row, fy_idx] = 1.0
|
||||
A_ineq[row, fz_idx] = -MU
|
||||
row += 1
|
||||
A_ineq[row, fy_idx] = -1.0
|
||||
A_ineq[row, fz_idx] = -MU
|
||||
row += 1
|
||||
|
||||
A_ineq = A_ineq[:row]
|
||||
u_ineq = u_ineq[:row]
|
||||
|
||||
# Stack constraints
|
||||
A_full = np.vstack([A_eq, A_ineq])
|
||||
l_full = np.concatenate([b_eq, -np.inf * np.ones(row)])
|
||||
u_full = np.concatenate([b_eq, u_ineq])
|
||||
|
||||
# --- Box constraints on forces (as identity rows in A) ---
|
||||
A_box = np.zeros((N * NU, nvars))
|
||||
l_box = -np.inf * np.ones(N * NU)
|
||||
u_box = np.inf * np.ones(N * NU)
|
||||
|
||||
for k in range(N):
|
||||
u_base = N * NX + k * NU
|
||||
for leg in range(4):
|
||||
idx = u_base + leg * 3
|
||||
box_row = k * NU + leg * 3
|
||||
# Identity rows for fx, fy, fz
|
||||
for j in range(3):
|
||||
A_box[box_row + j, idx + j] = 1.0
|
||||
|
||||
if contact_schedule[leg, k] == 1:
|
||||
# Stance: fz bounded
|
||||
l_box[box_row + 2] = FZ_MIN
|
||||
u_box[box_row + 2] = FZ_MAX
|
||||
else:
|
||||
# Swing: all forces = 0
|
||||
l_box[box_row:box_row+3] = 0.0
|
||||
u_box[box_row:box_row+3] = 0.0
|
||||
|
||||
# Final constraint matrix
|
||||
A_full = np.vstack([A_full, A_box])
|
||||
l_full = np.concatenate([l_full, l_box])
|
||||
u_full = np.concatenate([u_full, u_box])
|
||||
|
||||
# --- Solve with OSQP ---
|
||||
A_sparse = sparse.csc_matrix(A_full)
|
||||
H_sparse = sparse.triu(H, format='csc')
|
||||
|
||||
solver = osqp.OSQP()
|
||||
solver.setup(H_sparse, g, A_sparse, l_full, u_full,
|
||||
eps_abs=1e-4, eps_rel=1e-4,
|
||||
max_iter=500, polish=True, verbose=False,
|
||||
warm_start=True)
|
||||
|
||||
# Warm start with previous solution
|
||||
if self._last_forces is not None:
|
||||
x_warm = np.zeros(nvars)
|
||||
x_warm[N*NX:N*NX+NU] = self._last_forces
|
||||
solver.warm_start(x=x_warm)
|
||||
|
||||
result = solver.solve()
|
||||
|
||||
if result.info.status == 'solved' or result.info.status == 'solved_inaccurate':
|
||||
# Extract first timestep forces
|
||||
forces = result.x[N*NX:N*NX+NU]
|
||||
self._last_forces = forces.copy()
|
||||
else:
|
||||
forces = self._last_forces
|
||||
|
||||
return forces
|
||||
|
||||
def _discretize_dynamics(self, x0, foot_positions):
|
||||
"""Build discrete-time centroidal dynamics.
|
||||
|
||||
State: [pos, rpy, vel, omega] (12)
|
||||
Continuous: dx/dt = Ac @ x + Bc @ u + gc
|
||||
Discrete: x_{k+1} = Ad @ x + Bd @ u + gd
|
||||
"""
|
||||
m = self.mass
|
||||
I_inv = np.linalg.inv(self.I_body)
|
||||
dt = self.dt
|
||||
yaw = x0[5]
|
||||
cy, sy = np.cos(yaw), np.sin(yaw)
|
||||
|
||||
# Rotation for rpy rate ≈ R_z^T @ omega
|
||||
R_zT = np.array([[cy, sy, 0], [-sy, cy, 0], [0, 0, 1]])
|
||||
|
||||
# Ac (12×12)
|
||||
Ac = np.zeros((NX, NX))
|
||||
Ac[0:3, 6:9] = np.eye(3) # pos_dot = vel
|
||||
Ac[3:6, 9:12] = R_zT # rpy_dot ≈ R_z^T @ omega
|
||||
|
||||
# Ad = I + Ac*dt (first-order)
|
||||
Ad = np.eye(NX) + Ac * dt
|
||||
|
||||
# Bc varies per timestep (foot positions change contact point)
|
||||
Bd_list = []
|
||||
for k in range(self.N):
|
||||
Bc = np.zeros((NX, NU))
|
||||
for leg in range(4):
|
||||
r = foot_positions[leg]
|
||||
# vel_dot += f/m
|
||||
Bc[6:9, leg*3:(leg+1)*3] = np.eye(3) / m
|
||||
# omega_dot += I^{-1} @ (r × f)
|
||||
Bc[9:12, leg*3:(leg+1)*3] = I_inv @ _skew(r)
|
||||
Bd = Bc * dt
|
||||
Bd_list.append(Bd)
|
||||
|
||||
# Gravity contribution
|
||||
gd = np.zeros(NX)
|
||||
gd[6:9] = self.gravity * dt # vel += g*dt
|
||||
|
||||
return Ad, Bd_list, gd
|
||||
@@ -1,261 +0,0 @@
|
||||
"""MPC controller integration for wheeled-legged robot.
|
||||
|
||||
Integrates: gait scheduler + reference trajectory + ConvexMPC solver +
|
||||
swing leg control + stance force mapping + wheel drive.
|
||||
|
||||
Architecture (following go2-convex-mpc):
|
||||
- MPC runs at ~50 Hz (every MPC_DECIMATION control steps)
|
||||
- Swing/stance leg controller runs at control rate (250 Hz)
|
||||
- Wheel drive: stance legs use differential drive, swing legs coast
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
from robot import Robot, RobotState
|
||||
from dynamics import Dynamics
|
||||
from mpc import ConvexMPC, MPC_DT
|
||||
from config import (
|
||||
LEG_NAMES, WHEEL_RADIUS, WHEEL_TRACK, WHEEL_VEL_MAX,
|
||||
CTRL_DT, GAIT_FREQ, GAIT_DUTY, SWING_HEIGHT, PHASE_OFFSETS,
|
||||
DEFAULT_JOINT_ANGLES, ROBOT_MASS,
|
||||
)
|
||||
|
||||
# MPC update decimation (relative to control loop)
|
||||
MPC_DECIMATION = max(1, int(MPC_DT / CTRL_DT)) # ~5 steps at 250Hz
|
||||
|
||||
|
||||
class MPCController:
|
||||
"""Convex MPC locomotion controller for wheeled-legged robot."""
|
||||
|
||||
def __init__(self, robot: Robot):
|
||||
self.robot = robot
|
||||
self.dynamics = Dynamics()
|
||||
self.mpc = ConvexMPC(mass=ROBOT_MASS)
|
||||
|
||||
# User commands
|
||||
self.vel_x = 0.0
|
||||
self.vel_y = 0.0
|
||||
self.yaw_rate = 0.0
|
||||
self.height = 0.35 # actual standing height with default joint angles
|
||||
|
||||
# Gait state - start at phase 0 with all legs in stance (duty=0.6)
|
||||
self._gait_phase = 0.0
|
||||
self._step_count = 0
|
||||
self._initialized = False
|
||||
|
||||
# MPC solution cache - initialize with gravity compensation
|
||||
self._mpc_forces = np.zeros(12)
|
||||
self._init_gravity_comp()
|
||||
|
||||
# Swing trajectory state
|
||||
self._swing_start_foot = {leg: np.zeros(3) for leg in LEG_NAMES}
|
||||
self._swing_start_time = {leg: 0.0 for leg in LEG_NAMES}
|
||||
self._last_contact = {leg: True for leg in LEG_NAMES}
|
||||
|
||||
# Smoothed commands
|
||||
self._smooth_vx = 0.0
|
||||
self._smooth_vy = 0.0
|
||||
self._smooth_yaw = 0.0
|
||||
|
||||
def _init_gravity_comp(self):
|
||||
"""Pre-fill MPC forces with static gravity compensation."""
|
||||
fz_per_leg = ROBOT_MASS * 9.81 / 4.0
|
||||
for i in range(4):
|
||||
self._mpc_forces[i*3 + 2] = fz_per_leg
|
||||
|
||||
def compute(self, state: RobotState, dt: float):
|
||||
"""Main MPC control loop.
|
||||
|
||||
Uses MIT motor protocol: tau = kp*(q_des-q) + kd*(dq_des-dq) + tau_ff
|
||||
where tau_ff comes from MPC force mapping via Jacobian transpose.
|
||||
|
||||
Returns:
|
||||
tau_legs: (12,) feedforward torques for MIT mode
|
||||
wheel_targets: (4,) wheel velocity targets
|
||||
"""
|
||||
# Smooth commands
|
||||
alpha = min(dt * 3.0, 1.0)
|
||||
self._smooth_vx += alpha * (self.vel_x - self._smooth_vx)
|
||||
self._smooth_vy += alpha * (self.vel_y - self._smooth_vy)
|
||||
self._smooth_yaw += alpha * (self.yaw_rate - self._smooth_yaw)
|
||||
|
||||
# Update Pinocchio
|
||||
q_pin, dq_pin = self.robot.get_qpos_qvel_for_pinocchio()
|
||||
self.dynamics.update(q_pin, dq_pin)
|
||||
|
||||
# Initialize foot positions on first call
|
||||
if not self._initialized:
|
||||
for leg in LEG_NAMES:
|
||||
self._swing_start_foot[leg] = self.dynamics.get_foot_pos(leg)
|
||||
self._initialized = True
|
||||
|
||||
# Decide if we should trot or just stand
|
||||
moving = (abs(self._smooth_vx) > 0.02 or
|
||||
abs(self._smooth_vy) > 0.02 or
|
||||
abs(self._smooth_yaw) > 0.05)
|
||||
|
||||
if moving:
|
||||
self._gait_phase = (self._gait_phase + dt * GAIT_FREQ) % 1.0
|
||||
else:
|
||||
self._gait_phase = 0.0 # all legs in stance
|
||||
|
||||
# Contact schedule
|
||||
contacts = {}
|
||||
for leg in LEG_NAMES:
|
||||
phase = (self._gait_phase + PHASE_OFFSETS[leg]) % 1.0
|
||||
contacts[leg] = phase < GAIT_DUTY
|
||||
|
||||
# Get foot positions relative to CoM
|
||||
foot_positions = np.zeros((4, 3))
|
||||
for i, leg in enumerate(LEG_NAMES):
|
||||
foot_positions[i] = self.dynamics.get_foot_pos(leg) - state.pos
|
||||
|
||||
# --- Run MPC at lower rate ---
|
||||
if self._step_count % MPC_DECIMATION == 0:
|
||||
x0 = self._build_state_vector(state)
|
||||
x_ref = self._build_reference(state)
|
||||
contact_table = self._build_contact_table()
|
||||
self._mpc_forces = self.mpc.solve(x0, x_ref, foot_positions, contact_table)
|
||||
|
||||
self._step_count += 1
|
||||
|
||||
# --- Compute feedforward torques and desired joint positions ---
|
||||
tau_ff = np.zeros(12)
|
||||
q_des = np.zeros(12)
|
||||
dq_des = np.zeros(12)
|
||||
kp = np.zeros(12)
|
||||
kd = np.zeros(12)
|
||||
wheel_targets = np.zeros(4)
|
||||
|
||||
for i, leg in enumerate(LEG_NAMES):
|
||||
if contacts[leg]:
|
||||
# Stance: MPC force → feedforward torque, PD holds posture
|
||||
f_leg = self._mpc_forces[i*3:(i+1)*3]
|
||||
J = self.dynamics.get_foot_jacobian_leg(leg)
|
||||
tau_ff[i*3:(i+1)*3] = J.T @ (-f_leg)
|
||||
|
||||
# PD target: default standing angles (posture hold)
|
||||
q_des[i*3] = DEFAULT_JOINT_ANGLES["hip_abduction"]
|
||||
q_des[i*3+1] = DEFAULT_JOINT_ANGLES["hip_pitch"]
|
||||
q_des[i*3+2] = DEFAULT_JOINT_ANGLES["knee"]
|
||||
kp[i*3:(i+1)*3] = [40.0, 40.0, 40.0]
|
||||
kd[i*3:(i+1)*3] = [3.0, 3.0, 3.0]
|
||||
|
||||
# Record foot position
|
||||
self._swing_start_foot[leg] = self.dynamics.get_foot_pos(leg)
|
||||
self._last_contact[leg] = True
|
||||
|
||||
# Wheel drive
|
||||
wheel_targets[i] = self._wheel_cmd(leg)
|
||||
else:
|
||||
# Swing: IK target position, strong PD, no feedforward
|
||||
if self._last_contact[leg]:
|
||||
self._swing_start_foot[leg] = self.dynamics.get_foot_pos(leg)
|
||||
self._swing_start_time[leg] = state.time
|
||||
self._last_contact[leg] = False
|
||||
|
||||
q_ik = self._swing_leg_ik(leg, state, q_pin)
|
||||
q_des[i*3:(i+1)*3] = q_ik
|
||||
kp[i*3:(i+1)*3] = [60.0, 60.0, 60.0] # strong PD for swing
|
||||
kd[i*3:(i+1)*3] = [3.0, 3.0, 3.0]
|
||||
# tau_ff stays 0 for swing
|
||||
|
||||
wheel_targets[i] = 0.0
|
||||
|
||||
# Use MIT protocol via robot interface
|
||||
self.robot.set_ctrl_mit(q_des, dq_des, kp, kd, tau_ff, wheel_targets)
|
||||
# Return dummy (actual ctrl is set directly above)
|
||||
return None, None
|
||||
|
||||
def _build_state_vector(self, state: RobotState):
|
||||
"""Build MPC state: [pos, rpy, vel, omega]."""
|
||||
return np.concatenate([state.pos, state.rpy, state.lin_vel, state.ang_vel])
|
||||
|
||||
def _build_reference(self, state: RobotState):
|
||||
"""Build reference trajectory over MPC horizon."""
|
||||
N = self.mpc.N
|
||||
x_ref = np.zeros((12, N))
|
||||
|
||||
yaw = state.rpy[2]
|
||||
cy, sy = np.cos(yaw), np.sin(yaw)
|
||||
R_z = np.array([[cy, -sy, 0], [sy, cy, 0], [0, 0, 1]])
|
||||
vel_world = R_z @ np.array([self._smooth_vx, self._smooth_vy, 0.0])
|
||||
|
||||
for k in range(N):
|
||||
t = (k + 1) * self.mpc.dt
|
||||
# Position: integrate from current
|
||||
x_ref[0, k] = state.pos[0] + vel_world[0] * t
|
||||
x_ref[1, k] = state.pos[1] + vel_world[1] * t
|
||||
x_ref[2, k] = self.height
|
||||
# RPY: keep roll/pitch zero, integrate yaw
|
||||
x_ref[3, k] = 0.0
|
||||
x_ref[4, k] = 0.0
|
||||
x_ref[5, k] = yaw + self._smooth_yaw * t
|
||||
# Velocity
|
||||
x_ref[6, k] = vel_world[0]
|
||||
x_ref[7, k] = vel_world[1]
|
||||
x_ref[8, k] = 0.0
|
||||
# Angular velocity
|
||||
x_ref[9, k] = 0.0
|
||||
x_ref[10, k] = 0.0
|
||||
x_ref[11, k] = self._smooth_yaw
|
||||
|
||||
return x_ref
|
||||
|
||||
def _build_contact_table(self):
|
||||
"""Build contact schedule over MPC horizon."""
|
||||
N = self.mpc.N
|
||||
table = np.zeros((4, N), dtype=int)
|
||||
for k in range(N):
|
||||
future_phase = (self._gait_phase + (k + 1) * self.mpc.dt * GAIT_FREQ) % 1.0
|
||||
for i, leg in enumerate(LEG_NAMES):
|
||||
leg_phase = (future_phase + PHASE_OFFSETS[leg]) % 1.0
|
||||
table[i, k] = 1 if leg_phase < GAIT_DUTY else 0
|
||||
return table
|
||||
|
||||
def _swing_leg_ik(self, leg: str, state: RobotState, q_pin: np.ndarray):
|
||||
"""Swing leg: compute IK target joint angles for trajectory."""
|
||||
swing_phase = self._get_swing_phase(leg)
|
||||
|
||||
p_start = self._swing_start_foot[leg]
|
||||
p_end = self._compute_touchdown(leg, state)
|
||||
|
||||
s = swing_phase
|
||||
s_mj = 10*s**3 - 15*s**4 + 6*s**5
|
||||
|
||||
pos_des = p_start + (p_end - p_start) * s_mj
|
||||
# Z lift
|
||||
z_lift = 64.0 * s**3 * (1.0 - s)**3
|
||||
pos_des[2] = p_start[2] + SWING_HEIGHT * z_lift
|
||||
|
||||
# IK to get joint angles
|
||||
q_ik = self.dynamics.inverse_kinematics(leg, pos_des, q_pin)
|
||||
return q_ik
|
||||
|
||||
def _get_swing_phase(self, leg: str) -> float:
|
||||
phase = (self._gait_phase + PHASE_OFFSETS[leg]) % 1.0
|
||||
if phase < GAIT_DUTY:
|
||||
return 0.0
|
||||
return (phase - GAIT_DUTY) / (1.0 - GAIT_DUTY)
|
||||
|
||||
def _compute_touchdown(self, leg: str, state: RobotState) -> np.ndarray:
|
||||
"""Raibert heuristic for touchdown position."""
|
||||
td = self._swing_start_foot[leg].copy()
|
||||
t_stance = GAIT_DUTY / GAIT_FREQ
|
||||
|
||||
yaw = state.rpy[2]
|
||||
cy, sy = np.cos(yaw), np.sin(yaw)
|
||||
R_z = np.array([[cy, -sy, 0], [sy, cy, 0], [0, 0, 1]])
|
||||
cmd_vel_world = R_z @ np.array([self._smooth_vx, self._smooth_vy, 0.0])
|
||||
|
||||
td[0] += cmd_vel_world[0] * t_stance * 0.5
|
||||
td[1] += cmd_vel_world[1] * t_stance * 0.5
|
||||
td[2] = WHEEL_RADIUS
|
||||
return td
|
||||
|
||||
def _wheel_cmd(self, leg: str) -> float:
|
||||
"""Differential drive for a single wheel."""
|
||||
if leg[1] == "l":
|
||||
v = (self._smooth_vx - 0.5 * WHEEL_TRACK * self._smooth_yaw) / WHEEL_RADIUS
|
||||
else:
|
||||
v = (self._smooth_vx + 0.5 * WHEEL_TRACK * self._smooth_yaw) / WHEEL_RADIUS
|
||||
return np.clip(v, -WHEEL_VEL_MAX, WHEEL_VEL_MAX)
|
||||
@@ -1,229 +0,0 @@
|
||||
"""MuJoCo interface for the wheeled-legged robot.
|
||||
|
||||
Configures actuators as proper PD controllers at runtime:
|
||||
- Leg joints: force = kp*(ctrl - qpos) - kd*qvel, ctrl = target angle
|
||||
- Wheel joints: force = gain*(ctrl - qvel), ctrl = target velocity (rad/s)
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import mujoco
|
||||
from dataclasses import dataclass
|
||||
from config import (SCENE_XML, LEG_NAMES, LEG_JOINTS, WHEEL_JOINT,
|
||||
DEFAULT_JOINT_ANGLES, WHEEL_RADIUS, WHEEL_TRACK)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RobotState:
|
||||
"""Robot state from MuJoCo."""
|
||||
pos: np.ndarray # (3,) world position
|
||||
quat: np.ndarray # (4,) quaternion (w,x,y,z) MuJoCo convention
|
||||
rot: np.ndarray # (3,3) body→world rotation
|
||||
rpy: np.ndarray # (3,) roll, pitch, yaw
|
||||
lin_vel: np.ndarray # (3,) world frame linear velocity
|
||||
ang_vel: np.ndarray # (3,) body frame angular velocity
|
||||
joint_pos: np.ndarray # (16,) all joint positions [fl3+wheel, fr3+wheel, rl3+wheel, rr3+wheel]
|
||||
joint_vel: np.ndarray # (16,) all joint velocities
|
||||
time: float
|
||||
|
||||
|
||||
class Robot:
|
||||
"""MuJoCo simulation interface with proper PD actuator configuration."""
|
||||
|
||||
# Leg PD gains (tuned for 12.3kg robot)
|
||||
LEG_KP = 60.0
|
||||
LEG_KD = 3.0
|
||||
# Wheel velocity gain
|
||||
WHEEL_KP = 2.0
|
||||
|
||||
def __init__(self, xml_path=None):
|
||||
self.model = mujoco.MjModel.from_xml_path(str(xml_path or SCENE_XML))
|
||||
self.data = mujoco.MjData(self.model)
|
||||
|
||||
# Cache IDs
|
||||
self._base_bid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "base_link")
|
||||
self._actuator_ids = {} # name → actuator index
|
||||
self._joint_qpos_adr = {} # name → qpos address
|
||||
self._joint_qvel_adr = {} # name → qvel address
|
||||
|
||||
# Build joint/actuator maps
|
||||
self._ctrl_order = []
|
||||
for leg in LEG_NAMES:
|
||||
for jt in (*LEG_JOINTS, WHEEL_JOINT):
|
||||
name = f"{leg}_{jt}"
|
||||
aid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, name)
|
||||
jid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, name)
|
||||
self._actuator_ids[name] = aid
|
||||
self._joint_qpos_adr[name] = self.model.jnt_qposadr[jid]
|
||||
self._joint_qvel_adr[name] = self.model.jnt_dofadr[jid]
|
||||
self._ctrl_order.append(name)
|
||||
|
||||
# Configure actuators as proper PD controllers
|
||||
self._configure_actuators()
|
||||
|
||||
def _configure_actuators(self):
|
||||
"""Set actuators to proper PD mode.
|
||||
|
||||
Leg joints: force = kp*(ctrl - qpos) - kd*qvel
|
||||
Wheels: force = gain*(ctrl - qvel) (velocity tracking)
|
||||
"""
|
||||
for i in range(self.model.nu):
|
||||
name = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, i)
|
||||
self.model.actuator_biastype[i] = 1 # affine bias
|
||||
self.model.actuator_gaintype[i] = 0 # fixed gain
|
||||
self.model.actuator_forcelimited[i] = 0 # no force clamp (17Nm is in actuatorfrcrange)
|
||||
|
||||
if 'wheel' not in name:
|
||||
self.model.actuator_gainprm[i, 0] = self.LEG_KP
|
||||
self.model.actuator_biasprm[i, 0] = 0.0
|
||||
self.model.actuator_biasprm[i, 1] = -self.LEG_KP
|
||||
self.model.actuator_biasprm[i, 2] = -self.LEG_KD
|
||||
self.model.actuator_ctrlrange[i] = [-3.14, 3.14]
|
||||
else:
|
||||
self.model.actuator_gainprm[i, 0] = self.WHEEL_KP
|
||||
self.model.actuator_biasprm[i, 0] = 0.0
|
||||
self.model.actuator_biasprm[i, 1] = 0.0
|
||||
self.model.actuator_biasprm[i, 2] = -self.WHEEL_KP
|
||||
self.model.actuator_ctrlrange[i] = [-20.0, 20.0]
|
||||
|
||||
@property
|
||||
def dt(self):
|
||||
return self.model.opt.timestep
|
||||
|
||||
def reset(self):
|
||||
"""Reset to standing pose at correct height for default joint angles."""
|
||||
mujoco.mj_resetData(self.model, self.data)
|
||||
|
||||
# Set default leg angles
|
||||
for leg in LEG_NAMES:
|
||||
for jt, key in zip(LEG_JOINTS, ("hip_abduction", "hip_pitch", "knee")):
|
||||
name = f"{leg}_{jt}"
|
||||
adr = self._joint_qpos_adr[name]
|
||||
self.data.qpos[adr] = DEFAULT_JOINT_ANGLES[key]
|
||||
|
||||
# Compute correct base height from default angles using 2R FK
|
||||
# leg_length = sqrt(L1^2 + L2^2 - 2*L1*L2*cos(pi + knee))
|
||||
import math
|
||||
L1, L2 = 0.25, 0.20
|
||||
knee = DEFAULT_JOINT_ANGLES["knee"]
|
||||
leg_length = math.sqrt(L1**2 + L2**2 - 2*L1*L2*math.cos(math.pi + knee))
|
||||
# base_z = wheel_radius + leg_length - hip_z_offset
|
||||
base_z = 0.10 + leg_length - 0.054
|
||||
self.data.qpos[2] = base_z
|
||||
self.data.qpos[3] = 1.0 # quat w
|
||||
|
||||
mujoco.mj_forward(self.model, self.data)
|
||||
|
||||
# Set ctrl to match initial pose (so PD doesn't jerk)
|
||||
for leg in LEG_NAMES:
|
||||
for jt, key in zip(LEG_JOINTS, ("hip_abduction", "hip_pitch", "knee")):
|
||||
name = f"{leg}_{jt}"
|
||||
self.data.ctrl[self._actuator_ids[name]] = DEFAULT_JOINT_ANGLES[key]
|
||||
# Wheels: zero velocity
|
||||
self.data.ctrl[self._actuator_ids[f"{leg}_{WHEEL_JOINT}"]] = 0.0
|
||||
|
||||
def get_state(self) -> RobotState:
|
||||
"""Extract robot state."""
|
||||
pos = self.data.xpos[self._base_bid].copy()
|
||||
quat = self.data.xquat[self._base_bid].copy() # (w,x,y,z)
|
||||
rot = self.data.xmat[self._base_bid].reshape(3, 3).copy()
|
||||
|
||||
rpy = np.array([
|
||||
np.arctan2(rot[2, 1], rot[2, 2]),
|
||||
np.arctan2(-rot[2, 0], np.sqrt(rot[2, 1]**2 + rot[2, 2]**2)),
|
||||
np.arctan2(rot[1, 0], rot[0, 0]),
|
||||
])
|
||||
|
||||
# Base velocity (world frame)
|
||||
lin_vel = self.data.qvel[0:3].copy()
|
||||
ang_vel = self.data.qvel[3:6].copy()
|
||||
|
||||
# Joint states (16 joints: 4 legs × 4 joints each)
|
||||
joint_pos = np.zeros(16)
|
||||
joint_vel = np.zeros(16)
|
||||
for i, name in enumerate(self._ctrl_order):
|
||||
joint_pos[i] = self.data.qpos[self._joint_qpos_adr[name]]
|
||||
joint_vel[i] = self.data.qvel[self._joint_qvel_adr[name]]
|
||||
|
||||
return RobotState(
|
||||
pos=pos, quat=quat, rot=rot, rpy=rpy,
|
||||
lin_vel=lin_vel, ang_vel=ang_vel,
|
||||
joint_pos=joint_pos, joint_vel=joint_vel,
|
||||
time=self.data.time,
|
||||
)
|
||||
|
||||
def set_ctrl(self, leg_targets: np.ndarray, wheel_targets: np.ndarray):
|
||||
"""Set actuator commands (position PD mode).
|
||||
|
||||
Args:
|
||||
leg_targets: (12,) target joint angles for legs [fl3, fr3, rl3, rr3]
|
||||
wheel_targets: (4,) target wheel velocities [fl, fr, rl, rr] in rad/s
|
||||
"""
|
||||
for i, leg in enumerate(LEG_NAMES):
|
||||
for j, jt in enumerate(LEG_JOINTS):
|
||||
name = f"{leg}_{jt}"
|
||||
self.data.ctrl[self._actuator_ids[name]] = leg_targets[i * 3 + j]
|
||||
name = f"{leg}_{WHEEL_JOINT}"
|
||||
self.data.ctrl[self._actuator_ids[name]] = wheel_targets[i]
|
||||
|
||||
def set_ctrl_mit(self, q_des: np.ndarray, dq_des: np.ndarray,
|
||||
kp: np.ndarray, kd: np.ndarray, tau_ff: np.ndarray,
|
||||
wheel_targets: np.ndarray):
|
||||
"""MIT motor protocol: tau = kp*(q_des-q) + kd*(dq_des-dq) + tau_ff.
|
||||
|
||||
Computes torque in software, sends to actuators in torque mode.
|
||||
Call enable_torque_mode() first.
|
||||
|
||||
Args:
|
||||
q_des: (12,) desired joint angles
|
||||
dq_des: (12,) desired joint velocities
|
||||
kp: (12,) position gains (0 for pure torque)
|
||||
kd: (12,) velocity gains
|
||||
tau_ff: (12,) feedforward torques
|
||||
wheel_targets: (4,) wheel velocity targets
|
||||
"""
|
||||
for i, leg in enumerate(LEG_NAMES):
|
||||
for j, jt in enumerate(LEG_JOINTS):
|
||||
name = f"{leg}_{jt}"
|
||||
aid = self._actuator_ids[name]
|
||||
idx = i * 3 + j
|
||||
q = self.data.qpos[self._joint_qpos_adr[name]]
|
||||
dq = self.data.qvel[self._joint_qvel_adr[name]]
|
||||
tau = (kp[idx] * (q_des[idx] - q)
|
||||
+ kd[idx] * (dq_des[idx] - dq)
|
||||
+ tau_ff[idx])
|
||||
self.data.ctrl[aid] = np.clip(tau, -17.0, 17.0)
|
||||
name = f"{leg}_{WHEEL_JOINT}"
|
||||
self.data.ctrl[self._actuator_ids[name]] = wheel_targets[i]
|
||||
|
||||
def enable_torque_mode(self):
|
||||
"""Switch leg actuators to direct torque mode (for MPC/MIT)."""
|
||||
for i in range(self.model.nu):
|
||||
name = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, i)
|
||||
if 'wheel' not in name:
|
||||
self.model.actuator_gainprm[i, 0] = 1.0
|
||||
self.model.actuator_biasprm[i, :3] = [0, 0, 0]
|
||||
self.model.actuator_biastype[i] = 0
|
||||
self.model.actuator_ctrlrange[i] = [-17.0, 17.0]
|
||||
|
||||
def enable_position_mode(self):
|
||||
"""Switch leg actuators back to position PD mode."""
|
||||
self._configure_actuators()
|
||||
|
||||
def step(self):
|
||||
"""Advance one simulation timestep."""
|
||||
mujoco.mj_step(self.model, self.data)
|
||||
|
||||
def get_qpos_qvel_for_pinocchio(self):
|
||||
"""Get full qpos/qvel for Pinocchio (reorder quaternion)."""
|
||||
qpos = self.data.qpos.copy()
|
||||
qvel = self.data.qvel.copy()
|
||||
# MuJoCo quat: (w,x,y,z) → Pinocchio: (x,y,z,w)
|
||||
w, x, y, z = qpos[3], qpos[4], qpos[5], qpos[6]
|
||||
q_pin = np.concatenate([qpos[0:3], [x, y, z, w], qpos[7:]])
|
||||
# MuJoCo vel is already [lin_world(3), ang_body(3), joints(16)]
|
||||
# Pinocchio wants [lin_body(3), ang_body(3), joints(16)]
|
||||
from scipy.spatial.transform import Rotation
|
||||
R = Rotation.from_quat([x, y, z, w]).as_matrix()
|
||||
v_body = R.T @ qvel[0:3]
|
||||
dq_pin = np.concatenate([v_body, qvel[3:]])
|
||||
return q_pin, dq_pin
|
||||
@@ -1,111 +0,0 @@
|
||||
"""Main entry point: wheeled-legged robot simulation.
|
||||
|
||||
Controls:
|
||||
Mode: wheel (default) - differential drive + posture hold
|
||||
trot - quadruped gait with wheel assist
|
||||
mpc - convex MPC locomotion (torque control)
|
||||
|
||||
Keyboard (in MuJoCo viewer):
|
||||
W/S: vel_x ±0.1
|
||||
A/D: yaw_rate ±0.2
|
||||
Q/E: height ±0.02
|
||||
1: wheel mode
|
||||
2: trot mode
|
||||
3: MPC mode
|
||||
4: prone toggle
|
||||
Z: reset commands
|
||||
"""
|
||||
|
||||
import time
|
||||
import numpy as np
|
||||
import mujoco.viewer as mjv
|
||||
|
||||
from robot import Robot
|
||||
from controller import Controller
|
||||
from gui import GUI
|
||||
from config import CTRL_DECIMATION
|
||||
|
||||
|
||||
def main():
|
||||
robot = Robot()
|
||||
robot.reset()
|
||||
ctrl = Controller(robot)
|
||||
gui = GUI(ctrl)
|
||||
|
||||
step = 0
|
||||
sim_steps_per_ctrl = CTRL_DECIMATION
|
||||
|
||||
def key_callback(keycode):
|
||||
"""Called from MuJoCo render thread - only modify ctrl directly, not tkinter."""
|
||||
try:
|
||||
c = chr(keycode).lower()
|
||||
except (ValueError, OverflowError):
|
||||
return
|
||||
if c == 'w':
|
||||
ctrl.vel_x = min(ctrl.vel_x + 0.1, 1.5)
|
||||
elif c == 's':
|
||||
ctrl.vel_x = max(ctrl.vel_x - 0.1, -1.5)
|
||||
elif c == 'a':
|
||||
ctrl.yaw_rate = min(ctrl.yaw_rate + 0.2, 2.0)
|
||||
elif c == 'd':
|
||||
ctrl.yaw_rate = max(ctrl.yaw_rate - 0.2, -2.0)
|
||||
elif c == 'q':
|
||||
ctrl.height = min(ctrl.height + 0.02, 0.45)
|
||||
elif c == 'e':
|
||||
ctrl.height = max(ctrl.height - 0.02, 0.15)
|
||||
elif c == '1':
|
||||
ctrl.mode = "wheel"; ctrl.prone = False
|
||||
elif c == '2':
|
||||
ctrl.mode = "trot"; ctrl.prone = False
|
||||
elif c == '3':
|
||||
ctrl.mode = "mpc"; ctrl.prone = False
|
||||
elif c == '4':
|
||||
ctrl.prone = not ctrl.prone
|
||||
elif c == 'z':
|
||||
ctrl.vel_x = 0.0; ctrl.vel_y = 0.0; ctrl.yaw_rate = 0.0
|
||||
|
||||
with mjv.launch_passive(robot.model, robot.data, key_callback=key_callback) as viewer:
|
||||
viewer.cam.distance = 2.5
|
||||
viewer.cam.elevation = -20
|
||||
viewer.cam.azimuth = 135
|
||||
|
||||
last_time = robot.data.time
|
||||
|
||||
while viewer.is_running() and not gui.closed:
|
||||
t_start = time.perf_counter()
|
||||
|
||||
# Detect viewer reset (Backspace) - time jumps back to 0
|
||||
if robot.data.time < last_time:
|
||||
robot.reset()
|
||||
last_time = robot.data.time
|
||||
|
||||
# Get state and compute control
|
||||
state = robot.get_state()
|
||||
leg_targets, wheel_targets = ctrl.compute(state, robot.dt * sim_steps_per_ctrl)
|
||||
|
||||
# Apply control and step simulation
|
||||
# MPC mode sets ctrl directly via set_ctrl_mit, skip set_ctrl
|
||||
if ctrl.mode != "mpc":
|
||||
robot.set_ctrl(leg_targets, wheel_targets)
|
||||
for _ in range(sim_steps_per_ctrl):
|
||||
robot.step()
|
||||
|
||||
viewer.sync()
|
||||
step += 1
|
||||
|
||||
# Update GUI every 25 steps (~10 Hz)
|
||||
if step % 25 == 0:
|
||||
state = robot.get_state()
|
||||
gui.update_status(state, step)
|
||||
if not gui.tick():
|
||||
break
|
||||
|
||||
# Real-time sync
|
||||
elapsed = time.perf_counter() - t_start
|
||||
target_dt = robot.dt * sim_steps_per_ctrl
|
||||
if elapsed < target_dt:
|
||||
time.sleep(target_dt - elapsed)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -25,7 +25,7 @@ class MuJoCoIO:
|
||||
print("[MuJoCoIO] Loading MuJoCo model...")
|
||||
spec = mujoco.MjSpec.from_file(str(temp_xml))
|
||||
|
||||
# Override actuators to match mjlab exactly
|
||||
# Override actuators to match training/runtime actuator semantics.
|
||||
self._rebuild_actuators(spec)
|
||||
|
||||
self.m = spec.compile()
|
||||
@@ -91,8 +91,10 @@ class MuJoCoIO:
|
||||
for act in actuators_to_delete:
|
||||
spec.delete(act)
|
||||
|
||||
KP_LEG, KD_LEG = 40.0, 1.0
|
||||
KD_WHEEL = 0.5
|
||||
# Keep sim2sim aligned with the training robot config and sim2real runtime:
|
||||
# leg position PD = (50.0, 1.5), wheel velocity damping = 1.0.
|
||||
KP_LEG, KD_LEG = 50.0, 1.5
|
||||
KD_WHEEL = 1.0
|
||||
EFFORT_LIMIT = 17.0
|
||||
|
||||
leg_jnames = [
|
||||
|
||||
@@ -77,6 +77,7 @@ class PolicyRunner:
|
||||
|
||||
# Load both policy networks
|
||||
self.policies = {}
|
||||
self.policy_obs_dims = {}
|
||||
for name, path in self.policy_paths.items():
|
||||
print(f"[PolicyRunner] Loading {name} policy from: {path}")
|
||||
if Path(path).exists():
|
||||
@@ -84,6 +85,7 @@ class PolicyRunner:
|
||||
else:
|
||||
print(f"[PolicyRunner] WARNING: {name} policy file not found! Falling back to rough.")
|
||||
self.policies[name] = load_policy(self.policy_paths["rough"], device)
|
||||
self.policy_obs_dims[name] = int(self.policies[name].obs_mean.numel())
|
||||
|
||||
# Default DOF positions for each policy
|
||||
self.default_dof_poses = {
|
||||
@@ -226,11 +228,23 @@ class PolicyRunner:
|
||||
|
||||
raw_actions_out = {}
|
||||
for name in active_policies:
|
||||
# Flatten observation history
|
||||
obs_history_array = np.array(self.obs_histories[name])
|
||||
term_dims = [3, 3, 3, 12, 12, 4, 16]
|
||||
term_histories = np.split(obs_history_array, np.cumsum(term_dims)[:-1], axis=1)
|
||||
flat_obs = np.concatenate([h.flatten() for h in term_histories])
|
||||
expected_obs_dim = self.policy_obs_dims[name]
|
||||
|
||||
if expected_obs_dim == current_obs_53d.shape[0]:
|
||||
# Newer policies consume the current 53D observation directly.
|
||||
flat_obs = self.obs_histories[name][-1]
|
||||
elif expected_obs_dim == current_obs_53d.shape[0] * self.history_length:
|
||||
# Legacy policies expect 6-step history stacking grouped by term.
|
||||
obs_history_array = np.array(self.obs_histories[name])
|
||||
term_dims = [3, 3, 3, 12, 12, 4, 16]
|
||||
term_histories = np.split(obs_history_array, np.cumsum(term_dims)[:-1], axis=1)
|
||||
flat_obs = np.concatenate([h.flatten() for h in term_histories])
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Policy '{name}' expects obs dim {expected_obs_dim}, "
|
||||
f"but sim2sim can only provide {current_obs_53d.shape[0]} or "
|
||||
f"{current_obs_53d.shape[0] * self.history_length}."
|
||||
)
|
||||
|
||||
obs_tensor = torch.tensor(flat_obs, device=self.device, dtype=torch.float32).unsqueeze(0)
|
||||
|
||||
|
||||
@@ -74,7 +74,11 @@ from ..mdp.rewards import (
|
||||
joint_mirror,
|
||||
feet_contact_without_cmd,
|
||||
upright_roll_only,
|
||||
pitch_control_penalty,
|
||||
upward,
|
||||
joint_power,
|
||||
ang_vel_xy_l2,
|
||||
undesired_contacts,
|
||||
contact_forces,
|
||||
)
|
||||
from ..mdp.curriculums import terrain_levels_vel_strict
|
||||
from ..mdp.commands import UniformThresholdVelocityCommandCfg
|
||||
@@ -174,7 +178,7 @@ def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
|
||||
params={"asset_cfg": SceneEntityCfg("robot", joint_names=(".*_wheel_joint",))},
|
||||
scale=0.05, noise=Unoise(n_min=-1.0, n_max=1.0),
|
||||
),
|
||||
"actions": ObservationTermCfg(func=velocity_mdp.last_action, history_length=1),
|
||||
"actions": ObservationTermCfg(func=velocity_mdp.last_action),
|
||||
}
|
||||
|
||||
critic_terms = {
|
||||
@@ -192,7 +196,7 @@ def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
|
||||
observations = {
|
||||
"actor": ObservationGroupCfg(
|
||||
terms=actor_terms, concatenate_terms=True,
|
||||
enable_corruption=True, history_length=6,
|
||||
enable_corruption=True,
|
||||
),
|
||||
"critic": ObservationGroupCfg(
|
||||
terms=critic_terms, concatenate_terms=True, enable_corruption=False,
|
||||
@@ -239,8 +243,15 @@ def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
|
||||
"reset_base": EventTermCfg(
|
||||
func=envs_mdp.reset_root_state_uniform, mode="reset",
|
||||
params={
|
||||
"pose_range": {"z": (0.30, 0.50), "yaw": (-math.pi, math.pi)},
|
||||
"velocity_range": {"x": (-0.5, 0.5), "y": (-0.15, 0.15), "yaw": (-0.35, 0.35)},
|
||||
"pose_range": {"x": (-0.5, 0.5), "y": (-0.5, 0.5), "yaw": (-math.pi, math.pi)},
|
||||
"velocity_range": {
|
||||
"x": (-0.5, 0.5),
|
||||
"y": (-0.5, 0.5),
|
||||
"z": (-0.5, 0.5),
|
||||
"roll": (-0.5, 0.5),
|
||||
"pitch": (-0.5, 0.5),
|
||||
"yaw": (-0.5, 0.5),
|
||||
},
|
||||
"asset_cfg": SceneEntityCfg("robot"),
|
||||
},
|
||||
),
|
||||
@@ -254,13 +265,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 +277,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 +285,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 +334,8 @@ def _make_base_env_cfg() -> ManagerBasedRlEnvCfg:
|
||||
),
|
||||
commands=commands, actions=actions, observations=observations,
|
||||
rewards=rewards, terminations=terminations, events=events,
|
||||
metrics=metrics, curriculum=curriculum, decimation=10, episode_length_s=20.0,
|
||||
sim=SimulationCfg(mujoco=MujocoCfg(impratio=100, cone="elliptic")),
|
||||
metrics=metrics, curriculum=curriculum, decimation=4, episode_length_s=20.0,
|
||||
sim=SimulationCfg(mujoco=MujocoCfg(timestep=0.005, impratio=100, cone="elliptic")),
|
||||
viewer=ViewerConfig(body_name="base_link", distance=3.0, elevation=-20.0, azimuth=45.0),
|
||||
)
|
||||
|
||||
@@ -380,33 +368,36 @@ def rough_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
|
||||
size=(8.0, 8.0), border_width=20.0, num_rows=10, num_cols=20, curriculum=True,
|
||||
sub_terrains={
|
||||
"flat": BoxFlatTerrainCfg(proportion=0.05, size=(8.0, 8.0)),
|
||||
"pyramid_stairs": BoxPyramidStairsTerrainCfg(proportion=0.25, step_height_range=(0.0, 0.3), step_width=0.30, size=(8.0, 8.0)),
|
||||
"pyramid_stairs_inv": BoxInvertedPyramidStairsTerrainCfg(proportion=0.10, step_height_range=(0.0, 0.3), step_width=0.30, size=(8.0, 8.0)),
|
||||
"random_grid": BoxRandomGridTerrainCfg(proportion=0.1, grid_width=0.45, grid_height_range=(0.0, 0.3), size=(8.0, 8.0)),
|
||||
"random_rough": HfRandomUniformTerrainCfg(proportion=0.05, noise_range=(0.0, 0.06), noise_step=0.01, horizontal_scale=0.20, downsampled_scale=0.20, border_width=0.25, base_thickness_ratio=100.0, size=(8.0, 8.0)),
|
||||
"perlin_noise": HfPerlinNoiseTerrainCfg(proportion=0.05, height_range=(0.0, 0.06), octaves=2, persistence=0.4, lacunarity=2.0, horizontal_scale=0.20, resolution=0.20, border_width=0.50, base_thickness_ratio=100.0, size=(8.0, 8.0)),
|
||||
"rc_wall": RCWallTerrainCfg(proportion=0.25, wall_height_range=(0.0, 0.45), size=(8.0, 8.0)),
|
||||
"sloped_terrain": HfPyramidSlopedTerrainCfg(proportion=0.15, slope_range=(0.052, 0.325), platform_width=2.0, border_width=0.25, base_thickness_ratio=100.0, horizontal_scale=0.20, size=(8.0, 8.0)),
|
||||
"pyramid_stairs": BoxPyramidStairsTerrainCfg(proportion=0.05, step_height_range=(0.0, 0.3), step_width=0.30, size=(8.0, 8.0)),
|
||||
"pyramid_stairs_inv": BoxInvertedPyramidStairsTerrainCfg(proportion=0.45, step_height_range=(0.0, 0.3), step_width=0.30, size=(8.0, 8.0)),
|
||||
"random_grid": BoxRandomGridTerrainCfg(proportion=0.27, grid_width=0.45, grid_height_range=(0.0, 0.3), size=(8.0, 8.0)),
|
||||
"random_rough": HfRandomUniformTerrainCfg(proportion=0.01, noise_range=(0.0, 0.06), noise_step=0.01, horizontal_scale=0.20, downsampled_scale=0.20, border_width=0.25, base_thickness_ratio=100.0, size=(8.0, 8.0)),
|
||||
"perlin_noise": HfPerlinNoiseTerrainCfg(proportion=0.01, height_range=(0.0, 0.06), octaves=2, persistence=0.4, lacunarity=2.0, horizontal_scale=0.20, resolution=0.20, border_width=0.50, base_thickness_ratio=100.0, size=(8.0, 8.0)),
|
||||
"rc_wall": RCWallTerrainCfg(proportion=0.15, wall_height_range=(0.0, 0.45), size=(8.0, 8.0)),
|
||||
"sloped_terrain": HfPyramidSlopedTerrainCfg(proportion=0.01, slope_range=(0.052, 0.325), platform_width=2.0, border_width=0.25, base_thickness_ratio=100.0, horizontal_scale=0.20, size=(8.0, 8.0)),
|
||||
},
|
||||
),
|
||||
max_init_terrain_level=0,
|
||||
max_init_terrain_level=5,
|
||||
)
|
||||
|
||||
# Disable default velocity stages command and bind strict velocity terrain curriculum
|
||||
# Keep the custom terrain set, but align command/curriculum behavior with go2w rough.
|
||||
cfg.curriculum.pop("command_vel", None)
|
||||
cfg.curriculum["terrain_levels"] = CurriculumTermCfg(func=terrain_levels_vel_strict, params={"command_name": "twist"})
|
||||
cfg.curriculum["terrain_levels"] = CurriculumTermCfg(func=velocity_mdp.terrain_levels_vel, params={"command_name": "twist"})
|
||||
|
||||
cfg.commands["twist"].heading_command = True
|
||||
cfg.commands["twist"].rel_heading_envs = 0.5
|
||||
cfg.commands["twist"].heading_control_stiffness = 0.6
|
||||
cfg.commands["twist"].rel_heading_envs = 1.0
|
||||
cfg.commands["twist"].heading_control_stiffness = 0.5
|
||||
cfg.commands["twist"].ranges.heading = (-math.pi, math.pi)
|
||||
cfg.commands["twist"].rel_standing_envs = 0.2
|
||||
cfg.commands["twist"].rel_standing_envs = 0.02
|
||||
cfg.commands["twist"].ranges.lin_vel_x = (-1.0, 1.0)
|
||||
cfg.commands["twist"].ranges.lin_vel_y = (-0.6, 0.6)
|
||||
cfg.commands["twist"].ranges.ang_vel_z = (-1.0, 1.0)
|
||||
|
||||
# ------------------
|
||||
# Startup & Reset Randomizations
|
||||
# ------------------
|
||||
cfg.events["joint_friction"] = EventTermCfg(func=envs_dr.joint_friction, mode="startup", params={"asset_cfg": SceneEntityCfg("robot"), "ranges": (0.7, 1.3), "operation": "scale"})
|
||||
cfg.events["reset_joints"] = EventTermCfg(func=envs_mdp.reset_joints_by_offset, mode="reset", params={"position_range": (0.0, 0.1), "velocity_range": (0.0, 0.0), "asset_cfg": SceneEntityCfg("robot", joint_names=(".*",))})
|
||||
cfg.events.pop("joint_friction", None)
|
||||
cfg.events["reset_joints"] = EventTermCfg(func=envs_mdp.reset_joints_by_offset, mode="reset", params={"position_range": (0.0, 0.0), "velocity_range": (0.0, 0.0), "asset_cfg": SceneEntityCfg("robot", joint_names=(".*",))})
|
||||
|
||||
cfg.events["reset_base"] = EventTermCfg(
|
||||
func=envs_mdp.reset_root_state_uniform, mode="reset",
|
||||
@@ -426,42 +417,36 @@ def rough_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
|
||||
# Rewards Integration
|
||||
# ------------------
|
||||
cfg.rewards["track_lin_vel"] = RewardTermCfg(
|
||||
func=track_linear_velocity_l1,
|
||||
weight=4.5,
|
||||
func=track_linear_velocity,
|
||||
weight=3.0,
|
||||
params={"std": 0.5, "command_name": "twist"}
|
||||
)
|
||||
cfg.rewards["track_ang_vel"] = RewardTermCfg(
|
||||
func=track_angular_velocity,
|
||||
weight=1.5,
|
||||
params={"std": 0.5, "command_name": "twist"}
|
||||
)
|
||||
cfg.rewards["track_ang_vel"].weight = 2.0
|
||||
cfg.rewards["track_ang_vel"].params["std"] = 0.5
|
||||
|
||||
cfg.rewards["lin_vel_z"] = RewardTermCfg(func=lin_vel_z_l2, weight=-0.5) # 🌟 增强垂直速度惩罚，抑制越障后的惯性暴冲
|
||||
cfg.rewards["ang_vel_xy"] = RewardTermCfg(func=velocity_mdp.body_angular_velocity_penalty, weight=-0.3, params={"asset_cfg": SceneEntityCfg("robot", body_names=("base_link",))}) # 🌟 增强角速度惩罚，防止突发性翻转/后仰
|
||||
cfg.rewards["lin_vel_z"] = RewardTermCfg(func=lin_vel_z_l2, weight=-2.0)
|
||||
cfg.rewards["ang_vel_xy"] = RewardTermCfg(func=ang_vel_xy_l2, weight=-0.05, params={"asset_cfg": SceneEntityCfg("robot")})
|
||||
|
||||
cfg.rewards.pop("upright", None)
|
||||
cfg.rewards["roll_penalty"] = RewardTermCfg(
|
||||
func=upright_roll_only,
|
||||
weight=-1.0,
|
||||
params={"asset_cfg": SceneEntityCfg("robot")}
|
||||
)
|
||||
cfg.rewards.pop("roll_penalty", None)
|
||||
|
||||
# 🌟 限制俯仰角死区（Pitch Dead-zone）：允许正常爬坡时有最大 29 度（0.50 rad）的仰角，但严厉惩罚超过该仰角的“前轮悬空暴冲/后翻”
|
||||
cfg.rewards["pitch_penalty"] = RewardTermCfg(
|
||||
func=pitch_control_penalty,
|
||||
weight=-1.5,
|
||||
params={"max_pitch_rad": 0.50, "asset_cfg": SceneEntityCfg("robot")}
|
||||
)
|
||||
|
||||
# 动态课程奖励与动作惩罚衰减
|
||||
cfg.rewards.pop("terrain_level_bonus", None)
|
||||
cfg.rewards.pop("action_rate", None)
|
||||
cfg.rewards["action_rate_curriculum"] = RewardTermCfg(func=action_rate_curriculum_l2, weight=-0.005)
|
||||
cfg.rewards.pop("action_rate_curriculum", None)
|
||||
cfg.rewards["action_rate"].weight = -0.01
|
||||
|
||||
cfg.rewards["joint_torques"].weight = -1e-4
|
||||
cfg.rewards["joint_torques"].weight = -2.5e-5
|
||||
cfg.rewards["joint_power"] = RewardTermCfg(func=joint_power, weight=-2.0e-5)
|
||||
cfg.rewards.pop("joint_acc", None)
|
||||
cfg.rewards["leg_joint_acc_l2"] = RewardTermCfg(func=envs_mdp.joint_acc_l2, weight=-2.5e-7, params={"asset_cfg": SceneEntityCfg("robot", joint_names=(".*_hip_abduction_joint", ".*_hip_pitch_joint", ".*_knee_joint"))})
|
||||
cfg.rewards["wheel_joint_acc_l2"] = RewardTermCfg(func=envs_mdp.joint_acc_l2, weight=-2.5e-9, params={"asset_cfg": SceneEntityCfg("robot", joint_names=(".*_wheel_joint",))})
|
||||
|
||||
|
||||
cfg.rewards["joint_pos_limits"].weight = -0.2
|
||||
cfg.rewards["joint_pos_limits"].weight = -5.0
|
||||
cfg.rewards.pop("leg_motion_penalty", None)
|
||||
cfg.rewards["is_terminated"].weight = 0.0
|
||||
cfg.rewards.pop("leg_symmetry", None)
|
||||
@@ -487,7 +472,7 @@ def rough_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
|
||||
cfg.rewards.pop("joint_deviation_l2", None)
|
||||
cfg.rewards["joint_pos_penalty"] = RewardTermCfg(
|
||||
func=joint_pos_penalty,
|
||||
weight=-0.8,
|
||||
weight=-1.0,
|
||||
params={
|
||||
"stand_still_scale": 5.0,
|
||||
"velocity_threshold": 0.5,
|
||||
@@ -502,37 +487,29 @@ def rough_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
|
||||
weight=0.1,
|
||||
params={"command_name": "twist", "sensor_name": "feet_ground_contact"}
|
||||
)
|
||||
cfg.rewards["feet_air_time"].weight = 0.0
|
||||
cfg.rewards["upward"] = RewardTermCfg(func=upward, weight=1.0)
|
||||
|
||||
cfg.rewards["base_height_l2"].weight = -0.5
|
||||
cfg.rewards["base_height_l2"].weight = 0.0
|
||||
cfg.rewards["base_height_l2"].params["target_height"] = 0.40
|
||||
cfg.rewards["base_height_l2"].params["sensor_cfg"] = SceneEntityCfg("height_scanner")
|
||||
|
||||
# 恢复机身碰撞惩罚为-1.0，逼迫机器人高抬腿跨越障碍，防止拖地
|
||||
cfg.rewards.pop("body_collision", None)
|
||||
cfg.rewards["body_collision"] = RewardTermCfg(func=velocity_mdp.self_collision_cost, weight=-1.0, params={"sensor_name": "body_collision"})
|
||||
cfg.rewards["undesired_contacts"] = RewardTermCfg(func=undesired_contacts, weight=-1.0, params={"sensor_name": "body_collision", "threshold": 1.0})
|
||||
cfg.rewards["contact_forces"] = RewardTermCfg(func=contact_forces, weight=-1.5e-4, params={"sensor_name": "feet_ground_contact", "threshold": 100.0})
|
||||
|
||||
# 🌟 严厉惩罚机身/胸部碰撞（防止硬撞高墙），逼迫机器人学会用前轮触墙并主动抬腿攀爬的“触觉反射”
|
||||
cfg.rewards["base_collision"] = RewardTermCfg(
|
||||
func=velocity_mdp.self_collision_cost,
|
||||
weight=-5.0,
|
||||
params={"sensor_name": "base_ground_contact"}
|
||||
)
|
||||
|
||||
# 彻底移除机身俯仰约束，允许机器人抬头爬高? cfg.rewards.pop("flat_orientation", None)
|
||||
|
||||
cfg.rewards.pop("feet_air_time", None)
|
||||
|
||||
# Remove non-applicable rewards
|
||||
for key in ("wheel_roll_tracking", "wheel_contact_bonus", "body_ang_vel"):
|
||||
for key in ("wheel_roll_tracking", "wheel_contact_bonus", "body_ang_vel", "terrain_level_bonus", "action_rate_curriculum"):
|
||||
cfg.rewards.pop(key, None)
|
||||
|
||||
cfg.episode_length_s = 30.0
|
||||
cfg.sim = SimulationCfg(contact_sensor_maxmatch=128, mujoco=MujocoCfg(impratio=100, cone="elliptic", ccd_iterations=80))
|
||||
cfg.episode_length_s = 20.0
|
||||
cfg.sim = SimulationCfg(contact_sensor_maxmatch=128, mujoco=MujocoCfg(timestep=0.005, impratio=100, cone="elliptic", ccd_iterations=80))
|
||||
|
||||
# 移除 orientation 终止，允许机器人翻倒以学习回复
|
||||
cfg.terminations.pop("bad_orientation", None)
|
||||
# 移除 base_ground_contact 终止，越障时机身会碰到障碍物
|
||||
cfg.terminations.pop("base_ground_contact", None)
|
||||
|
||||
cfg.seed = 42
|
||||
if cfg.scene.terrain is not None:
|
||||
cfg.scene.terrain.num_envs = 2048
|
||||
@@ -663,7 +640,7 @@ def crawl_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
|
||||
cfg.rewards.pop(key, None)
|
||||
|
||||
cfg.episode_length_s = 30.0
|
||||
cfg.sim = SimulationCfg(contact_sensor_maxmatch=128, mujoco=MujocoCfg(impratio=100, cone="elliptic", ccd_iterations=80))
|
||||
cfg.sim = SimulationCfg(contact_sensor_maxmatch=128, mujoco=MujocoCfg(timestep=0.005, impratio=100, cone="elliptic", ccd_iterations=80))
|
||||
|
||||
# Loosen orientation bad threshold to 80 degrees for steep crawling tilts
|
||||
cfg.terminations["bad_orientation"].params["limit_angle"] = math.radians(80.0)
|
||||
|
||||
@@ -33,10 +33,10 @@ def rough_ppo_runner_cfg() -> RslRlOnPolicyRunnerCfg:
|
||||
value_loss_coef=1.0,
|
||||
use_clipped_value_loss=True,
|
||||
clip_param=0.2,
|
||||
entropy_coef=0.002,
|
||||
entropy_coef=0.001,#第一轮为0.003
|
||||
num_learning_epochs=5,
|
||||
num_mini_batches=4,
|
||||
learning_rate=2.0e-4,
|
||||
learning_rate=8.0e-4,
|
||||
schedule="adaptive",
|
||||
gamma=0.99,
|
||||
lam=0.95,
|
||||
@@ -44,9 +44,9 @@ def rough_ppo_runner_cfg() -> RslRlOnPolicyRunnerCfg:
|
||||
max_grad_norm=1.0,
|
||||
),
|
||||
experiment_name="robot_rough",
|
||||
save_interval=50,
|
||||
save_interval=100,
|
||||
num_steps_per_env=24,
|
||||
max_iterations=15_000,
|
||||
max_iterations=20_000,
|
||||
)
|
||||
|
||||
|
||||
@@ -55,5 +55,3 @@ def crawl_ppo_runner_cfg() -> RslRlOnPolicyRunnerCfg:
|
||||
cfg.experiment_name = "robot_crawl"
|
||||
return cfg
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -28,6 +28,7 @@ def track_linear_velocity(
|
||||
actual = asset.data.root_link_lin_vel_b
|
||||
xy_error = torch.sum(torch.square(command[:, :2] - actual[:, :2]), dim=1)
|
||||
reward = torch.exp(-xy_error / std**2)
|
||||
reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
@@ -45,6 +46,7 @@ def track_angular_velocity(
|
||||
actual = asset.data.root_link_ang_vel_b
|
||||
z_error = torch.square(command[:, 2] - actual[:, 2])
|
||||
reward = torch.exp(-z_error / std**2)
|
||||
reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
@@ -94,6 +96,7 @@ def base_height_l2(
|
||||
error = root_z - target_height
|
||||
|
||||
reward = torch.square(error)
|
||||
reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
@@ -304,6 +307,7 @@ def stand_still(
|
||||
angular_norm = torch.abs(command[:, 2])
|
||||
inactive = (linear_norm + angular_norm < command_threshold).float()
|
||||
reward = cost * inactive
|
||||
reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
def hip_deviation(
|
||||
@@ -355,6 +359,20 @@ def lin_vel_z_l2(
|
||||
asset_cfg = SceneEntityCfg("robot")
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
reward = torch.square(asset.data.root_link_lin_vel_b[:, 2])
|
||||
reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def ang_vel_xy_l2(
|
||||
env: ManagerBasedRlEnv,
|
||||
asset_cfg: SceneEntityCfg | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Penalize xy-axis base angular velocity using the go2w kernel."""
|
||||
if asset_cfg is None:
|
||||
asset_cfg = SceneEntityCfg("robot")
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
reward = torch.sum(torch.square(asset.data.root_link_ang_vel_b[:, :2]), dim=1)
|
||||
reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
@@ -650,6 +668,8 @@ def feet_contact_without_cmd(env, command_name: str, sensor_name: str) -> torch.
|
||||
linear_norm = torch.norm(cmd[:, :2], dim=1)
|
||||
angular_norm = torch.abs(cmd[:, 2])
|
||||
reward *= (linear_norm + angular_norm) < 0.1
|
||||
asset = env.scene["robot"]
|
||||
reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
def joint_pos_penalty(
|
||||
@@ -674,6 +694,7 @@ def joint_pos_penalty(
|
||||
running_reward,
|
||||
stand_still_scale * running_reward,
|
||||
)
|
||||
reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
def joint_mirror(env, mirror_joints: list[list[str]], asset_cfg: SceneEntityCfg | None = None) -> torch.Tensor:
|
||||
@@ -694,6 +715,50 @@ def joint_mirror(env, mirror_joints: list[list[str]], asset_cfg: SceneEntityCfg
|
||||
)
|
||||
reward += diff
|
||||
reward *= 1 / len(mirror_joints) if len(mirror_joints) > 0 else 0
|
||||
reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def undesired_contacts(
|
||||
env: ManagerBasedRlEnv,
|
||||
sensor_name: str,
|
||||
threshold: float = 1.0,
|
||||
) -> torch.Tensor:
|
||||
"""Penalize non-wheel contacts above a force threshold."""
|
||||
from mjlab.sensor import ContactSensor
|
||||
|
||||
sensor: ContactSensor = env.scene[sensor_name]
|
||||
data = sensor.data
|
||||
if data.force_history is not None:
|
||||
force_mag = torch.norm(data.force_history, dim=-1)
|
||||
is_contact = torch.max(force_mag, dim=2)[0] > threshold
|
||||
else:
|
||||
force_mag = torch.norm(data.force, dim=-1)
|
||||
is_contact = force_mag > threshold
|
||||
reward = torch.sum(is_contact, dim=1).float()
|
||||
asset = env.scene["robot"]
|
||||
reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
def contact_forces(
|
||||
env: ManagerBasedRlEnv,
|
||||
sensor_name: str,
|
||||
threshold: float = 100.0,
|
||||
) -> torch.Tensor:
|
||||
"""Penalize foot contact forces above threshold."""
|
||||
from mjlab.sensor import ContactSensor
|
||||
|
||||
sensor: ContactSensor = env.scene[sensor_name]
|
||||
data = sensor.data
|
||||
if data.force_history is not None:
|
||||
force_mag = torch.norm(data.force_history, dim=-1)
|
||||
peak_force = torch.max(force_mag, dim=2)[0]
|
||||
else:
|
||||
peak_force = torch.norm(data.force, dim=-1)
|
||||
reward = torch.sum(torch.clamp(peak_force - threshold, min=0.0), dim=1)
|
||||
asset = env.scene["robot"]
|
||||
reward *= torch.clamp(-asset.data.projected_gravity_b[:, 2], 0.0, 0.7) / 0.7
|
||||
return reward
|
||||
|
||||
|
||||
@@ -705,6 +770,14 @@ def upward(env, asset_cfg=None):
|
||||
reward = torch.square(1 - asset.data.projected_gravity_b[:, 2])
|
||||
return reward
|
||||
|
||||
|
||||
def joint_power(env: ManagerBasedRlEnv, asset_cfg: SceneEntityCfg | None = None) -> torch.Tensor:
|
||||
"""Penalty for total joint mechanical power: sum(|tau * dq|)."""
|
||||
if asset_cfg is None:
|
||||
asset_cfg = SceneEntityCfg("robot")
|
||||
asset: Entity = env.scene[asset_cfg.name]
|
||||
return torch.sum(torch.abs(asset.data.qfrc_actuator * asset.data.joint_vel), dim=1)
|
||||
|
||||
def upright_roll_only(env, asset_cfg=None):
|
||||
if asset_cfg is None:
|
||||
from mjlab.envs.manager_based_rl_env import SceneEntityCfg
|
||||
|
||||
@@ -1,41 +0,0 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
|
||||
|
||||
"""
|
||||
HimLoco RSL-RL implementation with history-informed models.
|
||||
"""
|
||||
|
||||
# Export HIM implementations
|
||||
from .algorithms.him_ppo import *
|
||||
from .modules.him_actor_critic import *
|
||||
from .modules.him_estimator import *
|
||||
from .storage.him_rollout_storage import *
|
||||
from .runners.him_on_policy_runner import *
|
||||
from .wrappers import *
|
||||
@@ -1,31 +0,0 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
|
||||
|
||||
from .him_ppo import HIMPPO
|
||||
@@ -1,192 +0,0 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
|
||||
from ..modules import HIMActorCritic
|
||||
from ..storage import HIMRolloutStorage
|
||||
|
||||
class HIMPPO:
|
||||
actor_critic: HIMActorCritic
|
||||
def __init__(self,
|
||||
actor_critic,
|
||||
num_learning_epochs=1,
|
||||
num_mini_batches=1,
|
||||
clip_param=0.2,
|
||||
gamma=0.998,
|
||||
lam=0.95,
|
||||
value_loss_coef=1.0,
|
||||
entropy_coef=0.0,
|
||||
learning_rate=1e-3,
|
||||
max_grad_norm=1.0,
|
||||
use_clipped_value_loss=True,
|
||||
schedule="fixed",
|
||||
desired_kl=0.01,
|
||||
device='cpu',
|
||||
):
|
||||
|
||||
self.device = device
|
||||
|
||||
self.desired_kl = desired_kl
|
||||
self.schedule = schedule
|
||||
self.learning_rate = learning_rate
|
||||
|
||||
# PPO components
|
||||
self.actor_critic = actor_critic
|
||||
self.actor_critic.to(self.device)
|
||||
self.storage = None # initialized later
|
||||
self.optimizer = optim.Adam(self.actor_critic.parameters(), lr=learning_rate)
|
||||
self.transition = HIMRolloutStorage.Transition()
|
||||
|
||||
# PPO parameters
|
||||
self.clip_param = clip_param
|
||||
self.num_learning_epochs = num_learning_epochs
|
||||
self.num_mini_batches = num_mini_batches
|
||||
self.value_loss_coef = value_loss_coef
|
||||
self.entropy_coef = entropy_coef
|
||||
self.gamma = gamma
|
||||
self.lam = lam
|
||||
self.max_grad_norm = max_grad_norm
|
||||
self.use_clipped_value_loss = use_clipped_value_loss
|
||||
|
||||
def init_storage(self, num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, action_shape):
|
||||
self.storage = HIMRolloutStorage(num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, action_shape, self.device)
|
||||
|
||||
def test_mode(self):
|
||||
self.actor_critic.test()
|
||||
|
||||
def train_mode(self):
|
||||
self.actor_critic.train()
|
||||
|
||||
def act(self, obs, critic_obs):
|
||||
# Compute the actions and values
|
||||
self.transition.actions = self.actor_critic.act(obs).detach()
|
||||
self.transition.values = self.actor_critic.evaluate(critic_obs).detach()
|
||||
self.transition.actions_log_prob = self.actor_critic.get_actions_log_prob(self.transition.actions).detach()
|
||||
self.transition.action_mean = self.actor_critic.action_mean.detach()
|
||||
self.transition.action_sigma = self.actor_critic.action_std.detach()
|
||||
# need to record obs and critic_obs before env.step()
|
||||
self.transition.observations = obs
|
||||
self.transition.critic_observations = critic_obs
|
||||
return self.transition.actions
|
||||
|
||||
def process_env_step(self, rewards, dones, infos, next_critic_obs):
|
||||
self.transition.next_critic_observations = next_critic_obs.clone()
|
||||
self.transition.rewards = rewards.clone()
|
||||
self.transition.dones = dones
|
||||
# Bootstrapping on time outs
|
||||
if 'time_outs' in infos:
|
||||
self.transition.rewards += self.gamma * torch.squeeze(self.transition.values * infos['time_outs'].unsqueeze(1).to(self.device), 1)
|
||||
|
||||
# Record the transition
|
||||
self.storage.add_transitions(self.transition)
|
||||
self.transition.clear()
|
||||
self.actor_critic.reset(dones)
|
||||
|
||||
def compute_returns(self, last_critic_obs):
|
||||
last_values= self.actor_critic.evaluate(last_critic_obs).detach()
|
||||
self.storage.compute_returns(last_values, self.gamma, self.lam)
|
||||
|
||||
def update(self):
|
||||
mean_value_loss = 0
|
||||
mean_surrogate_loss = 0
|
||||
mean_estimation_loss = 0
|
||||
mean_swap_loss = 0
|
||||
|
||||
generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs)
|
||||
|
||||
for obs_batch, critic_obs_batch, actions_batch, next_critic_obs_batch, target_values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, \
|
||||
old_mu_batch, old_sigma_batch in generator:
|
||||
|
||||
self.actor_critic.act(obs_batch)
|
||||
actions_log_prob_batch = self.actor_critic.get_actions_log_prob(actions_batch)
|
||||
value_batch = self.actor_critic.evaluate(critic_obs_batch)
|
||||
mu_batch = self.actor_critic.action_mean
|
||||
sigma_batch = self.actor_critic.action_std
|
||||
entropy_batch = self.actor_critic.entropy
|
||||
|
||||
# KL
|
||||
if self.desired_kl != None and self.schedule == 'adaptive':
|
||||
with torch.inference_mode():
|
||||
kl = torch.sum(
|
||||
torch.log(sigma_batch / old_sigma_batch + 1.e-5) + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) / (2.0 * torch.square(sigma_batch)) - 0.5, axis=-1)
|
||||
kl_mean = torch.mean(kl)
|
||||
|
||||
if kl_mean > self.desired_kl * 2.0:
|
||||
self.learning_rate = max(1e-5, self.learning_rate / 1.5)
|
||||
elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0:
|
||||
self.learning_rate = min(1e-2, self.learning_rate * 1.5)
|
||||
|
||||
for param_group in self.optimizer.param_groups:
|
||||
param_group['lr'] = self.learning_rate
|
||||
|
||||
#Estimator Update
|
||||
estimation_loss, swap_loss = self.actor_critic.estimator.update(obs_batch, next_critic_obs_batch, lr=self.learning_rate)
|
||||
|
||||
# Surrogate loss
|
||||
ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch))
|
||||
surrogate = -torch.squeeze(advantages_batch) * ratio
|
||||
surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp(ratio, 1.0 - self.clip_param,
|
||||
1.0 + self.clip_param)
|
||||
surrogate_loss = torch.max(surrogate, surrogate_clipped).mean()
|
||||
|
||||
# Value function loss
|
||||
if self.use_clipped_value_loss:
|
||||
value_clipped = target_values_batch + (value_batch - target_values_batch).clamp(-self.clip_param,
|
||||
self.clip_param)
|
||||
value_losses = (value_batch - returns_batch).pow(2)
|
||||
value_losses_clipped = (value_clipped - returns_batch).pow(2)
|
||||
value_loss = torch.max(value_losses, value_losses_clipped).mean()
|
||||
else:
|
||||
value_loss = (returns_batch - value_batch).pow(2).mean()
|
||||
|
||||
loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean()
|
||||
|
||||
# Gradient step
|
||||
self.optimizer.zero_grad()
|
||||
loss.backward()
|
||||
nn.utils.clip_grad_norm_(self.actor_critic.parameters(), self.max_grad_norm)
|
||||
self.optimizer.step()
|
||||
|
||||
mean_value_loss += value_loss.item()
|
||||
mean_surrogate_loss += surrogate_loss.item()
|
||||
mean_estimation_loss += estimation_loss
|
||||
mean_swap_loss += swap_loss
|
||||
|
||||
num_updates = self.num_learning_epochs * self.num_mini_batches
|
||||
mean_value_loss /= num_updates
|
||||
mean_surrogate_loss /= num_updates
|
||||
mean_estimation_loss /= num_updates
|
||||
mean_swap_loss /= num_updates
|
||||
self.storage.clear()
|
||||
|
||||
return mean_value_loss, mean_surrogate_loss, estimation_loss, swap_loss
|
||||
@@ -1,31 +0,0 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
#
|
||||
# Redistribution and use in source and binary forms, with or without
|
||||
# modification, are permitted provided that the following conditions are met:
|
||||
#
|
||||
# 1. Redistributions of source code must retain the above copyright notice, this
|
||||
# list of conditions and the following disclaimer.
|
||||
#
|
||||
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
# this list of conditions and the following disclaimer in the documentation
|
||||
# and/or other materials provided with the distribution.
|
||||
#
|
||||
# 3. Neither the name of the copyright holder nor the names of its
|
||||
# contributors may be used to endorse or promote products derived from
|
||||
# this software without specific prior written permission.
|
||||
#
|
||||
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
#
|
||||
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
|
||||
|
||||
from .rl_cfg import *
|
||||
@@ -1,180 +0,0 @@
|
||||
from dataclasses import MISSING
|
||||
from isaaclab.utils import configclass
|
||||
from typing import Literal
|
||||
|
||||
|
||||
@configclass
|
||||
class HIMBaseRunnerCfg:
|
||||
"""Base configuration of the runner."""
|
||||
|
||||
seed: int = 1
|
||||
"""The seed for the experiment. Default is 1."""
|
||||
|
||||
device: str = "cuda:0"
|
||||
"""The device for the rl-agent. Default is cuda:0."""
|
||||
|
||||
num_steps_per_env: int = MISSING
|
||||
"""The number of steps per environment per update."""
|
||||
|
||||
max_iterations: int = MISSING
|
||||
"""The maximum number of iterations."""
|
||||
|
||||
empirical_normalization: bool | None = None
|
||||
"""This parameter is deprecated and will be removed in the future.
|
||||
|
||||
Use `actor_obs_normalization` and `critic_obs_normalization` instead.
|
||||
"""
|
||||
|
||||
# obs_groups: dict[str, list[str]] = MISSING
|
||||
# """A mapping from observation groups to observation sets.
|
||||
|
||||
# The keys of the dictionary are predefined observation sets used by the underlying algorithm
|
||||
# and values are lists of observation groups provided by the environment.
|
||||
|
||||
# For instance, if the environment provides a dictionary of observations with groups "policy", "images",
|
||||
# and "privileged", these can be mapped to algorithmic observation sets as follows:
|
||||
|
||||
# .. code-block:: python
|
||||
|
||||
# obs_groups = {
|
||||
# "policy": ["policy", "images"],
|
||||
# "critic": ["policy", "privileged"],
|
||||
# }
|
||||
|
||||
# This way, the policy will receive the "policy" and "images" observations, and the critic will
|
||||
# receive the "policy" and "privileged" observations.
|
||||
|
||||
# For more details, please check ``vec_env.py`` in the rsl_rl library.
|
||||
# """
|
||||
|
||||
# clip_actions: float | None = None
|
||||
# """The clipping value for actions. If None, then no clipping is done. Defaults to None.
|
||||
|
||||
# .. note::
|
||||
# This clipping is performed inside the :class:`RslRlVecEnvWrapper` wrapper.
|
||||
# """
|
||||
|
||||
save_interval: int = MISSING
|
||||
"""The number of iterations between saves."""
|
||||
|
||||
experiment_name: str = MISSING
|
||||
"""The experiment name."""
|
||||
|
||||
run_name: str = ""
|
||||
"""The run name. Default is empty string.
|
||||
|
||||
The name of the run directory is typically the time-stamp at execution. If the run name is not empty,
|
||||
then it is appended to the run directory's name, i.e. the logging directory's name will become
|
||||
``{time-stamp}_{run_name}``.
|
||||
"""
|
||||
|
||||
logger: Literal["tensorboard", "neptune", "wandb"] = "tensorboard"
|
||||
"""The logger to use. Default is tensorboard."""
|
||||
|
||||
neptune_project: str = "isaaclab"
|
||||
"""The neptune project name. Default is "isaaclab"."""
|
||||
|
||||
wandb_project: str = "isaaclab"
|
||||
"""The wandb project name. Default is "isaaclab"."""
|
||||
|
||||
resume: bool = False
|
||||
"""Whether to resume a previous training. Default is False.
|
||||
|
||||
This flag will be ignored for distillation.
|
||||
"""
|
||||
|
||||
load_run: str = ".*"
|
||||
"""The run directory to load. Default is ".*" (all).
|
||||
|
||||
If regex expression, the latest (alphabetical order) matching run will be loaded.
|
||||
"""
|
||||
|
||||
load_checkpoint: str = "model_.*.pt"
|
||||
"""The checkpoint file to load. Default is ``"model_.*.pt"`` (all).
|
||||
|
||||
If regex expression, the latest (alphabetical order) matching file will be loaded.
|
||||
"""
|
||||
|
||||
|
||||
@configclass
|
||||
class HIMPPOActorCriticCfg:
|
||||
"""Configuration of the HIM PPO actor-critic."""
|
||||
|
||||
actor_hidden_dims: list[int] = [512, 256, 128]
|
||||
"""The hidden dimensions of the actor network."""
|
||||
|
||||
critic_hidden_dims: list[int] = [512, 256, 128]
|
||||
"""The hidden dimensions of the critic network."""
|
||||
|
||||
activation: str = "elu"
|
||||
"""The activation function to use. Default is 'elu'."""
|
||||
|
||||
init_noise_std: float = 1.0
|
||||
"""The initial noise standard deviation for the actor. Default is 1.0."""
|
||||
|
||||
normalize_obs: bool = False
|
||||
"""Whether to normalize observations. Default is False."""
|
||||
|
||||
@configclass
|
||||
class HIMPPPOAlgorithmCfg:
|
||||
"""Configuration of the HIM PPO algorithm."""
|
||||
num_learning_epochs: int = 1
|
||||
"""The number of learning epochs per update. Default is 1."""
|
||||
|
||||
num_mini_batches: int = 1
|
||||
"""The number of mini-batches per update. Default is 1."""
|
||||
|
||||
clip_param: float = 0.2
|
||||
"""The clipping parameter for PPO. Default is 0.2."""
|
||||
|
||||
gamma: float = 0.998
|
||||
"""The discount factor. Default is 0.998."""
|
||||
|
||||
lam: float = 0.95
|
||||
"""The GAE lambda parameter. Default is 0.95."""
|
||||
|
||||
value_loss_coef: float = 1.0
|
||||
"""The coefficient for the value loss. Default is 1.0."""
|
||||
|
||||
entropy_coef: float = 0.0
|
||||
"""The coefficient for the entropy bonus. Default is 0.0."""
|
||||
|
||||
learning_rate: float = 1.0e-3
|
||||
"""The learning rate. Default is 1.0e-3."""
|
||||
|
||||
max_grad_norm: float = 1.0
|
||||
"""The maximum gradient norm for clipping. Default is 1.0."""
|
||||
|
||||
use_clipped_value_loss: bool = True
|
||||
"""Whether to use clipped value loss. Default is True."""
|
||||
|
||||
schedule: str = "fixed"
|
||||
"""The learning rate schedule. Default is 'fixed'."""
|
||||
|
||||
desired_kl: float = 0.01
|
||||
"""The desired KL divergence for adaptive learning rate. Default is 0.01."""
|
||||
|
||||
@configclass
|
||||
class HIMOnPolicyRunnerCfg(HIMBaseRunnerCfg):
|
||||
"""Configuration of the runner for on-policy algorithms."""
|
||||
|
||||
class_name: str = "HIMOnPolicyRunner"
|
||||
"""The runner class name. Default is OnPolicyRunner."""
|
||||
|
||||
policy_class_name: str = "HIMActorCritic"
|
||||
"""The policy class name. Default is HIMActorCritic."""
|
||||
|
||||
algorithm_class_name: str = "HIMPPO"
|
||||
"""The algorithm class name. Default is HIMPPO."""
|
||||
|
||||
policy: HIMPPOActorCriticCfg = MISSING
|
||||
"""The policy configuration."""
|
||||
|
||||
algorithm: HIMPPPOAlgorithmCfg = MISSING
|
||||
"""The algorithm configuration."""
|
||||
|
||||
history_length: int = 0
|
||||
"""Number of historical time steps to stack with current observation (0 means current only). Default is 0."""
|
||||
|
||||
privileged_history_length: int = 0
|
||||
"""Number of historical time steps to stack with current privileged observation. Default is 0."""
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user