Transformers Provably Implement In-Context Reinforcement Learning with Policy Improvement

• URL: http://arxiv.org/abs/2605.05755v1
• Date: Thu, 07 May 2026 06:50:49 GMT
• Title: Transformers Provably Implement In-Context Reinforcement Learning with Policy Improvement
• Abstract summary: We investigate the ability of transformers to perform in-context reinforcement learning.<n>We show that a linear self-attention transformer block can provably implement policy-improvement methods.<n>We establish the first convergence guarantee in the ICRL literature.
• Score: 23.24714915251971
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• Transformers Efficiently Perform In-Context Logistic Regression via Normalized Gradient Descent [9.440916748352722]
  Transformer with softmax attention perform in-context learning on linear classification data.<n>We first construct a class of multi-layer transformers that can perform in-context logistic regression.<n>We show that our constructed transformer can be obtained through (i) training a single self-attention layer supervised by one-step gradient descent, and (ii) recurrently applying the trained layer to obtain a looped model.
  arXiv  Detail & Related papers  (2026-05-07T17:27:55Z)
• In-Context Reinforcement Learning From Suboptimal Historical Data [56.60512975858003]
  Transformer models have achieved remarkable empirical successes, largely due to their in-context learning capabilities.<n>We propose the Decision Importance Transformer framework, which emulates the actor-critic algorithm in an in-context manner.<n>Our results show that DIT achieves superior performance, particularly when the offline dataset contains suboptimal historical data.
  arXiv  Detail & Related papers  (2026-01-27T23:13:06Z)
• Heuristic Transformer: Belief Augmented In-Context Reinforcement Learning [1.8791091507292152]
  Heuristic Transformer (HT) is an in-context reinforcement learning approach that augments the in-context dataset with a belief distribution over rewards to achieve better decision-making.<n>We show that HT consistently surpasses comparable baselines in terms of both effectiveness and generalization.<n>Our method presents a promising direction to bridge the gap between belief-based augmentations and transformer-based decision-making.
  arXiv  Detail & Related papers  (2025-11-13T12:32:36Z)
• In-Context Learning for Gradient-Free Receiver Adaptation: Principles, Applications, and Theory [54.92893355284945]
  Deep learning-based wireless receivers offer the potential to dynamically adapt to varying channel environments.<n>Current adaptation strategies, including joint training, hypernetwork-based methods, and meta-learning, either demonstrate limited flexibility or necessitate explicit optimization through gradient descent.<n>This paper presents gradient-free adaptation techniques rooted in the emerging paradigm of in-context learning (ICL)
  arXiv  Detail & Related papers  (2025-06-18T06:43:55Z)
• Scalable In-Context Q-Learning [68.9917436397079]
  We propose textbfScalable textbfIn-textbfContext textbfQ-textbfLearning (textbfSICQL) to steer in-context reinforcement learning.<n>textbfSICQL harnesses dynamic programming and world modeling to steer ICRL toward efficient reward and task generalization.
  arXiv  Detail & Related papers  (2025-06-02T04:21:56Z)
• Distilling Reinforcement Learning Algorithms for In-Context Model-Based Planning [39.53836535326121]
  We propose Distillation for In-Context Planning (DICP), an in-context model-based RL framework where Transformers simultaneously learn environment dynamics and improve policy in-context.<n>Our results show that DICP achieves state-of-the-art performance while requiring significantly fewer environment interactions than baselines.
  arXiv  Detail & Related papers  (2025-02-26T10:16:57Z)
• Q-value Regularized Transformer for Offline Reinforcement Learning [70.13643741130899]
  We propose a Q-value regularized Transformer (QT) to enhance the state-of-the-art in offline reinforcement learning (RL) QT learns an action-value function and integrates a term maximizing action-values into the training loss of Conditional Sequence Modeling (CSM) Empirical evaluations on D4RL benchmark datasets demonstrate the superiority of QT over traditional DP and CSM methods.
  arXiv  Detail & Related papers  (2024-05-27T12:12:39Z)
• Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining [25.669038513039357]
  This paper provides a theoretical framework that analyzes supervised pretraining for in-context reinforcement learning. We show transformers with ReLU attention can efficiently approximate near-optimal online reinforcement learning algorithms.
  arXiv  Detail & Related papers  (2023-10-12T17:55:02Z)
• Supervised Pretraining Can Learn In-Context Reinforcement Learning [96.62869749926415]
  In this paper, we study the in-context learning capabilities of transformers in decision-making problems. We introduce and study Decision-Pretrained Transformer (DPT), a supervised pretraining method where the transformer predicts an optimal action. We find that the pretrained transformer can be used to solve a range of RL problems in-context, exhibiting both exploration online and conservatism offline.
  arXiv  Detail & Related papers  (2023-06-26T17:58:50Z)
• Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection [88.23337313766353]
  This work first provides a comprehensive statistical theory for transformers to perform ICL. We show that transformers can implement a broad class of standard machine learning algorithms in context. A emphsingle transformer can adaptively select different base ICL algorithms.
  arXiv  Detail & Related papers  (2023-06-07T17:59:31Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.