PAPER DEEP DIVE
分块式基于模型的策略优化
CMPO 将策略与世界模型的时间粒度统一为动作块和潜状态块，使多步想象轨迹只需一次前向计算即可完成，同时保留可微价值展开与闭环控制。方法在 HumanoidBench 的 19 个高维人形任务上验证训练效率和梯度稳定性，并在 Unitree Go2 双足站立任务上完成仿真到实机部署。
论文地址：arXiv:2610.14564。项目页：CMPO Project。代码：song-wx/CMPO。本文关注一个看似简单却会改变训练图结构的问题：既然动作分块已经能减少策略调用，为什么世界模型仍然要一次预测一步，再沿着时间反复调用？CMPO 的答案是让策略和动力学模型同时输出长度为 H 的分块，使整段想象轨迹只经过一次模型前向。
这篇工作建立在 TD-MPC2 的潜空间模型之上，但它并不是简单地把动作向量拼接得更长。作者把“动作分块”提升为策略与模型的联合建模约束：策略一次给出未来动作序列，模型一次给出对应未来潜状态序列。训练时，这个分块模型充当可微评估器；部署时，策略仍只执行分块中的第一个动作，以保留每一控制周期重新观测、重新决策的闭环能力。
1. 问题背景：递归想象为何成为瓶颈
基于模型的强化学习通常学习环境转移和奖励的近似，再利用学习到的模型生成想象轨迹。与纯无模型方法相比，这种方法能提高样本效率，因为它把真实环境交互之外的计算变成了可重复使用的训练信号。但是，一条长度为 H 的想象轨迹如果由单步模型递归生成，就必须串行执行 H 次模型调用，时间维上的依赖使这些调用无法完全并行。
递归不只带来延迟。若一步预测存在误差，预测状态又成为下一步输入，误差会沿轨迹累积。策略梯度若要穿过这条轨迹，还需要反向传播经过一串转移雅可比矩阵的乘积。这种深层复合算子容易产生梯度振荡，甚至让策略更新被长时程 BPTT 的数值不稳定性主导。CMPO 的目标不是回避未来预测，而是缩短计算图中沿时间方向的递归深度。
作者的联合分块预测可以概括为：
$$a_{t:t+H-1}=\pi_\theta(z_t),\qquad z_{t+1:t+H}=f_\phi(z_t,a_{t:t+H-1})$$
这里，策略输入当前潜状态，输出整个动作块；世界模型输入同一个当前潜状态和整段动作，输出所有未来潜状态。模型仍然学习未来，但不再让第 k 步输出成为第 k+1 次网络调用的输入。
2. 方法总览：一次调用生成整段价值展开
标准价值展开先基于模型生成轨迹，再累加多步奖励并用终端价值截断。CMPO 保留这个目标，但改变轨迹的生成方式。给定当前潜状态 z_t，分块策略产生动作，分块模型产生后续潜状态，奖励头继续给出逐步奖励，Q 网络对末端状态动作对估值。最终训练信号仍是多步价值展开，只是整段前向图由一次模型调用构造。
论文中的基础目标写为：
$$\hat G_t^{(H)}=\sum_{k=0}^{H-1}\gamma^k R(z_{t+k},a_{t+k})+\gamma^H Q(z_{t+H},\pi(z_{t+H}))$$
代码实现采用少一步的代理目标。动作块含 H 个动作，模型输出 H 个后续潜状态，策略使用前 H-1 个动作对应的奖励，并用第 H-1 个状态动作对作为 Q 终端：
$$G_t^{(H-1)}=\sum_{k=0}^{H-2}\gamma^kR(z_{t+k},a_{t+k})+\gamma^{H-1}Q(z_{t+H-1},a_{t+H-1})$$
两者不是矛盾，而是论文目标与公开代码在索引上的实现差异。理解这一点很重要：阅读结果时应看代码实际的 rollout 形状，而不是只按公式里的终端状态编号推断。
flowchart LR
S["真实观测 s_t"] --> E["编码器 h"]
E --> Z["当前潜状态 z_t"]
Z --> P["分块策略 π"]
P --> A["动作块 a_t 至 a_t+H-1"]
Z --> D["分块动力学 f"]
A --> D
D --> ZS["未来潜状态块 z_t+1 至 z_t+H"]
ZS --> R["奖励头 R"]
A --> R
ZS --> Q["评论家 Q"]
A --> Q
R --> G["价值展开目标"]
Q --> G
G --> U["更新分块策略"]
3. 误差与梯度：算法真正解决的是什么
论文用 Lipschitz 假设描述递归误差。若单步模型的局部误差上界为 epsilon，动力学相对潜状态的 Lipschitz 常数为 L_f，那么第 k 步预测误差可被放大为：
$$\left\|\hat z_{t+k}-z_{t+k}\right\|\le\epsilon\sum_{i=0}^{k-1}L_f^i$$
这个式子说明，递归结构的风险不是“多算了几步”这么简单。只要 L_f 大于 1，误差上界就随预测时域迅速增大。分块模型没有消除预测误差，但它把多步预测变成一次条件映射，减少误差通过网络调用链逐级反馈的机会，同时也改变了梯度沿时间传播的方式。
在单步递归模型中，后续潜状态对策略参数的梯度包含转移雅可比与策略雅可比的乘积：
$$\frac{d\hat z_{t+k+1}}{d\theta}=\left(\frac{\partial f}{\partial z}+\frac{\partial f}{\partial a}\frac{\partial\pi_\theta}{\partial z}\right)\frac{d\hat z_{t+k}}{d\theta}$$
当时域变长时，这种复合雅可比乘积容易出现范数爆炸或消失。CMPO 的分块动力学仍允许动作梯度影响未来状态，因此不是把模型梯度完全截断。代码通过参数 detach 让模型参数不吸收策略损失梯度，同时保留动作输入到模型输出的可微路径。这个细节使“模型作为评估器”与“策略通过预测价值学习”能够同时成立。
4. 三个损失与一次联合更新
世界模型训练包含潜状态一致性、奖励预测和价值预测。论文的目标可概括为：
$$\mathcal L_{\text{model}}=\mathbb E\left[\sum_{t=0}^{H}\lambda^t\left(\left\|z_t-\operatorname{sg}(h(s_t))\right\|_2^2+\operatorname{CE}(\hat r_t,r_{t+1})\right)\right]$$
其中 sg 表示停止梯度，分类交叉熵用于分布式的奖励预测。评论家沿分块时域拟合目标值：
$$\mathcal L_{\text{critics}}=\mathbb E\left[\sum_{t=0}^{H}\lambda^t\operatorname{CE}(\hat q_t,q_t)\right]$$
策略损失则直接最大化价值展开，并加入熵正则：
$$\mathcal L_{\text{policy}}=\mathbb E\left[-\alpha\left(\sum_{i=0}^{H-2}\gamma^iR(z_i,a_i)+\gamma^{H-1}Q(z_{H-1},a_{H-1})\right)-\beta\mathcal H(\pi(\cdot|s))\right]$$
公开配置中，一致性系数为 20，奖励与价值系数均为 0.1，时序衰减 rho 为 0.5。潜维度和 MLP 维度都是 512，使用 5 个 Q 头。较大的一致性系数表明，潜空间预测质量是后续价值展开稳定性的基础。
| 参数 | 数值 | 作用 |
|---|---|---|
| 潜状态维度 | 512 | 编码器、动力学和策略共享的潜表示宽度 |
| MLP 维度 | 512 | 动力学、奖励、策略和 Q 头隐藏层宽度 |
| Q 集成数量 | 5 | 降低单头价值估计偏差 |
| 一致性系数 | 20 | 强化多步潜状态预测约束 |
| 奖励 / 价值系数 | 0.1 / 0.1 | 平衡辅助预测与价值学习 |
| 时序衰减 rho | 0.5 | 对更远的预测误差施加折扣 |
5. 代码实现：分块是网络输出维度，也是控制接口
官方仓库提交 380ef58 的 common/world_model.py 把分块直接写进动力学与策略的输入输出维度。动力学输入是当前潜状态加全部动作，输出是全部未来潜状态；策略输出均值和对数标准差，因此最后一维还要乘 2：
# common/world_model.py:15-26
self._dynamics = layers.mlp(
cfg.latent_dim + cfg.action_dim * cfg.horizon,
2 * [cfg.mlp_dim],
cfg.latent_dim * cfg.horizon,
act=layers.SimNorm(cfg),
)
self._pi = layers.mlp(
cfg.latent_dim, 2 * [cfg.mlp_dim], 2 * cfg.action_dim * cfg.horizon
)真正的一次调用发生在 next()。动作块先被展平，与当前潜状态拼接，然后只调用一次动力学 MLP；输出再按时域 reshape。这与“反复执行单步动力学”有本质区别：时间维是同一层网络的输出维度，而不是网络调用次数。
# common/world_model.py:122-127
def next(self, z, a, detach=False):
"""Predicts all horizon latent states from one latent state and action chunk."""
z = torch.cat([z, a], dim=-1)
return self._dynamics(z) if not detach else self._detach_dynamics(z)部署接口 cmpo.act() 仍然只返回第一个动作。即便策略预测了整块动作，机器人每个控制周期都会重新编码当前观测、重新生成动作块，再取其首项执行。因此 CMPO 的“快”来自网络调用更少，而不是把机器人变成长时间开环执行。
# cmpo.py:76-86
@torch.no_grad()
def act(self, obs, t0=False, eval_mode=False):
obs = obs.to(self.device, non_blocking=True).unsqueeze(0)
z = self.model.encode(obs)
action, info = self.model.pi(z)
if eval_mode:
action = info["mean"][..., : self.cfg.action_dim]
return action[0][..., : self.cfg.action_dim].cpu(), None策略更新中的 detach 逻辑也值得单独解释。_estimate_value(..., detach=True) 让模型参数不接收策略梯度的累计，但动作通过动力学输入仍然可微，所以策略仍能获得“改变动作会怎样改变未来状态和回报”的路径导数。若误以为 detach 只是简单地切断整张图，就会低估该设计的精细程度。
# cmpo.py:170-184
action, info = self.model.pi(zs[0])
chunk_action = action.reshape(
-1, self.cfg.horizon, self.cfg.action_dim
).permute(1, 0, 2)
# Detached model parameters still propagate gradients through actions.
scores = self._estimate_value(zs[0], chunk_action, detach=True)
self.scale.update(scores)
scores = self.scale(scores)
pi_loss = (-(self.cfg.entropy_coef * info["scaled_entropy"] + scores)).mean()
pi_loss.backward()6. HumanoidBench：19 个高维任务上的结果
实验使用 HumanoidBench 中 19 个全身控制任务，机器人为 Unitree H1hand，动作空间为 61 维。任务覆盖行走、奔跑、站立、楼梯、滑面、跨栏、迷宫、平衡、开门、搬箱子、书架、篮球、窗户和勺子等场景。这个集合既有连续运动，也有躯干与手部协同操作，能检验同一算法在长时程控制和精细交互上的适应性。
论文的结论不是“每个任务都绝对第一”，而是 CMPO 在多数任务上具有更好的平均性能和样本效率。这个表述更可信，也更符合高维人形控制任务的差异性：不同任务对接触、稀疏奖励和动作平滑的要求并不一致，一个只使用直接策略推理的模型梯度方法很难在全部场景中无条件领先。
作者在一张 RTX 4090 上比较了 20 万训练步的训练墙钟时间和单次推理延迟，且未使用 Torch Compile 等优化。结果如下：
| 指标 | Walk / TD-MPC2 | Walk / CMPO | Run / TD-MPC2 | Run / CMPO |
|---|---|---|---|---|
| 20 万训练步耗时 | 3.57 小时 | 1.74 小时 | 3.54 小时 | 1.72 小时 |
| 单次推理耗时 | 40.8 毫秒 | 1.4 毫秒 | 39.5 毫秒 | 1.3 毫秒 |
训练时间约减半，推理时间约降低到原来的三十分之一。需要区分这两个收益的来源：训练端来自分块模型减少递归调用与相关优化开销；推理端则主要来自不再执行 TD-MPC2 式 MPPI 在线规划，而是直接输出动作。CMPO 的主要实验默认不加测试时规划，因此这个比较同时包含了“分块模型”和“无规划推理”的差异。
7. 消融：只把策略做成分块还不够
作者强调策略和世界模型必须协同分块。如果策略输出动作块，但动力学仍然只是单步模型，模型就必须沿时间递归调用，之前希望去掉的深层计算图又回来了。反过来，如果动力学一次输出未来状态，而策略每次只能给出一个动作，则未来状态块缺少完整动作条件，问题本身就不成立。
| 方法 | Walk | Run | Stair | Slide |
|---|---|---|---|---|
| CMPO | 954.1 | 629.2 | 60.9 | 265.8 |
| CMPO + MPC | 931.1 | 864.5 | 71.9 | 391.2 |
| 去掉分块世界模型 | 22.2 | 9.6 | 37.1 | 36.1 |
去掉分块模型后，四个任务都出现明显下降。尤其 Walk 和 Run 从 954.1、629.2 掉到 22.2、9.6，说明问题不是策略输出不够长，而是单步模型仍迫使训练和评估穿过递归复合。Stair 和 Slide 的降幅较小，说明不同任务对 BPTT 病理的敏感程度并不相同，这也提醒读者不要把这个消融理解成在所有环境中都有同等幅度。
加入 MPC 后，Run、Stair 和 Slide 有提升，Walk 略降。论文对此的解释是，MPPI 的采样式探索对尚未饱和的任务有帮助，而较简单任务仅靠分块模型的梯度信息就能达到饱和。CMPO 与 MPC 是正交设计，联合使用不是“算法失灵后的补丁”，而是以额外推理计算换取探索能力的另一种权衡。
时域消融给出一个必要的限制：更大的 H 并不自动更好。分块模型仍然要做直接的多步预测，时域越长，单个网络承担的未来关系越复杂。论文报告最佳性能出现在 H=2 与 H=4，主实验为与 TD-MPC2 公平比较统一取 H=3，而增大到 H=5 时性能下降。这个结果准确揭示了方法的边界：它去掉递归时间深度，但没有消除长时程预测本身。
8. Unitree Go2 实机：从仿真站立到真实部署
仿真之外，作者在 Unitree Go2 四足机器人上训练双足站立与速度跟踪。策略控制全部 12 个驱动关节，通过 PD 控制器执行残差关节位置命令。奖励包含线速度与角速度跟踪、机身高度、直立姿态、关节加速度、功率、碰撞、动作变化、动作平滑和髋关节正则项。
实机评估要求机器人维持双足站立 20 秒。CMPO 在 10 次试验中成功 7 次，SAC 成功 4 次。作者选择 SAC 作为硬件对比基线，是因为 TD-MPC2 的迭代式测试时规划不适合同样的低延迟 TorchScript 前馈部署接口。因此，这组对比证明了 CMPO 的直接推理链路能落地，但不能被解释为“CMPO 已经全面优于所有模型方法”。
从工程角度看，这项实机结果最值得关注的是接口形态。训练时模型负责生成未来、评估动作块；部署时策略是一条直接前馈路径，不需要每个控制周期在线采样数百条候选轨迹。对需要低延迟反馈的机器人，这种训练与部署的计算解耦比单纯报告一个成功率更重要。
9. 局限与批判性判断
作者明确承认，第一，模型主体架构基本继承 TD-MPC2，真正针对分块建模重新设计的网络结构仍未探索。当前方法更多是在已有强基线上改变时间组织方式，而不是从头提出一套分块专用架构。第二，策略虽然输出动作块，但在动态环境中只执行第一个动作，剩余动作被丢弃；在低动态环境中，这可能造成明显的动作块利用率不足。
独立来看，还有三点需要谨慎解读。实机证据只有 Go2 双足站立这一项任务，且每类方法 10 次试验，样本规模不足以覆盖不同地面、负载和扰动。仿真结果虽然覆盖 19 个任务，但主要实验训练预算统一为 20 万步，在更长训练后基线是否追平仍需观察。最后，分块预测降低了递归深度，却没有从根本上解决接触突变、分布外状态和长时程预测误差，时域增大时性能退化就是直接证据。
因此，更准确的定位是：CMPO 为“可微模型策略优化”提供了一种更浅、更快、更稳定的时间计算图。它适合直接前馈部署，也能与 MPC 组合；但它不是对世界模型学习难点的终结，也不意味着动作块内部所有动作都应被开环执行。
10. 结论
CMPO 的核心贡献是把策略和世界模型的时间粒度统一起来，让一个动作块和一个潜状态块在一次网络调用中形成可微价值展开。这个设计同时改善了训练效率和推理延迟，并在 HumanoidBench 与 Unitree Go2 上给出了实证支持。代码也清楚展示了两个关键边界：模型参数被 detach 以避免策略梯度污染，但动作路径仍保留梯度；部署时只执行动作块首项，以保持闭环控制。
如果要把这项工作继续推进，最自然的方向不是盲目增大分块时域，而是研究分块专用网络结构、随环境动态调整执行时域，以及在高接触变化任务中更可靠地约束多步预测。论文最有价值的启示可以浓缩为一句话：CMPO 去掉的是循环计算图，不是对未来建模的需求。
