Next Embedding Prediction Makes World Models Stronger

📄 arXiv: 2603.02765v1 📥 PDF

作者: George Bredis, Nikita Balagansky, Daniil Gavrilov, Ruslan Rakhimov

分类: cs.LG, cs.AI

发布日期: 2026-03-03

💡 一句话要点

NE-Dreamer：基于Transformer的下一嵌入预测增强世界模型

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 模型强化学习 世界模型 时间Transformer 表征学习 嵌入预测

📋 核心要点

  1. 在部分可观察、高维环境中，捕获时间依赖性对于基于模型的强化学习至关重要，现有方法难以有效建模。
  2. NE-Dreamer通过时间Transformer预测下一步的编码器嵌入，直接优化表征空间中的时间预测对齐，无需重建损失。
  3. 实验表明，NE-Dreamer在DeepMind Control Suite和DMLab任务上表现出色，尤其在涉及记忆和空间推理的任务中提升显著。

📝 摘要（中文）

本文提出了一种名为NE-Dreamer的无解码器模型强化学习（MBRL）智能体，它利用时间Transformer从潜在状态序列中预测下一步的编码器嵌入，从而直接优化表征空间中的时间预测对齐。这种方法使NE-Dreamer能够在没有重建损失或辅助监督的情况下学习连贯的、可预测的状态表征。在DeepMind Control Suite上，NE-Dreamer的性能与DreamerV3和领先的无解码器智能体相匹配或超过。在涉及记忆和空间推理的DMLab任务的具有挑战性的子集上，NE-Dreamer取得了显著的提升。这些结果表明，使用时间Transformer进行下一步嵌入预测是复杂、部分可观察环境中MBRL的有效且可扩展的框架。

🔬 方法详解

问题定义：在基于模型的强化学习（MBRL）中，尤其是在部分可观察的环境中，如何有效地学习和利用时间依赖关系是一个关键问题。传统的MBRL方法通常依赖于解码器来重建观测，这可能引入额外的复杂性和信息损失。此外，辅助监督信号也可能限制模型的泛化能力。因此，需要一种更直接、更有效的方法来学习连贯且可预测的状态表征。

核心思路：NE-Dreamer的核心思路是通过预测下一步的嵌入（next embedding prediction）来学习状态表征，而不是像传统方法那样重建观测。这种方法直接在表征空间中优化时间预测对齐，避免了重建过程中的信息损失。使用Transformer来建模时间依赖关系，能够更好地捕捉长期的时间相关性。

技术框架：NE-Dreamer的整体框架包括一个编码器、一个时间Transformer和一个策略网络。编码器将观测映射到潜在状态空间。时间Transformer接收潜在状态序列作为输入，并预测下一步的嵌入。策略网络基于潜在状态做出决策。整个框架通过最大化累积奖励进行端到端训练。

关键创新：NE-Dreamer的关键创新在于使用时间Transformer进行下一步嵌入预测。与传统的基于解码器的MBRL方法相比，NE-Dreamer避免了重建损失，从而能够更有效地学习状态表征。此外，NE-Dreamer不需要额外的辅助监督信号，从而提高了模型的泛化能力。

关键设计：NE-Dreamer使用Transformer作为时间建模的核心组件，Transformer的输入是过去一段时间的embedding序列，输出是下一步的embedding预测。损失函数采用的是预测embedding和真实embedding之间的距离度量，例如余弦相似度或均方误差。策略网络通常是一个简单的多层感知机，输入是当前的状态embedding，输出是动作的概率分布。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

NE-Dreamer在DeepMind Control Suite上取得了与DreamerV3和领先的无解码器智能体相媲美的性能。在DMLab任务的具有挑战性的子集上，NE-Dreamer取得了显著的提升，尤其是在涉及记忆和空间推理的任务中。这些结果表明，NE-Dreamer能够有效地学习状态表征，并利用时间依赖关系做出更好的决策。

🎯 应用场景

NE-Dreamer具有广泛的应用前景，例如机器人导航、游戏AI、自动驾驶等。该方法能够有效地处理部分可观察的环境，并学习长期的时间依赖关系，从而使智能体能够更好地理解和适应复杂环境。此外，NE-Dreamer的无解码器设计使其更加高效和可扩展，适用于资源受限的场景。

📄 摘要（原文）

Capturing temporal dependencies is critical for model-based reinforcement learning (MBRL) in partially observable, high-dimensional domains. We introduce NE-Dreamer, a decoder-free MBRL agent that leverages a temporal transformer to predict next-step encoder embeddings from latent state sequences, directly optimizing temporal predictive alignment in representation space. This approach enables NE-Dreamer to learn coherent, predictive state representations without reconstruction losses or auxiliary supervision. On the DeepMind Control Suite, NE-Dreamer matches or exceeds the performance of DreamerV3 and leading decoder-free agents. On a challenging subset of DMLab tasks involving memory and spatial reasoning, NE-Dreamer achieves substantial gains. These results establish next-embedding prediction with temporal transformers as an effective, scalable framework for MBRL in complex, partially observable environments.