Research map / Reinforcement learning / Offline RL

Transformers Provably Implement In-Context Reinforcement Learning with Policy Improvement

NeurIPS 2026 · Reinforcement learning › Offline RL · approach: Decision Transformer sequence models · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Shows a one-layer linear-attention block can implement batch semi-gradient SARSA and actor-critic updates via explicit constructions, and proves local exponential convergence of gradient flow to the optimal parameter manifold under richness assumptions, validated on random tabular MDPs.

Problem

Whether transformers can implement policy-improvement RL algorithms (SARSA, actor-critic) in context, and whether gradient training provably reaches the parameters that realize them, not just the existence of such parameters.

Approach

Extends prior policy-evaluation (TD) analysis of linear attention to policy improvement by having the block directly output updated critic (and actor) parameters, then proves a local Polyak-Lojasiewicz inequality on a tube around the scaling-symmetric optimal manifold, giving exponential convergence of gradient flow under richness assumptions on training MDPs.

Key results

• Learned P and V matrices of the trained one-layer transformer recover the non-zero structure of the theoretical SARSA and actor-critic constructions (up to scaling) while other blocks stay near zero.
• Training loss decays approximately exponentially to zero in both SARSA and actor-critic experiments.
• On 100 held-out MDPs the transformer's closed-loop return closely tracks the analytical teacher and approaches the oracle greedy policy, with the remaining gap shared by the teacher.

Surprising findings

• The remaining gap to the oracle is attributed to linear function approximation, since the analytical teacher shows the same gap.

Limitations

• Richness assumptions on training MDPs (Assumption on excitation and no destructive cancellation) are hard to verify and only argued generically; verification deferred.
• Convergence is local, requiring initialization near the optimal manifold; global landscape not analyzed.
• Analysis is for deterministic gradient flow in a single linear-attention block, not stochastic optimization or softmax/multi-layer architectures.

Experimental setup

Datasets: randomly generated tabular MDPs (Dirichlet transitions, uniform rewards)

Benchmarks: 100 held-out random tabular MDPs

Metrics: discounted return parameter structure recovery (learned P and V blocks) training loss decay

Baselines: analytical SARSA teacher analytical actor-critic teacher oracle greedy policy (value iteration) uniform random policy

Builds on: Adam Transformers learn in-context by gradient descent linear attention policy evaluation via TD (prior ICRL theory) Algorithm distillation

Closest papers by meaning

• Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining (ICLR 2024 · similarity 0.92)
• Transformers Can Learn Temporal Difference Methods for In-Context Reinforcement Learning (ICLR 2025 · similarity 0.92)
• In-Context Reinforcement Learning From Suboptimal Historical Data (ICML 2025 · similarity 0.91)
• Towards Provable Emergence of In-Context Reinforcement Learning (NeurIPS 2025 · similarity 0.90)
• Distilling Reinforcement Learning Algorithms for In-Context Model-Based Planning (ICLR 2025 · similarity 0.89)
• Optimality and NP-Hardness of Transformers in Learning Markovian Dynamical Functions (NeurIPS 2025 · similarity 0.89)