Init
This commit is contained in:
wty-yy
168
.gitignore
vendored
Normal file
168
.gitignore
vendored
Normal file
@@ -0,0 +1,168 @@
|
||||
*.pt
|
||||
MUJOCO_LOG.TXT
|
||||
*.mp4
|
||||
videos
|
||||
images
|
||||
mytest*
|
||||
data*
|
||||
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
share/python-wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
|
||||
# PyInstaller
|
||||
# Usually these files are written by a python script from a template
|
||||
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Installer logs
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.nox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
*.py,cover
|
||||
.hypothesis/
|
||||
.pytest_cache/
|
||||
cover/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Django stuff:
|
||||
*.log
|
||||
local_settings.py
|
||||
db.sqlite3
|
||||
db.sqlite3-journal
|
||||
|
||||
# Flask stuff:
|
||||
instance/
|
||||
.webassets-cache
|
||||
|
||||
# Scrapy stuff:
|
||||
.scrapy
|
||||
|
||||
# Sphinx documentation
|
||||
docs/_build/
|
||||
|
||||
# PyBuilder
|
||||
.pybuilder/
|
||||
target/
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# IPython
|
||||
profile_default/
|
||||
ipython_config.py
|
||||
|
||||
# pyenv
|
||||
# For a library or package, you might want to ignore these files since the code is
|
||||
# intended to run in multiple environments; otherwise, check them in:
|
||||
# .python-version
|
||||
|
||||
# pipenv
|
||||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
||||
# install all needed dependencies.
|
||||
#Pipfile.lock
|
||||
|
||||
# poetry
|
||||
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
||||
#poetry.lock
|
||||
|
||||
# pdm
|
||||
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
||||
#pdm.lock
|
||||
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
|
||||
# in version control.
|
||||
# https://pdm.fming.dev/latest/usage/project/#working-with-version-control
|
||||
.pdm.toml
|
||||
.pdm-python
|
||||
.pdm-build/
|
||||
|
||||
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
||||
__pypackages__/
|
||||
|
||||
# Celery stuff
|
||||
celerybeat-schedule
|
||||
celerybeat.pid
|
||||
|
||||
# SageMath parsed files
|
||||
*.sage.py
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.venv
|
||||
venv/
|
||||
ENV/
|
||||
venv.bak/
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
.dmypy.json
|
||||
dmypy.json
|
||||
|
||||
# Pyre type checker
|
||||
.pyre/
|
||||
|
||||
# pytype static type analyzer
|
||||
.pytype/
|
||||
|
||||
# Cython debug symbols
|
||||
cython_debug/
|
||||
|
||||
# Editors
|
||||
.idea/
|
||||
.vscode
|
||||
|
||||
# logs
|
||||
logs
|
||||
29
LICENSE
29
LICENSE
@@ -1,6 +1,6 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2025 wty
|
||||
Copyright (c) 2025 Wu Tianyang
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
@@ -19,3 +19,30 @@ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
|
||||
This project is based on or incorporates code from https://github.com/unitreerobotics/unitree_rl_gym,
|
||||
which is licensed under the BSD 3-Clause License:
|
||||
|
||||
BSD 3-Clause License
|
||||
Copyright (c) 2016-2023 HangZhou YuShu TECHNOLOGY CO.,LTD. ("Unitree Robotics")
|
||||
All rights reserved.
|
||||
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.
|
||||
|
||||
152
README.md
Normal file
152
README.md
Normal file
@@ -0,0 +1,152 @@
|
||||
<div align="center">
|
||||
<h1 align="center">Go2 RL GYM</h1>
|
||||
<p align="center">
|
||||
<a href="README_en.md">🌎 English</a> | <span>🇨🇳 中文</span>
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<p align="center">
|
||||
<strong>本仓库基于<a href="https://github.com/unitreerobotics/unitree_rl_gym">unitree_rl_gym</a>，使用强化学习训练Go2机器狗。</strong>
|
||||
</p>
|
||||
|
||||
<div align="center">
|
||||
|
||||
| <div align="center"> Isaac Gym </div> | <div align="center"> Mujoco </div> | <div align="center"> Physical </div> |
|
||||
|--- | --- | --- |
|
||||
| TODO | TODO | TODO |
|
||||
|
||||
</div>
|
||||
|
||||
## 📦 安装配置
|
||||
|
||||
安装和配置步骤请参考 [setup.md](/doc/setup_zh.md)
|
||||
|
||||
## 🛠️ 使用指南
|
||||
|
||||
### 1. 训练
|
||||
|
||||
运行以下命令进行训练：
|
||||
|
||||
```bash
|
||||
python legged_gym/scripts/train.py --task=xxx
|
||||
```
|
||||
|
||||
#### ⚙️ 参数说明
|
||||
- `--task`: 必选参数，值可选(go2, go2_cts, go2_moe_cts)
|
||||
- `--headless`: 默认启动图形界面，设为 true 时不渲染图形界面（效率更高）
|
||||
- `--resume`: 从日志中选择 checkpoint 继续训练
|
||||
- `--experiment_name`: 运行/加载的 experiment 名称
|
||||
- `--run_name`: 运行/加载的 run 名称
|
||||
- `--load_run`: 加载运行的名称，默认加载最后一次运行
|
||||
- `--checkpoint`: checkpoint 编号，默认加载最新一次文件
|
||||
- `--num_envs`: 并行训练的环境个数
|
||||
- `--seed`: 随机种子
|
||||
- `--max_iterations`: 训练的最大迭代次数
|
||||
- `--sim_device`: 仿真计算设备，指定 CPU 为 `--sim_device=cpu`
|
||||
- `--rl_device`: 强化学习计算设备，指定 CPU 为 `--rl_device=cpu`
|
||||
|
||||
**默认保存训练结果**：`logs/<experiment_name>/<date_time>_<run_name>/model_<iteration>.pt`
|
||||
|
||||
---
|
||||
|
||||
### 2. Play
|
||||
|
||||
如果想要在 Gym 中查看训练效果，可以运行以下命令：
|
||||
|
||||
```bash
|
||||
python legged_gym/scripts/play.py --task=xxx
|
||||
```
|
||||
|
||||
**说明**：
|
||||
|
||||
- Play 启动参数为随机地形，难度在7到9之间。
|
||||
- 默认加载实验文件夹最新训练的一个模型。
|
||||
- 可通过 `experiment_name` 和 `checkpoint` 指定其他模型，例如
|
||||
```bash
|
||||
python legged_gym/scripts/play.py --task=go2_cts --num_envs 100 --experiment_name go2_cts_hard_terrain --checkpoint 100000
|
||||
```
|
||||
|
||||
#### 💾 导出网络
|
||||
|
||||
Play 会导出 Actor 网络，保存于 `logs/{experiment_name}/exported/policies` 中：
|
||||
- `policy.pt`: torch script模型，用于Sim2Sim。
|
||||
- `policy.onnx`: onnx模型，用于Sim2Real。
|
||||
- `policy.pkl`: 模型权重。
|
||||
|
||||
#### Play 效果
|
||||
|
||||
| Go2 |
|
||||
|--- |
|
||||
| TODO |
|
||||
|
||||
---
|
||||
|
||||
### 3. Sim2Sim (Mujoco)
|
||||
|
||||
支持在 Mujoco 仿真器中运行 Sim2Sim：
|
||||
|
||||
```bash
|
||||
python deploy/deploy_mujoco/deploy_go2.py
|
||||
```
|
||||
|
||||
如果有xbox协议的手柄接入主机，自动切换为手柄控制，否则只会保持默认指令前进。
|
||||
|
||||
- 替换网络模型：默认模型位于 `deploy/pre_train/go2/go2_cts_150k.pt`；自己训练模型保存于`logs/{experiment_name}/exported/policies/policy.pt`，只需替换 yaml 配置文件中 `policy_path`。
|
||||
- 替换环境地形：默认地形为 `resources/robots/go2/stairs.xml`，其他可选地形，平地 `flat.xml`，赛道 `race_track.xml`，地形使用[terrain_generator.py](resources/robots/go2/terrain_generator.py)生成，参考[unitree_mujoco/terrain_tool](https://github.com/unitreerobotics/unitree_mujoco/tree/main/terrain_tool)。
|
||||
|
||||
#### 运行效果
|
||||
|
||||
| 平地 | 台阶 | 赛道 |
|
||||
|--- | --- | --- |
|
||||
| | | |
|
||||
|
||||
---
|
||||
|
||||
### 4. Sim2Real
|
||||
|
||||
#### 4.1 Python实物部署 （需要安装 unitree_sdk2_python）
|
||||
|
||||
先用app进入设备→服务状态→点击运控服务，关闭`mcf`，打开`ota_box`服务。
|
||||
|
||||
假设和下位机连接的网卡名称为`eth0`，执行
|
||||
```bash
|
||||
cd deploy/deploy_real
|
||||
python deploy_real_go2.py eth0
|
||||
```
|
||||
`start`站立，`A`启动控制
|
||||
|
||||
#### 4.2 C++实物部署（需要安装 unitree_cpp_deploy）
|
||||
|
||||
参考[unitree_cpp_deploy](https://github.com/wty-yy-mini/unitree_cpp_deploy)使用说明。
|
||||
|
||||
#### 运行效果
|
||||
|
||||
| Python部署 | C++部署 |
|
||||
| --- | --- |
|
||||
| | |
|
||||
|
||||
---
|
||||
|
||||
## 🎉 致谢
|
||||
|
||||
本仓库开发离不开以下开源项目的支持与贡献，特此感谢：
|
||||
|
||||
- [unitree_rl_gym](https://github.com/unitreerobotics/unitree_rl_gym)：宇树机器人强化学习训练基础框架。
|
||||
- [legged\_gym](https://github.com/leggedrobotics/legged_gym)：构建基础训练环境。
|
||||
- [rsl\_rl](https://github.com/leggedrobotics/rsl_rl.git)：强化学习算法实现。
|
||||
- [mujoco](https://github.com/google-deepmind/mujoco.git)：提供强大CPU仿真功能。
|
||||
- [unitree\_sdk2\_python](https://github.com/unitreerobotics/unitree_sdk2_python.git)：实物部署硬件Python通信接口。
|
||||
- [unitree_sdk2](https://github.com/unitreerobotics/unitree_sdk2)：实物部署硬件C++通信接口。
|
||||
|
||||
本仓库实现包含以下论文，特此感谢：
|
||||
- [CTS: Concurrent Teacher-Student Reinforcement Learning for Legged Locomotion](https://arxiv.org/pdf/2405.10830)
|
||||
|
||||
---
|
||||
|
||||
## 🔖 许可证
|
||||
|
||||
新增内容根据 [MIT License](./LICENSE) 授权，原仓库unitree_rl_gym根据 [BSD 3-Clause License](./LICENSE) 授权。
|
||||
|
||||
详情请阅读完整 [LICENSE 文件](./LICENSE)。
|
||||
|
||||
|
||||
11
README_en.md
Normal file
11
README_en.md
Normal file
@@ -0,0 +1,11 @@
|
||||
<div align="center">
|
||||
<h1 align="center">Go2 RL GYM</h1>
|
||||
<p align="center">
|
||||
<span> 🌎English </span> | <a href="README.md"> 🇨🇳中文 </a>
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<p align="center">
|
||||
<strong>This is a repository for reinforcement learning implementation based on Unitree Go2. Based on <a href="https://github.com/unitreerobotics/unitree_rl_gym">unitree_rl_gym</a>.
|
||||
</strong>
|
||||
</p>
|
||||
126
UPDATE.md
Normal file
126
UPDATE.md
Normal file
@@ -0,0 +1,126 @@
|
||||
# 20251221
|
||||
1. 修改最大地形速度限制, y在所有地形上最大为1.0, z只有平地最大为2.0, x最大为2.0
|
||||
2. 上调难度9地形难度 (都是moe-cts 100k能通过的难度):
|
||||
1. Slope: 19.8度 -> 29.6度
|
||||
2. Stairs height: 0.212m -> 0.257m
|
||||
3. Obstacles height: 0.23m -> 0.275m
|
||||
# 20251204
|
||||
1. 添加不同地形的最大速度限制
|
||||
2. 加入`legs_distance`惩罚两腿横向距离过近的问题, 通过robogauge指标验证发现没有提升, 删除
|
||||
# 20251128
|
||||
1. go2.xml添加侧向和背面初始化 (需手动解开注释)
|
||||
2. cts, moe-cts算法中添加norm_layer选项, 将latent_dim提升到512, norm_layer改为simnorm (dim=8)
|
||||
3. 默认的layer norm还是用l2norm, latent_dim=32, 地形使用偏平地的地形, 因为开悟测评的分数还是更高, 平地也能走的更好
|
||||
# 20251127
|
||||
1. 修改deploy_go2.py支持xbox手柄指令输入
|
||||
2. 修改turn_over开关, 默认关闭
|
||||
# 20251126
|
||||
## v0.1.11
|
||||
1. 加入`turn_over`具体内容如下 (翻身roll为180度, 侧身roll为90度, 随机两个方向):
|
||||
1. `init_state.turn_over_proportions`: 支持初始化时, 处于翻身/侧身的概率分别为0.1/0.3
|
||||
2. `init_state.turn_over_init_heights`: 翻身/侧身的机的初始高度范围, 在该范围中随机采样
|
||||
3. `commands.turn_over_zero_time`: 翻身和侧身翻转过来预计用时, 在这段时间内给出的command都是全零
|
||||
4. `rewards.turn_over_roll_threshold`: 翻身奖励计算的roll阈值, 当roll.abs小于该阈值时转为计算原奖励
|
||||
5. `reward.turn_over_scales`: 翻身奖励配置, 目前只使用了`_reward_upright`奖励计算负base的重力投影z分量和z轴单位向量做差, 系数1.0
|
||||
2. 在环境数为4096时, 如果想达到8192的训练效果(尤其是上楼梯), 需要增大num_steps_per_env, 训4096时先加到48
|
||||
3. 发现部分地形分数较低: wave, stairs up, down, 修改地形比例配置, 增大这三个地形, 降低flat, slop
|
||||
4. 一开始从翻身开始是学不到, 只能从侧身开始学
|
||||
# 20251125
|
||||
## v0.1.10
|
||||
1. 添加robot_properties的低频域随机化更新 (后来发现无法使用, 会导致机器人本体信息错误, 删除)
|
||||
2. 消融修改num_steps_per_env=48,72,96
|
||||
# 20251124
|
||||
## v0.1.9
|
||||
1. 修改action delay从0/1进行延迟, 改为{0,5,10,15}ms上进行延迟
|
||||
2. 修改motor_offset和pd的域随机化从开始一次, 改为每次reset环境时随机
|
||||
# 20251120
|
||||
## v0.1.8
|
||||
1. friction为0-2有提升, 在瓷砖地上可以上楼梯了
|
||||
2. 将dynamic_sigma线速度最大值调到1.5, max_sigma等比例增大
|
||||
3. rsl_rl支持保存全部的env_cfg和train_cfg配置文件, 便于参数检查
|
||||
4. 加入MoE-CTS算法, 新任务go2_moe_cts, 消融load_balance_loss正则损失
|
||||
# 20251119
|
||||
## v0.1.7
|
||||
1. 修复rsl_rl中剩余训练时间计算的小bug
|
||||
2. 发现correct_base_height的奖励过小, 只有1e-3级别, 其他都是1e-2的, 可能导致滑倒时的重心偏移, 加入0-1500训练步的奖励课程系数1-10变化
|
||||
3. 总训练步数调到50k
|
||||
4. 域随机化: restitutionb不变0-0.5(isaacgym计算弹性系数是(地形0+机械人弹性系数)/2, 所以也不会很大), friction扩大到0.3-1.7 -> 0-2 (isaacgym中的摩擦计算方法是(地形1+机械人摩擦)/2, 因此可以开的范围更大点模仿瓷砖地)
|
||||
5. 边界速度对向速度变换: 当连续两次切换边界速度时, 直接将边界速度取反
|
||||
6. exporter加入新的pkl模型参数导出功能
|
||||
# 20251118
|
||||
## v0.1.6
|
||||
1. 修改边界速度采样, 启用x正/负方向和y正/负方向最大值采样
|
||||
2. 将heading_commands的角速度最大速度设置为当前的角速度范围 (原来是-1,1)
|
||||
3. 加入command_range_curriculum在20000步时将最大速度范围修改为-1,1
|
||||
4. 加入dynamic_sigma动态调整追踪奖励系数:
|
||||
设动态调整$\sigma$的速度绝对值范围为$[v_{min}, v_{max}]$（角速度同理），第$i$种地形在最大速度下的速度追踪系数记为$\sigma_{max}^{terrain_i}$，则当指令速度为$v_{x}$（y轴方向速度同理）在第$i$种地形下时，当前系数为
|
||||
$$\sigma_{vel} = \begin{cases}
|
||||
\sigma,&\quad v_x\in[0,v_{min}),\\
|
||||
\sigma(v_x-v_{min})+\sigma_{max}^{terrain_i}(v_{max}-v_{xy}),&\quad v_x\in[v_{min},v_{max}),\\
|
||||
\sigma_{max}^{terrain_i},&\quad v_x\in(v_{max},\infty).
|
||||
\end{cases}$$
|
||||
$$\sigma_{now}=\sigma+\min(e^{\frac{\text{level}_i+1}{10}}-1,1)(\sigma_{vel}-\sigma)$$
|
||||
5. 上真机发现在1.0速度时突变速度可能发生打滑的问题, 参考CTS论文给到`[0.3,1.7]`, 并消融`[0.5, 3.0]`
|
||||
6. 真机的原地转向效果很差, 没有到1 rad/s, 关闭heading_command, 并将原来的limit_lin_vel改为limit_vel, 其中加入对最大角速度的最大最小和0的组合; 在原地禁止时, 加入20%的概率已边界角速度原地旋转
|
||||
# 20251117
|
||||
## v0.1.5
|
||||
1. 加入动态调整resample command逻辑, 保证xy线速度能够出地形的一半距离:
|
||||
设当前是$n_r$次采样指令，第$i$次的xy方向命令为$c_i^{xy}$，$T_r$ 为采样间隔时间，$T$为episode时间，则第$n_r+1$次命令最小xy速度为
|
||||
$$v_{min} = \text{clip}\left(\frac{5-||\sum_{i}^{n_r}c_i^{xy}||_2T_r}{T-n_rT_r},0,v_{max}^{x\ or\ y}\right)$$
|
||||
第$n_r+1$次有$p_{zero}$概率(课程增加到0.1)使xy零速度，其持续时长为
|
||||
$$T^{zero} = \text{clip}\left(T-n_rT_r-\frac{5-||\sum_{i}^{n_r}c_i^{xy}||_2T_r}{0.8\times\max(v^{max}_x,v^{max}_y)},0,T_r\right)$$
|
||||
2. 移除当翻滚和俯仰导致的异常终止
|
||||
3. 修改地形等级变化的distance, 为当前episodex距离中心点最大的distance
|
||||
4. 加入地形之间的terrain_spacing=0.5m
|
||||
5. 之前训练中将上下楼梯比例反了
|
||||
6. 重新打开heading_command范围(-pi,pi), 并在reset时随机朝向
|
||||
7. 调整地形比例降低wave, 更偏向平地斜坡
|
||||
8. 加入一种边界线速度采样, 对x正方向和y正负方向最大值采样, 概率0.1, 当才到边界速度时取消角速度和heading指令
|
||||
9. 加入域随机化:
|
||||
1. `randomize_restitution=[0, 0.2]`机器人的弹性系数, 仅在初始化机器人时计算
|
||||
2. `randomize_motor_strength=[0.8, 1.2]`机器人随机电机力矩系数, 仅在env.reset时计算
|
||||
3. `randomize_action_delay`以0.5概率使用last_action, 每个env.step都计算
|
||||
10. 提高域随机化`friction_range=[1, 5]`能明显提高真机前进效果 (消融了`[0.5, 3]`, 结果接近, 在变速时可能摔倒)
|
||||
11. obs中加入`feet_contact_forces, dof_torques, dof_acceleration: 4, 12, 12`
|
||||
# 20251116
|
||||
## v0.1.4
|
||||
1. 用平地奖励训练出来的模型会有x速度过大产生高抬腿动作, 同时给xy也会产生高抬腿, 重心偏移摔倒的问题, 说明加入更多的域随机化也没解决
|
||||
2. 发现用CTS训练的模型能在mujoco中上20cm台阶, 但是由于双腿之间间距过小导致真机迁移效果非常差, 完全无法正常移动, 但是也能爬箱子
|
||||
3. 考虑用CTS奖励进行训练, 加上similar_to_default或hip_to_default, 消融两种奖励, 以及对应的奖励系数, 0.01还是0.05
|
||||
4. 真机上有明显的跺脚动作, 加上feet_contact_forces惩罚接触力大小, 设置接触力阈值大小为go2的重量15*9.8=147, 系数为-1 (参考collision)
|
||||
5. 修复resume时第一个训练step前奖励系数没更新的bug
|
||||
6. 关闭在x,y速度小于0.2时变为0, 加入zero_commands的课程, 0->1500训练步, 逐渐从0.0->0.1
|
||||
7. 降低训练的command采样范围`x,y:[-1,1]->[-0.5,0.5]`
|
||||
8. resample command降低`10->5`
|
||||
# 20251115
|
||||
## v0.1.3
|
||||
1. 修改`terrain_level`计算方法, 取当前全部环境等级取平均
|
||||
2. 仍然不使用`heading_command`, 虽然能够更稳定的提升环境等级, 并可以采样到更多的角速度指令, 但是无法以一个恒定的角速度进行移动, 和实际操作中不符
|
||||
3. 修复`torch.jit`导出问题
|
||||
4. 分别记录每个地形的奖励
|
||||
5. 加载训练模型时, 支持环境奖励系数课程加载
|
||||
6. 参考[yusongmin1-My_unitree_go2_gym](https://github.com/yusongmin1/My_unitree_go2_gym)加入域随机化`randomize_link_mass, randomize_base_com, randomize_pd_gains, randomize_motor_zero_offset`, 降低base_mass的最大值`3->1`, 降低`push_robots`xy方向速度`1.0->0.4`, 加入角速度推力`0.6`, 增大摩擦最小值`0.1->0.2`
|
||||
7. 修改范围`lin_vel_y: 0.5->1.0`, 修改命令采样时间`resampling_time: 30->10`
|
||||
8. 修改`base_height`计算使用的点云范围, 长宽`0.7x0.5->0.4x0.3`
|
||||
9. 修改`wave, stairs`训练配置, 其中都加入`0.1`的平地
|
||||
10. 关闭对`correct_base_height, dof_power`奖励的课程降低
|
||||
11. 修改obs中`height_measurements`系数`5 -> 2.5`
|
||||
# 20251114
|
||||
## v0.1.2
|
||||
1. 解决CTS`rollout_storage_cts.py`中学生教授数据采样混乱的问题
|
||||
# 20251113
|
||||
## v0.1.1
|
||||
1. 上下楼梯环境训练7h13min完成, 纯PPO无法学到上楼梯动作, 下楼梯基本能完成
|
||||
2. 测试10k的wave, slope, rough_slope训练, 修复地形提升问题
|
||||
3. play.py中加入onnx模型导出功能
|
||||
4. 机身高度稍微有点低, 提升3cm, `base_height_target: 0.35 -> 0.38`
|
||||
5. 加入CTS算法替代PPO:
|
||||
- 新文件: `on_policy_runner_cts.py, actor_critic_cts.py, cts.py, rollout_storage_cts.py`
|
||||
- 新配置: `LeggedRobotCfgCTS`
|
||||
- 新任务: `go2_cts`
|
||||
- 新导出: 修改torch.script和onnx导出代码, onnx模型的输入是按照IsaacLab的按照item堆叠的结果, 部署C++代码的帧堆叠[obsevation_manager.h](https://github.com/unitreerobotics/unitree_rl_lab/blob/61bfba15d35f1a93e3bacab85fe06b31643c83b7/deploy/include/isaaclab/manager/observation_manager.h#L63)
|
||||
## v0.1
|
||||
1. 添加地形选择`wave, slope, rough_slope, stairs down, stairs up, obstacles, stepping_stones, gap, flat`
|
||||
2. 添加高度特征
|
||||
3. 修改base_height计算方法, 通过高度特征平均值计算
|
||||
4. 支持课程奖励系数`rewards.curriculum_rewards`
|
||||
15
cmd.md
Normal file
15
cmd.md
Normal file
@@ -0,0 +1,15 @@
|
||||
## Train
|
||||
```bash
|
||||
python legged_gym/scripts/train.py --task=go2 --num_envs 4096 --headless
|
||||
python legged_gym/scripts/train.py --task=go2 --num_envs 128 --resume --load_run Nov13_11-14-22_wave_slope_rough_slope
|
||||
python legged_gym/scripts/train.py --task=go2 --num_envs 8 # DEBUG
|
||||
# CTS
|
||||
python legged_gym/scripts/train.py --task=go2_cts --num_envs 8096 --headless
|
||||
```
|
||||
## Play
|
||||
```bash
|
||||
python legged_gym/scripts/play.py --task=go2 --num_envs 8 # load latest
|
||||
python legged_gym/scripts/play.py --task=go2 --num_envs 8 --load_run Nov13_00-00-05_ # load specified run
|
||||
# CTS
|
||||
python legged_gym/scripts/play.py --task=go2_cts --num_envs 8
|
||||
```
|
||||
61
deploy/deploy_mujoco/configs/go2.yaml
Normal file
61
deploy/deploy_mujoco/configs/go2.yaml
Normal file
@@ -0,0 +1,61 @@
|
||||
policy_path: "{LEGGED_GYM_ROOT_DIR}/deploy/pre_train/go2/go2_cts_150k.pt"
|
||||
|
||||
# xml_path: "{LEGGED_GYM_ROOT_DIR}/resources/robots/go2/flat.xml"
|
||||
# xml_path: "{LEGGED_GYM_ROOT_DIR}/resources/robots/go2/race_track.xml" # change go2 init pos to pos="-5 2 0.445" in resources/robots/go2/go2.xml
|
||||
xml_path: "{LEGGED_GYM_ROOT_DIR}/resources/robots/go2/stairs.xml"
|
||||
|
||||
# Total simulation time
|
||||
simulation_duration: 60000000.0
|
||||
# Simulation time step
|
||||
simulation_dt: 0.002
|
||||
# Controller update frequency (meets the requirement of simulation_dt * controll_decimation=0.02; 50Hz)
|
||||
control_decimation: 10
|
||||
|
||||
kps: [20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0, 20.0]
|
||||
kds: [0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5]
|
||||
|
||||
default_angles: [0.1, 0.8, -1.5, -0.1, 0.8, -1.5,
|
||||
0.1, 1.0, -1.5, -0.1, 1.0, -1.5]
|
||||
|
||||
lin_vel_scale: 2.0
|
||||
ang_vel_scale: 0.25
|
||||
dof_pos_scale: 1.0
|
||||
dof_vel_scale: 0.05
|
||||
action_scale: 0.25
|
||||
cmd_scale: [2.0, 2.0, 0.25]
|
||||
num_actions: 12
|
||||
num_obs: 45
|
||||
max_cmd: [2.0, 1.5, 2.5]
|
||||
|
||||
cmd_init: [1.0, 0.0, 0.0]
|
||||
|
||||
mujoco_joint_names: [
|
||||
"FL_hip_joint",
|
||||
"FL_thigh_joint",
|
||||
"FL_calf_joint",
|
||||
"FR_hip_joint",
|
||||
"FR_thigh_joint",
|
||||
"FR_calf_joint",
|
||||
"RL_hip_joint",
|
||||
"RL_thigh_joint",
|
||||
"RL_calf_joint",
|
||||
"RR_hip_joint",
|
||||
"RR_thigh_joint",
|
||||
"RR_calf_joint"
|
||||
]
|
||||
|
||||
# IsaacGym is same as Mujoco joint names
|
||||
model_joint_names: [
|
||||
"FL_hip_joint",
|
||||
"FL_thigh_joint",
|
||||
"FL_calf_joint",
|
||||
"FR_hip_joint",
|
||||
"FR_thigh_joint",
|
||||
"FR_calf_joint",
|
||||
"RL_hip_joint",
|
||||
"RL_thigh_joint",
|
||||
"RL_calf_joint",
|
||||
"RR_hip_joint",
|
||||
"RR_thigh_joint",
|
||||
"RR_calf_joint"
|
||||
]
|
||||
251
deploy/deploy_mujoco/deploy_go2.py
Normal file
251
deploy/deploy_mujoco/deploy_go2.py
Normal file
@@ -0,0 +1,251 @@
|
||||
import time
|
||||
import mujoco.viewer
|
||||
import mujoco
|
||||
import numpy as np
|
||||
from legged_gym import LEGGED_GYM_ROOT_DIR
|
||||
import torch
|
||||
import yaml
|
||||
import os
|
||||
import imageio
|
||||
from pathlib import Path
|
||||
from argparse import ArgumentParser
|
||||
import pygame
|
||||
from matplotlib import pyplot as plt
|
||||
|
||||
|
||||
def get_gravity_orientation(quaternion):
|
||||
qw = quaternion[0]
|
||||
qx = quaternion[1]
|
||||
qy = quaternion[2]
|
||||
qz = quaternion[3]
|
||||
|
||||
gravity_orientation = np.zeros(3)
|
||||
|
||||
gravity_orientation[0] = 2 * (-qz * qx + qw * qy)
|
||||
gravity_orientation[1] = -2 * (qz * qy + qw * qx)
|
||||
gravity_orientation[2] = 1 - 2 * (qw * qw + qz * qz)
|
||||
|
||||
return gravity_orientation
|
||||
|
||||
|
||||
def pd_control(target_q, q, kp, target_dq, dq, kd):
|
||||
"""Calculates torques from position commands"""
|
||||
return (target_q - q) * kp + (target_dq - dq) * kd
|
||||
|
||||
def get_xbox_command(joystick, max_cmd):
|
||||
pygame.event.pump()
|
||||
dead_zone = 0.1
|
||||
lx = joystick.get_axis(0)
|
||||
ly = joystick.get_axis(1)
|
||||
rx = joystick.get_axis(3)
|
||||
if abs(lx) < dead_zone: lx = 0
|
||||
if abs(ly) < dead_zone: ly = 0
|
||||
if abs(rx) < dead_zone: rx = 0
|
||||
cmd_x = -ly * max_cmd[0]
|
||||
cmd_y = -lx * max_cmd[1]
|
||||
cmd_yaw = -rx * max_cmd[2]
|
||||
return np.array([cmd_x, cmd_y, cmd_yaw], dtype=np.float32)
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument("--save-video", action="store_true", help="Whether to save video of the simulation.")
|
||||
parser.add_argument("--visualize-moe-weights", action="store_true", help="Whether to visualize mixture of experts weights.")
|
||||
parser.add_argument("--save-moe-latent", action="store_true", help="Whether to save mixture of experts latent vectors.")
|
||||
args = parser.parse_args()
|
||||
save_video = args.save_video
|
||||
visualize_moe_weights = args.visualize_moe_weights
|
||||
save_moe_latent = args.save_moe_latent
|
||||
config_file = "go2.yaml"
|
||||
|
||||
pygame.init()
|
||||
use_joystick = False
|
||||
joystick = None
|
||||
if pygame.joystick.get_count() > 0:
|
||||
joystick = pygame.joystick.Joystick(0)
|
||||
joystick.init()
|
||||
use_joystick = True
|
||||
print(f"Detected Joystick: {joystick.get_name()}")
|
||||
else:
|
||||
print("No Joystick detected. Using default commands from config.")
|
||||
|
||||
with open(f"{LEGGED_GYM_ROOT_DIR}/deploy/deploy_mujoco/configs/{config_file}", "r") as f:
|
||||
config = yaml.load(f, Loader=yaml.FullLoader)
|
||||
policy_path = config["policy_path"].replace("{LEGGED_GYM_ROOT_DIR}", LEGGED_GYM_ROOT_DIR)
|
||||
xml_path = config["xml_path"].replace("{LEGGED_GYM_ROOT_DIR}", LEGGED_GYM_ROOT_DIR)
|
||||
|
||||
simulation_duration = config["simulation_duration"]
|
||||
simulation_dt = config["simulation_dt"]
|
||||
control_decimation = config["control_decimation"]
|
||||
|
||||
kps = np.array(config["kps"], dtype=np.float32)
|
||||
kds = np.array(config["kds"], dtype=np.float32)
|
||||
|
||||
default_angles = np.array(config["default_angles"], dtype=np.float32)
|
||||
|
||||
lin_vel_scale = config["lin_vel_scale"]
|
||||
ang_vel_scale = config["ang_vel_scale"]
|
||||
dof_pos_scale = config["dof_pos_scale"]
|
||||
dof_vel_scale = config["dof_vel_scale"]
|
||||
action_scale = config["action_scale"]
|
||||
cmd_scale = np.array(config["cmd_scale"], dtype=np.float32)
|
||||
|
||||
num_actions = config["num_actions"]
|
||||
num_obs = config["num_obs"]
|
||||
|
||||
cmd = np.array(config["cmd_init"], dtype=np.float32)
|
||||
|
||||
idx_model2mj = idx_mj2model = list(range(num_actions))
|
||||
if 'mujoco_joint_names' in config and 'model_joint_names' in config:
|
||||
mujoco_joint_names = config["mujoco_joint_names"]
|
||||
model_joint_names = config["model_joint_names"]
|
||||
idx_model2mj = [model_joint_names.index(joint) for joint in mujoco_joint_names]
|
||||
idx_mj2model = [mujoco_joint_names.index(joint) for joint in model_joint_names]
|
||||
|
||||
video_save_dir = str(Path(__file__).parent / "videos")
|
||||
os.makedirs(video_save_dir, exist_ok=True)
|
||||
|
||||
model_name = os.path.basename(policy_path).split('.')[0]
|
||||
cmd_str = f"cmd_{cmd[0]}_{cmd[1]}_{cmd[2]}"
|
||||
|
||||
# define context variables
|
||||
action = np.zeros(num_actions, dtype=np.float32)
|
||||
last_action = np.zeros(num_actions, dtype=np.float32)
|
||||
target_dof_pos = default_angles.copy()
|
||||
obs = np.zeros(num_obs, dtype=np.float32)
|
||||
|
||||
counter = 0
|
||||
|
||||
# Load robot model
|
||||
m = mujoco.MjModel.from_xml_path(xml_path)
|
||||
d = mujoco.MjData(m)
|
||||
m.opt.timestep = simulation_dt
|
||||
|
||||
renderer = mujoco.Renderer(m, height=360, width=640)
|
||||
|
||||
# load policy
|
||||
policy = torch.jit.load(policy_path)
|
||||
|
||||
if save_video:
|
||||
video_filename = f"{model_name}_{cmd_str}.mp4"
|
||||
video_path = os.path.join(video_save_dir, video_filename)
|
||||
print(f"Video recording will be saved to: {video_path}")
|
||||
video_fps = 50
|
||||
sim_fps = 1.0 / m.opt.timestep
|
||||
frame_skip = int(sim_fps / video_fps)
|
||||
if frame_skip < 1:
|
||||
frame_skip = 1
|
||||
writer = imageio.get_writer(video_path, fps=video_fps)
|
||||
print(f"Sim FPS: {sim_fps:.2f}, Video FPS: {video_fps}, Frame Skip: {frame_skip}, Save at: {video_path}")
|
||||
|
||||
if visualize_moe_weights:
|
||||
plt.ion()
|
||||
fig, ax = plt.subplots(figsize=(5,3))
|
||||
ax.set_title(f"Command: Vx={cmd[0]:.2f}, Vy={cmd[1]:.2f}, Wz={cmd[2]:.2f}")
|
||||
bars = None
|
||||
|
||||
if save_moe_latent:
|
||||
latent_save_dir = str(Path(__file__).parent / "data_latents")
|
||||
os.makedirs(latent_save_dir, exist_ok=True)
|
||||
latent_filename = f"{model_name}_{cmd_str}_latents.npy"
|
||||
latent_path = os.path.join(latent_save_dir, latent_filename)
|
||||
all_latents = []
|
||||
|
||||
with mujoco.viewer.launch_passive(m, d) as viewer:
|
||||
|
||||
# set viewer.camera to follow robot
|
||||
viewer.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
|
||||
viewer.cam.trackbodyid = 1
|
||||
viewer.cam.distance = 3.0
|
||||
viewer.cam.elevation = -30.0
|
||||
viewer.cam.azimuth = 0.0
|
||||
|
||||
# Close the viewer automatically after simulation_duration wall-seconds.
|
||||
start = time.time()
|
||||
while viewer.is_running() and time.time() - start < simulation_duration:
|
||||
step_start = time.time()
|
||||
|
||||
if use_joystick and counter % control_decimation == 0:
|
||||
cmd = get_xbox_command(joystick, config["max_cmd"])
|
||||
print(f"Cmd: Vx={cmd[0]:.2f}, Vy={cmd[1]:.2f}, Wz={cmd[2]:.2f}", end='\r')
|
||||
|
||||
tau = pd_control(target_dof_pos, d.qpos[7:], kps, np.zeros_like(kds), d.qvel[6:], kds)
|
||||
d.ctrl[:] = tau
|
||||
# mj_step can be replaced with code that also evaluates
|
||||
# a policy and applies a control signal before stepping the physics.
|
||||
mujoco.mj_step(m, d)
|
||||
|
||||
if save_video and counter % frame_skip == 0:
|
||||
try:
|
||||
renderer.update_scene(d, camera=viewer.cam)
|
||||
frame = renderer.render()
|
||||
writer.append_data(frame)
|
||||
except Exception as e:
|
||||
print(f"Error rendering frame: {e}")
|
||||
|
||||
counter += 1
|
||||
if counter % control_decimation == 0:
|
||||
# Apply control signal here.
|
||||
|
||||
# create observation
|
||||
qj = d.qpos[7:]
|
||||
dqj = d.qvel[6:]
|
||||
quat = d.qpos[3:7]
|
||||
lin_vel = d.qvel[:3]
|
||||
ang_vel = d.qvel[3:6]
|
||||
|
||||
qj = (qj - default_angles) * dof_pos_scale
|
||||
|
||||
dqj = dqj * dof_vel_scale
|
||||
gravity_orientation = get_gravity_orientation(quat)
|
||||
lin_vel = lin_vel * lin_vel_scale
|
||||
ang_vel = ang_vel * ang_vel_scale
|
||||
|
||||
obs[:3] = ang_vel
|
||||
obs[3:6] = gravity_orientation
|
||||
obs[6:9] = cmd * cmd_scale
|
||||
obs[9 : 9 + num_actions] = qj[idx_mj2model]
|
||||
obs[9 + num_actions : 9 + 2 * num_actions] = dqj[idx_mj2model]
|
||||
obs[9 + 2 * num_actions : 9 + 3 * num_actions] = action[idx_mj2model]
|
||||
obs_tensor = torch.from_numpy(obs).unsqueeze(0)
|
||||
# policy inference
|
||||
last_action = action
|
||||
result = policy(obs_tensor)
|
||||
if isinstance(result, tuple):
|
||||
action, (weights, latent) = result # moe
|
||||
action = action.detach().numpy().squeeze()[idx_model2mj]
|
||||
weights = weights.detach().numpy().squeeze()
|
||||
latent = latent.detach().numpy().squeeze()
|
||||
if visualize_moe_weights:
|
||||
if bars is None:
|
||||
x = np.arange(len(weights))
|
||||
bars = ax.bar(x, weights)
|
||||
ax.set_ylim(0, 1)
|
||||
else:
|
||||
for bar, w in zip(bars, weights):
|
||||
bar.set_height(w)
|
||||
|
||||
plt.draw()
|
||||
plt.pause(0.001) # 这会造成大约 1ms 的延迟
|
||||
if save_moe_latent:
|
||||
all_latents.append(latent)
|
||||
else:
|
||||
action = result.detach().cpu().numpy().squeeze()[idx_model2mj]
|
||||
# transform action to target_dof_pos
|
||||
target_dof_pos = action * action_scale + default_angles
|
||||
|
||||
# Pick up changes to the physics state, apply perturbations, update options from GUI.
|
||||
viewer.sync()
|
||||
|
||||
# Rudimentary time keeping, will drift relative to wall clock.
|
||||
# time_until_next_step = m.opt.timestep - (time.time() - step_start) - 0.1
|
||||
# if time_until_next_step > 0:
|
||||
# time.sleep(time_until_next_step)
|
||||
|
||||
# writer.close()
|
||||
if save_video:
|
||||
print(f"Video saved successfully to {video_path}")
|
||||
writer.close()
|
||||
if save_moe_latent and len(all_latents) > 0:
|
||||
all_latents = np.array(all_latents)
|
||||
np.save(latent_path, all_latents)
|
||||
print(f"Latent vectors saved successfully to {latent_path}")
|
||||
126
deploy/deploy_mujoco/deploy_go2_lab.py
Normal file
126
deploy/deploy_mujoco/deploy_go2_lab.py
Normal file
@@ -0,0 +1,126 @@
|
||||
import time
|
||||
|
||||
import mujoco.viewer
|
||||
import mujoco
|
||||
import numpy as np
|
||||
from legged_gym import LEGGED_GYM_ROOT_DIR
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
|
||||
def get_gravity_orientation(quaternion):
|
||||
qw = quaternion[0]
|
||||
qx = quaternion[1]
|
||||
qy = quaternion[2]
|
||||
qz = quaternion[3]
|
||||
|
||||
gravity_orientation = np.zeros(3)
|
||||
|
||||
gravity_orientation[0] = 2 * (-qz * qx + qw * qy)
|
||||
gravity_orientation[1] = -2 * (qz * qy + qw * qx)
|
||||
gravity_orientation[2] = 1 - 2 * (qw * qw + qz * qz)
|
||||
|
||||
return gravity_orientation
|
||||
|
||||
|
||||
def pd_control(target_q, q, kp, target_dq, dq, kd):
|
||||
"""Calculates torques from position commands"""
|
||||
return (target_q - q) * kp + (target_dq - dq) * kd
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# get config file name from command line
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("config_file", type=str, help="config file name in the config folder")
|
||||
args = parser.parse_args()
|
||||
config_file = args.config_file
|
||||
with open(f"{LEGGED_GYM_ROOT_DIR}/deploy/deploy_mujoco/configs/{config_file}", "r") as f:
|
||||
config = yaml.load(f, Loader=yaml.FullLoader)
|
||||
policy_path = config["policy_path"].replace("{LEGGED_GYM_ROOT_DIR}", LEGGED_GYM_ROOT_DIR)
|
||||
xml_path = config["xml_path"].replace("{LEGGED_GYM_ROOT_DIR}", LEGGED_GYM_ROOT_DIR)
|
||||
|
||||
simulation_duration = config["simulation_duration"]
|
||||
simulation_dt = config["simulation_dt"]
|
||||
control_decimation = config["control_decimation"]
|
||||
|
||||
kps = np.array(config["kps"], dtype=np.float32)
|
||||
kds = np.array(config["kds"], dtype=np.float32)
|
||||
|
||||
default_angles = np.array(config["default_angles"], dtype=np.float32)
|
||||
|
||||
joint_ids_map = config["joint_ids_map"]
|
||||
ang_vel_scale = config["ang_vel_scale"]
|
||||
dof_pos_scale = config["dof_pos_scale"]
|
||||
dof_vel_scale = config["dof_vel_scale"]
|
||||
action_scale = config["action_scale"]
|
||||
cmd_scale = np.array(config["cmd_scale"], dtype=np.float32)
|
||||
|
||||
num_actions = config["num_actions"]
|
||||
num_obs = config["num_obs"]
|
||||
|
||||
cmd = np.array(config["cmd_init"], dtype=np.float32)
|
||||
|
||||
# define context variables
|
||||
action = np.zeros(num_actions, dtype=np.float32)
|
||||
target_dof_pos = default_angles.copy()
|
||||
obs = np.zeros(num_obs, dtype=np.float32)
|
||||
|
||||
counter = 0
|
||||
|
||||
# Load robot model
|
||||
m = mujoco.MjModel.from_xml_path(xml_path)
|
||||
d = mujoco.MjData(m)
|
||||
m.opt.timestep = simulation_dt
|
||||
|
||||
# load policy
|
||||
policy = torch.jit.load(policy_path)
|
||||
|
||||
with mujoco.viewer.launch_passive(m, d) as viewer:
|
||||
# Close the viewer automatically after simulation_duration wall-seconds.
|
||||
start = time.time()
|
||||
while viewer.is_running() and time.time() - start < simulation_duration:
|
||||
step_start = time.time()
|
||||
temp = target_dof_pos[[0,4,8,1,5,9,2,6,10,3,7,11]]
|
||||
tau = pd_control(temp, d.qpos[7:], kps, np.zeros_like(kds), d.qvel[6:], kds)
|
||||
d.ctrl[:] = tau
|
||||
# mj_step can be replaced with code that also evaluates
|
||||
# a policy and applies a control signal before stepping the physics.
|
||||
mujoco.mj_step(m, d)
|
||||
|
||||
counter += 1
|
||||
if counter % control_decimation == 0:
|
||||
# Apply control signal here.
|
||||
|
||||
# create observation
|
||||
qj = d.qpos[7:]
|
||||
dqj = d.qvel[6:]
|
||||
quat = d.qpos[3:7]
|
||||
ang_vel = d.qvel[3:6]
|
||||
|
||||
qj = (qj - default_angles) * dof_pos_scale
|
||||
|
||||
dqj = dqj * dof_vel_scale
|
||||
gravity_orientation = get_gravity_orientation(quat)
|
||||
ang_vel = ang_vel * ang_vel_scale
|
||||
|
||||
obs[:3] = ang_vel
|
||||
obs[3:6] = gravity_orientation
|
||||
obs[6:9] = cmd * cmd_scale
|
||||
obs[9 : 9 + num_actions] = qj[joint_ids_map]
|
||||
obs[9 + num_actions : 9 + 2 * num_actions] = dqj[joint_ids_map]
|
||||
obs[9 + 2 * num_actions : 9 + 3 * num_actions] = action
|
||||
obs_tensor = torch.from_numpy(obs).unsqueeze(0)
|
||||
# policy inference
|
||||
action = policy(obs_tensor).detach().numpy().squeeze()
|
||||
# transform action to target_dof_pos
|
||||
target_dof_pos = action * action_scale + default_angles[joint_ids_map]
|
||||
|
||||
# Pick up changes to the physics state, apply perturbations, update options from GUI.
|
||||
viewer.sync()
|
||||
|
||||
# Rudimentary time keeping, will drift relative to wall clock.
|
||||
time_until_next_step = m.opt.timestep - (time.time() - step_start)
|
||||
if time_until_next_step > 0:
|
||||
time.sleep(time_until_next_step)
|
||||
272
deploy/deploy_mujoco/deploy_go2_moe.py
Normal file
272
deploy/deploy_mujoco/deploy_go2_moe.py
Normal file
@@ -0,0 +1,272 @@
|
||||
import time
|
||||
import mujoco.viewer
|
||||
import mujoco
|
||||
import numpy as np
|
||||
from legged_gym import LEGGED_GYM_ROOT_DIR
|
||||
import torch
|
||||
import yaml
|
||||
import os
|
||||
import imageio
|
||||
from pathlib import Path
|
||||
from argparse import ArgumentParser
|
||||
import pygame
|
||||
# from matplotlib import pyplot as plt # 移除 matplotlib
|
||||
|
||||
def get_gravity_orientation(quaternion):
|
||||
qw = quaternion[0]
|
||||
qx = quaternion[1]
|
||||
qy = quaternion[2]
|
||||
qz = quaternion[3]
|
||||
|
||||
gravity_orientation = np.zeros(3)
|
||||
|
||||
gravity_orientation[0] = 2 * (-qz * qx + qw * qy)
|
||||
gravity_orientation[1] = -2 * (qz * qy + qw * qx)
|
||||
gravity_orientation[2] = 1 - 2 * (qw * qw + qz * qz)
|
||||
|
||||
return gravity_orientation
|
||||
|
||||
|
||||
def pd_control(target_q, q, kp, target_dq, dq, kd):
|
||||
"""Calculates torques from position commands"""
|
||||
return (target_q - q) * kp + (target_dq - dq) * kd
|
||||
|
||||
def get_xbox_command(joystick, max_cmd):
|
||||
# 注意：如果开启了 Pygame 显示窗口，这里 event.pump 也是必要的
|
||||
pygame.event.pump()
|
||||
dead_zone = 0.1
|
||||
if joystick is not None:
|
||||
lx = joystick.get_axis(0)
|
||||
ly = joystick.get_axis(1)
|
||||
rx = joystick.get_axis(3)
|
||||
if abs(lx) < dead_zone: lx = 0
|
||||
if abs(ly) < dead_zone: ly = 0
|
||||
if abs(rx) < dead_zone: rx = 0
|
||||
cmd_x = -ly * max_cmd[0]
|
||||
cmd_y = -lx * max_cmd[1]
|
||||
cmd_yaw = -rx * max_cmd[2]
|
||||
return np.array([cmd_x, cmd_y, cmd_yaw], dtype=np.float32)
|
||||
return np.zeros(3, dtype=np.float32)
|
||||
|
||||
def draw_moe_weights(screen, weights, width, height):
|
||||
"""使用 Pygame 绘制 MoE 权重"""
|
||||
screen.fill((255, 255, 255)) # 白底
|
||||
|
||||
num_experts = len(weights)
|
||||
if num_experts == 0:
|
||||
return
|
||||
|
||||
# 设置边距
|
||||
margin = 5
|
||||
bar_width = (width - 2 * margin) / num_experts
|
||||
max_bar_height = height - 2 * margin
|
||||
|
||||
for i, w in enumerate(weights):
|
||||
# 限制 w 在 [0, 1] 之间用于显示
|
||||
w_clamped = max(0.0, min(1.0, w))
|
||||
bar_height = int(w_clamped * max_bar_height)
|
||||
|
||||
# 计算矩形位置 (Pygame 坐标原点在左上角)
|
||||
# left, top, width, height
|
||||
x = margin + i * bar_width
|
||||
y = height - margin - bar_height # 从底部向上长
|
||||
|
||||
# 绘制矩形 (蓝色)
|
||||
# 在 bar 之间留一点空隙 (width - 2)
|
||||
pygame.draw.rect(screen, (50, 100, 255), (x, y, bar_width - 2, bar_height))
|
||||
|
||||
pygame.display.flip()
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument("--save-video", action="store_true", help="Whether to save video of the simulation.")
|
||||
parser.add_argument("--visualize-moe-weights", action="store_true", help="Whether to visualize mixture of experts weights.")
|
||||
args = parser.parse_args()
|
||||
save_video = args.save_video
|
||||
visualize_moe_weights = args.visualize_moe_weights
|
||||
config_file = "go2.yaml"
|
||||
|
||||
# Pygame 初始化
|
||||
pygame.init()
|
||||
|
||||
use_joystick = False
|
||||
joystick = None
|
||||
if pygame.joystick.get_count() > 0:
|
||||
joystick = pygame.joystick.Joystick(0)
|
||||
joystick.init()
|
||||
use_joystick = True
|
||||
print(f"Detected Joystick: {joystick.get_name()}")
|
||||
else:
|
||||
print("No Joystick detected. Using default commands from config.")
|
||||
|
||||
# 如果需要可视化权重，设置 Pygame 窗口
|
||||
screen = None
|
||||
win_width, win_height = 400, 200
|
||||
if visualize_moe_weights:
|
||||
# 创建一个独立的窗口用于显示权重
|
||||
screen = pygame.display.set_mode((win_width, win_height))
|
||||
pygame.display.set_caption("MoE Weights Visualization")
|
||||
|
||||
with open(f"{LEGGED_GYM_ROOT_DIR}/deploy/deploy_mujoco/configs/{config_file}", "r") as f:
|
||||
config = yaml.load(f, Loader=yaml.FullLoader)
|
||||
policy_path = config["policy_path"].replace("{LEGGED_GYM_ROOT_DIR}", LEGGED_GYM_ROOT_DIR)
|
||||
xml_path = config["xml_path"].replace("{LEGGED_GYM_ROOT_DIR}", LEGGED_GYM_ROOT_DIR)
|
||||
|
||||
simulation_duration = config["simulation_duration"]
|
||||
simulation_dt = config["simulation_dt"]
|
||||
control_decimation = config["control_decimation"]
|
||||
|
||||
kps = np.array(config["kps"], dtype=np.float32)
|
||||
kds = np.array(config["kds"], dtype=np.float32)
|
||||
|
||||
default_angles = np.array(config["default_angles"], dtype=np.float32)
|
||||
|
||||
lin_vel_scale = config["lin_vel_scale"]
|
||||
ang_vel_scale = config["ang_vel_scale"]
|
||||
dof_pos_scale = config["dof_pos_scale"]
|
||||
dof_vel_scale = config["dof_vel_scale"]
|
||||
action_scale = config["action_scale"]
|
||||
cmd_scale = np.array(config["cmd_scale"], dtype=np.float32)
|
||||
|
||||
num_actions = config["num_actions"]
|
||||
num_obs = config["num_obs"]
|
||||
|
||||
cmd = np.array(config["cmd_init"], dtype=np.float32)
|
||||
|
||||
idx_model2mj = idx_mj2model = list(range(num_actions))
|
||||
if 'mujoco_joint_names' in config and 'model_joint_names' in config:
|
||||
mujoco_joint_names = config["mujoco_joint_names"]
|
||||
model_joint_names = config["model_joint_names"]
|
||||
idx_model2mj = [model_joint_names.index(joint) for joint in mujoco_joint_names]
|
||||
idx_mj2model = [mujoco_joint_names.index(joint) for joint in model_joint_names]
|
||||
|
||||
video_save_dir = str(Path(__file__).parent / "videos")
|
||||
os.makedirs(video_save_dir, exist_ok=True)
|
||||
|
||||
model_name = os.path.basename(policy_path).split('.')[0]
|
||||
cmd_str = f"cmd_{cmd[0]}_{cmd[1]}_{cmd[2]}"
|
||||
video_filename = f"{model_name}_{cmd_str}.mp4"
|
||||
video_path = os.path.join(video_save_dir, video_filename)
|
||||
print(f"Video recording will be saved to: {video_path}")
|
||||
|
||||
# define context variables
|
||||
action = np.zeros(num_actions, dtype=np.float32)
|
||||
last_action = np.zeros(num_actions, dtype=np.float32)
|
||||
target_dof_pos = default_angles.copy()
|
||||
obs = np.zeros(num_obs, dtype=np.float32)
|
||||
|
||||
counter = 0
|
||||
|
||||
# Load robot model
|
||||
m = mujoco.MjModel.from_xml_path(xml_path)
|
||||
d = mujoco.MjData(m)
|
||||
m.opt.timestep = simulation_dt
|
||||
|
||||
renderer = mujoco.Renderer(m, height=360, width=640)
|
||||
|
||||
# load policy
|
||||
policy = torch.jit.load(policy_path)
|
||||
|
||||
if save_video:
|
||||
video_fps = 50
|
||||
sim_fps = 1.0 / m.opt.timestep
|
||||
frame_skip = int(sim_fps / video_fps)
|
||||
if frame_skip < 1:
|
||||
frame_skip = 1
|
||||
writer = imageio.get_writer(video_path, fps=video_fps)
|
||||
print(f"Sim FPS: {sim_fps:.2f}, Video FPS: {video_fps}, Frame Skip: {frame_skip}, Save at: {video_path}")
|
||||
|
||||
# 移除了 plt 初始化逻辑
|
||||
|
||||
with mujoco.viewer.launch_passive(m, d) as viewer:
|
||||
|
||||
# set viewer.camera to follow robot
|
||||
viewer.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
|
||||
viewer.cam.trackbodyid = 1
|
||||
viewer.cam.distance = 3.0
|
||||
viewer.cam.elevation = -30.0
|
||||
viewer.cam.azimuth = 0.0
|
||||
|
||||
# Close the viewer automatically after simulation_duration wall-seconds.
|
||||
start = time.time()
|
||||
while viewer.is_running() and time.time() - start < simulation_duration:
|
||||
step_start = time.time()
|
||||
|
||||
if use_joystick and counter % control_decimation == 0:
|
||||
cmd = get_xbox_command(joystick, config["max_cmd"])
|
||||
print(f"Cmd: Vx={cmd[0]:.2f}, Vy={cmd[1]:.2f}, Wz={cmd[2]:.2f}", end='\r')
|
||||
elif visualize_moe_weights and counter % control_decimation == 0:
|
||||
# 如果没有手柄但开了可视化，也需要 pump 事件，防止窗口卡死
|
||||
pygame.event.pump()
|
||||
|
||||
tau = pd_control(target_dof_pos, d.qpos[7:], kps, np.zeros_like(kds), d.qvel[6:], kds)
|
||||
d.ctrl[:] = tau
|
||||
|
||||
mujoco.mj_step(m, d)
|
||||
|
||||
if save_video and counter % frame_skip == 0:
|
||||
try:
|
||||
renderer.update_scene(d, camera=viewer.cam)
|
||||
frame = renderer.render()
|
||||
writer.append_data(frame)
|
||||
except Exception as e:
|
||||
print(f"Error rendering frame: {e}")
|
||||
|
||||
counter += 1
|
||||
if counter % control_decimation == 0:
|
||||
# Apply control signal here.
|
||||
|
||||
# create observation
|
||||
qj = d.qpos[7:]
|
||||
dqj = d.qvel[6:]
|
||||
quat = d.qpos[3:7]
|
||||
lin_vel = d.qvel[:3]
|
||||
ang_vel = d.qvel[3:6]
|
||||
|
||||
qj = (qj - default_angles) * dof_pos_scale
|
||||
|
||||
dqj = dqj * dof_vel_scale
|
||||
gravity_orientation = get_gravity_orientation(quat)
|
||||
lin_vel = lin_vel * lin_vel_scale
|
||||
ang_vel = ang_vel * ang_vel_scale
|
||||
|
||||
obs[:3] = ang_vel
|
||||
obs[3:6] = gravity_orientation
|
||||
obs[6:9] = cmd * cmd_scale
|
||||
obs[9 : 9 + num_actions] = qj[idx_mj2model]
|
||||
obs[9 + num_actions : 9 + 2 * num_actions] = dqj[idx_mj2model]
|
||||
obs[9 + 2 * num_actions : 9 + 3 * num_actions] = action[idx_mj2model]
|
||||
obs_tensor = torch.from_numpy(obs).unsqueeze(0)
|
||||
# policy inference
|
||||
last_action = action
|
||||
result = policy(obs_tensor)
|
||||
|
||||
# 处理 MoE 和 绘图
|
||||
if isinstance(result, tuple):
|
||||
action, weights = result # moe
|
||||
action = action.detach().numpy().squeeze()[idx_model2mj]
|
||||
weights = weights.detach().numpy().squeeze()
|
||||
|
||||
if visualize_moe_weights and screen is not None:
|
||||
draw_moe_weights(screen, weights, win_width, win_height)
|
||||
|
||||
else:
|
||||
action = result.detach().numpy().squeeze()[idx_model2mj]
|
||||
|
||||
# transform action to target_dof_pos
|
||||
target_dof_pos = action * action_scale + default_angles
|
||||
|
||||
# Pick up changes to the physics state, apply perturbations, update options from GUI.
|
||||
viewer.sync()
|
||||
|
||||
# 如果需要严格同步时间，可以解开下面的注释
|
||||
# time_until_next_step = m.opt.timestep - (time.time() - step_start)
|
||||
# if time_until_next_step > 0:
|
||||
# time.sleep(time_until_next_step)
|
||||
|
||||
if save_video:
|
||||
writer.close()
|
||||
|
||||
# 退出时清理 Pygame
|
||||
pygame.quit()
|
||||
print(f"Video saved successfully to {video_path}")
|
||||
61
deploy/deploy_real/common/command_helper.py
Normal file
61
deploy/deploy_real/common/command_helper.py
Normal file
@@ -0,0 +1,61 @@
|
||||
from unitree_sdk2py.idl.unitree_go.msg.dds_ import LowCmd_ as LowCmdGo
|
||||
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowCmd_ as LowCmdHG
|
||||
from typing import Union
|
||||
|
||||
|
||||
class MotorMode:
|
||||
PR = 0 # Series Control for Pitch/Roll Joints
|
||||
AB = 1 # Parallel Control for A/B Joints
|
||||
|
||||
|
||||
def create_damping_cmd(cmd: Union[LowCmdGo, LowCmdHG]):
|
||||
size = len(cmd.motor_cmd)
|
||||
for i in range(size):
|
||||
cmd.motor_cmd[i].q = 0
|
||||
cmd.motor_cmd[i].qd = 0
|
||||
cmd.motor_cmd[i].kp = 0
|
||||
cmd.motor_cmd[i].kd = 8
|
||||
cmd.motor_cmd[i].tau = 0
|
||||
|
||||
|
||||
def create_zero_cmd(cmd: Union[LowCmdGo, LowCmdHG]):
|
||||
size = len(cmd.motor_cmd)
|
||||
for i in range(size):
|
||||
cmd.motor_cmd[i].q = 0
|
||||
cmd.motor_cmd[i].qd = 0
|
||||
cmd.motor_cmd[i].kp = 0
|
||||
cmd.motor_cmd[i].kd = 0
|
||||
cmd.motor_cmd[i].tau = 0
|
||||
|
||||
|
||||
def init_cmd_hg(cmd: LowCmdHG, mode_machine: int, mode_pr: int):
|
||||
cmd.mode_machine = mode_machine
|
||||
cmd.mode_pr = mode_pr
|
||||
size = len(cmd.motor_cmd)
|
||||
for i in range(size):
|
||||
cmd.motor_cmd[i].mode = 1
|
||||
cmd.motor_cmd[i].q = 0
|
||||
cmd.motor_cmd[i].qd = 0
|
||||
cmd.motor_cmd[i].kp = 0
|
||||
cmd.motor_cmd[i].kd = 0
|
||||
cmd.motor_cmd[i].tau = 0
|
||||
|
||||
|
||||
def init_cmd_go(cmd: LowCmdGo, weak_motor: list):
|
||||
cmd.head[0] = 0xFE
|
||||
cmd.head[1] = 0xEF
|
||||
cmd.level_flag = 0xFF
|
||||
cmd.gpio = 0
|
||||
PosStopF = 2.146e9
|
||||
VelStopF = 16000.0
|
||||
size = len(cmd.motor_cmd)
|
||||
for i in range(size):
|
||||
if i in weak_motor:
|
||||
cmd.motor_cmd[i].mode = 1
|
||||
else:
|
||||
cmd.motor_cmd[i].mode = 0x0A
|
||||
cmd.motor_cmd[i].q = PosStopF
|
||||
cmd.motor_cmd[i].qd = VelStopF
|
||||
cmd.motor_cmd[i].kp = 0
|
||||
cmd.motor_cmd[i].kd = 0
|
||||
cmd.motor_cmd[i].tau = 0
|
||||
39
deploy/deploy_real/common/remote_controller.py
Normal file
39
deploy/deploy_real/common/remote_controller.py
Normal file
@@ -0,0 +1,39 @@
|
||||
import struct
|
||||
|
||||
|
||||
class KeyMap:
|
||||
R1 = 0
|
||||
L1 = 1
|
||||
start = 2
|
||||
select = 3
|
||||
R2 = 4
|
||||
L2 = 5
|
||||
F1 = 6
|
||||
F2 = 7
|
||||
A = 8
|
||||
B = 9
|
||||
X = 10
|
||||
Y = 11
|
||||
up = 12
|
||||
right = 13
|
||||
down = 14
|
||||
left = 15
|
||||
|
||||
|
||||
class RemoteController:
|
||||
def __init__(self):
|
||||
self.lx = 0
|
||||
self.ly = 0
|
||||
self.rx = 0
|
||||
self.ry = 0
|
||||
self.button = [0] * 16
|
||||
|
||||
def set(self, data):
|
||||
# wireless_remote
|
||||
keys = struct.unpack("H", data[2:4])[0]
|
||||
for i in range(16):
|
||||
self.button[i] = (keys & (1 << i)) >> i
|
||||
self.lx = struct.unpack("f", data[4:8])[0]
|
||||
self.rx = struct.unpack("f", data[8:12])[0]
|
||||
self.ry = struct.unpack("f", data[12:16])[0]
|
||||
self.ly = struct.unpack("f", data[20:24])[0]
|
||||
25
deploy/deploy_real/common/rotation_helper.py
Normal file
25
deploy/deploy_real/common/rotation_helper.py
Normal file
@@ -0,0 +1,25 @@
|
||||
import numpy as np
|
||||
from scipy.spatial.transform import Rotation as R
|
||||
|
||||
|
||||
def get_gravity_orientation(quaternion):
|
||||
qw = quaternion[0]
|
||||
qx = quaternion[1]
|
||||
qy = quaternion[2]
|
||||
qz = quaternion[3]
|
||||
|
||||
gravity_orientation = np.zeros(3)
|
||||
|
||||
gravity_orientation[0] = 2 * (-qz * qx + qw * qy)
|
||||
gravity_orientation[1] = -2 * (qz * qy + qw * qx)
|
||||
gravity_orientation[2] = 1 - 2 * (qw * qw + qz * qz)
|
||||
|
||||
return gravity_orientation
|
||||
|
||||
|
||||
def transform_imu_data(waist_yaw, waist_yaw_omega, imu_quat, imu_omega):
|
||||
RzWaist = R.from_euler("z", waist_yaw).as_matrix()
|
||||
R_torso = R.from_quat([imu_quat[1], imu_quat[2], imu_quat[3], imu_quat[0]]).as_matrix()
|
||||
R_pelvis = np.dot(R_torso, RzWaist.T)
|
||||
w = np.dot(RzWaist, imu_omega[0]) - np.array([0, 0, waist_yaw_omega])
|
||||
return R.from_matrix(R_pelvis).as_quat()[[3, 0, 1, 2]], w
|
||||
35
deploy/deploy_real/config_go2.py
Normal file
35
deploy/deploy_real/config_go2.py
Normal file
@@ -0,0 +1,35 @@
|
||||
from legged_gym import LEGGED_GYM_ROOT_DIR
|
||||
import numpy as np
|
||||
import yaml
|
||||
|
||||
|
||||
class Config:
|
||||
def __init__(self, file_path) -> None:
|
||||
with open(file_path, "r") as f:
|
||||
config = yaml.load(f, Loader=yaml.FullLoader)
|
||||
|
||||
self.control_dt = config["control_dt"]
|
||||
|
||||
self.joint2motor_idx = config["joint2motor_idx"]
|
||||
|
||||
self.msg_type = config["msg_type"]
|
||||
self.imu_type = config["imu_type"]
|
||||
|
||||
self.lowcmd_topic = config["lowcmd_topic"]
|
||||
self.lowstate_topic = config["lowstate_topic"]
|
||||
|
||||
self.policy_path = config["policy_path"].replace("{LEGGED_GYM_ROOT_DIR}", LEGGED_GYM_ROOT_DIR)
|
||||
|
||||
self.kps = np.array(config["kps"],dtype=np.float32)
|
||||
self.kds = np.array(config["kds"],dtype=np.float32)
|
||||
self.default_angles = np.array(config["default_angles"], dtype=np.float32)
|
||||
|
||||
self.obs_scales_ang_vel = config["obs_scales_ang_vel"]
|
||||
self.obs_scales_dof_pos = config["obs_scales_dof_pos"]
|
||||
self.obs_scales_dof_vel = config["obs_scales_dof_vel"]
|
||||
|
||||
self.command_scale = config["command_scale"]
|
||||
self.action_scale = config["action_scale"]
|
||||
|
||||
self.num_actions = config["num_actions"]
|
||||
self.num_obs = config["num_obs"]
|
||||
27
deploy/deploy_real/configs/go2.yaml
Normal file
27
deploy/deploy_real/configs/go2.yaml
Normal file
@@ -0,0 +1,27 @@
|
||||
control_dt: 0.02
|
||||
|
||||
msg_type: "go" # "hg" or "go"
|
||||
imu_type: "torso" # "torso" or "pelvis"
|
||||
|
||||
lowcmd_topic: "rt/lowcmd"
|
||||
lowstate_topic: "rt/lowstate"
|
||||
|
||||
policy_path: "{LEGGED_GYM_ROOT_DIR}/deploy/pre_train/go2/go2_cts_150k.pt"
|
||||
|
||||
joint2motor_idx: [3,4,5,0,1,2,9,10,11,6,7,8]
|
||||
|
||||
kps: [20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20]
|
||||
kds: [0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5]
|
||||
default_angles: [ 0.1, 0.8, -1.5,
|
||||
-0.1, 0.8, -1.5,
|
||||
0.1, 1.0, -1.5,
|
||||
-0.1, 1.0, -1.5]
|
||||
obs_scales_ang_vel: 0.25
|
||||
obs_scales_dof_pos: 1.0
|
||||
obs_scales_dof_vel: 0.05
|
||||
command_scale: [3.0, 2.0, 0.5]
|
||||
|
||||
action_scale: 0.25
|
||||
num_actions: 12
|
||||
num_obs: 45
|
||||
|
||||
218
deploy/deploy_real/deploy_real_go2.py
Normal file
218
deploy/deploy_real/deploy_real_go2.py
Normal file
@@ -0,0 +1,218 @@
|
||||
from legged_gym import LEGGED_GYM_ROOT_DIR
|
||||
import numpy as np
|
||||
import time
|
||||
import torch
|
||||
|
||||
from unitree_sdk2py.core.channel import ChannelPublisher,ChannelSubscriber,ChannelFactoryInitialize
|
||||
from unitree_sdk2py.idl.default import unitree_go_msg_dds__LowCmd_,unitree_go_msg_dds__LowState_
|
||||
from unitree_sdk2py.idl.unitree_go.msg.dds_ import LowCmd_ as LowCmdGo
|
||||
from unitree_sdk2py.idl.unitree_go.msg.dds_ import LowState_ as LowStateGo
|
||||
|
||||
from unitree_sdk2py.utils.crc import CRC
|
||||
|
||||
from common.command_helper import create_zero_cmd,create_damping_cmd
|
||||
from common.rotation_helper import get_gravity_orientation
|
||||
from common.remote_controller import RemoteController, KeyMap
|
||||
from config_go2 import Config
|
||||
|
||||
HIGHLEVEL = 0xEE
|
||||
LOWLEVEL = 0xFF
|
||||
TRIGERLEVEL = 0xF0
|
||||
PosStopF = 2.146e9
|
||||
VelStopF = 16000.0
|
||||
|
||||
|
||||
def init_cmd_go2(cmd:LowCmdGo):
|
||||
cmd.head[0] = 0xFE
|
||||
cmd.head[1] = 0xEF
|
||||
cmd.level_flag = 0xFF
|
||||
cmd.gpio = 0
|
||||
for i in range(12):
|
||||
cmd.motor_cmd[i].mode = 0x0A # 0x01
|
||||
cmd.motor_cmd[i].q = PosStopF
|
||||
cmd.motor_cmd[i].dq = VelStopF # or qd
|
||||
cmd.motor_cmd[i].kp = 0.0
|
||||
cmd.motor_cmd[i].kd = 0.0
|
||||
cmd.motor_cmd[i].tau = 0.0
|
||||
|
||||
|
||||
class Controller:
|
||||
def __init__(self,config:Config) -> None:
|
||||
self.config = config
|
||||
self.remote_controller = RemoteController()
|
||||
self.use_remote_controller = True
|
||||
|
||||
self.policy = torch.jit.load(config.policy_path)
|
||||
self._warm_up()
|
||||
|
||||
self.qj = np.zeros(config.num_actions,dtype=np.float32)
|
||||
self.dqj = np.zeros(config.num_actions,dtype=np.float32)
|
||||
self.action = np.zeros(config.num_actions,dtype=np.float32)
|
||||
self.target_dof_pos = config.default_angles.copy()
|
||||
self.obs = np.zeros(config.num_obs,dtype=np.float32)
|
||||
self.cmd = np.array([0.8, 0, 0],dtype=np.float32)
|
||||
self.counter = 0
|
||||
|
||||
self.low_cmd = unitree_go_msg_dds__LowCmd_()
|
||||
self.low_state = unitree_go_msg_dds__LowState_()
|
||||
self.lowcmd_publisher = ChannelPublisher(config.lowcmd_topic,LowCmdGo)
|
||||
self.lowcmd_publisher.Init()
|
||||
self.lowstate_subscriber = ChannelSubscriber(config.lowstate_topic,LowStateGo)
|
||||
self.lowstate_subscriber.Init(self.LowStateHandler,10)
|
||||
# self.replay_buffer = ReplayBuffer(max_replay_buffer_size=200,flag='real_new')
|
||||
|
||||
self.wait_for_low_state()
|
||||
init_cmd_go2(self.low_cmd)
|
||||
|
||||
def _warm_up(self):
|
||||
obs = torch.ones((1,45))
|
||||
for _ in range(10):
|
||||
_ = self.policy(obs)
|
||||
print('Network has been warmed up.')
|
||||
|
||||
def wait_for_low_state(self):
|
||||
while self.low_state.tick == 0:
|
||||
time.sleep(self.config.control_dt)
|
||||
print("Successfully connected to the robot.")
|
||||
|
||||
def LowStateHandler(self,msg:LowStateGo):
|
||||
self.low_state = msg
|
||||
self.remote_controller.set(self.low_state.wireless_remote)
|
||||
|
||||
def send_cmd(self,cmd:LowCmdGo):
|
||||
cmd.crc = CRC().Crc(cmd)
|
||||
self.lowcmd_publisher.Write(cmd)
|
||||
|
||||
def zero_torque_state(self):
|
||||
print("Enter zero torque state.")
|
||||
print("Waiting for the start signal...")
|
||||
while self.remote_controller.button[KeyMap.start] != 1:
|
||||
create_zero_cmd(self.low_cmd)
|
||||
self.send_cmd(self.low_cmd)
|
||||
time.sleep(self.config.control_dt)
|
||||
|
||||
|
||||
def move_to_default_pos(self):
|
||||
print('Moving to default pos.')
|
||||
total_time = 2
|
||||
num_step = int(total_time / self.config.control_dt)
|
||||
|
||||
dof_idx = self.config.joint2motor_idx
|
||||
default_pos = self.config.default_angles
|
||||
|
||||
|
||||
init_dof_pos = np.zeros(12,dtype=np.float32)
|
||||
for i in range(12):
|
||||
init_dof_pos[i] = self.low_state.motor_state[dof_idx[i]].q
|
||||
|
||||
for i in range(num_step):
|
||||
alpha = i / num_step
|
||||
for j in range(12):
|
||||
motor_idx = dof_idx[j]
|
||||
target_pos = default_pos[j]
|
||||
self.low_cmd.motor_cmd[motor_idx].q = init_dof_pos[j] * (1 - alpha) + target_pos * alpha
|
||||
self.low_cmd.motor_cmd[motor_idx].dq = 0.0 # qd
|
||||
self.low_cmd.motor_cmd[motor_idx].kp = 40.0
|
||||
self.low_cmd.motor_cmd[motor_idx].kd = 0.6
|
||||
self.low_cmd.motor_cmd[motor_idx].tau = 0.0
|
||||
self.send_cmd(self.low_cmd)
|
||||
time.sleep(self.config.control_dt)
|
||||
|
||||
|
||||
def default_pos_state(self):
|
||||
print("Enter default pos state.")
|
||||
print("Waiting for the Button A signal...")
|
||||
while self.remote_controller.button[KeyMap.A] != 1:
|
||||
for i in range(12):
|
||||
motor_idx = self.config.joint2motor_idx[i]
|
||||
self.low_cmd.motor_cmd[motor_idx].q = self.config.default_angles[i]
|
||||
self.low_cmd.motor_cmd[motor_idx].dq = 0.0 # qd
|
||||
self.low_cmd.motor_cmd[motor_idx].kp = 40.0
|
||||
self.low_cmd.motor_cmd[motor_idx].kd = 0.6
|
||||
self.low_cmd.motor_cmd[motor_idx].tau = 0
|
||||
self.send_cmd(self.low_cmd)
|
||||
time.sleep(self.config.control_dt)
|
||||
|
||||
|
||||
def run(self):
|
||||
self.counter += 1
|
||||
for i in range(12):
|
||||
self.qj[i] = self.low_state.motor_state[self.config.joint2motor_idx[i]].q
|
||||
self.dqj[i] = self.low_state.motor_state[self.config.joint2motor_idx[i]].dq
|
||||
|
||||
ang_vel = np.array([self.low_state.imu_state.gyroscope], dtype=np.float32) * self.config.obs_scales_ang_vel
|
||||
quat = self.low_state.imu_state.quaternion
|
||||
gravity_orientation = get_gravity_orientation(quat) # imu_state quaternion: w, x, y, z
|
||||
|
||||
if self.use_remote_controller:
|
||||
self.cmd[0] = self.remote_controller.ly
|
||||
self.cmd[1] = self.remote_controller.lx * -1
|
||||
self.cmd[2] = self.remote_controller.rx * -1
|
||||
|
||||
qj_obs = self.qj.copy()
|
||||
qj_obs = (qj_obs - self.config.default_angles) * self.config.obs_scales_dof_pos
|
||||
dqj_obs = self.dqj.copy()
|
||||
dqj_obs = dqj_obs * self.config.obs_scales_dof_vel
|
||||
|
||||
self.obs[:3] = ang_vel
|
||||
self.obs[3:6] = gravity_orientation
|
||||
self.obs[6:9] = self.cmd * self.config.command_scale
|
||||
self.obs[9:21] = qj_obs
|
||||
self.obs[21:33] = dqj_obs
|
||||
self.obs[33:45] = self.action
|
||||
|
||||
obs_tensor = torch.from_numpy(self.obs).unsqueeze(0)
|
||||
self.action = self.policy(obs_tensor).detach().numpy().squeeze()
|
||||
|
||||
target_dof_pos = self.config.default_angles + self.action * self.config.action_scale
|
||||
# target_dof_pos = self.config.default_angles
|
||||
|
||||
for i in range(12):
|
||||
motor_idx = self.config.joint2motor_idx[i]
|
||||
self.low_cmd.motor_cmd[motor_idx].q = target_dof_pos[i]
|
||||
self.low_cmd.motor_cmd[motor_idx].dq = 0.0
|
||||
self.low_cmd.motor_cmd[motor_idx].kp = 20.0
|
||||
self.low_cmd.motor_cmd[motor_idx].kd = 0.5
|
||||
self.low_cmd.motor_cmd[motor_idx].tau = 0
|
||||
|
||||
self.send_cmd(self.low_cmd)
|
||||
time.sleep(self.config.control_dt)
|
||||
|
||||
# === 调试：遥控器 & 模型输出 ===
|
||||
# print(f"RC: lx={self.remote_controller.lx:+.2f} ly={self.remote_controller.ly:+.2f} "
|
||||
# f"rx={self.remote_controller.rx:+.2f}")
|
||||
# print(f"OBS cmd: {self.obs[6:9]}") # 遥控器信号在 obs 的位置
|
||||
# print(f"RAW action: {self.action[:4]}...") # 只看前 4 个，防止刷屏
|
||||
# print(f"TARGET Q: {target_dof_pos[::3]}") # 每 3 个关节抽 1 个，易读
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("net", type=str, help="network interface")
|
||||
args = parser.parse_args()
|
||||
|
||||
config_path = f"{LEGGED_GYM_ROOT_DIR}/deploy/deploy_real/configs/go2.yaml"
|
||||
config = Config(config_path)
|
||||
|
||||
ChannelFactoryInitialize(0, args.net)
|
||||
|
||||
controller = Controller(config)
|
||||
|
||||
controller.zero_torque_state()
|
||||
controller.move_to_default_pos()
|
||||
controller.default_pos_state()
|
||||
|
||||
while True:
|
||||
try:
|
||||
controller.run()
|
||||
if controller.remote_controller.button[KeyMap.select] == 1:
|
||||
break
|
||||
except KeyboardInterrupt:
|
||||
break
|
||||
|
||||
|
||||
create_damping_cmd(controller.low_cmd)
|
||||
controller.send_cmd(controller.low_cmd)
|
||||
print('Exit')
|
||||
|
||||
BIN
deploy/pre_train/go2/go2_cts_150k.pt
Normal file
BIN
deploy/pre_train/go2/go2_cts_150k.pt
Normal file
Binary file not shown.
119
doc/setup_en.md
Normal file
119
doc/setup_en.md
Normal file
@@ -0,0 +1,119 @@
|
||||
# Installation Guide
|
||||
|
||||
## System Requirements
|
||||
|
||||
- **Operating System**: Recommended Ubuntu 18.04 or later
|
||||
- **GPU**: Nvidia GPU
|
||||
- **Driver Version**: Recommended version 525 or later
|
||||
|
||||
---
|
||||
|
||||
## 1. Creating a Virtual Environment
|
||||
|
||||
It is recommended to run training or deployment programs in a virtual environment. Conda is recommended for creating virtual environments. If Conda is already installed on your system, you can skip step 1.1.
|
||||
|
||||
### 1.1 Download and Install MiniConda
|
||||
|
||||
MiniConda is a lightweight distribution of Conda, suitable for creating and managing virtual environments. Use the following commands to download and install:
|
||||
|
||||
```bash
|
||||
mkdir -p ~/miniconda3
|
||||
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh
|
||||
bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
|
||||
rm ~/miniconda3/miniconda.sh
|
||||
```
|
||||
|
||||
After installation, initialize Conda:
|
||||
|
||||
```bash
|
||||
~/miniconda3/bin/conda init --all
|
||||
source ~/.bashrc
|
||||
```
|
||||
|
||||
### 1.2 Create a New Environment
|
||||
|
||||
Use the following command to create a virtual environment:
|
||||
|
||||
```bash
|
||||
conda create -n unitree-rl python=3.8
|
||||
```
|
||||
|
||||
### 1.3 Activate the Virtual Environment
|
||||
|
||||
```bash
|
||||
conda activate unitree-rl
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 2. Installing Dependencies
|
||||
|
||||
### 2.1 Install PyTorch
|
||||
|
||||
PyTorch is a neural network computation framework used for model training and inference. Install it using the following command:
|
||||
|
||||
```bash
|
||||
conda install pytorch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 pytorch-cuda=12.1 -c pytorch -c nvidia
|
||||
```
|
||||
|
||||
### 2.2 Install Isaac Gym
|
||||
|
||||
Isaac Gym is a rigid body simulation and training framework provided by Nvidia.
|
||||
|
||||
#### 2.2.1 Download
|
||||
|
||||
Download [Isaac Gym](https://developer.nvidia.com/isaac-gym) from Nvidia’s official website.
|
||||
|
||||
#### 2.2.2 Install
|
||||
|
||||
After extracting the package, navigate to the `isaacgym/python` folder and install it using the following commands:
|
||||
|
||||
```bash
|
||||
cd isaacgym/python
|
||||
pip install -e .
|
||||
```
|
||||
|
||||
#### 2.2.3 Verify Installation
|
||||
|
||||
Run the following command. If a window opens displaying 1080 balls falling, the installation was successful:
|
||||
|
||||
```bash
|
||||
cd examples
|
||||
python 1080_balls_of_solitude.py
|
||||
```
|
||||
|
||||
If you encounter any issues, refer to the official documentation at `isaacgym/docs/index.html`.
|
||||
|
||||
### 2.3 Install rsl_rl
|
||||
|
||||
`rsl_rl` is a library implementing reinforcement learning algorithms.
|
||||
|
||||
#### 2.3.1 Install
|
||||
|
||||
```bash
|
||||
cd rsl_rl
|
||||
pip install -e .
|
||||
```
|
||||
|
||||
### 2.4 Install go2_rl_gym
|
||||
|
||||
#### 2.4.1 Download
|
||||
|
||||
Clone the repository using Git:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/unitreerobotics/go2_rl_gym.git
|
||||
```
|
||||
|
||||
#### 2.4.2 Install
|
||||
|
||||
Navigate to the directory and install it:
|
||||
|
||||
```bash
|
||||
cd go2_rl_gym
|
||||
pip install -e .
|
||||
```
|
||||
|
||||
### 2.5 Install unitree_cpp_deploy (Optional)
|
||||
|
||||
Refer to our C++ deployment repository, which is based on unitree_rl_lab and specifically designed for deploying models trained in this repository: [unitree_cpp_deploy](https://github.com/wty-yy-mini/unitree_cpp_deploy).
|
||||
137
doc/setup_zh.md
Normal file
137
doc/setup_zh.md
Normal file
@@ -0,0 +1,137 @@
|
||||
# 安装配置文档
|
||||
|
||||
## 系统要求
|
||||
|
||||
- **操作系统**：推荐使用 Ubuntu 18.04 或更高版本
|
||||
- **显卡**：Nvidia 显卡
|
||||
- **驱动版本**：建议使用 525 或更高版本
|
||||
|
||||
---
|
||||
|
||||
## 1. 创建虚拟环境
|
||||
|
||||
建议在虚拟环境中运行训练或部署程序，推荐使用 Conda 创建虚拟环境。如果您的系统中已经安装了 Conda，可以跳过步骤 1.1。
|
||||
|
||||
### 1.1 下载并安装 MiniConda
|
||||
|
||||
MiniConda 是 Conda 的轻量级发行版，适用于创建和管理虚拟环境。使用以下命令下载并安装：
|
||||
|
||||
```bash
|
||||
mkdir -p ~/miniconda3
|
||||
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh
|
||||
bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
|
||||
rm ~/miniconda3/miniconda.sh
|
||||
```
|
||||
|
||||
安装完成后，初始化 Conda：
|
||||
|
||||
```bash
|
||||
~/miniconda3/bin/conda init --all
|
||||
source ~/.bashrc
|
||||
```
|
||||
|
||||
### 1.2 创建新环境
|
||||
|
||||
使用以下命令创建虚拟环境：
|
||||
|
||||
```bash
|
||||
conda create -n unitree-rl python=3.8
|
||||
```
|
||||
|
||||
### 1.3 激活虚拟环境
|
||||
|
||||
```bash
|
||||
conda activate unitree-rl
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 2. 安装依赖
|
||||
|
||||
### 2.1 安装 PyTorch
|
||||
|
||||
PyTorch 是一个神经网络计算框架，用于模型训练和推理。使用以下命令安装：
|
||||
|
||||
```bash
|
||||
conda install pytorch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 pytorch-cuda=12.1 -c pytorch -c nvidia
|
||||
```
|
||||
|
||||
### 2.2 安装 Isaac Gym
|
||||
|
||||
Isaac Gym 是 Nvidia 提供的刚体仿真和训练框架。
|
||||
|
||||
#### 2.2.1 下载
|
||||
|
||||
从 Nvidia 官网下载 [Isaac Gym](https://developer.nvidia.com/isaac-gym)。
|
||||
|
||||
#### 2.2.2 安装
|
||||
|
||||
解压后进入 `isaacgym/python` 文件夹，执行以下命令安装：
|
||||
|
||||
```bash
|
||||
cd isaacgym/python
|
||||
pip install -e .
|
||||
```
|
||||
|
||||
#### 2.2.3 验证安装
|
||||
|
||||
运行以下命令，若弹出窗口并显示 1080 个球下落，则安装成功：
|
||||
|
||||
```bash
|
||||
cd examples
|
||||
python 1080_balls_of_solitude.py
|
||||
```
|
||||
|
||||
如有问题，可参考 `isaacgym/docs/index.html` 中的官方文档。
|
||||
|
||||
### 2.3 安装 rsl_rl
|
||||
|
||||
`rsl_rl` 是一个强化学习算法库。
|
||||
|
||||
#### 2.3.1 安装
|
||||
|
||||
```bash
|
||||
cd rsl_rl
|
||||
pip install -e .
|
||||
```
|
||||
|
||||
### 2.4 安装 go2_rl_gym
|
||||
|
||||
#### 2.4.1 下载
|
||||
|
||||
通过 Git 克隆仓库：
|
||||
|
||||
```bash
|
||||
git clone https://github.com/wty-yy/go2_rl_gym.git
|
||||
```
|
||||
|
||||
#### 2.4.2 安装
|
||||
|
||||
进入目录并安装：
|
||||
|
||||
```bash
|
||||
cd go2_rl_gym
|
||||
pip install -e .
|
||||
```
|
||||
|
||||
### 2.5 真机部署（可选）
|
||||
|
||||
#### 2.5.1 unitree_sdk2
|
||||
|
||||
C++ sdk, 编译参考[官方教程](https://github.com/unitreerobotics/unitree_sdk2?tab=readme-ov-file#environment-setup)
|
||||
|
||||
#### 2.5.2 unitree_sdk2_python（选择用Python部署）
|
||||
|
||||
```bash
|
||||
conda create -n kaiwu python=3.8
|
||||
conda activate kaiwu
|
||||
pip3 install pytorch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 pytorch-cuda=12.1 -c pytorch -c nvidia
|
||||
|
||||
git clone https://github.com/unitreerobotics/unitree_sdk2_python.git
|
||||
cd unitree_sdk2_python
|
||||
pip install -e .
|
||||
```
|
||||
|
||||
### 2.5.3 安装 unitree_cpp_deploy（选择用C++部署）
|
||||
|
||||
我们基于unitree_rl_lab修改的C++部署仓库，专门用于部署本仓库训练的模型 [unitree_cpp_deploy](https://github.com/wty-yy-mini/unitree_cpp_deploy)
|
||||
31
legged_gym/LICENSE
Normal file
31
legged_gym/LICENSE
Normal file
@@ -0,0 +1,31 @@
|
||||
Copyright (c) 2021, ETH Zurich, Nikita Rudin
|
||||
Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES
|
||||
All rights reserved.
|
||||
|
||||
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.
|
||||
|
||||
See licenses/assets for license information for assets included in this repository.
|
||||
See licenses/dependencies for license information of dependencies of this package.
|
||||
4
legged_gym/__init__.py
Normal file
4
legged_gym/__init__.py
Normal file
@@ -0,0 +1,4 @@
|
||||
import os
|
||||
|
||||
LEGGED_GYM_ROOT_DIR = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
|
||||
LEGGED_GYM_ENVS_DIR = os.path.join(LEGGED_GYM_ROOT_DIR, 'legged_gym', 'envs')
|
||||
11
legged_gym/envs/__init__.py
Normal file
11
legged_gym/envs/__init__.py
Normal file
@@ -0,0 +1,11 @@
|
||||
from legged_gym import LEGGED_GYM_ROOT_DIR, LEGGED_GYM_ENVS_DIR
|
||||
|
||||
from legged_gym.envs.go2.go2_env import Go2Robot
|
||||
from legged_gym.envs.go2.go2_config import GO2Cfg, GO2CfgPPO, GO2CfgCTS, GO2CfgMoECTS
|
||||
from .base.legged_robot import LeggedRobot
|
||||
|
||||
from legged_gym.utils.task_registry import task_registry
|
||||
|
||||
task_registry.register( "go2", Go2Robot, GO2Cfg(), GO2CfgPPO())
|
||||
task_registry.register( "go2_cts", Go2Robot, GO2Cfg(), GO2CfgCTS())
|
||||
task_registry.register( "go2_moe_cts", Go2Robot, GO2Cfg(), GO2CfgMoECTS())
|
||||
25
legged_gym/envs/base/base_config.py
Normal file
25
legged_gym/envs/base/base_config.py
Normal file
@@ -0,0 +1,25 @@
|
||||
import inspect
|
||||
|
||||
class BaseConfig:
|
||||
def __init__(self) -> None:
|
||||
""" Initializes all member classes recursively. Ignores all namse starting with '__' (buit-in methods)."""
|
||||
self.init_member_classes(self)
|
||||
|
||||
@staticmethod
|
||||
def init_member_classes(obj):
|
||||
# iterate over all attributes names
|
||||
for key in dir(obj):
|
||||
# disregard builtin attributes
|
||||
# if key.startswith("__"):
|
||||
if key=="__class__":
|
||||
continue
|
||||
# get the corresponding attribute object
|
||||
var = getattr(obj, key)
|
||||
# check if it the attribute is a class
|
||||
if inspect.isclass(var):
|
||||
# instantate the class
|
||||
i_var = var()
|
||||
# set the attribute to the instance instead of the type
|
||||
setattr(obj, key, i_var)
|
||||
# recursively init members of the attribute
|
||||
BaseConfig.init_member_classes(i_var)
|
||||
115
legged_gym/envs/base/base_task.py
Normal file
115
legged_gym/envs/base/base_task.py
Normal file
@@ -0,0 +1,115 @@
|
||||
|
||||
import sys
|
||||
from isaacgym import gymapi
|
||||
from isaacgym import gymutil
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
# Base class for RL tasks
|
||||
class BaseTask():
|
||||
|
||||
def __init__(self, cfg, sim_params, physics_engine, sim_device, headless):
|
||||
self.gym = gymapi.acquire_gym()
|
||||
|
||||
self.sim_params = sim_params
|
||||
self.physics_engine = physics_engine
|
||||
self.sim_device = sim_device
|
||||
sim_device_type, self.sim_device_id = gymutil.parse_device_str(self.sim_device)
|
||||
self.headless = headless
|
||||
|
||||
# env device is GPU only if sim is on GPU and use_gpu_pipeline=True, otherwise returned tensors are copied to CPU by physX.
|
||||
if sim_device_type=='cuda' and sim_params.use_gpu_pipeline:
|
||||
self.device = self.sim_device
|
||||
else:
|
||||
self.device = 'cpu'
|
||||
|
||||
# graphics device for rendering, -1 for no rendering
|
||||
self.graphics_device_id = self.sim_device_id
|
||||
if self.headless == True:
|
||||
self.graphics_device_id = -1
|
||||
|
||||
self.num_envs = cfg.env.num_envs
|
||||
self.num_obs = cfg.env.num_observations
|
||||
self.num_privileged_obs = cfg.env.num_privileged_obs
|
||||
self.num_actions = cfg.env.num_actions
|
||||
|
||||
# optimization flags for pytorch JIT
|
||||
torch._C._jit_set_profiling_mode(False)
|
||||
torch._C._jit_set_profiling_executor(False)
|
||||
|
||||
# allocate buffers
|
||||
self.obs_buf = torch.zeros(self.num_envs, self.num_obs, device=self.device, dtype=torch.float)
|
||||
self.rew_buf = torch.zeros(self.num_envs, device=self.device, dtype=torch.float)
|
||||
self.reset_buf = torch.ones(self.num_envs, device=self.device, dtype=torch.long)
|
||||
self.episode_length_buf = torch.zeros(self.num_envs, device=self.device, dtype=torch.long)
|
||||
self.time_out_buf = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool)
|
||||
if self.num_privileged_obs is not None:
|
||||
self.privileged_obs_buf = torch.zeros(self.num_envs, self.num_privileged_obs, device=self.device, dtype=torch.float)
|
||||
else:
|
||||
self.privileged_obs_buf = None
|
||||
# self.num_privileged_obs = self.num_obs
|
||||
|
||||
self.extras = {}
|
||||
|
||||
# create envs, sim and viewer
|
||||
self.create_sim()
|
||||
self.gym.prepare_sim(self.sim)
|
||||
|
||||
# todo: read from config
|
||||
self.enable_viewer_sync = True
|
||||
self.viewer = None
|
||||
|
||||
# if running with a viewer, set up keyboard shortcuts and camera
|
||||
if self.headless == False:
|
||||
# subscribe to keyboard shortcuts
|
||||
self.viewer = self.gym.create_viewer(
|
||||
self.sim, gymapi.CameraProperties())
|
||||
self.gym.subscribe_viewer_keyboard_event(
|
||||
self.viewer, gymapi.KEY_ESCAPE, "QUIT")
|
||||
self.gym.subscribe_viewer_keyboard_event(
|
||||
self.viewer, gymapi.KEY_V, "toggle_viewer_sync")
|
||||
|
||||
def get_observations(self):
|
||||
return self.obs_buf
|
||||
|
||||
def get_privileged_observations(self):
|
||||
return self.privileged_obs_buf
|
||||
|
||||
def reset_idx(self, env_ids):
|
||||
"""Reset selected robots"""
|
||||
raise NotImplementedError
|
||||
|
||||
def reset(self):
|
||||
""" Reset all robots"""
|
||||
self.reset_idx(torch.arange(self.num_envs, device=self.device))
|
||||
obs, privileged_obs, _, _, _ = self.step(torch.zeros(self.num_envs, self.num_actions, device=self.device, requires_grad=False))
|
||||
return obs, privileged_obs
|
||||
|
||||
def step(self, actions):
|
||||
raise NotImplementedError
|
||||
|
||||
def render(self, sync_frame_time=True):
|
||||
if self.viewer:
|
||||
# check for window closed
|
||||
if self.gym.query_viewer_has_closed(self.viewer):
|
||||
sys.exit()
|
||||
|
||||
# check for keyboard events
|
||||
for evt in self.gym.query_viewer_action_events(self.viewer):
|
||||
if evt.action == "QUIT" and evt.value > 0:
|
||||
sys.exit()
|
||||
elif evt.action == "toggle_viewer_sync" and evt.value > 0:
|
||||
self.enable_viewer_sync = not self.enable_viewer_sync
|
||||
|
||||
# fetch results
|
||||
if self.device != 'cpu':
|
||||
self.gym.fetch_results(self.sim, True)
|
||||
|
||||
# step graphics
|
||||
if self.enable_viewer_sync:
|
||||
self.gym.step_graphics(self.sim)
|
||||
self.gym.draw_viewer(self.viewer, self.sim, True)
|
||||
if sync_frame_time:
|
||||
self.gym.sync_frame_time(self.sim)
|
||||
else:
|
||||
self.gym.poll_viewer_events(self.viewer)
|
||||
1409
legged_gym/envs/base/legged_robot.py
Normal file
1409
legged_gym/envs/base/legged_robot.py
Normal file
File diff suppressed because it is too large
Load Diff
361
legged_gym/envs/base/legged_robot_config.py
Normal file
361
legged_gym/envs/base/legged_robot_config.py
Normal file
@@ -0,0 +1,361 @@
|
||||
import math
|
||||
from .base_config import BaseConfig
|
||||
|
||||
class LeggedRobotCfg(BaseConfig):
|
||||
class env:
|
||||
num_envs = 4096
|
||||
num_observations = 48
|
||||
num_privileged_obs = None # if not None a priviledge_obs_buf will be returned by step() (critic obs for assymetric training). None is returned otherwise
|
||||
num_actions = 12
|
||||
env_spacing = 3. # not used with heightfields/trimeshes
|
||||
send_timeouts = True # send time out information to the algorithm
|
||||
episode_length_s = 20 # episode length in seconds
|
||||
test = False
|
||||
|
||||
class terrain:
|
||||
mesh_type = 'trimesh' # none, plane, heightfield or trimesh
|
||||
horizontal_scale = 0.1 # [m]
|
||||
vertical_scale = 0.005 # [m]
|
||||
border_size = 25 # [m]
|
||||
curriculum = True
|
||||
static_friction = 1.0
|
||||
dynamic_friction = 1.0
|
||||
restitution = 0.
|
||||
# rough terrain only:
|
||||
measure_heights = True
|
||||
measured_points_x = [-0.8, -0.7, -0.6, -0.5, -0.4, -0.3, -0.2, -0.1, 0., 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8] # 1mx1.6m rectangle (without center line)
|
||||
measured_points_y = [-0.5, -0.4, -0.3, -0.2, -0.1, 0., 0.1, 0.2, 0.3, 0.4, 0.5]
|
||||
selected = False # select a unique terrain type and pass all arguments
|
||||
terrain_kwargs = None # Dict of arguments for selected terrain
|
||||
max_init_terrain_level = 5 # starting curriculum state
|
||||
terrain_length = 8.
|
||||
terrain_width = 8.
|
||||
num_rows= 10 # number of terrain rows (levels)
|
||||
num_cols = 20 # number of terrain cols (types)
|
||||
terrain_spacing = 0.5 # spacing between different terrain types [m]
|
||||
|
||||
# [wave, slope, rough slope, stairs down, stairs up, obstacles, stepping stones, gap, flat]
|
||||
terrain_proportions = [0.1, 0.1, 0.1, 0.2, 0.2, 0.1, 0.1, 0.1, 0.0]
|
||||
# trimesh only:
|
||||
slope_treshold = 0.75 # slopes above this threshold will be corrected to vertical surfaces
|
||||
|
||||
class commands:
|
||||
curriculum = False
|
||||
max_curriculum = 1.
|
||||
num_commands = 4 # default: lin_vel_x, lin_vel_y, ang_vel_yaw, heading (in heading mode ang_vel_yaw is recomputed from heading error)
|
||||
resampling_time = 10. # time before command are changed[s]
|
||||
heading_command = False # if true: compute ang vel command from heading error
|
||||
zero_command_curriculum = None
|
||||
# start training with zero commands and then gradually increase zero command probability
|
||||
# eg. {'start_iter': 0, 'end_iter': 1500, 'start_value': 0.0, 'end_value': 0.1}
|
||||
limit_ang_vel_at_zero_command_prob = 0.0 # probability of add limiting angular velocity commands when zero command is sampled
|
||||
limit_vel_prob = 0.0 # probability of limiting linear velocity command
|
||||
limit_vel_invert_when_continuous = True # invert the limit logic when using continuous sample limit velocity commands
|
||||
limit_vel = {"lin_vel_x": [-1, 1], "lin_vel_y": [-1, 1], "ang_vel_yaw": [-1, 0, 1]} # sample vel commands from min [-1] or zero [0] or max [1] range only
|
||||
stop_heading_at_limit = True # stop heading updates when vel is limited
|
||||
command_range_curriculum = [] # list for command range curriculums at specific training iterations
|
||||
# eg: [{
|
||||
# 'iter': 20000, # training iteration at which the command ranges are updated
|
||||
# 'lin_vel_x': [-1.0, 1.0], # min max [m/s]
|
||||
# 'lin_vel_y': [-1.0, 1.0], # min max [m/s]
|
||||
# 'ang_vel_yaw': [-2.0, 2.0], # min max [rad/s]
|
||||
# 'heading': [-1.57, 1.57], # min max [rad]
|
||||
# }]
|
||||
turn_over_zero_time = { # if turn_over is true, time robot must be stable before sampling new commands after a turn over
|
||||
"backflip": 5.0,
|
||||
"sideflip": 3.0,
|
||||
}
|
||||
# [wave, slope, rough slope, stairs up, stairs down, obstacles, stepping stones, gap, flat]
|
||||
terrain_max_command_ranges = [
|
||||
{'lin_vel_x': [-1.5, 1.5], 'lin_vel_y': [-1.5, 1.5], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]}, # wave
|
||||
{'lin_vel_x': [-1.5, 1.5], 'lin_vel_y': [-1.5, 1.5], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]}, # slope
|
||||
{'lin_vel_x': [-1.5, 1.5], 'lin_vel_y': [-1.5, 1.5], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]}, # rough slope
|
||||
{'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]}, # stairs up
|
||||
{'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]}, # stairs down
|
||||
{'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]}, # obstacles
|
||||
{'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]}, # stepping stones
|
||||
{'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]}, # gap
|
||||
{'lin_vel_x': [-2.0, 2.0], 'lin_vel_y': [-1.5, 1.5], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]}, # flat
|
||||
]
|
||||
|
||||
class ranges:
|
||||
lin_vel_x = [-1.0, 1.0] # min max [m/s]
|
||||
lin_vel_y = [-0.5, 0.5] # min max [m/s]
|
||||
ang_vel_yaw = [-1, 1] # min max [rad/s]
|
||||
heading = [-3.14, 3.14]
|
||||
|
||||
class init_state:
|
||||
pos = [0.0, 0.0, 1.] # x,y,z [m]
|
||||
rot = [0.0, 0.0, 0.0, 1.0] # x,y,z,w [quat]
|
||||
lin_vel = [0.0, 0.0, 0.0] # x,y,z [m/s]
|
||||
ang_vel = [0.0, 0.0, 0.0] # x,y,z [rad/s]
|
||||
default_joint_angles = { # target angles when action = 0.0
|
||||
"joint_a": 0.,
|
||||
"joint_b": 0.}
|
||||
turn_over = False # if true, initialize the robot in a flipped over position
|
||||
turn_over_proportions = [0.0, 0.2, 0.8] # proportions for backflip, sideflip, no flip
|
||||
turn_over_init_heights = { # initial heights range for each flip type
|
||||
'backflip': [0.10, 0.15],
|
||||
'sideflip': [0.16, 0.21],
|
||||
}
|
||||
|
||||
class control:
|
||||
control_type = 'P' # P: position, V: velocity, T: torques
|
||||
# PD Drive parameters:
|
||||
stiffness = {'joint_a': 10.0, 'joint_b': 15.} # [N*m/rad]
|
||||
damping = {'joint_a': 1.0, 'joint_b': 1.5} # [N*m*s/rad]
|
||||
# action scale: target angle = actionScale * action + defaultAngle
|
||||
action_scale = 0.5
|
||||
# decimation: Number of control action updates @ sim DT per policy DT
|
||||
decimation = 4
|
||||
|
||||
class asset:
|
||||
file = ""
|
||||
name = "legged_robot" # actor name
|
||||
foot_name = "None" # name of the feet bodies, used to index body state and contact force tensors
|
||||
penalize_contacts_on = []
|
||||
terminate_after_contacts_on = []
|
||||
disable_gravity = False
|
||||
collapse_fixed_joints = True # merge bodies connected by fixed joints. Specific fixed joints can be kept by adding " <... dont_collapse="true">
|
||||
fix_base_link = False # fixe the base of the robot
|
||||
default_dof_drive_mode = 3 # see GymDofDriveModeFlags (0 is none, 1 is pos tgt, 2 is vel tgt, 3 effort)
|
||||
self_collisions = 0 # 1 to disable, 0 to enable...bitwise filter
|
||||
replace_cylinder_with_capsule = True # replace collision cylinders with capsules, leads to faster/more stable simulation
|
||||
flip_visual_attachments = True # Some .obj meshes must be flipped from y-up to z-up
|
||||
|
||||
density = 0.001
|
||||
angular_damping = 0.
|
||||
linear_damping = 0.
|
||||
max_angular_velocity = 1000.
|
||||
max_linear_velocity = 1000.
|
||||
armature = 0.
|
||||
thickness = 0.01
|
||||
|
||||
class domain_rand:
|
||||
### Robot properties ###
|
||||
robot_properties_update = None
|
||||
# eg: {'start_iter': 5000, 'interval': 5000}
|
||||
|
||||
randomize_friction = True
|
||||
friction_range = [0.2, 1.25]
|
||||
|
||||
randomize_base_mass = True
|
||||
added_mass_range = [-1., 1.]
|
||||
|
||||
randomize_link_mass = True
|
||||
multiplied_link_mass_range = [0.9, 1.1]
|
||||
|
||||
randomize_base_com = True
|
||||
added_base_com_range = [-0.03, 0.03]
|
||||
|
||||
randomize_restitution = False # restitution to robot links (Robot init)
|
||||
restitution_range = [0.0, 0.2]
|
||||
|
||||
### Environment reset ###
|
||||
randomize_pd_gains = True
|
||||
stiffness_multiplier_range = [0.9, 1.1]
|
||||
damping_multiplier_range = [0.9, 1.1]
|
||||
|
||||
randomize_motor_zero_offset = True
|
||||
motor_zero_offset_range = [-0.035, 0.035]
|
||||
|
||||
randomize_motor_strength = False # (Env reset)
|
||||
motor_strength_range = [0.8, 1.2]
|
||||
|
||||
### Environment step ###
|
||||
push_robots = True
|
||||
push_interval_s = 4
|
||||
max_push_vel_xy = 0.4
|
||||
max_push_ang_vel = 0.6
|
||||
|
||||
randomize_action_delay = False # use last_action with 0~20 ms delay, 4 decimation
|
||||
|
||||
class rewards:
|
||||
class scales:
|
||||
termination = -0.0
|
||||
tracking_lin_vel = 1.0
|
||||
tracking_ang_vel = 0.5
|
||||
lin_vel_z = -2.0
|
||||
ang_vel_xy = -0.05
|
||||
orientation = -0.
|
||||
torques = -0.00001
|
||||
dof_vel = -0.
|
||||
dof_acc = -2.5e-7
|
||||
base_height = -0.
|
||||
feet_air_time = 1.0
|
||||
collision = -1.
|
||||
feet_stumble = -0.0
|
||||
action_rate = -0.01
|
||||
stand_still = -0.
|
||||
|
||||
class turn_over_scales:
|
||||
upright = 1.0
|
||||
|
||||
only_positive_rewards = True # if true negative total rewards are clipped at zero (avoids early termination problems)
|
||||
tracking_sigma = 0.25 # tracking reward = exp(-error^2/sigma)
|
||||
soft_dof_pos_limit = 1. # percentage of urdf limits, values above this limit are penalized
|
||||
soft_dof_vel_limit = 1.
|
||||
soft_torque_limit = 1.
|
||||
base_height_target = 1.
|
||||
max_contact_force = 100. # forces above this value are penalized
|
||||
curriculum_rewards = None # reward names to apply curriculum scaling to, List[dict]
|
||||
# eg: [{'reward_name': 'lin_vel_z', 'start_iter': 0, 'end_iter': 1500, 'start_value': 1.0, 'end_value': 0.0}]
|
||||
dynamic_sigma = None # linear interpolation of sigma based on command velocity, **Must start terrain curriculum first**
|
||||
# eg: {
|
||||
# "min_vel": 0.5, # min abs velocity to have default sigma
|
||||
# "max_vel": 1.0, # max abs velocity to have max sigma
|
||||
# # wave, slope, rough_slope, stairs up, stairs down, obstacles, stepping_stones, gap, flat]
|
||||
# "max_sigma": [1/3, 1/4, 1/4, 1/2.7, 1/2.7, 1/2, 1, 1, 1/4]
|
||||
# }
|
||||
turn_over_roll_threshold = math.pi / 4 # threshold on roll to use turn over rewards
|
||||
min_legs_distance = 0.1 # min distance between legs to not be considered stumbling
|
||||
|
||||
class normalization:
|
||||
class obs_scales:
|
||||
lin_vel = 2.0
|
||||
ang_vel = 0.25
|
||||
dof_pos = 1.0
|
||||
dof_vel = 0.05
|
||||
height_measurements = 2.5
|
||||
clip_observations = 100.
|
||||
clip_actions = 100.
|
||||
|
||||
class noise:
|
||||
add_noise = True
|
||||
noise_level = 1.0 # scales other values
|
||||
class noise_scales:
|
||||
dof_pos = 0.01
|
||||
dof_vel = 1.5
|
||||
lin_vel = 0.1
|
||||
ang_vel = 0.2
|
||||
gravity = 0.05
|
||||
height_measurements = 0.1
|
||||
|
||||
# viewer camera:
|
||||
class viewer:
|
||||
ref_env = 0
|
||||
pos = [10, 0, 6] # [m]
|
||||
lookat = [11., 5, 3.] # [m]
|
||||
|
||||
class sim:
|
||||
dt = 0.005
|
||||
substeps = 1
|
||||
gravity = [0., 0. ,-9.81] # [m/s^2]
|
||||
up_axis = 1 # 0 is y, 1 is z
|
||||
|
||||
class physx:
|
||||
num_threads = 10
|
||||
solver_type = 1 # 0: pgs, 1: tgs
|
||||
num_position_iterations = 4
|
||||
num_velocity_iterations = 0
|
||||
contact_offset = 0.01 # [m]
|
||||
rest_offset = 0.0 # [m]
|
||||
bounce_threshold_velocity = 0.5 #0.5 [m/s]
|
||||
max_depenetration_velocity = 1.0
|
||||
max_gpu_contact_pairs = 2**23 #2**24 -> needed for 8000 envs and more
|
||||
default_buffer_size_multiplier = 5
|
||||
contact_collection = 2 # 0: never, 1: last sub-step, 2: all sub-steps (default=2)
|
||||
|
||||
class LeggedRobotCfgPPO(BaseConfig):
|
||||
seed = 1
|
||||
runner_class_name = 'OnPolicyRunner'
|
||||
class policy:
|
||||
init_noise_std = 1.0
|
||||
actor_hidden_dims = [512, 256, 128]
|
||||
critic_hidden_dims = [512, 256, 128]
|
||||
activation = 'elu' # can be elu, relu, selu, crelu, lrelu, tanh, sigmoid
|
||||
# only for 'ActorCriticRecurrent':
|
||||
# rnn_type = 'lstm'
|
||||
# rnn_hidden_size = 512
|
||||
# rnn_num_layers = 1
|
||||
|
||||
class algorithm:
|
||||
# training params
|
||||
value_loss_coef = 1.0
|
||||
use_clipped_value_loss = True
|
||||
clip_param = 0.2
|
||||
entropy_coef = 0.01
|
||||
num_learning_epochs = 5
|
||||
num_mini_batches = 4 # mini batch size = num_envs*nsteps / nminibatches
|
||||
learning_rate = 1.e-3 #5.e-4
|
||||
schedule = 'adaptive' # could be adaptive, fixed
|
||||
gamma = 0.99
|
||||
lam = 0.95
|
||||
desired_kl = 0.01
|
||||
max_grad_norm = 1.
|
||||
|
||||
class runner:
|
||||
policy_class_name = 'ActorCritic'
|
||||
algorithm_class_name = 'PPO'
|
||||
num_steps_per_env = 24 # per iteration
|
||||
max_iterations = 1500 # number of policy updates
|
||||
|
||||
# logging
|
||||
save_interval = 50 # check for potential saves every this many iterations
|
||||
experiment_name = 'test'
|
||||
run_name = ''
|
||||
# load and resume
|
||||
resume = False
|
||||
load_run = -1 # -1 = last run
|
||||
checkpoint = -1 # -1 = last saved model
|
||||
resume_path = None # updated from load_run and chkpt
|
||||
|
||||
class LeggedRobotCfgCTS(BaseConfig):
|
||||
seed = 0
|
||||
runner_class_name = "OnPolicyRunnerCTS"
|
||||
history_length = 5
|
||||
class policy:
|
||||
init_noise_std = 1.0
|
||||
actor_hidden_dims = [512, 256, 128]
|
||||
critic_hidden_dims = [512, 256, 128]
|
||||
teacher_encoder_hidden_dims = [512, 256]
|
||||
student_encoder_hidden_dims = [512, 256]
|
||||
activation = 'elu' # can be elu, relu, selu, crelu, lrelu, tanh, sigmoid
|
||||
latent_dim = 32
|
||||
norm_type = 'l2norm' # normalization type for encoders: l2norm, simnorm
|
||||
|
||||
class algorithm:
|
||||
# training params
|
||||
value_loss_coef = 1.0
|
||||
use_clipped_value_loss = True
|
||||
clip_param = 0.2
|
||||
entropy_coef = 0.01
|
||||
num_learning_epochs = 5
|
||||
num_mini_batches = 4 # mini batch size = num_envs*nsteps / nminibatches
|
||||
learning_rate = 1.e-3 #5.e-4
|
||||
student_encoder_learning_rate = 1e-3
|
||||
schedule = 'adaptive' # could be adaptive, fixed
|
||||
gamma = 0.99
|
||||
lam = 0.95
|
||||
desired_kl = 0.01
|
||||
max_grad_norm = 1.
|
||||
teacher_env_ratio = 0.75 # percentage of envs assigned to teacher
|
||||
# teacher_env_ratio = 1.00 # percentage of envs assigned to teacher
|
||||
|
||||
class runner:
|
||||
policy_class_name = 'ActorCriticCTS'
|
||||
algorithm_class_name = 'CTS'
|
||||
num_steps_per_env = 24 # per iteration
|
||||
max_iterations = 1500 # number of policy updates
|
||||
|
||||
# logging
|
||||
save_interval = 50 # check for potential saves every this many iterations
|
||||
experiment_name = 'test'
|
||||
run_name = ''
|
||||
# load and resume
|
||||
resume = False
|
||||
load_run = -1 # -1 = last run
|
||||
checkpoint = -1 # -1 = last saved model
|
||||
resume_path = None # updated from load_run and chkpt
|
||||
|
||||
class LeggedRobotCfgMoECTS(LeggedRobotCfgCTS):
|
||||
class policy(LeggedRobotCfgCTS.policy):
|
||||
obs_no_goal_mask = None # mask for observation without goal inputs
|
||||
student_expert_num = 8 # number of experts in the student model
|
||||
|
||||
class algorithm(LeggedRobotCfgCTS.algorithm):
|
||||
load_balance_coef = 0.01 # coefficient for load balance loss
|
||||
|
||||
class runner(LeggedRobotCfgCTS.runner):
|
||||
policy_class_name = 'ActorCriticMoECTS'
|
||||
algorithm_class_name = 'MoECTS'
|
||||
266
legged_gym/envs/go2/go2_config.py
Normal file
266
legged_gym/envs/go2/go2_config.py
Normal file
@@ -0,0 +1,266 @@
|
||||
import math
|
||||
from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg, LeggedRobotCfgPPO, LeggedRobotCfgCTS, LeggedRobotCfgMoECTS, LeggedRobotCfgMoECTS
|
||||
|
||||
class GO2Cfg(LeggedRobotCfg):
|
||||
class init_state(LeggedRobotCfg.init_state):
|
||||
pos = [0.0, 0.0, 0.42] # x,y,z [m]
|
||||
default_joint_angles = { # = target angles [rad] when action = 0.0
|
||||
'FL_hip_joint': 0.1, # [rad]
|
||||
'RL_hip_joint': 0.1, # [rad]
|
||||
'FR_hip_joint': -0.1 , # [rad]
|
||||
'RR_hip_joint': -0.1, # [rad]
|
||||
|
||||
'FL_thigh_joint': 0.8, # [rad]
|
||||
'RL_thigh_joint': 1., # [rad]
|
||||
'FR_thigh_joint': 0.8, # [rad]
|
||||
'RR_thigh_joint': 1., # [rad]
|
||||
|
||||
'FL_calf_joint': -1.5, # [rad]
|
||||
'RL_calf_joint': -1.5, # [rad]
|
||||
'FR_calf_joint': -1.5, # [rad]
|
||||
'RR_calf_joint': -1.5, # [rad]
|
||||
}
|
||||
turn_over = False # initialize the robot in a flipped over position
|
||||
# turn_over_proportions = [0.1, 0.3, 0.6] # proportions for backflip, sideflip, noflip
|
||||
turn_over_proportions = [0.0, 0.2, 0.8] # proportions for backflip, sideflip, noflip
|
||||
turn_over_init_heights = { # initial heights range for each flip type
|
||||
'backflip': [0.10, 0.15],
|
||||
'sideflip': [0.16, 0.21],
|
||||
}
|
||||
# turn_over_proportions = [0.0, 1.0, 0.0] # proportions for backflip, sideflip, noflip
|
||||
|
||||
class env(LeggedRobotCfg.env):
|
||||
num_envs = 8192
|
||||
num_observations = 45
|
||||
# obs(45) + base_lin_vel(3) + height_measurements(187)
|
||||
num_privileged_obs = 45 + 3 + 4 + 12 + 12 + 187 # 263
|
||||
# num_privileged_obs = 45 + 3 + 187 # 235
|
||||
# num_privileged_obs = 48 # without height measurements
|
||||
episode_length_s = 25
|
||||
|
||||
class domain_rand(LeggedRobotCfg.domain_rand):
|
||||
### Robot properties ###
|
||||
randomize_friction = True
|
||||
friction_range = [0.0, 2.0]
|
||||
|
||||
randomize_base_mass = True
|
||||
added_mass_range = [-1., 1.]
|
||||
|
||||
randomize_link_mass = True
|
||||
multiplied_link_mass_range = [0.9, 1.1]
|
||||
|
||||
randomize_base_com = True
|
||||
added_base_com_range = [-0.03, 0.03]
|
||||
|
||||
randomize_restitution = True # restitution to robot links (Robot init)
|
||||
restitution_range = [0.0, 0.5]
|
||||
|
||||
### Environment reset ###
|
||||
randomize_pd_gains = True
|
||||
stiffness_multiplier_range = [0.9, 1.1]
|
||||
damping_multiplier_range = [0.9, 1.1]
|
||||
|
||||
randomize_motor_zero_offset = True
|
||||
motor_zero_offset_range = [-0.035, 0.035]
|
||||
|
||||
randomize_motor_strength = True # (Env reset)
|
||||
motor_strength_range = [0.8, 1.2]
|
||||
|
||||
### Environment step ###
|
||||
push_robots = True
|
||||
push_interval_s = 4
|
||||
max_push_vel_xy = 0.4
|
||||
max_push_ang_vel = 0.6
|
||||
|
||||
randomize_action_delay = True # use last_action with 0~20 ms delay, 4 decimation
|
||||
|
||||
class control(LeggedRobotCfg.control):
|
||||
# PD Drive parameters:
|
||||
control_type = 'P'
|
||||
stiffness = {'joint': 20.0} # [N*m/rad]
|
||||
damping = {'joint': 0.5} # [N*m*s/rad]
|
||||
# action scale: target angle = actionScale * action + defaultAngle
|
||||
action_scale = 0.25
|
||||
# decimation: Number of control action updates @ sim DT per policy DT
|
||||
decimation = 4
|
||||
|
||||
class terrain(LeggedRobotCfg.terrain):
|
||||
max_init_terrain_level = 5
|
||||
# [wave, slope, rough_slope, stairs up, stairs down, obstacles, stepping_stones, gap, flat]
|
||||
# terrain_proportions = [0.2, 0.05, 0.05, 0.30, 0.05, 0.25, 0.0, 0.0, 0.1] # 更偏向wave
|
||||
terrain_proportions = [0.05, 0.20, 0.05, 0.25, 0.10, 0.20, 0.0, 0.0, 0.15] # 这个更偏向平地斜坡
|
||||
# terrain_proportions = [0.20, 0.05, 0.05, 0.30, 0.15, 0.20, 0.0, 0.0, 0.05] # 更偏向wave和stairs
|
||||
# terrain_proportions = [0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0]
|
||||
# terrain_proportions = [0.3, 0.3, 0.3, 0.0, 0.0, 0.0, 0.0, 0.0, 0.1]
|
||||
# terrain_proportions = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]
|
||||
|
||||
class commands(LeggedRobotCfg.commands):
|
||||
curriculum = False
|
||||
max_curriculum = 1.
|
||||
num_commands = 4 # default: lin_vel_x, lin_vel_y, ang_vel_yaw (in heading mode ang_vel_yaw is recomputed from heading error)
|
||||
resampling_time = 5. # time before command are changed[s]
|
||||
heading_command = False # if true: compute ang vel command from heading error
|
||||
# start training with zero commands and then gradually increase zero command probability
|
||||
zero_command_curriculum = {'start_iter': 0, 'end_iter': 1500, 'start_value': 0.0, 'end_value': 0.1}
|
||||
limit_ang_vel_at_zero_command_prob = 0.2 # probability of add limiting angular velocity commands when zero command is sampled
|
||||
limit_vel_prob = 0.2 # probability of limiting linear velocity command
|
||||
limit_vel_invert_when_continuous = True # invert the limit logic when using continuous sample limit velocity commands
|
||||
limit_vel = {"lin_vel_x": [-1, 1], "lin_vel_y": [-1, 1], "ang_vel_yaw": [-1, 0, 1]} # sample vel commands from min [-1] or zero [0] or max [1] range only
|
||||
stop_heading_at_limit = True # stop heading updates when vel is limited
|
||||
command_range_curriculum = [{ # list for command range curriculums at specific training iterations
|
||||
'iter': 20000, # training iteration at which the command ranges are updated
|
||||
'lin_vel_x': [-1.0, 1.0], # min max [m/s]
|
||||
'lin_vel_y': [-1.0, 1.0], # min max [m/s]
|
||||
'ang_vel_yaw': [-1.5, 1.5], # min max [rad/s]
|
||||
'heading': [-1.57, 1.57], # min max [rad]
|
||||
}, { # list for command range curriculums at specific training iterations
|
||||
'iter': 50000, # training iteration at which the command ranges are updated
|
||||
'lin_vel_x': [-2.0, 2.0], # min max [m/s]
|
||||
'lin_vel_y': [-1.0, 1.0], # min max [m/s]
|
||||
'ang_vel_yaw': [-2.0, 2.0], # min max [rad/s]
|
||||
'heading': [-1.57, 1.57], # min max [rad]
|
||||
}]
|
||||
turn_over_zero_time = { # if turn_over is true, time robot must be stable before sampling new commands after a turn over
|
||||
"backflip": 5.0,
|
||||
"sideflip": 3.0,
|
||||
}
|
||||
# [wave, slope, rough slope, stairs up, stairs down, obstacles, stepping stones, gap, flat]
|
||||
terrain_max_command_ranges = [
|
||||
{'lin_vel_x': [-1.5, 1.5], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]}, # wave
|
||||
{'lin_vel_x': [-1.5, 1.5], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]}, # slope
|
||||
{'lin_vel_x': [-1.5, 1.5], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]}, # rough slope
|
||||
{'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]}, # stairs up
|
||||
{'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]}, # stairs down
|
||||
{'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]}, # obstacles
|
||||
{'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]}, # stepping stones
|
||||
{'lin_vel_x': [-1.0, 1.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-1.5, 1.5], 'heading': [-1.57, 1.57]}, # gap
|
||||
{'lin_vel_x': [-2.0, 2.0], 'lin_vel_y': [-1.0, 1.0], 'ang_vel_yaw': [-2.0, 2.0], 'heading': [-1.57, 1.57]}, # flat
|
||||
]
|
||||
|
||||
class ranges:
|
||||
lin_vel_x = [-0.5, 0.5] # min max [m/s]
|
||||
lin_vel_y = [-0.5, 0.5] # min max [m/s]
|
||||
ang_vel_yaw = [-1.0, 1.0] # min max [rad/s]
|
||||
heading = [-1.57, 1.57] # min max [rad]
|
||||
|
||||
class asset(LeggedRobotCfg.asset):
|
||||
file = '{LEGGED_GYM_ROOT_DIR}/resources/robots/go2/urdf/go2.urdf'
|
||||
name = "go2"
|
||||
foot_name = "foot"
|
||||
penalize_contacts_on = ["thigh", "calf"]
|
||||
terminate_after_contacts_on = ["base"]
|
||||
self_collisions = 1 # 1 to disable, 0 to enable...bitwise filter
|
||||
|
||||
class rewards(LeggedRobotCfg.rewards):
|
||||
soft_dof_pos_limit = 0.9
|
||||
base_height_target = 0.38
|
||||
only_positive_rewards = False
|
||||
max_contact_force = 147. # forces above this value are penalized, go2 weight 15kg
|
||||
curriculum_rewards = [
|
||||
{'reward_name': 'lin_vel_z', 'start_iter': 0, 'end_iter': 1500, 'start_value': 1.0, 'end_value': 0.0},
|
||||
{'reward_name': 'correct_base_height', 'start_iter': 0, 'end_iter': 5000, 'start_value': 1.0, 'end_value': 10.0},
|
||||
# {'reward_name': 'dof_power', 'start_iter': 0, 'end_iter': 3000, 'start_value': 1.0, 'end_value': 0.1},
|
||||
# {'reward_name': 'upright', 'start_iter': 0, 'end_iter': 1500, 'start_value': 1.0, 'end_value': 0.0},
|
||||
]
|
||||
tracking_sigma = 0.25 # tracking reward = exp(-error^2/sigma)
|
||||
dynamic_sigma = { # linear interpolation of sigma based on command velocity, **Must start terrain curriculum first**
|
||||
"min_lin_vel": 0.5, # min abs linear velocity to have default sigma
|
||||
"max_lin_vel": 1.5, # max abs linear velocity to have max sigma
|
||||
"min_ang_vel": 1.0, # min abs angular velocity to have default sigma
|
||||
"max_ang_vel": 2.0, # max abs angular velocity to have max sigma
|
||||
# wave, slope, rough_slope, stairs up, stairs down, obstacles, stepping_stones, gap, flat]
|
||||
# "max_sigma": [1/3, 1/4, 1/4, 1/2.7, 1/2.7, 1/2, 1, 1, 1/4]
|
||||
"max_sigma": [5/12, 1/4, 1/4, 1/2, 1/2, 3/4, 1, 1, 1/4]
|
||||
}
|
||||
min_legs_distance = 0.1 # min distance between legs to not be considered stumbling
|
||||
class scales:
|
||||
# tracking_lin_vel = 1.0
|
||||
# tracking_ang_vel = 0.2
|
||||
# lin_vel_z = -10.0
|
||||
# base_height = -50.0
|
||||
# action_rate = -0.005
|
||||
# similar_to_default = -0.1
|
||||
# dof_power = -1e-3 # 能够明显抑制跳跃
|
||||
# dof_acc = -3e-7
|
||||
|
||||
# tracking_lin_vel = 1.0
|
||||
# tracking_ang_vel = 0.5
|
||||
# lin_vel_z = -2.0
|
||||
# ang_vel_xy = -0.05
|
||||
# dof_acc = -2.5e-7
|
||||
# dof_power = -1e-3 # 能够明显抑制跳跃
|
||||
# # torques = -1e-4 # 无用会走着走着倒了
|
||||
# correct_base_height = -10.0
|
||||
# action_rate = -0.01
|
||||
# action_smoothness = -0.01
|
||||
# collision = -1.0
|
||||
# dof_pos_limits = -2.0
|
||||
# feet_regulation = -0.05
|
||||
# hip_to_default = -0.1
|
||||
# similar_to_default = -0.05
|
||||
|
||||
# CTS reward
|
||||
tracking_lin_vel = 1.0
|
||||
tracking_ang_vel = 0.5
|
||||
lin_vel_z = -2.0
|
||||
ang_vel_xy = -0.05
|
||||
dof_acc = -2.5e-7
|
||||
dof_power = -2e-5
|
||||
torques = -1e-4
|
||||
correct_base_height = -1.0
|
||||
action_rate = -0.01
|
||||
action_smoothness = -0.01
|
||||
collision = -1.0
|
||||
dof_pos_limits = -2.0
|
||||
feet_regulation = -0.05
|
||||
# CTS奖励训出来双脚距离非常近, 真机效果很差, 但是sim2sim能上20cm楼梯, 尝试加入hip_to_default奖励或similar_to_default奖励
|
||||
hip_to_default = -0.05 # 在训练到y=1.5时, 双脚会明显碰撞, 为避免该问题提升hip, 效果更差, 还是保持0.05 (y最大也只到0.1了)
|
||||
# legs_distance = -1.5 # 奖励双脚距离, 避免CTS训练出来双脚距离过近, 尝试加入后robogauge flat验证效果变差, 删除
|
||||
# similar_to_default = -0.01
|
||||
# feet_contact_forces = -1.0 # 尝试加入但并没有起到任何效果, 删除
|
||||
|
||||
turn_over_roll_threshold = math.pi / 4 # threshold on roll to use turn over rewards
|
||||
class turn_over_scales:
|
||||
upright = 1.0
|
||||
# dof_acc = -2.5e-7
|
||||
# dof_power = -2e-5
|
||||
# action_rate = -0.001
|
||||
# action_smoothness = -0.001
|
||||
|
||||
class noise(LeggedRobotCfg.noise):
|
||||
add_noise = True
|
||||
|
||||
class GO2CfgPPO(LeggedRobotCfgPPO):
|
||||
class algorithm(LeggedRobotCfgPPO.algorithm):
|
||||
entropy_coef = 0.01
|
||||
class runner(LeggedRobotCfgPPO.runner):
|
||||
run_name = ''
|
||||
experiment_name = 'go2_ppo'
|
||||
max_iterations = 100000
|
||||
save_interval = 500
|
||||
|
||||
class GO2CfgCTS(LeggedRobotCfgCTS):
|
||||
class runner(LeggedRobotCfgCTS.runner):
|
||||
num_steps_per_env = 24
|
||||
run_name = ''
|
||||
experiment_name = 'go2_cts'
|
||||
max_iterations = 100000
|
||||
save_interval = 500
|
||||
|
||||
class policy(LeggedRobotCfgCTS.policy):
|
||||
latent_dim = 32
|
||||
norm_type = 'l2norm'
|
||||
|
||||
class GO2CfgMoECTS(LeggedRobotCfgMoECTS):
|
||||
class policy(LeggedRobotCfgMoECTS.policy):
|
||||
obs_no_goal_mask = [True] * 6 + [False] * 3 + [True] * 36 # mask for obs without command info
|
||||
student_expert_num = 8 # number of experts in the student model
|
||||
|
||||
class algorithm(LeggedRobotCfgMoECTS.algorithm):
|
||||
load_balance_coef = 0.01
|
||||
|
||||
class runner(LeggedRobotCfgMoECTS.runner):
|
||||
run_name = ''
|
||||
experiment_name = 'go2_moe_cts'
|
||||
max_iterations = 100000
|
||||
save_interval = 500
|
||||
53
legged_gym/envs/go2/go2_env.py
Normal file
53
legged_gym/envs/go2/go2_env.py
Normal file
@@ -0,0 +1,53 @@
|
||||
|
||||
from legged_gym.envs.base.legged_robot import LeggedRobot
|
||||
|
||||
from isaacgym.torch_utils import *
|
||||
from isaacgym import gymtorch, gymapi, gymutil
|
||||
import torch
|
||||
|
||||
class Go2Robot(LeggedRobot):
|
||||
def _get_noise_scale_vec(self, cfg):
|
||||
noise_vec = torch.zeros_like(self.obs_buf[0])
|
||||
self.add_noise = self.cfg.noise.add_noise
|
||||
noise_scales = self.cfg.noise.noise_scales
|
||||
noise_level = self.cfg.noise.noise_level
|
||||
noise_vec[:3] = noise_scales.ang_vel * noise_level * self.obs_scales.ang_vel
|
||||
noise_vec[3:6] = noise_scales.gravity * noise_level
|
||||
noise_vec[6:9] = 0. # commands
|
||||
noise_vec[9:9+self.num_actions] = noise_scales.dof_pos * noise_level * self.obs_scales.dof_pos
|
||||
noise_vec[9+self.num_actions:9+2*self.num_actions] = noise_scales.dof_vel * noise_level * self.obs_scales.dof_vel
|
||||
noise_vec[9+2*self.num_actions:9+3*self.num_actions] = 0. # previous actions
|
||||
|
||||
return noise_vec
|
||||
|
||||
def compute_observations(self):
|
||||
""" Computes observations
|
||||
"""
|
||||
self.obs_buf = torch.cat((self.base_ang_vel * self.obs_scales.ang_vel,
|
||||
self.projected_gravity,
|
||||
self.commands[:, :3] * self.commands_scale,
|
||||
(self.dof_pos - self.default_dof_pos) * self.obs_scales.dof_pos,
|
||||
self.dof_vel * self.obs_scales.dof_vel,
|
||||
self.actions,
|
||||
),dim=-1)
|
||||
|
||||
heights = torch.clip(self.root_states[:, 2].unsqueeze(1) - 0.5 - self.measured_heights, -1, 1.0) * self.obs_scales.height_measurements
|
||||
|
||||
self.privileged_obs_buf = torch.cat(( self.base_lin_vel * self.obs_scales.lin_vel,
|
||||
self.base_ang_vel * self.obs_scales.ang_vel,
|
||||
self.projected_gravity,
|
||||
self.commands[:, :3] * self.commands_scale,
|
||||
(self.dof_pos - self.default_dof_pos) * self.obs_scales.dof_pos,
|
||||
self.dof_vel * self.obs_scales.dof_vel,
|
||||
self.actions,
|
||||
torch.norm(self.contact_forces[:, self.feet_indices, :], dim=-1) * 1e-3, # foot contact forces (4,)
|
||||
self.torques / self.torque_limits, # motor torques (12,)
|
||||
(self.last_dof_vel - self.dof_vel) / self.dt * 1e-4, # motor accelerations (12,)
|
||||
heights, # height measurements (187,)
|
||||
),dim=-1)
|
||||
# print(f"foot contact: {self.privileged_obs_buf[:,48:48+4].min(), self.privileged_obs_buf[:,48:48+4].max()}")
|
||||
# print(f"torques: {self.privileged_obs_buf[:,48+4:48+4+12].min(), self.privileged_obs_buf[:,48+4:48+4+12].max()}")
|
||||
# print(f"acc: {self.privileged_obs_buf[:,48+4+12:48+4+12+12].min(), self.privileged_obs_buf[:,48+4+12:48+4+12+12].max()}")
|
||||
|
||||
if self.add_noise:
|
||||
self.obs_buf += (2 * torch.rand_like(self.obs_buf) - 1) * self.noise_scale_vec
|
||||
70
legged_gym/scripts/play.py
Normal file
70
legged_gym/scripts/play.py
Normal file
@@ -0,0 +1,70 @@
|
||||
import sys
|
||||
from legged_gym import LEGGED_GYM_ROOT_DIR
|
||||
import os
|
||||
import sys
|
||||
from legged_gym import LEGGED_GYM_ROOT_DIR
|
||||
|
||||
import isaacgym
|
||||
from legged_gym.envs import *
|
||||
from legged_gym.utils import get_args, task_registry, Logger
|
||||
from legged_gym.utils.exporter import export_policy_as_jit, export_policy_as_onnx, export_policy_as_pkl
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
def play(args):
|
||||
env_cfg, train_cfg = task_registry.get_cfgs(name=args.task)
|
||||
# override some parameters for testing
|
||||
env_cfg.env.num_envs = min(env_cfg.env.num_envs, 100)
|
||||
# env_cfg.terrain.mesh_type = 'plane'
|
||||
env_cfg.terrain.num_rows = 7
|
||||
env_cfg.terrain.num_cols = 7
|
||||
env_cfg.terrain.curriculum = False
|
||||
env_cfg.noise.add_noise = False
|
||||
env_cfg.domain_rand.randomize_friction = False
|
||||
env_cfg.domain_rand.push_robots = False
|
||||
env_cfg.domain_rand.randomize_base_mass = False
|
||||
env_cfg.domain_rand.randomize_link_mass = False
|
||||
env_cfg.domain_rand.randomize_base_com = False
|
||||
env_cfg.domain_rand.randomize_pd_gains = False
|
||||
env_cfg.domain_rand.randomize_motor_zero_offset = False
|
||||
|
||||
env_cfg.env.test = True
|
||||
|
||||
# prepare environment
|
||||
env, _ = task_registry.make_env(name=args.task, args=args, env_cfg=env_cfg)
|
||||
obs = env.get_observations()
|
||||
# load policy
|
||||
train_cfg.runner.resume = True
|
||||
runner, train_cfg = task_registry.make_alg_runner(env=env, name=args.task, args=args, train_cfg=train_cfg)
|
||||
policy = runner.get_inference_policy(device=env.device)
|
||||
|
||||
# export policy as a jit module (used to run it from C++)
|
||||
if EXPORT_POLICY:
|
||||
path = os.path.join(LEGGED_GYM_ROOT_DIR, 'logs', train_cfg.runner.experiment_name, 'exported', 'policies')
|
||||
if hasattr(runner.alg, 'actor_critic'):
|
||||
model = runner.alg.actor_critic
|
||||
else:
|
||||
model = runner.alg.model
|
||||
export_policy_as_jit(model, path)
|
||||
export_policy_as_onnx(model, path)
|
||||
export_policy_as_pkl(model, path)
|
||||
print('Exported policy as jit script / onnx to: ', path)
|
||||
|
||||
for i in range(10*int(env.max_episode_length)):
|
||||
actions = policy(obs.detach())
|
||||
|
||||
if FIX_COMMAND:
|
||||
env.commands[:, 0] = 1.0
|
||||
env.commands[:, 1] = 0.0
|
||||
env.commands[:, 2] = 0.0
|
||||
|
||||
obs, _, rews, dones, infos = env.step(actions.detach())
|
||||
|
||||
if __name__ == '__main__':
|
||||
EXPORT_POLICY = True
|
||||
RECORD_FRAMES = False
|
||||
MOVE_CAMERA = False
|
||||
FIX_COMMAND = True
|
||||
args = get_args()
|
||||
play(args)
|
||||
20
legged_gym/scripts/train.py
Normal file
20
legged_gym/scripts/train.py
Normal file
@@ -0,0 +1,20 @@
|
||||
import os
|
||||
import numpy as np
|
||||
from datetime import datetime
|
||||
import sys
|
||||
|
||||
import isaacgym
|
||||
from legged_gym.envs import *
|
||||
from legged_gym.utils import get_args, task_registry
|
||||
import torch
|
||||
|
||||
def train(args):
|
||||
env, env_cfg = task_registry.make_env(name=args.task, args=args)
|
||||
runner, train_cfg = task_registry.make_alg_runner(env=env, name=args.task, args=args)
|
||||
env.common_step_counter = runner.current_learning_iteration * env.num_steps_per_env # resume env step counter
|
||||
env.update_reward_curriculum(force_update=True) # force update reward curriculum at start
|
||||
runner.learn(num_learning_iterations=train_cfg.runner.max_iterations, init_at_random_ep_len=True)
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = get_args()
|
||||
train(args)
|
||||
5
legged_gym/utils/__init__.py
Normal file
5
legged_gym/utils/__init__.py
Normal file
@@ -0,0 +1,5 @@
|
||||
from .helpers import class_to_dict, get_load_path, get_args, set_seed, update_class_from_dict
|
||||
from .task_registry import task_registry
|
||||
from .logger import Logger
|
||||
from .math import *
|
||||
from .terrain import Terrain
|
||||
268
legged_gym/utils/exporter.py
Normal file
268
legged_gym/utils/exporter.py
Normal file
@@ -0,0 +1,268 @@
|
||||
# Copyright (c) 2022-2025, The Isaac Lab Project Developers.
|
||||
# All rights reserved.
|
||||
#
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
import copy
|
||||
import os
|
||||
import torch
|
||||
from torch import nn
|
||||
from typing import Optional
|
||||
|
||||
|
||||
def export_policy_as_jit(policy: object, path: str, normalizer: Optional[object] = None, filename="policy.pt"):
|
||||
"""Export policy into a Torch JIT file.
|
||||
|
||||
Args:
|
||||
policy: The policy torch module.
|
||||
normalizer: The empirical normalizer module. If None, Identity is used.
|
||||
path: The path to the saving directory.
|
||||
filename: The name of exported JIT file. Defaults to "policy.pt".
|
||||
"""
|
||||
policy_exporter = _TorchPolicyExporter(policy, normalizer)
|
||||
policy_exporter.export(path, filename)
|
||||
|
||||
|
||||
def export_policy_as_onnx(
|
||||
policy: object, path: str, normalizer: Optional[object] = None, filename="policy.onnx", verbose=False
|
||||
):
|
||||
"""Export policy into a Torch ONNX file.
|
||||
|
||||
Args:
|
||||
policy: The policy torch module.
|
||||
normalizer: The empirical normalizer module. If None, Identity is used.
|
||||
path: The path to the saving directory.
|
||||
filename: The name of exported ONNX file. Defaults to "policy.onnx".
|
||||
verbose: Whether to print the model summary. Defaults to False.
|
||||
"""
|
||||
if not os.path.exists(path):
|
||||
os.makedirs(path, exist_ok=True)
|
||||
policy_exporter = _OnnxPolicyExporter(policy, normalizer, verbose)
|
||||
policy_exporter.export(path, filename)
|
||||
|
||||
|
||||
|
||||
def export_policy_as_pkl(
|
||||
policy: nn.Module, path: str, filename="policy.pkl"
|
||||
):
|
||||
"""Export policy into a Torch pkl file.
|
||||
|
||||
Args:
|
||||
policy: The policy torch module.
|
||||
normalizer: The empirical normalizer module. If None, Identity is used.
|
||||
path: The path to the saving directory.
|
||||
filename: The name of exported pkl file. Defaults to "policy.pkl".
|
||||
"""
|
||||
if not os.path.exists(path):
|
||||
os.makedirs(path, exist_ok=True)
|
||||
model_dict = policy.state_dict()
|
||||
torch.save(model_dict, os.path.join(path, filename))
|
||||
|
||||
|
||||
"""
|
||||
Helper Classes - Private.
|
||||
"""
|
||||
|
||||
|
||||
class _TorchPolicyExporter(torch.nn.Module):
|
||||
"""Exporter of actor-critic into JIT file."""
|
||||
|
||||
def __init__(self, policy, normalizer=None):
|
||||
super().__init__()
|
||||
self.is_recurrent = policy.is_recurrent
|
||||
# copy policy parameters
|
||||
if hasattr(policy, "student_encoder"):
|
||||
self.student_encoder = copy.deepcopy(policy.student_encoder).cpu()
|
||||
self.history = torch.zeros([1, policy.history.shape[1], policy.history.shape[2]], device='cpu')
|
||||
self.forward = self.forward_cts
|
||||
if hasattr(policy, "student_moe_encoder"):
|
||||
self.student_moe_encoder = copy.deepcopy(policy.student_moe_encoder).cpu()
|
||||
self.obs_no_goal_mask = copy.deepcopy(policy.obs_no_goal_mask).cpu()
|
||||
self.history_length = policy.history.shape[1]
|
||||
self.history = torch.zeros([1, policy.history.shape[1], policy.history.shape[2]], device='cpu')
|
||||
self.forward = self.forward_moe_cts
|
||||
if hasattr(policy, "actor"):
|
||||
self.actor = copy.deepcopy(policy.actor)
|
||||
if self.is_recurrent:
|
||||
self.rnn = copy.deepcopy(policy.memory_a.rnn)
|
||||
elif hasattr(policy, "student"):
|
||||
self.actor = copy.deepcopy(policy.student)
|
||||
if self.is_recurrent:
|
||||
self.rnn = copy.deepcopy(policy.memory_s.rnn)
|
||||
else:
|
||||
raise ValueError("Policy does not have an actor/student module.")
|
||||
# set up recurrent network
|
||||
if self.is_recurrent:
|
||||
self.rnn.cpu()
|
||||
self.register_buffer("hidden_state", torch.zeros(self.rnn.num_layers, 1, self.rnn.hidden_size))
|
||||
self.register_buffer("cell_state", torch.zeros(self.rnn.num_layers, 1, self.rnn.hidden_size))
|
||||
self.forward = self.forward_lstm
|
||||
self.reset = self.reset_memory
|
||||
# copy normalizer if exists
|
||||
if normalizer:
|
||||
self.normalizer = copy.deepcopy(normalizer)
|
||||
else:
|
||||
self.normalizer = torch.nn.Identity()
|
||||
|
||||
def forward_lstm(self, x):
|
||||
x = self.normalizer(x)
|
||||
x, (h, c) = self.rnn(x.unsqueeze(0), (self.hidden_state, self.cell_state))
|
||||
self.hidden_state[:] = h
|
||||
self.cell_state[:] = c
|
||||
x = x.squeeze(0)
|
||||
return self.actor(x)
|
||||
|
||||
def forward(self, x):
|
||||
return self.actor(self.normalizer(x))
|
||||
|
||||
def forward_cts(self, x): # x is single observations
|
||||
x = self.normalizer(x)
|
||||
self.history = torch.cat([self.history[:, 1:], x.unsqueeze(1)], dim=1)
|
||||
latent = self.student_encoder(self.history.flatten(1))
|
||||
x = torch.cat([latent, x], dim=1)
|
||||
return self.actor(x)
|
||||
|
||||
def forward_moe_cts(self, x): # x is single observations
|
||||
x = self.normalizer(x)
|
||||
self.history = torch.cat([self.history[:, 1:], x.unsqueeze(1)], dim=1)
|
||||
history_no_goal = self.history.reshape(1, self.history_length, -1)[:, :, self.obs_no_goal_mask].reshape(1, -1)
|
||||
latent, weights = self.student_moe_encoder(self.history.flatten(1), history_no_goal)
|
||||
x = torch.cat([latent, x], dim=1)
|
||||
return self.actor(x), (weights, latent)
|
||||
|
||||
@torch.jit.export
|
||||
def reset(self):
|
||||
if hasattr(self, 'history'):
|
||||
self.history = torch.zeros_like(self.history)
|
||||
|
||||
def reset_memory(self):
|
||||
self.hidden_state[:] = 0.0
|
||||
self.cell_state[:] = 0.0
|
||||
|
||||
def export(self, path, filename):
|
||||
os.makedirs(path, exist_ok=True)
|
||||
path = os.path.join(path, filename)
|
||||
self.to("cpu")
|
||||
traced_script_module = torch.jit.script(self)
|
||||
traced_script_module.save(path)
|
||||
|
||||
|
||||
class _OnnxPolicyExporter(torch.nn.Module):
|
||||
"""Exporter of actor-critic into ONNX file."""
|
||||
|
||||
def __init__(self, policy, normalizer=None, verbose=False):
|
||||
super().__init__()
|
||||
self.verbose = verbose
|
||||
self.input_dim = None
|
||||
self.num_actions = 12
|
||||
self.normalizer = torch.nn.Identity()
|
||||
|
||||
# copy policy parameters
|
||||
if hasattr(policy, 'student_encoder'):
|
||||
self.student_encoder = copy.deepcopy(policy.student_encoder)
|
||||
self.forward = self.forward_cts
|
||||
self.input_dim = self.student_encoder[0].in_features
|
||||
|
||||
elif hasattr(policy, "student_moe_encoder"):
|
||||
self.student_moe_encoder = copy.deepcopy(policy.student_moe_encoder)
|
||||
self.obs_no_goal_mask = copy.deepcopy(policy.obs_no_goal_mask).cpu()
|
||||
self.history_length = policy.history.shape[1]
|
||||
self.forward = self.forward_moe_cts
|
||||
self.input_dim = self.history_length * policy.history.shape[2]
|
||||
|
||||
if hasattr(policy, "actor"):
|
||||
self.actor = copy.deepcopy(policy.actor)
|
||||
if hasattr(self, 'is_recurrent') and self.is_recurrent:
|
||||
self.rnn = copy.deepcopy(policy.memory_a.rnn)
|
||||
if self.input_dim is None:
|
||||
self.input_dim = self.actor[0].in_features
|
||||
else:
|
||||
raise ValueError("Policy does not have an actor/student module.")
|
||||
|
||||
def forward_cts(self, x): # x is stack observations by terms
|
||||
x = self.normalizer(x)
|
||||
term_dims = [3, 3, 3, self.num_actions, self.num_actions, self.num_actions]
|
||||
obs_dim = sum(term_dims)
|
||||
if x.shape[1] % obs_dim != 0:
|
||||
raise ValueError(f"x.shape[1] ({x.shape[1]}) 不是 obs_dim ({obs_dim}) 的整数倍")
|
||||
|
||||
frames = x.shape[1] // obs_dim
|
||||
split_sizes = [dim * frames for dim in term_dims]
|
||||
# [B, dim0*frames], [B, dim1*frames], ...
|
||||
term_chunks = torch.split(x, split_sizes, dim=1)
|
||||
|
||||
# [ [B, frames, dim0], [B, frames, dim1], ... ]
|
||||
frame_terms_reshaped = [
|
||||
chunk.view(-1, frames, dim)
|
||||
for chunk, dim in zip(term_chunks, term_dims)
|
||||
]
|
||||
|
||||
history_by_frame = []
|
||||
for i in range(frames):
|
||||
# [ [B, dim0], [B, dim1], ... ]
|
||||
terms_for_this_frame = [ftr[:, i, :] for ftr in frame_terms_reshaped]
|
||||
history_by_frame.append(torch.cat(terms_for_this_frame, dim=1))
|
||||
# [B, (Frame0_AllTerms), (Frame1_AllTerms), ...]
|
||||
history = torch.cat(history_by_frame, dim=1)
|
||||
|
||||
last_obs = history[:, -obs_dim:]
|
||||
latent = self.student_encoder(history)
|
||||
x = torch.cat([latent, last_obs], dim=1)
|
||||
|
||||
return self.actor(x)
|
||||
|
||||
def forward_moe_cts(self, x):
|
||||
x = self.normalizer(x)
|
||||
term_dims = [3, 3, 3, self.num_actions, self.num_actions, self.num_actions]
|
||||
obs_dim = sum(term_dims)
|
||||
if x.shape[1] % obs_dim != 0:
|
||||
raise ValueError(f"x.shape[1] ({x.shape[1]}) 不是 obs_dim ({obs_dim}) 的整数倍")
|
||||
|
||||
frames = x.shape[1] // obs_dim
|
||||
split_sizes = [dim * frames for dim in term_dims]
|
||||
# [B, dim0*frames], [B, dim1*frames], ...
|
||||
term_chunks = torch.split(x, split_sizes, dim=1)
|
||||
|
||||
# [ [B, frames, dim0], [B, frames, dim1], ... ]
|
||||
frame_terms_reshaped = [
|
||||
chunk.view(-1, frames, dim)
|
||||
for chunk, dim in zip(term_chunks, term_dims)
|
||||
]
|
||||
|
||||
history_by_frame = []
|
||||
for i in range(frames):
|
||||
# [ [B, dim0], [B, dim1], ... ]
|
||||
terms_for_this_frame = [ftr[:, i, :] for ftr in frame_terms_reshaped]
|
||||
history_by_frame.append(torch.cat(terms_for_this_frame, dim=1))
|
||||
# [B, (Frame0_AllTerms), (Frame1_AllTerms), ...]
|
||||
history = torch.cat(history_by_frame, dim=1)
|
||||
|
||||
last_obs = history[:, -obs_dim:]
|
||||
history_3d = history.view(-1, self.history_length, obs_dim)
|
||||
history_no_goal = history_3d[:, :, self.obs_no_goal_mask].reshape(x.shape[0], -1)
|
||||
|
||||
latent, weights = self.student_moe_encoder(history, history_no_goal)
|
||||
x = torch.cat([latent, last_obs], dim=1)
|
||||
|
||||
return self.actor(x), weights, latent
|
||||
|
||||
def export(self, path, filename):
|
||||
self.to("cpu")
|
||||
obs = torch.zeros(1, self.input_dim)
|
||||
|
||||
output_names = ["actions"]
|
||||
if self.forward == self.forward_moe_cts:
|
||||
output_names.append("weights")
|
||||
output_names.append("latent")
|
||||
|
||||
torch.onnx.export(
|
||||
self,
|
||||
obs,
|
||||
os.path.join(path, filename),
|
||||
export_params=True,
|
||||
opset_version=11,
|
||||
verbose=self.verbose,
|
||||
input_names=["obs"],
|
||||
output_names=output_names,
|
||||
dynamic_axes={},
|
||||
)
|
||||
193
legged_gym/utils/helpers.py
Normal file
193
legged_gym/utils/helpers.py
Normal file
@@ -0,0 +1,193 @@
|
||||
import os
|
||||
import copy
|
||||
import torch
|
||||
import numpy as np
|
||||
import random
|
||||
from pathlib import Path
|
||||
from isaacgym import gymapi
|
||||
from isaacgym import gymutil
|
||||
|
||||
from legged_gym import LEGGED_GYM_ROOT_DIR, LEGGED_GYM_ENVS_DIR
|
||||
|
||||
def class_to_dict(obj) -> dict:
|
||||
if not hasattr(obj,"__dict__"):
|
||||
return obj
|
||||
result = {}
|
||||
for key in dir(obj):
|
||||
if key.startswith("_"):
|
||||
continue
|
||||
element = []
|
||||
val = getattr(obj, key)
|
||||
if isinstance(val, list):
|
||||
for item in val:
|
||||
element.append(class_to_dict(item))
|
||||
else:
|
||||
element = class_to_dict(val)
|
||||
result[key] = element
|
||||
return result
|
||||
|
||||
def update_class_from_dict(obj, dict):
|
||||
for key, val in dict.items():
|
||||
attr = getattr(obj, key, None)
|
||||
if isinstance(attr, type):
|
||||
update_class_from_dict(attr, val)
|
||||
else:
|
||||
setattr(obj, key, val)
|
||||
return
|
||||
|
||||
def set_seed(seed):
|
||||
if seed == -1:
|
||||
seed = np.random.randint(0, 10000)
|
||||
print("Setting seed: {}".format(seed))
|
||||
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
os.environ['PYTHONHASHSEED'] = str(seed)
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
def parse_sim_params(args, cfg):
|
||||
# code from Isaac Gym Preview 2
|
||||
# initialize sim params
|
||||
sim_params = gymapi.SimParams()
|
||||
|
||||
# set some values from args
|
||||
if args.physics_engine == gymapi.SIM_FLEX:
|
||||
if args.device != "cpu":
|
||||
print("WARNING: Using Flex with GPU instead of PHYSX!")
|
||||
elif args.physics_engine == gymapi.SIM_PHYSX:
|
||||
sim_params.physx.use_gpu = args.use_gpu
|
||||
sim_params.physx.num_subscenes = args.subscenes
|
||||
sim_params.use_gpu_pipeline = args.use_gpu_pipeline
|
||||
|
||||
# if sim options are provided in cfg, parse them and update/override above:
|
||||
if "sim" in cfg:
|
||||
gymutil.parse_sim_config(cfg["sim"], sim_params)
|
||||
|
||||
# Override num_threads if passed on the command line
|
||||
if args.physics_engine == gymapi.SIM_PHYSX and args.num_threads > 0:
|
||||
sim_params.physx.num_threads = args.num_threads
|
||||
|
||||
return sim_params
|
||||
|
||||
def get_load_path(root, load_run=-1, checkpoint=-1):
|
||||
try:
|
||||
runs = os.listdir(root)
|
||||
runs = [run for run in runs if len(list((Path(root) / run).glob('model_*.pt'))) > 0]
|
||||
#TODO sort by date to handle change of month
|
||||
runs.sort()
|
||||
if 'exported' in runs: runs.remove('exported')
|
||||
last_run = os.path.join(root, runs[-1])
|
||||
except:
|
||||
raise ValueError("No runs in this directory: " + root)
|
||||
if load_run==-1:
|
||||
load_run = last_run
|
||||
else:
|
||||
load_run = os.path.join(root, load_run)
|
||||
|
||||
if checkpoint==-1:
|
||||
models = [file for file in os.listdir(load_run) if 'model' in file]
|
||||
models.sort(key=lambda m: '{0:0>15}'.format(m))
|
||||
model = models[-1]
|
||||
else:
|
||||
model = "model_{}.pt".format(checkpoint)
|
||||
|
||||
load_path = os.path.join(load_run, model)
|
||||
return load_path
|
||||
|
||||
def update_cfg_from_args(env_cfg, cfg_train, args):
|
||||
# seed
|
||||
if env_cfg is not None:
|
||||
# num envs
|
||||
if args.num_envs is not None:
|
||||
env_cfg.env.num_envs = args.num_envs
|
||||
if cfg_train is not None:
|
||||
if args.seed is not None:
|
||||
cfg_train.seed = args.seed
|
||||
# alg runner parameters
|
||||
if args.max_iterations is not None:
|
||||
cfg_train.runner.max_iterations = args.max_iterations
|
||||
if args.resume:
|
||||
cfg_train.runner.resume = args.resume
|
||||
if args.experiment_name is not None:
|
||||
cfg_train.runner.experiment_name = args.experiment_name
|
||||
if args.run_name is not None:
|
||||
cfg_train.runner.run_name = args.run_name
|
||||
if args.load_run is not None:
|
||||
cfg_train.runner.load_run = args.load_run
|
||||
if args.checkpoint is not None:
|
||||
cfg_train.runner.checkpoint = args.checkpoint
|
||||
|
||||
return env_cfg, cfg_train
|
||||
|
||||
def get_args():
|
||||
custom_parameters = [
|
||||
{"name": "--task", "type": str, "default": "go2", "help": "Resume training or start testing from a checkpoint. Overrides config file if provided."},
|
||||
{"name": "--resume", "action": "store_true", "default": False, "help": "Resume training from a checkpoint"},
|
||||
{"name": "--experiment_name", "type": str, "help": "Name of the experiment to run or load. Overrides config file if provided."},
|
||||
{"name": "--run_name", "type": str, "help": "Name of the run. Overrides config file if provided."},
|
||||
{"name": "--load_run", "type": str, "help": "Name of the run to load when resume=True. If -1: will load the last run. Overrides config file if provided."},
|
||||
{"name": "--checkpoint", "type": int, "help": "Saved model checkpoint number. If -1: will load the last checkpoint. Overrides config file if provided."},
|
||||
|
||||
{"name": "--headless", "action": "store_true", "default": False, "help": "Force display off at all times"},
|
||||
{"name": "--horovod", "action": "store_true", "default": False, "help": "Use horovod for multi-gpu training"},
|
||||
{"name": "--rl_device", "type": str, "default": "cuda:0", "help": 'Device used by the RL algorithm, (cpu, gpu, cuda:0, cuda:1 etc..)'},
|
||||
{"name": "--num_envs", "type": int, "help": "Number of environments to create. Overrides config file if provided."},
|
||||
{"name": "--seed", "type": int, "help": "Random seed. Overrides config file if provided."},
|
||||
{"name": "--max_iterations", "type": int, "help": "Maximum number of training iterations. Overrides config file if provided."},
|
||||
]
|
||||
# parse arguments
|
||||
args = gymutil.parse_arguments(
|
||||
description="RL Policy",
|
||||
custom_parameters=custom_parameters)
|
||||
|
||||
# name allignment
|
||||
args.sim_device_id = args.compute_device_id
|
||||
args.sim_device = args.sim_device_type
|
||||
if args.sim_device=='cuda':
|
||||
args.sim_device += f":{args.sim_device_id}"
|
||||
return args
|
||||
|
||||
# def export_policy_as_jit(actor_critic, path):
|
||||
# if hasattr(actor_critic, 'memory_a'):
|
||||
# # assumes LSTM: TODO add GRU
|
||||
# exporter = PolicyExporterLSTM(actor_critic)
|
||||
# exporter.export(path)
|
||||
# else:
|
||||
# os.makedirs(path, exist_ok=True)
|
||||
# path = os.path.join(path, 'policy_1.pt')
|
||||
# model = copy.deepcopy(actor_critic.actor).to('cpu')
|
||||
# traced_script_module = torch.jit.script(model)
|
||||
# traced_script_module.save(path)
|
||||
|
||||
|
||||
# class PolicyExporterLSTM(torch.nn.Module):
|
||||
# def __init__(self, actor_critic):
|
||||
# super().__init__()
|
||||
# self.actor = copy.deepcopy(actor_critic.actor)
|
||||
# self.is_recurrent = actor_critic.is_recurrent
|
||||
# self.memory = copy.deepcopy(actor_critic.memory_a.rnn)
|
||||
# self.memory.cpu()
|
||||
# self.register_buffer(f'hidden_state', torch.zeros(self.memory.num_layers, 1, self.memory.hidden_size))
|
||||
# self.register_buffer(f'cell_state', torch.zeros(self.memory.num_layers, 1, self.memory.hidden_size))
|
||||
|
||||
# def forward(self, x):
|
||||
# out, (h, c) = self.memory(x.unsqueeze(0), (self.hidden_state, self.cell_state))
|
||||
# self.hidden_state[:] = h
|
||||
# self.cell_state[:] = c
|
||||
# return self.actor(out.squeeze(0))
|
||||
|
||||
# @torch.jit.export
|
||||
# def reset_memory(self):
|
||||
# self.hidden_state[:] = 0.
|
||||
# self.cell_state[:] = 0.
|
||||
|
||||
# def export(self, path):
|
||||
# os.makedirs(path, exist_ok=True)
|
||||
# path = os.path.join(path, 'policy_lstm_1.pt')
|
||||
# self.to('cpu')
|
||||
# traced_script_module = torch.jit.script(self)
|
||||
# traced_script_module.save(path)
|
||||
|
||||
|
||||
47
legged_gym/utils/isaacgym_utils.py
Normal file
47
legged_gym/utils/isaacgym_utils.py
Normal file
@@ -0,0 +1,47 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import random
|
||||
import torch
|
||||
|
||||
@torch.jit.script
|
||||
def copysign(a, b):
|
||||
# type: (float, Tensor) -> Tensor
|
||||
a = torch.tensor(a, device=b.device, dtype=torch.float).repeat(b.shape[0])
|
||||
return torch.abs(a) * torch.sign(b)
|
||||
def get_euler_xyz(q):
|
||||
qx, qy, qz, qw = 0, 1, 2, 3
|
||||
# roll (x-axis rotation)
|
||||
sinr_cosp = 2.0 * (q[:, qw] * q[:, qx] + q[:, qy] * q[:, qz])
|
||||
cosr_cosp = q[:, qw] * q[:, qw] - q[:, qx] * \
|
||||
q[:, qx] - q[:, qy] * q[:, qy] + q[:, qz] * q[:, qz]
|
||||
roll = torch.atan2(sinr_cosp, cosr_cosp)
|
||||
|
||||
# pitch (y-axis rotation)
|
||||
sinp = 2.0 * (q[:, qw] * q[:, qy] - q[:, qz] * q[:, qx])
|
||||
pitch = torch.where(
|
||||
torch.abs(sinp) >= 1, copysign(np.pi / 2.0, sinp), torch.asin(sinp))
|
||||
|
||||
# yaw (z-axis rotation)
|
||||
siny_cosp = 2.0 * (q[:, qw] * q[:, qz] + q[:, qx] * q[:, qy])
|
||||
cosy_cosp = q[:, qw] * q[:, qw] + q[:, qx] * \
|
||||
q[:, qx] - q[:, qy] * q[:, qy] - q[:, qz] * q[:, qz]
|
||||
yaw = torch.atan2(siny_cosp, cosy_cosp)
|
||||
|
||||
return torch.stack((roll, pitch, yaw), dim=-1)
|
||||
|
||||
def sample_disjoint_intervals(env_ids, limit_bound, cfg_min, cfg_max, device):
|
||||
"""
|
||||
sample uniform distribution from [cfg_min, -limit_bound] U [limit_bound, cfg_max]
|
||||
"""
|
||||
width_neg = torch.nn.functional.relu(-limit_bound - cfg_min)
|
||||
width_pos = torch.nn.functional.relu(cfg_max - limit_bound)
|
||||
|
||||
total_width = width_neg + width_pos + 1e-6 # 加极小值防除零
|
||||
u = torch.rand(len(env_ids), device=device) * total_width
|
||||
|
||||
samples = torch.where(
|
||||
u < width_neg,
|
||||
cfg_min + u,
|
||||
cfg_max - width_pos + (u - width_neg)
|
||||
)
|
||||
return samples
|
||||
39
legged_gym/utils/logger.py
Normal file
39
legged_gym/utils/logger.py
Normal file
@@ -0,0 +1,39 @@
|
||||
import numpy as np
|
||||
from collections import defaultdict
|
||||
from multiprocessing import Process, Value
|
||||
|
||||
class Logger:
|
||||
def __init__(self, dt):
|
||||
self.state_log = defaultdict(list)
|
||||
self.rew_log = defaultdict(list)
|
||||
self.dt = dt
|
||||
self.num_episodes = 0
|
||||
self.plot_process = None
|
||||
|
||||
def log_state(self, key, value):
|
||||
self.state_log[key].append(value)
|
||||
|
||||
def log_states(self, dict):
|
||||
for key, value in dict.items():
|
||||
self.log_state(key, value)
|
||||
|
||||
def log_rewards(self, dict, num_episodes):
|
||||
for key, value in dict.items():
|
||||
if 'rew' in key:
|
||||
self.rew_log[key].append(value.item() * num_episodes)
|
||||
self.num_episodes += num_episodes
|
||||
|
||||
def reset(self):
|
||||
self.state_log.clear()
|
||||
self.rew_log.clear()
|
||||
|
||||
def print_rewards(self):
|
||||
print("Average rewards per second:")
|
||||
for key, values in self.rew_log.items():
|
||||
mean = np.sum(np.array(values)) / self.num_episodes
|
||||
print(f" - {key}: {mean}")
|
||||
print(f"Total number of episodes: {self.num_episodes}")
|
||||
|
||||
def __del__(self):
|
||||
if self.plot_process is not None:
|
||||
self.plot_process.kill()
|
||||
26
legged_gym/utils/math.py
Normal file
26
legged_gym/utils/math.py
Normal file
@@ -0,0 +1,26 @@
|
||||
import torch
|
||||
from torch import Tensor
|
||||
import numpy as np
|
||||
from isaacgym.torch_utils import quat_apply, normalize
|
||||
from typing import Tuple
|
||||
|
||||
# @ torch.jit.script
|
||||
def quat_apply_yaw(quat, vec):
|
||||
quat_yaw = quat.clone().view(-1, 4)
|
||||
quat_yaw[:, :2] = 0.
|
||||
quat_yaw = normalize(quat_yaw)
|
||||
return quat_apply(quat_yaw, vec)
|
||||
|
||||
# @ torch.jit.script
|
||||
def wrap_to_pi(angles):
|
||||
angles %= 2*np.pi
|
||||
angles -= 2*np.pi * (angles > np.pi)
|
||||
return angles
|
||||
|
||||
# @ torch.jit.script
|
||||
def torch_rand_sqrt_float(lower, upper, shape, device):
|
||||
# type: (float, float, Tuple[int, int], str) -> Tensor
|
||||
r = 2*torch.rand(*shape, device=device) - 1
|
||||
r = torch.where(r<0., -torch.sqrt(-r), torch.sqrt(r))
|
||||
r = (r + 1.) / 2.
|
||||
return (upper - lower) * r + lower
|
||||
129
legged_gym/utils/task_registry.py
Normal file
129
legged_gym/utils/task_registry.py
Normal file
@@ -0,0 +1,129 @@
|
||||
import os
|
||||
from datetime import datetime
|
||||
from typing import Tuple
|
||||
import torch
|
||||
import numpy as np
|
||||
import sys
|
||||
|
||||
from rsl_rl.env import VecEnv
|
||||
from rsl_rl.runners import OnPolicyRunner, OnPolicyRunnerCTS
|
||||
|
||||
from legged_gym import LEGGED_GYM_ROOT_DIR, LEGGED_GYM_ENVS_DIR
|
||||
from .helpers import get_args, update_cfg_from_args, class_to_dict, get_load_path, set_seed, parse_sim_params
|
||||
from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg, LeggedRobotCfgPPO
|
||||
|
||||
class TaskRegistry():
|
||||
def __init__(self):
|
||||
self.task_classes = {}
|
||||
self.env_cfgs = {}
|
||||
self.train_cfgs = {}
|
||||
|
||||
def register(self, name: str, task_class: VecEnv, env_cfg: LeggedRobotCfg, train_cfg: LeggedRobotCfgPPO):
|
||||
self.task_classes[name] = task_class
|
||||
self.env_cfgs[name] = env_cfg
|
||||
self.train_cfgs[name] = train_cfg
|
||||
|
||||
def get_task_class(self, name: str) -> VecEnv:
|
||||
return self.task_classes[name]
|
||||
|
||||
def get_cfgs(self, name) -> Tuple[LeggedRobotCfg, LeggedRobotCfgPPO]:
|
||||
train_cfg = self.train_cfgs[name]
|
||||
env_cfg = self.env_cfgs[name]
|
||||
# copy seed
|
||||
env_cfg.seed = train_cfg.seed
|
||||
return env_cfg, train_cfg
|
||||
|
||||
def make_env(self, name, args=None, env_cfg=None) -> Tuple[VecEnv, LeggedRobotCfg]:
|
||||
""" Creates an environment either from a registered namme or from the provided config file.
|
||||
|
||||
Args:
|
||||
name (string): Name of a registered env.
|
||||
args (Args, optional): Isaac Gym comand line arguments. If None get_args() will be called. Defaults to None.
|
||||
env_cfg (Dict, optional): Environment config file used to override the registered config. Defaults to None.
|
||||
|
||||
Raises:
|
||||
ValueError: Error if no registered env corresponds to 'name'
|
||||
|
||||
Returns:
|
||||
isaacgym.VecTaskPython: The created environment
|
||||
Dict: the corresponding config file
|
||||
"""
|
||||
# if no args passed get command line arguments
|
||||
if args is None:
|
||||
args = get_args()
|
||||
# check if there is a registered env with that name
|
||||
if name in self.task_classes:
|
||||
task_class = self.get_task_class(name)
|
||||
else:
|
||||
raise ValueError(f"Task with name: {name} was not registered")
|
||||
if env_cfg is None:
|
||||
# load config files
|
||||
env_cfg, _ = self.get_cfgs(name)
|
||||
# override cfg from args (if specified)
|
||||
env_cfg, _ = update_cfg_from_args(env_cfg, None, args)
|
||||
set_seed(env_cfg.seed)
|
||||
# parse sim params (convert to dict first)
|
||||
sim_params = {"sim": class_to_dict(env_cfg.sim)}
|
||||
sim_params = parse_sim_params(args, sim_params)
|
||||
env = task_class( cfg=env_cfg,
|
||||
sim_params=sim_params,
|
||||
physics_engine=args.physics_engine,
|
||||
sim_device=args.sim_device,
|
||||
headless=args.headless)
|
||||
return env, env_cfg
|
||||
|
||||
def make_alg_runner(self, env, name=None, args=None, train_cfg=None, log_root="default") -> Tuple[OnPolicyRunner, LeggedRobotCfgPPO]:
|
||||
""" Creates the training algorithm either from a registered namme or from the provided config file.
|
||||
|
||||
Args:
|
||||
env (isaacgym.VecTaskPython): The environment to train (TODO: remove from within the algorithm)
|
||||
name (string, optional): Name of a registered env. If None, the config file will be used instead. Defaults to None.
|
||||
args (Args, optional): Isaac Gym comand line arguments. If None get_args() will be called. Defaults to None.
|
||||
train_cfg (Dict, optional): Training config file. If None 'name' will be used to get the config file. Defaults to None.
|
||||
log_root (str, optional): Logging directory for Tensorboard. Set to 'None' to avoid logging (at test time for example).
|
||||
Logs will be saved in <log_root>/<date_time>_<run_name>. Defaults to "default"=<path_to_LEGGED_GYM>/logs/<experiment_name>.
|
||||
|
||||
Raises:
|
||||
ValueError: Error if neither 'name' or 'train_cfg' are provided
|
||||
Warning: If both 'name' or 'train_cfg' are provided 'name' is ignored
|
||||
|
||||
Returns:
|
||||
PPO: The created algorithm
|
||||
Dict: the corresponding config file
|
||||
"""
|
||||
# if no args passed get command line arguments
|
||||
if args is None:
|
||||
args = get_args()
|
||||
# if config files are passed use them, otherwise load from the name
|
||||
if train_cfg is None:
|
||||
if name is None:
|
||||
raise ValueError("Either 'name' or 'train_cfg' must be not None")
|
||||
# load config files
|
||||
_, train_cfg = self.get_cfgs(name)
|
||||
else:
|
||||
if name is not None:
|
||||
print(f"'train_cfg' provided -> Ignoring 'name={name}'")
|
||||
# override cfg from args (if specified)
|
||||
_, train_cfg = update_cfg_from_args(None, train_cfg, args)
|
||||
|
||||
if log_root=="default":
|
||||
log_root = os.path.join(LEGGED_GYM_ROOT_DIR, 'logs', train_cfg.runner.experiment_name)
|
||||
log_dir = os.path.join(log_root, datetime.now().strftime('%b%d_%H-%M-%S') + '_' + train_cfg.runner.run_name)
|
||||
elif log_root is None:
|
||||
log_dir = None
|
||||
else:
|
||||
log_dir = os.path.join(log_root, datetime.now().strftime('%b%d_%H-%M-%S') + '_' + train_cfg.runner.run_name)
|
||||
|
||||
train_cfg_dict = class_to_dict(train_cfg)
|
||||
runner = eval(train_cfg.runner_class_name)(env, train_cfg_dict, log_dir, device=args.rl_device)
|
||||
#save resume path before creating a new log_dir
|
||||
resume = train_cfg.runner.resume
|
||||
if resume:
|
||||
# load previously trained model
|
||||
resume_path = get_load_path(log_root, load_run=train_cfg.runner.load_run, checkpoint=train_cfg.runner.checkpoint)
|
||||
print(f"Loading model from: {resume_path}")
|
||||
runner.load(resume_path)
|
||||
return runner, train_cfg
|
||||
|
||||
# make global task registry
|
||||
task_registry = TaskRegistry()
|
||||
188
legged_gym/utils/terrain.py
Normal file
188
legged_gym/utils/terrain.py
Normal file
@@ -0,0 +1,188 @@
|
||||
from collections import defaultdict
|
||||
import numpy as np
|
||||
from numpy.random import choice
|
||||
from scipy import interpolate
|
||||
|
||||
from isaacgym import terrain_utils
|
||||
from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg
|
||||
|
||||
class Terrain:
|
||||
def __init__(self, cfg: LeggedRobotCfg.terrain, num_robots) -> None:
|
||||
|
||||
self.cfg = cfg
|
||||
self.num_robots = num_robots
|
||||
self.type = cfg.mesh_type
|
||||
if self.type in ["none", 'plane']:
|
||||
return
|
||||
self.env_length = cfg.terrain_length
|
||||
self.env_width = cfg.terrain_width
|
||||
self.proportions = [np.sum(cfg.terrain_proportions[:i+1]) for i in range(len(cfg.terrain_proportions))]
|
||||
|
||||
self.cfg.num_sub_terrains = cfg.num_rows * cfg.num_cols
|
||||
self.env_origins = np.zeros((cfg.num_rows, cfg.num_cols, 3))
|
||||
|
||||
self.width_per_env_pixels = int(self.env_width / cfg.horizontal_scale)
|
||||
self.length_per_env_pixels = int(self.env_length / cfg.horizontal_scale)
|
||||
|
||||
self.spacing = cfg.terrain_spacing
|
||||
self.spacing_pixels = int(self.spacing / cfg.horizontal_scale)
|
||||
|
||||
self.border = int(cfg.border_size/self.cfg.horizontal_scale)
|
||||
self.tot_cols = int(cfg.num_cols * self.width_per_env_pixels + max(0, cfg.num_cols-1) * self.spacing_pixels) + 2 * self.border
|
||||
self.tot_rows = int(cfg.num_rows * self.length_per_env_pixels + max(0, cfg.num_rows-1) * self.spacing_pixels) + 2 * self.border
|
||||
self.name2cols = defaultdict(set) # terrain type to column index
|
||||
self.cols2id = [] # column index to terrain id
|
||||
|
||||
self.height_field_raw = np.zeros((self.tot_rows , self.tot_cols), dtype=np.int16)
|
||||
if cfg.curriculum:
|
||||
self.curiculum()
|
||||
elif cfg.selected:
|
||||
self.selected_terrain()
|
||||
else:
|
||||
self.randomized_terrain()
|
||||
|
||||
self.heightsamples = self.height_field_raw
|
||||
if self.type=="trimesh":
|
||||
self.vertices, self.triangles = terrain_utils.convert_heightfield_to_trimesh( self.height_field_raw,
|
||||
self.cfg.horizontal_scale,
|
||||
self.cfg.vertical_scale,
|
||||
self.cfg.slope_treshold)
|
||||
|
||||
def randomized_terrain(self):
|
||||
for k in range(self.cfg.num_sub_terrains):
|
||||
# Env coordinates in the world
|
||||
(i, j) = np.unravel_index(k, (self.cfg.num_rows, self.cfg.num_cols))
|
||||
|
||||
choice = np.random.uniform(0, 1)
|
||||
difficulty = np.random.choice([0.5, 0.75, 0.9])
|
||||
terrain = self.make_terrain(choice, difficulty)
|
||||
self.add_terrain_to_map(terrain, i, j)
|
||||
|
||||
def curiculum(self):
|
||||
for j in range(self.cfg.num_cols):
|
||||
for i in range(self.cfg.num_rows):
|
||||
difficulty = i / self.cfg.num_rows
|
||||
choice = j / self.cfg.num_cols + 0.001
|
||||
|
||||
terrain = self.make_terrain(choice, difficulty)
|
||||
self.add_terrain_to_map(terrain, i, j)
|
||||
self.name2cols[terrain.terrain_name].add(j)
|
||||
self.cols2id.append(terrain.terrain_id)
|
||||
|
||||
def selected_terrain(self):
|
||||
terrain_type = self.cfg.terrain_kwargs.pop('type')
|
||||
for k in range(self.cfg.num_sub_terrains):
|
||||
# Env coordinates in the world
|
||||
(i, j) = np.unravel_index(k, (self.cfg.num_rows, self.cfg.num_cols))
|
||||
|
||||
terrain = terrain_utils.SubTerrain("terrain",
|
||||
width=self.width_per_env_pixels,
|
||||
length=self.width_per_env_pixels,
|
||||
vertical_scale=self.vertical_scale,
|
||||
horizontal_scale=self.horizontal_scale)
|
||||
|
||||
eval(terrain_type)(terrain, **self.cfg.terrain_kwargs.terrain_kwargs)
|
||||
self.add_terrain_to_map(terrain, i, j)
|
||||
|
||||
def make_terrain(self, choice, difficulty):
|
||||
terrain = terrain_utils.SubTerrain("terrain",
|
||||
width=self.width_per_env_pixels,
|
||||
length=self.width_per_env_pixels,
|
||||
vertical_scale=self.cfg.vertical_scale,
|
||||
horizontal_scale=self.cfg.horizontal_scale)
|
||||
slope = 0.1 + difficulty * 0.52
|
||||
step_height = 0.05 + 0.23 * difficulty
|
||||
discrete_obstacles_height = 0.05 + difficulty * 0.25
|
||||
stepping_stones_size = 1.5 * (1.05 - difficulty)
|
||||
stone_distance = 0.05 if difficulty==0 else 0.1
|
||||
gap_size = 1. * difficulty
|
||||
# pit_depth = 1. * difficulty # 删除坑洞
|
||||
amplitude = 0.1 + 0.2 * difficulty
|
||||
|
||||
if choice < self.proportions[0]:
|
||||
terrain.terrain_name = "wave"
|
||||
terrain.terrain_id = 0
|
||||
terrain_utils.wave_terrain(terrain, num_waves=5, amplitude=amplitude)
|
||||
terrain_utils.random_uniform_terrain(terrain, min_height=-0.05, max_height=0.05, step=0.005, downsampled_scale=0.2)
|
||||
elif choice < self.proportions[1]: # 平滑坡
|
||||
terrain.terrain_name = "slope"
|
||||
terrain.terrain_id = 1
|
||||
if choice < (self.proportions[0] + self.proportions[1])/ 2: # 一半正坡, 一半负坡
|
||||
slope *= -1
|
||||
terrain_utils.pyramid_sloped_terrain(terrain, slope=slope, platform_size=3.)
|
||||
elif choice < self.proportions[2]: # 粗糙坡
|
||||
terrain.terrain_name = "rough_slope"
|
||||
terrain.terrain_id = 2
|
||||
terrain_utils.pyramid_sloped_terrain(terrain, slope=slope, platform_size=3.)
|
||||
terrain_utils.random_uniform_terrain(terrain, min_height=-0.05, max_height=0.05, step=0.005, downsampled_scale=0.2)
|
||||
elif choice < self.proportions[4]: # 下楼梯
|
||||
terrain.terrain_name = "stairs_down"
|
||||
terrain.terrain_id = 4
|
||||
if choice<self.proportions[3]: # 上楼梯
|
||||
terrain.terrain_name = "stairs_up"
|
||||
terrain.terrain_id = 3
|
||||
step_height *= -1
|
||||
terrain_utils.pyramid_stairs_terrain(terrain, step_width=0.31, step_height=step_height, platform_size=3.)
|
||||
elif choice < self.proportions[5]: # 障碍物
|
||||
terrain.terrain_name = "obstacles"
|
||||
terrain.terrain_id = 5
|
||||
num_rectangles = 20
|
||||
rectangle_min_size = 1.
|
||||
rectangle_max_size = 2.
|
||||
terrain_utils.discrete_obstacles_terrain(terrain, discrete_obstacles_height, rectangle_min_size, rectangle_max_size, num_rectangles, platform_size=3.)
|
||||
elif choice < self.proportions[6]: # 梅花桩
|
||||
terrain.terrain_name = "stepping_stones"
|
||||
terrain.terrain_id = 6
|
||||
terrain_utils.stepping_stones_terrain(terrain, stone_size=stepping_stones_size, stone_distance=stone_distance, max_height=0., platform_size=4.)
|
||||
elif choice < self.proportions[7]: # 间隙
|
||||
terrain.terrain_name = "gap"
|
||||
terrain.terrain_id = 7
|
||||
gap_terrain(terrain, gap_size=gap_size, platform_size=3.)
|
||||
else: # 平地
|
||||
terrain.terrain_name = "flat"
|
||||
terrain.terrain_id = 8
|
||||
pit_terrain(terrain, depth=0.0, platform_size=4.)
|
||||
|
||||
return terrain
|
||||
|
||||
def add_terrain_to_map(self, terrain, row, col):
|
||||
i = row
|
||||
j = col
|
||||
# map coordinate system
|
||||
start_x = self.border + i * (self.length_per_env_pixels + self.spacing_pixels)
|
||||
end_x = start_x + self.length_per_env_pixels
|
||||
start_y = self.border + j * (self.width_per_env_pixels + self.spacing_pixels)
|
||||
end_y = start_y + self.width_per_env_pixels
|
||||
self.height_field_raw[start_x: end_x, start_y:end_y] = terrain.height_field_raw
|
||||
|
||||
env_origin_x = (i + 0.5) * self.env_length + i * self.spacing
|
||||
env_origin_y = (j + 0.5) * self.env_width + j * self.spacing
|
||||
x1 = int((self.env_length/2. - 1) / terrain.horizontal_scale)
|
||||
x2 = int((self.env_length/2. + 1) / terrain.horizontal_scale)
|
||||
y1 = int((self.env_width/2. - 1) / terrain.horizontal_scale)
|
||||
y2 = int((self.env_width/2. + 1) / terrain.horizontal_scale)
|
||||
env_origin_z = np.max(terrain.height_field_raw[x1:x2, y1:y2])*terrain.vertical_scale
|
||||
self.env_origins[i, j] = [env_origin_x, env_origin_y, env_origin_z]
|
||||
|
||||
def gap_terrain(terrain, gap_size, platform_size=1.):
|
||||
gap_size = int(gap_size / terrain.horizontal_scale)
|
||||
platform_size = int(platform_size / terrain.horizontal_scale)
|
||||
|
||||
center_x = terrain.length // 2
|
||||
center_y = terrain.width // 2
|
||||
x1 = (terrain.length - platform_size) // 2
|
||||
x2 = x1 + gap_size
|
||||
y1 = (terrain.width - platform_size) // 2
|
||||
y2 = y1 + gap_size
|
||||
|
||||
terrain.height_field_raw[center_x-x2 : center_x + x2, center_y-y2 : center_y + y2] = -1000
|
||||
terrain.height_field_raw[center_x-x1 : center_x + x1, center_y-y1 : center_y + y1] = 0
|
||||
|
||||
def pit_terrain(terrain, depth, platform_size=1.):
|
||||
depth = int(depth / terrain.vertical_scale)
|
||||
platform_size = int(platform_size / terrain.horizontal_scale / 2)
|
||||
x1 = terrain.length // 2 - platform_size
|
||||
x2 = terrain.length // 2 + platform_size
|
||||
y1 = terrain.width // 2 - platform_size
|
||||
y2 = terrain.width // 2 + platform_size
|
||||
terrain.height_field_raw[x1:x2, y1:y2] = -depth
|
||||
36346
resources/robots/go2/assets/base_0.obj
Normal file
36346
resources/robots/go2/assets/base_0.obj
Normal file
File diff suppressed because it is too large
Load Diff
23131
resources/robots/go2/assets/base_1.obj
Normal file
23131
resources/robots/go2/assets/base_1.obj
Normal file
File diff suppressed because it is too large
Load Diff
8367
resources/robots/go2/assets/base_2.obj
Normal file
8367
resources/robots/go2/assets/base_2.obj
Normal file
File diff suppressed because it is too large
Load Diff
10932
resources/robots/go2/assets/base_3.obj
Normal file
10932
resources/robots/go2/assets/base_3.obj
Normal file
File diff suppressed because it is too large
Load Diff
199205
resources/robots/go2/assets/base_4.obj
Normal file
199205
resources/robots/go2/assets/base_4.obj
Normal file
File diff suppressed because it is too large
Load Diff
24148
resources/robots/go2/assets/calf_0.obj
Normal file
24148
resources/robots/go2/assets/calf_0.obj
Normal file
File diff suppressed because it is too large
Load Diff
9243
resources/robots/go2/assets/calf_1.obj
Normal file
9243
resources/robots/go2/assets/calf_1.obj
Normal file
File diff suppressed because it is too large
Load Diff
24145
resources/robots/go2/assets/calf_mirror_0.obj
Normal file
24145
resources/robots/go2/assets/calf_mirror_0.obj
Normal file
File diff suppressed because it is too large
Load Diff
9252
resources/robots/go2/assets/calf_mirror_1.obj
Normal file
9252
resources/robots/go2/assets/calf_mirror_1.obj
Normal file
File diff suppressed because it is too large
Load Diff
30786
resources/robots/go2/assets/foot.obj
Normal file
30786
resources/robots/go2/assets/foot.obj
Normal file
File diff suppressed because it is too large
Load Diff
BIN
resources/robots/go2/assets/height_field.png
Normal file
BIN
resources/robots/go2/assets/height_field.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 3.2 KiB |
76960
resources/robots/go2/assets/hip_0.obj
Normal file
76960
resources/robots/go2/assets/hip_0.obj
Normal file
File diff suppressed because it is too large
Load Diff
73961
resources/robots/go2/assets/hip_1.obj
Normal file
73961
resources/robots/go2/assets/hip_1.obj
Normal file
File diff suppressed because it is too large
Load Diff
75706
resources/robots/go2/assets/thigh_0.obj
Normal file
75706
resources/robots/go2/assets/thigh_0.obj
Normal file
File diff suppressed because it is too large
Load Diff
40483
resources/robots/go2/assets/thigh_1.obj
Normal file
40483
resources/robots/go2/assets/thigh_1.obj
Normal file
File diff suppressed because it is too large
Load Diff
75572
resources/robots/go2/assets/thigh_mirror_0.obj
Normal file
75572
resources/robots/go2/assets/thigh_mirror_0.obj
Normal file
File diff suppressed because it is too large
Load Diff
40450
resources/robots/go2/assets/thigh_mirror_1.obj
Normal file
40450
resources/robots/go2/assets/thigh_mirror_1.obj
Normal file
File diff suppressed because it is too large
Load Diff
BIN
resources/robots/go2/assets/wood.png
Normal file
BIN
resources/robots/go2/assets/wood.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 919 KiB |
217
resources/robots/go2/dae/base.dae
Normal file
217
resources/robots/go2/dae/base.dae
Normal file
File diff suppressed because one or more lines are too long
130
resources/robots/go2/dae/calf.dae
Normal file
130
resources/robots/go2/dae/calf.dae
Normal file
File diff suppressed because one or more lines are too long
130
resources/robots/go2/dae/calf_mirror.dae
Normal file
130
resources/robots/go2/dae/calf_mirror.dae
Normal file
File diff suppressed because one or more lines are too long
89
resources/robots/go2/dae/foot.dae
Normal file
89
resources/robots/go2/dae/foot.dae
Normal file
File diff suppressed because one or more lines are too long
130
resources/robots/go2/dae/hip.dae
Normal file
130
resources/robots/go2/dae/hip.dae
Normal file
File diff suppressed because one or more lines are too long
130
resources/robots/go2/dae/thigh.dae
Normal file
130
resources/robots/go2/dae/thigh.dae
Normal file
File diff suppressed because one or more lines are too long
130
resources/robots/go2/dae/thigh_mirror.dae
Normal file
130
resources/robots/go2/dae/thigh_mirror.dae
Normal file
File diff suppressed because one or more lines are too long
22
resources/robots/go2/flat.xml
Normal file
22
resources/robots/go2/flat.xml
Normal file
@@ -0,0 +1,22 @@
|
||||
<mujoco model="go2 scene">
|
||||
<include file="go2.xml"/>
|
||||
|
||||
<statistic center="1.0 0.7 1.0" extent="0.8"/>
|
||||
|
||||
<visual>
|
||||
<headlight diffuse="0.6 0.6 0.6" ambient="0.1 0.1 0.1" specular="0.9 0.9 0.9"/>
|
||||
<rgba haze="0.15 0.25 0.35 1"/>
|
||||
<global azimuth="-140" elevation="-20"/>
|
||||
</visual>
|
||||
|
||||
<asset>
|
||||
<texture type="skybox" builtin="flat" rgb1="0 0 0" 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"/>
|
||||
</asset>
|
||||
|
||||
<worldbody>
|
||||
<light pos="1 0 3.5" dir="0 0 -1" directional="true"/>
|
||||
<geom name="floor" size="0 0 0.05" type="plane" material="groundplane"/>
|
||||
</worldbody>
|
||||
</mujoco>
|
||||
184
resources/robots/go2/go2.xml
Normal file
184
resources/robots/go2/go2.xml
Normal file
@@ -0,0 +1,184 @@
|
||||
<mujoco model="go2_description">
|
||||
|
||||
<compiler angle="radian" meshdir="assets"/>
|
||||
|
||||
|
||||
<statistic meansize="0.144785" extent="1.23314" center="0.025392 2.0634e-05 -0.245975"/>
|
||||
<default>
|
||||
<joint damping="0.001" armature="0.01" frictionloss="0.1"/>
|
||||
<default class="visual">
|
||||
<geom type="mesh" contype="0" conaffinity="0" group="2"/>
|
||||
</default>
|
||||
</default>
|
||||
|
||||
<asset>
|
||||
<material name="metal" rgba=".9 .95 .95 1" />
|
||||
<material name="black" rgba="0 0 0 1" />
|
||||
<material name="white" rgba="1 1 1 1" />
|
||||
<material name="gray" rgba="0.671705 0.692426 0.774270 1" />
|
||||
|
||||
<mesh file="base_0.obj" />
|
||||
<mesh file="base_1.obj" />
|
||||
<mesh file="base_2.obj" />
|
||||
<mesh file="base_3.obj" />
|
||||
<mesh file="base_4.obj" />
|
||||
<mesh file="hip_0.obj" />
|
||||
<mesh file="hip_1.obj" />
|
||||
<mesh file="thigh_0.obj" />
|
||||
<mesh file="thigh_1.obj" />
|
||||
<mesh file="thigh_mirror_0.obj" />
|
||||
<mesh file="thigh_mirror_1.obj" />
|
||||
<mesh file="calf_0.obj" />
|
||||
<mesh file="calf_1.obj" />
|
||||
<mesh file="calf_mirror_0.obj" />
|
||||
<mesh file="calf_mirror_1.obj" />
|
||||
<mesh file="foot.obj" />
|
||||
</asset>
|
||||
|
||||
|
||||
<worldbody>
|
||||
|
||||
<body name="base" pos="0 0 0.445">
|
||||
<!-- <body name="base" pos="-5 2 0.445"> -->
|
||||
<!-- <body name="base" pos="-5 2 0.30" quat="0 0 0.706 0.706"> -->
|
||||
<!-- <body name="base" pos="-5 2 0.445" quat="0 0 1 0"> -->
|
||||
<!-- <body name="base" pos="0 0 0.16" quat="0 0 0.706 0.706"> -->
|
||||
<!-- <body name="base" pos="0 0 0.1" quat="0 0 1 0"> -->
|
||||
<inertial pos="0.021112 0 -0.005366"
|
||||
mass="6.921"
|
||||
diaginertia="0.02448 0.098077 0.107" />
|
||||
<freejoint name="root"/>
|
||||
<geom mesh="base_0" material="black" class="visual" />
|
||||
<geom mesh="base_1" material="black" class="visual" />
|
||||
<geom mesh="base_2" material="black" class="visual" />
|
||||
<geom mesh="base_3" material="white" class="visual" />
|
||||
<geom mesh="base_4" material="gray" class="visual" />
|
||||
|
||||
<geom name="base_box" type="box" size="0.1881 0.04675 0.057"/>
|
||||
<geom name="base_cyl1" type="cylinder" size="0.05 0.045" pos="0.285 0 0.01"/>
|
||||
<geom name="base_sphere1" type="sphere" size="0.047" pos="0.293 0 -0.06"/>
|
||||
|
||||
<body name="FL_hip" pos="0.1934 0.0465 0">
|
||||
<inertial pos="-0.0054 0.00194 -0.000105" quat="0.497014 0.499245 0.505462 0.498237" mass="0.678" diaginertia="0.00088403 0.000596003 0.000479967"/>
|
||||
<joint name="FL_hip_joint" pos="0 0 0" axis="1 0 0" range="-1.0472 1.0472" actuatorfrcrange="-23.7 23.7"/>
|
||||
<geom mesh="hip_0" material="metal" class="visual" />
|
||||
<geom mesh="hip_1" material="gray" class="visual" />
|
||||
<geom size="0.046 0.02" pos="0 0.08 0" quat="0.707107 0.707107 0 0" type="cylinder" rgba="1 1 1 1"/>
|
||||
|
||||
<body name="FL_thigh" pos="0 0.0955 0">
|
||||
<inertial pos="-0.00374 -0.0223 -0.0327" quat="0.829533 0.0847635 -0.0200632 0.551623" mass="1.152" diaginertia="0.00594973 0.00584149 0.000878787"/>
|
||||
<joint name="FL_thigh_joint" pos="0 0 0" axis="0 1 0" range="-1.5708 3.4907" actuatorfrcrange="-23.7 23.7"/>
|
||||
<geom mesh="thigh_0" material="metal" class="visual" />
|
||||
<geom mesh="thigh_1" material="gray" class="visual" />
|
||||
<geom size="0.055 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0" type="box" rgba="1 1 1 1"/>
|
||||
<body name="FL_calf" pos="0 0 -0.213">
|
||||
<inertial pos="0.00629595 -0.000622121 -0.141417" quat="0.710672 0.00154099 -0.00450087 0.703508" mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05"/>
|
||||
<joint name="FL_calf_joint" pos="0 0 0" axis="0 1 0" range="-2.7227 -0.83776" actuatorfrcrange="-35.55 35.55"/>
|
||||
<geom mesh="calf_0" material="gray" class="visual" />
|
||||
<geom mesh="calf_1" material="black" class="visual" />
|
||||
<geom size="0.012 0.06" pos="0.008 0 -0.06" quat="0.994493 0 -0.104807 0" type="cylinder" rgba="1 1 1 1"/>
|
||||
<geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder"/>
|
||||
<geom size="0.0155 0.015" pos="0.00801333 0 -0.18745" quat="0.965093 0 0.261909 0" type="cylinder"/>
|
||||
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
|
||||
<geom size="0.022" pos="-0.002 0 -0.213" rgba="1 1 1 1"/>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
<body name="FR_hip" pos="0.1934 -0.0465 0">
|
||||
<inertial pos="-0.0054 -0.00194 -0.000105" quat="0.498237 0.505462 0.499245 0.497014" mass="0.678" diaginertia="0.00088403 0.000596003 0.000479967"/>
|
||||
<joint name="FR_hip_joint" pos="0 0 0" axis="1 0 0" range="-1.0472 1.0472" actuatorfrcrange="-23.7 23.7"/>
|
||||
<geom mesh="hip_0" material="metal" class="visual" quat="4.63268e-05 1 0 0" />
|
||||
<geom mesh="hip_1" material="gray" class="visual" quat="4.63268e-05 1 0 0" />
|
||||
<geom size="0.046 0.02" pos="0 -0.08 0" quat="0.707107 0.707107 0 0" type="cylinder" rgba="1 1 1 1"/>
|
||||
<body name="FR_thigh" pos="0 -0.0955 0">
|
||||
<inertial pos="-0.00374 0.0223 -0.0327" quat="0.551623 -0.0200632 0.0847635 0.829533" mass="1.152" diaginertia="0.00594973 0.00584149 0.000878787"/>
|
||||
<joint name="FR_thigh_joint" pos="0 0 0" axis="0 1 0" range="-1.5708 3.4907" actuatorfrcrange="-23.7 23.7"/>
|
||||
<geom mesh="thigh_mirror_0" material="metal" class="visual" />
|
||||
<geom mesh="thigh_mirror_1" material="gray" class="visual" />
|
||||
<geom size="0.055 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0" type="box" rgba="1 1 1 1"/>
|
||||
<body name="FR_calf" pos="0 0 -0.213">
|
||||
<inertial pos="0.00629595 0.000622121 -0.141417" quat="0.703508 -0.00450087 0.00154099 0.710672" mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05"/>
|
||||
<joint name="FR_calf_joint" pos="0 0 0" axis="0 1 0" range="-2.7227 -0.83776" actuatorfrcrange="-35.55 35.55"/>
|
||||
<geom mesh="calf_mirror_0" material="gray" class="visual" />
|
||||
<geom mesh="calf_mirror_1" material="black" class="visual" />
|
||||
<geom size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0" type="cylinder" rgba="1 1 1 1"/>
|
||||
<geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder"/>
|
||||
<geom size="0.0155 0.015" pos="0.00801333 0 -0.18745" quat="0.965093 0 0.261909 0" type="cylinder"/>
|
||||
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
|
||||
<geom size="0.022" pos="-0.002 0 -0.213" rgba="1 1 1 1"/>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
<body name="RL_hip" pos="-0.1934 0.0465 0">
|
||||
<inertial pos="0.0054 0.00194 -0.000105" quat="0.505462 0.498237 0.497014 0.499245" mass="0.678" diaginertia="0.00088403 0.000596003 0.000479967"/>
|
||||
<joint name="RL_hip_joint" pos="0 0 0" axis="1 0 0" range="-1.0472 1.0472" actuatorfrcrange="-23.7 23.7"/>
|
||||
<geom mesh="hip_0" material="metal" class="visual" quat="4.63268e-05 0 1 0" />
|
||||
<geom mesh="hip_1" material="gray" class="visual" quat="4.63268e-05 0 1 0" />
|
||||
<geom size="0.046 0.02" pos="0 0.08 0" quat="0.707107 0.707107 0 0" type="cylinder" rgba="1 1 1 1"/>
|
||||
<body name="RL_thigh" pos="0 0.0955 0">
|
||||
<inertial pos="-0.00374 -0.0223 -0.0327" quat="0.829533 0.0847635 -0.0200632 0.551623" mass="1.152" diaginertia="0.00594973 0.00584149 0.000878787"/>
|
||||
<joint name="RL_thigh_joint" pos="0 0 0" axis="0 1 0" range="-0.5236 4.5379" actuatorfrcrange="-23.7 23.7"/>
|
||||
<geom mesh="thigh_0" material="metal" class="visual" />
|
||||
<geom mesh="thigh_1" material="gray" class="visual" />
|
||||
<geom size="0.055 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0" type="box" rgba="1 1 1 1"/>
|
||||
<body name="RL_calf" pos="0 0 -0.213">
|
||||
<inertial pos="0.00629595 -0.000622121 -0.141417" quat="0.710672 0.00154099 -0.00450087 0.703508" mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05"/>
|
||||
<joint name="RL_calf_joint" pos="0 0 0" axis="0 1 0" range="-2.7227 -0.83776" actuatorfrcrange="-35.55 35.55"/>
|
||||
<geom mesh="calf_0" material="gray" class="visual" />
|
||||
<geom mesh="calf_1" material="black" class="visual" />
|
||||
<geom size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0" type="cylinder" rgba="1 1 1 1"/>
|
||||
<geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder"/>
|
||||
<geom size="0.0155 0.015" pos="0.00801333 0 -0.18745" quat="0.965093 0 0.261909 0" type="cylinder"/>
|
||||
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
|
||||
<geom size="0.022" pos="-0.002 0 -0.213" rgba="1 1 1 1"/>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
|
||||
<body name="RR_hip" pos="-0.1934 -0.0465 0">
|
||||
<inertial pos="0.0054 -0.00194 -0.000105" quat="0.499245 0.497014 0.498237 0.505462" mass="0.678" diaginertia="0.00088403 0.000596003 0.000479967"/>
|
||||
<joint name="RR_hip_joint" pos="0 0 0" axis="1 0 0" range="-1.0472 1.0472" actuatorfrcrange="-23.7 23.7"/>
|
||||
<geom mesh="hip_0" material="metal" class="visual"
|
||||
quat="2.14617e-09 4.63268e-05 4.63268e-05 -1" />
|
||||
<geom mesh="hip_1" material="gray" class="visual"
|
||||
quat="2.14617e-09 4.63268e-05 4.63268e-05 -1" />
|
||||
<geom size="0.046 0.02" pos="0 -0.08 0" quat="0.707107 0.707107 0 0" type="cylinder" rgba="1 1 1 1"/>
|
||||
<body name="RR_thigh" pos="0 -0.0955 0">
|
||||
<inertial pos="-0.00374 0.0223 -0.0327" quat="0.551623 -0.0200632 0.0847635 0.829533" mass="1.152" diaginertia="0.00594973 0.00584149 0.000878787"/>
|
||||
<joint name="RR_thigh_joint" pos="0 0 0" axis="0 1 0" range="-0.5236 4.5379" actuatorfrcrange="-23.7 23.7"/>
|
||||
<geom mesh="thigh_mirror_0" material="metal" class="visual" />
|
||||
<geom mesh="thigh_mirror_1" material="gray" class="visual" />
|
||||
<geom size="0.055 0.01225 0.017" pos="0 0 -0.1065" quat="0.707107 0 0.707107 0" type="box" rgba="1 1 1 1"/>
|
||||
<body name="RR_calf" pos="0 0 -0.213">
|
||||
<inertial pos="0.00629595 0.000622121 -0.141417" quat="0.703508 -0.00450087 0.00154099 0.710672" mass="0.241352" diaginertia="0.0014901 0.00146356 5.31397e-05"/>
|
||||
<joint name="RR_calf_joint" pos="0 0 0" axis="0 1 0" range="-2.7227 -0.83776" actuatorfrcrange="-35.55 35.55"/>
|
||||
<geom mesh="calf_mirror_0" material="gray" class="visual" />
|
||||
<geom mesh="calf_mirror_1" material="black" class="visual" />
|
||||
<geom size="0.013 0.06" pos="0.01 0 -0.06" quat="0.995004 0 -0.0998334 0" type="cylinder" rgba="1 1 1 1"/>
|
||||
<geom size="0.011 0.0325" pos="0.02 0 -0.148" quat="0.999688 0 0.0249974 0" type="cylinder"/>
|
||||
<geom size="0.0155 0.015" pos="0.00801333 0 -0.18745" quat="0.965093 0 0.261909 0" type="cylinder"/>
|
||||
<geom pos="0 0 -0.213" mesh="foot" class="visual" material="black" />
|
||||
<geom size="0.022" pos="-0.002 0 -0.213" rgba="1 1 1 1"/>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
|
||||
</body>
|
||||
</worldbody>
|
||||
|
||||
<actuator>
|
||||
<motor name="FL_hip" joint="FL_hip_joint"/>
|
||||
<motor name="FL_thigh" joint="FL_thigh_joint"/>
|
||||
<motor name="FL_calf" joint="FL_calf_joint"/>
|
||||
<motor name="FR_hip" joint="FR_hip_joint"/>
|
||||
<motor name="FR_thigh" joint="FR_thigh_joint"/>
|
||||
<motor name="FR_calf" joint="FR_calf_joint"/>
|
||||
<motor name="RL_hip" joint="RL_hip_joint"/>
|
||||
<motor name="RL_thigh" joint="RL_thigh_joint"/>
|
||||
<motor name="RL_calf" joint="RL_calf_joint"/>
|
||||
<motor name="RR_hip" joint="RR_hip_joint"/>
|
||||
<motor name="RR_thigh" joint="RR_thigh_joint"/>
|
||||
<motor name="RR_calf" joint="RR_calf_joint"/>
|
||||
</actuator>
|
||||
|
||||
</mujoco>
|
||||
22
resources/robots/go2/race_track.xml
Normal file
22
resources/robots/go2/race_track.xml
Normal file
File diff suppressed because one or more lines are too long
33
resources/robots/go2/stairs.xml
Normal file
33
resources/robots/go2/stairs.xml
Normal file
File diff suppressed because one or more lines are too long
807
resources/robots/go2/terrain_generator.py
Normal file
807
resources/robots/go2/terrain_generator.py
Normal file
@@ -0,0 +1,807 @@
|
||||
import xml.etree.ElementTree as xml_et
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
import cv2
|
||||
import noise
|
||||
import os
|
||||
|
||||
ROBOT = "go2"
|
||||
INPUT_SCENE_PATH = os.path.join(os.path.dirname(__file__), "scene.xml")
|
||||
OUTPUT_SCENE_PATH = os.path.join(os.path.dirname(__file__), "scene_terrain.xml")
|
||||
PATH_DIR = Path(__file__).parent.absolute()
|
||||
# zyx euler angle to quaternion
|
||||
def euler_to_quat(roll, pitch, yaw):
|
||||
cx = np.cos(roll / 2)
|
||||
sx = np.sin(roll / 2)
|
||||
cy = np.cos(pitch / 2)
|
||||
sy = np.sin(pitch / 2)
|
||||
cz = np.cos(yaw / 2)
|
||||
sz = np.sin(yaw / 2)
|
||||
|
||||
return np.array(
|
||||
[
|
||||
cx * cy * cz + sx * sy * sz,
|
||||
sx * cy * cz - cx * sy * sz,
|
||||
cx * sy * cz + sx * cy * sz,
|
||||
cx * cy * sz - sx * sy * cz,
|
||||
],
|
||||
dtype=np.float64,
|
||||
)
|
||||
|
||||
|
||||
# zyx euler angle to rotation matrix
|
||||
def euler_to_rot(roll, pitch, yaw):
|
||||
rot_x = np.array(
|
||||
[
|
||||
[1, 0, 0],
|
||||
[0, np.cos(roll), -np.sin(roll)],
|
||||
[0, np.sin(roll), np.cos(roll)],
|
||||
],
|
||||
dtype=np.float64,
|
||||
)
|
||||
|
||||
rot_y = np.array(
|
||||
[
|
||||
[np.cos(pitch), 0, np.sin(pitch)],
|
||||
[0, 1, 0],
|
||||
[-np.sin(pitch), 0, np.cos(pitch)],
|
||||
],
|
||||
dtype=np.float64,
|
||||
)
|
||||
rot_z = np.array(
|
||||
[
|
||||
[np.cos(yaw), -np.sin(yaw), 0],
|
||||
[np.sin(yaw), np.cos(yaw), 0],
|
||||
[0, 0, 1],
|
||||
],
|
||||
dtype=np.float64,
|
||||
)
|
||||
return rot_z @ rot_y @ rot_x
|
||||
|
||||
|
||||
# 2d rotate
|
||||
def rot2d(x, y, yaw):
|
||||
nx = x * np.cos(yaw) - y * np.sin(yaw)
|
||||
ny = x * np.sin(yaw) + y * np.cos(yaw)
|
||||
return nx, ny
|
||||
|
||||
|
||||
# 3d rotate
|
||||
def rot3d(pos, euler):
|
||||
R = euler_to_rot(euler[0], euler[1], euler[2])
|
||||
return R @ pos
|
||||
|
||||
|
||||
def list_to_str(vec):
|
||||
return " ".join(str(s) for s in vec)
|
||||
|
||||
|
||||
class TerrainGenerator:
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.scene = xml_et.parse(INPUT_SCENE_PATH)
|
||||
self.root = self.scene.getroot()
|
||||
self.worldbody = self.root.find("worldbody")
|
||||
self.asset = self.root.find("asset")
|
||||
self._add_wood_material()
|
||||
self._add_sponge_material()
|
||||
|
||||
def _add_wood_material(self):
|
||||
# 1. 添加纹理 (Texture)
|
||||
# 这里使用内置的 "flat" 类型配上棕色，模拟一种简单的木头颜色。
|
||||
# 如果你有真实的木纹图片(比如 wood.png)，请将 type="2d" builtin="flat"
|
||||
# 改为 type="2d" file="../wood.png"
|
||||
tex = xml_et.SubElement(self.asset, "texture")
|
||||
tex.attrib["name"] = "wood_tex"
|
||||
tex.attrib["type"] = "2d"
|
||||
tex.attrib["file"] = "./assets/wood.png"
|
||||
tex.attrib["rgb1"] = "0.6 0.4 0.2" # 棕色 (RGB)
|
||||
tex.attrib["width"] = "512"
|
||||
tex.attrib["height"] = "512"
|
||||
|
||||
# 2. 添加材质 (Material)
|
||||
mat = xml_et.SubElement(self.asset, "material")
|
||||
mat.attrib["name"] = "wood_mat" # 材质名称，后面 AddBox 要用
|
||||
mat.attrib["texture"] = "wood_tex" # 关联上面的纹理
|
||||
mat.attrib["specular"] = "0.2" # 木头反光度较低
|
||||
mat.attrib["shininess"] = "0.1" # 亮度较低
|
||||
mat.attrib["rgba"] = "1 1 1 1"
|
||||
|
||||
def _add_sponge_material(self):
|
||||
# 1. 添加纹理 (Texture)
|
||||
# 因为没有图片文件，我们使用 builtin="flat" 来生成纯色纹理
|
||||
tex = xml_et.SubElement(self.asset, "texture")
|
||||
tex.attrib["name"] = "sponge_tex"
|
||||
tex.attrib["type"] = "2d"
|
||||
tex.attrib["builtin"] = "flat" # 使用内置平面纹理，不需要 file 路径
|
||||
tex.attrib["rgb1"] = "1.0 0.7 0.7" # 设置颜色：粉色 (参考图片颜色)
|
||||
tex.attrib["width"] = "512"
|
||||
tex.attrib["height"] = "512"
|
||||
|
||||
# 2. 添加材质 (Material)
|
||||
mat = xml_et.SubElement(self.asset, "material")
|
||||
mat.attrib["name"] = "mat_sponge" # 材质名称，AddBox 中调用这个名字
|
||||
mat.attrib["texture"] = "sponge_tex"
|
||||
|
||||
# 海绵的关键视觉特性：不反光、不油亮
|
||||
mat.attrib["specular"] = "0.1" # 几乎没有镜面反射 (相比木头的0.2要低很多)
|
||||
mat.attrib["shininess"] = "0.1" # 几乎没有光泽
|
||||
mat.attrib["rgba"] = "1 0.7 0.7 1" # 叠加颜色，保持原样
|
||||
|
||||
# Add Box to scene
|
||||
def AddBox(self,
|
||||
position=[1.0, 0.0, 0.0],
|
||||
euler=[0.0, 0.0, 0.0],
|
||||
size=[0.1, 0.1, 0.1],
|
||||
sponge=False):
|
||||
geo = xml_et.SubElement(self.worldbody, "geom")
|
||||
geo.attrib["pos"] = list_to_str(position)
|
||||
geo.attrib["type"] = "box"
|
||||
geo.attrib["size"] = list_to_str(
|
||||
0.5 * np.array(size)) # half size of box for mujoco
|
||||
quat = euler_to_quat(euler[0], euler[1], euler[2])
|
||||
geo.attrib["quat"] = list_to_str(quat)
|
||||
# === 修改部分开始 ===
|
||||
if sponge:
|
||||
# 1. 视觉：使用海绵材质 (假设你在 asset 中定义的名字叫 mat_sponge)
|
||||
geo.attrib["material"] = "mat_sponge"
|
||||
# 2. 物理：solref 时间常数越大越软 (0.02 比较软, 默认约 0.002)
|
||||
geo.attrib["solref"] = "0.03 1"
|
||||
geo.attrib["priority"] = "1"
|
||||
geo.attrib["solmix"] = "1"
|
||||
# 3. 摩擦：海绵通常摩擦力较大 (可选)
|
||||
geo.attrib["friction"] = "1.2 0.005 0.0001"
|
||||
else:
|
||||
geo.attrib["material"] = "wood_mat"
|
||||
geo.attrib["friction"] = "0.5 0.005 0.0001"
|
||||
|
||||
def AddGeometry(self,
|
||||
position=[1.0, 0.0, 0.0],
|
||||
euler=[0.0, 0.0, 0.0],
|
||||
size=[0.1, 0.1],geo_type="box"):
|
||||
|
||||
# geo_type supports "plane", "sphere", "capsule", "ellipsoid", "cylinder", "box"
|
||||
geo = xml_et.SubElement(self.worldbody, "geom")
|
||||
geo.attrib["pos"] = list_to_str(position)
|
||||
geo.attrib["type"] = geo_type
|
||||
geo.attrib["size"] = list_to_str(
|
||||
0.5 * np.array(size)) # half size of box for mujoco
|
||||
quat = euler_to_quat(euler[0], euler[1], euler[2])
|
||||
geo.attrib["quat"] = list_to_str(quat)
|
||||
geo.attrib["material"] = "wood_mat"
|
||||
|
||||
def AddStairs(self,
|
||||
init_pos=[1.0, 0.0, 0.0],
|
||||
yaw=0.0,
|
||||
width=0.2,
|
||||
height=0.15,
|
||||
length=1.5,
|
||||
stair_nums=10):
|
||||
|
||||
local_pos = [0.0, 0.0, -0.5 * height]
|
||||
for i in range(stair_nums):
|
||||
local_pos[0] += width
|
||||
local_pos[2] += height
|
||||
x, y = rot2d(local_pos[0], local_pos[1], yaw)
|
||||
self.AddBox([x + init_pos[0], y + init_pos[1], local_pos[2]],
|
||||
[0.0, 0.0, yaw], [width, length, height])
|
||||
|
||||
|
||||
def AddDownStairs(self,
|
||||
init_pos=[1.0, 0.0, 0.0],
|
||||
yaw=0.0,
|
||||
width=0.3,
|
||||
height=0.15,
|
||||
length=1.5,
|
||||
stair_nums=10):
|
||||
|
||||
# 从上向下生成台阶：第一个台阶中心在 +0.5*height，随后每步降低 height
|
||||
local_pos = [0.0, 0.0, 0.5 * height]
|
||||
for i in range(stair_nums):
|
||||
local_pos[0] += width
|
||||
local_pos[2] -= height
|
||||
x, y = rot2d(local_pos[0], local_pos[1], yaw)
|
||||
self.AddBox([x + init_pos[0], y + init_pos[1], local_pos[2]],
|
||||
[0.0, 0.0, yaw], [width, length, height])
|
||||
|
||||
|
||||
def AddStairsSeries(self,
|
||||
init_pos=[1.0, 0.0, 0.0],
|
||||
yaw=0.0,
|
||||
width=0.3,
|
||||
length=1.5,
|
||||
stair_nums_up=6,
|
||||
stair_nums_down=6,
|
||||
start_height=0.05,
|
||||
step_inc=0.03,
|
||||
max_height=0.2,
|
||||
flat_length=0.5):
|
||||
"""生成一系列台阶：先上台阶（stair_nums_up），顶部有一段平地（flat_length），再下台阶（stair_nums_down）。
|
||||
每完成一对上/下台阶后，单步高度增加 step_inc，直到达到 max_height。
|
||||
|
||||
参数说明：
|
||||
- init_pos: 底层起点（列表），序列沿局部 x 方向展开
|
||||
- yaw: 台阶朝向
|
||||
- width: 每级台阶在 x 方向的深度（步幅）
|
||||
- length: 台阶在 y 方向的宽度（和 AddStairs 一致）
|
||||
- stair_nums_up/down: 上/下台阶的级数
|
||||
- start_height: 第一对台阶的每级高度
|
||||
- step_inc: 每对增加的高度
|
||||
- max_height: 最大每级高度（包含）
|
||||
- flat_length: 顶部平地长度（沿 x）
|
||||
"""
|
||||
|
||||
# 保持 init_pos 不变（上/下台阶成对结束后回到同一基准高度）
|
||||
base_pos = np.array(init_pos, dtype=float)
|
||||
height = start_height
|
||||
|
||||
# 平台长度至少为 1.0 米
|
||||
platform_length = max(flat_length, 1.0)
|
||||
|
||||
# 迭代每一对台阶直到高度超限
|
||||
while height <= max_height + 1e-8:
|
||||
# --- 上台阶 ---
|
||||
# 上台阶第 i 级中心 z = base_z + (-0.5 + i) * height, i = 1..stair_nums_up
|
||||
for i in range(1, stair_nums_up + 1):
|
||||
center_x_local = i * width
|
||||
center_z = base_pos[2] + (-0.5 + i) * height
|
||||
x, y = rot2d(center_x_local, 0.0, yaw)
|
||||
self.AddBox([x + base_pos[0], y + base_pos[1], center_z],
|
||||
[0.0, 0.0, yaw], [width, length, height])
|
||||
|
||||
# 顶部平地：放在最后一级之后，平台长度至少 platform_length
|
||||
last_up_center_x = stair_nums_up * width
|
||||
last_up_top_surface = base_pos[2] + stair_nums_up * height # 顶部平面高度
|
||||
flat_thickness = height # 平地厚度，使用与台阶同高度以保证接触
|
||||
flat_center_local_x = last_up_center_x + width / 2.0 + platform_length / 2.0
|
||||
flat_center_z = last_up_top_surface + flat_thickness / 2.0
|
||||
x, y = rot2d(flat_center_local_x, 0.0, yaw)
|
||||
# 尺寸：在 x 方向用 platform_length, y 用 length, z 用 flat_thickness
|
||||
self.AddBox([x + base_pos[0], y + base_pos[1], flat_center_z],
|
||||
[0.0, 0.0, yaw], [platform_length, length, flat_thickness])
|
||||
|
||||
# --- 下台阶 ---
|
||||
# 平地末端 x
|
||||
flat_end_x = last_up_center_x + width / 2.0 + platform_length
|
||||
for j in range(1, stair_nums_down + 1):
|
||||
center_x_local = flat_end_x + width / 2.0 + (j - 1) * width
|
||||
# 第 j 级下台阶的中心 z = last_up_top_surface - 0.5*height - (j-1)*height
|
||||
center_z = last_up_top_surface - 0.5 * height - (j - 1) * height
|
||||
x, y = rot2d(center_x_local, 0.0, yaw)
|
||||
self.AddBox([x + base_pos[0], y + base_pos[1], center_z],
|
||||
[0.0, 0.0, yaw], [width, length, height])
|
||||
|
||||
# 下台阶之后也添加一段平地（连接到下一组上台阶），长度至少 platform_length
|
||||
seq_end_local_x = flat_end_x + stair_nums_down * width # 这是最后一个下台阶的前缘 x
|
||||
post_flat_center_local_x = seq_end_local_x + platform_length / 2.0
|
||||
# 该平地应与下一组上台阶的起始高度对齐：其顶面与 base_z + height 对齐
|
||||
post_flat_center_z = base_pos[2] + height / 2.0
|
||||
x, y = rot2d(post_flat_center_local_x, 0.0, yaw)
|
||||
self.AddBox([x + base_pos[0], y + base_pos[1], post_flat_center_z],
|
||||
[0.0, 0.0, yaw], [platform_length, length, flat_thickness])
|
||||
|
||||
# 为下一对台阶准备：把 base_pos 在 x 方向平移到当前序列末端（post flat 末端），保持 z 不变
|
||||
seq_total_end_local_x = seq_end_local_x + platform_length
|
||||
# 计算下一组基准位移，使下一组上台阶第一级的前缘与当前 post-flat 的末端无缝对接
|
||||
# base_shift_local_x 为相对于当前 base 的局部 x 偏移
|
||||
overlap = 1e-3 # 以米为单位，微小重叠以避免可视缝隙
|
||||
base_shift_local_x = seq_total_end_local_x - width / 2.0 - overlap
|
||||
dx, dy = rot2d(base_shift_local_x, 0.0, yaw)
|
||||
base_pos[0] = base_pos[0] + dx
|
||||
base_pos[1] = base_pos[1] + dy
|
||||
|
||||
# 增加单级高度
|
||||
height = round(height + step_inc, 8)
|
||||
|
||||
|
||||
def AddSuspendStairs(self,
|
||||
init_pos=[1.0, 0.0, 0.0],
|
||||
yaw=1.0,
|
||||
width=0.2,
|
||||
height=0.15,
|
||||
length=1.5,
|
||||
gap=0.1,
|
||||
stair_nums=10):
|
||||
|
||||
local_pos = [0.0, 0.0, -0.5 * height]
|
||||
for i in range(stair_nums):
|
||||
local_pos[0] += width
|
||||
local_pos[2] += height
|
||||
x, y = rot2d(local_pos[0], local_pos[1], yaw)
|
||||
self.AddBox([x + init_pos[0], y + init_pos[1], local_pos[2]],
|
||||
[0.0, 0.0, yaw],
|
||||
[width, length, abs(height - gap)])
|
||||
|
||||
def AddRoughGround(self,
|
||||
init_pos=[1.0, 0.0, 0.0],
|
||||
euler=[0.0, -0.0, 0.0],
|
||||
nums=[10, 10],
|
||||
box_size=[0.5, 0.5, 0.5],
|
||||
box_euler=[0.0, 0.0, 0.0],
|
||||
separation=[0.2, 0.2],
|
||||
box_size_rand=[0.05, 0.05, 0.05],
|
||||
box_euler_rand=[0.2, 0.2, 0.2],
|
||||
separation_rand=[0.05, 0.05]):
|
||||
|
||||
local_pos = [0.0, 0.0, -0.5 * box_size[2]]
|
||||
new_separation = np.array(separation) + np.array(
|
||||
separation_rand) * np.random.uniform(-1.0, 1.0, 2)
|
||||
for i in range(nums[0]):
|
||||
local_pos[0] += new_separation[0]
|
||||
local_pos[1] = 0.0
|
||||
for j in range(nums[1]):
|
||||
new_box_size = np.array(box_size) + np.array(
|
||||
box_size_rand) * np.random.uniform(-1.0, 1.0, 3)
|
||||
new_box_euler = np.array(box_euler) + np.array(
|
||||
box_euler_rand) * np.random.uniform(-1.0, 1.0, 3)
|
||||
new_separation = np.array(separation) + np.array(
|
||||
separation_rand) * np.random.uniform(-1.0, 1.0, 2)
|
||||
|
||||
local_pos[1] += new_separation[1]
|
||||
pos = rot3d(local_pos, euler) + np.array(init_pos)
|
||||
self.AddBox(pos, new_box_euler, new_box_size)
|
||||
|
||||
def AddPerlinHeighField(
|
||||
self,
|
||||
position=[1.0, 0.0, 0.0], # position
|
||||
euler=[0.0, -0.0, 0.0], # attitude
|
||||
size=[1.0, 1.0], # width and length
|
||||
height_scale=0.2, # max height
|
||||
negative_height=0.2, # height in the negative direction of z axis
|
||||
image_width=128, # height field image size
|
||||
img_height=128,
|
||||
smooth=100.0, # smooth scale
|
||||
perlin_octaves=6, # perlin noise parameter
|
||||
perlin_persistence=0.5,
|
||||
perlin_lacunarity=2.0,
|
||||
output_hfield_image="height_field.png"):
|
||||
|
||||
# Generating height field based on perlin noise
|
||||
terrain_image = np.zeros((img_height, image_width), dtype=np.uint8)
|
||||
for y in range(image_width):
|
||||
for x in range(image_width):
|
||||
# Perlin noise
|
||||
noise_value = noise.pnoise2(x / smooth,
|
||||
y / smooth,
|
||||
octaves=perlin_octaves,
|
||||
persistence=perlin_persistence,
|
||||
lacunarity=perlin_lacunarity)
|
||||
terrain_image[y, x] = int((noise_value + 1) / 2 * 255)
|
||||
|
||||
cv2.imwrite(str(PATH_DIR / "assets" / output_hfield_image), terrain_image)
|
||||
|
||||
hfield = xml_et.SubElement(self.asset, "hfield")
|
||||
hfield.attrib["name"] = "perlin_hfield"
|
||||
hfield.attrib["size"] = list_to_str(
|
||||
[size[0] / 2.0, size[1] / 2.0, height_scale, negative_height])
|
||||
hfield.attrib["file"] = "../" + output_hfield_image
|
||||
|
||||
geo = xml_et.SubElement(self.worldbody, "geom")
|
||||
geo.attrib["type"] = "hfield"
|
||||
geo.attrib["hfield"] = "perlin_hfield"
|
||||
geo.attrib["pos"] = list_to_str(position)
|
||||
quat = euler_to_quat(euler[0], euler[1], euler[2])
|
||||
geo.attrib["quat"] = list_to_str(quat)
|
||||
|
||||
def AddHeighFieldFromImage(
|
||||
self,
|
||||
position=[1.0, 0.0, 0.0], # position
|
||||
euler=[0.0, -0.0, 0.0], # attitude
|
||||
size=[2.0, 1.6], # width and length
|
||||
height_scale=0.02, # max height
|
||||
negative_height=0.1, # height in the negative direction of z axis
|
||||
input_img=None,
|
||||
output_hfield_image="height_field.png",
|
||||
image_scale=[1.0, 1.0], # reduce image resolution
|
||||
invert_gray=False):
|
||||
|
||||
input_image = cv2.imread(input_img) # 替换为你的图像文件路径
|
||||
|
||||
width = int(input_image.shape[1] * image_scale[0])
|
||||
height = int(input_image.shape[0] * image_scale[1])
|
||||
resized_image = cv2.resize(input_image, (width, height),
|
||||
interpolation=cv2.INTER_AREA)
|
||||
terrain_image = cv2.cvtColor(resized_image, cv2.COLOR_BGR2GRAY)
|
||||
if invert_gray:
|
||||
terrain_image = 255 - position
|
||||
cv2.imwrite(str(PATH_DIR / "assets" / output_hfield_image), terrain_image)
|
||||
|
||||
hfield = xml_et.SubElement(self.asset, "hfield")
|
||||
hfield.attrib["name"] = "image_hfield"
|
||||
hfield.attrib["size"] = list_to_str(
|
||||
[size[0] / 2.0, size[1] / 2.0, height_scale, negative_height])
|
||||
hfield.attrib["file"] = "../" + output_hfield_image
|
||||
|
||||
geo = xml_et.SubElement(self.worldbody, "geom")
|
||||
geo.attrib["type"] = "hfield"
|
||||
geo.attrib["hfield"] = "image_hfield"
|
||||
geo.attrib["pos"] = list_to_str(position)
|
||||
quat = euler_to_quat(euler[0], euler[1], euler[2])
|
||||
geo.attrib["quat"] = list_to_str(quat)
|
||||
|
||||
def Save(self):
|
||||
self.scene.write(OUTPUT_SCENE_PATH)
|
||||
|
||||
def AddSlope(self,
|
||||
position=[1.0, 0.0, 0.0],
|
||||
yaw=0.0,
|
||||
length=0.575,
|
||||
width=0.55,
|
||||
height=0.15,
|
||||
thickness=0,
|
||||
add_baffle=True,
|
||||
add_baffle_height=False,
|
||||
sponge=False):
|
||||
"""
|
||||
生成带后背板和两侧挡板的斜坡。
|
||||
侧挡板上沿与斜坡面对齐。
|
||||
"""
|
||||
if not thickness:
|
||||
thickness = 0.05 # 斜坡面板厚度
|
||||
side_thickness = 0.05 # 侧板厚度
|
||||
|
||||
# ===========================
|
||||
# 1. 计算公共几何参数
|
||||
# ===========================
|
||||
|
||||
# 计算坡度角和斜边长
|
||||
angle = np.arctan2(height, length)
|
||||
ramp_len = np.sqrt(length**2 + height**2)
|
||||
|
||||
# 法线向量 (nx, nz)
|
||||
nx = np.sin(angle)
|
||||
nz = -np.cos(angle)
|
||||
|
||||
# 斜坡面板的几何中心 (局部坐标)
|
||||
# 这里的逻辑是将面板中心沿着法线向下偏移厚度的一半，保证上表面对齐理想斜面
|
||||
mid_x = length / 2.0
|
||||
mid_z = height / 2.0
|
||||
lx = mid_x + nx * (thickness / 2.0)
|
||||
lz = mid_z + nz * (thickness / 2.0)
|
||||
|
||||
# 统一的旋转角度 (Pitch: -angle, Yaw: yaw)
|
||||
final_euler = [0.0, -angle, yaw]
|
||||
|
||||
# ===========================
|
||||
# 2. 生成主斜坡面 (Ramp)
|
||||
# ===========================
|
||||
gx, gy = rot2d(lx, 0, yaw)
|
||||
ramp_pos = [position[0] + gx, position[1] + gy, position[2] + lz]
|
||||
self.AddBox(ramp_pos, final_euler, [ramp_len, width, thickness],sponge=sponge)
|
||||
|
||||
# ===========================
|
||||
# 3. 生成垂直背板 (Back Wall)
|
||||
# ===========================
|
||||
# 位于斜坡末端，高度为 height
|
||||
if add_baffle:
|
||||
back_lx = length + (thickness / 2.0)
|
||||
if add_baffle_height:
|
||||
back_lz = height
|
||||
height *= 2
|
||||
else:
|
||||
back_lz = height / 2.0
|
||||
bgx, bgy = rot2d(back_lx, 0, yaw)
|
||||
back_pos = [position[0] + bgx, position[1] + bgy, position[2] + back_lz]
|
||||
# 背板竖直放置，只受Yaw影响
|
||||
self.AddBox(back_pos, [0, 0, yaw], [thickness, width, height])
|
||||
|
||||
# 生成两侧挡板 (Side Walls)
|
||||
side_h = height / 2
|
||||
|
||||
# 计算侧板的中心 Z 坐标 (side_lz)
|
||||
# 目标：侧板的上表面 Z = 斜坡的上表面 Z
|
||||
# 斜坡上表面 Z (局部) = lz + thickness/2
|
||||
# 侧板上表面 Z (局部) = side_lz + side_h/2
|
||||
# 等式：lz + thickness/2 = side_lz + side_h/2
|
||||
# 解得：
|
||||
side_lz = lz + (thickness / 2.0) - (side_h / 2.0)
|
||||
if add_baffle_height:
|
||||
side_lz += side_h
|
||||
# side_h *= 2
|
||||
# lx += height * 2 / length / 2
|
||||
# 侧板的 X 坐标与斜坡中心一致 (lx)
|
||||
|
||||
# 侧板的 Y 偏移量
|
||||
# 放在斜坡宽度的两侧：(斜坡宽/2) + (侧板厚/2)
|
||||
y_shift = (width / 2.0) + (side_thickness / 2.0)
|
||||
|
||||
# 生成左右两个侧板
|
||||
for sign in [-1, 1]: # -1:左侧, 1:右侧
|
||||
if add_baffle_height and sign == 1: continue
|
||||
local_y = sign * y_shift
|
||||
|
||||
# 将 (lx, local_y) 旋转 Yaw 角到全局
|
||||
sgx, sgy = rot2d(lx, local_y, yaw)
|
||||
|
||||
side_pos = [
|
||||
position[0] + sgx,
|
||||
position[1] + sgy,
|
||||
position[2] + side_lz
|
||||
]
|
||||
|
||||
# 侧板的旋转角度与斜坡完全一致，这样上边缘才会平行
|
||||
self.AddBox(side_pos, final_euler, [ramp_len, side_thickness, side_h])
|
||||
|
||||
def AddSlopeGroup(self, position=[0.0, 0.0, 0.0], yaw=0.0, add_baffle=True, add_baffle_height=False,):
|
||||
|
||||
L = 0.6 # 坡长 (爬升方向)
|
||||
W = 0.6 # 坡宽 (侧向)
|
||||
H = 0.164
|
||||
|
||||
p1_local = [-L / 2, -W]
|
||||
yaw1 = np.pi / 2
|
||||
|
||||
p2_local = [-L, W/2]
|
||||
yaw2 = 0
|
||||
|
||||
p3_local = [L, -W/2]
|
||||
yaw3 = np.pi
|
||||
|
||||
p4_local = [L / 2, W]
|
||||
yaw4 = np.pi * 3 / 2
|
||||
|
||||
blocks = [
|
||||
(p1_local, yaw1),
|
||||
(p2_local, yaw2),
|
||||
(p3_local, yaw3),
|
||||
(p4_local, yaw4)
|
||||
]
|
||||
|
||||
for pos_local, local_yaw in blocks:
|
||||
off_x, off_y = rot2d(pos_local[0], pos_local[1], yaw)
|
||||
|
||||
abs_pos = [
|
||||
position[0] + off_x,
|
||||
position[1] + off_y,
|
||||
position[2]
|
||||
]
|
||||
|
||||
abs_yaw = yaw + local_yaw
|
||||
|
||||
self.AddSlope(position=abs_pos,
|
||||
yaw=abs_yaw,
|
||||
length=L,
|
||||
width=W,
|
||||
height=H,
|
||||
add_baffle=add_baffle,
|
||||
add_baffle_height=add_baffle_height)
|
||||
|
||||
def AddBlockyHeightField(
|
||||
self,
|
||||
position=[1.0, 0.0, 0.0],
|
||||
euler=[0.0, -0.0, 0.0],
|
||||
size=[1.0, 1.0],
|
||||
height_scale=0.2, # 高度差幅度
|
||||
negative_height=0.1,
|
||||
image_width=128,
|
||||
img_height=128,
|
||||
smooth=50.0, # 注意：如果要完全随机，把这个数改得很小（如 2.0）
|
||||
pixels_per_block=16, # <--- 新参数：决定方块的大小
|
||||
output_hfield_image="height_field.png"):
|
||||
|
||||
# 1. 准备图像数据
|
||||
terrain_image = np.zeros((img_height, image_width), dtype=np.uint8)
|
||||
|
||||
# 2. 预先生成一个随机种子偏移，保证每次地形不一样
|
||||
seed_offset_x = np.random.randint(0, 10000)
|
||||
seed_offset_y = np.random.randint(0, 10000)
|
||||
|
||||
for y in range(img_height):
|
||||
for x in range(image_width):
|
||||
|
||||
# === 核心修改开始 ===
|
||||
# 这里的整除逻辑 (//) 是制造“方块感”的关键
|
||||
# 它将坐标强制归整，例如 x=0到15 都会变成 0，x=16到31 都会变成 16
|
||||
# 这样这 16 个像素取到的噪声值就是一模一样的，形成平坦的台阶
|
||||
block_x = (x // pixels_per_block) * pixels_per_block
|
||||
block_y = (y // pixels_per_block) * pixels_per_block
|
||||
|
||||
# 使用归整后的 block_x, block_y 来生成噪声
|
||||
noise_value = noise.pnoise2((block_x + seed_offset_x) / smooth,
|
||||
(block_y + seed_offset_y) / smooth,
|
||||
octaves=1, # 减少细节，让方块表面平整
|
||||
persistence=0.5,
|
||||
lacunarity=2.0)
|
||||
# === 核心修改结束 ===
|
||||
|
||||
# 映射到 0-255
|
||||
terrain_image[y, x] = int((noise_value + 1) / 2 * 255)
|
||||
|
||||
hfield = xml_et.SubElement(self.asset, "hfield")
|
||||
hfield.attrib["name"] = "perlin_hfield"
|
||||
hfield.attrib["size"] = list_to_str(
|
||||
[size[0] / 2.0, size[1] / 2.0, height_scale, negative_height])
|
||||
hfield.attrib["file"] = output_hfield_image
|
||||
|
||||
geo = xml_et.SubElement(self.worldbody, "geom")
|
||||
geo.attrib["type"] = "hfield"
|
||||
geo.attrib["hfield"] = "perlin_hfield"
|
||||
geo.attrib["pos"] = list_to_str(position)
|
||||
quat = euler_to_quat(euler[0], euler[1], euler[2])
|
||||
geo.attrib["quat"] = list_to_str(quat)
|
||||
|
||||
def unit1_sponge(self):
|
||||
self.AddSlope(position=[0.0, 1.2, 0.0], yaw=np.pi, length=2.32, width=1.2, height=0.6, add_baffle_height=True)
|
||||
self.AddSlope(position=[-2.32, 0.0, 0.0], yaw=0, length=2.32, width=1.2, height=0.6, add_baffle_height=True)
|
||||
|
||||
self.AddBox(position=[0.6, -0.6, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
self.AddBox(position=[-2.92, -0.6, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
self.AddBox(position=[0.6, 1.8, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
self.AddBox(position=[-2.92, 1.8, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
self.AddBox(position=[-3.52, 0.0, 0.3], size=[0.6, 1.2, 0.03], euler=[0, np.pi / 2, 0])
|
||||
|
||||
self.AddBox(position=[-2.92, 0.6, 0.025], size=[1.2, 2.4, 0.05], sponge=True)
|
||||
self.AddBox(position=[0.6, 0.6, 0.025], size=[1.2, 2.4, 0.05], sponge=True)
|
||||
|
||||
self.AddSlope(position=[0.0, 1.2, 0.05], yaw=np.pi, length=2.32, width=1.2, thickness=0.05, height=0.6, sponge=True, add_baffle=False)
|
||||
self.AddSlope(position=[-2.32, 0.0, 0.05], yaw=0, length=2.32, width=1.2, thickness=0.05, height=0.6, sponge=True, add_baffle=False)
|
||||
|
||||
|
||||
|
||||
def unit2_slopes(self):
|
||||
self.AddBox(position=[1.225, 0.6, 0.0], size=[2.4, 0.05, 0.1], euler=[0, 0, np.pi / 2])
|
||||
self.AddSlopeGroup(position=[1.85, 0.0, 0.0], yaw=0.0, add_baffle=True)
|
||||
self.AddSlopeGroup(position=[1.85, 1.2, 0.0], yaw=0.0, add_baffle=True)
|
||||
self.AddBox(position=[2.475, 0.6, 0.0], size=[2.4, 0.05, 0.1], euler=[0, 0, np.pi / 2])
|
||||
self.AddSlopeGroup(position=[3.1, 0.0, 0.0], yaw=0.0, add_baffle=True)
|
||||
self.AddSlopeGroup(position=[3.1, 1.2, 0.0], yaw=0.0, add_baffle=True)
|
||||
self.AddSlopeGroup(position=[4.3, 0.0, 0.0], yaw=0.0, add_baffle=True)
|
||||
self.AddSlopeGroup(position=[4.3, 1.2, 0.0], yaw=0.0, add_baffle=True)
|
||||
self.AddBox(position=[3.7, 0.575, 0.0], size=[2.4, 0.05, 0.1])
|
||||
self.AddBox(position=[3.7, 0.625, 0.0], size=[2.4, 0.05, 0.1])
|
||||
self.AddBox(position=[4.925, 0.6, 0.0], size=[2.4, 0.05, 0.1], euler=[0, 0, np.pi / 2])
|
||||
self.AddSlopeGroup(position=[5.55, 0.0, 0.0], yaw=0.0, add_baffle=True)
|
||||
self.AddSlopeGroup(position=[5.55, 1.2, 0.0], yaw=0.0, add_baffle=True)
|
||||
self.AddBox(position=[6.175, 0.6, 0.0], size=[2.4, 0.05, 0.1], euler=[0, 0, np.pi / 2])
|
||||
self.AddBox(position=[1.9, 1.8, 0.3], size=[1.4, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
self.AddBox(position=[3.2, 1.8, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
self.AddBox(position=[4.4, 1.8, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
self.AddBox(position=[5.6, 1.8, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
|
||||
self.AddBox(position=[1.9, -0.6, 0.3], size=[1.4, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
self.AddBox(position=[3.2, -0.6, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
self.AddBox(position=[4.4, -0.6, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
self.AddBox(position=[5.6, -0.6, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
|
||||
self.AddBox(position=[1.2, 0.0, 0.3], size=[0.6, 1.2, 0.03], euler=[0, np.pi / 2, 0])
|
||||
self.AddBox(position=[2.5, 1.2, 0.3], size=[0.6, 1.2, 0.03], euler=[0, np.pi / 2, 0])
|
||||
self.AddBox(position=[4.95, 0.0, 0.3], size=[0.6, 1.2, 0.03], euler=[0, np.pi / 2, 0])
|
||||
self.AddBox(position=[6.2, 1.2, 0.3], size=[0.6, 1.2, 0.03], euler=[0, np.pi / 2, 0])
|
||||
|
||||
|
||||
def unit3_stairs(self):
|
||||
self.AddBox(position = [6.8, 1.2, 0.1], size = [1.2, 1.2, 0.2])
|
||||
self.AddBox(position = [8.0, 0.0, 0.1], size = [1.2, 1.2, 0.2])
|
||||
self.AddBox(position = [9.2, 1.2, 0.1], size = [1.2, 1.2, 0.2])
|
||||
self.AddBox(position = [10.4, 0.0, 0.1], size = [1.2, 1.2, 0.2])
|
||||
|
||||
self.AddBox(position=[6.8, 1.8, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
self.AddBox(position=[8.0, 1.8, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
self.AddBox(position=[9.2, 1.8, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
self.AddBox(position=[10.4, 1.8, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
|
||||
self.AddBox(position=[6.8, -0.6, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
self.AddBox(position=[8.0, -0.6, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
self.AddBox(position=[9.2, -0.6, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
self.AddBox(position=[10.4, -0.6, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
|
||||
self.AddBox(position=[11.0, 1.2, 0.3], size=[0.6, 1.2, 0.03], euler=[0, np.pi / 2, 0])
|
||||
|
||||
def unit4_diagonal(self):
|
||||
self.AddSlope(position=[12.8, 1.8, 0.0], yaw=np.pi / 2, length=2.32, width=1.2, height=0.6, add_baffle_height=True)
|
||||
self.AddSlope(position=[12.8, 1.8, 0.05], yaw=np.pi / 2, length=2.32/48, width=1.2, height=0.6/48, thickness=0.05, add_baffle=False)
|
||||
self.AddSlope(position=[12.8, 4.07, 0.6375], yaw=np.pi / 2, length=2.32/48, width=1.2, height=0.6/48, thickness=0.05, add_baffle=False)
|
||||
self.AddSlope(position=[12.225, 1.8, 0.05], yaw=np.pi / 2, length=2.32, width=0.05, height=0.6, thickness=0.05, add_baffle=False)
|
||||
self.AddSlope(position=[13.375, 1.8, 0.05], yaw=np.pi / 2, length=2.32, width=0.05, height=0.6, thickness=0.05, add_baffle=False)
|
||||
self.AddBox(position=[12.8, 3.54, 0.475], size=[1.64, 0.05, 0.05], euler=[0.0, -0.1779, 0.7685])
|
||||
self.AddBox(position=[12.8, 2.38, 0.175], size=[1.64, 0.05, 0.05], euler=[0.0, 0.1779, -0.7685])
|
||||
|
||||
self.AddSlope(position=[11.6, 4.12, 0.0], yaw=-np.pi / 2, length=2.32, width=1.2, height=0.6, add_baffle_height=True)
|
||||
self.AddSlope(position=[11.6, 4.12, 0.05], yaw=-np.pi / 2, length=2.32/48, width=1.2, height=0.6/48, thickness=0.05, add_baffle=False)
|
||||
self.AddSlope(position=[11.6, 1.85, 0.6375], yaw=-np.pi / 2, length=2.32/48, width=1.2, height=0.6/48, thickness=0.05, add_baffle=False)
|
||||
self.AddSlope(position=[12.175, 4.12, 0.05], yaw=-np.pi / 2, length=2.32, width=0.05, height=0.6, thickness=0.05, add_baffle=False)
|
||||
self.AddSlope(position=[11.025, 4.12, 0.05], yaw=-np.pi / 2, length=2.32, width=0.05, height=0.6, thickness=0.05, add_baffle=False)
|
||||
self.AddBox(position=[11.6, 3.54, 0.175], size=[1.64, 0.05, 0.05], euler=[0.0, -0.1779, -0.7685])
|
||||
self.AddBox(position=[11.6, 2.38, 0.475], size=[1.64, 0.05, 0.05], euler=[0.0, 0.1779, 0.7685])
|
||||
|
||||
self.AddBox(position=[11.6, -0.6, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
self.AddBox(position=[12.2, 0.0, 0.3], size=[0.6, 1.2, 0.03], euler=[0, np.pi / 2, 0])
|
||||
self.AddBox(position=[12.8, 0.6, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
self.AddBox(position=[13.4, 1.2, 0.3], size=[0.6, 1.2, 0.03], euler=[0, np.pi / 2, 0])
|
||||
|
||||
self.AddBox(position=[12.2, 0.625, 0.025], size=[2.4, 0.05, 0.05])
|
||||
self.AddBox(position=[12.2, 1.775, 0.025], size=[2.4, 0.05, 0.05])
|
||||
self.AddBox(position=[11.025, 1.2, 0.025], size=[1.2, 0.05, 0.05], euler=[0, 0, np.pi / 2])
|
||||
self.AddBox(position=[13.375, 1.2, 0.025], size=[1.2, 0.05, 0.05], euler=[0, 0, np.pi / 2])
|
||||
self.AddBox(position=[11.6, 1.2, 0.025], size=[1.64, 0.05, 0.05], euler=[0, 0, -np.pi / 4])
|
||||
self.AddBox(position=[12.8, 1.2, 0.025], size=[1.64, 0.05, 0.05], euler=[0, 0, np.pi / 4])
|
||||
|
||||
self.AddBox(position=[12.2, 4.145, 0.025], size=[2.4, 0.05, 0.05])
|
||||
self.AddBox(position=[12.2, 5.295, 0.025], size=[2.4, 0.05, 0.05])
|
||||
self.AddBox(position=[11.025, 4.72, 0.025], size=[1.2, 0.05, 0.05], euler=[0, 0, np.pi / 2])
|
||||
self.AddBox(position=[13.375, 4.72, 0.025], size=[1.2, 0.05, 0.05], euler=[0, 0, np.pi / 2])
|
||||
self.AddBox(position=[11.6, 4.72, 0.025], size=[1.64, 0.05, 0.05], euler=[0, 0, np.pi / 4])
|
||||
self.AddBox(position=[12.8, 4.72, 0.025], size=[1.64, 0.05, 0.05], euler=[0, 0, -np.pi / 4])
|
||||
|
||||
self.AddBox(position=[11.0, 4.72, 0.3], size=[0.6, 1.2, 0.03], euler=[0, np.pi / 2, 0])
|
||||
self.AddBox(position=[13.4, 4.72, 0.3], size=[0.6, 1.2, 0.03], euler=[0, np.pi / 2, 0])
|
||||
self.AddBox(position=[11.6, 5.32, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
|
||||
self.AddBox(position=[11.6, 0.0, 0.0], size=[1.2, 1.2, 0.03])
|
||||
self.AddBox(position=[12.2, 1.2, 0.0], size=[2.4, 1.2, 0.03])
|
||||
self.AddBox(position=[12.2, 4.72, 0.0], size=[2.4, 1.2, 0.03])
|
||||
|
||||
|
||||
def unit5_sandstone(self):
|
||||
|
||||
self.AddBox(position=[11.0, 5.92, 0.3], size=[0.6, 1.2, 0.03], euler=[0, np.pi / 2, 0])
|
||||
self.AddBox(position=[11.0, 7.12, 0.3], size=[0.6, 1.2, 0.03], euler=[0, np.pi / 2, 0])
|
||||
self.AddBox(position=[11.0, 8.32, 0.3], size=[0.6, 1.2, 0.03], euler=[0, np.pi / 2, 0])
|
||||
self.AddBox(position=[11.0, 9.52, 0.3], size=[0.6, 1.2, 0.03], euler=[0, np.pi / 2, 0])
|
||||
self.AddBox(position=[13.4, 5.92, 0.3], size=[0.6, 1.2, 0.03], euler=[0, np.pi / 2, 0])
|
||||
self.AddBox(position=[13.4, 7.12, 0.3], size=[0.6, 1.2, 0.03], euler=[0, np.pi / 2, 0])
|
||||
self.AddBox(position=[13.4, 8.32, 0.3], size=[0.6, 1.2, 0.03], euler=[0, np.pi / 2, 0])
|
||||
self.AddBox(position=[13.4, 9.52, 0.3], size=[0.6, 1.2, 0.03], euler=[0, np.pi / 2, 0])
|
||||
|
||||
self.AddBox(position=[12.8, 6.52, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
self.AddBox(position=[12.8, 10.12, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
self.AddBox(position=[11.6, 8.92, 0.3], size=[1.2, 0.6, 0.03], euler=[np.pi / 2, 0, 0])
|
||||
|
||||
self.AddBox(position=[12.2, 6.545, 0.05], size=[2.4, 0.05, 0.05])
|
||||
self.AddBox(position=[12.2, 8.895, 0.05], size=[2.4, 0.05, 0.05])
|
||||
self.AddBox(position=[11.025, 7.72, 0.05], size=[2.4, 0.05, 0.05], euler=[0, 0, np.pi / 2])
|
||||
self.AddBox(position=[13.375, 7.72, 0.05], size=[2.4, 0.05, 0.05], euler=[0, 0, np.pi / 2])
|
||||
self.AddBox(position=[12.2, 7.72, 0.05], size=[3.28, 0.05, 0.05], euler=[0, 0, -np.pi / 4])
|
||||
self.AddBox(position=[12.2, 7.72, 0.05], size=[3.28, 0.05, 0.05], euler=[0, 0, np.pi / 4])
|
||||
self.AddBox(position=[11.6, 7.12, 0.05], size=[1.64, 0.05, 0.05], euler=[0, 0, -np.pi / 4])
|
||||
self.AddBox(position=[12.8, 7.12, 0.05], size=[1.64, 0.05, 0.05], euler=[0, 0, np.pi / 4])
|
||||
self.AddBox(position=[11.6, 8.32, 0.05], size=[1.64, 0.05, 0.05], euler=[0, 0, np.pi / 4])
|
||||
self.AddBox(position=[12.8, 8.32, 0.05], size=[1.64, 0.05, 0.05], euler=[0, 0, -np.pi / 4])
|
||||
|
||||
self.AddBlockyHeightField(position=[12.2, 7.72, -0.00], size=[2.4, 4.8], height_scale=0.08)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tg = TerrainGenerator()
|
||||
|
||||
# # Box obstacle
|
||||
# tg.AddBox(position=[0.55 / 2 + 2, 0, 0.075], size=[np.sqrt(0.55 * 0.55 + 0.15*0.15), 0.1, 0.01], euler=[np.pi / 2, 0, 0])
|
||||
# tg.AddBox(position=[3.5, 0, 0.075], size=[1.5, 0.5, 0.01], euler=[0, 0, 0])
|
||||
|
||||
# # Geometry obstacle
|
||||
# # geo_type supports "plane", "sphere", "capsule", "ellipsoid", "cylinder", "box"
|
||||
# tg.AddGeometry(position=[1.5, 0.0, 0.25], euler=[0, 0, 0.0], size=[1.0,0.5,0.5],geo_type="cylinder")
|
||||
|
||||
# # Slope
|
||||
# tg.AddBox(position=[2.0, 2.0, 0.5],
|
||||
# euler=[0.0, -0.5, 0.0],
|
||||
# size=[3, 1.5, 0.1])
|
||||
|
||||
# # Stairs
|
||||
# tg.AddStairs(init_pos=[1.0, 4.0, 0.0], yaw=0.0)
|
||||
|
||||
# # Suspend stairs
|
||||
# tg.AddSuspendStairs(init_pos=[1.0, 6.0, 0.0], yaw=0.0)
|
||||
|
||||
# # Rough ground
|
||||
# tg.AddRoughGround(init_pos=[-2.5, 5.0, 0.0],
|
||||
# euler=[0, 0, 0.0],
|
||||
# nums=[10, 8])
|
||||
|
||||
# # Perlin heigh field
|
||||
# tg.AddPerlinHeighField(position=[-1.5, 4.0, 0.0], size=[2.0, 1.5])
|
||||
|
||||
# # Heigh field from image
|
||||
# tg.AddHeighFieldFromImage(position=[-1.5, 2.0, 0.0],
|
||||
# euler=[0, 0, -1.57],
|
||||
# size=[2.0,2.0],
|
||||
# input_img="./unitree_robot.jpeg",
|
||||
# image_scale=[1.0, 1.0],
|
||||
# output_hfield_image="unitree_hfield.png")
|
||||
tg.unit1_sponge()
|
||||
tg.unit2_slopes()
|
||||
tg.unit3_stairs()
|
||||
tg.unit4_diagonal()
|
||||
tg.unit5_sandstone()
|
||||
|
||||
tg.Save()
|
||||
1217
resources/robots/go2/urdf/go2.urdf
Normal file
1217
resources/robots/go2/urdf/go2.urdf
Normal file
File diff suppressed because it is too large
Load Diff
12
rsl_rl/.gitignore
vendored
Normal file
12
rsl_rl/.gitignore
vendored
Normal file
@@ -0,0 +1,12 @@
|
||||
# IDEs
|
||||
.idea
|
||||
|
||||
# builds
|
||||
*.egg-info
|
||||
|
||||
# cache
|
||||
__pycache__
|
||||
.pytest_cache
|
||||
|
||||
# vs code
|
||||
.vscode
|
||||
30
rsl_rl/LICENSE
Normal file
30
rsl_rl/LICENSE
Normal file
@@ -0,0 +1,30 @@
|
||||
Copyright (c) 2021, ETH Zurich, Nikita Rudin
|
||||
Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES
|
||||
All rights reserved.
|
||||
|
||||
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.
|
||||
|
||||
See licenses/dependencies for license information of dependencies of this package.
|
||||
26
rsl_rl/README.md
Normal file
26
rsl_rl/README.md
Normal file
@@ -0,0 +1,26 @@
|
||||
# RSL RL
|
||||
Fast and simple implementation of RL algorithms, designed to run fully on GPU.
|
||||
This code is an evolution of `rl-pytorch` provided with NVIDIA's Isaac GYM.
|
||||
|
||||
Only PPO is implemented for now. More algorithms will be added later.
|
||||
Contributions are welcome.
|
||||
|
||||
## Setup
|
||||
|
||||
```
|
||||
git clone https://github.com/leggedrobotics/rsl_rl
|
||||
cd rsl_rl
|
||||
pip install -e .
|
||||
```
|
||||
|
||||
### Useful Links ###
|
||||
Example use case: https://github.com/leggedrobotics/legged_gym
|
||||
Project website: https://leggedrobotics.github.io/legged_gym/
|
||||
Paper: https://arxiv.org/abs/2109.11978
|
||||
|
||||
**Maintainer**: Nikita Rudin
|
||||
**Affiliation**: Robotic Systems Lab, ETH Zurich & NVIDIA
|
||||
**Contact**: rudinn@ethz.ch
|
||||
|
||||
|
||||
|
||||
30
rsl_rl/licenses/dependencies/numpy_license.txt
Normal file
30
rsl_rl/licenses/dependencies/numpy_license.txt
Normal file
@@ -0,0 +1,30 @@
|
||||
Copyright (c) 2005-2021, NumPy Developers.
|
||||
All rights reserved.
|
||||
|
||||
Redistribution and use in source and binary forms, with or without
|
||||
modification, are permitted provided that the following conditions are
|
||||
met:
|
||||
|
||||
* Redistributions of source code must retain the above copyright
|
||||
notice, this list of conditions and the following disclaimer.
|
||||
|
||||
* 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.
|
||||
|
||||
* Neither the name of the NumPy Developers nor the names of any
|
||||
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
|
||||
OWNER 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.
|
||||
73
rsl_rl/licenses/dependencies/torch_license.txt
Normal file
73
rsl_rl/licenses/dependencies/torch_license.txt
Normal file
@@ -0,0 +1,73 @@
|
||||
From PyTorch:
|
||||
|
||||
Copyright (c) 2016- Facebook, Inc (Adam Paszke)
|
||||
Copyright (c) 2014- Facebook, Inc (Soumith Chintala)
|
||||
Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
|
||||
Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu)
|
||||
Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
|
||||
Copyright (c) 2011-2013 NYU (Clement Farabet)
|
||||
Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
|
||||
Copyright (c) 2006 Idiap Research Institute (Samy Bengio)
|
||||
Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
|
||||
|
||||
From Caffe2:
|
||||
|
||||
Copyright (c) 2016-present, Facebook Inc. All rights reserved.
|
||||
|
||||
All contributions by Facebook:
|
||||
Copyright (c) 2016 Facebook Inc.
|
||||
|
||||
All contributions by Google:
|
||||
Copyright (c) 2015 Google Inc.
|
||||
All rights reserved.
|
||||
|
||||
All contributions by Yangqing Jia:
|
||||
Copyright (c) 2015 Yangqing Jia
|
||||
All rights reserved.
|
||||
|
||||
All contributions by Kakao Brain:
|
||||
Copyright 2019-2020 Kakao Brain
|
||||
|
||||
All contributions from Caffe:
|
||||
Copyright(c) 2013, 2014, 2015, the respective contributors
|
||||
All rights reserved.
|
||||
|
||||
All other contributions:
|
||||
Copyright(c) 2015, 2016 the respective contributors
|
||||
All rights reserved.
|
||||
|
||||
Caffe2 uses a copyright model similar to Caffe: each contributor holds
|
||||
copyright over their contributions to Caffe2. The project versioning records
|
||||
all such contribution and copyright details. If a contributor wants to further
|
||||
mark their specific copyright on a particular contribution, they should
|
||||
indicate their copyright solely in the commit message of the change when it is
|
||||
committed.
|
||||
|
||||
All rights reserved.
|
||||
|
||||
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 names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America
|
||||
and IDIAP Research Institute 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 OWNER 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.
|
||||
29
rsl_rl/rsl_rl/__init__.py
Normal file
29
rsl_rl/rsl_rl/__init__.py
Normal file
@@ -0,0 +1,29 @@
|
||||
# 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
|
||||
33
rsl_rl/rsl_rl/algorithms/__init__.py
Normal file
33
rsl_rl/rsl_rl/algorithms/__init__.py
Normal file
@@ -0,0 +1,33 @@
|
||||
# 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 .ppo import PPO
|
||||
from .cts import CTS
|
||||
from .moe_cts import MoECTS
|
||||
286
rsl_rl/rsl_rl/algorithms/cts.py
Normal file
286
rsl_rl/rsl_rl/algorithms/cts.py
Normal file
@@ -0,0 +1,286 @@
|
||||
# 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
|
||||
|
||||
import itertools
|
||||
from rsl_rl.modules import ActorCriticCTS
|
||||
from rsl_rl.storage import RolloutStorageCTS
|
||||
|
||||
class CTS:
|
||||
model: ActorCriticCTS
|
||||
def __init__(self,
|
||||
model,
|
||||
num_envs,
|
||||
history_length,
|
||||
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,
|
||||
student_encoder_learning_rate=1e-3,
|
||||
max_grad_norm=1.0,
|
||||
use_clipped_value_loss=True,
|
||||
schedule="fixed",
|
||||
desired_kl=0.01,
|
||||
teacher_env_ratio=0.75,
|
||||
device='cpu',
|
||||
):
|
||||
|
||||
self.device = device
|
||||
|
||||
self.desired_kl = desired_kl
|
||||
self.schedule = schedule
|
||||
self.learning_rate = learning_rate
|
||||
self.history_length = history_length
|
||||
|
||||
# CTS components
|
||||
self.model = model
|
||||
self.model.to(self.device)
|
||||
self.storage = None # initialized later
|
||||
params1 = [
|
||||
{"params": self.model.teacher_encoder.parameters()},
|
||||
{"params": self.model.critic.parameters()},
|
||||
{"params": self.model.actor.parameters()},
|
||||
{"params": self.model.std}
|
||||
]
|
||||
self.optimizer1 = optim.Adam(params1, lr=learning_rate)
|
||||
self.optimizer2 = optim.Adam(self.model.student_encoder.parameters(), lr=student_encoder_learning_rate)
|
||||
self.transition = RolloutStorageCTS.Transition()
|
||||
|
||||
# CTS 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
|
||||
self.teacher_num_envs = max(int(num_envs * teacher_env_ratio), 1)
|
||||
self.student_num_envs = num_envs - self.teacher_num_envs
|
||||
student_env_ratio = 1 - teacher_env_ratio
|
||||
self.teacher_env_idxs = torch.tensor([i for i in range(num_envs) if i % int(1/student_env_ratio) != 0], device=self.device)
|
||||
self.student_env_idxs = torch.tensor([i for i in range(num_envs) if i % int(1/student_env_ratio) == 0], device=self.device)
|
||||
# self.teacher_env_idxs = torch.arange(num_envs, device=self.device)
|
||||
# self.student_env_idxs = []
|
||||
assert len(self.teacher_env_idxs) == self.teacher_num_envs, f"{len(self.teacher_env_idxs)=} != {self.teacher_num_envs=}"
|
||||
assert len(self.student_env_idxs) == self.student_num_envs, f"{len(self.student_env_idxs)=} != {self.student_num_envs=}"
|
||||
|
||||
def init_storage(self, num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, action_shape):
|
||||
self.storage = RolloutStorageCTS(num_envs, self.teacher_num_envs, self.history_length, num_transitions_per_env, actor_obs_shape, critic_obs_shape, action_shape, self.device)
|
||||
|
||||
def test_mode(self):
|
||||
self.model.test()
|
||||
|
||||
def train_mode(self):
|
||||
self.model.train()
|
||||
|
||||
def act(self, obs, privileged_obs, history):
|
||||
history = history.clone()
|
||||
def get_results(obs, privileged_obs, history, is_teacher):
|
||||
actions = self.model.act(obs, privileged_obs, history, is_teacher).detach()
|
||||
return (
|
||||
actions,
|
||||
self.model.evaluate(privileged_obs, history, is_teacher).detach(),
|
||||
self.model.get_actions_log_prob(actions).detach(),
|
||||
self.model.action_mean.detach(),
|
||||
self.model.action_std.detach(),
|
||||
)
|
||||
ti, si = self.teacher_env_idxs, self.student_env_idxs
|
||||
teacher_results = get_results(obs[ti], privileged_obs[ti], history[ti], True)
|
||||
student_results = get_results(obs[si], privileged_obs[si], history[si], False)
|
||||
results = []
|
||||
for x1, x2 in zip(teacher_results, student_results):
|
||||
results.append(torch.cat([x1, x2], dim=0))
|
||||
# Compute the actions and values
|
||||
self.transition.actions = results[0]
|
||||
self.transition.values = results[1]
|
||||
self.transition.actions_log_prob = results[2]
|
||||
self.transition.action_mean = results[3]
|
||||
self.transition.action_sigma = results[4]
|
||||
# need to record obs and critic_obs before env.step()
|
||||
self.transition.history = torch.cat([history[ti], history[si]], dim=0)
|
||||
self.transition.observations = torch.cat([obs[ti], obs[si]], dim=0)
|
||||
self.transition.critic_observations = torch.cat([privileged_obs[ti], privileged_obs[si]], dim=0)
|
||||
real_actions = torch.zeros_like(self.transition.actions)
|
||||
real_actions[ti] = self.transition.actions[:self.teacher_num_envs]
|
||||
real_actions[si] = self.transition.actions[self.teacher_num_envs:]
|
||||
return real_actions
|
||||
|
||||
def process_env_step(self, rewards, dones, infos):
|
||||
ti, si = self.teacher_env_idxs, self.student_env_idxs
|
||||
rewards = rewards.clone()
|
||||
self.transition.rewards = torch.cat([rewards[ti], rewards[si]], dim=0)
|
||||
self.transition.dones = torch.cat([dones[ti], dones[si]], dim=0)
|
||||
# Bootstrapping on time outs
|
||||
if 'time_outs' in infos:
|
||||
time_outs = torch.cat([infos['time_outs'][ti], infos['time_outs'][si]], dim=0)
|
||||
self.transition.rewards += self.gamma * torch.squeeze(self.transition.values * time_outs.unsqueeze(1).to(self.device), 1)
|
||||
|
||||
# Record the transition
|
||||
self.storage.add_transitions(self.transition)
|
||||
self.transition.clear()
|
||||
self.model.reset(dones)
|
||||
|
||||
def compute_returns(self, last_privileged_obs, last_history):
|
||||
ti, si = self.teacher_env_idxs, self.student_env_idxs
|
||||
last_values = torch.cat([
|
||||
self.model.evaluate(last_privileged_obs[ti], last_history[ti], True).detach(),
|
||||
self.model.evaluate(last_privileged_obs[si], last_history[si], False).detach(),
|
||||
], dim=0)
|
||||
self.storage.compute_returns(last_values, self.gamma, self.lam)
|
||||
|
||||
def update(self):
|
||||
mean_value_loss = 0
|
||||
mean_surrogate_loss = 0
|
||||
mean_entropy_loss = 0
|
||||
mean_latent_loss = 0
|
||||
assert not self.model.is_recurrent
|
||||
data = list(self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs))
|
||||
teacher_samples = self.teacher_num_envs * self.storage.num_transitions_per_env // self.num_mini_batches
|
||||
student_samples = self.student_num_envs * self.storage.num_transitions_per_env // self.num_mini_batches
|
||||
for sample in data:
|
||||
(
|
||||
obs_batch, privileged_obs_batch, actions_batch, history_batch,
|
||||
target_values_batch, advantages_batch, returns_batch,
|
||||
old_actions_log_prob_batch, old_mu_batch, old_sigma_batch,
|
||||
hid_states_batch, masks_batch
|
||||
) = sample
|
||||
def get_results(start, end, is_teacher):
|
||||
self.model.act(obs_batch[start:end], privileged_obs_batch[start:end], history_batch[start:end], is_teacher)
|
||||
actions_log_prob = self.model.get_actions_log_prob(actions_batch[start:end])
|
||||
value = self.model.evaluate(privileged_obs_batch[start:end], history_batch[start:end], is_teacher)
|
||||
mu = self.model.action_mean
|
||||
sigma = self.model.action_std
|
||||
entropy = self.model.entropy
|
||||
return actions_log_prob, value, mu, sigma, entropy
|
||||
teacher_results = get_results(0, teacher_samples, True)
|
||||
student_results = get_results(teacher_samples, teacher_samples + student_samples, False)
|
||||
results = []
|
||||
for x1, x2 in zip(teacher_results, student_results):
|
||||
results.append(torch.cat([x1, x2], dim=0))
|
||||
actions_log_prob_batch = results[0]
|
||||
value_batch = results[1]
|
||||
mu_batch = results[2]
|
||||
sigma_batch = results[3]
|
||||
entropy_batch = results[4]
|
||||
|
||||
# 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.optimizer1.param_groups:
|
||||
param_group['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_losses = torch.max(surrogate, surrogate_clipped)
|
||||
teacher_surrogate_loss = surrogate_losses[:teacher_samples].mean()
|
||||
student_surrogate_loss = surrogate_losses[teacher_samples:].mean()
|
||||
surrogate_loss = teacher_surrogate_loss + student_surrogate_loss
|
||||
# surrogate_loss = teacher_surrogate_loss
|
||||
|
||||
# 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()
|
||||
# teacher_value_loss = value_losses[:teacher_samples].mean()
|
||||
# student_value_loss = value_losses[teacher_samples:].mean()
|
||||
# value_loss = teacher_value_loss # + student_value_loss
|
||||
|
||||
loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean()
|
||||
|
||||
# Gradient step
|
||||
self.optimizer1.zero_grad()
|
||||
loss.backward()
|
||||
params_to_clip = itertools.chain.from_iterable(g['params'] for g in self.optimizer1.param_groups)
|
||||
nn.utils.clip_grad_norm_(params_to_clip, self.max_grad_norm)
|
||||
self.optimizer1.step()
|
||||
|
||||
mean_value_loss += value_loss.item()
|
||||
mean_surrogate_loss += surrogate_loss.item()
|
||||
mean_entropy_loss += entropy_batch.mean().item()
|
||||
|
||||
for sample in data:
|
||||
(
|
||||
obs_batch, privileged_obs_batch, actions_batch, history_batch,
|
||||
target_values_batch, advantages_batch, returns_batch,
|
||||
old_actions_log_prob_batch, old_mu_batch, old_sigma_batch,
|
||||
hid_states_batch, masks_batch
|
||||
) = sample
|
||||
# Student encoder update
|
||||
student_latent = self.model.student_encoder(history_batch[teacher_samples:])
|
||||
with torch.no_grad():
|
||||
teacher_latent = self.model.teacher_encoder(privileged_obs_batch[teacher_samples:])
|
||||
latent_loss = (teacher_latent - student_latent).pow(2).mean()
|
||||
|
||||
self.optimizer2.zero_grad()
|
||||
latent_loss.backward()
|
||||
nn.utils.clip_grad_norm_(self.model.student_encoder.parameters(), self.max_grad_norm)
|
||||
self.optimizer2.step()
|
||||
|
||||
mean_latent_loss += latent_loss.item()
|
||||
|
||||
num_updates = self.num_learning_epochs * self.num_mini_batches
|
||||
mean_value_loss /= num_updates
|
||||
mean_surrogate_loss /= num_updates
|
||||
mean_entropy_loss /= num_updates
|
||||
mean_latent_loss /= num_updates
|
||||
self.storage.clear()
|
||||
|
||||
return mean_value_loss, mean_surrogate_loss, mean_entropy_loss, mean_latent_loss
|
||||
234
rsl_rl/rsl_rl/algorithms/moe_cts.py
Normal file
234
rsl_rl/rsl_rl/algorithms/moe_cts.py
Normal file
@@ -0,0 +1,234 @@
|
||||
# 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
|
||||
|
||||
import itertools
|
||||
from rsl_rl.modules import ActorCriticMoECTS
|
||||
from rsl_rl.storage import RolloutStorageCTS
|
||||
from rsl_rl.algorithms.cts import CTS
|
||||
|
||||
class MoECTS(CTS):
|
||||
model: ActorCriticMoECTS
|
||||
def __init__(self,
|
||||
model,
|
||||
num_envs,
|
||||
history_length,
|
||||
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,
|
||||
load_balance_coef=0.01,
|
||||
learning_rate=1e-3,
|
||||
student_encoder_learning_rate=1e-3,
|
||||
max_grad_norm=1.0,
|
||||
use_clipped_value_loss=True,
|
||||
schedule="fixed",
|
||||
desired_kl=0.01,
|
||||
teacher_env_ratio=0.75,
|
||||
device='cpu',
|
||||
):
|
||||
|
||||
self.device = device
|
||||
|
||||
self.desired_kl = desired_kl
|
||||
self.schedule = schedule
|
||||
self.learning_rate = learning_rate
|
||||
self.history_length = history_length
|
||||
|
||||
# CTS components
|
||||
self.model = model
|
||||
self.model.to(self.device)
|
||||
self.storage = None # initialized later
|
||||
params1 = [
|
||||
{"params": self.model.teacher_encoder.parameters()},
|
||||
{"params": self.model.critic.parameters()},
|
||||
{"params": self.model.actor.parameters()},
|
||||
{"params": self.model.std}
|
||||
]
|
||||
self.optimizer1 = optim.Adam(params1, lr=learning_rate)
|
||||
self.optimizer2 = optim.Adam(self.model.student_moe_encoder.parameters(), lr=student_encoder_learning_rate)
|
||||
self.transition = RolloutStorageCTS.Transition()
|
||||
|
||||
# CTS 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.load_balance_coef = load_balance_coef
|
||||
self.gamma = gamma
|
||||
self.lam = lam
|
||||
self.max_grad_norm = max_grad_norm
|
||||
self.use_clipped_value_loss = use_clipped_value_loss
|
||||
self.teacher_num_envs = max(int(num_envs * teacher_env_ratio), 1)
|
||||
self.student_num_envs = num_envs - self.teacher_num_envs
|
||||
student_env_ratio = 1 - teacher_env_ratio
|
||||
self.teacher_env_idxs = torch.tensor([i for i in range(num_envs) if i % int(1/student_env_ratio) != 0], device=self.device)
|
||||
self.student_env_idxs = torch.tensor([i for i in range(num_envs) if i % int(1/student_env_ratio) == 0], device=self.device)
|
||||
assert len(self.teacher_env_idxs) == self.teacher_num_envs, f"{len(self.teacher_env_idxs)=} != {self.teacher_num_envs=}"
|
||||
assert len(self.student_env_idxs) == self.student_num_envs, f"{len(self.student_env_idxs)=} != {self.student_num_envs=}"
|
||||
|
||||
def update(self):
|
||||
mean_value_loss = 0
|
||||
mean_surrogate_loss = 0
|
||||
mean_entropy_loss = 0
|
||||
mean_latent_loss = 0
|
||||
mean_load_balance_loss = 0
|
||||
assert not self.model.is_recurrent
|
||||
data = list(self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs))
|
||||
teacher_samples = self.teacher_num_envs * self.storage.num_transitions_per_env // self.num_mini_batches
|
||||
student_samples = self.student_num_envs * self.storage.num_transitions_per_env // self.num_mini_batches
|
||||
for sample in data:
|
||||
(
|
||||
obs_batch, privileged_obs_batch, actions_batch, history_batch,
|
||||
target_values_batch, advantages_batch, returns_batch,
|
||||
old_actions_log_prob_batch, old_mu_batch, old_sigma_batch,
|
||||
hid_states_batch, masks_batch
|
||||
) = sample
|
||||
def get_results(start, end, is_teacher):
|
||||
self.model.act(obs_batch[start:end], privileged_obs_batch[start:end], history_batch[start:end], is_teacher)
|
||||
actions_log_prob = self.model.get_actions_log_prob(actions_batch[start:end])
|
||||
value = self.model.evaluate(privileged_obs_batch[start:end], history_batch[start:end], is_teacher)
|
||||
mu = self.model.action_mean
|
||||
sigma = self.model.action_std
|
||||
entropy = self.model.entropy
|
||||
return actions_log_prob, value, mu, sigma, entropy
|
||||
teacher_results = get_results(0, teacher_samples, True)
|
||||
student_results = get_results(teacher_samples, teacher_samples + student_samples, False)
|
||||
results = []
|
||||
for x1, x2 in zip(teacher_results, student_results):
|
||||
results.append(torch.cat([x1, x2], dim=0))
|
||||
actions_log_prob_batch = results[0]
|
||||
value_batch = results[1]
|
||||
mu_batch = results[2]
|
||||
sigma_batch = results[3]
|
||||
entropy_batch = results[4]
|
||||
|
||||
# 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.optimizer1.param_groups:
|
||||
param_group['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_losses = torch.max(surrogate, surrogate_clipped)
|
||||
teacher_surrogate_loss = surrogate_losses[:teacher_samples].mean()
|
||||
student_surrogate_loss = surrogate_losses[teacher_samples:].mean()
|
||||
surrogate_loss = teacher_surrogate_loss + student_surrogate_loss
|
||||
# surrogate_loss = teacher_surrogate_loss
|
||||
|
||||
# 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()
|
||||
# teacher_value_loss = value_losses[:teacher_samples].mean()
|
||||
# student_value_loss = value_losses[teacher_samples:].mean()
|
||||
# value_loss = teacher_value_loss # + student_value_loss
|
||||
|
||||
loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean()
|
||||
|
||||
# Gradient step
|
||||
self.optimizer1.zero_grad()
|
||||
loss.backward()
|
||||
params_to_clip = itertools.chain.from_iterable(g['params'] for g in self.optimizer1.param_groups)
|
||||
nn.utils.clip_grad_norm_(params_to_clip, self.max_grad_norm)
|
||||
self.optimizer1.step()
|
||||
|
||||
mean_value_loss += value_loss.item()
|
||||
mean_surrogate_loss += surrogate_loss.item()
|
||||
mean_entropy_loss += entropy_batch.mean().item()
|
||||
|
||||
for sample in data:
|
||||
(
|
||||
obs_batch, privileged_obs_batch, actions_batch, history_batch,
|
||||
target_values_batch, advantages_batch, returns_batch,
|
||||
old_actions_log_prob_batch, old_mu_batch, old_sigma_batch,
|
||||
hid_states_batch, masks_batch
|
||||
) = sample
|
||||
# Student encoder update
|
||||
student_latent, gating_weights = self.model.get_student_latent_and_weights(history_batch[teacher_samples:])
|
||||
with torch.no_grad():
|
||||
teacher_latent = self.model.teacher_encoder(privileged_obs_batch[teacher_samples:])
|
||||
latent_loss = (teacher_latent - student_latent).pow(2).mean()
|
||||
|
||||
# Load balance loss
|
||||
mean_usage = torch.mean(gating_weights, dim=0)
|
||||
target_usage = torch.full_like(mean_usage, 1.0 / gating_weights.shape[1])
|
||||
load_balance_loss = torch.mean((mean_usage - target_usage).pow(2))
|
||||
# load_balance_loss = torch.sum(mean_usage.pow(2)) * gating_weights.shape[1] # Switch Transformer style
|
||||
|
||||
student_loss = latent_loss + self.load_balance_coef * load_balance_loss
|
||||
|
||||
self.optimizer2.zero_grad()
|
||||
student_loss.backward()
|
||||
nn.utils.clip_grad_norm_(self.model.student_moe_encoder.parameters(), self.max_grad_norm)
|
||||
self.optimizer2.step()
|
||||
|
||||
mean_latent_loss += latent_loss.item()
|
||||
mean_load_balance_loss += load_balance_loss.item()
|
||||
|
||||
num_updates = self.num_learning_epochs * self.num_mini_batches
|
||||
mean_value_loss /= num_updates
|
||||
mean_surrogate_loss /= num_updates
|
||||
mean_entropy_loss /= num_updates
|
||||
mean_latent_loss /= num_updates
|
||||
mean_load_balance_loss /= num_updates
|
||||
self.storage.clear()
|
||||
|
||||
return mean_value_loss, mean_surrogate_loss, mean_entropy_loss, mean_latent_loss, mean_load_balance_loss
|
||||
187
rsl_rl/rsl_rl/algorithms/ppo.py
Normal file
187
rsl_rl/rsl_rl/algorithms/ppo.py
Normal file
@@ -0,0 +1,187 @@
|
||||
# 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 rsl_rl.modules import ActorCritic
|
||||
from rsl_rl.storage import RolloutStorage
|
||||
|
||||
class PPO:
|
||||
actor_critic: ActorCritic
|
||||
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 = RolloutStorage.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 = RolloutStorage(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):
|
||||
if self.actor_critic.is_recurrent:
|
||||
self.transition.hidden_states = self.actor_critic.get_hidden_states()
|
||||
# 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):
|
||||
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
|
||||
if self.actor_critic.is_recurrent:
|
||||
generator = self.storage.reccurent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs)
|
||||
else:
|
||||
generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs)
|
||||
for obs_batch, critic_obs_batch, actions_batch, target_values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, \
|
||||
old_mu_batch, old_sigma_batch, hid_states_batch, masks_batch in generator:
|
||||
|
||||
|
||||
self.actor_critic.act(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[0])
|
||||
actions_log_prob_batch = self.actor_critic.get_actions_log_prob(actions_batch)
|
||||
value_batch = self.actor_critic.evaluate(critic_obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1])
|
||||
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
|
||||
|
||||
|
||||
# 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()
|
||||
|
||||
num_updates = self.num_learning_epochs * self.num_mini_batches
|
||||
mean_value_loss /= num_updates
|
||||
mean_surrogate_loss /= num_updates
|
||||
self.storage.clear()
|
||||
|
||||
return mean_value_loss, mean_surrogate_loss
|
||||
31
rsl_rl/rsl_rl/env/__init__.py
vendored
Normal file
31
rsl_rl/rsl_rl/env/__init__.py
vendored
Normal file
@@ -0,0 +1,31 @@
|
||||
# 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 .vec_env import VecEnv
|
||||
60
rsl_rl/rsl_rl/env/vec_env.py
vendored
Normal file
60
rsl_rl/rsl_rl/env/vec_env.py
vendored
Normal file
@@ -0,0 +1,60 @@
|
||||
# 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 abc import ABC, abstractmethod
|
||||
import torch
|
||||
from typing import Tuple, Union
|
||||
|
||||
# minimal interface of the environment
|
||||
class VecEnv(ABC):
|
||||
num_envs: int
|
||||
num_obs: int
|
||||
num_privileged_obs: int
|
||||
num_actions: int
|
||||
max_episode_length: int
|
||||
privileged_obs_buf: torch.Tensor
|
||||
obs_buf: torch.Tensor
|
||||
rew_buf: torch.Tensor
|
||||
reset_buf: torch.Tensor
|
||||
episode_length_buf: torch.Tensor # current episode duration
|
||||
extras: dict
|
||||
device: torch.device
|
||||
@abstractmethod
|
||||
def step(self, actions: torch.Tensor) -> Tuple[torch.Tensor, Union[torch.Tensor, None], torch.Tensor, torch.Tensor, dict]:
|
||||
pass
|
||||
@abstractmethod
|
||||
def reset(self, env_ids: Union[list, torch.Tensor]):
|
||||
pass
|
||||
@abstractmethod
|
||||
def get_observations(self) -> torch.Tensor:
|
||||
pass
|
||||
@abstractmethod
|
||||
def get_privileged_observations(self) -> Union[torch.Tensor, None]:
|
||||
pass
|
||||
34
rsl_rl/rsl_rl/modules/__init__.py
Normal file
34
rsl_rl/rsl_rl/modules/__init__.py
Normal file
@@ -0,0 +1,34 @@
|
||||
# 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 .actor_critic import ActorCritic
|
||||
from .actor_critic_recurrent import ActorCriticRecurrent
|
||||
from .actor_critic_cts import ActorCriticCTS
|
||||
from .actor_critic_moe_cts import ActorCriticMoECTS
|
||||
155
rsl_rl/rsl_rl/modules/actor_critic.py
Normal file
155
rsl_rl/rsl_rl/modules/actor_critic.py
Normal file
@@ -0,0 +1,155 @@
|
||||
# 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 numpy as np
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.distributions import Normal
|
||||
from torch.nn.modules import rnn
|
||||
|
||||
class ActorCritic(nn.Module):
|
||||
is_recurrent = False
|
||||
def __init__(self, num_actor_obs,
|
||||
num_critic_obs,
|
||||
num_actions,
|
||||
actor_hidden_dims=[256, 256, 256],
|
||||
critic_hidden_dims=[256, 256, 256],
|
||||
activation='elu',
|
||||
init_noise_std=1.0,
|
||||
**kwargs):
|
||||
if kwargs:
|
||||
print("ActorCritic.__init__ got unexpected arguments, which will be ignored: " + str([key for key in kwargs.keys()]))
|
||||
super(ActorCritic, self).__init__()
|
||||
|
||||
activation = get_activation(activation)
|
||||
|
||||
mlp_input_dim_a = num_actor_obs
|
||||
mlp_input_dim_c = num_critic_obs
|
||||
|
||||
# Policy
|
||||
actor_layers = []
|
||||
actor_layers.append(nn.Linear(mlp_input_dim_a, actor_hidden_dims[0]))
|
||||
actor_layers.append(activation)
|
||||
for l in range(len(actor_hidden_dims)):
|
||||
if l == len(actor_hidden_dims) - 1:
|
||||
actor_layers.append(nn.Linear(actor_hidden_dims[l], num_actions))
|
||||
else:
|
||||
actor_layers.append(nn.Linear(actor_hidden_dims[l], actor_hidden_dims[l + 1]))
|
||||
actor_layers.append(activation)
|
||||
self.actor = nn.Sequential(*actor_layers)
|
||||
|
||||
# Value function
|
||||
critic_layers = []
|
||||
critic_layers.append(nn.Linear(mlp_input_dim_c, critic_hidden_dims[0]))
|
||||
critic_layers.append(activation)
|
||||
for l in range(len(critic_hidden_dims)):
|
||||
if l == len(critic_hidden_dims) - 1:
|
||||
critic_layers.append(nn.Linear(critic_hidden_dims[l], 1))
|
||||
else:
|
||||
critic_layers.append(nn.Linear(critic_hidden_dims[l], critic_hidden_dims[l + 1]))
|
||||
critic_layers.append(activation)
|
||||
self.critic = nn.Sequential(*critic_layers)
|
||||
|
||||
print(f"Actor MLP: {self.actor}")
|
||||
print(f"Critic MLP: {self.critic}")
|
||||
|
||||
# Action noise
|
||||
self.std = nn.Parameter(init_noise_std * torch.ones(num_actions))
|
||||
self.distribution = None
|
||||
# disable args validation for speedup
|
||||
Normal.set_default_validate_args = False
|
||||
|
||||
# seems that we get better performance without init
|
||||
# self.init_memory_weights(self.memory_a, 0.001, 0.)
|
||||
# self.init_memory_weights(self.memory_c, 0.001, 0.)
|
||||
|
||||
@staticmethod
|
||||
# not used at the moment
|
||||
def init_weights(sequential, scales):
|
||||
[torch.nn.init.orthogonal_(module.weight, gain=scales[idx]) for idx, module in
|
||||
enumerate(mod for mod in sequential if isinstance(mod, nn.Linear))]
|
||||
|
||||
|
||||
def reset(self, dones=None):
|
||||
pass
|
||||
|
||||
def forward(self):
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def action_mean(self):
|
||||
return self.distribution.mean
|
||||
|
||||
@property
|
||||
def action_std(self):
|
||||
return self.distribution.stddev
|
||||
|
||||
@property
|
||||
def entropy(self):
|
||||
return self.distribution.entropy().sum(dim=-1)
|
||||
|
||||
def update_distribution(self, observations):
|
||||
mean = self.actor(observations)
|
||||
self.distribution = Normal(mean, mean*0. + self.std)
|
||||
|
||||
def act(self, observations, **kwargs):
|
||||
self.update_distribution(observations)
|
||||
return self.distribution.sample()
|
||||
|
||||
def get_actions_log_prob(self, actions):
|
||||
return self.distribution.log_prob(actions).sum(dim=-1)
|
||||
|
||||
def act_inference(self, observations):
|
||||
actions_mean = self.actor(observations)
|
||||
return actions_mean
|
||||
|
||||
def evaluate(self, critic_observations, **kwargs):
|
||||
value = self.critic(critic_observations)
|
||||
return value
|
||||
|
||||
def get_activation(act_name):
|
||||
if act_name == "elu":
|
||||
return nn.ELU()
|
||||
elif act_name == "selu":
|
||||
return nn.SELU()
|
||||
elif act_name == "relu":
|
||||
return nn.ReLU()
|
||||
elif act_name == "crelu":
|
||||
return nn.ReLU()
|
||||
elif act_name == "lrelu":
|
||||
return nn.LeakyReLU()
|
||||
elif act_name == "tanh":
|
||||
return nn.Tanh()
|
||||
elif act_name == "sigmoid":
|
||||
return nn.Sigmoid()
|
||||
else:
|
||||
print("invalid activation function!")
|
||||
return None
|
||||
243
rsl_rl/rsl_rl/modules/actor_critic_cts.py
Normal file
243
rsl_rl/rsl_rl/modules/actor_critic_cts.py
Normal file
@@ -0,0 +1,243 @@
|
||||
# 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 numpy as np
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.distributions import Normal
|
||||
|
||||
class ActorCriticCTS(nn.Module):
|
||||
is_recurrent = False
|
||||
def __init__(self, num_actor_obs,
|
||||
num_critic_obs,
|
||||
num_actions,
|
||||
num_envs,
|
||||
history_length,
|
||||
actor_hidden_dims=[512, 256, 128],
|
||||
critic_hidden_dims=[512, 256, 128],
|
||||
teacher_encoder_hidden_dims=[512, 256],
|
||||
student_encoder_hidden_dims=[512, 256],
|
||||
activation='elu',
|
||||
init_noise_std=1.0,
|
||||
latent_dim=32,
|
||||
norm_type='l2norm',
|
||||
**kwargs):
|
||||
if kwargs:
|
||||
print("ActorCritic.__init__ got unexpected arguments, which will be ignored: " + str([key for key in kwargs.keys()]))
|
||||
assert norm_type in ['l2norm', 'simnorm'], f"Normalization type {norm_type} not supported!"
|
||||
super(ActorCriticCTS, self).__init__()
|
||||
self.num_actions = num_actions
|
||||
|
||||
activation = get_activation(activation)
|
||||
|
||||
mlp_input_dim_t = num_critic_obs
|
||||
mlp_input_dim_s = num_actor_obs * history_length
|
||||
mlp_input_dim_a = latent_dim + num_actor_obs
|
||||
mlp_input_dim_c = latent_dim + num_critic_obs
|
||||
|
||||
# History
|
||||
self.history = torch.zeros((num_envs, history_length, num_actor_obs), device='cuda')
|
||||
|
||||
# Teacher encoder
|
||||
encoder_layers = []
|
||||
encoder_layers.append(nn.Linear(mlp_input_dim_t, teacher_encoder_hidden_dims[0]))
|
||||
encoder_layers.append(activation)
|
||||
for l in range(len(teacher_encoder_hidden_dims)):
|
||||
if l == len(teacher_encoder_hidden_dims) - 1:
|
||||
encoder_layers.append(nn.Linear(teacher_encoder_hidden_dims[l], latent_dim))
|
||||
if norm_type == 'l2norm':
|
||||
encoder_layers.append(L2Norm())
|
||||
elif norm_type == 'simnorm':
|
||||
encoder_layers.append(SimNorm())
|
||||
else:
|
||||
encoder_layers.append(nn.Linear(teacher_encoder_hidden_dims[l], teacher_encoder_hidden_dims[l + 1]))
|
||||
encoder_layers.append(activation)
|
||||
self.teacher_encoder = nn.Sequential(*encoder_layers)
|
||||
|
||||
# Student encoder
|
||||
encoder_layers = []
|
||||
encoder_layers.append(nn.Linear(mlp_input_dim_s, student_encoder_hidden_dims[0]))
|
||||
encoder_layers.append(activation)
|
||||
for l in range(len(student_encoder_hidden_dims)):
|
||||
if l == len(student_encoder_hidden_dims) - 1:
|
||||
encoder_layers.append(nn.Linear(student_encoder_hidden_dims[l], latent_dim))
|
||||
if norm_type == 'l2norm':
|
||||
encoder_layers.append(L2Norm())
|
||||
elif norm_type == 'simnorm':
|
||||
encoder_layers.append(SimNorm())
|
||||
else:
|
||||
encoder_layers.append(nn.Linear(student_encoder_hidden_dims[l], student_encoder_hidden_dims[l + 1]))
|
||||
encoder_layers.append(activation)
|
||||
self.student_encoder = nn.Sequential(*encoder_layers)
|
||||
|
||||
# Policy
|
||||
actor_layers = []
|
||||
actor_layers.append(nn.Linear(mlp_input_dim_a, actor_hidden_dims[0]))
|
||||
actor_layers.append(activation)
|
||||
for l in range(len(actor_hidden_dims)):
|
||||
if l == len(actor_hidden_dims) - 1:
|
||||
actor_layers.append(nn.Linear(actor_hidden_dims[l], num_actions))
|
||||
else:
|
||||
actor_layers.append(nn.Linear(actor_hidden_dims[l], actor_hidden_dims[l + 1]))
|
||||
actor_layers.append(activation)
|
||||
self.actor = nn.Sequential(*actor_layers)
|
||||
|
||||
# Value function
|
||||
critic_layers = []
|
||||
critic_layers.append(nn.Linear(mlp_input_dim_c, critic_hidden_dims[0]))
|
||||
critic_layers.append(activation)
|
||||
for l in range(len(critic_hidden_dims)):
|
||||
if l == len(critic_hidden_dims) - 1:
|
||||
critic_layers.append(nn.Linear(critic_hidden_dims[l], 1))
|
||||
else:
|
||||
critic_layers.append(nn.Linear(critic_hidden_dims[l], critic_hidden_dims[l + 1]))
|
||||
critic_layers.append(activation)
|
||||
self.critic = nn.Sequential(*critic_layers)
|
||||
|
||||
print(f"Actor MLP: {self.actor}")
|
||||
print(f"Critic MLP: {self.critic}")
|
||||
print(f"Teacher Encoder MLP: {self.teacher_encoder}")
|
||||
print(f"Student Encoder MLP: {self.student_encoder}")
|
||||
|
||||
# Action noise
|
||||
self.std = nn.Parameter(init_noise_std * torch.ones(num_actions))
|
||||
self.distribution = None
|
||||
# disable args validation for speedup
|
||||
Normal.set_default_validate_args = False
|
||||
|
||||
# seems that we get better performance without init
|
||||
# self.init_memory_weights(self.memory_a, 0.001, 0.)
|
||||
# self.init_memory_weights(self.memory_c, 0.001, 0.)
|
||||
|
||||
@staticmethod
|
||||
# not used at the moment
|
||||
def init_weights(sequential, scales):
|
||||
[torch.nn.init.orthogonal_(module.weight, gain=scales[idx]) for idx, module in
|
||||
enumerate(mod for mod in sequential if isinstance(mod, nn.Linear))]
|
||||
|
||||
|
||||
def reset(self, dones=None):
|
||||
self.history[dones > 0] = 0.0
|
||||
|
||||
def forward(self):
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def action_mean(self):
|
||||
return self.distribution.mean
|
||||
|
||||
@property
|
||||
def action_std(self):
|
||||
return self.distribution.stddev
|
||||
|
||||
@property
|
||||
def entropy(self):
|
||||
return self.distribution.entropy().sum(dim=-1)
|
||||
|
||||
def update_distribution(self, latent_and_obs):
|
||||
mean = self.actor(latent_and_obs)
|
||||
self.distribution = Normal(mean, mean*0. + self.std)
|
||||
|
||||
def act(self, obs, privileged_obs, history, is_teacher, **kwargs):
|
||||
if is_teacher:
|
||||
latent = self.teacher_encoder(privileged_obs)
|
||||
else:
|
||||
latent = self.student_encoder(history).detach()
|
||||
x = torch.cat([latent, obs], dim=1)
|
||||
self.update_distribution(x)
|
||||
return self.distribution.sample()
|
||||
|
||||
def get_actions_log_prob(self, actions):
|
||||
return self.distribution.log_prob(actions).sum(dim=-1)
|
||||
|
||||
def act_inference(self, obs):
|
||||
self.history = torch.cat([self.history[:, 1:], obs.unsqueeze(1)], dim=1)
|
||||
latent = self.student_encoder(self.history.flatten(1))
|
||||
x = torch.cat([latent, obs], dim=1)
|
||||
actions_mean = self.actor(x)
|
||||
return actions_mean
|
||||
|
||||
def evaluate(self, privileged_obs, history, is_teacher, **kwargs):
|
||||
if is_teacher:
|
||||
latent = self.teacher_encoder(privileged_obs)
|
||||
else:
|
||||
latent = self.student_encoder(history)
|
||||
x = torch.cat([latent.detach(), privileged_obs], dim=1)
|
||||
value = self.critic(x)
|
||||
return value
|
||||
|
||||
def get_activation(act_name):
|
||||
if act_name == "elu":
|
||||
return nn.ELU()
|
||||
elif act_name == "selu":
|
||||
return nn.SELU()
|
||||
elif act_name == "relu":
|
||||
return nn.ReLU()
|
||||
elif act_name == "crelu":
|
||||
return nn.ReLU()
|
||||
elif act_name == "lrelu":
|
||||
return nn.LeakyReLU()
|
||||
elif act_name == "tanh":
|
||||
return nn.Tanh()
|
||||
elif act_name == "sigmoid":
|
||||
return nn.Sigmoid()
|
||||
else:
|
||||
print("invalid activation function!")
|
||||
return None
|
||||
|
||||
class L2Norm(nn.Module):
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, x):
|
||||
return F.normalize(x, p=2.0, dim=-1)
|
||||
|
||||
class SimNorm(nn.Module):
|
||||
"""
|
||||
Simplicial normalization.
|
||||
Adapted from https://arxiv.org/abs/2204.00616.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.dim = 8 # for latent dim 512
|
||||
|
||||
def forward(self, x):
|
||||
shp = x.shape
|
||||
target_shape = list(shp[:-1]) + [-1, self.dim]
|
||||
x = x.view(target_shape)
|
||||
x = F.softmax(x, dim=-1)
|
||||
return x.view(shp)
|
||||
|
||||
def __repr__(self):
|
||||
return f"SimNorm(dim={self.dim})"
|
||||
301
rsl_rl/rsl_rl/modules/actor_critic_moe_cts.py
Normal file
301
rsl_rl/rsl_rl/modules/actor_critic_moe_cts.py
Normal file
@@ -0,0 +1,301 @@
|
||||
# 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 numpy as np
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.distributions import Normal
|
||||
|
||||
class ActorCriticMoECTS(nn.Module):
|
||||
is_recurrent = False
|
||||
def __init__(self, num_obs,
|
||||
num_critic_obs,
|
||||
num_actions,
|
||||
num_envs,
|
||||
history_length,
|
||||
obs_no_goal_mask,
|
||||
actor_hidden_dims=[512, 256, 128],
|
||||
critic_hidden_dims=[512, 256, 128],
|
||||
teacher_encoder_hidden_dims=[512, 256],
|
||||
student_encoder_hidden_dims=[512, 256],
|
||||
student_expert_num=8,
|
||||
activation='elu',
|
||||
init_noise_std=1.0,
|
||||
latent_dim=32,
|
||||
norm_type='l2norm',
|
||||
**kwargs):
|
||||
if kwargs:
|
||||
print("ActorCritic.__init__ got unexpected arguments, which will be ignored: " + str([key for key in kwargs.keys()]))
|
||||
assert norm_type in ['l2norm', 'simnorm'], f"Normalization type {norm_type} not supported!"
|
||||
super(ActorCriticMoECTS, self).__init__()
|
||||
self.num_actions = num_actions
|
||||
self.history_length = history_length
|
||||
self.register_buffer("obs_no_goal_mask", torch.tensor(obs_no_goal_mask, dtype=torch.bool), persistent=False)
|
||||
|
||||
activation_str = activation
|
||||
activation = get_activation(activation)
|
||||
|
||||
mlp_input_dim_t = num_critic_obs
|
||||
mlp_input_dim_e = torch.sum(self.obs_no_goal_mask).item() * history_length # exclude command inputs for expert
|
||||
mlp_input_dim_g = num_obs * history_length # all obs for gating
|
||||
mlp_input_dim_a = latent_dim + num_obs
|
||||
mlp_input_dim_c = latent_dim + num_critic_obs
|
||||
|
||||
# History
|
||||
self.register_buffer("history", torch.zeros((num_envs, history_length, num_obs)), persistent=False)
|
||||
|
||||
# Teacher encoder
|
||||
encoder_layers = []
|
||||
encoder_layers.append(nn.Linear(mlp_input_dim_t, teacher_encoder_hidden_dims[0]))
|
||||
encoder_layers.append(activation)
|
||||
for l in range(len(teacher_encoder_hidden_dims)):
|
||||
if l == len(teacher_encoder_hidden_dims) - 1:
|
||||
encoder_layers.append(nn.Linear(teacher_encoder_hidden_dims[l], latent_dim))
|
||||
if norm_type == 'l2norm':
|
||||
encoder_layers.append(L2Norm())
|
||||
elif norm_type == 'simnorm':
|
||||
encoder_layers.append(SimNorm())
|
||||
else:
|
||||
encoder_layers.append(nn.Linear(teacher_encoder_hidden_dims[l], teacher_encoder_hidden_dims[l + 1]))
|
||||
encoder_layers.append(activation)
|
||||
self.teacher_encoder = nn.Sequential(*encoder_layers)
|
||||
|
||||
# Student MoE encoder
|
||||
self.student_moe_encoder = StudentMoEEncoder(
|
||||
expert_dim=mlp_input_dim_e,
|
||||
gating_dim=mlp_input_dim_g,
|
||||
hidden_dims=student_encoder_hidden_dims,
|
||||
expert_num=student_expert_num,
|
||||
latent_dim=latent_dim,
|
||||
activation=activation_str
|
||||
)
|
||||
|
||||
# Policy
|
||||
actor_layers = []
|
||||
actor_layers.append(nn.Linear(mlp_input_dim_a, actor_hidden_dims[0]))
|
||||
actor_layers.append(activation)
|
||||
for l in range(len(actor_hidden_dims)):
|
||||
if l == len(actor_hidden_dims) - 1:
|
||||
actor_layers.append(nn.Linear(actor_hidden_dims[l], num_actions))
|
||||
else:
|
||||
actor_layers.append(nn.Linear(actor_hidden_dims[l], actor_hidden_dims[l + 1]))
|
||||
actor_layers.append(activation)
|
||||
self.actor = nn.Sequential(*actor_layers)
|
||||
|
||||
# Value function
|
||||
critic_layers = []
|
||||
critic_layers.append(nn.Linear(mlp_input_dim_c, critic_hidden_dims[0]))
|
||||
critic_layers.append(activation)
|
||||
for l in range(len(critic_hidden_dims)):
|
||||
if l == len(critic_hidden_dims) - 1:
|
||||
critic_layers.append(nn.Linear(critic_hidden_dims[l], 1))
|
||||
else:
|
||||
critic_layers.append(nn.Linear(critic_hidden_dims[l], critic_hidden_dims[l + 1]))
|
||||
critic_layers.append(activation)
|
||||
self.critic = nn.Sequential(*critic_layers)
|
||||
|
||||
print(f"Actor MLP: {self.actor}")
|
||||
print(f"Critic MLP: {self.critic}")
|
||||
print(f"Teacher Encoder: {self.teacher_encoder}")
|
||||
print(f"Student MoE Encoder: {self.student_moe_encoder}")
|
||||
|
||||
# Action noise
|
||||
self.std = nn.Parameter(init_noise_std * torch.ones(num_actions))
|
||||
self.distribution = None
|
||||
# disable args validation for speedup
|
||||
Normal.set_default_validate_args = False
|
||||
|
||||
# seems that we get better performance without init
|
||||
# self.init_memory_weights(self.memory_a, 0.001, 0.)
|
||||
# self.init_memory_weights(self.memory_c, 0.001, 0.)
|
||||
|
||||
@staticmethod
|
||||
# not used at the moment
|
||||
def init_weights(sequential, scales):
|
||||
[torch.nn.init.orthogonal_(module.weight, gain=scales[idx]) for idx, module in
|
||||
enumerate(mod for mod in sequential if isinstance(mod, nn.Linear))]
|
||||
|
||||
|
||||
def reset(self, dones=None):
|
||||
self.history[dones > 0] = 0.0
|
||||
|
||||
def forward(self):
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def action_mean(self):
|
||||
return self.distribution.mean
|
||||
|
||||
@property
|
||||
def action_std(self):
|
||||
return self.distribution.stddev
|
||||
|
||||
@property
|
||||
def entropy(self):
|
||||
return self.distribution.entropy().sum(dim=-1)
|
||||
|
||||
def update_distribution(self, latent_and_obs):
|
||||
mean = self.actor(latent_and_obs)
|
||||
self.distribution = Normal(mean, mean*0. + self.std)
|
||||
|
||||
def act(self, obs, privileged_obs, history, is_teacher, **kwargs):
|
||||
if is_teacher:
|
||||
latent = self.teacher_encoder(privileged_obs)
|
||||
else:
|
||||
with torch.no_grad():
|
||||
latent, _ = self.get_student_latent_and_weights(history)
|
||||
x = torch.cat([latent, obs], dim=1)
|
||||
self.update_distribution(x)
|
||||
return self.distribution.sample()
|
||||
|
||||
def get_actions_log_prob(self, actions):
|
||||
return self.distribution.log_prob(actions).sum(dim=-1)
|
||||
|
||||
def act_inference(self, obs):
|
||||
self.history = torch.cat([self.history[:, 1:], obs.unsqueeze(1)], dim=1)
|
||||
latent, _ = self.get_student_latent_and_weights(self.history.flatten(1))
|
||||
x = torch.cat([latent, obs], dim=1)
|
||||
actions_mean = self.actor(x)
|
||||
return actions_mean
|
||||
|
||||
def evaluate(self, privileged_obs, history, is_teacher, **kwargs):
|
||||
if is_teacher:
|
||||
latent = self.teacher_encoder(privileged_obs)
|
||||
else:
|
||||
latent, _ = self.get_student_latent_and_weights(history)
|
||||
x = torch.cat([latent.detach(), privileged_obs], dim=1)
|
||||
value = self.critic(x)
|
||||
return value
|
||||
|
||||
def get_student_latent_and_weights(self, history):
|
||||
B = history.shape[0]
|
||||
history_no_goal = history.reshape(B, self.history_length, -1)[:, :, self.obs_no_goal_mask].reshape(B, -1)
|
||||
return self.student_moe_encoder(history, history_no_goal)
|
||||
|
||||
class StudentMoEEncoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
expert_dim,
|
||||
gating_dim,
|
||||
hidden_dims=[512, 256],
|
||||
expert_num=8,
|
||||
expert_hidden_dim=128,
|
||||
latent_dim=32,
|
||||
activation='elu',
|
||||
norm_type='l2norm',
|
||||
):
|
||||
super().__init__()
|
||||
self.expert_num = expert_num
|
||||
self.latent_dim = latent_dim
|
||||
self.norm_layer = L2Norm() if norm_type == 'l2norm' else SimNorm()
|
||||
activation = get_activation(activation)
|
||||
|
||||
# Expert networks
|
||||
experts_layers = []
|
||||
last_dim = expert_dim
|
||||
for l in hidden_dims:
|
||||
experts_layers.append(nn.Linear(last_dim, l))
|
||||
experts_layers.append(activation)
|
||||
last_dim = l
|
||||
self.experts_backbone = nn.Sequential(*experts_layers)
|
||||
self.experts_hidden = nn.Sequential(
|
||||
nn.Linear(last_dim, expert_num * expert_hidden_dim),
|
||||
activation
|
||||
)
|
||||
self.experts_out = nn.Linear(expert_hidden_dim, latent_dim)
|
||||
|
||||
# Gating network
|
||||
gating_layers = []
|
||||
last_dim = gating_dim
|
||||
for l in hidden_dims:
|
||||
gating_layers.append(nn.Linear(last_dim, l))
|
||||
gating_layers.append(activation)
|
||||
last_dim = l
|
||||
gating_layers.append(nn.Linear(last_dim, expert_num))
|
||||
gating_layers.append(nn.Softmax(dim=-1))
|
||||
self.gating_network = nn.Sequential(*gating_layers)
|
||||
|
||||
def forward(self, obs, obs_no_goal):
|
||||
weights = self.gating_network(obs) # (batch, expert_num)
|
||||
shared_features = self.experts_backbone(obs_no_goal)
|
||||
expert_hidden = self.experts_hidden(shared_features)
|
||||
expert_hidden = expert_hidden.view(-1, self.expert_num, expert_hidden.shape[-1] // self.expert_num)
|
||||
expert_latent = self.experts_out(expert_hidden) # (batch, expert_num, latent_dim)
|
||||
latent = torch.sum(weights.unsqueeze(-1) * expert_latent, dim=1) # (batch, latent_dim)
|
||||
latent = self.norm_layer(latent)
|
||||
return latent, weights
|
||||
|
||||
def get_activation(act_name):
|
||||
if act_name == "elu":
|
||||
return nn.ELU()
|
||||
elif act_name == "selu":
|
||||
return nn.SELU()
|
||||
elif act_name == "relu":
|
||||
return nn.ReLU()
|
||||
elif act_name == "crelu":
|
||||
return nn.ReLU()
|
||||
elif act_name == "lrelu":
|
||||
return nn.LeakyReLU()
|
||||
elif act_name == "tanh":
|
||||
return nn.Tanh()
|
||||
elif act_name == "sigmoid":
|
||||
return nn.Sigmoid()
|
||||
else:
|
||||
print("invalid activation function!")
|
||||
return None
|
||||
|
||||
class L2Norm(nn.Module):
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, x):
|
||||
return F.normalize(x, p=2.0, dim=-1)
|
||||
|
||||
class SimNorm(nn.Module):
|
||||
"""
|
||||
Simplicial normalization.
|
||||
Adapted from https://arxiv.org/abs/2204.00616.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.dim = 8 # for latent dim 512
|
||||
|
||||
def forward(self, x):
|
||||
shp = x.shape
|
||||
x = x.view(*shp[:-1], -1, self.dim)
|
||||
x = F.softmax(x, dim=-1)
|
||||
return x.view(*shp)
|
||||
|
||||
def __repr__(self):
|
||||
return f"SimNorm(dim={self.dim})"
|
||||
116
rsl_rl/rsl_rl/modules/actor_critic_recurrent.py
Normal file
116
rsl_rl/rsl_rl/modules/actor_critic_recurrent.py
Normal file
@@ -0,0 +1,116 @@
|
||||
# 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 numpy as np
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.distributions import Normal
|
||||
from torch.nn.modules import rnn
|
||||
from .actor_critic import ActorCritic, get_activation
|
||||
from rsl_rl.utils import unpad_trajectories
|
||||
|
||||
class ActorCriticRecurrent(ActorCritic):
|
||||
is_recurrent = True
|
||||
def __init__(self, num_actor_obs,
|
||||
num_critic_obs,
|
||||
num_actions,
|
||||
actor_hidden_dims=[256, 256, 256],
|
||||
critic_hidden_dims=[256, 256, 256],
|
||||
activation='elu',
|
||||
rnn_type='lstm',
|
||||
rnn_hidden_size=256,
|
||||
rnn_num_layers=1,
|
||||
init_noise_std=1.0,
|
||||
**kwargs):
|
||||
if kwargs:
|
||||
print("ActorCriticRecurrent.__init__ got unexpected arguments, which will be ignored: " + str(kwargs.keys()),)
|
||||
|
||||
super().__init__(num_actor_obs=rnn_hidden_size,
|
||||
num_critic_obs=rnn_hidden_size,
|
||||
num_actions=num_actions,
|
||||
actor_hidden_dims=actor_hidden_dims,
|
||||
critic_hidden_dims=critic_hidden_dims,
|
||||
activation=activation,
|
||||
init_noise_std=init_noise_std)
|
||||
|
||||
activation = get_activation(activation)
|
||||
|
||||
self.memory_a = Memory(num_actor_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_size)
|
||||
self.memory_c = Memory(num_critic_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_size)
|
||||
|
||||
print(f"Actor RNN: {self.memory_a}")
|
||||
print(f"Critic RNN: {self.memory_c}")
|
||||
|
||||
def reset(self, dones=None):
|
||||
self.memory_a.reset(dones)
|
||||
self.memory_c.reset(dones)
|
||||
|
||||
def act(self, observations, masks=None, hidden_states=None):
|
||||
input_a = self.memory_a(observations, masks, hidden_states)
|
||||
return super().act(input_a.squeeze(0))
|
||||
|
||||
def act_inference(self, observations):
|
||||
input_a = self.memory_a(observations)
|
||||
return super().act_inference(input_a.squeeze(0))
|
||||
|
||||
def evaluate(self, critic_observations, masks=None, hidden_states=None):
|
||||
input_c = self.memory_c(critic_observations, masks, hidden_states)
|
||||
return super().evaluate(input_c.squeeze(0))
|
||||
|
||||
def get_hidden_states(self):
|
||||
return self.memory_a.hidden_states, self.memory_c.hidden_states
|
||||
|
||||
|
||||
class Memory(torch.nn.Module):
|
||||
def __init__(self, input_size, type='lstm', num_layers=1, hidden_size=256):
|
||||
super().__init__()
|
||||
# RNN
|
||||
rnn_cls = nn.GRU if type.lower() == 'gru' else nn.LSTM
|
||||
self.rnn = rnn_cls(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers)
|
||||
self.hidden_states = None
|
||||
|
||||
def forward(self, input, masks=None, hidden_states=None):
|
||||
batch_mode = masks is not None
|
||||
if batch_mode:
|
||||
# batch mode (policy update): need saved hidden states
|
||||
if hidden_states is None:
|
||||
raise ValueError("Hidden states not passed to memory module during policy update")
|
||||
out, _ = self.rnn(input, hidden_states)
|
||||
out = unpad_trajectories(out, masks)
|
||||
else:
|
||||
# inference mode (collection): use hidden states of last step
|
||||
out, self.hidden_states = self.rnn(input.unsqueeze(0), self.hidden_states)
|
||||
return out
|
||||
|
||||
def reset(self, dones=None):
|
||||
# When the RNN is an LSTM, self.hidden_states_a is a list with hidden_state and cell_state
|
||||
for hidden_state in self.hidden_states:
|
||||
hidden_state[..., dones, :] = 0.0
|
||||
32
rsl_rl/rsl_rl/runners/__init__.py
Normal file
32
rsl_rl/rsl_rl/runners/__init__.py
Normal file
@@ -0,0 +1,32 @@
|
||||
# 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 .on_policy_runner import OnPolicyRunner
|
||||
from .on_policy_runner_cts import OnPolicyRunnerCTS
|
||||
233
rsl_rl/rsl_rl/runners/on_policy_runner.py
Normal file
233
rsl_rl/rsl_rl/runners/on_policy_runner.py
Normal file
@@ -0,0 +1,233 @@
|
||||
# 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 time
|
||||
import os
|
||||
from collections import deque
|
||||
import statistics
|
||||
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
import torch
|
||||
|
||||
from rsl_rl.algorithms import PPO
|
||||
from rsl_rl.modules import ActorCritic, ActorCriticRecurrent
|
||||
from rsl_rl.env import VecEnv
|
||||
|
||||
|
||||
class OnPolicyRunner:
|
||||
|
||||
def __init__(self,
|
||||
env: VecEnv,
|
||||
train_cfg,
|
||||
log_dir=None,
|
||||
device='cpu'):
|
||||
|
||||
self.cfg=train_cfg["runner"]
|
||||
self.alg_cfg = train_cfg["algorithm"]
|
||||
self.policy_cfg = train_cfg["policy"]
|
||||
self.device = device
|
||||
self.env = env
|
||||
if self.env.num_privileged_obs is not None:
|
||||
num_critic_obs = self.env.num_privileged_obs
|
||||
else:
|
||||
num_critic_obs = self.env.num_obs
|
||||
actor_critic_class = eval(self.cfg["policy_class_name"]) # ActorCritic
|
||||
actor_critic: ActorCritic = actor_critic_class( self.env.num_obs,
|
||||
num_critic_obs,
|
||||
self.env.num_actions,
|
||||
**self.policy_cfg).to(self.device)
|
||||
alg_class = eval(self.cfg["algorithm_class_name"]) # PPO
|
||||
self.alg: PPO = alg_class(actor_critic, device=self.device, **self.alg_cfg)
|
||||
self.num_steps_per_env = self.cfg["num_steps_per_env"]
|
||||
self.save_interval = self.cfg["save_interval"]
|
||||
|
||||
# init storage and model
|
||||
self.alg.init_storage(self.env.num_envs, self.num_steps_per_env, [self.env.num_obs], [self.env.num_privileged_obs], [self.env.num_actions])
|
||||
|
||||
# Log
|
||||
self.log_dir = log_dir
|
||||
self.writer = None
|
||||
self.tot_timesteps = 0
|
||||
self.tot_time = 0
|
||||
self.current_learning_iteration = 0
|
||||
|
||||
_, _ = self.env.reset()
|
||||
|
||||
def learn(self, num_learning_iterations, init_at_random_ep_len=False):
|
||||
# initialize writer
|
||||
if self.log_dir is not None and self.writer is None:
|
||||
self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10)
|
||||
if init_at_random_ep_len:
|
||||
self.env.episode_length_buf = torch.randint_like(self.env.episode_length_buf, high=int(self.env.max_episode_length))
|
||||
obs = self.env.get_observations()
|
||||
privileged_obs = self.env.get_privileged_observations()
|
||||
critic_obs = privileged_obs if privileged_obs is not None else obs
|
||||
obs, critic_obs = obs.to(self.device), critic_obs.to(self.device)
|
||||
self.alg.actor_critic.train() # switch to train mode (for dropout for example)
|
||||
|
||||
ep_infos = []
|
||||
rewbuffer = deque(maxlen=100)
|
||||
lenbuffer = deque(maxlen=100)
|
||||
cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device)
|
||||
cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device)
|
||||
|
||||
tot_iter = self.current_learning_iteration + num_learning_iterations
|
||||
for it in range(self.current_learning_iteration, tot_iter):
|
||||
start = time.time()
|
||||
# Rollout
|
||||
with torch.inference_mode():
|
||||
for i in range(self.num_steps_per_env):
|
||||
actions = self.alg.act(obs, critic_obs)
|
||||
obs, privileged_obs, rewards, dones, infos = self.env.step(actions)
|
||||
critic_obs = privileged_obs if privileged_obs is not None else obs
|
||||
obs, critic_obs, rewards, dones = obs.to(self.device), critic_obs.to(self.device), rewards.to(self.device), dones.to(self.device)
|
||||
self.alg.process_env_step(rewards, dones, infos)
|
||||
|
||||
if self.log_dir is not None:
|
||||
# Book keeping
|
||||
if 'episode' in infos:
|
||||
ep_infos.append(infos['episode'])
|
||||
cur_reward_sum += rewards
|
||||
cur_episode_length += 1
|
||||
new_ids = (dones > 0).nonzero(as_tuple=False)
|
||||
rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist())
|
||||
lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist())
|
||||
cur_reward_sum[new_ids] = 0
|
||||
cur_episode_length[new_ids] = 0
|
||||
|
||||
stop = time.time()
|
||||
collection_time = stop - start
|
||||
|
||||
# Learning step
|
||||
start = stop
|
||||
self.alg.compute_returns(critic_obs)
|
||||
|
||||
mean_value_loss, mean_surrogate_loss = self.alg.update()
|
||||
stop = time.time()
|
||||
learn_time = stop - start
|
||||
if self.log_dir is not None:
|
||||
self.log(locals())
|
||||
if it % self.save_interval == 0:
|
||||
self.save(os.path.join(self.log_dir, 'model_{}.pt'.format(it)))
|
||||
ep_infos.clear()
|
||||
|
||||
self.current_learning_iteration += num_learning_iterations
|
||||
self.save(os.path.join(self.log_dir, 'model_{}.pt'.format(self.current_learning_iteration)))
|
||||
|
||||
def log(self, locs, width=80, pad=35):
|
||||
self.tot_timesteps += self.num_steps_per_env * self.env.num_envs
|
||||
self.tot_time += locs['collection_time'] + locs['learn_time']
|
||||
iteration_time = locs['collection_time'] + locs['learn_time']
|
||||
|
||||
ep_string = f''
|
||||
if locs['ep_infos']:
|
||||
for key in locs['ep_infos'][0]:
|
||||
infotensor = torch.tensor([], device=self.device)
|
||||
for ep_info in locs['ep_infos']:
|
||||
# handle scalar and zero dimensional tensor infos
|
||||
if not isinstance(ep_info[key], torch.Tensor):
|
||||
ep_info[key] = torch.Tensor([ep_info[key]])
|
||||
if len(ep_info[key].shape) == 0:
|
||||
ep_info[key] = ep_info[key].unsqueeze(0)
|
||||
infotensor = torch.cat((infotensor, ep_info[key].to(self.device)))
|
||||
value = torch.mean(infotensor)
|
||||
self.writer.add_scalar('Episode/' + key, value, locs['it'])
|
||||
ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n"""
|
||||
mean_std = self.alg.actor_critic.std.mean()
|
||||
fps = int(self.num_steps_per_env * self.env.num_envs / (locs['collection_time'] + locs['learn_time']))
|
||||
|
||||
self.writer.add_scalar('Loss/value_function', locs['mean_value_loss'], locs['it'])
|
||||
self.writer.add_scalar('Loss/surrogate', locs['mean_surrogate_loss'], locs['it'])
|
||||
self.writer.add_scalar('Loss/learning_rate', self.alg.learning_rate, locs['it'])
|
||||
self.writer.add_scalar('Policy/mean_noise_std', mean_std.item(), locs['it'])
|
||||
self.writer.add_scalar('Perf/total_fps', fps, locs['it'])
|
||||
self.writer.add_scalar('Perf/collection time', locs['collection_time'], locs['it'])
|
||||
self.writer.add_scalar('Perf/learning_time', locs['learn_time'], locs['it'])
|
||||
if len(locs['rewbuffer']) > 0:
|
||||
self.writer.add_scalar('Train/mean_reward', statistics.mean(locs['rewbuffer']), locs['it'])
|
||||
self.writer.add_scalar('Train/mean_episode_length', statistics.mean(locs['lenbuffer']), locs['it'])
|
||||
self.writer.add_scalar('Train/mean_reward/time', statistics.mean(locs['rewbuffer']), self.tot_time)
|
||||
self.writer.add_scalar('Train/mean_episode_length/time', statistics.mean(locs['lenbuffer']), self.tot_time)
|
||||
|
||||
str = f" \033[1m Learning iteration {locs['it']}/{self.current_learning_iteration + locs['num_learning_iterations']} \033[0m "
|
||||
|
||||
if len(locs['rewbuffer']) > 0:
|
||||
log_string = (f"""{'#' * width}\n"""
|
||||
f"""{str.center(width, ' ')}\n\n"""
|
||||
f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[
|
||||
'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n"""
|
||||
f"""{'Value function loss:':>{pad}} {locs['mean_value_loss']:.4f}\n"""
|
||||
f"""{'Surrogate loss:':>{pad}} {locs['mean_surrogate_loss']:.4f}\n"""
|
||||
f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n"""
|
||||
f"""{'Mean reward:':>{pad}} {statistics.mean(locs['rewbuffer']):.2f}\n"""
|
||||
f"""{'Mean episode length:':>{pad}} {statistics.mean(locs['lenbuffer']):.2f}\n""")
|
||||
# f"""{'Mean reward/step:':>{pad}} {locs['mean_reward']:.2f}\n"""
|
||||
# f"""{'Mean episode length/episode:':>{pad}} {locs['mean_trajectory_length']:.2f}\n""")
|
||||
else:
|
||||
log_string = (f"""{'#' * width}\n"""
|
||||
f"""{str.center(width, ' ')}\n\n"""
|
||||
f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[
|
||||
'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n"""
|
||||
f"""{'Value function loss:':>{pad}} {locs['mean_value_loss']:.4f}\n"""
|
||||
f"""{'Surrogate loss:':>{pad}} {locs['mean_surrogate_loss']:.4f}\n"""
|
||||
f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""")
|
||||
# f"""{'Mean reward/step:':>{pad}} {locs['mean_reward']:.2f}\n"""
|
||||
# f"""{'Mean episode length/episode:':>{pad}} {locs['mean_trajectory_length']:.2f}\n""")
|
||||
|
||||
log_string += ep_string
|
||||
log_string += (f"""{'-' * width}\n"""
|
||||
f"""{'Total timesteps:':>{pad}} {self.tot_timesteps}\n"""
|
||||
f"""{'Iteration time:':>{pad}} {iteration_time:.2f}s\n"""
|
||||
f"""{'Total time:':>{pad}} {self.tot_time:.2f}s\n"""
|
||||
f"""{'ETA:':>{pad}} {self.tot_time / (locs['it'] + 1) * (
|
||||
locs['num_learning_iterations'] - locs['it']):.1f}s\n""")
|
||||
print(log_string)
|
||||
|
||||
def save(self, path, infos=None):
|
||||
torch.save({
|
||||
'model_state_dict': self.alg.actor_critic.state_dict(),
|
||||
'optimizer_state_dict': self.alg.optimizer.state_dict(),
|
||||
'iter': self.current_learning_iteration,
|
||||
'infos': infos,
|
||||
}, path)
|
||||
|
||||
def load(self, path, load_optimizer=True):
|
||||
loaded_dict = torch.load(path)
|
||||
self.alg.actor_critic.load_state_dict(loaded_dict['model_state_dict'])
|
||||
if load_optimizer:
|
||||
self.alg.optimizer.load_state_dict(loaded_dict['optimizer_state_dict'])
|
||||
self.current_learning_iteration = loaded_dict['iter']
|
||||
return loaded_dict['infos']
|
||||
|
||||
def get_inference_policy(self, device=None):
|
||||
self.alg.actor_critic.eval() # switch to evaluation mode (dropout for example)
|
||||
if device is not None:
|
||||
self.alg.actor_critic.to(device)
|
||||
return self.alg.actor_critic.act_inference
|
||||
286
rsl_rl/rsl_rl/runners/on_policy_runner_cts.py
Normal file
286
rsl_rl/rsl_rl/runners/on_policy_runner_cts.py
Normal file
@@ -0,0 +1,286 @@
|
||||
# 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 time
|
||||
import os
|
||||
from collections import deque
|
||||
import statistics
|
||||
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
import torch
|
||||
|
||||
from rsl_rl.algorithms import CTS, MoECTS
|
||||
from rsl_rl.modules import ActorCriticCTS, ActorCriticMoECTS
|
||||
from rsl_rl.env import VecEnv
|
||||
|
||||
import yaml
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
from legged_gym.utils.helpers import class_to_dict
|
||||
from typing import Union
|
||||
|
||||
def numpy_representer(dumper, data):
|
||||
return dumper.represent_float(float(data))
|
||||
|
||||
def numpy_int_representer(dumper, data):
|
||||
return dumper.represent_int(int(data))
|
||||
|
||||
# Add the numpy representer to yaml
|
||||
yaml.add_representer(np.float32, numpy_representer, Dumper=yaml.SafeDumper)
|
||||
yaml.add_representer(np.float64, numpy_representer, Dumper=yaml.SafeDumper)
|
||||
yaml.add_representer(np.int32, numpy_int_representer, Dumper=yaml.SafeDumper)
|
||||
yaml.add_representer(np.int64, numpy_int_representer, Dumper=yaml.SafeDumper)
|
||||
|
||||
|
||||
class OnPolicyRunnerCTS:
|
||||
|
||||
def __init__(self,
|
||||
env: VecEnv,
|
||||
train_cfg,
|
||||
log_dir=None,
|
||||
device='cpu'):
|
||||
|
||||
self.cfg=train_cfg["runner"]
|
||||
self.alg_cfg = train_cfg["algorithm"]
|
||||
self.policy_cfg = train_cfg["policy"]
|
||||
self.device = device
|
||||
self.env = env
|
||||
if self.env.num_privileged_obs is not None:
|
||||
num_critic_obs = self.env.num_privileged_obs
|
||||
else:
|
||||
num_critic_obs = self.env.num_obs
|
||||
history_length = train_cfg["history_length"]
|
||||
actor_critic_class = eval(self.cfg["policy_class_name"])
|
||||
model: Union[ActorCriticCTS, ActorCriticMoECTS] = actor_critic_class(
|
||||
self.env.num_obs,
|
||||
num_critic_obs,
|
||||
self.env.num_actions,
|
||||
self.env.num_envs,
|
||||
history_length,
|
||||
**self.policy_cfg).to(self.device)
|
||||
alg_class = eval(self.cfg["algorithm_class_name"])
|
||||
self.alg: Union[CTS, MoECTS] = alg_class(model, self.env.num_envs, history_length, device=self.device, **self.alg_cfg)
|
||||
self.num_steps_per_env = self.cfg["num_steps_per_env"]
|
||||
self.save_interval = self.cfg["save_interval"]
|
||||
|
||||
# init storage and model
|
||||
self.alg.init_storage(self.env.num_envs, self.num_steps_per_env, [self.env.num_obs], [self.env.num_privileged_obs], [self.env.num_actions])
|
||||
|
||||
# init history
|
||||
self.history = torch.zeros((self.env.num_envs, history_length, self.env.num_obs), device=self.device)
|
||||
|
||||
# Log
|
||||
self.log_dir = log_dir
|
||||
self.writer = None
|
||||
self.tot_timesteps = 0
|
||||
self.tot_time = 0
|
||||
self.current_learning_iteration = 0
|
||||
|
||||
_, _ = self.env.reset()
|
||||
if self.log_dir is not None and self.env.cfg.env.test is False:
|
||||
Path(self.log_dir).mkdir(parents=True, exist_ok=True)
|
||||
all_cfg = {"train_cfg": train_cfg, "env_cfg": class_to_dict(self.env.cfg)}
|
||||
yaml.safe_dump(all_cfg, open(os.path.join(self.log_dir, 'config.yaml'), 'w'))
|
||||
|
||||
def learn(self, num_learning_iterations, init_at_random_ep_len=False):
|
||||
# initialize writer
|
||||
if self.log_dir is not None and self.writer is None:
|
||||
self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10)
|
||||
if init_at_random_ep_len:
|
||||
self.env.episode_length_buf = torch.randint_like(self.env.episode_length_buf, high=int(self.env.max_episode_length))
|
||||
obs = self.env.get_observations()
|
||||
privileged_obs = self.env.get_privileged_observations()
|
||||
assert privileged_obs is not None
|
||||
obs, privileged_obs = obs.to(self.device), privileged_obs.to(self.device)
|
||||
self.history = torch.cat([self.history[:, 1:], obs.unsqueeze(1)], dim=1)
|
||||
self.alg.model.train() # switch to train mode (for dropout for example)
|
||||
|
||||
ep_infos = []
|
||||
teacher_rewbuffer = deque(maxlen=100)
|
||||
teacher_lenbuffer = deque(maxlen=100)
|
||||
student_rewbuffer = deque(maxlen=100)
|
||||
student_lenbuffer = deque(maxlen=100)
|
||||
|
||||
cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device)
|
||||
cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device)
|
||||
|
||||
self.start_learning_iteration = self.current_learning_iteration
|
||||
tot_iter = self.current_learning_iteration + num_learning_iterations
|
||||
for it in range(self.current_learning_iteration, tot_iter):
|
||||
start = time.time()
|
||||
# Rollout
|
||||
with torch.inference_mode():
|
||||
for i in range(self.num_steps_per_env):
|
||||
actions = self.alg.act(obs, privileged_obs, self.history.flatten(1))
|
||||
obs, privileged_obs, rewards, dones, infos = self.env.step(actions)
|
||||
obs, privileged_obs, rewards, dones = obs.to(self.device), privileged_obs.to(self.device), rewards.to(self.device), dones.to(self.device)
|
||||
self.history[dones > 0] = 0.0
|
||||
self.history = torch.cat([self.history[:, 1:], obs.unsqueeze(1)], dim=1)
|
||||
self.alg.process_env_step(rewards, dones, infos)
|
||||
|
||||
if self.log_dir is not None:
|
||||
# Book keeping
|
||||
if 'episode' in infos:
|
||||
ep_infos.append(infos['episode'])
|
||||
cur_reward_sum += rewards
|
||||
cur_episode_length += 1
|
||||
new_ids = (dones > 0).nonzero(as_tuple=False)
|
||||
if new_ids.shape[0]:
|
||||
ti = self.alg.teacher_env_idxs
|
||||
teacher_ids = new_ids[torch.isin(new_ids, ti)]
|
||||
student_ids = new_ids[~torch.isin(new_ids, ti)]
|
||||
teacher_rewbuffer.extend(cur_reward_sum[teacher_ids].cpu().numpy().tolist())
|
||||
teacher_lenbuffer.extend(cur_episode_length[teacher_ids].cpu().numpy().tolist())
|
||||
student_rewbuffer.extend(cur_reward_sum[student_ids].cpu().numpy().tolist())
|
||||
student_lenbuffer.extend(cur_episode_length[student_ids].cpu().numpy().tolist())
|
||||
cur_reward_sum[new_ids] = 0
|
||||
cur_episode_length[new_ids] = 0
|
||||
|
||||
stop = time.time()
|
||||
collection_time = stop - start
|
||||
|
||||
# Learning step
|
||||
start = stop
|
||||
self.alg.compute_returns(privileged_obs, self.history.flatten(1))
|
||||
|
||||
if self.cfg["algorithm_class_name"] == "CTS":
|
||||
mean_value_loss, mean_surrogate_loss, mean_entropy_loss, mean_latent_loss = self.alg.update()
|
||||
elif self.cfg["algorithm_class_name"] == "MoECTS":
|
||||
mean_value_loss, mean_surrogate_loss, mean_entropy_loss, mean_latent_loss, mean_load_balance_loss = self.alg.update()
|
||||
stop = time.time()
|
||||
learn_time = stop - start
|
||||
self.current_learning_iteration += 1
|
||||
if self.log_dir is not None:
|
||||
self.log(locals())
|
||||
if it % self.save_interval == 0:
|
||||
self.save(os.path.join(self.log_dir, 'model_{}.pt'.format(it)))
|
||||
ep_infos.clear()
|
||||
|
||||
self.save(os.path.join(self.log_dir, 'model_{}.pt'.format(self.current_learning_iteration)))
|
||||
|
||||
def log(self, locs, width=80, pad=35):
|
||||
self.tot_timesteps += self.num_steps_per_env * self.env.num_envs
|
||||
self.tot_time += locs['collection_time'] + locs['learn_time']
|
||||
iteration_time = locs['collection_time'] + locs['learn_time']
|
||||
|
||||
ep_string = f''
|
||||
if locs['ep_infos']:
|
||||
for key in locs['ep_infos'][0]:
|
||||
infotensor = torch.tensor([], device=self.device)
|
||||
for ep_info in locs['ep_infos']:
|
||||
# handle scalar and zero dimensional tensor infos
|
||||
if not isinstance(ep_info[key], torch.Tensor):
|
||||
ep_info[key] = torch.Tensor([ep_info[key]])
|
||||
if len(ep_info[key].shape) == 0:
|
||||
ep_info[key] = ep_info[key].unsqueeze(0)
|
||||
infotensor = torch.cat((infotensor, ep_info[key].to(self.device)))
|
||||
value = torch.mean(infotensor)
|
||||
if 'terrain' in key:
|
||||
self.writer.add_scalar('Terrain/' + key, value, locs['it'])
|
||||
else:
|
||||
self.writer.add_scalar('Episode/' + key, value, locs['it'])
|
||||
ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n"""
|
||||
mean_std = self.alg.model.std.mean()
|
||||
fps = int(self.num_steps_per_env * self.env.num_envs / (locs['collection_time'] + locs['learn_time']))
|
||||
|
||||
self.writer.add_scalar('Loss/value_function', locs['mean_value_loss'], locs['it'])
|
||||
self.writer.add_scalar('Loss/surrogate', locs['mean_surrogate_loss'], locs['it'])
|
||||
self.writer.add_scalar('Loss/entropy', locs['mean_entropy_loss'], locs['it'])
|
||||
self.writer.add_scalar('Loss/latent', locs['mean_latent_loss'], locs['it'])
|
||||
if 'mean_load_balance_loss' in locs:
|
||||
self.writer.add_scalar('Loss/load_balance', locs['mean_load_balance_loss'], locs['it'])
|
||||
self.writer.add_scalar('Loss/learning_rate', self.alg.learning_rate, locs['it'])
|
||||
self.writer.add_scalar('Policy/mean_noise_std', mean_std.item(), locs['it'])
|
||||
self.writer.add_scalar('Perf/total_fps', fps, locs['it'])
|
||||
self.writer.add_scalar('Perf/collection time', locs['collection_time'], locs['it'])
|
||||
self.writer.add_scalar('Perf/learning_time', locs['learn_time'], locs['it'])
|
||||
if len(locs['teacher_rewbuffer']) > 0:
|
||||
self.writer.add_scalar('Train/mean_teacher_reward', statistics.mean(locs['teacher_rewbuffer']), locs['it'])
|
||||
self.writer.add_scalar('Train/mean_teacher_episode_length', statistics.mean(locs['teacher_lenbuffer']), locs['it'])
|
||||
self.writer.add_scalar('Train/mean_teacher_reward/time', statistics.mean(locs['teacher_rewbuffer']), self.tot_time)
|
||||
self.writer.add_scalar('Train/mean_teacher_episode_length/time', statistics.mean(locs['teacher_lenbuffer']), self.tot_time)
|
||||
if len(locs['student_rewbuffer']) > 0:
|
||||
self.writer.add_scalar('Train/mean_student_reward', statistics.mean(locs['student_rewbuffer']), locs['it'])
|
||||
self.writer.add_scalar('Train/mean_student_episode_length', statistics.mean(locs['student_lenbuffer']), locs['it'])
|
||||
self.writer.add_scalar('Train/mean_student_reward/time', statistics.mean(locs['student_rewbuffer']), self.tot_time)
|
||||
self.writer.add_scalar('Train/mean_student_episode_length/time', statistics.mean(locs['student_lenbuffer']), self.tot_time)
|
||||
|
||||
str = f" \033[1m Learning iteration {self.current_learning_iteration}/{locs['tot_iter']} \033[0m "
|
||||
|
||||
log_string = (f"""{'#' * width}\n"""
|
||||
f"""{str.center(width, ' ')}\n\n"""
|
||||
f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[
|
||||
'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n"""
|
||||
f"""{'Value function loss:':>{pad}} {locs['mean_value_loss']:.4f}\n"""
|
||||
f"""{'Surrogate loss:':>{pad}} {locs['mean_surrogate_loss']:.4f}\n"""
|
||||
f"""{'Entropy loss:':>{pad}} {locs['mean_entropy_loss']:.4f}\n"""
|
||||
f"""{'Latent loss:':>{pad}} {locs['mean_latent_loss']:.4f}\n""")
|
||||
if 'mean_load_balance_loss' in locs:
|
||||
log_string += f"""{'Load balance loss:':>{pad}} {locs['mean_load_balance_loss']:.4f}\n"""
|
||||
log_string += f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n"""
|
||||
if len(locs['teacher_rewbuffer']):
|
||||
log_string += (f"""{'Mean teacher reward:':>{pad}} {statistics.mean(locs['teacher_rewbuffer']):.2f}\n"""
|
||||
f"""{'Mean teacher episode length:':>{pad}} {statistics.mean(locs['teacher_lenbuffer']):.2f}\n""")
|
||||
if len(locs['student_rewbuffer']):
|
||||
log_string += (f"""{'Mean student reward:':>{pad}} {statistics.mean(locs['student_rewbuffer']):.2f}\n"""
|
||||
f"""{'Mean student episode length:':>{pad}} {statistics.mean(locs['student_lenbuffer']):.2f}\n""")
|
||||
|
||||
log_string += ep_string
|
||||
log_string += (f"""{'-' * width}\n"""
|
||||
f"""{'Total timesteps:':>{pad}} {self.tot_timesteps}\n"""
|
||||
f"""{'Iteration time:':>{pad}} {iteration_time:.2f}s\n"""
|
||||
f"""{'Total time:':>{pad}} {self.tot_time:.2f}s\n"""
|
||||
f"""{'ETA:':>{pad}} {self.tot_time / (self.current_learning_iteration - self.start_learning_iteration) * (
|
||||
locs['tot_iter'] - locs['it']):.1f}s\n""")
|
||||
print(log_string)
|
||||
|
||||
def save(self, path, infos=None):
|
||||
torch.save({
|
||||
'model_state_dict': self.alg.model.state_dict(),
|
||||
'optimizer1_state_dict': self.alg.optimizer1.state_dict(),
|
||||
'optimizer2_state_dict': self.alg.optimizer2.state_dict(),
|
||||
'iter': self.current_learning_iteration,
|
||||
'infos': infos,
|
||||
}, path)
|
||||
|
||||
def load(self, path, load_optimizer=True):
|
||||
loaded_dict = torch.load(path)
|
||||
self.alg.model.load_state_dict(loaded_dict['model_state_dict'])
|
||||
if load_optimizer:
|
||||
self.alg.optimizer1.load_state_dict(loaded_dict['optimizer1_state_dict'])
|
||||
self.alg.optimizer2.load_state_dict(loaded_dict['optimizer2_state_dict'])
|
||||
self.current_learning_iteration = loaded_dict['iter']
|
||||
return loaded_dict['infos']
|
||||
|
||||
def get_inference_policy(self, device=None):
|
||||
self.alg.model.eval() # switch to evaluation mode (dropout for example)
|
||||
if device is not None:
|
||||
self.alg.model.to(device)
|
||||
return self.alg.model.act_inference
|
||||
5
rsl_rl/rsl_rl/storage/__init__.py
Normal file
5
rsl_rl/rsl_rl/storage/__init__.py
Normal file
@@ -0,0 +1,5 @@
|
||||
# Copyright 2021 ETH Zurich, NVIDIA CORPORATION
|
||||
# SPDX-License-Identifier: BSD-3-Clause
|
||||
|
||||
from .rollout_storage import RolloutStorage
|
||||
from .rollout_storage_cts import RolloutStorageCTS
|
||||
235
rsl_rl/rsl_rl/storage/rollout_storage.py
Normal file
235
rsl_rl/rsl_rl/storage/rollout_storage.py
Normal file
@@ -0,0 +1,235 @@
|
||||
# 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 numpy as np
|
||||
|
||||
from rsl_rl.utils import split_and_pad_trajectories
|
||||
|
||||
class RolloutStorage:
|
||||
class Transition:
|
||||
def __init__(self):
|
||||
self.observations = None
|
||||
self.critic_observations = None
|
||||
self.actions = None
|
||||
self.rewards = None
|
||||
self.dones = None
|
||||
self.values = None
|
||||
self.actions_log_prob = None
|
||||
self.action_mean = None
|
||||
self.action_sigma = None
|
||||
self.hidden_states = None
|
||||
|
||||
def clear(self):
|
||||
self.__init__()
|
||||
|
||||
def __init__(self, num_envs, num_transitions_per_env, obs_shape, privileged_obs_shape, actions_shape, device='cpu'):
|
||||
|
||||
self.device = device
|
||||
|
||||
self.obs_shape = obs_shape
|
||||
self.privileged_obs_shape = privileged_obs_shape
|
||||
self.actions_shape = actions_shape
|
||||
|
||||
# Core
|
||||
self.observations = torch.zeros(num_transitions_per_env, num_envs, *obs_shape, device=self.device)
|
||||
if privileged_obs_shape[0] is not None:
|
||||
self.privileged_observations = torch.zeros(num_transitions_per_env, num_envs, *privileged_obs_shape, device=self.device)
|
||||
else:
|
||||
self.privileged_observations = None
|
||||
self.rewards = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
|
||||
self.dones = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device).byte()
|
||||
|
||||
# For PPO
|
||||
self.actions_log_prob = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.values = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.returns = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.advantages = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.mu = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
|
||||
self.sigma = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
|
||||
|
||||
self.num_transitions_per_env = num_transitions_per_env
|
||||
self.num_envs = num_envs
|
||||
|
||||
# rnn
|
||||
self.saved_hidden_states_a = None
|
||||
self.saved_hidden_states_c = None
|
||||
|
||||
self.step = 0
|
||||
|
||||
def add_transitions(self, transition: Transition):
|
||||
if self.step >= self.num_transitions_per_env:
|
||||
raise AssertionError("Rollout buffer overflow")
|
||||
self.observations[self.step].copy_(transition.observations)
|
||||
if self.privileged_observations is not None: self.privileged_observations[self.step].copy_(transition.critic_observations)
|
||||
self.actions[self.step].copy_(transition.actions)
|
||||
self.rewards[self.step].copy_(transition.rewards.view(-1, 1))
|
||||
self.dones[self.step].copy_(transition.dones.view(-1, 1))
|
||||
self.values[self.step].copy_(transition.values)
|
||||
self.actions_log_prob[self.step].copy_(transition.actions_log_prob.view(-1, 1))
|
||||
self.mu[self.step].copy_(transition.action_mean)
|
||||
self.sigma[self.step].copy_(transition.action_sigma)
|
||||
self._save_hidden_states(transition.hidden_states)
|
||||
self.step += 1
|
||||
|
||||
def _save_hidden_states(self, hidden_states):
|
||||
if hidden_states is None or hidden_states==(None, None):
|
||||
return
|
||||
# make a tuple out of GRU hidden state sto match the LSTM format
|
||||
hid_a = hidden_states[0] if isinstance(hidden_states[0], tuple) else (hidden_states[0],)
|
||||
hid_c = hidden_states[1] if isinstance(hidden_states[1], tuple) else (hidden_states[1],)
|
||||
|
||||
# initialize if needed
|
||||
if self.saved_hidden_states_a is None:
|
||||
self.saved_hidden_states_a = [torch.zeros(self.observations.shape[0], *hid_a[i].shape, device=self.device) for i in range(len(hid_a))]
|
||||
self.saved_hidden_states_c = [torch.zeros(self.observations.shape[0], *hid_c[i].shape, device=self.device) for i in range(len(hid_c))]
|
||||
# copy the states
|
||||
for i in range(len(hid_a)):
|
||||
self.saved_hidden_states_a[i][self.step].copy_(hid_a[i])
|
||||
self.saved_hidden_states_c[i][self.step].copy_(hid_c[i])
|
||||
|
||||
|
||||
def clear(self):
|
||||
self.step = 0
|
||||
|
||||
def compute_returns(self, last_values, gamma, lam):
|
||||
advantage = 0
|
||||
for step in reversed(range(self.num_transitions_per_env)):
|
||||
if step == self.num_transitions_per_env - 1:
|
||||
next_values = last_values
|
||||
else:
|
||||
next_values = self.values[step + 1]
|
||||
next_is_not_terminal = 1.0 - self.dones[step].float()
|
||||
delta = self.rewards[step] + next_is_not_terminal * gamma * next_values - self.values[step]
|
||||
advantage = delta + next_is_not_terminal * gamma * lam * advantage
|
||||
self.returns[step] = advantage + self.values[step]
|
||||
|
||||
# Compute and normalize the advantages
|
||||
self.advantages = self.returns - self.values
|
||||
self.advantages = (self.advantages - self.advantages.mean()) / (self.advantages.std() + 1e-8)
|
||||
|
||||
def get_statistics(self):
|
||||
done = self.dones
|
||||
done[-1] = 1
|
||||
flat_dones = done.permute(1, 0, 2).reshape(-1, 1)
|
||||
done_indices = torch.cat((flat_dones.new_tensor([-1], dtype=torch.int64), flat_dones.nonzero(as_tuple=False)[:, 0]))
|
||||
trajectory_lengths = (done_indices[1:] - done_indices[:-1])
|
||||
return trajectory_lengths.float().mean(), self.rewards.mean()
|
||||
|
||||
def mini_batch_generator(self, num_mini_batches, num_epochs=8):
|
||||
batch_size = self.num_envs * self.num_transitions_per_env
|
||||
mini_batch_size = batch_size // num_mini_batches
|
||||
indices = torch.randperm(num_mini_batches*mini_batch_size, requires_grad=False, device=self.device)
|
||||
|
||||
observations = self.observations.flatten(0, 1)
|
||||
if self.privileged_observations is not None:
|
||||
critic_observations = self.privileged_observations.flatten(0, 1)
|
||||
else:
|
||||
critic_observations = observations
|
||||
|
||||
actions = self.actions.flatten(0, 1)
|
||||
values = self.values.flatten(0, 1)
|
||||
returns = self.returns.flatten(0, 1)
|
||||
old_actions_log_prob = self.actions_log_prob.flatten(0, 1)
|
||||
advantages = self.advantages.flatten(0, 1)
|
||||
old_mu = self.mu.flatten(0, 1)
|
||||
old_sigma = self.sigma.flatten(0, 1)
|
||||
|
||||
for epoch in range(num_epochs):
|
||||
for i in range(num_mini_batches):
|
||||
|
||||
start = i*mini_batch_size
|
||||
end = (i+1)*mini_batch_size
|
||||
batch_idx = indices[start:end]
|
||||
|
||||
obs_batch = observations[batch_idx]
|
||||
critic_observations_batch = critic_observations[batch_idx]
|
||||
actions_batch = actions[batch_idx]
|
||||
target_values_batch = values[batch_idx]
|
||||
returns_batch = returns[batch_idx]
|
||||
old_actions_log_prob_batch = old_actions_log_prob[batch_idx]
|
||||
advantages_batch = advantages[batch_idx]
|
||||
old_mu_batch = old_mu[batch_idx]
|
||||
old_sigma_batch = old_sigma[batch_idx]
|
||||
yield obs_batch, critic_observations_batch, actions_batch, target_values_batch, advantages_batch, returns_batch, \
|
||||
old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, (None, None), None
|
||||
|
||||
# for RNNs only
|
||||
def reccurent_mini_batch_generator(self, num_mini_batches, num_epochs=8):
|
||||
|
||||
padded_obs_trajectories, trajectory_masks = split_and_pad_trajectories(self.observations, self.dones)
|
||||
if self.privileged_observations is not None:
|
||||
padded_critic_obs_trajectories, _ = split_and_pad_trajectories(self.privileged_observations, self.dones)
|
||||
else:
|
||||
padded_critic_obs_trajectories = padded_obs_trajectories
|
||||
|
||||
mini_batch_size = self.num_envs // num_mini_batches
|
||||
for ep in range(num_epochs):
|
||||
first_traj = 0
|
||||
for i in range(num_mini_batches):
|
||||
start = i*mini_batch_size
|
||||
stop = (i+1)*mini_batch_size
|
||||
|
||||
dones = self.dones.squeeze(-1)
|
||||
last_was_done = torch.zeros_like(dones, dtype=torch.bool)
|
||||
last_was_done[1:] = dones[:-1]
|
||||
last_was_done[0] = True
|
||||
trajectories_batch_size = torch.sum(last_was_done[:, start:stop])
|
||||
last_traj = first_traj + trajectories_batch_size
|
||||
|
||||
masks_batch = trajectory_masks[:, first_traj:last_traj]
|
||||
obs_batch = padded_obs_trajectories[:, first_traj:last_traj]
|
||||
critic_obs_batch = padded_critic_obs_trajectories[:, first_traj:last_traj]
|
||||
|
||||
actions_batch = self.actions[:, start:stop]
|
||||
old_mu_batch = self.mu[:, start:stop]
|
||||
old_sigma_batch = self.sigma[:, start:stop]
|
||||
returns_batch = self.returns[:, start:stop]
|
||||
advantages_batch = self.advantages[:, start:stop]
|
||||
values_batch = self.values[:, start:stop]
|
||||
old_actions_log_prob_batch = self.actions_log_prob[:, start:stop]
|
||||
|
||||
# reshape to [num_envs, time, num layers, hidden dim] (original shape: [time, num_layers, num_envs, hidden_dim])
|
||||
# then take only time steps after dones (flattens num envs and time dimensions),
|
||||
# take a batch of trajectories and finally reshape back to [num_layers, batch, hidden_dim]
|
||||
last_was_done = last_was_done.permute(1, 0)
|
||||
hid_a_batch = [ saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj].transpose(1, 0).contiguous()
|
||||
for saved_hidden_states in self.saved_hidden_states_a ]
|
||||
hid_c_batch = [ saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj].transpose(1, 0).contiguous()
|
||||
for saved_hidden_states in self.saved_hidden_states_c ]
|
||||
# remove the tuple for GRU
|
||||
hid_a_batch = hid_a_batch[0] if len(hid_a_batch)==1 else hid_a_batch
|
||||
hid_c_batch = hid_c_batch[0] if len(hid_c_batch)==1 else hid_a_batch
|
||||
|
||||
yield obs_batch, critic_obs_batch, actions_batch, values_batch, advantages_batch, returns_batch, \
|
||||
old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, (hid_a_batch, hid_c_batch), masks_batch
|
||||
|
||||
first_traj = last_traj
|
||||
216
rsl_rl/rsl_rl/storage/rollout_storage_cts.py
Normal file
216
rsl_rl/rsl_rl/storage/rollout_storage_cts.py
Normal file
@@ -0,0 +1,216 @@
|
||||
# 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 numpy as np
|
||||
|
||||
from rsl_rl.utils import split_and_pad_trajectories
|
||||
|
||||
class RolloutStorageCTS:
|
||||
class Transition:
|
||||
def __init__(self):
|
||||
self.observations = None
|
||||
self.critic_observations = None
|
||||
self.actions = None
|
||||
self.rewards = None
|
||||
self.dones = None
|
||||
self.values = None
|
||||
self.actions_log_prob = None
|
||||
self.action_mean = None
|
||||
self.action_sigma = None
|
||||
self.hidden_states = None
|
||||
self.history = None
|
||||
|
||||
def clear(self):
|
||||
self.__init__()
|
||||
|
||||
def __init__(self, num_envs, teacher_num_envs, history_length, num_transitions_per_env, obs_shape, privileged_obs_shape, actions_shape, device='cpu'):
|
||||
|
||||
self.device = device
|
||||
|
||||
self.obs_shape = obs_shape
|
||||
self.privileged_obs_shape = privileged_obs_shape
|
||||
self.actions_shape = actions_shape
|
||||
self.teacher_num_envs = teacher_num_envs
|
||||
self.student_num_envs = num_envs - teacher_num_envs
|
||||
self.history_length = history_length
|
||||
|
||||
# Core
|
||||
self.observations = torch.zeros(num_transitions_per_env, num_envs, *obs_shape, device=self.device)
|
||||
if privileged_obs_shape[0] is not None:
|
||||
self.privileged_observations = torch.zeros(num_transitions_per_env, num_envs, *privileged_obs_shape, device=self.device)
|
||||
else:
|
||||
self.privileged_observations = None
|
||||
self.rewards = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
|
||||
self.dones = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device).byte()
|
||||
self.history = torch.zeros(num_transitions_per_env, num_envs, self.history_length * obs_shape[0], device=self.device)
|
||||
|
||||
# For PPO
|
||||
self.actions_log_prob = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.values = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.returns = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.advantages = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
|
||||
self.mu = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
|
||||
self.sigma = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
|
||||
|
||||
self.num_transitions_per_env = num_transitions_per_env
|
||||
self.num_envs = num_envs
|
||||
|
||||
# rnn
|
||||
self.saved_hidden_states_a = None
|
||||
self.saved_hidden_states_c = None
|
||||
|
||||
self.step = 0
|
||||
|
||||
def add_transitions(self, transition: Transition):
|
||||
if self.step >= self.num_transitions_per_env:
|
||||
raise AssertionError("Rollout buffer overflow")
|
||||
self.observations[self.step].copy_(transition.observations)
|
||||
if self.privileged_observations is not None: self.privileged_observations[self.step].copy_(transition.critic_observations)
|
||||
self.actions[self.step].copy_(transition.actions)
|
||||
self.history[self.step].copy_(transition.history)
|
||||
self.rewards[self.step].copy_(transition.rewards.view(-1, 1))
|
||||
self.dones[self.step].copy_(transition.dones.view(-1, 1))
|
||||
self.values[self.step].copy_(transition.values)
|
||||
self.actions_log_prob[self.step].copy_(transition.actions_log_prob.view(-1, 1))
|
||||
self.mu[self.step].copy_(transition.action_mean)
|
||||
self.sigma[self.step].copy_(transition.action_sigma)
|
||||
self._save_hidden_states(transition.hidden_states)
|
||||
self.step += 1
|
||||
|
||||
def _save_hidden_states(self, hidden_states):
|
||||
if hidden_states is None or hidden_states==(None, None):
|
||||
return
|
||||
# make a tuple out of GRU hidden state sto match the LSTM format
|
||||
hid_a = hidden_states[0] if isinstance(hidden_states[0], tuple) else (hidden_states[0],)
|
||||
hid_c = hidden_states[1] if isinstance(hidden_states[1], tuple) else (hidden_states[1],)
|
||||
|
||||
# initialize if needed
|
||||
if self.saved_hidden_states_a is None:
|
||||
self.saved_hidden_states_a = [torch.zeros(self.observations.shape[0], *hid_a[i].shape, device=self.device) for i in range(len(hid_a))]
|
||||
self.saved_hidden_states_c = [torch.zeros(self.observations.shape[0], *hid_c[i].shape, device=self.device) for i in range(len(hid_c))]
|
||||
# copy the states
|
||||
for i in range(len(hid_a)):
|
||||
self.saved_hidden_states_a[i][self.step].copy_(hid_a[i])
|
||||
self.saved_hidden_states_c[i][self.step].copy_(hid_c[i])
|
||||
|
||||
|
||||
def clear(self):
|
||||
self.step = 0
|
||||
|
||||
def compute_returns(self, last_values, gamma, lam):
|
||||
advantage = 0
|
||||
for step in reversed(range(self.num_transitions_per_env)):
|
||||
if step == self.num_transitions_per_env - 1:
|
||||
next_values = last_values
|
||||
else:
|
||||
next_values = self.values[step + 1]
|
||||
next_is_not_terminal = 1.0 - self.dones[step].float()
|
||||
delta = self.rewards[step] + next_is_not_terminal * gamma * next_values - self.values[step]
|
||||
advantage = delta + next_is_not_terminal * gamma * lam * advantage
|
||||
self.returns[step] = advantage + self.values[step]
|
||||
|
||||
# Compute and normalize the advantages
|
||||
self.advantages = self.returns - self.values
|
||||
self.advantages = (self.advantages - self.advantages.mean()) / (self.advantages.std() + 1e-8)
|
||||
|
||||
def get_statistics(self):
|
||||
done = self.dones
|
||||
done[-1] = 1
|
||||
flat_dones = done.permute(1, 0, 2).reshape(-1, 1)
|
||||
done_indices = torch.cat((flat_dones.new_tensor([-1], dtype=torch.int64), flat_dones.nonzero(as_tuple=False)[:, 0]))
|
||||
trajectory_lengths = (done_indices[1:] - done_indices[:-1])
|
||||
return trajectory_lengths.float().mean(), self.rewards.mean()
|
||||
|
||||
def mini_batch_generator(self, num_mini_batches, num_epochs=8):
|
||||
teacher_samples_num = self.teacher_num_envs * self.num_transitions_per_env
|
||||
student_samples_num = self.student_num_envs * self.num_transitions_per_env
|
||||
teacher_mini_batch_size = teacher_samples_num // num_mini_batches
|
||||
student_mini_batch_size = student_samples_num // num_mini_batches
|
||||
teacher_indices = torch.randperm(teacher_samples_num, requires_grad=False, device=self.device)
|
||||
student_indices = teacher_samples_num + torch.randperm(student_samples_num, requires_grad=False, device=self.device)
|
||||
|
||||
# observations = self.observations.flatten(0, 1)
|
||||
# if self.privileged_observations is not None:
|
||||
# critic_observations = self.privileged_observations.flatten(0, 1)
|
||||
# else:
|
||||
# critic_observations = observations
|
||||
|
||||
# actions = self.actions.flatten(0, 1)
|
||||
# history = self.history.flatten(0, 1)
|
||||
# values = self.values.flatten(0, 1)
|
||||
# returns = self.returns.flatten(0, 1)
|
||||
# old_actions_log_prob = self.actions_log_prob.flatten(0, 1)
|
||||
# advantages = self.advantages.flatten(0, 1)
|
||||
# old_mu = self.mu.flatten(0, 1)
|
||||
# old_sigma = self.sigma.flatten(0, 1)
|
||||
obs_dims = list(range(2, self.observations.dim()))
|
||||
observations = self.observations.permute(1, 0, *obs_dims).flatten(0, 1)
|
||||
if self.privileged_observations is not None:
|
||||
critic_dims = list(range(2, self.privileged_observations.dim()))
|
||||
critic_observations = self.privileged_observations.permute(1, 0, *critic_dims).flatten(0, 1)
|
||||
else:
|
||||
critic_observations = observations
|
||||
|
||||
action_dims = list(range(2, self.actions.dim()))
|
||||
actions = self.actions.permute(1, 0, *action_dims).flatten(0, 1)
|
||||
old_mu = self.mu.permute(1, 0, *action_dims).flatten(0, 1)
|
||||
old_sigma = self.sigma.permute(1, 0, *action_dims).flatten(0, 1)
|
||||
hist_dims = list(range(2, self.history.dim()))
|
||||
history = self.history.permute(1, 0, *hist_dims).flatten(0, 1)
|
||||
values = self.values.permute(1, 0, 2).flatten(0, 1)
|
||||
returns = self.returns.permute(1, 0, 2).flatten(0, 1)
|
||||
old_actions_log_prob = self.actions_log_prob.permute(1, 0, 2).flatten(0, 1)
|
||||
advantages = self.advantages.permute(1, 0, 2).flatten(0, 1)
|
||||
|
||||
def get_teacher_student_samples(data, slice):
|
||||
(i1, i2), (j1, j2) = slice
|
||||
return torch.cat([data[teacher_indices[i1:i2]], data[student_indices[j1:j2]]], 0).detach()
|
||||
|
||||
for _ in range(num_epochs):
|
||||
for i in range(num_mini_batches):
|
||||
slice = (
|
||||
(i * teacher_mini_batch_size, (i+1) * teacher_mini_batch_size),
|
||||
(i * student_mini_batch_size, (i+1) * student_mini_batch_size),
|
||||
)
|
||||
|
||||
obs_batch = get_teacher_student_samples(observations, slice)
|
||||
critic_observations_batch = get_teacher_student_samples(critic_observations, slice)
|
||||
actions_batch = get_teacher_student_samples(actions, slice)
|
||||
target_values_batch = get_teacher_student_samples(values, slice)
|
||||
returns_batch = get_teacher_student_samples(returns, slice)
|
||||
old_actions_log_prob_batch = get_teacher_student_samples(old_actions_log_prob, slice)
|
||||
advantages_batch = get_teacher_student_samples(advantages, slice)
|
||||
old_mu_batch = get_teacher_student_samples(old_mu, slice)
|
||||
old_sigma_batch = get_teacher_student_samples(old_sigma, slice)
|
||||
history_batch = get_teacher_student_samples(history, slice)
|
||||
yield obs_batch, critic_observations_batch, actions_batch, history_batch, target_values_batch, advantages_batch, returns_batch, \
|
||||
old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, (None, None), None
|
||||
31
rsl_rl/rsl_rl/utils/__init__.py
Normal file
31
rsl_rl/rsl_rl/utils/__init__.py
Normal file
@@ -0,0 +1,31 @@
|
||||
# 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 .utils import split_and_pad_trajectories, unpad_trajectories
|
||||
71
rsl_rl/rsl_rl/utils/utils.py
Normal file
71
rsl_rl/rsl_rl/utils/utils.py
Normal file
@@ -0,0 +1,71 @@
|
||||
# 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
|
||||
|
||||
def split_and_pad_trajectories(tensor, dones):
|
||||
""" Splits trajectories at done indices. Then concatenates them and padds with zeros up to the length og the longest trajectory.
|
||||
Returns masks corresponding to valid parts of the trajectories
|
||||
Example:
|
||||
Input: [ [a1, a2, a3, a4 | a5, a6],
|
||||
[b1, b2 | b3, b4, b5 | b6]
|
||||
]
|
||||
|
||||
Output:[ [a1, a2, a3, a4], | [ [True, True, True, True],
|
||||
[a5, a6, 0, 0], | [True, True, False, False],
|
||||
[b1, b2, 0, 0], | [True, True, False, False],
|
||||
[b3, b4, b5, 0], | [True, True, True, False],
|
||||
[b6, 0, 0, 0] | [True, False, False, False],
|
||||
] | ]
|
||||
|
||||
Assumes that the inputy has the following dimension order: [time, number of envs, aditional dimensions]
|
||||
"""
|
||||
dones = dones.clone()
|
||||
dones[-1] = 1
|
||||
# Permute the buffers to have order (num_envs, num_transitions_per_env, ...), for correct reshaping
|
||||
flat_dones = dones.transpose(1, 0).reshape(-1, 1)
|
||||
|
||||
# Get length of trajectory by counting the number of successive not done elements
|
||||
done_indices = torch.cat((flat_dones.new_tensor([-1], dtype=torch.int64), flat_dones.nonzero()[:, 0]))
|
||||
trajectory_lengths = done_indices[1:] - done_indices[:-1]
|
||||
trajectory_lengths_list = trajectory_lengths.tolist()
|
||||
# Extract the individual trajectories
|
||||
trajectories = torch.split(tensor.transpose(1, 0).flatten(0, 1),trajectory_lengths_list)
|
||||
padded_trajectories = torch.nn.utils.rnn.pad_sequence(trajectories)
|
||||
|
||||
|
||||
trajectory_masks = trajectory_lengths > torch.arange(0, tensor.shape[0], device=tensor.device).unsqueeze(1)
|
||||
return padded_trajectories, trajectory_masks
|
||||
|
||||
def unpad_trajectories(trajectories, masks):
|
||||
""" Does the inverse operation of split_and_pad_trajectories()
|
||||
"""
|
||||
# Need to transpose before and after the masking to have proper reshaping
|
||||
return trajectories.transpose(1, 0)[masks.transpose(1, 0)].view(-1, trajectories.shape[0], trajectories.shape[-1]).transpose(1, 0)
|
||||
16
rsl_rl/setup.py
Normal file
16
rsl_rl/setup.py
Normal file
@@ -0,0 +1,16 @@
|
||||
from setuptools import setup, find_packages
|
||||
|
||||
setup(name='rsl_rl',
|
||||
version='1.0.2',
|
||||
author='Nikita Rudin',
|
||||
author_email='rudinn@ethz.ch',
|
||||
license="BSD-3-Clause",
|
||||
packages=find_packages(),
|
||||
description='Fast and simple RL algorithms implemented in pytorch',
|
||||
python_requires='>=3.6',
|
||||
install_requires=[
|
||||
"torch>=1.4.0",
|
||||
"torchvision>=0.5.0",
|
||||
"numpy>=1.16.4"
|
||||
],
|
||||
)
|
||||
21
setup.py
Normal file
21
setup.py
Normal file
@@ -0,0 +1,21 @@
|
||||
from setuptools import find_packages
|
||||
from distutils.core import setup
|
||||
|
||||
setup(name='unitree_rl_gym',
|
||||
version='1.0.0',
|
||||
author='Unitree Robotics',
|
||||
license="BSD-3-Clause",
|
||||
packages=find_packages(),
|
||||
author_email='support@unitree.com',
|
||||
description='Template RL environments for Unitree Robots',
|
||||
install_requires=[
|
||||
'isaacgym',
|
||||
'rsl-rl',
|
||||
'matplotlib',
|
||||
'numpy==1.20',
|
||||
'tensorboard',
|
||||
'mujoco==3.2.3',
|
||||
'pyyaml',
|
||||
'onnx==1.17.0',
|
||||
'pygame'
|
||||
])
|
||||
Reference in New Issue
Block a user