Research map / Reinforcement learning / Visual and representation RL
Transformers Can Learn Temporal Difference Methods for In-Context Reinforcement Learning
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 theoretically and empirically that linear-attention transformers trained with multi-task nonlinear TD on random policy-evaluation tasks learn to implement TD (and TD(λ), residual gradient, average-reward TD) in their forward pass, with an invariant-set analysis for the one-layer case.
Problem
Understanding how in-context reinforcement learning (ICRL) for policy evaluation arises from reinforcement pretraining, i.e., which algorithm a pretrained transformer implements in its forward pass and why TD-based pretraining produces it.
Approach
Constructs linear-attention transformer weights that exactly implement batch TD(0) (and other TD variants) layer by layer, then shows multi-task TD pretraining of the transformer on random Boyan-chain MRPs converges toward these weights, and proves a parameter set implementing in-context TD is invariant under the expected TD update for one layer.
Key results
- Learned P0 and Q0 after 4000 MRPs closely match the TD construction with C_l = I_d (P0 bottom-right entry increases while other entries stay near zero).
- Trace of upper-left block of Q0 approaches -d and upper-right block trace approaches d, matching the in-context TD construction.
- Multi-task TD pretraining with L=2,3,4 converges to in-context TD rather than in-context residual gradient, even though residual gradient is also representable.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| train with multi-task TD on randomly generated Boyan's chain policy-evaluation tasks | linear transformer · Boyan's chain | in-context TD emergence | Helped |
| use a task distribution with fixed transition probabilities or fixed features | multi-task TD pretraining · Boyan's chain | in-context TD emergence | Hurt |
| learn through multi-task TD pretraining rather than explicitly implementing residual gradient | linear transformer · Boyan's chain policy-evaluation tasks | learned in-context algorithm | Mixed |
| use a normal transformer without parameter sharing (ablation) | multi-task TD-pretrained transformer · policy-evaluation tasks | in-context TD emergence | No clear effect |
| use softmax-based attention instead of linear attention (ablation) | multi-task TD-pretrained transformer · policy-evaluation tasks | in-context TD emergence | No clear effect |
| train on CartPole-based tasks instead of Boyan's chain-based tasks (ablation) | multi-task TD-pretrained transformer · CartPole | in-context TD emergence | No clear effect |
Surprising findings
- Although transformers can implement residual gradient, pretraining consistently converges to the TD implementation.
- For in-context TD to emerge, the task distribution must be sufficiently diverse; fixing (p0, p) or phi prevented its emergence.
Limitations
- Only policy evaluation is studied, not control.
- Theory relies on linear attention and assumptions (independence and symmetry of features); the invariant set result is for the single-layer case and does not prove convergence.
- Requires the ability to randomly generate policy evaluation tasks, and experiments are limited to Boyan's chain and CartPole.
Experimental setup
Datasets: Boyan's chain (randomly generated MRPs) CartPole
Benchmarks: Boyan's chain CartPole
Metrics: value approximation error parameter convergence (P0, Q0) value difference implicit weight similarity sensitivity similarity
Baselines: residual gradient TD(lambda) average-reward TD
Builds on: in-context learning as gradient descent Boyan's chain TD(lambda) residual gradient
Closest papers by meaning
- Towards Provable Emergence of In-Context Reinforcement Learning (NeurIPS 2025 · similarity 0.94)
- Transformers Provably Implement In-Context Reinforcement Learning with Policy Improvement (NeurIPS 2026 · similarity 0.92)
- Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining (ICLR 2024 · similarity 0.91)
- In-Context Reinforcement Learning From Suboptimal Historical Data (ICML 2025 · similarity 0.88)
- Distilling Reinforcement Learning Algorithms for In-Context Model-Based Planning (ICLR 2025 · similarity 0.88)
- Transformers as Game Players: Provable In-context Game-playing Capabilities of Pre-trained Models (NeurIPS 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.