下一嵌入预测使世界模型更强大 Next Embedding Prediction Makes World Models Stronger
用因果Transformer预测下一时刻嵌入，替代像素重建学习世界模型
前置知识
模型基强化学习 (MBRL)
模型基强化学习是一种学习环境动力学模型的范式。与直接从交互中学习策略的无模型方法不同，MBRL 先学习一个世界模型来预测环境的状态转移、奖励等，然后在学到的世界模型中想象和规划，从而大幅提高样本效率。Dreamer 系列是该领域的代表性工作，通过在潜在空间中进行想象式 rollout 来训练 actor-critic。
NE-Dreamer 是一种 MBRL 方法，理解 MBRL 的基本框架（学习世界模型 + 在想象中训练策略）是理解本文的基础。
循环状态空间模型 (RSSM)
RSSM 是 Dreamer 系列世界模型的核心架构。它维护一个确定性循环状态 $h_t$ 和一个随机潜在变量 $z_t$。确定性状态通过 $h_t = f_{\text{rec}}(h_{t-1}, z_{t-1}, a_{t-1})$ 递推更新，潜在变量则有先验 $p_\phi(z_t | h_t)$ 和后验 $q_\phi(z_t | h_t, e_t)$ 两个分布，分别用于想象和训练时的推断。
NE-Dreamer 保留了 RSSM 作为动力学骨架，只替换了表示学习目标，因此理解 RSSM 的结构是理解本文改动的前提。
Barlow Twins 冗余减少损失
Barlow Twins 是一种自监督学习的冗余减少目标。它计算两个嵌入向量之间的交叉相关矩阵 $C$，然后最小化对角线元素与 1 的偏差（鼓励表达性），同时最小化非对角线元素（减少冗余）。损失函数为 $\mathcal{L}_{BT} = \sum_i (1 - C_{ii})^2 + \lambda_{BT} \sum_{i \neq j} C_{ij}^2$。它不需要负样本即可避免表示坍塌。
NE-Dreamer 使用 Barlow Twins 作为下一嵌入预测的对齐损失，这是防止表示坍塌的关键机制。
因果 Transformer
因果 Transformer 是一种序列模型，其中注意力掩码确保每个位置只能关注之前（含自身）的位置，不能看到未来。这种结构使其适合时间序列预测任务，因为预测 $t+1$ 时刻的输出时只能使用 $t$ 及之前的信息。在本文中，它被用于从历史序列 $(h_{\leq t}, z_{\leq t}, a_{\leq t})$ 中预测下一时刻的编码器嵌入。
因果 Transformer 是本文的核心技术创新之一，它使世界模型能够从历史中学习时间预测性表示，而非仅做同时刻匹配。
无解码器世界模型
传统 Dreamer 系列使用像素解码器重建观测图像来学习表示，但重建目标可能让模型把容量浪费在与任务无关的视觉细节上。无解码器方法去掉像素解码器，用其他自监督目标（如对比学习、冗余减少）直接在潜在空间中学习表示，简化了优化管线并提高效率。
NE-Dreamer 属于无解码器家族，理解这类方法的动机和挑战对于把握本文的贡献至关重要。
研究动机
在部分可观测、高维度的强化学习环境中，模型基方法需要学习能支持长程预测和控制的紧凑潜在状态。传统 Dreamer 系列通过像素重建来学习表示，但重建目标引入了沉重的生成负担，优化复杂，且容易将模型容量分配给视觉上详细但与任务无关的方面（如纹理、背景）。去掉解码器后，已有方法（如 R2-Dreamer、DreamerPro）主要使用同时刻匹配目标，这些目标只保证同一时间步的即时一致性，而不保证状态在时间上的预测性。在需要记忆和空间推理的任务中（如 DMLab Rooms），这种不足会导致潜在状态漂移或坍塌，长期结构弱化，严重损害性能。
本文的目标是本文的目标是提出一种新的无解码器世界模型表示学习目标，使学到的潜在状态天然具有时间预测性。具体而言，作者希望在不使用像素重建的情况下，让世界模型从历史信息中预测未来的潜在表示，并通过冗余减少目标对齐预测与真实嵌入，从而在部分可观测环境中获得更强的长程记忆和导航能力。
与已有工作不同的是，已有无解码器方法的核心问题在于它们的表示学习信号只关注同一时刻：要么做同时刻的编码器-解码器匹配，要么做同一帧内的自监督对齐。这种设计在部分可观测环境下是不够的——一个在时间 $t$ 与当前观测对齐的状态，不一定能预测 $t+1$ 会发生什么。NE-Dreamer 的独特切入角度是将表示学习重新定义为因果下一步预测问题：用一个因果 Transformer 从历史中预测下一时刻的编码器嵌入，并用 Barlow Twins 损失对齐预测与停止梯度的目标。这把自监督学习中的冗余减少思想从同帧不变性扩展到了未来预测，显式地对潜在空间施加了时间一致性约束。
核心方法
NE-Dreamer 的整体思路可以分两层来理解。直觉层面：传统 Dreamer 像素重建要求模型记住每一帧长什么样，但很多视觉细节（如背景纹理）与控制无关；NE-Dreamer 改为要求模型预测下一帧的抽象表示长什么样，这样模型自然会保留那些对未来预测有用的、持久的、与决策相关的结构。技术路线层面：NE-Dreamer 保留了 Dreamer 的 RSSM 动力学骨架和基于想象的 actor-critic 训练流程，只替换了世界模型中的表示学习目标。具体来说，它在 RSSM 之上加了一个因果时间 Transformer，从历史序列中预测下一时刻的编码器嵌入 $\hat{e}_{t+1}$，然后用 Barlow Twins 损失将这个预测与真实的下一时刻嵌入 $e^*_{t+1}$ 对齐。世界模型的总损失变为 $\mathcal{L}_{wm} = \mathcal{L}_{rew} + \mathcal{L}_{cont} + \beta_{kl}\mathcal{L}_{kl} + \beta_{ne}\mathcal{L}_{NE}$，其中最后一项就是新的下一嵌入预测损失。
NE-Dreamer 的核心创新是将无解码器世界模型的表示学习从同时刻匹配转变为因果下一步预测。已有方法（如 R2-Dreamer）用 Barlow Twins 在同一时间步上对齐编码器嵌入和潜在状态，这只能保证此刻的一致性；NE-Dreamer 则要求模型从历史 $h_{\leq t}, z_{\leq t}, a_{\leq t}$ 中预测出下一时刻的编码器嵌入 $\hat{e}_{t+1}$，然后与停止梯度的真实嵌入 $e^*_{t+1} = \text{sg}(f_{\text{enc}}(x_{t+1}))$ 对齐。这个看似简单的目标偏移有深刻的影响：它迫使潜在状态保留对预测未来观测有用的信息，从而在潜在空间中引入了时间连贯性。本质上，这把自监督学习中下一词预测的思想从语言模型引入了世界模型的表示学习。
方法步骤详情
NE-Dreamer 的方法分为三个主要步骤。第一步是编码和潜在推断：编码器将观测映射为嵌入 $e_t = f_{\text{enc}}(x_t)$，RSSM 更新确定性状态 $h_t = f_{\text{rec}}(h_{t-1}, z_{t-1}, a_{t-1})$，并定义先验 $p_\phi(z_t | h_t)$ 和后验 $q_\phi(z_t | h_t, e_t)$ 两个分布。第二步是下一嵌入预测和对齐：一个因果时间 Transformer $\mathcal{T}_\theta$ 接收历史序列 $(h_{\leq t}, z_{\leq t}, a_{\leq t})$，输出下一时刻嵌入预测 $\hat{e}_{t+1}$，然后计算 Barlow Twins 损失 $\mathcal{L}_{NE}$ 将预测与停止梯度目标 $e^*_{t+1}$ 对齐。该损失先对小批量内的嵌入做零均值单位方差归一化，计算交叉相关矩阵 $C_{ij}$，然后最小化 $\mathcal{L}_{NE} = \sum_i (1 - C_{ii})^2 + \lambda_{BT}\sum_{i \neq j} C_{ij}^2$。第三步是 actor-critic 训练：与 DreamerV3 相同，在 RSSM 生成的想象轨迹上（horizon $H=15$）训练策略和价值函数，策略梯度通过世界模型反向传播。
技术新颖性
NE-Dreamer 的技术新颖性体现在几个方面。首先是下一步目标偏移的视角：已有无解码器方法（R2-Dreamer、DreamerPro 等）都作用于同一时间步，NE-Dreamer 首次在 Dreamer 框架内引入了跨时间步的预测目标，将表示学习重新定义为因果序列预测问题。其次是因果 Transformer 的引入方式：它不是一个独立的序列模型，而是轻量级地嵌入到 RSSM 训练循环中，复用了已有的状态序列，无需额外的 rollout 或数据增强。第三是将 Barlow Twins 从同帧不变性扩展到未来预测：原始 Barlow Twins 用于同一图像的两个增强视图，NE-Dreamer 则将其用于预测嵌入和真实嵌入的对齐，这是冗余减少思想在时间维度上的自然推广。最后，整个设计极其简洁——只增加了因果 Transformer 和目标偏移，不引入重建、数据增强或辅助正则化，但效果显著。
实验结果
NE-Dreamer 在两个主流 RL 基准上进行了全面评估。在 DMLab Rooms 四个记忆/导航任务上，NE-Dreamer 在匹配的计算和模型容量下（50M 环境步，5 种子，12M 参数）显著优于所有基线。具体而言，在 Rooms Collect Good Objects Train 任务上，NE-Dreamer 最终回报约 10，而 DreamerV3 约 5，R2-Dreamer 约 4；在 Rooms Deferred Effects Train 上，NE-Dreamer 约 50，DreamerV3 约 25；在 Rooms Nonmatching Object 上，NE-Dreamer 约 60，DreamerV3 约 30。消融实验表明，去掉因果 Transformer（w/o transformer）或去掉下一步目标偏移（w/o shift）都会导致性能大幅下降，验证了两个组件的不可或缺性。去掉投影头（w/o projectors）只轻微影响优化速度。在 DMC 20 个任务上（1M 环境步，5 种子），NE-Dreamer 与 DreamerV3、R2-Dreamer、DreamerPro 持平或略优，说明去掉重建不会损害标准连续控制性能。后验解码器诊断显示，NE-Dreamer 的潜在表示在时间上保持一致——物体身份和空间布局随时间稳定保留；而 Dreamer 和 R2-Dreamer 的表示存在时间不一致问题，任务相关属性会在后续时间步消失。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| DMLab Rooms Collect Good Objects Train | 平均回报 (50M env steps) | ~10 | DreamerV3 ~5, R2-Dreamer ~4 | 约 2x 提升 |
| DMLab Rooms Deferred Effects Train | 平均回报 (50M env steps) | ~50 | DreamerV3 ~25 | 约 2x 提升 |
| DMLab Rooms Nonmatching Object | 平均回报 (50M env steps) | ~60 | DreamerV3 ~30 | 约 2x 提升 |
| DMC 20 tasks 平均 | 任务均值回报 (1M env steps) | 与 DreamerV3 持平 | DreamerV3, R2-Dreamer, DreamerPro | 无退化，略有提升 |
局限与改进
作者在论文中明确指出了几个局限性。首先，实验聚焦于长期结构而非精细视觉细节为主要挑战的环境，尚不清楚无解码器预测目标能否在视觉复杂的高保真任务中匹敌重建方法。其次，NE-Dreamer 使用的 Barlow Twins 只是众多冗余减少目标中的一种，论文指出任何鼓励表达性和非退化解的对齐损失都可以替代，但未系统探索其他选择。从独立观察来看，DMLab 的提升主要集中在记忆密集型任务上，在简单控制任务上优势不明显；因果 Transformer 增加了一定的计算开销；实验规模（12M 参数，50M 步）相对较小，更大规模下的表现有待验证；此外，论文未讨论超参数敏感性分析。
独立分析的弱点
NE-Dreamer 存在几个可以改进的弱点。第一，实验验证范围有限：DMLab 只测试了 4 个 Rooms 任务，DMC 20 个任务中很多已接近饱和，难以体现差异；未来应在一个更广泛、更有挑战性的基准上测试。第二，Barlow Twins 损失的选择是经验性的，作者也承认其他对齐损失可能同样有效；系统比较 VICReg、BYOL、SimSiam 等替代目标是有价值的研究方向。第三，因果 Transformer 的设计（2 层，256 维，4 头）是固定的，缺乏对模型容量的敏感性分析；在更大模型或更长序列上，架构选择可能需要重新调整。第四，论文未与其他非 Transformer 的序列模型（如 LSTM、SSM）进行对比，因果 Transformer 的必要性未被完全隔离。第五，计算开销的量化缺失——Transformer 增加了多少训练时间、推理延迟？这些对实际部署很重要。
未来方向
作者提出了一些未来方向：探索替代的对齐损失，测试在视觉复杂领域中预测目标能否匹敌重建。基于现有成果，还可以延伸多个方向：(1) 将下一嵌入预测推广到多步预测（latent overshooting），论文提到因果 Transformer 天然支持这一点，但未充分实验；(2) 在更大规模的世界模型（如百亿参数级）中验证该方法的可扩展性；(3) 探索与数据增强的结合——NE-Dreamer 不需要增强，但增强是否能进一步提升？(4) 将该框架应用于真实机器人控制或自动驾驶等实际场景；(5) 研究预测性表示是否能改善探索行为——如果状态能预测未来，agent 是否能更有效地发现新状态？
复现评估
从复现角度看，有几个积极信号和挑战。积极方面：论文基于 R2-Dreamer 的统一 PyTorch 代码库，所有超参数与 DreamerV3 一致（Table 1 完整列出），作者承诺使用统一协议和 5 个随机种子。训练设置清晰——DMC 1M 步，DMLab 50M 步，12M 参数的 Dreamer-S 架构，16 个并行环境实例。挑战方面：论文未明确说明代码是否开源；DMLab 环境的安装和配置相对复杂；50M 步的训练在单 GPU 上可能需要数天；Barlow Twins 的实现细节（归一化方式、梯度停止位置）对结果可能敏感。总体而言，复现难度中等偏高，主要受限于计算资源和环境配置。
论文图表