diff --git a/README.md b/README.md index 30b88ff..8dda9ce 100644 --- a/README.md +++ b/README.md @@ -39,8 +39,11 @@ Demo of the evaluation process (4 terrains, 2 difficulty levels): | Robot | Type | Velocity Tracking | | - | - | - | +| [Unitree go1](./assets/docs/go1_policy_io_zh.md) | Quadruped | ✅ | | Unitree go2 | Quadruped | ✅ | +> Go1 Isaac Gym and RobotLab exports use different ONNX contracts: Gym uses a 5-frame, 225-value history, while RobotLab uses a 10-frame, 450-value history. Both TorchScript exports consume the current 45-value frame and maintain history internally. See the linked Go1 policy I/O guide for the exact layout. + ## Installation Recommended environment: `python=3.8`. Install any PyTorch version >= 2.0 following the [official PyTorch guide](https://pytorch.org/get-started/locally/), then install this repo: diff --git a/README_zh.md b/README_zh.md index c13164b..0f9781b 100644 --- a/README_zh.md +++ b/README_zh.md @@ -38,6 +38,8 @@ | Unitree go2 | 四足 | ✅ | | [Unitree go1](./assets/docs/go1_policy_io_zh.md) | 四足 | ✅ | +> Go1 的 Isaac Gym 与 RobotLab 导出模型接口不同:Gym ONNX 使用 5 帧、225 维历史,RobotLab ONNX 使用 10 帧、450 维历史。TorchScript 均接收当前 45 维单帧观测并在模型内部维护历史。精确排列顺序见 [Go1 策略输入输出与部署说明](./assets/docs/go1_policy_io_zh.md)。 + ## 安装 推荐安装`python=3.8`, 参考[PyTorch官网](https://pytorch.org/get-started/locally/)安装2.0以上任意版本, 安装本仓库 diff --git a/assets/docs/go1_policy_io_zh.md b/assets/docs/go1_policy_io_zh.md index 0a1344d..7b0f95f 100644 --- a/assets/docs/go1_policy_io_zh.md +++ b/assets/docs/go1_policy_io_zh.md @@ -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 Gym:5 帧 + Student MoE encoder 使用 5 个连续控制帧: ```text @@ -137,8 +159,22 @@ Student MoE encoder 使用 5 个连续控制帧: 调用 `model.reset()` 后,历史缓冲为全零。随后每次调用模型都会丢弃最旧 一帧并追加当前观测。这与训练和 RoboGauge 中的行为一致。 +### 4.2 RobotLab:10 帧 + +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 ONNX:5 帧、225 维 + `policy_15k.onnx` 是无状态模型,不在模型内部保存历史。它与 TorchScript 的输入接口不同: @@ -201,7 +250,69 @@ ONNX 输入不是直接对 `[1, 5, 45]` 执行 `reshape`。225 维输入按观 调用方需要维护 5 帧历史,并在 episode reset 时将历史清零。RoboGauge 当前 直接使用 TorchScript;ONNX 主要用于其他推理后端。 -## 7. 最小调用示例 +### 6.2 RobotLab ONNX:10 帧、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]`。 diff --git a/resources/models/go1/policy_robotlab_3500.onnx b/resources/models/go1/policy_robotlab_3500.onnx new file mode 100644 index 0000000..9fd1824 Binary files /dev/null and b/resources/models/go1/policy_robotlab_3500.onnx differ diff --git a/resources/models/go1/policy_robotlab_3500.pt b/resources/models/go1/policy_robotlab_3500.pt new file mode 100644 index 0000000..e35e913 Binary files /dev/null and b/resources/models/go1/policy_robotlab_3500.pt differ diff --git a/resources/models/go1/policy_robotlab_6000.onnx b/resources/models/go1/policy_robotlab_6000.onnx new file mode 100644 index 0000000..d36f539 Binary files /dev/null and b/resources/models/go1/policy_robotlab_6000.onnx differ diff --git a/resources/models/go1/policy_robotlab_6000.pt b/resources/models/go1/policy_robotlab_6000.pt new file mode 100644 index 0000000..074f809 Binary files /dev/null and b/resources/models/go1/policy_robotlab_6000.pt differ diff --git a/resources/models/go1/policy_robotlab_6500.onnx b/resources/models/go1/policy_robotlab_6500.onnx new file mode 100644 index 0000000..935123a Binary files /dev/null and b/resources/models/go1/policy_robotlab_6500.onnx differ diff --git a/resources/models/go1/policy_robotlab_6500.pt b/resources/models/go1/policy_robotlab_6500.pt new file mode 100644 index 0000000..962d11f Binary files /dev/null and b/resources/models/go1/policy_robotlab_6500.pt differ diff --git a/robogauge/tasks/__init__.py b/robogauge/tasks/__init__.py index 525db8b..3474c53 100644 --- a/robogauge/tasks/__init__.py +++ b/robogauge/tasks/__init__.py @@ -10,6 +10,8 @@ from robogauge.tasks.robots import ( Go2MoETerrainConfig, Go1MoEConfig, Go1MoETerrainConfig, + Go1LabConfig, + Go1LabTerrainConfig, ) from robogauge.tasks.pipeline import BasePipeline from robogauge.tasks.gauge import BaseGaugeConfig @@ -46,6 +48,15 @@ task_register.register('go1_moe.stairs_fd', BasePipeline, MujocoConfig, Go2Stair task_register.register('go1_moe.stairs_bd', BasePipeline, MujocoConfig, Go2StairsBackwardGaugeConfig, Go1MoETerrainConfig) task_register.register('go1_moe.obstacle', BasePipeline, MujocoConfig, Go2ObstacleGaugeConfig, Go1MoETerrainConfig) +# Go1 RobotLab +task_register.register('go1_lab.flat', BasePipeline, MujocoConfig, Go2FlatGaugeConfig, Go1LabConfig) +task_register.register('go1_lab.slope_fd', BasePipeline, MujocoConfig, Go2SlopeForwardGaugeConfig, Go1LabTerrainConfig) +task_register.register('go1_lab.slope_bd', BasePipeline, MujocoConfig, Go2SlopeBackwardGaugeConfig, Go1LabTerrainConfig) +task_register.register('go1_lab.wave', BasePipeline, MujocoConfig, Go2WaveGaugeConfig, Go1LabTerrainConfig) +task_register.register('go1_lab.stairs_fd', BasePipeline, MujocoConfig, Go2StairsForwardGaugeConfig, Go1LabTerrainConfig) +task_register.register('go1_lab.stairs_bd', BasePipeline, MujocoConfig, Go2StairsBackwardGaugeConfig, Go1LabTerrainConfig) +task_register.register('go1_lab.obstacle', BasePipeline, MujocoConfig, Go2ObstacleGaugeConfig, Go1LabTerrainConfig) + # Go2 Lab task_register.register('go2_lab.flat', BasePipeline, MujocoConfig, Go2FlatGaugeConfig, Go2LabConfig) task_register.register('go2_lab.slope_fd', BasePipeline, MujocoConfig, Go2SlopeForwardGaugeConfig, Go2LabTerrainConfig) diff --git a/robogauge/tasks/pipeline/base_pipeline.py b/robogauge/tasks/pipeline/base_pipeline.py index a53c767..e9bc9df 100644 --- a/robogauge/tasks/pipeline/base_pipeline.py +++ b/robogauge/tasks/pipeline/base_pipeline.py @@ -17,7 +17,7 @@ from copy import deepcopy from robogauge.utils.logger import logger from robogauge.tasks.simulator import MujocoSimulator, MujocoConfig, SimData from robogauge.tasks.robots import ( - BaseRobot, RobotConfig, Go2Config, Go2, Go2MoEConfig, Go2MoE, Go1MoE + BaseRobot, RobotConfig, Go2Config, Go2, Go2MoEConfig, Go2MoE, Go1, Go1MoE ) from robogauge.tasks.gauge import BaseGauge, BaseGaugeConfig from robogauge.tasks.gauge.goal_data import GoalData, VelocityGoal, PositionGoal diff --git a/robogauge/tasks/robots/__init__.py b/robogauge/tasks/robots/__init__.py index a6c4dbf..bdfee36 100644 --- a/robogauge/tasks/robots/__init__.py +++ b/robogauge/tasks/robots/__init__.py @@ -6,3 +6,4 @@ from .go2.go2 import Go2 from .go2.go2_moe_config import Go2MoEConfig, Go2MoETerrainConfig from .go2.go2_moe import Go2MoE from .go1 import Go1Config, Go1TerrainConfig, Go1, Go1MoEConfig, Go1MoETerrainConfig, Go1MoE +from .go1.go1_lab_config import Go1LabConfig, Go1LabTerrainConfig diff --git a/robogauge/tasks/robots/go1/go1_lab_config.py b/robogauge/tasks/robots/go1/go1_lab_config.py new file mode 100644 index 0000000..785e89f --- /dev/null +++ b/robogauge/tasks/robots/go1/go1_lab_config.py @@ -0,0 +1,18 @@ +# -*- coding: utf-8 -*- +"""Go1 RobotLab policy configuration for RoboGauge.""" + +from robogauge.tasks.robots.go1.go1_config import Go1Config, Go1TerrainConfig + + +class Go1LabConfig(Go1Config): + """Go1 configuration aligned with RobotLab observation scaling.""" + + class control(Go1Config.control): + model_path = "{ROBOGAUGE_ROOT_DIR}/resources/models/go1/policy_robotlab_3500.pt" + + class scales(Go1Config.control.scales): + cmd = [1.0, 1.0, 1.0] + + +class Go1LabTerrainConfig(Go1LabConfig, Go1TerrainConfig): + """Go1 RobotLab policy configuration for terrain tasks."""