110 lines
5.5 KiB
Markdown
110 lines
5.5 KiB
Markdown
# S10 SuperMimic
|
||
|
||
基于 mjlab 和 MuJoCo Warp 的 S10 轮足机器人强化学习项目。通过仿真数据训练世界模型，再利用世界模型生成想象轨迹来优化速度控制策略，支持平地、楼梯、平台、斜坡和碎石地形。
|
||
|
||
演示使用v7_50k策略，在仓库中附带提交导出onnx模型
|
||
|
||
<video src="docs/WorldModel_demo1.mp4" controls width="75%"></video>
|
||
|
||
<video src="docs/Mujoco_demo2.mp4" controls width="75%"></video>
|
||
|
||
## Requirements
|
||
|
||
Python 3.10–3.13、PyTorch ≥ 2.7、mjlab 1.6.0、MuJoCo Warp；参考配置需要 16 GB 显存，如果使用更高显存训练，可以提高 `--imag-batch`。依赖安装：`pip install -r requirements.txt`。
|
||
|
||
## Tasks
|
||
|
||
- `s10-v4`：平地激进策略。
|
||
- `s10-v7`：平地/台阶保守策略，可上 20 cm 台阶。
|
||
- `s10-v4-upward`：倒地起身训练，开发中。
|
||
|
||
## Train
|
||
|
||
```bash
|
||
export PYTHONPATH="$PWD"
|
||
python scripts/train.py --task s10-v4 --output logs/s10_v4
|
||
python scripts/train.py --task s10-v7 --output logs/s10_v7
|
||
python scripts/train.py --task s10-v4-upward --output logs/s10_v4_upward
|
||
```
|
||
|
||
`--task` 加载任务默认配置，命令行参数可覆盖。训练输出 `model_<step>.pt` 和 `metrics.jsonl`；`logs/` 保留 v4、v7 的 50k 最终权重。
|
||
|
||
## Play and Export
|
||
|
||
```bash
|
||
python scripts/export_onnx.py --checkpoint <pt> --output <onnx>
|
||
```
|
||
|
||
ONNX 只包含确定性 actor。`./scripts/play_gui.sh --checkpoint <pt>` 提供 GUI 交互演示，可可视化世界模型的未来预测，但不能替代 sim2sim；sim2sim 建议使用导出的 ONNX。
|
||
|
||
## World model
|
||
|
||
世界模型是带残差预测的前馈多层感知机（MLP），默认使用三层 512 维隐藏层和 SiLU 激活函数。它根据当前状态和动作预测下一时刻的状态变化，以及机身平面位移和偏航角变化。输入和输出使用运行统计量进行归一化。
|
||
|
||
| 输入 | 维度 | 内容 |
|
||
| --- | ---: | --- |
|
||
| 机器人状态 | 38 | 高度 1、投影重力 3、机身线速度 3、角速度 3、腿关节位置 12、腿关节速度 12、轮速 4 |
|
||
| 当前动作 | 16 | 12 维腿部位置控制和 4 维轮速控制 |
|
||
| 上一动作 | 16 | 上一控制周期的动作 |
|
||
| 物理参数 | 7 | 四个轮子的摩擦系数和机身质心偏移 |
|
||
| 地形扫描 | 273 | 复杂地形模型使用的 21 × 13 高度采样网格 |
|
||
|
||
平地世界模型输入为 **77 维**；复杂地形世界模型加入高度扫描后为 **350 维**。两者的动力学输出均为 **41 维**：38 维状态变化和 3 维平面位姿变化。复杂地形模型（revision ≥ 4）带有接触概率与接触力预测头，地形 v7 还带有非轮部位碰撞预测头。世界模型仍使用真实轮速作为训练时的特权状态，这与策略观测不同。
|
||
|
||
训练时从仿真回放中采样连续片段，世界模型按时间自回归预测并拟合实际轨迹；默认训练窗口为 8 步。随后使用世界模型展开想象轨迹，并结合奖励和 critic 更新策略。
|
||
|
||
## Actor
|
||
|
||
Actor 根据机器人自身观测和速度指令输出控制动作。当前 S10 Actor（v4 / v4-upward / v7）每帧接收 **53 维**观测：
|
||
|
||
| 观测 | 每帧维度 |
|
||
| --- | ---: |
|
||
| 机身角速度 | 3 |
|
||
| 投影重力 | 3 |
|
||
| 腿关节位置 | 12 |
|
||
| 腿关节速度 | 12 |
|
||
| 轮速 | 4 |
|
||
| 上一动作 | 16 |
|
||
| 速度指令（前后、横移、转向） | 3 |
|
||
|
||
默认保留 **5 帧**历史，因此 Actor 输入为 **265 维**；各观测项分别按时间从旧到新排列。控制周期为 0.02 秒，5 帧覆盖 5 个控制时刻。
|
||
|
||
Actor 是默认隐藏层为 **256 → 128 → 64** 的 MLP，使用输入归一化和 ELU 激活，最后通过 tanh 输出 **16 维动作**：12 维腿部位置控制和 4 维轮速控制。Actor 不接收地形扫描，实现盲走控制；演示时使用 Actor 生成动作。
|
||
|
||
## Critic
|
||
|
||
Critic 在训练中估计当前状态的未来回报，输出一个标量价值。它使用训练时可获得的特权信息，输入不需要堆叠 Actor 的历史帧：
|
||
|
||
| 输入 | 维度 | 内容 |
|
||
| --- | ---: | --- |
|
||
| 机器人状态 | 38 | 包含机身高度、线速度、角速度、关节状态和轮速 |
|
||
| 速度指令 | 3 | 前后、横移和转向指令 |
|
||
| 物理参数 | 7 | 四个轮子摩擦系数和机身质心偏移 |
|
||
| 上一动作 | 16 | 上一控制周期的动作 |
|
||
| 地形扫描 | 273 | 复杂地形的 21 × 13 高度采样网格 |
|
||
|
||
平地 Critic 输入为 **64 维**，复杂地形 Critic 输入为 **337 维**。Critic 默认同样使用 **256 → 128 → 64** 的 MLP、输入归一化和 ELU 激活，输出 **1 维**价值估计。Critic 和地形扫描只用于训练；运行策略时不需要 Critic。
|
||
|
||
训练时，世界模型生成想象轨迹，Actor 通过轨迹中的奖励和目标 Critic 的末端价值优化动作；Critic 使用想象轨迹计算的 TD(λ) 回报作为目标进行拟合。Actor 的 5 帧历史与世界模型默认的 8 步训练窗口是不同的参数，分别由 `--history` 和 `--wm-window` 设置。
|
||
|
||
## Structure
|
||
|
||
```text
|
||
s10_supermimic/
|
||
├── core/ # 世界模型、策略、想象训练、经验回放与演示界面
|
||
├── env/ # 平地及复杂地形环境
|
||
├── tasks/ # 任务注册与默认参数
|
||
└── robot/ # S10 机器人配置与模型资产
|
||
scripts/ # 训练、导出与演示入口
|
||
docs/ # 演示素材
|
||
tests/ # 模型、任务与奖励检查
|
||
logs/ # 保留的最终权重与训练日志
|
||
```
|
||
|
||
## References
|
||
|
||
- 仿真基于：https://github.com/mujocolab/mjlab
|
||
- 算法基于：
|
||
- https://github.com/zita-ch/supermimic_vibed
|
||
- https://github.com/uwvwko-zzz/go2w-supermimic
|