SuperMimic 深度解读：一个可微世界模型如何把人形机器人全身控制的训练从集群搬到单卡

Hermit_Rabbit

32人浏览 · 2026-11-07 11:59:49

Hermit_Rabbit · 2026-11-07 11:59:49 发布

0. 简介

SuperMimic 面向人形机器人的大规模全身动作跟踪，处理的是"动作库越大、训练算力越贵"这个具体矛盾。它没有继续沿着无模型强化学习加集群算力的路线，而是用一个从仿真数据拟合出来的可微世界模型承载所有动作共享的机器人动力学，再通过 32 步想象推演上的时间反向传播直接给策略解析梯度，让仿真器退回标签生成器的位置。从实验与代码实现看，最值得关注的是单张 H100 训练 3 小时就在 129,514 段动作的同一评测集上超过了 SONIC 公开检查点，失败动作 25 段对 179 段，而这份复现在一张 RTX PRO 6000 上 1.24 小时达到论文报告的 0.82 回报。

在这里插入图片描述

1. 研究问题：动作库越大，为什么算力越贵

1.1 大规模动作跟踪的既有路线

基于物理的动作模仿可以追溯到 2018 年的 DeepMimic：给定一段参考动作，在仿真器里让角色去跟，跟得越像奖励越高，用 PPO 这类无模型强化学习优化策略。后来的 PHC、MaskedMimic、BeyondMimic 和 SONIC 都沿着这条路扩大规模。经验很一致，动作库越多样，学出来的控制器能力越强。目前的两种扩展方式各有代价：一种先给每段动作训练一个专家，再蒸馏成一个通用策略，蒸馏会丢信息；另一种直接用集群算力硬训一个通用策略，SONIC 走的是后一条路，128 张 L40 训练约 168 小时，折合约 21,000 GPU 小时。
在这里插入图片描述

1.2 无模型梯度估计的真实代价

更深一层的原因在梯度估计方式，这也是所有无模型方法共同的瓶颈。PPO 计算的是似然比梯度，先随机采一个动作，看结果比平均好还是差，再据此调整这个动作出现的概率。整个过程只用到"这个动作好不好"这一个标量信息，用不到"动作往哪个方向微调会更好"这种局部结构信息：

$\nabla_\theta J(\theta)=\mathbb{E}\Big[\nabla_\theta\log\pi_\theta(a\mid s)\,\hat{A}(s,a)\Big]$

其中 $\theta$ 是策略参数， $\hat{A}$ 是优势估计。这个估计不需要知道环境怎么运作，代价是方差大、样本利用率低。对于 29 个关节、接触频繁切换的人形系统，要把方差压下去只能靠海量并行采样。进一步看，动作库扩大的时候，每一种新动作都要重新经历同样的采样过程，而它们共享的物理规律从未被显式利用，这正是算力会随覆盖范围线性上涨的根源。

直觉理解：这就像蒙着眼睛在陌生房间里找灯开关。你只能伸手碰一下、感觉亮不亮，然后调整下次伸手的方向，却不知道墙上的电线怎么走。碰一百次能大致找到位置，但每次都要真的伸手。PPO 只知道"这一步走得比平均好还是差"，不知道"往哪边挪一点会更好"；BPTT 两样都知道，代价是必须有一个可求导的模型。

在这里插入图片描述

2. 整体框架：一个循环，两类学习对象

2.1 三步交替，而不是两段式预训练

SuperMimic 的训练不是"先离线训好世界模型、再冻结它训练策略"的两阶段流程。每个迭代依次做三件事：当前策略在 MuJoCo Warp 里执行动作，动作上叠加高斯噪声 $\epsilon_t\sim\mathcal{N}(0,0.1^2)$ ，产生的状态转移存入固定大小的回放缓冲区；从缓冲区采样长度为 8 的轨迹窗口训练世界模型；暂时冻结世界模型权重，从缓冲区取真实访问过的状态作起点，想象 32 步，用 BPTT 更新策略并拟合 Critic。三个对象在整个训练过程中一起演进。

把这个循环落到代码上，supermimic_vibed/supermimic/train.py 的主循环只有十几行。更新数据比（UTD）0.4 的实现方式稍微绕一点：每次迭代平均推进 2.5 个向量化仿真步，用一个累加器把小数部分留到下一轮。每次迭代做 10 次世界模型更新，前 500 次迭代只训练世界模型、不更新策略，这是复现作者自己加的预热，论文没有规定。

# supermimic_vibed/supermimic/train.py
for it in range(start_it, args.iterations):
    step_acc += 1.0 / args.utd          # UTD=0.4，平均每次迭代采 2.5 个仿真步
    n = int(step_acc); step_acc -= n
    collect(n, stats)                    # ① 带噪声的策略在 MuJoCo Warp 里采样，写入回放缓冲区

    for _ in range(args.wm_batches):     # ② 每次迭代 10 个 batch，每个 batch 2048 条 8 步窗口
        l, e8 = wm_update()

    if it >= args.wm_only_iters:         # 复现自加：前 500 次迭代只训世界模型
        pl, vl, ir, gn, surv = policy_update()   # ③ 冻结世界模型，BPTT 更新策略并拟合 Critic

2.2 旧策略的数据为什么还能用

动力学和数据由哪个策略采集无关。同一个状态、同一个动作、同一组物理参数，下一时刻的状态就是确定的。这一点让回放缓冲区里的数据始终有效，也让世界模型能随策略探索新的动作区域持续补充训练标签，减轻离线模型被策略钻空子的风险。动作噪声的作用也只在这里：它扩大了世界模型见过的动作范围，策略本身是确定性的，论文明确说这个噪声是可选的。

2.3 想象起点取自真实状态

直接从任意初态推演上千步，预测误差会累积，策略也能学会利用模型的偏差。SuperMimic 的做法是从回放缓冲区取真实访问过的状态作为想象起点，只用世界模型展开 $H = 32$ 个控制步，末端再由 Critic 估计余下价值。这样做还有一个副作用：想象起点的分布始终贴着策略实际访问过的状态，模型不会被迫在自己从没见过的区域做长距离外推：

$\hat{J}(\theta)=\mathbb{E}_{s_0\sim\rho_0,\pi_\theta}\Big[\sum_{t=0}^{H-1}\gamma^t\hat{c}_t\hat{r}_{t+1}+\gamma^H\hat{c}_H V_\psi(\hat{s}_H,g_H,\xi)\Big]$

其中 $\rho_0$ 是回放缓冲区里的起点分布， $\gamma=0.99$ 是折扣因子， $\hat{c}_t$ 是软存活权重， $V_\psi$ 是 Critic。世界模型只负责可控长度内的状态与梯度预测，其余长期回报交给 Critic，两者一起缓解长时域模型偏差和策略短视。这个"用真实状态做起点"的写法来自 Sutton 的 Dyna 架构和后来的 MBPO。

3. 第一条核心机制：可微世界模型学什么

3.1 确定性预测与增量表示

世界模型 $f_\omega$ 的输入是当前状态、动作和仿真上下文 $\xi$ ，输出下一状态。 $\xi$ 就是这一回合的域随机化参数，包括脚地摩擦系数、29 个关节的标定零点偏差和躯干质心偏移，复现里拼成一个 33 维向量。因为仿真器的随机性来源都是已知的，模型可以保持确定性：给定三元组就输出一个确定的下一状态，不需要预测分布。传感器噪声被单独拿出来，观测写成 $o_t=\bar{o}_t+\eta_t$ ，模型只预测无噪声观测，策略推演时再从已知分布重新采样噪声加上去。

在这里插入图片描述

状态在复现里是六元组 $(p,R,v,\omega,q,\dot{q})$ ，展平后 71 维。直接让网络输出绝对位置有个问题：机器人朝东走和朝北走物理上完全等价，绝对坐标却完全不同，网络要浪费容量学这种对称性。复现的解法是把输入换成与朝向和水平位置无关的特征，输出则是在机体或航向坐标系下表达的帧内增量，再积分回世界坐标。

# supermimic_vibed/supermimic/networks.py
def step(self, S, a, xi_n):
    p, R, v, w, q, qd = S
    # 输入只含与航向、水平位置无关的特征，朝东走和朝北走对网络是同一个样本
    y = run_mlp(self.net, self.in_norm(self.inputs(S, a, xi_n)))
    d = y * self.out_norm.std() + self.out_norm.mean      # 反归一化得到帧内增量
    H = yaw_mat(R)                                         # 只含偏航角的旋转
    dp, dth, dv, dw, dq, dqd = d.split((3, 3, 3, 3, NUM_J, NUM_J), -1)
    p1 = p + (H @ dp.unsqueeze(-1)).squeeze(-1)            # 位置增量在航向系下表达
    R1 = R @ so3_exp(dth)                                  # 旋转增量用指数映射叠加
    v1 = v + (R @ dv.unsqueeze(-1)).squeeze(-1)            # 速度增量在机体系下表达
    w1 = w + (R @ dw.unsqueeze(-1)).squeeze(-1)
    return (p1, R1, v1, w1, q + dq, qd + dqd)

网络本身是三层 2048 宽的 MLP，带 LayerNorm 和 SiLU 激活。最后一层权重和偏置初始化为零，训练开始时模型预测"什么都不变"，这比随机输出一个大增量要稳得多，因为策略在最初几百次迭代里面对的至少是一个"静止的世界"，而不是一个乱跳的世界。输入输出各有一个滑动均值方差归一化器，统计量在采样阶段用真实转移更新，不参与梯度。这套初始化加归一化的组合，是复现作者把单次迭代从每秒 0.66 次提到 3.1 次之外，另一个让训练稳定收敛的关键细节。

3.2 8 步自回归：让模型学会纠正自己

训练世界模型时不是只做单步监督，这一点决定了长时域推演能不能用。如果只喂真实的上一帧，模型永远在"干净的输入"上被评估，一旦开始自己滚动，误差就会迅速累积而它毫无准备。所以它从回放缓冲区截取长度 $k = 8$ 的真实片段，用首帧真实状态初始化，按记录的动作让模型用自己上一步的预测作为下一步输入，连续推演 8 步，每步都和真实状态比较：

$\mathcal{L}_{\mathrm{WM}}(\omega)=\mathbb{E}_{\mathcal{B}}\sum_{j=1}^{k}\Big[\|\hat{s}_j-s_j\|^2+(\hat{r}_j-r_j)^2+\|\hat{\bar{o}}_j-\bar{o}_j\|^2+\mathrm{BCE}(\hat{d}_j,d_j)\Big]$

解析计算的输出头没有对应损失项，复现里实际只剩第一项。自回归训练的意义在于，单步预测误差再小，连续推演也会累积；让模型在训练时就面对"自己略有偏差的输入"，它会学到把轨迹拉回来的能力。论文观察到用 8 步窗口训练的模型可以稳定推演远长于 8 步的轨迹，策略更新用的 32 步想象没有出问题。这里要厘清的是，这是经验观察，8 步损失本身并不构成对任意长时域误差的数学保证。

窗口采样还有一条容易被忽略的边界条件：如果一个 8 步窗口中间跨过回合重置，后面的转移就不再是同一条轨迹，机器人会从摔倒状态瞬间"跳"回某段参考动作的起点。让世界模型拟合这种跳变，只会教它学一种物理上不存在的转移。复现用一个累乘掩码把回合结束之后的权重全部置零。

# supermimic_vibed/supermimic/buffer.py
d = out["done"].float()
# 第 j 个转移有效，当且仅当前 j-1 个转移都没有触发回合结束
mask = torch.cumprod(torch.cat((torch.ones(B, 1, device=self.device), 1 - d[:, :-1]), -1), -1)

3.3 随机性在训练里的分工

采样阶段执行的是带噪动作，噪声写进回放缓冲区，目的只是扩大模型拟合覆盖的控制输入范围。想象阶段策略是确定性的，传感噪声只用来构造带噪观测。把动作采集噪声当成 PPO 的概率策略，或者认为世界模型需要预测动作噪声，都会导致理解偏差。这套分工是 SuperMimic 与 PPO 最容易混淆的地方。

工程价值：噪声只进数据、不进模型，意味着世界模型不必为一个随机量建模，网络容量全部用来拟合确定性的物理转移。这是一个很小但很值钱的设计决定，它让同样的 MLP 规模在 264 小时数据上仍然够用。

4. 第二条核心机制：沿模型直接求导

4.1 软存活权重让终止可微

真实仿真里机器人摔倒就终止，这是一个 0/1 的硬事件，无法求导。世界模型给出的是终止概率 $\hat{d}_t\in[0,1]$ ，存活权重连续变化，预计会摔倒的轨迹后面的奖励被逐步打折。复现里的终止概率由硬阈值做 sigmoid 软化：根部高度偏差超过 0.25 米，或倾斜误差超过 0.8 时，存活概率迅速降到零。这是训练用的近似，真实系统的摔倒并不会因此变得连续。

# supermimic_vibed/supermimic/mdp.py
def soft_termination(self, S, f, tz: float = 0.02, tt: float = 0.04):
    e_z = (rp[..., 2] - p[..., 2]).abs()            # 根部高度偏差，硬阈值 0.25 m
    e_tilt = (rR[..., 2, 2] - R[..., 2, 2]).abs()   # 重力方向 z 分量差，硬阈值 0.8
    alive = (1 - torch.sigmoid((e_z - self.cfg.term_height) / tz)) \
          * (1 - torch.sigmoid((e_tilt - self.cfg.term_tilt) / tt))
    return 1 - alive                                 # 可微的终止概率

4.2 32 步展开，一次反向传播

有了目标函数，实现就是一个前向展开加一次 loss.backward()，没有重要性采样、没有裁剪比率、也没有优势归一化。展开前把世界模型和 Critic 的参数设为不需要梯度，反向传播会穿过它们的计算图，却不会更新它们的权重。换句话说，求导路径完整保留，只是这两个模块在这一次更新里"只读"。下面这段是复现代码的核心，几乎逐行对应式 (2)：

# supermimic_vibed/supermimic/train.py 的 policy_update()，有删节
set_requires_grad(wm, False)          # 冻结世界模型：梯度穿过它，但不改它
set_requires_grad(critic, False)
for t in range(HZ):                   # HZ = 32，在世界模型里向前想象
    S, a_prev, o_hist, f, g = imag_step(S, o_hist, a_prev, f, xi, xi_n, g)
    Ss.append(S); As.append(a_prev)   # 每一步都保留计算图

r, d = reward_term(S_all, f_all, flat(As), flat(Aps))   # 解析奖励与软终止
alive = torch.cumprod(1 - d, 0)                          # c_{t+1} = c_t (1 - d_{t+1})
c_t = torch.cat((torch.ones(1, B, device=dev), alive[:-1]), 0)
J = (disc * c_t * r).sum(0)                              # 折扣累计奖励
J = J + (args.gamma ** HZ) * alive[-1] * critic(_critic_input(S, f, xi_n, g))  # 时域外由 Critic 补足
loss = -J.mean()
loss.backward()                                          # 一次 BPTT，梯度沿 32 步传回策略
nn.utils.clip_grad_norm_(pi.parameters(), 1.0)           # 先按全局范数裁剪
nn.utils.clip_grad_value_(pi.parameters(), 0.1)          # 再逐元素裁剪

两层梯度裁剪是论文 Table A.5 规定的。BPTT 的梯度要穿过 32 次网络调用，偶尔会出现很大的值，先裁全局范数再逐元素裁剪，相当于给每步更新加了硬上限。复现还多一道保险：只有梯度范数有限时才执行优化器步进。这两道限制与 PPO 的裁剪比率作用不同，PPO 限制的是新旧策略的概率比，这里限制的是梯度本身的幅度，针对的是长链连乘带来的数值不稳定。

和 PPO 的差别在于梯度从哪来。PPO 对每个采样动作只得到一个标量信号，“这个动作比平均好多少”，然后靠大量样本平均出方向。BPTT 得到的是奖励对动作的完整导数向量，29 个关节各自该往哪边调、调多少，一次就算出来。前提是世界模型在局部足够准、想象时域足够短，否则长链上的 Jacobian 连乘会让梯度爆炸或失真，这也是只想象 32 步的原因。

4.3 Critic 负责时域之外的回报

只看 32 步（0.64 秒）的奖励，策略容易短视：被推一把之后，恢复平衡的收益可能要一两秒后才出现，落在想象窗口之外，而这段时间的收益在 BPTT 里根本没有对应的计算图。Critic 的任务就是补上这一段不可微的未来，它在想象轨迹上用广义优势估计的递推目标拟合：

$\hat{R}_t=\hat{r}_{t+1}+\gamma(1-\hat{d}_{t+1})\big[(1-\lambda)V_{t+1}+\lambda\hat{R}_{t+1}\big],\qquad \hat{R}_H=V_H,\ \lambda=0.95$

目标值每条轨迹只算一次，拟合期间保持不变；每次更新做 5 个 epoch，每个 epoch 4 个打乱的小批次。这里的关键是目标固定，如果边拟合边更新目标，回归问题本身会跟着漂移，Critic 就学不稳。这意味着短时域的可微控制和长期的估值判断是两种不同的能力，不能互相替代。这套"短时域 BPTT 加 Critic"的组合来自 SHAC，SuperMimic 改了两处：用软存活权重替代 SHAC 的硬终止，Critic 也没有使用指数滑动平均的目标网络。

在这里插入图片描述

5. 机器人接口：观测、动作与目标

5.1 策略看到什么，输出什么

策略输入分两部分。一是本体感知观测，单帧包括相对默认姿态的 29 个关节角、29 个关节速度、IMU 测得的 3 维基座角速度、基座坐标系下的 3 维投影重力，以及上一步的 29 维动作，合计 93 维，策略使用最近 10 帧历史。每一维都加均匀噪声，例如关节角 $±0.01 \pm0.01$ rad、关节速度 $±0.5 \pm0.5$ rad/s。二是目标，即未来 $K = 10$ 帧的参考动作；BeyondMimic 只用单帧目标，SuperMimic 改成未来序列，让策略能提前看到接下来要做什么。策略不接收任何全局真值状态，这是能部署到真机的前提。

动作本身是 29 维归一化关节设定值，不是直接输出力矩，也不是输出期望加速度，这一点常被按腿足控制的传统栈误读：读者容易以为策略后面还挂着一个 MPC 或全身控制求解器。它通过下式映射为关节位置目标，再由执行层的低增益 PD 控制器跟踪，而低增益又把位置目标"软化"成了力矩代理：

$q^{sp}=\alpha\,a_t+q^0$

其中 $\alpha$ 是每关节的动作缩放， $q^0$ 是默认姿态。复现沿用了 BeyondMimic 的驱动器参数，刚度按 10 Hz 自然频率和反射电机转动惯量计算，阻尼比为 2，缩放取 $\alpha=0.25\,\tau_{\max}/K_p$ 。这里有个容易误读的地方，论文明确说由于 PD 增益很低，这个设定值实际起到的是关节力矩代理的作用，而不是期望位置。

5.2 两个编码器，一个共享解码器

策略网络拆成两段：目标编码器把目标压成 32 维潜变量，用 tanh 限幅；共享解码器以潜变量和 10 帧观测历史为输入，输出动作。默认任务只用运动编码器，结构是 2048-1024-512-256 的 MLP；解码器同样规模；Critic 很小，只有 256-128。解码器最后一层权重乘 0.01，训练开始时输出接近零动作，也就是接近默认站姿。这种拆分是为下游任务准备的：换一种指令形式，只需要训练新编码器，解码器和世界模型都可以复用。

# supermimic_vibed/supermimic/networks.py
class Policy(nn.Module):
    """a = dec(enc(g), o_{t-9:t})；确定性策略，潜变量 z ∈ R^32 由 tanh 限幅"""
    def forward(self, obs_hist, goal):
        z = torch.tanh(run_mlp(self.enc, self.goal_norm(goal)))   # 目标 → 32 维潜变量
        o = self.obs_norm(obs_hist).flatten(-2)                  # 10 帧观测展平
        return run_mlp(self.dec, torch.cat((z, o), -1))          # 解码出 29 维动作