arXiv:2511.06946cs.LGcs.AI2025-11中稿 · Embodied World Mod…

给Transformer注意力加规则,让智能体更关注关键历史信息。

Learning to Focus: Prioritizing Informative Histories with Structured Attention Mechanisms in Partially Observable Reinforcement Learning

  • 用高斯分布平滑分配注意力,聚焦对决策有用的历史
  • 在Atari 100k上比基线提升77%的人类归一化得分
  • 适合需要高效利用稀疏观测的强化学习任务

Transformers在建模长时依赖方面表现强劲,被广泛用于部分可观测强化学习中的世界模型。然而,与自然语言语料不同,强化学习轨迹稀疏且以奖励驱动,标准自注意力机制会均匀分配权重,难以突出对控制至关重要的少数状态转移。为此,我们在UniZero的动态头中引入结构化归纳偏置:(i) 每头记忆长度先验,限制注意力于任务特定窗口;(ii) 分布先验,学习过去状态-动作对上的平滑高斯权重。实验表明,在Atari 100k基准上,高斯先验带来主要性能提升,其平滑分配注意力至关键过渡;而记忆长度先验常因过严截断导致有效信号丢失。具体地,高斯注意力使平均人类归一化得分相对UniZero提升77%。结果表明,在具有非平稳时间依赖性的部分可观测环境中,离散记忆窗口难以可靠学习,而平滑分布先验可灵活适应不同时间跨度,实现更强的数据效率。整体而言,将结构化时间先验直接编码进自注意力,能有效提升动态建模中对关键历史的优先级处理能力。

原文摘要 · Abstract (English)

Transformers have shown strong ability to model long-term dependencies and are increasingly adopted as world models in model-based reinforcement learning (RL) under partial observability. However, unlike natural language corpora, RL trajectories are sparse and reward-driven, making standard self-attention inefficient because it distributes weight uniformly across all past tokens rather than emphasizing the few transitions critical for control. To address this, we introduce structured inductive priors into the self-attention mechanism of the dynamics head: (i) per-head memory-length priors that constrain attention to task-specific windows, and (ii) distributional priors that learn smooth Gaussian weightings over past state-action pairs. We integrate these mechanisms into UniZero, a model-based RL agent with a Transformer-based world model that supports planning under partial observability. Experiments on the Atari 100k benchmark show that most efficiency gains arise from the Gaussian prior, which smoothly allocates attention to informative transitions, while memory-length priors often truncate useful signals with overly restrictive cut-offs. In particular, Gaussian Attention achieves a 77% relative improvement in mean human-normalized scores over UniZero. These findings suggest that in partially observable RL domains with non-stationary temporal dependencies, discrete memory windows are difficult to learn reliably, whereas smooth distributional priors flexibly adapt across horizons and yield more robust data efficiency. Overall, our results demonstrate that encoding structured temporal priors directly into self-attention improves the prioritization of informative histories for dynamics modeling under partial observability.

强化学习Transformer注意力机制世界模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。