Transformers Provably Implement In-Context Reinforcement Learning with Policy Improvement

📄 arXiv: 2605.05755v1 📥 PDF

作者: Haodong Liang, Lifeng Lai

分类: stat.ML, cs.AI, cs.LG

发布日期: 2026-05-07

备注: 25 pages, 4 figures

💡 一句话要点

证明Transformer可通过参数构造实现上下文强化学习并提供收敛性保证

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

关键词: 上下文强化学习 Transformer架构 策略改进 梯度流动力学 算法内化 少样本学习

📋 核心要点

  1. 核心问题：现有ICRL研究多依赖经验观察，缺乏对Transformer如何内化强化学习算法的理论解释及收敛性保证。
  2. 方法要点：通过显式参数构造证明Transformer可实现SARSA与Actor-Critic，并提出教师模仿训练流程以实现参数收敛。
  3. 实验效果：实验证实模型能学习到预期的参数结构，并在未见过的MDP任务中表现出优异的上下文控制能力。

📝 摘要（中文）

本文研究了Transformer执行上下文强化学习（ICRL）的能力，即模型如何在不更新参数的情况下，通过轨迹数据推断并执行学习算法。研究证明，线性自注意力Transformer块可通过显式参数构造实现策略改进方法，包括半梯度SARSA和Actor-Critic算法。除了存在性证明，作者还设计了一种教师模仿训练流程，并分析了其梯度流动力学，在ICRL文献中首次建立了收敛性保证：在训练MDP分布满足丰富性条件的假设下，梯度流可局部且指数级收敛到对应于目标RL更新的最优参数流形。在随机生成的表格型MDP上的实验验证了这些理论预测，模型不仅恢复了显式构造的参数结构，且在未见过的MDP上展现了强大的上下文控制性能。这些结果揭示了Transformer架构如何内化并执行经典强化学习算法，架起了机制理解与训练动力学之间的桥梁。

🔬 方法详解

问题定义：研究Transformer在无需参数更新的前提下，如何通过上下文轨迹数据实现强化学习算法（如SARSA、Actor-Critic）。现有研究缺乏对模型内部机制的数学解释，以及在训练动力学层面的收敛性证明。

核心思路：利用Transformer的注意力机制模拟经典RL算法的更新步骤。通过构造特定的权重矩阵，使Transformer的输出等价于策略改进算子，从而将学习过程转化为模型的前向推理过程。

技术框架：整体框架包含显式参数构造与教师模仿训练两部分。首先，通过数学推导证明线性自注意力块能够表示策略改进算子；其次，引入教师模仿训练，通过梯度流分析证明模型参数在特定分布下可收敛至最优解。

关键创新：首次在ICRL领域建立了理论收敛性保证，证明了在满足MDP分布丰富性条件下，梯度流可指数级收敛至最优参数流形，为理解Transformer的“算法内化”提供了严谨的数学基础。

关键设计：采用线性自注意力机制作为核心算子，通过特定的权重初始化与教师模仿损失函数（Teacher-mimicking loss），引导模型学习到与经典RL算法一致的参数结构，确保模型在推理阶段能够准确执行策略更新。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验在随机生成的表格型MDP上进行，结果表明：模型成功恢复了理论推导出的参数结构，验证了Transformer对RL算法的内化能力。在未见过的MDP测试中，模型表现出极强的泛化性能，且梯度流分析与实验观察到的收敛行为高度吻合，首次在ICRL领域实现了理论与实证的闭环。

🎯 应用场景

该研究为构建通用智能体提供了理论支撑，特别是在少样本学习、自适应控制及机器人在线决策领域具有重要价值。通过理解Transformer如何内化算法，未来可设计更高效的上下文学习模型，使其在复杂、动态的环境中无需微调即可快速适应新任务，显著提升AI系统的自主决策能力。

📄 摘要（原文）

We investigate the ability of transformers to perform in-context reinforcement learning (ICRL), where a model must infer and execute learning algorithms from trajectory data without parameter updates. We show that a linear self-attention transformer block can provably implement policy-improvement methods, including semi-gradient SARSA and actor-critic, via explicit parameter constructions. Beyond existence, we design a teacher-mimicking training procedure, analyze its gradient-flow dynamics, and establish the first convergence guarantee in the ICRL literature: under suitable richness conditions on the training MDP distribution, gradient flow converges locally and exponentially to an optimal parameter manifold corresponding to the desired RL update. Empirically, training transformers on randomly generated tabular MDPs confirms these predictions: the learned models recover the parameter structure of our explicit constructions and, when deployed on unseen MDPs, deliver strong in-context control performance. Together, these results illuminate how transformer architectures internalize and execute classical reinforcement learning algorithms in context, bridging mechanistic understanding and training dynamics in ICRL.