diff --git a/README.md b/README.md
index 500c239..9b6e1d2 100644
--- a/README.md
+++ b/README.md
@@ -1,152 +1,172 @@
-
Go2 RL GYM
-
- 🌎 English | 🇨🇳 中文
-
+
Go2 RL GYM
+
+ 🌎 English | 🇨🇳 中文
+
- 本仓库基于unitree_rl_gym,使用强化学习训练Go2机器狗。
+ This repository builds on unitree_rl_gym to train the Unitree Go2 quadruped with reinforcement learning.
-|
Isaac Gym
|
Mujoco
|
Physical
|
+|
Isaac Gym
|
Mujoco
|
Physical
|
|--- | --- | --- |
-| TODO | TODO | TODO |
+|  |  |  |
-## 📦 安装配置
+## 📦 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
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`
+#### ⚙️ Arguments
+- `--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`: Render viewer by default; set to `true` to disable rendering for higher throughput.
+- `--resume`: Resume training from a chosen checkpoint in the logs.
+- `--experiment_name`: Experiment folder to save/load from.
+- `--run_name`: Run subfolder name to save/load from.
+- `--load_run`: Name of the run to load (defaults to the most recent run).
+- `--checkpoint`: Checkpoint index to load (defaults to the latest file).
+- `--num_envs`: Number of parallel simulated environments.
+- `--seed`: Random seed.
+- `--max_iterations`: Maximum training iterations.
+- `--sim_device`: Physics simulation device. Use `--sim_device=cpu` to force 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//_/model_.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//_/model_.pt`
---
### 2. Play
-如果想要在 Gym 中查看训练效果,可以运行以下命令:
+Visualize policies inside Gym with:
```bash
python legged_gym/scripts/play.py --task=xxx
```
-**说明**:
+**Notes**
-- 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 launches on randomized terrain with difficulty between 7 and 9.
+- It automatically loads the latest checkpoint inside the experiment folder.
+- Override via `experiment_name` and `checkpoint`, for example:
+ ```bash
+ 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` 中:
-- `policy.pt`: torch script模型,用于Sim2Sim。
-- `policy.onnx`: onnx模型,用于Sim2Real。
-- `policy.pkl`: 模型权重。
+Play exports the Actor network to `logs/{experiment_name}/exported/policies`:
+- `policy.pt`: TorchScript model for Sim2Sim.
+- `policy.onnx`: ONNX model for Sim2Real.
+- `policy.pkl`: Raw weights.
-#### Play 效果
+#### Demonstration
-| Go2 |
-|--- |
-| TODO |
+
---
### 3. Sim2Sim (Mujoco)
-支持在 Mujoco 仿真器中运行 Sim2Sim:
+Run policies in the Mujoco simulator:
```bash
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`。
-- 替换环境地形:默认地形为 `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 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`.
+- **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 |
|--- | --- | --- |
-| | | |
+|  |  |  |
---
### 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
cd deploy/deploy_real
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 |
| --- | --- |
-| | |
+|  |  |
---
-## 🎉 致谢
+## 🎉 Acknowledgements
-本仓库开发离不开以下开源项目的支持与贡献,特此感谢:
+This repository would not exist without the following open-source projects:
-- [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++通信接口。
+- [unitree_rl_gym](https://github.com/unitreerobotics/unitree_rl_gym): Unitree's core RL training framework.
+- [legged_gym](https://github.com/leggedrobotics/legged_gym): Base locomotion environment.
+- [rsl_rl](https://github.com/leggedrobotics/rsl_rl.git): Reinforcement learning algorithms.
+- [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 hardware interface for deployment.
+- [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)
+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) 授权。
-
-详情请阅读完整 [LICENSE 文件](./LICENSE)。
+New contributions follow the [MIT License](LICENSE); the original unitree_rl_gym remains under the [BSD 3-Clause License](LICENSE).
+See the complete [LICENSE file](LICENSE) for details.
diff --git a/README_en.md b/README_en.md
deleted file mode 100644
index 6524d86..0000000
--- a/README_en.md
+++ /dev/null
@@ -1,11 +0,0 @@
-
-
Go2 RL GYM
-
- 🌎English | 🇨🇳中文
-
-
-
-
- This is a repository for reinforcement learning implementation based on Unitree Go2. Based on unitree_rl_gym.
-
-
diff --git a/README_zh.md b/README_zh.md
new file mode 100644
index 0000000..af7ae37
--- /dev/null
+++ b/README_zh.md
@@ -0,0 +1,171 @@
+
+
+
+ 本仓库基于unitree_rl_gym,使用强化学习训练Go2机器狗。
+
+
+
+
+|
Isaac Gym
|
Mujoco
|
Physical
|
+|--- | --- | --- |
+|  |  |  |
+
+
+
+## 📦 安装配置
+
+安装和配置步骤请参考 [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//_/model_.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 效果
+
+
+
+---
+
+### 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](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++部署 |
+| --- | --- |
+|  |  |
+
+---
+
+## 🎉 致谢
+
+本仓库开发离不开以下开源项目的支持与贡献,特此感谢:
+
+- [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)。
+
+
diff --git a/deploy/deploy_mujoco/deploy_go2.py b/deploy/deploy_mujoco/deploy_go2.py
index dfd1fdb..f1512db 100644
--- a/deploy/deploy_mujoco/deploy_go2.py
+++ b/deploy/deploy_mujoco/deploy_go2.py
@@ -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 mujoco.viewer
import mujoco
@@ -7,12 +14,10 @@ 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]
@@ -110,7 +115,7 @@ if __name__ == "__main__":
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")
+ video_save_dir = str(PATH_PARENT / "videos")
os.makedirs(video_save_dir, exist_ok=True)
model_name = os.path.basename(policy_path).split('.')[0]
@@ -134,17 +139,18 @@ if __name__ == "__main__":
# load policy
policy = torch.jit.load(policy_path)
+ video_fps = 50
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}")
+ mujoco_render_utils = MujocoRenderUtils(video_fps, m.opt.timestep)
if visualize_moe_weights:
plt.ion()
@@ -153,7 +159,7 @@ if __name__ == "__main__":
bars = None
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)
latent_filename = f"{model_name}_{cmd_str}_latents.npy"
latent_path = os.path.join(latent_save_dir, latent_filename)
@@ -164,9 +170,9 @@ if __name__ == "__main__":
# 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
+ viewer.cam.distance = 2.0
+ viewer.cam.elevation = -20.0
+ viewer.cam.azimuth = 60.0
# Close the viewer automatically after simulation_duration wall-seconds.
start = time.time()
@@ -188,10 +194,12 @@ if __name__ == "__main__":
# 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)
+ mujoco_render_utils.update(cmd, d)
if save_video and counter % frame_skip == 0:
try:
renderer.update_scene(d, camera=viewer.cam)
+ mujoco_render_utils.update_external_rendering(renderer, ctype='renderer')
frame = renderer.render()
writer.append_data(frame)
except Exception as e:
@@ -249,6 +257,7 @@ if __name__ == "__main__":
target_dof_pos = action * action_scale + default_angles
# Pick up changes to the physics state, apply perturbations, update options from GUI.
+ mujoco_render_utils.update_external_rendering(viewer, ctype='viewer')
viewer.sync()
# Rudimentary time keeping, will drift relative to wall clock.
diff --git a/deploy/deploy_mujoco/utils.py b/deploy/deploy_mujoco/utils.py
new file mode 100644
index 0000000..2badcd8
--- /dev/null
+++ b/deploy/deploy_mujoco/utils.py
@@ -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
diff --git a/deploy/deploy_real/config_go2.py b/deploy/deploy_real/config_go2.py
index a7e4fde..05e9fd2 100644
--- a/deploy/deploy_real/config_go2.py
+++ b/deploy/deploy_real/config_go2.py
@@ -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 yaml
diff --git a/deploy/deploy_real/configs/go2.yaml b/deploy/deploy_real/configs/go2.yaml
index c4b221c..f64ede8 100644
--- a/deploy/deploy_real/configs/go2.yaml
+++ b/deploy/deploy_real/configs/go2.yaml
@@ -6,7 +6,7 @@ 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"
+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]
diff --git a/deploy/deploy_real/deploy_real_go2.py b/deploy/deploy_real/deploy_real_go2.py
index aac67e7..56954c2 100644
--- a/deploy/deploy_real/deploy_real_go2.py
+++ b/deploy/deploy_real/deploy_real_go2.py
@@ -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 time
import torch
@@ -85,11 +86,13 @@ class Controller:
def zero_torque_state(self):
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:
create_zero_cmd(self.low_cmd)
self.send_cmd(self.low_cmd)
time.sleep(self.config.control_dt)
+ print("Start signal received.")
+ print("Press *select* button to exit.")
def move_to_default_pos(self):
@@ -121,7 +124,7 @@ class Controller:
def default_pos_state(self):
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:
for i in range(12):
motor_idx = self.config.joint2motor_idx[i]
@@ -162,7 +165,12 @@ class Controller:
self.obs[33:45] = self.action
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
diff --git a/doc/setup_en.md b/doc/setup_en.md
index 1ee8b54..2a29e94 100644
--- a/doc/setup_en.md
+++ b/doc/setup_en.md
@@ -1,20 +1,20 @@
-# Installation Guide
+# Installation and Configuration Guide
## System Requirements
-- **Operating System**: Recommended Ubuntu 18.04 or later
-- **GPU**: Nvidia GPU
-- **Driver Version**: Recommended version 525 or later
+- **OS**: Ubuntu 18.04 or higher is recommended
+- **GPU**: Nvidia GPU
+- **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
-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
mkdir -p ~/miniconda3
@@ -30,7 +30,7 @@ After installation, initialize Conda:
source ~/.bashrc
```
-### 1.2 Create a New Environment
+### 1.2 Create New 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
```
-### 1.3 Activate the Virtual Environment
+### 1.3 Activate Virtual Environment
```bash
conda activate unitree-rl
@@ -46,7 +46,7 @@ conda activate unitree-rl
---
-## 2. Installing Dependencies
+## 2. Install Dependencies
### 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
-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
-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
-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
cd isaacgym/python
@@ -75,18 +75,24 @@ pip install -e .
#### 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
cd examples
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
-`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
@@ -97,23 +103,41 @@ 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:
+Enter the directory and install:
```bash
cd go2_rl_gym
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 .
+```
diff --git a/doc/setup_zh.md b/doc/setup_zh.md
index eedce79..18e26c9 100644
--- a/doc/setup_zh.md
+++ b/doc/setup_zh.md
@@ -131,3 +131,11 @@ 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)
+
+### 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 .
+```