Research map / Reinforcement learning / Visual and representation RL

Towards Provable Emergence of In-Context Reinforcement Learning

NeurIPS 2025 · Reinforcement learning › Visual and representation RL · approach: Probabilistic and autonomous RL · 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.

Proves that the in-context TD weights of a linear Transformer are a (limiting) global minimizer of the NEU loss for multi-task TD and MC pretraining as depth goes to infinity, and verifies this with Boyan's chain experiments.

Problem

Why do Transformers pretrained with standard RL algorithms (TD or MC policy evaluation) acquire parameters that enable in-context reinforcement learning, and why do those parameters emerge?

Approach

Shows that the known in-context TD parameterization of a linear Transformer is a limiting global minimizer of the NEU loss for multi-task TD and a multi-task MC update, extending prior single-layer invariant-set analysis to multi-layer linear Transformers and showing convergence of the forward pass to the true value as depth grows.

Key results

• Theorem: the linear Transformer's forward pass with in-context TD parameters converges to the true value v(s_q) as depth L goes to infinity.
• Theorem and corollary: the in-context TD parameters have NEU loss J(theta)->0 (for multi-task TD) and J'(theta)->0 (for multi-task MC) as L->infinity under the looped parameter space.
• Experiments on Boyan's chain show MSVE decreases with context length for both multi-task TD and MC, and both converge to the same in-context TD parameter configuration.

Measured effects reported in the paper

Surprising findings

• Multi-task MC pretraining also yields the in-context TD parameter, suggesting the model learns to implement TD rather than imitating the pretraining algorithm.

Limitations

• Analysis restricted to linear attention without softmax or ReLU nonlinearities.
• Prompt construction assumes enumeration of states and expected next features, which is impractical for continuous or large state spaces.
• Empirical study is small-scale (Boyan's chain) and covers only policy evaluation, not control.

Experimental setup

Datasets: Boyan's chain MRPs (randomly generated)

Benchmarks: Boyan's chain

Metrics: mean squared value error (MSVE) vs context length parameter convergence to in-context TD weights

Baselines: multi-task TD

Builds on: in-context TD (Shi et al.) algorithm distillation TD learning gradient TD NEU loss

Closest papers by meaning

• Transformers Can Learn Temporal Difference Methods for In-Context Reinforcement Learning (ICLR 2025 · similarity 0.94)
• Transformers Provably Implement In-Context Reinforcement Learning with Policy Improvement (NeurIPS 2026 · similarity 0.90)
• Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining (ICLR 2024 · similarity 0.90)
• Vintix: Action Model via In-Context Reinforcement Learning (ICML 2025 · similarity 0.89)
• Mixture-of-Experts Meets In-Context Reinforcement Learning (NeurIPS 2025 · similarity 0.89)
• Towards Large-Scale In-Context Reinforcement Learning by Meta-Training in Randomized Worlds (NeurIPS 2025 · similarity 0.88)