# 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