5.5 KiB
S10 SuperMimic
基于 mjlab 和 MuJoCo Warp 的 S10 轮足机器人强化学习项目。通过仿真数据训练世界模型，再利用世界模型生成想象轨迹来优化速度控制策略，支持平地、楼梯、平台、斜坡和碎石地形。
演示使用v7_50k策略，在仓库中附带提交导出onnx模型
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
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
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
s10_supermimic/
├── core/ # 世界模型、策略、想象训练、经验回放与演示界面
├── env/ # 平地及复杂地形环境
├── tasks/ # 任务注册与默认参数
└── robot/ # S10 机器人配置与模型资产
scripts/ # 训练、导出与演示入口
docs/ # 演示素材
tests/ # 模型、任务与奖励检查
logs/ # 保留的最终权重与训练日志