update lab

This commit is contained in:
youyuan.chen
2026-07-26 19:11:42 +08:00
parent 217ef2247d
commit 82fbc97e14
13 changed files with 203 additions and 26 deletions

View File

@@ -1,8 +1,10 @@
# Go1 MoE 策略输入输出与部署说明
本文说明 `go1` 分支中 Go1 MoE locomotion 策略的输入、输出、状态管理和
RoboGauge/MuJoCo 接入方式。适用于由 `go2_rl_gym` 导出的 Go1 MoE CTS
TorchScript 模型,例如:
RoboGauge/MuJoCo 接入方式,同时覆盖 `go2_rl_gym`Isaac Gym
`go2_rl_robotlab`Isaac Lab/RobotLab导出的模型。
Isaac Gym 模型例如:
- `resources/models/go1/policy.pt`
- `resources/models/go1/policy_10k.pt`
@@ -11,23 +13,33 @@ TorchScript 模型,例如:
- `resources/models/go1/policy_25k.pt`
- `resources/models/go1/policy_25k.onnx`
RobotLab 模型例如:
- `resources/models/go1/policy_robotlab_3500.pt`
- `resources/models/go1/policy_robotlab_3500.onnx`
- `resources/models/go1/policy_robotlab_6000.pt`
- `resources/models/go1/policy_robotlab_6000.onnx`
- `resources/models/go1/policy_robotlab_6500.pt`
- `resources/models/go1/policy_robotlab_6500.onnx`
## 1. 接口概览
| 项目 | 规格 |
| - | - |
| 控制频率 | 50 Hz控制周期 `0.02 s` |
| TorchScript 输入 | `torch.float32`,形状 `[1, 45]` |
| 内部历史 | 5 帧,形状 `[1, 5, 45]` |
| 策略动作 | `torch.float32`,形状 `[1, 12]` |
| MoE 权重 | `torch.float32`,形状 `[1, 8]` |
| Student latent | `torch.float32`,形状 `[1, 32]` |
| 控制方式 | 关节位置 PD 控制 |
| 关节位置增量比例 | `action_scale = 0.25` |
| PD 参数 | `Kp = 28.0``Kd = 0.7` |
| 项目 | Isaac Gym | RobotLab |
| - | - | - |
| 控制频率 | 50 Hz`0.02 s` | 50 Hz`0.02 s` |
| TorchScript 输入 | `[1, 45]` 当前帧 | `[1, 45]` 当前帧 |
| TorchScript 内部历史 | 5 帧 | 10 帧 |
| TorchScript 输出 | action、MoE weights、latent | action |
| ONNX 输入 | `[1, 225]`5 帧 | `[1, 450]`10 帧 |
| ONNX 输出 | action、MoE weights、latent | action |
| 策略动作 | `[1, 12]` | `[1, 12]` |
| 控制方式 | 关节位置 PD | 关节位置 PD |
| `action_scale` | `0.25` | `0.25` |
| PD 参数 | `Kp=28.0, Kd=0.7` | `Kp=28.0, Kd=0.7` |
导出的 TorchScript 是有状态模型调用方每个控制周期只输入当前 45 维
观测,模型会在内部更新 5 帧历史。不要 TorchScript 输入端再次拼接成
225 维
所有输入和输出张量均为 `float32`。两类 TorchScript 是有状态模型调用方
每个控制周期只输入当前 45 维观测。不要 TorchScript 输入 225 或 450 维
历史。两类 ONNX 都是无状态模型,历史必须由调用方维护
## 2. 坐标系和关节顺序
@@ -72,7 +84,7 @@ Isaac Gym 中 Go1 资产的原始关节顺序不同,但训练环境已经转
| - | -: | - | - |
| `[0:3]` | 3 | 机身坐标系角速度 `[wx, wy, wz]` | `base_ang_vel * 0.25` |
| `[3:6]` | 3 | 世界重力方向在机身坐标系中的投影 | `projected_gravity` |
| `[6:9]` | 3 | 速度指令 `[vx, vy, yaw_rate]` | `command * [2.0, 2.0, 0.25]` |
| `[6:9]` | 3 | 速度指令 `[vx, vy, yaw_rate]` | Gym: `command * [2.0, 2.0, 0.25]`RobotLab: `command * [1.0, 1.0, 1.0]` |
| `[9:21]` | 12 | 关节位置 | `(q - q_default) * 1.0` |
| `[21:33]` | 12 | 关节速度 | `dq * 0.05` |
| `[33:45]` | 12 | 上一个控制周期的原始策略动作 | `last_action` |
@@ -99,7 +111,15 @@ projected_gravity = get_projected_gravity(base_quat_wxyz)
### 3.3 速度指令
输入的是期望机身速度,不是目标点坐标
输入的是期望机身速度,不是目标点坐标。Gym 与 RobotLab 的命令归一化不同,
不能共用同一个缩放数组:
```python
GYM_COMMAND_SCALE = np.array([2.0, 2.0, 0.25], dtype=np.float32)
ROBOTLAB_COMMAND_SCALE = np.array([1.0, 1.0, 1.0], dtype=np.float32)
```
指令语义均为:
```text
[forward_velocity, lateral_velocity, yaw_rate]
@@ -125,6 +145,8 @@ yaw_rate = [-1.0, 1.0] rad/s
## 4. 历史观测
### 4.1 Isaac Gym5 帧
Student MoE encoder 使用 5 个连续控制帧:
```text
@@ -137,8 +159,22 @@ Student MoE encoder 使用 5 个连续控制帧:
调用 `model.reset()` 后,历史缓冲为全零。随后每次调用模型都会丢弃最旧
一帧并追加当前观测。这与训练和 RoboGauge 中的行为一致。
### 4.2 RobotLab10 帧
RobotLab Student MoE encoder 使用 10 个连续控制帧:
```text
[obs(t-9), obs(t-8), ..., obs(t-1), obs(t)]
```
控制周期为 20 ms最旧帧与最新帧相隔 180 ms。RobotLab TorchScript 内部
维护的是 450 维、按观测项分组的历史缓冲,调用接口仍是当前 `[1, 45]`
单帧。每个 episode 开始时同样必须调用 `model.reset()`
## 5. TorchScript 输出
### 5.1 Isaac Gym TorchScript
MoE TorchScript 返回:
```python
@@ -171,8 +207,21 @@ Kd = 0.7
实际部署还应按电机能力限制输出力矩。
### 5.2 RobotLab TorchScript
RobotLab TorchScript 只返回动作:
```python
action = model(obs_tensor) # [1, 12]
```
它不会返回 MoE weights 和 latent。动作到 PD 目标位置的换算、关节顺序以及
`Kp/Kd` 与上文相同。
## 6. ONNX 输入输出
### 6.1 Isaac Gym ONNX5 帧、225 维
`policy_15k.onnx` 是无状态模型,不在模型内部保存历史。它与 TorchScript
的输入接口不同:
@@ -201,7 +250,69 @@ ONNX 输入不是直接对 `[1, 5, 45]` 执行 `reshape`。225 维输入按观
调用方需要维护 5 帧历史,并在 episode reset 时将历史清零。RoboGauge 当前
直接使用 TorchScriptONNX 主要用于其他推理后端。
## 7. 最小调用示例
### 6.2 RobotLab ONNX10 帧、450 维
`policy_robotlab_*.onnx` 同样是无状态模型,但历史长度为 10
| 项目 | ONNX 规格 |
| - | - |
| 输入名称 | `obs` |
| 输入 dtype | `float32` |
| 输入形状 | `[1, 450]` |
| 输出名称 | `actions` |
| 输出形状 | `[1, 12]` |
450 维输入采用 **term-major** 排列,即先放同一观测项的 10 帧历史,再放
下一个观测项。每个时间块内部从旧到新排列:
```text
[
angular_velocity(t-9:t), # 10 * 3 = 30
projected_gravity(t-9:t), # 10 * 3 = 30
command(t-9:t), # 10 * 3 = 30
dof_position_error(t-9:t), # 10 * 12 = 120
dof_velocity(t-9:t), # 10 * 12 = 120
last_action(t-9:t), # 10 * 12 = 120
] total = 450
```
因此以下写法是错误的:
```python
# 错误:这是 frame-major不是 RobotLab ONNX 需要的 term-major。
onnx_input = np.stack(last_10_observations).reshape(1, 450)
```
与导出器一致的历史更新实现如下:
```python
TERM_DIMS = (3, 3, 3, 12, 12, 12)
HISTORY_LENGTH = 10
onnx_history = np.zeros((1, 450), dtype=np.float32)
def push_robotlab_onnx_history(single_obs):
single_obs = np.asarray(single_obs, dtype=np.float32).reshape(45)
history_offset = 0
single_offset = 0
for dim in TERM_DIMS:
block_size = dim * HISTORY_LENGTH
block = onnx_history[:, history_offset:history_offset + block_size]
block[:, :-dim] = block[:, dim:].copy()
block[:, -dim:] = single_obs[None, single_offset:single_offset + dim]
history_offset += block_size
single_offset += dim
return onnx_history
def reset_robotlab_onnx_history():
onnx_history.fill(0.0)
```
首次推理前和每次 episode reset 时都必须清零 `onnx_history`。调用 ONNX
Runtime 时直接传入 `push_robotlab_onnx_history(single_obs)` 的返回值。
## 7. Isaac Gym TorchScript 最小调用示例
```python
import numpy as np
@@ -265,6 +376,13 @@ target_dof_pos = DEFAULT_DOF_POS + 0.25 * action
last_action = action.copy()
```
RobotLab TorchScript 的观测构建只需将示例中的命令缩放改为
`[1.0, 1.0, 1.0]`,推理返回值改为:
```python
action_tensor = model(obs_tensor)
```
## 8. Reset 要求
以下情况必须同时重置模型历史和上一动作:
@@ -275,27 +393,31 @@ last_action = action.copy()
- 控制器长时间暂停后重新启动
- 真机急停解除后重新接管
TorchScript 模型执行:
```python
model.reset()
last_action.fill(0.0)
```
只清零 `last_action` 而不调用 `model.reset()`,会残留上一段运行的历史状态。
ONNX 没有 `reset()` 方法,必须清零调用方维护的 225 或 450 维历史缓冲,
并同时清零 `last_action`
## 9. 不需要的部署信号
部署使用 Student 策略,不需要训练阶段 Teacher/Critic 的 263 维 privileged
observation。以下信号不进入部署策略
部署使用 Student 策略,不需要训练阶段 Teacher/Critic 的 privileged
observation(具体维度随 Gym/RobotLab 配置变化)。以下信号不进入部署策略:
- 机身线速度
- 足端接触力
- 电机力矩
- 关节加速度
- 187 维地形高度扫描
- 地形高度扫描(维度随训练配置变化)
- 深度图、相机图像或 LiDAR
这些信息仅在训练阶段用于 Teacher/Critic。Student 通过 5 帧本体观测历史估计
32 维 latent。
这些信息仅在训练阶段用于 Teacher/Critic。Student 通过 Gym 的 5 帧或
RobotLab 的 10 帧本体观测历史估计 32 维 latent。
## 10. RoboGauge 对应实现
@@ -303,6 +425,7 @@ RoboGauge 已实现上述接口:
- 观测构建:`robogauge/tasks/robots/go2/go2.py`
- Go1 参数和关节顺序:`robogauge/tasks/robots/go1/go1_config.py`
- RobotLab 命令缩放:`robogauge/tasks/robots/go1/go1_lab_config.py`
- MoE 输出处理:`robogauge/tasks/robots/go2/go2_moe.py`
- 投影重力:`robogauge/utils/math_utils.py`
@@ -319,9 +442,26 @@ python robogauge/scripts/run.py \
--goals target_pos_velocity
```
RobotLab 模型使用独立的 `go1_lab.*` 任务,确保命令观测使用 RobotLab 缩放:
```bash
MUJOCO_GL=glfw \
PYTHONPATH=/path/to/RoboGauge \
python robogauge/scripts/run.py \
--task-name go1_lab.stairs_fd \
--experiment-name go1_robotlab_test \
--model-path /path/to/policy_robotlab_6500.pt \
--level 3 \
--spawn-type level_eval \
--goals keyboard
```
不要用 `go1_moe.*` 任务加载 RobotLab 模型;该任务使用 Isaac Gym 的命令
缩放。RoboGauge 当前评估加载 TorchScript `.pt`,不会直接加载 ONNX。
## 11. 常见错误检查
1. 输入形状应为 `[1, 45]`,不是 `[45]` `[1, 225]`
1. TorchScript 输入形状应为 `[1, 45]`Gym ONNX 输入为 `[1, 225]`RobotLab ONNX 输入为 `[1, 450]`
2. 输入 dtype 必须为 `torch.float32`
3. 角速度必须在机身坐标系下。
4. 四元数必须是 `[w, x, y, z]`
@@ -329,3 +469,5 @@ python robogauge/scripts/run.py \
6. `last_action` 必须使用未乘 `0.25` 的原始模型输出。
7. 每次 reset 必须同时调用 `model.reset()` 并清零 `last_action`
8. 不要向部署模型输入 privileged observation 或高度扫描。
9. RobotLab ONNX 历史必须按 term-major 排列,不能直接 flatten 形状为 `[10, 45]` 的逐帧历史。
10. RobotLab 模型必须使用命令缩放 `[1.0, 1.0, 1.0]`,不能沿用 Gym 的 `[2.0, 2.0, 0.25]`