Title: Towards Provable Emergence of In-Context Reinforcement Learning
Venue: NeurIPS 2025
Accepted: yes
Keywords: in-context reinforcement learning

policy evaluation

transformers
Abstract: Typically, a modern reinforcement learning (RL) agent solves a task by updating its neural network parameters to adapt its policy to the task. Recently, it has been observed that some RL agents can solve a wide range of new out-of-distribution tasks without parameter updates after pretraining on some task distribution. When evaluated in a new task, instead of making parameter updates, the pretrained agent conditions its policy on additional input called the context, e.g., the agent's interaction history in the new task. The agent's performance increases as the information in the context increases, with the agent's parameters fixed. This phenomenon is typically called in-context RL (ICRL). The pretrained parameters of the agent network enable the remarkable ICRL phenomenon. However, many ICRL works perform the pretraining with standard RL algorithms. This raises the central question this paper aims to address: Why can the RL pretraining algorithm generate network parameters that enable ICRL? We hypothesize that the parameters capable of ICRL are minimizers of the pretraining loss. This work provides initial support for this hypothesis through a case study. In particular, we prove that when a Transformer is pretrained for policy evaluation, one of the global minimizers of the pretraining loss can enable in-context temporal difference learning.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates in-context reinforcement learning using Transformers with linear attention. Specifically, the authors establish two theoretical results: (1) In finite state spaces and under a particular sparse structure of the weight matrix, a specially parameterized prompt—equipped with known transition and reward functions—can compute the value function when passed through the Transformer (Theorem 1). (2) The desired weight matrix is shown to be a global minimizer of a pretraining loss function, where gradient descent is performed and the transition and reward components in the prompt are sampled from a prior distribution (Theorem 2). The paper also includes a synthetic experiment that numerically validates these two theorems.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\n- The paper is generally well-structured and easy to follow.\n- The theoretical investigation of in-context reinforcement learning is novel and addresses an emerging problem of interest.\n\nWeaknesses\n\n- The paper assumes a finite state space and full knowledge of the ground-truth transition and reward functions. Under these assumptions, the reinforcement learning task reduces to computing the value function of an infinite-horizon Markov reward process, which is arguably not representative of core challenges in modern RL and can be solved by relatively simple methods.\n- The writing contains some inconsistencies and relies on setups or results from prior work without sufficient explanation. For example, in line 197, the prompt appears to be derived from a trajectory of length $n$, whereas elsewhere $n$ denotes the size of the state space. Additionally, the \"multi-task TD\" setting is introduced and used in experiments without a clear formal definition.\n- As acknowledged in the limitations section, the Transformer architecture and the sparse parameterization considered in this work represent a highly simplified setting that lacks scalability to realistic applications.\n\nQUESTIONS:\n- I understand that by \"multi-tasking,\" the paper refers to sampling different ground-truth transition and reward functions and then computing the value function under each sampled environment. Is this interpretation correct? If so, does this find application in other reinforcement learning studies?\n- In Section 6.1, it appears that the trained Transformer is able to perform policy evaluation accurately even when provided with only a single realization (i.e., a sample trajectory) from the task environment. This seems to differ from the setting described in Equation (4), where the prompt consists of an enumeration of the full state space. I am curious whether this discrepancy has significant empirical or theoretical implications. Specifically, I would expect that unless the stationary distribution is uniform over the state space, such an approximation could introduce bias. Could the authors comment on this distinction?\n\nLIMITATIONS:\nSee weaknesses above.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-24"}

comments: SUMMARY:
The paper tackles the question of why a Transformer that is pretrained with a standard reinforcement-learning algorithm can later exhibit in-context reinforcement learning (ICRL). Building directly on the single-layer analysis of Wang et al.​ (2025), it generalizes the setting to multi-layer linear-attention Transformers used for policy evaluation. The authors (1) prove an inference-time result: with a particular “ICTD’’ weight configuration the value-prediction error shrinks to zero as depth grows, so the forward pass implements batched temporal-difference (TD) learning; (2) establish that the same weights are global minimizers of a norm-of-expected-update loss for both multi-task TD and a newly introduced multi-task Monte-Carlo objective, offering an optimization-based explanation for their empirical emergence; and (3) confirm the theory on randomly generated tasks, where pretrained models improve steadily with longer context lengths and converge numerically to the predicted ICTD weights.

STRENGTHS AND WEAKNESSES:
The theoretical development is sound. Both the inference-time convergence proof and the global-minimizer result are stated with clear assumptions. The experiments, although lightweight, are cleanly targeted at the two main claims and reproduce the predicted behavior. A weakness is the heavy reliance on linear attention and one-hot features; practical relevance to nonlinear Transformers or continuous domains is still open.

QUESTIONS:
How does MSVE change if you vary the number of layers? Can you elaborate on this with an empirical study?

LIMITATIONS:
All theory and experiments rely on linear attention; real-world models use softmax non-linear operator.
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-07-03

comments: SUMMARY:
This paper investigates the emergence of in-context reinforcement learning (ICRL) capabilities in Transformers pretrained with standard RL algorithms/objectives. The core question is: Why do RL-pretrained network parameters enable ICRL (i.e., solving new tasks without parameter updates by conditioning on context)? The authors hypothesize that such parameters are minimizers of the pretraining loss. Through a focused case study on policy evaluation, they prove the hypothesis.

STRENGTHS AND WEAKNESSES:
Strengths:

1. To the best of mt knowledge, this paper offers the first proof that ICRL parameters (specifically for ICPE) emerge as global minimizers of standard RL pretraining losses (NEU). This directly addresses the "why" of emergent ICRL.

2. Authors provide a unified analysis framework that shows $\theta^{TD}_{L}$ minimizes NEU for both TD and MC pretraining (Theorem 2, Corollary 1), suggesting robustness across algorithms.

3. I appreciate that authors clearly define the scope and limitations of the proposed theory: they explicitly acknowledge assumptions (linear attention, enumerated states) and scalability limits (Sec 7).

Weaknesses:

1. Limited Practical Applicability: Theoretical results rely on linear Transformers, while practical ICRL uses nonlinear models (e.g., ReLU), leading to a significant gap in real-world applicability.

2. Small-Scale Experiments: Experiments use tiny MRPs (Boyan’s chain with 13 states). Claims lack validation in complex environments. And the experiment misses a comparison to non-ICRL baselines (e.g., fine-tuning/offline-to-online RL).

QUESTIONS:
see above

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis is a theoretical study of the In-Context Policy Evaluation — a specific problem in the broader In-Context Reinforcement Learning setting. This setting allows for ground-up studying the emergence of ICRL capabilities in the forward passes of the transformer.\n\nThe authors develop a theory demonstrating the convergence of the TD and MC based approaches to be implemented in the forward pass of the linear transformer, as the number of layers goes to infinity (with empirical results covering depth 30 with looped weights).\n\nAuthors use Boyan’s chain to empirically verify their theoretical findings, demonstrating that convergence can happen in toy settings.\n\nSTRENGTHS AND WEAKNESSES:\n*Strengths*\n\n- **Rigorous theoretical extension:** Builds on Wang et al. 2025 by generalizing the single-layer linear-attention result to deep (number of layers to infinity) linear transformers, with a formal bound showing value-prediction error decays exponentially in depth.\n- **Unified global minimizer proof:** Introduces a novel analysis demonstrating that both Monte-Carlo and TD objectives share the same global minimizer under the NEU loss, tightening the link between the two training regimes.\n- **Sharper characteristic of the solution set:**  Shows that the broad invariant set identified by Wang et al. 2025 collapses to a single point in the deep-network limit, eliminating extraneous parameter configurations.\n\n*Weaknesses*\n- **Remaining convergence gap:**  The key open question from Wang et al. whether standard pre-training converges to the invariant/minimizer set remains unproven.  The paper offers only empirical evidence, leaving the theoretical guarantee unresolved.\n\n- **Narrow empirical scope:**  Verification is done in a single toy domain (5 state Boyans chain) and one architecture choice (30-layer looped linear Transformer).  No ablations on depth, context length, task scale are provided. The assertion that the experiments “empirically verify our theoretical insights” (in Introduction) is too strong given their limited breadth and the manually picked hyperparameters chosen to avoid numerical instability.\n\n- **Occasional over-claiming of novelty:**  Some phrasing suggests an ever first link between standard RL losses and ICPE, whereas Wang et al. 2025 already established the connection for single-layer models; the manuscript would benefit highly from cooling down the claims\n\nQUESTIONS:\nFirst, I believe this paper is timely and important. However, at the moment, this paper falls for me into the borderline accept category. I do not see how it could be qualified for Strong Accept as I do not qualify it “with groundbreaking impact on one or more areas of AI”. But, with a certain amount of modifications, I would be willing to put the paper into the Accept category.\n\n**What would make me to put the paper into the *Accept* category:**\n- A proof of convergence. I understand that this is not an easy task, but that would solidify it as accept for me.\n- An extended set of empirical investigation. Ablating the following would highly likely will result in moving your paper towards the accept:\n  - ablation to different depths, e.g.: 4, 8, 16, 32, 64, etc\n  - ablation for different context lengths: What happens for very small windows (m <= 5) or large ones (m>100)? Also is it possible to mismatch settings -- train with short context, evaluate with long?\n  - experiments beyond Boyan Chains, not sure what this could be, but going beyond just one environment is a must\n\n**What would make me to put the paper into the *Reject* category:**\n- As the paper is primarily theoretical, the correctness of the proofs is a necessary and critical condition for acceptance. I have been able to only partially verify the proofs, and have not found any mistakes. However, if other reviewers would highlight any mistakes (that are difficult to fix within this review period), I would be inclined to reject the paper regardless of any experimental evaluations.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-16"}