用于潜在世界模型长期预测的Rollout-Decoded Reconstruction

# 用于潜在世界模型中长时程预测的展开解码重建
来源：https://arxiv.org/html/2608.25017
Rishi Shah，Rishav Shrestha，机器学习工程师，首席技术官，E3A Healthcare，E3A Healthcare，rishishah994@gmail\.com，rishav@e3ahealth\.com

摘要

- 潜在世界模型在训练时，其解码器基于锚定于观测的潜在表示进行训练，然后在评估时，部署于模型自身的自由运行展开中，这可能超出最后观测数百步。展开解码重建通过一个单一的损失项弥合了这一差距：该项在训练期间以与评估完全相同的方式让模型自由运行，解码每一个展开的潜在状态，并针对真实值惩罚重建误差。该方法不增加任何参数，仅在训练时增加计算开销，且在权重为零时退化为标准目标。因此，本文中的所有比较都是仅改变一个标志位的A/B测试。在混沌的Kuramoto–Sivashinsky方程上，RDR将有效预测时间（即归一化误差首次超过0.5的时间）从3.87±0.23提升至6.97±0.42个时间单位，参数量相同（均为193,568），实现了1.80倍的提升，该结果在未用于选择的随机种子上得到验证，并且在10个预注册配置中的10个里均成立，提升比例为1.71–2.50倍。所有结果来自单一系统；一项扫描表明优势随潜在宽度增长而扩大是描述性的，并且在两个经典任务上的控制实验是初步的。

## 1 引言

潜在世界模型通过自由运行进行预测：编码器将最后的观测压缩为一个潜在状态，转换模型根据其自身的输出迭代该潜在状态向前推进，解码器将每个展开的潜在状态转换为预测的观测[9, 8, 10, 11]。然而，解码器仅在锚定于观测的潜在表示上进行训练：即编码器的输出，以及从这些输出单步教师强迫预测的输出。标准目标中没有任何内容将解码质量与长时间展开中漂移后的潜在状态联系起来，因此，一个在观测附近表现清晰的解码器，在其本应解码的展开中，可能在数步之后就失效。PlaNet[9]命名了直接的修复方法——在训练期间解码展开，但将其搁置；此后没有潜在世界模型将其作为研究对象（第2节）。

我们提出了展开解码重建，RDR，一种直接弥合此差距的损失项。在训练期间，模型从编码的初始状态开始自由运行，与评估时的展开方式完全相同；每个展开的潜在状态都被解码；并且针对真实值惩罚重建误差。该方法不增加任何参数，不改变任何架构，并且在其权重λ为零时退化为标准目标。因此，每个预注册的比较都是在固定数据、种子、预算和参数量下进行的A/B测试，仅在一个标志位上有所不同；增加的开销仅为训练时的计算（附录C）。

在Kuramoto–Sivashinsky方程上，一个具有精确真实值和标准长时程指标的混沌偏微分方程，RDR将有效预测时间（VPT，即归一化误差首次超过0.5的时间）从3.87±0.23提升至6.97±0.42个时间单位：在相同的193,568个参数下实现了1.80倍的提升，在全新的随机种子上以及在两个评估时长上均得到验证，并且方向在10个预注册配置中的10个里均成立，比例为1.71–2.50倍（第4.2节）。该效应与训练分布解释一致：一个在相同目标上额外增加25.7%参数的变体在10个配置中的7个表现更差，这排除了增加容量作为解释的可能性，并且两组的展开潜在表示都严重偏离了后验分布，而RDR在匹配的距离下进行解码，误差更低（第5节）。一项潜在宽度扫描表明，随着潜在宽度增加而标准目标丢失时程时，优势会增长；该趋势是描述性的（第5.3节）。初步的控制实验表明，RDR对规划器与训练展开不匹配的情况更稳健，并且一个优化器步数匹配的控制组基本消除了固定轮次的差距（第6节）。在匹配的预算下，一个观测空间预测器在这个完全可观察的系统上达到了相同时程；该比较限定了潜在瓶颈本身，而RDR对比始终在潜在空间内进行（第4.3节）。

RDR是一种训练目标。对称性约简、离散潜在表示、KL平衡和潜在过冲都保留了解码器，且其训练分布不变，因此RDR可应用于这些方法之上；其改进是否迁移到这些场景尚未测试。每个定量声明都追溯到在预注册协议下产生的归档评估制品，每个图表都通过一个脚本渲染，该脚本会重新核实所有引用的数字与这些制品一致。

## 2 背景

#### 潜在世界模型。

训练期间，观测数据可用，因此潜在状态是锚定于数据的*后验*状态；存在的多步训练信号仅在潜在空间中作用。部署时，模型进行*自由运行*：转换模型从编码的初始状态开始，根据其自身输出进行迭代，解码器应用于由此产生的展开潜在状态。RSSM系列[9, 8, 10, 11]始终在后验潜在状态上训练解码器；PlaNet[9]考虑了解码展开的变体（观测过冲），并认为在图像域中计算成本过高，而Dreamer系列[8, 10, 11]未再研究。无解码器模型（TD-MPC2[12]通过值预测，MuDreamer[3]通过删除重建损失）是RDR无法触及的唯一类别，因为它们移除了RDR训练的组件。

#### 有效预测时间。

混沌系统的长时程预测通过VPT进行评分：模型自由运行，解码预测的RMSE针对气候学标准差进行归一化，VPT是该归一化误差首次超过阈值（通常为0.5）的时间，时长以李雅普诺夫时间为单位校准。对于L=22的KS方程，我们使用文献值λmax≈0.043[7]，因此一个李雅普诺夫时间为23.26个时间单位（tu）。该领域已发表的结果（对称性约简流形模型[17]，全状态储层计算机[18]）是与研究目标正交的架构结果；附录D将所有组置于该绝对尺度上。

#### 相关工作。

暴露偏差文献处理训练/部署在*输入*端的不匹配：计划采样[1]将模型预测混合到教师强迫输入中，教授强迫[15]通过对抗方式对齐隐藏状态分布。观测空间PDE代理（前向传播技巧[2]，求解器在回路中的训练[23]）展开模型并在其自身输出上训练，但它们没有潜在空间，因此没有受不匹配影响的解码器。像素空间视频模型（Self Forcing[13]，Diffusion Forcing[4]，Next Forcing[24]）在自身展开上训练，没有编码器-潜在-解码器三元组。两篇2026年的同期预印本，NeuroWorld[6]和Koopman Dreamer[16]，在冻结动力学或辅助损失设置中解码展开状态；两者都没有将解码器的训练分布孤立为研究对象。最接近的是Rollout-LaSDI[22]，它在一个潜在降阶模型上，针对参考解训练解码器来解码展开的潜在状态，并将二维Burgers方程上的最大误差削减了三分之二。它在关键点上与此不同：它的潜在动力学是通过回归拟合的、针对每个参数的线性系数，其系统是光滑的，它未报告控制结果，并且它固定了潜在宽度；本研究研究的是在一个混沌系统上，一个学习到的非线性转换模型在自身输出上迭代。展开解码训练在潜在降阶模型中是确立的，据我们所知，在世界模型中尚未测量过；本研究孤立了它并测量了其价值。

## 3 方法

### 3.1 模型与基础目标

模型是一个精心设计的普通编码器-转换-解码器三元组。一个MLP编码器Eφ将场快照ut∈R映射为一个潜在表示zt∈R（两个隐藏层，每层256，GELU激活）。一个状态空间预测器fθ，一个状态大小为64的四层S5堆栈[21]，将潜在状态推进一步。一个解码器Dψ，一个隐藏层为512的MLP，将潜在表示映射回场空间。一个目标编码器̄E，是Eφ的指数移动平均（衰减0.999），提供潜在回归目标；我们写̄zt=̄E(ut)，并借用RSSM实践，将由观测数据计算的潜在状态称为*后验*。

四个标准项训练这个三元组。使用sg表示停止梯度，̂zt+k表示自由运行的展开（̂zt=Eφ(ut)，̂zt+1=fθ(̂zt)，梯度流经整个链）：

LTF = (1/T) ∑t ‖fθ(̄zt) - sg(̄zt+1)‖²  教师强迫的单步潜在预测，
LR = (1/K) ∑k=1K ‖̂zt+k - sg(̄zt+k)‖²  多步潜在展开一致性，
LOBS = (1/T) ∑t ‖Dψ(fθ(̄zt)) - ut+1‖²  解码教师强迫的预测，
Lrecon = (1/T) ∑t ‖Dψ(Eφ(ut)) - ut‖²  在线编码器上的重建锚定。
没有任何项让解码器接触到距离观测超过一个教师强迫步的潜在状态。

### 3.2 RDR目标

RDR解码与LR约束且与评估评分相同的自由运行展开，并惩罚其与真实场的误差：

LRDR = (1/K) ∑k=1K ‖Dψ(̂zt+k) - ut+k‖², L = LTF + αe LR + LOBS + Lrecon + λ LRDR。 (1)

梯度直接到达解码器，并通过展开链到达预测器和编码器，因此所有三个组件都受到模型将实际产生的轨迹的塑造。课程设置是展开损失的标准做法：αe在第2到5个epoch之间从0线性增加到1，展开分支（以及RDR项）在两个预热epoch之后开启，并且在整个过程中展开都是纯自由运行的。设置λ=0可以完全恢复基线；我们称该组为*仅后验*。操作权重是λ=0.3，结果在λ∈{0.1, 0.3, 0.6, 1.0}中是平坦的（第5.2节）。

该方法不增加任何参数。两组已经为LR计算了自由运行的展开，因此边际成本是每个窗口K次额外的解码器评估及其反向传播，仅在训练时产生；推理没有变化。附录C报告了每组的步数、解码器评估次数、挂钟时间和硬件。

## 4 预测结果

### 4.1 设置

我们使用ETDRK4在dt=0.1下对L=22的KS方程进行积分，使用64个网格点，丢弃2000步的瞬态过程，并从不相交的种子生成512条训练轨迹、64条验证轨迹和64条测试轨迹，每条轨迹有256个快照。训练时，每条轨迹抽取一个固定的160个快照的窗口，并将展开损失展开K=128步；批量大小为64（每个epoch优化器步数为8），使用Adam优化器，学习率为3e-4，每组三个随机种子。主要配置是潜在宽度32，解码器宽度512，λ=0.3，320个epoch。

评估时，从每条保留轨迹的第一个快照开始，对每个模型进行自由运行，使用训练中相同的展开路径，进行解码，并按照第2节的定义对VPT进行评分。规范时长是200步（20 tu）；长时长是1300步（130 tu），在单独生成的、去混叠的记录上进行。附录A给出了完整的指标、长记录生成过程、规范/长时长对背后的固定估计器构建，以及验证集重新排序，确认报告的获胜者可在验证集上选择。

### 4.2 主要结果

表2（附录B）和图1总结了预注册的扫描。本节中的每个数字都是[email protected]：解码自由运行展开的归一化RMSE首次超过0.5的时间，在所有64条保留轨迹上评分，并对每组三个随机种子取平均（第2、4.1节）。在全新的随机种子（10-12）上，未在选择过程中使用，仅后验组达到3.87±0.23 tu，而RDR组达到6.97±0.42 tu：在相同的193,568个可训练参数下实现了1.80倍的提升（在附录B中核算）。相同的检查点在长时长上的得分是3.77对6.90 tu，图2展示了一条保留轨迹。在十行预注册结果（九种配置加上获胜者的新种子重跑）中，RDR在容量匹配的情况下10胜10。