cyy/go2_rl_gym
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

v1.1.1-rc2; add README

Browse Source
This commit is contained in:
wty-yy
2026-01-27 21:36:44 +08:00
parent 064d0c90ea
commit 93c815b2d9
10 changed files with 468 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files

150
README.md
View File

@@ -1,152 +1,172 @@
<div align="center"> <div align="center">
<h1 align="center">Go2 RL GYM</h1> <h1 align="center">Go2 RL GYM</h1>
<p align="center"> <p align="center">
<a href="README_en.md">🌎 English</a> | <span>🇨🇳 中文</span> <span>🌎 English</span> | <a href="README_zh.md">🇨🇳 中文</a>
</p> </p>
</div> </div>
<p align="center"> <p align="center">
<strong>本仓库基于<a href="https://github.com/unitreerobotics/unitree_rl_gym">unitree_rl_gym</a>，使用强化学习训练Go2机器狗。</strong> <strong>This repository builds on <a href="https://github.com/unitreerobotics/unitree_rl_gym">unitree_rl_gym</a> to train the Unitree Go2 quadruped with reinforcement learning.</strong>
</p> </p>
<div align="center"> <div align="center">
| <div align="center"> Isaac Gym </div> | <div align="center"> Mujoco </div> | <div align="center"> Physical </div> | | <div align="center"> Isaac Gym </div> | <div align="center"> Mujoco </div> | <div align="center"> Physical </div> |
|--- | --- | --- | |--- | --- | --- |
| TODO | TODO | TODO | | ![isaacgym eval](https://raw.githubusercontent.com/robogauge/picture-bed/refs/heads/main/go2_rl_gym/isaacgym_eval.gif) | ![mujoco eval](https://raw.githubusercontent.com/robogauge/picture-bed/refs/heads/main/go2_rl_gym/mujoco_eval.gif) | ![real eval](https://raw.githubusercontent.com/robogauge/picture-bed/refs/heads/main/go2_rl_gym/real_eval.gif) |
</div> </div>
## 📦 安装配置 ## 📦 Installation
安装和配置步骤请参考 [setup.md](/doc/setup_zh.md) Follow the step-by-step setup guide in [setup.md](doc/setup_en.md).
## 🛠️ 使用指南 ## 🛠️ Usage Guide
### 1. 训练 ### 1. Train
运行以下命令进行训练： Run the following command to launch training:
```bash ```bash
python legged_gym/scripts/train.py --task=xxx python legged_gym/scripts/train.py --task=xxx
``` ```
#### ⚙️ 参数说明 #### ⚙️ Arguments
- `--task`: 必选参数，值可选(go2, go2_cts, go2_moe_cts) - `--task`: Required. Options include `go2`, `go2_cts`, `go2_moe_cts`, `go2_moe_ng_cts`, `go2_mcp_cts`, `go2_ac_moe_cts`, `go2_dual_moe_cts`; `go2_moe_cts` is the paper's final version.
- `--headless`: 默认启动图形界面，设为 true 时不渲染图形界面（效率更高） - `--headless`: Render viewer by default; set to `true` to disable rendering for higher throughput.
- `--resume`: 从日志中选择 checkpoint 继续训练 - `--resume`: Resume training from a chosen checkpoint in the logs.
- `--experiment_name`: 运行/加载的 experiment 名称 - `--experiment_name`: Experiment folder to save/load from.
- `--run_name`: 运行/加载的 run 名称 - `--run_name`: Run subfolder name to save/load from.
- `--load_run`: 加载运行的名称，默认加载最后一次运行 - `--load_run`: Name of the run to load (defaults to the most recent run).
- `--checkpoint`: checkpoint 编号，默认加载最新一次文件 - `--checkpoint`: Checkpoint index to load (defaults to the latest file).
- `--num_envs`: 并行训练的环境个数 - `--num_envs`: Number of parallel simulated environments.
- `--seed`: 随机种子 - `--seed`: Random seed.
- `--max_iterations`: 训练的最大迭代次数 - `--max_iterations`: Maximum training iterations.
- `--sim_device`: 仿真计算设备，指定 CPU 为 `--sim_device=cpu` - `--sim_device`: Physics simulation device. Use `--sim_device=cpu` to force CPU.
- `--rl_device`: 强化学习计算设备，指定 CPU 为 `--rl_device=cpu` - `--rl_device`: RL computation device. Use `--rl_device=cpu` to force CPU.
- `--robogauge`: Enable RoboGauge evaluation tool; disabled by default. Evaluation results are saved as `results_{it}.yaml` in `logs/{exp_name}/{date}/robogauge_results` and logged to TensorBoard.
- `--robogauge_port`: RoboGauge server port; default is 9973.
**默认保存训练结果**：`logs/<experiment_name>/<date_time>_<run_name>/model_<iteration>.pt` > RoboGauge evaluation requires a separate server to be started. Refer to the [RoboGauge documentation](https://github.com/wty-yy/RoboGauge).
**Default checkpoint path**: `logs/<experiment_name>/<date_time>_<run_name>/model_<iteration>.pt`
--- ---
### 2. Play ### 2. Play
如果想要在 Gym 中查看训练效果，可以运行以下命令： Visualize policies inside Gym with:
```bash ```bash
python legged_gym/scripts/play.py --task=xxx python legged_gym/scripts/play.py --task=xxx
``` ```
**说明**： **Notes**
- Play 启动参数为随机地形，难度在7到9之间。 - Play launches on randomized terrain with difficulty between 7 and 9.
- 默认加载实验文件夹最新训练的一个模型。 - It automatically loads the latest checkpoint inside the experiment folder.
- 可通过 `experiment_name` 和 `checkpoint` 指定其他模型，例如 - Override via `experiment_name` and `checkpoint`, for example:
```bash ```bash
python legged_gym/scripts/play.py --task=go2_cts --num_envs 100 --experiment_name go2_cts_hard_terrain --checkpoint 100000 python legged_gym/scripts/play.py --task=go2_cts --num_envs 100 --experiment_name go2_cts_hard_terrain --checkpoint 100000
``` ```
#### 💾 导出网络 #### 💾 Policy Export
Play 会导出 Actor 网络，保存于 `logs/{experiment_name}/exported/policies` 中： Play exports the Actor network to `logs/{experiment_name}/exported/policies`:
- `policy.pt`: torch script模型，用于Sim2Sim。 - `policy.pt`: TorchScript model for Sim2Sim.
- `policy.onnx`: onnx模型，用于Sim2Real。 - `policy.onnx`: ONNX model for Sim2Real.
- `policy.pkl`: 模型权重。 - `policy.pkl`: Raw weights.
#### Play 效果 #### Demonstration
| Go2 | ![isaacgym play](https://raw.githubusercontent.com/robogauge/picture-bed/refs/heads/main/go2_rl_gym/isaacgym_play.gif)
|--- |
| TODO |
--- ---
### 3. Sim2Sim (Mujoco) ### 3. Sim2Sim (Mujoco)
支持在 Mujoco 仿真器中运行 Sim2Sim： Run policies in the Mujoco simulator:
```bash ```bash
python deploy/deploy_mujoco/deploy_go2.py python deploy/deploy_mujoco/deploy_go2.py
``` ```
如果有xbox协议的手柄接入主机，自动切换为手柄控制，否则只会保持默认指令前进。 Connect an Xbox-compatible gamepad to enable teleoperation; otherwise, the agent keeps a default forward command.
- 替换网络模型：默认模型位于 `deploy/pre_train/go2/go2_cts_150k.pt`；自己训练模型保存于`logs/{experiment_name}/exported/policies/policy.pt`，只需替换 yaml 配置文件中 `policy_path`。 - **Swap the policy**: The default checkpoint is `deploy/pre_train/go2/go2_cts_150k.pt`. Replace `policy_path` in the YAML config with your own `logs/{experiment_name}/exported/policies/policy.pt`.
- 替换环境地形：默认地形为 `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)。 - **Swap terrains**: Default terrain is `resources/robots/go2/stairs.xml`. Alternatives include `flat.xml` and `race_track.xml`. Generate new terrains with [terrain_generator.py](resources/robots/go2/terrain_generator.py) (see also [unitree_mujoco/terrain_tool](https://github.com/unitreerobotics/unitree_mujoco/tree/main/terrain_tool)).
#### 运行效果 #### Results
| 平地 | 台阶 | 赛道 | | Flat | Stairs | Race Track |
|--- | --- | --- | |--- | --- | --- |
| | | | | ![eval flat](https://raw.githubusercontent.com/robogauge/picture-bed/refs/heads/main/go2_rl_gym/mujoco_eval_flat.gif) | ![eval stairs](https://raw.githubusercontent.com/robogauge/picture-bed/refs/heads/main/go2_rl_gym/mujoco_eval.gif) | ![eval track](https://raw.githubusercontent.com/robogauge/picture-bed/refs/heads/main/go2_rl_gym/mujoco_eval_track.gif) |
--- ---
### 4. Sim2Real ### 4. Sim2Real
#### 4.1 Python实物部署 （需要安装 unitree_sdk2_python） #### 4.1 Python Deployment (requires [unitree_sdk2_python](https://github.com/unitreerobotics/unitree_sdk2_python))
先用app进入设备→服务状态→点击运控服务，关闭`mcf`，打开`ota_box`服务。 ```bash
# Onboard Jetson: pick Python by JetPack version
# JetPack 6: Python 3.10
# JetPack 5: Python 3.8
conda create -n deploy python=3.10
conda activate deploy
# Install the matching PyTorch wheel for your Jetson
# https://forums.developer.nvidia.com/t/pytorch-for-jetson/72048
git clone https://github.com/unitreerobotics/unitree_sdk2_python.git
cd unitree_sdk2_python
pip3 install -e .
```
In the Unitree app, open Device → Service, disable `mcf/*`, and enable the `ota_box` service.
Assuming the interface to the low-level controller is `eth0`:
假设和下位机连接的网卡名称为`eth0`，执行
```bash ```bash
cd deploy/deploy_real cd deploy/deploy_real
python deploy_real_go2.py eth0 python deploy_real_go2.py eth0
``` ```
`start`站立，`A`启动控制
#### 4.2 C++实物部署（需要安装 unitree_cpp_deploy） Press `start` to stand and `A` to engage the controller.
参考[unitree_cpp_deploy](https://github.com/wty-yy-mini/unitree_cpp_deploy)使用说明。 #### 4.2 C++ Deployment (requires unitree_cpp_deploy)
#### 运行效果 Follow the usage described in [unitree_cpp_deploy](https://github.com/wty-yy-mini/unitree_cpp_deploy).
| Python部署 | C++部署 | #### Demonstration
| Python Deploy | C++ Deploy |
| --- | --- | | --- | --- |
| | | | ![python deploy](https://raw.githubusercontent.com/robogauge/picture-bed/refs/heads/main/deploy/py_deploy_with_commands.gif) | ![cpp deploy](https://raw.githubusercontent.com/robogauge/picture-bed/refs/heads/main/deploy/cpp_deploy_with_commands.gif) |
--- ---
## 🎉 致谢 ## 🎉 Acknowledgements
本仓库开发离不开以下开源项目的支持与贡献，特此感谢： This repository would not exist without the following open-source projects:
- [unitree_rl_gym](https://github.com/unitreerobotics/unitree_rl_gym)：宇树机器人强化学习训练基础框架。 - [unitree_rl_gym](https://github.com/unitreerobotics/unitree_rl_gym): Unitree's core RL training framework.
- [legged\_gym](https://github.com/leggedrobotics/legged_gym)：构建基础训练环境。 - [legged_gym](https://github.com/leggedrobotics/legged_gym): Base locomotion environment.
- [rsl\_rl](https://github.com/leggedrobotics/rsl_rl.git)：强化学习算法实现。 - [rsl_rl](https://github.com/leggedrobotics/rsl_rl.git): Reinforcement learning algorithms.
- [mujoco](https://github.com/google-deepmind/mujoco.git)：提供强大CPU仿真功能。 - [mujoco](https://github.com/google-deepmind/mujoco.git): High-performance CPU physics simulator.
- [unitree\_sdk2\_python](https://github.com/unitreerobotics/unitree_sdk2_python.git)：实物部署硬件Python通信接口。 - [unitree_sdk2_python](https://github.com/unitreerobotics/unitree_sdk2_python.git): Python hardware interface for deployment.
- [unitree_sdk2](https://github.com/unitreerobotics/unitree_sdk2)：实物部署硬件C++通信接口。 - [unitree_sdk2](https://github.com/unitreerobotics/unitree_sdk2): C++ hardware interface for deployment.
本仓库实现包含以下论文，特此感谢： Related publications implemented in this repo:
- [CTS: Concurrent Teacher-Student Reinforcement Learning for Legged Locomotion](https://arxiv.org/pdf/2405.10830) - [CTS: Concurrent Teacher-Student Reinforcement Learning for Legged Locomotion](https://arxiv.org/pdf/2405.10830)
Contributors:
- [@windigal](https://github.com/windigal): CTS algorithm reproduction, video editing
- [@wertyuilife2](https://github.com/wertyuilife2): CTS algorithm reproduction
--- ---
## 🔖 许可证 ## 🔖 License
新增内容根据 [MIT License](./LICENSE) 授权，原仓库unitree_rl_gym根据 [BSD 3-Clause License](./LICENSE) 授权。 New contributions follow the [MIT License](LICENSE); the original unitree_rl_gym remains under the [BSD 3-Clause License](LICENSE).
详情请阅读完整 [LICENSE 文件](./LICENSE)。
See the complete [LICENSE file](LICENSE) for details.

11
README_en.md
View File

@@ -1,11 +0,0 @@
<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>

171
README_zh.md Normal file
View File

@@ -0,0 +1,171 @@
<div align="center">
<h1 align="center">Go2 RL GYM</h1>
<p align="center">
<a href="README.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> |
|--- | --- | --- |
| ![isaacgym eval](https://raw.githubusercontent.com/robogauge/picture-bed/refs/heads/main/go2_rl_gym/isaacgym_eval.gif) | ![mujoco eval](https://raw.githubusercontent.com/robogauge/picture-bed/refs/heads/main/go2_rl_gym/mujoco_eval.gif) | ![real eval](https://raw.githubusercontent.com/robogauge/picture-bed/refs/heads/main/go2_rl_gym/real_eval.gif) |
</div>
## 📦 安装配置
安装和配置步骤请参考 [setup.md](/doc/setup_zh.md)
## 🛠️ 使用指南
### 1. 训练
运行以下命令进行训练：
```bash
python legged_gym/scripts/train.py --task=xxx
```
#### ⚙️ 参数说明
- `--task`: 必选参数，值可选(go2, go2_cts, go2_moe_cts, go2_moe_ng_cts, go2_mcp_cts, go2_ac_moe_cts, go2_dual_moe_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`
- `--robogauge`: 是否启用 RoboGauge 评估工具，默认关闭，评估结果会以 `results_{it}.yaml` 保存在 `logs/{exp_name}/{date}/robogauge_results` 下，并记录在 TensorBoard 中
- `--robogauge_port`: RoboGauge 服务端端口，默认 9973
> RoboGauge 评估还需单独启动服务端，使用方法参考 [RoboGauge 文档](https://github.com/wty-yy/RoboGauge)
**默认保存训练结果**：`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 效果
![isaacgym play](https://raw.githubusercontent.com/robogauge/picture-bed/refs/heads/main/go2_rl_gym/isaacgym_play.gif)
---
### 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)。
#### 运行效果
| 平地 | 台阶 | 赛道 |
|--- | --- | --- |
| ![eval flat](https://raw.githubusercontent.com/robogauge/picture-bed/refs/heads/main/go2_rl_gym/mujoco_eval_flat.gif) | ![eval stairs](https://raw.githubusercontent.com/robogauge/picture-bed/refs/heads/main/go2_rl_gym/mujoco_eval.gif) | ![eval track](https://raw.githubusercontent.com/robogauge/picture-bed/refs/heads/main/go2_rl_gym/mujoco_eval_track.gif) |
---
### 4. Sim2Real
#### 4.1 Python实物部署 （需要安装 [unitree_sdk2_python](https://github.com/unitreerobotics/unitree_sdk2_python)）
```bash
# 如果机载电脑部署，根据Jetson版本选择Python版本
# JetPack 6: Python 3.10
# JetPack 5: Python 3.8
conda create -n deploy python=3.10
conda activate deploy
# 下载并安装对应Jetson设备和Python的PyTorch whl包
# https://forums.developer.nvidia.com/t/pytorch-for-jetson/72048
git clone https://github.com/unitreerobotics/unitree_sdk2_python.git
cd unitree_sdk2_python
pip3 install -e .
```
先用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++部署 |
| --- | --- |
| ![python deploy](https://raw.githubusercontent.com/robogauge/picture-bed/refs/heads/main/deploy/py_deploy_with_commands.gif) | ![cpp deploy](https://raw.githubusercontent.com/robogauge/picture-bed/refs/heads/main/deploy/cpp_deploy_with_commands.gif) |
---
## 🎉 致谢
本仓库开发离不开以下开源项目的支持与贡献，特此感谢：
- [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)
贡献者：
- [@windigal](https://github.com/windigal)：复现CTS算法，剪辑视频
- [@wertyuilife2](https://github.com/wertyuilife2)：复现CTS算法
---
## 🔖 许可证
新增内容根据 [MIT License](./LICENSE) 授权，原仓库unitree_rl_gym根据 [BSD 3-Clause License](./LICENSE) 授权。
详情请阅读完整 [LICENSE 文件](./LICENSE)。

25
deploy/deploy_mujoco/deploy_go2.py
View File

@@ -1,3 +1,10 @@
import sys
from pathlib import Path
PATH_PARENT = Path(__file__).parent
sys.path.append(str(PATH_PARENT))
from utils import MujocoRenderUtils
import os
import time import time
import mujoco.viewer import mujoco.viewer
import mujoco import mujoco
@@ -7,12 +14,10 @@ import torch
import yaml import yaml
import os import os
import imageio import imageio
from pathlib import Path
from argparse import ArgumentParser from argparse import ArgumentParser
import pygame import pygame
from matplotlib import pyplot as plt from matplotlib import pyplot as plt
def get_gravity_orientation(quaternion): def get_gravity_orientation(quaternion):
qw = quaternion[0] qw = quaternion[0]
qx = quaternion[1] qx = quaternion[1]
@@ -110,7 +115,7 @@ if __name__ == "__main__":
idx_model2mj = [model_joint_names.index(joint) for joint in mujoco_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] idx_mj2model = [mujoco_joint_names.index(joint) for joint in model_joint_names]
video_save_dir = str(Path(__file__).parent / "videos") video_save_dir = str(PATH_PARENT / "videos")
os.makedirs(video_save_dir, exist_ok=True) os.makedirs(video_save_dir, exist_ok=True)
model_name = os.path.basename(policy_path).split('.')[0] model_name = os.path.basename(policy_path).split('.')[0]
@@ -134,17 +139,18 @@ if __name__ == "__main__":
# load policy # load policy
policy = torch.jit.load(policy_path) policy = torch.jit.load(policy_path)
video_fps = 50
if save_video: if save_video:
video_filename = f"{model_name}_{cmd_str}.mp4" video_filename = f"{model_name}_{cmd_str}.mp4"
video_path = os.path.join(video_save_dir, video_filename) video_path = os.path.join(video_save_dir, video_filename)
print(f"Video recording will be saved to: {video_path}") print(f"Video recording will be saved to: {video_path}")
video_fps = 50
sim_fps = 1.0 / m.opt.timestep sim_fps = 1.0 / m.opt.timestep
frame_skip = int(sim_fps / video_fps) frame_skip = int(sim_fps / video_fps)
if frame_skip < 1: if frame_skip < 1:
frame_skip = 1 frame_skip = 1
writer = imageio.get_writer(video_path, fps=video_fps) 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}") print(f"Sim FPS: {sim_fps:.2f}, Video FPS: {video_fps}, Frame Skip: {frame_skip}, Save at: {video_path}")
mujoco_render_utils = MujocoRenderUtils(video_fps, m.opt.timestep)
if visualize_moe_weights: if visualize_moe_weights:
plt.ion() plt.ion()
@@ -153,7 +159,7 @@ if __name__ == "__main__":
bars = None bars = None
if save_moe_latent: if save_moe_latent:
latent_save_dir = str(Path(__file__).parent / "data_latents") latent_save_dir = str(PATH_PARENT / "data_latents")
os.makedirs(latent_save_dir, exist_ok=True) os.makedirs(latent_save_dir, exist_ok=True)
latent_filename = f"{model_name}_{cmd_str}_latents.npy" latent_filename = f"{model_name}_{cmd_str}_latents.npy"
latent_path = os.path.join(latent_save_dir, latent_filename) latent_path = os.path.join(latent_save_dir, latent_filename)
@@ -164,9 +170,9 @@ if __name__ == "__main__":
# set viewer.camera to follow robot # set viewer.camera to follow robot
viewer.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING viewer.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
viewer.cam.trackbodyid = 1 viewer.cam.trackbodyid = 1
viewer.cam.distance = 3.0 viewer.cam.distance = 2.0
viewer.cam.elevation = -30.0 viewer.cam.elevation = -20.0
viewer.cam.azimuth = 0.0 viewer.cam.azimuth = 60.0
# Close the viewer automatically after simulation_duration wall-seconds. # Close the viewer automatically after simulation_duration wall-seconds.
start = time.time() start = time.time()
@@ -188,10 +194,12 @@ if __name__ == "__main__":
# mj_step can be replaced with code that also evaluates # mj_step can be replaced with code that also evaluates
# a policy and applies a control signal before stepping the physics. # a policy and applies a control signal before stepping the physics.
mujoco.mj_step(m, d) mujoco.mj_step(m, d)
mujoco_render_utils.update(cmd, d)
if save_video and counter % frame_skip == 0: if save_video and counter % frame_skip == 0:
try: try:
renderer.update_scene(d, camera=viewer.cam) renderer.update_scene(d, camera=viewer.cam)
mujoco_render_utils.update_external_rendering(renderer, ctype='renderer')
frame = renderer.render() frame = renderer.render()
writer.append_data(frame) writer.append_data(frame)
except Exception as e: except Exception as e:
@@ -249,6 +257,7 @@ if __name__ == "__main__":
target_dof_pos = action * action_scale + default_angles target_dof_pos = action * action_scale + default_angles
# Pick up changes to the physics state, apply perturbations, update options from GUI. # Pick up changes to the physics state, apply perturbations, update options from GUI.
mujoco_render_utils.update_external_rendering(viewer, ctype='viewer')
viewer.sync() viewer.sync()
# Rudimentary time keeping, will drift relative to wall clock. # Rudimentary time keeping, will drift relative to wall clock.

112
deploy/deploy_mujoco/utils.py Normal file
View File

@@ -0,0 +1,112 @@
from typing import Union, Literal
import numpy as np
import mujoco
import mujoco.viewer
class MujocoRenderUtils:
def __init__(self, render_fps, sim_dt):
self.target_velocity = None
self.vis_smooth_factor = 1.0
self.ren_smooth_factor = 1.0
self.vis_cur_vel = np.zeros(3)
self.ren_cur_vel = np.zeros(3)
self.mj_data = None
def update(self, target_velocity, mj_data):
self.target_velocity = target_velocity
self.mj_data = mj_data
def update_external_rendering(self,
handle: Union[mujoco.viewer.Handle, mujoco.Renderer],
ctype: Literal['viewer', 'renderer'],
):
""" Update external rendering handle (viewer or renderer). """
def add_thick_arrow(geom_elem, pos, vec, rgba, scale=0.7):
vel_norm = np.linalg.norm(vec)
display_norm = min(vel_norm * scale, 1.0)
if display_norm < 0.10:
mujoco.mjv_initGeom(
geom_elem,
type=mujoco.mjtGeom.mjGEOM_NONE,
size=[0,0,0], pos=pos, mat=np.eye(3).flatten(), rgba=[0,0,0,0]
)
return
mat = np.zeros(9)
target_quat = np.zeros(4)
vec_normalized = vec / vel_norm
mujoco.mju_quatZ2Vec(target_quat, vec_normalized)
mujoco.mju_quat2Mat(mat, target_quat)
mat = mat.reshape(3, 3)
mat[:, 2] *= display_norm
mujoco.mjv_initGeom(
geom_elem,
type=mujoco.mjtGeom.mjGEOM_ARROW,
size=[0.02, 0.02, display_norm], # [height, width, length]
pos=pos,
mat=mat.flatten(),
rgba=rgba
)
viewer_geom_idx = 0
if ctype == 'viewer':
handle.user_scn.ngeom = 0 # reset user scene geometry
if self.target_velocity is not None:
base_pos_world = self.mj_data.qpos[:3]
base_quat = self.mj_data.qpos[3:7]
# rendering arrows start position
offset_body = np.array([0.0, 0.0, 0.2])
offset_world = np.zeros(3)
mujoco.mju_rotVecQuat(offset_world, offset_body, base_quat)
start_pos = base_pos_world + offset_world
tgt_vel_body = np.array([self.target_velocity[0], self.target_velocity[1], 0.0])
raw_cur_vel_world = self.mj_data.qvel[:3]
raw_cur_vel = np.zeros(3)
neg_quat = np.zeros(4)
mujoco.mju_negQuat(neg_quat, base_quat)
mujoco.mju_rotVecQuat(raw_cur_vel, raw_cur_vel_world, neg_quat)
cur_vel_body = np.array([raw_cur_vel[0], raw_cur_vel[1], 0.0])
# EMA: v_smooth = alpha * v_new + (1 - alpha) * v_old
# alpha = self.vis_smooth_factor if ctype == 'viewer' else self.ren_smooth_factor
self.vis_cur_vel = cur_vel_body
self.ren_cur_vel = cur_vel_body
tgt_vel_world = np.zeros(3)
cur_vel_world = np.zeros(3)
mujoco.mju_rotVecQuat(tgt_vel_world, tgt_vel_body, base_quat)
if ctype == 'viewer':
mujoco.mju_rotVecQuat(cur_vel_world, self.vis_cur_vel, base_quat)
else:
mujoco.mju_rotVecQuat(cur_vel_world, self.ren_cur_vel, base_quat)
COLOR_CMD = [0, 1, 0, 1] # Green 0x00ff00
COLOR_REAL = [0, 0, 1, 1] # Blue 0x0000ff
if ctype == 'viewer':
# Cmd Arrow
add_thick_arrow(handle.user_scn.geoms[viewer_geom_idx], start_pos, tgt_vel_world, COLOR_CMD)
viewer_geom_idx += 1
# Real Arrow
add_thick_arrow(handle.user_scn.geoms[viewer_geom_idx], start_pos, cur_vel_world, COLOR_REAL)
viewer_geom_idx += 1
else:
# Renderer Append
handle.scene.ngeom += 1
add_thick_arrow(handle.scene.geoms[handle.scene.ngeom - 1], start_pos, tgt_vel_world, COLOR_CMD)
handle.scene.ngeom += 1
add_thick_arrow(handle.scene.geoms[handle.scene.ngeom - 1], start_pos, cur_vel_world, COLOR_REAL)
if ctype == 'viewer':
handle.user_scn.ngeom = viewer_geom_idx

3
deploy/deploy_real/config_go2.py
View File

@@ -1,4 +1,5 @@
from legged_gym import LEGGED_GYM_ROOT_DIR from pathlib import Path
LEGGED_GYM_ROOT_DIR = str(Path(__file__).parents[2])
import numpy as np import numpy as np
import yaml import yaml

2
deploy/deploy_real/configs/go2.yaml
View File

@@ -6,7 +6,7 @@ imu_type: "torso" # "torso" or "pelvis"
lowcmd_topic: "rt/lowcmd" lowcmd_topic: "rt/lowcmd"
lowstate_topic: "rt/lowstate" lowstate_topic: "rt/lowstate"
policy_path: "{LEGGED_GYM_ROOT_DIR}/deploy/pre_train/go2/go2_cts_150k.pt" policy_path: "{LEGGED_GYM_ROOT_DIR}/deploy/pre_train/go2/go2_moe_cts_137000_0.6365.pt"
joint2motor_idx: [3,4,5,0,1,2,9,10,11,6,7,8] joint2motor_idx: [3,4,5,0,1,2,9,10,11,6,7,8]

16
deploy/deploy_real/deploy_real_go2.py
View File

@@ -1,4 +1,5 @@
from legged_gym import LEGGED_GYM_ROOT_DIR from pathlib import Path
LEGGED_GYM_ROOT_DIR = str(Path(__file__).parents[2])
import numpy as np import numpy as np
import time import time
import torch import torch
@@ -85,11 +86,13 @@ class Controller:
def zero_torque_state(self): def zero_torque_state(self):
print("Enter zero torque state.") print("Enter zero torque state.")
print("Waiting for the start signal...") print("Waiting for the *start* signal...")
while self.remote_controller.button[KeyMap.start] != 1: while self.remote_controller.button[KeyMap.start] != 1:
create_zero_cmd(self.low_cmd) create_zero_cmd(self.low_cmd)
self.send_cmd(self.low_cmd) self.send_cmd(self.low_cmd)
time.sleep(self.config.control_dt) time.sleep(self.config.control_dt)
print("Start signal received.")
print("Press *select* button to exit.")
def move_to_default_pos(self): def move_to_default_pos(self):
@@ -121,7 +124,7 @@ class Controller:
def default_pos_state(self): def default_pos_state(self):
print("Enter default pos state.") print("Enter default pos state.")
print("Waiting for the Button A signal...") print("Waiting for the *Button A* signal...")
while self.remote_controller.button[KeyMap.A] != 1: while self.remote_controller.button[KeyMap.A] != 1:
for i in range(12): for i in range(12):
motor_idx = self.config.joint2motor_idx[i] motor_idx = self.config.joint2motor_idx[i]
@@ -162,7 +165,12 @@ class Controller:
self.obs[33:45] = self.action self.obs[33:45] = self.action
obs_tensor = torch.from_numpy(self.obs).unsqueeze(0) obs_tensor = torch.from_numpy(self.obs).unsqueeze(0)
self.action = self.policy(obs_tensor).detach().numpy().squeeze() results = self.policy(obs_tensor)
if isinstance(results, tuple):
self.action = results[0]
else:
self.action = results
self.action = self.action.detach().numpy().squeeze()
target_dof_pos = self.config.default_angles + self.action * self.config.action_scale target_dof_pos = self.config.default_angles + self.action * self.config.action_scale
# target_dof_pos = self.config.default_angles # target_dof_pos = self.config.default_angles

80
doc/setup_en.md
View File

@@ -1,20 +1,20 @@
# Installation Guide # Installation and Configuration Guide
## System Requirements ## System Requirements
- **Operating System**: Recommended Ubuntu 18.04 or later - **OS**: Ubuntu 18.04 or higher is recommended
- **GPU**: Nvidia GPU - **GPU**: Nvidia GPU
- **Driver Version**: Recommended version 525 or later - **Driver Version**: Version 525 or higher is recommended
--- ---
## 1. Creating a Virtual Environment ## 1. Create 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. It is recommended to run training or deployment programs within a virtual environment. Conda is recommended for creating and managing virtual environments. If Conda is already installed on your system, you can skip step 1.1.
### 1.1 Download and Install MiniConda ### 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: MiniConda is a lightweight distribution of Conda suitable for creating and managing virtual environments. Use the following commands to download and install:
```bash ```bash
mkdir -p ~/miniconda3 mkdir -p ~/miniconda3
@@ -30,7 +30,7 @@ After installation, initialize Conda:
source ~/.bashrc source ~/.bashrc
``` ```
### 1.2 Create a New Environment ### 1.2 Create New Environment
Use the following command to create a virtual environment: Use the following command to create a virtual environment:
@@ -38,7 +38,7 @@ Use the following command to create a virtual environment:
conda create -n unitree-rl python=3.8 conda create -n unitree-rl python=3.8
``` ```
### 1.3 Activate the Virtual Environment ### 1.3 Activate Virtual Environment
```bash ```bash
conda activate unitree-rl conda activate unitree-rl
@@ -46,7 +46,7 @@ conda activate unitree-rl
--- ---
## 2. Installing Dependencies ## 2. Install Dependencies
### 2.1 Install PyTorch ### 2.1 Install PyTorch
@@ -58,15 +58,15 @@ conda install pytorch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 pytorch-cuda=
### 2.2 Install Isaac Gym ### 2.2 Install Isaac Gym
Isaac Gym is a rigid body simulation and training framework provided by Nvidia. Isaac Gym is Nvidia's rigid body simulation and training framework.
#### 2.2.1 Download #### 2.2.1 Download
Download [Isaac Gym](https://developer.nvidia.com/isaac-gym) from Nvidia’s official website. Download [Isaac Gym](https://developer.nvidia.com/isaac-gym) from the Nvidia official website.
#### 2.2.2 Install #### 2.2.2 Install
After extracting the package, navigate to the `isaacgym/python` folder and install it using the following commands: Unzip the file, enter the `isaacgym/python` folder, and execute the following command to install:
```bash ```bash
cd isaacgym/python cd isaacgym/python
@@ -75,18 +75,24 @@ pip install -e .
#### 2.2.3 Verify Installation #### 2.2.3 Verify Installation
Run the following command. If a window opens displaying 1080 balls falling, the installation was successful: Run the following commands. If a window pops up showing 1080 balls falling, the installation is successful:
```bash ```bash
cd examples cd examples
python 1080_balls_of_solitude.py python 1080_balls_of_solitude.py
``` ```
If you encounter any issues, refer to the official documentation at `isaacgym/docs/index.html`. If there are any issues, please refer to the official documentation in `isaacgym/docs/index.html`.
### 2.3 Install rsl_rl ### 2.3 Install rsl_rl
`rsl_rl` is a library implementing reinforcement learning algorithms. `rsl_rl` is a reinforcement learning algorithm library.
Our repository includes `rsl_rl` with new algorithms. Clone the Git repository:
```bash
git clone https://github.com/wty-yy/go2_rl_gym.git
```
#### 2.3.1 Install #### 2.3.1 Install
@@ -97,23 +103,41 @@ pip install -e .
### 2.4 Install go2_rl_gym ### 2.4 Install go2_rl_gym
#### 2.4.1 Download Enter the directory and install:
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 ```bash
cd go2_rl_gym cd go2_rl_gym
pip install -e . pip install -e .
``` ```
### 2.5 Install unitree_cpp_deploy (Optional) ### 2.5 Real Robot Deployment (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). #### 2.5.1 unitree_sdk2
C++ SDK. For compilation, please refer to the [official tutorial](https://github.com/unitreerobotics/unitree_sdk2?tab=readme-ov-file#environment-setup).
#### 2.5.2 unitree_sdk2_python (Choose for Python Deployment)
```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 Install unitree_cpp_deploy (Choose for C++ Deployment)
We use a modified C++ deployment repository based on `unitree_rl_lab`, specifically designed for deploying models trained in this repository. See [unitree_cpp_deploy](https://github.com/wty-yy-mini/unitree_cpp_deploy).
### 2.6 RoboGauge Evaluation (Optional)
RoboGauge is a project for evaluating quadruped robot performance via Sim2Sim in Mujoco. It performs asynchronous evaluation on the CPU during training. For specific details, refer to the [README](https://github.com/wty-yy/RoboGauge).
```bash
git clone [https://github.com/wty-yy/RoboGauge.git](https://github.com/wty-yy/RoboGauge.git)
cd RoboGauge
pip install -e .
```

8
doc/setup_zh.md
View File

@@ -131,3 +131,11 @@ pip install -e .
### 2.5.3 安装 unitree_cpp_deploy（选择用C++部署） ### 2.5.3 安装 unitree_cpp_deploy（选择用C++部署）
我们基于unitree_rl_lab修改的C++部署仓库，专门用于部署本仓库训练的模型 [unitree_cpp_deploy](https://github.com/wty-yy-mini/unitree_cpp_deploy) 我们基于unitree_rl_lab修改的C++部署仓库，专门用于部署本仓库训练的模型 [unitree_cpp_deploy](https://github.com/wty-yy-mini/unitree_cpp_deploy)
### 2.6 RoboGauge评估（可选）
RoboGauge是一个Mujoco中通过Sim2Sim评估四足机器人性能的项目，在训练同时中异步地在cpu上进行评估，具体细节参考[README](https://github.com/wty-yy/RoboGauge)，安装方法
```bash
git clone https://github.com/wty-yy/RoboGauge.git
cd RoboGauge
pip install -e .
```