Towards Provable Emergence of In-Context Reinforcement Learning
Authors: Jiuqi Wang
Rohan Chandra
Shangtong Zhang
Venue: NeurIPS 2025 poster ()
Keywords: in-context reinforcement learning
policy evaluation
transformers
Primary area: reinforcement_learning

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.

Towards Provable Emergence of In-Contex: Reinforcement Learning
Jiuqi Wang
Rohan Chandra
Department of Computer Science University of Virginia Charlottesville, VA 22903 j siuuilemaii..vrriinii... www.wiiiiiii.... edu
Department of Computer Science
University of Virginia
Charlottesville, VA 22903
rchanchandra@virrinni..
Shangtong Zhang
Department of Computer Science
University of Virginia Charlottesville, VA 22903 shangtong@virginna.. edu
Abstract
Typically, a modern reinforcement learning (RL) agent solves a task by updating neural network parameters adapt its policy the task. Recently, it has been observed that some RL agents can solve a wide range of new outoof-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 '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 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 prove that when
support for this hypothesis through a Transformer is pretrained for policy evaluation, one of the global minimizers the pretraining loss enable in-context temporal difference learning.
study. In particular,
Introduction
Reinforcement learning (RL,, Sutton and Barto [2018]) is powerful machine learning paradigm for solving sequential peciisionmmaknng problems via trial and error. deep RL [Mnih al., 2015, Schulman et al., 2015, 2017], a classic workflow is to first train an agent with an RL algorithm in same task. Once the task changes, the skills acquired by the agent will
task and then deploy be obsolete, and a new agent must typically be trained for the new task by repeating the expensive procedure above. In-context RL (ICRL, Moeini et al. [2025b]), first coined by Laskin al. [2023], has the potential to address this limitation. In short, ICRL enables the agent to learn in the forward pass of their neural network by adding some additional data, called the context, to the network's input. This learning in the forward pass the RL agent is pretrained on a set of pretraining tasks. After pretraining, the RL agent network's parameters stay fixed, and the agent gets tested
without any parameter updates. Specifically, in ICRL,
test tasks, which can significantly
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

differ from the pretraining tasks. For example, in Laskin et al. [2023], the test task is bandit problem having opposite optimal arms the bandit problems in the pretraining tasks. Despite this discrepancy between pretraining and test tasks, one observes that the agent's performance in the test task increases as the context grows, with the agent's parameters frozen. We recall that the context serves as the additional input to the agent network, and a representative example is the agent's s historical observations and actions in the new task until the current time step [Laskin et al., 2023]. This performance improvement cannot be due to the hypothesis that the fixed pretrained parameters memorize the pretraining tasks because the test task can drastically differ from the pretraining tasks, The only plausible explanation seems to be that the pretrained parameters enable some reinforcement learning process in the forward pass of the network to process the information in the context and then output good action. This inference-ting time RL process is called ICRL. Notably, by tasks, we include both policy evaluation and control problems. The former predicts the value function, while the latter searches for performant policy.
ICRL has a diverse array of pretraining algorithms, which can be divided into supervised pretraining and reinforcement pretraining. The key idea of supervised pretraining is to perform imitation learning to train the agent network to imitate the behaviour of some existing RL algorithms. It is thus not surprising that the pretrained agent network behaves like an RL algorithm in the forward pass. We defer more discussion about supervised pretraining to Section 2. What is surprising is the reinforcement pretraining, where the agent network is some sequence model (e.g., Transformer [Vaswani et al., 2017]) and the pretraining algorithm is merely (a variant of) standard deep RL algorithms [Duan et al., 2016, Wang et al., 2016, Kirsch et al., 2022, Lu et al., 2023, Bauer et al., 2023, Grigsby et al., b024a,b, Park et al., 2025, Xu et al., 2024]. Essentially, in reinforcement pretraining, the network is trained to output good action (for control tasks) to make good value predictions (for policy evaluation tasks). It, however, turns out that the forward pass of the pretrained network behaves like an RL algorithm. We now give a few examples of reinforcement pretraining. Duan et al. [2016] use recurrent neural networks (RNNs) as the agent network and use trust region policy optimization [Schulman et al., 2015] as the pretraining algorithm. The proposed RL algorithm performs impressively in multi- aumti- bandits, tabular environments, and visually rich domains. A closely related and concurrent work is by Wang et al. [2016], where they train RRNs with A3C [Mnih et al., 2016]. Grigsby et al. [2024a] propose Amago that trains variant of Transformer (Vaswani et al., 2017] with customized REDQ [Chen al., 2021], which works competitively challenging robot benchmarks. Wang et al. [2025] train Transformers with temporal difference learning (TD, Sutton [1988]) and find that the pretrained Transformers perform well on unseen policy evaluation tasks. The pretraining algorithm is just a standard RL algorithm in all these examples. It gives rise to two important questions about the learned parameters:
How can the learned parameters enable ICRL in the network's forward pass?
Why do those parameters emerge after using the standard RL algorithm for pretraining?
Both questions are challenging since answering them requires white- the learned neural network's internals and understanding the dynamics of the pretraining algorithm. To our knowledge. no prior work can answer both questions. In addition, given the diversity of the pretraining algorithms, having a single answer to all pretraining algorithms unlikely. We present a case study following Wang et al. [2025] in this work.
Wang et al. [2025] answer Question (i) for policy evaluation tasks. In particular, they use TD to pretrain a Transformer for policy evaluation tasks on multiple randomly generated MDPs. After the pretraining converges, they study the converged parameters. Surprisingly, Wang al. [2025] prove that with the converged parameters, the layer by layer forward pass of their Transformer is equivalent to iteration by iteration updates of TD. In other words, the converged Transformer performs policy evaluation in its forward pass by implementing TD in the forward pass. However, Wang et al. [2025] only partially answer Question (ii). In particular, they show thay the converged parameters belong to an invariant set of the the pretraining algorithm (i.e., TD). By an invariant set, we mean a set of parameters such that once the pretraining algorithm enters, it remains in the set forever (in expectation). For example, if the parameter is in Rd then a trivial invariant set is just IR itself. The invariant set characterized by Wang et al. [2025] is much smaller than Rd (so it is nontrival and contains valuable information). But Wang et al. [2025] fall short in three aspect. First, their analysis considers only single-layer Transformer. Second, their invariant set contains more parameters than the converged parameters. Third, whether the pretraining algorithm will always converge to the
2

---

invariant set remains unclear. This work builds upon Wang et al. [2025] and provides finer answers to Questions (i) and (ii) by making the following three contributions.
1. (Inference Time Convergence) We prove that when Transformer parameterized with the converged parameters observed by Wang et al. [2025], the value prediction error made by the Transformer diminishes the depth of the Transformer grows infinity.
2. (Global Minimizer of Pretraining Loss) We prove that the converged parameters observed by Wang et al. [2025] is a global minimizer of a few pretraining algorithms, including both TD and Monte Carlo.
3. We empirically verify our theoretical insights with controlled environments mirroring our setup in theory.
2 Related Works
One can categorize the pretraining of ICRL into supervised pretraining and reinforcement pretraining. The key idea of supervised pretraining is imitation learning, algorithm distillation [Laskin et al., 2023]. In supervised pretraining, a dataset is first constructed by running a few existing known RL algorithms on the pretraining tasks and then collecting all the trajectories generated during the execution of those RL algorithms. Notably, a trajectory starts from the initialization of the RL algorithm and runs until the RL algorithm converges, thus spanning multiple episodes. As result, the rewards near the beginning of the trajectory are lower, and those near the end are higher. The trajectory thus demonstrates the RL policy's 's learning progress. Then, a sequence model is used fit the trajectories in the dataset. Namely, for a trajectory at time t, the input to the sequence model is all the data the trajectory before time The target output the sequence model the action taken at time I. The sequence model is trained via an imitation learning loss to output the target action. The sequence model thus mimics the behaviour of RL algorithms. Notable ICRL works with supervised pretraining include Laskin et al. [2023], Liu and Abbeel [2023], Zisman et al. [2024], Shi et al. [2024), Huang et al. [2024a,b]. Dai et al. [2024], Zisman et al. [2025], Polubarov et al. [2025]. Overall, we argue that in supervised pretraining. it is not surprising that the sequence model behaves like an RL algorithm in the forward pass because it was trained to do so. Lin et al. [2024] provide a theoretical analysis of supervised pretraining. In particular, they prove that with supervised pretraining, the pretrained sequence model behaves like the RL algorithms used to generate the dataset indeed. Lin et al. [2024] also provide some parameter construction of Transformers, such that the Transformer with those parameters can indeed implement some RL algorithms in the forward pass, including LinUCB [Chu al., 2011], Thompson sampling [Russo et al., 2018], and UCB-VI [Azar al., 2017]. However, the parameter construction in Lin et al. [2024] is overly complicated, and no evidence shows the constructed parameters are learnable. Our work is different from Lin al. [2024] in that (1) we study reinforcement pretraining, whose dynamics are entirely different from supervised pretraining; (2) the Transformer parameters we study are learnable by practical and standard RL algorithms, which we empirically demonstrate.
In terms of reinforcement pretraining, the closest our work is Park et al. [2025], which studies ICRL with large language models (LLMs) and proposes a regret-loss for reinforcement pretraining. They prove by construction that the global optimizer of their regret-loss implements the FTRL algorithm (Shalev-Shward and Singer, 2007]. While we classify Park et al. [2025] as reinforcement pretraining, it certainly deviates from standard RL. algorithms (cf. the reinforcement pretraining algorithms in Section 1). Our work differs from Park et al. [2025] in that (1) the pretraining algorithm different (so the pretraining dynamics are entirely different); (2) Park al. [2025] only study single-layer Transformers, while our results apply to multi-layer Transformers; (3) No evidence exists to support that the parameters studied by Park al. [2025] are learnable empirically, while the Transformer parameters we study are learnable by practical and standard RL algorithms, which we empirically demonstrate, Tarasov et al. [2025] build on top of algorithm distillation [Laskin et al., 2023] and propose to optimize the Transformer with RL objective. They report considerable improvement (30%) over the original algorithm distillation baseline on the testbeds. However, their work is entirely empirical and lacks theoretical insights. Moeini et al. [2025a] investigate ICRL with a safety constraint. They observe that reinforcement pretraining produces a better policy than the supervised pretraining counterpart in the safe ICRL domain.
3

---

ICRL falls into the category of black box meta RL. See Beck et al. [2025] for more comprehensive treatment of meta RL. ICRL is also a special case in-context learning (ICL, Brown al. [2020]) we use ICI to denote any learning paradigm that occurs in the inference time of the neural network. More commonly, ICI. exclusively refers to inference-time supervised learning. studies the provable emergence of in-context supervised learning [Ahn et al., 2024, Zhang et al., 2024, Gatmiry et al., 2024]. The techniques we use in this paper are entirely different from those works because the dynamics of supervised learning and reinforcement learning are completely different.
line of works
3 Preliminaries
We use In denote an n identity matrix and Omxn to denote an all-zero matrix in R=xx in this work. Given two vectors I, we will use (z, 1) and x yy to denote their inner product interchangeably.
3.1 Policy Evaluation
We consider anfinite-horizor Markov decision process (MDP, Puterman [2014]) with finite state space S and action space A, a bounded reward function r S x AA - IR, a state transition probability function p: S x A -> [0, 1], an initial distribution po : - [0, 1]. and discount factor [0, 1). An agent implements a policy π A x S - [0, 1]. Suppose the agent is at state St at time step t, outputs action A, x- S.), receives a reward R4+1 = r(S,, Ad) and transitions to the next state St ~ p(- St, At). When the policy æ is fixed, the MDP reduces to a Markov reward process (MRP), characterized by the tuple (S, p. Po, 7), where p(s' 8) = EneA Fa s)p(s' 8, s,a) and r(s) = E.EA m[a]s)r(s, a). The goal of policy evaluation estimating the value function R. defined as v(s) E -
So = =s. Since we will only consider policy evaluation tasks in this work, we will work with MRPs for simplicity. It is common to approximate with some parametric model in where w E Rd the weight. One of the simplest function approximators is arguably the linear model. Suppose there exists some feature mapping é S - Rd, a linear function approximation of v is then i.. (s) = x(s) T w vss). Note that p defines a Markov chain. Assuming the chain is ergodic, there exists unique and
stationary distribution a S + [0, 1]. Then, a common measure of the approximation error called the mean square value error (MSVE) can be well-defined as MSVE(u, v²) = E.es (v())((s)) - r((s))² for v', v: S + RR Thus, one typically wishes to find w" such that MSVE(U, (w-)) is minimized.
Monte Carlo (MC, Sutton and Barto [2018)) is straightforward policy evaluation method. MC trains the function approximator towards the return, defined as G(s) 2 7 Re+1 S. We note that [+](())] v(s) by definition. MC updates the weight vector
Wi-1 =w + (G((()))) - Ow. (S:)) Cw, (St),
where } is sequence of learning rates. is known that MC directly minimizes the MSVE [Sutton and Barto, 2018].
TD [Sutton, 1988, Sutton and Barto, 2018] is another fundamental policy evaluation algorithm. With linear function approximation, TD updates the weight vector iteratively as
W(+1 =w; + at + Tw.. (S++1) ow, (S.) (S.)
=w + a, (Re- + qwo ((S++1) - w ((S)))
where { is a sequence of learning rates. The term (R++1 + yow, (S.+1) tw, (St)) is the TD error responsible for correcting the estimate. Unlike MC, TD forms the training target using its estimate, known as "bootstrapping" [Sutton and Barto, 2018].
The norm of the expected update (NEU) loss is an objective that gradient TD methods [Sutton et al., 2008, 2009] optimize. While it was originally defined exclusively for linear function approximation, we generalize it in our work to arbitrary function approximation. Given the weights w, we define NEU as (EU = |[R[6V i (.))) (S)]]I, where & = G(S) in (S) for MC and 5 R + "0. - 6. (S) for TD.

---

3.2 Linear Attention and Transformer
canonical single-head self-attention [Vaswani al., 2017] computes Attmw,, W., w. (2) = W,Z softmax (Z W W,Z), where WE, Wg. W, denote the key, query, and value matrices, respec- tively, and Z is the imput prompt. The recent theoretical advancements of in-context reinforcement learning consider a linear attention variant, where the softmax activation is replaced with an identity mapping [Park al., 2025, Wang et al., 2025]. We consider the same linear aelf-attention structure and define LinAttn Z; Q) = PZM(2' TZ),, where Z E ((22+1)x(m+1) is the input prompt, P E ((24+1)((2++1) is the reparameterization of W.. Q E [222++1)x(2++) is the reparameterization
of W Wg, and M [(m+1)x(m+1) is fixed mask defined as M = Im 0mx1
We delegate
01xm 0
the first m columns of Z as the context and the (m 1)-th column for the query. The purpose of M is to separate the query from the context during forward propagation. Having defined the linear attention, we can define the linear Transformer as stack of linear self-attention layers. In an L-layer Transformer with parameters {Pis Q1) the input prompt Zo evolves layer by layer
Ze = Z LinAttm (2: Z4 + P.Z.M (z)
By the convention of Wang et al. [2025], we define TF, (Pi, Q1) -1)=
7 7,22+1,,
as the output of the L-layer Transformer, given input prompt Zy. Notably, 2(2d+1, (12d+1.++))) is the bottom-right element of the output embedding matrix, which is a scalar. The negative sign is in place be consistent with previous work and simplify the notations in the proof.
3.3 In-Context Temporal Difference Learning
Wang et al. [2025] show that under certain parameterization of P/ and Q₁, the layer-by-layer forward propagation of TF, is equivalent an algorithm on the context. In particular, let So, R2 Si,...; Rm, Sn be a trajectory unrolled from task. Suppose we wish estimate v(s,) for some query state Sq E S. Then, using shorthands 6, = (S,) and 0 = o(sg), we define for = 0, 1, L.
execution of batched version of TD
Qm-1
0q 0dx1 0
Zo -
701
2 Pm
E R 112441)x(m+3)
R1
IIa
Ia
0axı
PTTT
O2dx2d
02dx1 1 1
QT
Oaxd
Oaxd
0axı
R. (2(24+))((2411)
01x2d
O1xd
01xd
Let 0TD = (P.TD. Q J --- represent the parameters compactly. Theorem of Wang et al.
[2025] proves that, given Zo defined in (2), it holds that TF, ODD) = (d), wi), where {w) is defined as wo 0 and
W++1 wi
(Re41 ywo Φi+1 1i Φ1) 6i-
(3)
4 Inference Time Convergence
In this section, we establish the convergence and inference optimality of ICTD implemented by the L-layer linear Transformer L + xo. Before we proceed to our analysis, we first need to redefine the input prompt Zo and the feature function é to facilitate our proof the rest of the paper. Let = [S| and 80, 81,..., 8n-1 be the states in S. We employ a one-hot feature function : S > {0, 1)" that maps each state s to an n-dimensional vector where the n-th dimension is one and zero everywhere else. Using shorthands o, = ő(s.) and R.41 = =(s.), redefine Zo
d p)p(s) so) R
Φn-1
Φg
Zo =
1-1 E
7
Σ n- == , 8n-1)
0xx1 0
E
(4)
5

---

Now, the top rows of Zu are the enumerations of the states in S, represented in one-hot encodings. Rows n to 2n are now expected next features w.r.t p. discounted by y. Hence, given Zo defined in (4), it holds that TF, Zo; 0) (dq: wr), where {wr) } is defined as wi 0 and
W1+1 = +
R1+1 + yuu Σ p)p(s} s.) - w/ Øs Øz
(5)
by the virtue of (3). We employ P E [0, 1 to denote the transition probability matrix correspond- ing to p. where P(i, j) = p(s, 8,), and o E {0, 10,1) to indicate the feature matrix, where the i-th row equals d(s:). It is obvious that = In. Likewise, we use r E IR" to denote the vector representation of the reward function when it does not introduce ambiguity, where the i-th dimension r(s)). Then, we have the following theorem.
Theorem 1. Given query state 8q S and constructing Zo as (4), it holds that lim TFr (Zo; (2) v((8))
The proof is in Appendix B.1. Theorem demonstrates that, under our formulation of Zo in (4), I D converges the true value v(sy) of the query state Sy as the Transformer gets infinitely deep. This finding hints at the superiority of the parameterization @ID from the policy evaluation point of view.
Global Minimizer of Pretraining Loss
Wang al. [2025] discover that 0 emerges after pretraining on some task and feature distributions using the proposed multi-task TD update. Their effort in explaining the emergence of 0F. includes proof that 0% belongs an invariant set of the multi-task TD update. However, their analysis is restricted to the single-layer case, yet the phenomenon is also observed in multi-layer linear Trans- formers. Previously, we justify that 8% is desirable for in-context policy evaluation in Theorem In this section, we investigate it from an optimization perspective and show that 8T, is one of the global minimizers of the pretraining loss.
Let (r) denote the distribution of reward functions and (p) denote the distribution of transition probabilities. Then, the tuple (A(r,, Ap)) characterizes the MRP distribution. Let 0 denote the (vectorization and concatenation of) parameters {(P), Qi) of a Transformer of L layers. We recall that in (4), the prompt Zo depends on the query feature o.. as well as p and r. To this end, write Zo Zo,prr, to emphasize this dependency and use (p, r, ((s)) to denote the prompt by replacing the query oy with the feature (5s) for some E S. Wang et al. [2025) update iteratively as 0x+1 = 0g + a TD (0&), where
ID (0) En,-je . -p(p,pps))))))))))))))) (p) ((((s) + +TFF - (Zo V,TH, 0)]
Here, uef denotes the stationary distribution of the transition dynamics p. We further define Z/j = (o r, ((,))) and Zo = Zo r, o(s))). This motivates us to consider the Norm of Expected Update (NEU) loss, defined as J(Ø)= = || A (T) Similar to Gatmiry et al. [2024], we enforce sparsity assumption on the parameter space. Namely, we search 9 over θ defined as
o (P= O2nx 2ne u
0gnx1
Q
A 01x2n
O 1) Donxe))
0znx1
u E
Ix2n A 22222n
L
Notably, the parameter space + different Transformer layers can have different parameters (i.c., different and A). One special case of θ is called the looped (autoregressive) Transformer, where all Transformer layers are forced to share the same parameters. We use O Looped defined as
{{(== O2n
0znx 1)=
t poppd
{
0znx11
Q
R 1x2n A 2
0lx2n
to denote the parameter space (with sparsity constraints) of the looped Transformer we consider. It is easy to see that both 0HD E and 0,D E Looped hold. Although the evaluation of TF, (Zn; 0) and TF (Z;; 0) does not depend on the choice of the parameter space, the evaluation of V TFLF (Zo:0)

---

does. When & E θ the gradient is taken w.r.t. L different pairs of (u, A). On the other hand, when 0 E @Lopped the gradient take w.r.t. single pair of (u, A).
The theorem below confirms that 0 is the global optimizer of the NEU loss for the looped Trans- former when the number of layers grows to 2c. Theorem 2. Let the parameter space be OLopped, Then, it holds that lim (0 = 0.
few lemmas are in the sequl to prepare us to prove the above theorem.
Lemma For any query state Sq. transition probability p and reward function r, it holds that (egp)(()) (r(8)) 7TF Zoy 04.) TF, (Zo; 022)] ?|rll
The proof is in Appendix B.2. Lemma 1 shows that the absolute value of the expected TD error decays at an exponential rate.
Define X₁ 2 (1222,11.)) E _2=* the top left 2n by n block of Z,, and Yi z z(2n+1.lm) R 1x the bottom row of Zi except for the (n + 1)-th entry simplify notations. We furtherm define
dg 1nx1
a =
E 2 to represent the last column of Z except for the last dimension. Likewise, we
Xi Yi
write of to denote the bottom right entry of Z, Then, we can write Zi as Zi
uf
Lemma 2. Let Zi be the 1-th embedding evolved by an L-layer linear Transformer parameterized by 0 following (1), where Zo is the input prompt defined in (4). Then, for 1,...... L 1,
holds that I. wi' S (1 + +1)l 2. [|ill, S (7 + 1/lll for all query state Sy S, transition probability p and reward function r.
We leave the proof in Appendix B.3. Essentially, Lemma 2 implies that the norm of the evolving part of Zi, i.e., the norm of Yi and vi is uniformly bounded because of S for ( 2 O.
Lemma 3. V (Zo 0 S 42+L v LE, for some constants v and independent of L.
The proof can be found in Appendix B.4. Lemma 3 shows the norm of the gradient term grows at polynomial rate. We now have everything we need to prove Theorem 2 in Appendix B.5.
Recall that in multi-task TD, Wang et al. [2025] use the target r(s)) + rTF (Z6 0) to update the parameters of the Transformer, where TF& (Z6 8) can be viewed as a (biased) proxy for u ((s4)) As E E(((sg) + yr((s) = v(sg), we conjecture that using the Monte Carlo return G(ss) as the target can also enable in-context policy evaluation capabilities in Transformers. light of this, we propose the multi-task MC update that updates 0 iteratively as lut1 = 0% + ay MC (0%), where
A MC (0) = Ex,-H" Bss.""" ppp)) [(G(s.) TF, (Zo; 0) Vo 0)].
We again consider the NEU loss, defined as (0) A (0) Remarkably, we also prove that OP is the global optimizer of the NEU loss for our multi-task MC update as the Transforner's depth grows infinity.
Corollary 1. Let the parameter space be OLooped, Then it holds that lim J'(PP) 0.
need the following lemma to assist us in proving corollary.
Lemma For any query state 8q transition probability p and reward function r, it holds that Ep.r (G((s,) TF (Zo: 0T)) S 11
The proof is in Appendix B.6. Lemma 4 is analogous to Lemma in that it proves the norm of the expected upproximation error of MC decays an exponential rate. We can now prove Corollary in Appendix B.7.
6 Experiments
We conduct empirical studies to verify our theoretical results, In particular, mirroring our theoretical setup, investigate the following two questions:

---

1. Does the pretraining yield a Transformer that can perform in-context policy evaluation?
2. Do the converged parameters align with OTD?
We answer both questions for both multi-task TD and multi-task MC.
Following Wang et al. [2025], we employ Boyan's 's chain [Boyan, 1999] as our environment to con- struct the MRPs. Boyan's chain allows to have full information and control over the environment, including analytically solving for the true values and the stationary distributions. Figure 3 shows an example of S-state Boyan's 's chain. Adapting the technique of Wang et al. [2025], we randomly generate the reward and transition probability functions, preserving the topology of the chain to ensure its ergodicity. Our only simplifying modification is setting the stationary distribution as the initial distribution. The details of the task generation, including the distributions to sample and r, can be found in Algorithm
Since our convergence results require L - oc, we construct a looped Transformer TFz with = 30. We find this depth yields consistent results without rumning into numerical instability. For each trial of the experiment, we generate 20,000 tasks for training with y = 0.9. Within each task characterized by p, r, generate mini-batch of size b 64. For multi-task TD, unroll the MRP generate trajectory So, R1, S1, Re..... ., Sh Rs, Sb. Then, we update the parameter @ of the Transformer by
Ou+1 = 02 +
8, A (SS))):
where 6, R4+1 + TTF (Zo r, (00,,,,,((S111)) (((())));;)))))))))))11))))))))) 0x) TF2 (Zo(p, r, (5.))) See Algorithm 1 for a more detailed description. Similarly, for multi-task MC, we unroll the MRP to generate a trajectory So, R1, ,, R2..... .2........................ S 2, Rg Sh For each S,, we independently unroll the MRP for 200 steps obtain truncated return G(S.) beginning from that state. Admittedly, G(S)) introduces bias due to truncation. However, the bias is negligible as 7.200 x 7 10-10 parameter by
Multi-task MC updates the
b-1
0x+1 = 0% +
(ā(S)) - TF (Z),, VoTT)) (oTFF(Za((r,,665))
See Algorithm 2 for a more comprehensive illustration. We inherit most the experimental settings from Wang al. [2025]. See Table for a comprchensive list of hyperparameters and training details.
6.1 In-Context Policy Evaluation Verification
Recall that one property of ICRL that the agent's performance improves as the context gets longer without parameter updates. Based on this principle, we investigate the MSVE as a dependent variable of the context length for the converged Transformers. After each trial, we sample & = 10 novel tasks validation set. Notably, Wang et al. [2025] theoretically
Ti...., Th from the task distribution and empirically proved that Transformers parameterized by (TD can perform policy evaluation in context unrolled from MRP. Thus, we employ the same context construction as in Wang al. [2025] for evaluation instead of using the enumerated states and expected features that we used for training the Transformers. Suppose we need a context with length m from T,, unroll T; for m + 1 steps and get s(()) (((() s(()
,..., B
s(t)
s(i) m
R.(() sli) Then, we form the context as
( (((
(
(())
m-1
C(i,m) = 7o ((t)
To
S
((())
R (2n+ 12m+1)xm
R(
Given a query state Sq. we then define the input prompt Z(i, m, = [C(i, zyg] () where Ig [o(s,)T 01x(=+1)] Then, given parameter 8, we can T. with context length m by MSVE(i, m) =
evaluate the MSVE of the Transformer Eees uP( (s) m, s); 0) v(s))², where wdd means the stationary distribution with re- spect to the transition probabilities of T,. We average the MSVEs across the validation tasks to get an

---

accurate estimation of the expected MSVE with context length m. We compute the averaged MSVEs with respect an array of increasing m's. We repeat our experiments for 20 trials for both multi-tast TD and MC using the last checkpoint after pretraining. Figure plots the mean and standard error of the averaged MSVE against the context length for multi-task TD and MC. The MSVE exhibits clear decreasing trend as the context gets longer for both pretraining schemes. Therefore, we can confidently conclude that multi-task TD and MC can produce Transformers capable of in-context policy evaluation. The answer to question I is affirmative.
60-
msve
50
30-
20-
context length
Figure 1: Mean and standard error of the averaged MSVEs against context lengths. The curves averaged over 20 random trials. The shaded areas represent the standard errors.
6.2 Weight Convergence
In this section, inspect the parameters of the Transformers after pretraining with multi-task TD and MC. Figure 2 displays the final P and Q matrices after pretraining. Despite mild noise, both methods produce the parameter configuration 04D Hence, the answer to question 2 is also affirmative. Surprisingly, even their noise patterns look very similar. Therefore, we speculate that the weights converge the same place with multi-task TD and MC. Notably, the Transformer learns to implement in-context TD even after pretraining with multi-tast task MC. This observation testifies to the fact that the model does not merely imitate the reinforcement pretraining algorithm like in algorithm distillation for supervised pretraining. Instead, the model itself decides how to best learn from the context.
(a) Multi-task TD
(b) Multi-task MC
Figure 2: Mean Transformer parameters after pretraining. The parameters averaged over 20 trials and normalized to stay in the range 11,
7 Limitations and Future Works
While our work dives deeper into the question of why reinforcement pretraining enables ICTD, has several limitations. Firstly, our analysis assumes a linear Transforner. Although linear attentions are usual in theoretical analysis for in-context learning [Ahn et al., 2024, Gatmiry et al., 2024, Park et al., 2025, Wang et al., 2025], it leaves a gap between theory and practice. Future works can look into extending the results to simple non-linear activations, such as ReLUU. Secondly, our prompt construction assumes enumeration of the states and expected next features. While this assumption

---

facilitates our analysis, it can hardly hold in practice, especially in MDPs with a continuous state space. Future works may include generalizing our results sampled states and/or next features. Thirdly, our empirical study is relatively small in scale. While it is intended for a clear proof of concept, future work can scale up our experiments. In this work, we focus only on the prediction problem. Exciting future directions include studying the emergence of in-context control via reinforcement pretraining.
8 Conclusion
This work presents an in-depth analysis of the emergence of the in-context TD parameter through reinforcement pretraining. We show that under some prompt construction, the in-context prediction of the Transformer with ICTD weights converges to the true value as it gets deeper. Furthermore, prove that the ICTD parameter is minimizer of the NEU loss for both multi-task TD and MC pretraining as the number of attention layers goes to infinity. Through controlled experiments, we show that (1) both multi-task TD and MC induce in-context policy evaluation capabilities in linear Transformers, and (2) both pretraining schemes produce the same final parameter - the ICTD weights. We acknowledge the limitations of our work, including the linear attention assumption put on the Transformer, the practicality issue of our prompt construction, and the small scale of our empirical study. We believe addressing these limitations will make admirable contributions the advancement of ICRL.
Acknowledgements
This work is supported in part by the US National Science Foundation under the awards 2I121180019, SLES-2331 1944 and CAREER-2442098, the Commonwealth Cyber Initiative's Central Virginia Node under the award VV-1Q26-001, and an NVIDIA academic grant program award.
References
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement precon- ditioned gradient descent for in-context learning. In Advances in Neural Information Processing Systems. 2024.
Jason Ansel, Edward Yang. Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesensky, Bin Bao, Peter Bell. David Berard, Evgeni Burovski, Geeta Chauhan, Anjali Chourdia, Will Constable, Alban Desmaison. Zachary Vito,o Elias Ellison, Will Feng. Jiong Gong, Michael Gschwind, Brian Hirsh, Sherlock Huang. Kshiteej Kalamharkar, Laurent Kirsch, Michael Lazos, Mario Lezcano, Yanho Liang, Jason Liang, Yinghai Lu, CK Luk, Bert Maher, Yumjie Pan, Christian Puhrsch, Matthias Reso, Mark Saroufim, Marcos Yukio Siraichi, Helen Suk, Michacl Suo, Phil Tillet, Eikan Wang, Xiaodong Wang, William Wen, Shunting Zhang, Xu Zhao, Keren Zhou, Richard Zou, Ajit Mathews. Gregory Chanan, Peng Wu, and Soumith Chintala. PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation. In Proceedings of the ACM International Conference on Architectural Support for Programming Languages and Operating Systems, 2024.
Mohammad Gheshlaghi Azar, Ian Osband, and Rémi Munos. Minimax regret bounds for reinforcement learning. In Proceedings of the International Conference on Machine Learning. 2017.
Jakob Bauer, Kate Baumli, Feryal Behbahani, Avishkar Bhoopchand, Nathalie Bradley-Schmieg. Michael Chang, Natalie Clay, Adrian Collister, Vibhavari Dasagi. Lucy Gonzalez, et al. Human-timescale adaptation in an open-ended task space. In Proceedines of the International Conference on Machine Learning. 2023.
Jacob Beck, Risto Vuorio, Evan Zheran Liu, Zheng Xiong, Luisa Zintgraf, Chelsea Finn, and Shimon Whiteson. A tutorial on meta-reinnforcement learning. Foundations and Trends@ in Machine Learning, 2025.
Justin A. Boyan. Least-squares temporal difference learning. In Proceedings of the International Conference on Machine Learning, 1999.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-VVon Voss, Gretchen Krueger. Tom Henighan. Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few shot learners. In Advances in Neural Information Processing Systems, 2020.
10

---

Xinyue Chen. Che Wang. Zijian Zhou, and Keith W. Ross. Randomized ensembled double q-learning: Learning fast without model. In Proceedings of the International Conference Learning Representations. 2021.
Wei Chu, Lihong Li, Lev Reyzin, and Robert Schapire. Contextual bandits with linear payoff functions. In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, 2011.
Zhenwen Dai, Federico Tomasi, and Sina Ghiassian. In-context exploration-exploiatiie for reinforcement learning. In Proceedings of the International Conference on Learning Representations, 2024.
Yan Duan, John Schulman, Xi Chen, Peter L Bartleu, Ilya Sutskever, and Pieter Abbeel. RI2: fast reinforcement learning via slow reinforcement learning. ArXiv Preprint, 2016.
Khashayar Gatmiry, Nikunj Saunshi, Sashank J. Reddi, Stefanie Jegelka, and Sanjiv Kumar. Can looped transformers learn to implement multi-step gradient descent for in-context learning? In Proceedings of the International Conference on Machine Learning, 2024.
Jake Grigsby, Linxi Fan, and Yuke Zhu. AMAGO: Scalable in-context reinforcement learning for adaptive agents. In Proceedings of the International Conference Learning Representations, 2024a.
Jake Grigshy, Justin Sasek, Samyak Parajuli, Tkechukwu Adebi, Amy Zhang, and Yuke Zhu. Amago-2:
Breaking the multi-task barrier in meta-reinnfrrcement learning with transformers. In Advances in Neural Information Processing Systems, 2024b.
Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtamen, David Coumnapea,, Eric Wieser, Julian Taylor, Sehastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Haldane, Jaime Fernández del Río, Mark Wiebe, Pearu Peter- son, Pierre -arahhann Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Arry programming with NumPy. Nature, 2020.
Sili Huang,
Jifeng Hu, Hechang Chen, Lichao Sun, and Bo Yang. In-context decision transformer: reinforcement learning via hicrarchical chain-of of-thought In Proceedings of the International Conference on Machine Learning, 2024a.
Sili Huang, Jifeng Hu, Zhejian Yang, Liwei Yang, Tao Luo, Hechang Chen, Lichao Sun, and Bo Yang. Decision mamba: Reinforcement learning via hybrid selective sequence modeling. In Advances in Neural Information Processing Systems, 2024b.
J. Hunter. Matplodib: A 2d graphics environment. Computing in Science & Engineering, 2007.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Proceedings of the International Conference on Learning Representations. 2015.
Louis Kirsch, Sebastian Flennerhag, Hado van Hasselt, Abram Friesen, Junhyuk Oh, and Yutian Chen. Intro- ducing symmetries to black box meta reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, 2022.
Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steigerwald, DJ Strouse, Steven Stenberg Hamsen, Angelos Filos, Ethan Brooks, maxime gazeau, Himanshu Sahni, Satinder Singh, and Volodymyr Mnih. In-contex. reinforcement learning with algorithm distillation. In Proceedings of the International Conference Learning Representations, 2023.
Licong Lin, Yu Bai, and Song Mei. Transformers as decision makers: Provable in-context reinforcement learning via supervised pretraining. In Proceedings of the International Conference on Learning Representations. 2024.
Hao Liu and Pieter Abbeel. Emergent agentic transformer from chain of hindsight experience. In Proceedings of the International Conference on Machine Learning. 2023.
Chris Lu, Yannick Schroecker. Albert Gu, Emilio Parisotto, Jakob Foerster, Satinder Singh, and Feryal Behbahani. Structured state space models for in-context reinforcement learning. In Advances in Neural Information Processing Systems. 2023.
Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Bellemare, Alex Graves. Martin A. Riedmiller, Andreas Fidjeland, Georg Ostrovski. Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 2015.
Volodymyr Mnih Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap. Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In Proceedings of the International Conference on Machine Learning, 2016.
11

---

Amir Moeini, Minjae Kwon, Alper Kamil Bozkurt. Yuichi Motai, Rohan Chandra, Lu Feng, and Shangtong Zhang. Safe in-context reinforcement learning. ArXiv Preprint, 2025a.
Amir Mocini, Jiuqi Wang, Jacoh Beck, Ethan Blaser, Shimon Whiteson, Rohan Chandra, and Shangtong Zhang. A survey of in-context reinforcement learning. ArXiv Preprint, 2025b.
Chanwoo Park, Xiangyu Liu, Asuman E. Ozdaglar, and Kaiqing Zhang. Do LLM agents have regret? case study in online learning and games. In Proceedings of the International Conference on Learning Representations, 2025.
Andrei Polubarov, Lyubaykin Nikita, Alexander Derevyagin, llya Zisman, Denis Tarasov, Alexander Nikulin, and Vladislav Kurenkov. Vintix: Action model via in-context reinforcement learning. In Proceedings of the International Conference on Machine Learning. 2025.
Martin Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Daniel Russo, Benjamin Van Roy, Abbas Kazerouni, Ian Osband, Zheng Wen, et al. A tutorial on thompson sampling. Foundations and Trends@ in Machine Learning, 2018.
John Schulman, Sergey Levine, Pieter Abbrel, Michael I. Jordan, and Philipp Moritz. Trust region policy optimization. In Proceedings of the International Conference on Machine Learning, 2015.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimoy. Proximal policy optimization algorithms. ArXiv Preprint, 2017.
Shai Shalev-Shwartz and Yoram Singer. A primal-dual perspective of online learning algorithms. Machine Learning, 2007.
Lucy Xiaoyang Shi, Yunfan Jiang, Jake Grigshy, Linxi Fan, and Yuke Zhu, Cross-episodic curriculum for transformer agents. In Advances in Neural Information Pracessing Systems, 2024.
Richard S. Sutton. Learning to predict by the methods of temporal differences. Machine Learning, 1988.
Richard Sutton and Andrew G Barto. Reinforcement Learning: An Introduction (2nd Edition). MIT press. 2018.
-empora--diferrnnce Richard Sutton, Csaba Szepesvári, and Hamid Reza Maci convergent O(n) temporal- difference algorithm for off-policy learning with linear function approximation. In Advances in Neural Information Processing Systems, 2008
Richard S. Sutton, Hamid Reza Maei, Doina Precup. Shalabh Bhatnagar, David Silver, Csaba Szepesvári, and Eric Wiewiora. Fast cradient-deseent methods for -hrpprr--iifffrence learning with lincar function approximation. In Proceedings of the International Conference on Machine Learning, 2009.
Denis Tarasov. Alexander Nikulin, 1lya Zisman, Albina Klepach. Andrei Polubarov, Lyubaykin Nikita, Alexander Derevyagin, Igor Kiselev, and Vladislav Kurenkov. Yes. q-learning helps offline in-context RL. In Scaling Self-Improving Foundation Models without Human Supervision, 2025.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, & ukasz Kaiser, and Illia Polosukhin. Attention all you need. In Advances in Neural Information Processing Systems, 2017.
Jane X Wang. Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Z Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement Iearn. ArXiv Preprint, 2016.
Jiuqi Wang, Ethan Blaser, Hadi Daneshmand, and Shangtong Zhang. Transformers can learn temporal difference methods for in-context reinforcement learning. In Proceedings of the International Conference on Learning Representations, 2025.
Tengye Xu, Zihao Li, and Qinyuan Ren. Meta rettareeiffreemeet learning robust to distributional shift via performing lifelong in- in-context learning. In Proceedings of the International Conference on Machine Learning, 2024.
Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 2024.
Ilya Zisman, Vladislav Kurenkov, Alexander Nikulin, Viacheslav Sinii, and Sergey Kolesnikov. Emergence of in context reinforcement learning from noise distillation. In Proceedings of the International Conference on Machine Learning, 2024.
12

---

Ilya Zisman, Alexander Nikulin, Viacheslav Sinii. Denis Tarasov, Lyubaykin Nikita, Andrei Polubarov, Igor Kiselev, and Vladislav Kurenkov. N N-gram induction heads for in-context RL: Improving stability and reducing data needs. First Workshop on Scalable Optimization for Efficient Adaptive Foundation Models, 2025.
13

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We clearly define the scope of our work in the abstract and intro. Each claimed contribution is reflected in one more of the subsequent sections.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
Section summarizes the limitations of our work.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-speeifficttion, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
14

---

Answer: [Yes]
Justification: We state the assumptions in the theory statement or in the context of the theory. We put our complete proof in the Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We include the important information in Section We put more experimental details in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
15

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: We still need to clean up the codebase and create instructions to run our code. We will submit our code part of the supplementary material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. cc/
ceaminisssooooooooiioooo cesssonooliii olicyyyy)))lo for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / //nips. cc/ wwcp/bllc///aeeeccceeeceecceelee http//uudee//eeeee CCtpes//ooSSubimeecc cessonnoolcyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We specify the experimental details in Section 6 well the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification:
We plot the standard errors for our first experiment.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi-
dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for
example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
16

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We disclose the compute resources and wall clock times for running the experiments in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: / http///eeeipsppeeeeeeee wwterriss..cppupiiee cc/public/ Ceteii////iieeeee csQuidel ines?
Answer: [Yes]
Justification: We reviewed the Code of Ethics and confirm that our research conforms with all them.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No. they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work is mainly theoretical and fundamental. The perceived societal impacts are minimal.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
17

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work does
rely existing data.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We cite the tools use implement experiments and generate our figures in the Appendix.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
18

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: Our paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
19

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: LLM is not involved in the core method development of our paper.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: /ttp///errpp/////ee for what should or should be described.
/LLM)
20

---

A Multi-task TD and MC
We provide the pseudocode of multi-task TD and MC in this section.
Algorithm 1: Multi-Tasss Temporal Difference Learning (adapted from Algorithm 1 of Wang et al. [2025])
1: Input: state number MRP distribution A(p, one-hot encoder ø, batch size b number of training tasks k, learning rate a, discount factor y, transformer parameters
0 = (Pi, -00......-1
2: for i f- to do
3: p,r ~ A(p.r)
4: Po
distribution(p) Rb, S. 5: Unroll So. R1. S1, S
from the MRP (po. p. <+- // to aggregate semi-gradients
6:
7:
for f
0 0....,
do
((s))
Φ ((5
((S.)
8:
Zo
7 Σ 1 p(5j 80))(88)
j=0 p(sj 88 ({b(88)
r(so)
r(nn-1)
((50)
665n-1
((S++1
7 Σド= p(8j 80)((88)
7 Σうー p(sj 8n (8((8))
0
r(0))
0
9: < ( + a Ru+1 TF Zg: 0)
TFz Zo: 8)) TFL Zo:
end for
11: 0 f- 0+0++ : // apply update
12: end for
Algorithm 2: Multi-Task Monte Carlo Learning
1: Input: state number n, MRP distribution A(p, r). one-hot encoder Φ, MC unroll budget т, batch size b, number of training tasks k, learning rate discount factor y, transformer parameters {P,, Q1)
2: for i +-
do 3: p.r (p)
4: Po f
5: Unroll So, R1, S1, ......, Sb-2, Ro-1, S. from the MRP (po.
6:
<+0 // to aggregate semi -emdiieas
7:
for f
0 b do
((50)
n- ((nn-1)
((S.)
8:
Zo +
? E ((8j 80) (0))(8)) r(so)
EE= p(sj | 8 1 18(88)
0
((
// Monte Carlo unrolling G+0
9:
10:
1 1
11:
8 + S.
12:
for 1,2, do
13:
G + wr(s) + G
+ vi ~ p(- s)
15:
s +- s'
17: end for
<+ aGG TF, <
19: end for
20: 0- 0 + &
// apply update
21: end for
21

---

B
Proof of Theoretical Results
B.1 Theorem
Proof. We can rewrite (5) compactly as
Wt+1 w + ① T + 7 Pww: T dwr
w + ① (r
T + (7P (PPP-1yuup 1)
Substituting + In. we have
W141 =w: + r + (TP L)
=r y Per-
Unrolling the recursive definition, we get
Wp =0
W1 =r
W2 = + Pr
W3 -r + Pr + (xP)3r
:
1-1
Wt
(7p))r
(6)
k=0
Therefore, we have
L-1
lim 1 +
wp
lim
Σ
( (1-
With slight abuse of notation, employ IR" denote the value function in vector form, where the s-th dimension us)) When writing out v as a Bellman equation, we have
+ PPo
(1 - p U r
tr. v= (I 7P) -t.
Thus, we have shown that lim WL V. Finally, recall that
TF TF P (d WL).
We therefore have
TF (Z; 0DD) 1
lim
lim 1.
(0g, w1.)
(0m 1.
lim wr, ) (0,, 1) =((g))
B.2
Lemma
Proof.
Beirp(en) (r(3)) + TF,, (Zo; 02D) TF Zo: 022]
(((g) + (Bgpp((x,) TF2 (Zo ]TD] TF (Zo (TD))
(0)rr) + BBB . ( ( wr)]
(Bgg wi)
1--1
1-1
0 r + -TE [6/]
Σ (x 0 Σ (x
(eq. (6))
k=0
k=0
22

---

1-1
0 r + 76 P)
Σ
(xp)) - 0g
Σ
(
A=
4
4-1
9 + 9
T
Σ
Pp))
T
P)
k=1
k=0
10r oy (1P)
4T 0g
= Kon p(r)
(Hölder's inequality)
= plrll
(lø, 1)
4 |rllo
(1 1lac 1)
B.3 Lemma 2
Proof. We first note that of
-T, 1 (oo: (
(bbg wi). Therefore, we get
1x/1 110g un}i
(Hölder's inequality)
|
(1ا4 1)
P)**
(eq. (6))
= (I (xP)) (I -rl
(I (xP)))oll
(eq. (7))
= ||v - /
SI|v|l00 + |'ollo
Sl + +
=(1 +1
( 1)
We have proved 1. of Lemma 2.
Under the structure of Pt in 0., it holds that the first 2n rows of Z remain unchanged. Thus, drop the subscript of X; and write
X x
Zi
Yi
We then proceed to compute the recursive forms of Yi and 8/-
We first have
O2nx2n 01x2n
02ax1
X
2:9
In
Oaxi
P.Z.M =
Yi
01xn
O2nx2n
02nx
X
Oznx
01x2n
Yi
02nx1 Xn
Yi
Next, we have
27 QiZi
TXT
Y T
A
02nx1 1
X
17
01x2n
Vi
T
Y,T
AxX
Apr9
01xn
23

---

XTX A.X
XT Aar"
x9
29" Arr9
Putting the two parts together, we have
PIZ. (ZT
= 02axn
02nx1
XT A,X
XT A1a99
Yi
x AiX
x² Axr9
0znxn
02nx1 YIX A129
=
YiX A,X
We therefore get By induction, we prove that
Y++ = + YY A,X;
124 = v8 + YiX A
Y44 Yo + Σ A.X.
A,X.
Similarly, we can show that
v/t1 vū >
Y A₃29.
i=0
Since vă 0, we get
vif1 = Σ Y A,²9.
(8)
2m0
Let vt (k) denote the k-th component of Y,. is clear that
Vt+1
(k)
Vo
(M)
YiX TYXTA, (Az(1),
(9)
==0
where (A) denotes the k-th column of X. Recall that o/ = - (0g. wi). Substituting (k) in (9) into x:9 in (8), we get
(k) - (Oxx, wi).
V1
Therefore, we have
Yi Yo w/ In T wi
We finally get
|7
Y7 T
Σ +
r (I
| - (I - (:P)))olls
(eq. (7))
=|(TT [[]TP)(T 7 -1
=|[TT - -Ppp (I (I- (1(yP)))ollo
<|((P))o ~ IyPull x
P' 11/loc + r r|l
-(1'+7)llels>>
(P'l =1,1|P||cc 1)
This concludes the proof of 2. of Lemma 2.
24

---

B.4 Lemma 3
Proof. Since F E @loped we can drop the subseripts of the linear Transformer parame- ters because of parameter sharing across the layers. We further have (Z. =
vec (V TFy, (Zo:0P))
(2 1)(2m), where vec denotes the vectorization operation. We then
V. LFL (Zo: DD)
have
Y1+1 + (u. + =x AX + YI(T ([I+XT' AXAX);; = vi + (uX + x
(10)
wit1
Ar = + xX Az + Y Ax
We begin by solving for V A (1)
where 4/+1 denotes the i-th component of Yi+1. Define f(A) I + X AX simplify notations. By (10), we have
(6)
=
fuxX AX
AA
AA
T AA +
d
dEE Di ())
f(A)ss
dA
where 200)) indicates the i-th column of X. We analyze the summands separately. Since
VM (a Mb) ab given matrix M and vectors and b, it holds that
d[uXXT Ar
XX u (0)
dA
d An
Plugging in the condition that 0 in D. we get
0.
dA
d Σ 2 {(A)5a)
Next, tackle the gradient term we get
By the summation rule and the product rule,
AA
Σラ= vi f(A)ss
d
Ml f(A))1)
dy:
+ ( d{(}ss (11)
Σ
[(A))
dA
j=1 dA
dA
AA
1=1
Employing the structure of A, first note that
f(A) =1 + X AX
1 1
+1+ !T yP]
0
PPT
=1+1-1 I
YPT
=1
=YPT.
also note that
d alls his + x!1) Acli)
af(A);;
x00
dA
dA
Therefore, plugging them in (11), we have
Σ 1 11 (1)
{())
Σ
PP, Vaul
(3) + w 11 10000000)))
dA
1=1
m Σ
=XYTTx(1)
Pr. P.,yAAl 1. 1
25

---

Consequently, we have
V A3141 =XY, 200)
(6)
+ Σ
()
==1
We now bound the gradient term norm
V
(
We first have
V 1
XY,
201
Σ
Pc. V
A3
XY, 200'|1
11, +
Σ P.. V All
||||| |,, 11
|1,
+ ?
>
P(.
V n1
(3)
1
j=1
= -1[//--
1
| 00
[00
+ Σ
Pio
V
11
j=1
= -1+ - 1)illl +7 Σ
Pis V A3ll
(5)
/=1
<n(1+ +7)lle11 +
Σ Proj
V x
(Lemma 2)
j=c
S En(1+)) +7
Psa
V AB (3)
Let B = nn1 + 7) Poll Having observed this pattern, we claim
V A3/
S 1/3
(12)
for all i = 1,2, ......... n and prove it by induction. Since Yo is not dependent on A, we have VAYo 0.
Therefore, the hase case
VAyD (6)
< 0 trivially holds. Then, suppose (12) holds for I. We have
V AU Autf
|1,
SB+ +7 ? Σ
Pi.
V
Tu
(3)
j=1
n
S 7 Σ
j=1
B + +13
(1 71)
< (1 + 1)33
The bound in (12) is proved.
We then proceed with V VAVLLT Define 9(A) = XT Az to simplify notations. Similar to V by (10), we have
d d Ax
d[Yig(A)]
V A -V =VAV
dA
dA
d d Ax
Ek= Vi (K)
9(A)
=VAn +
dA
dA
26

---

Firstly, we get
d [ Ax
XXX u x9
dA
dfuxxT Az?
Substituting O into the equation, get
0. Then, we have
dA
d Σ k= Vi g())a
(1)
(K) (AA)
dup (k)
(2) (g(A))
g(A)
+ y1
(13)
dA
AA
AA
dA
by the summation rule and the product rule. We note that, by the structure of A in 03 we have
9(A) =XT Ar
I
=11 PPI
0
0
0
=[1 >PI
Qq
= Q4
(14)
We further note that
d d12(1)) Ax."
dg(AA)x
x (8)
dA dA
Combining these and plugging into (13), we get
d Σ (K) 9(A)
Σ
(K) (k) - -Ooll V Ab
111 x
dA
=XY] x
1ogle Aul (t)
Since og a one-hot vector, we let w be the index where [0ylw 1. We then have
th
(k)
g(A){
Vi
XY, x VAyl
dA
Hence, we have
V VAVL+1
V V + XY 29 Vavl (w)
Now, we derive the bound for |
A || We first have
V Au8 + XY, Ta - -vau
=
<
V A3
+
XY,T 29
11, V ||
We note that
XXY] 2
1 5,5
1
1 29 '1,
TP
(1 >)|lo
< n(1 + >)((' +
(Lemma 2)
s n(1 + 12)loollls> =
Therefore, by (12), we have
VAM < | +1.3
27

---

With this form, we claim and prove by induction that
12+1_ 2 /
V S
(15)
for 2 O. Since v does not depend on A, we have VADi = 0. Therefore, the base case VV38l1, S trivially holds. Now suppose (15) holds for I. We get
|| AV 11 - V AUl
|| + + 1)/3
1a + 1
S
3 + (1 1)33
2
(1 + 1)² (1 (1+1).
B.
2
The bound in (15) is proved.
Next, we solve for V. By (10), we first have
d E Pijsi (1)
d d Ar(
du
du
=XX Ax
P.o Vuyl 1
+
Define V,Yi = Vuyl() (1)
Vayl
Vuvl")) R 2unn
We thus have
xxx V.Yi+1 XX AX VX+7V.YPT.
We now bound VuY++1111 We get
XXX AX +V
$ || XXX AX + V
|||
$ |X]
||x AX]I, + 1
PT
|| 11 1 >P)
0 >PT +>IlV.Yi
(1 + 11+ - -||, + +VIl
S1 + 7) VVVVl
Let Ç = (1 + ?)3. We claim
| V SK
(16)
and prove it by induction. We first have V.Yo = because Y, does not depend on u. Hence, the base case Voll I 0 holds trivially. Assume (16) holds for 1. We have
[V.Y4
S Ç + V
$ Ç + 775
$ (1 + 70)ç
(1 + I)ç.
The bound in (16) holds. Finally, we compute V V.yL.1- By (10), we have
x Ax
d dlYxXTT Ar
V Vuwi+ =V +
du
du
=Vuff + XXXT Ar (V.Y))X Ard,
Then, we bound V VuDi+zll, by
V W.
=|| V + xxx Azd T +
(V.X))x Ar
28

---

V 13 + |XX Ar + +[ Ar
<|V +
$ VV + +
(eq. (16))
|V (1++)
(eq. (14))
| w 11 +1+))++
Define = (1 + 7)) We prove the bound
12 + + -
I/ 11
(17)
by induction. We first have the base case Vuvill, I 0 holds because vó is independent from u and thus V.y0 = Now, suppose (17) holds for I. We have
S
V w 11 17 7 15
12 +1
( + in + 17
12 +1
<
+ (1 1)5 + (1+ 1)n
12 + +
( + (1+ 1)m
2
(1 + 1)2 (1+ 1)
+ 1)n-
The bound holds for (17).
Since TFL (Zo; (T
= -/6 we have
V TF P | VAvill < 1244B B5
V TFL Zo: 07D) = V.vill S 2
22 Lnp.
We then have bounds
vee V xTTF, (2:.TF//oo0:0000 (20:0 < 2n 1244 3
V TFx 0 - 2n
Litbe + Lnj
We therefore get
vec V vec(V TF, Zz; RT.)
VATF, (Zo: 0
0 ||,
vec V A TF Zo: TTD) ||,
| TFL TF2 (Zu: |||
1 + L
(2n ( + ()) + I.(2n)).
Let v 2n(B + 5) and & = 2ny. We finally arrive at
L2+ + L
V (Zo; 01) <
LE.
2
B.5
Theorem
Proof. We have
29

---

Bo ppp.a ((r(s,) >TFF (Z)F 0 TZ: 0/ V 02) (Zo: |10
Bemp + (ZTT, TF (Zo:
Vo (2)
(o: 0
-m1 (r() (f) TF (20:
Zo: 0
<B Mpp pl)p
(7E TF (Zo: (2)) 0
-Beg pr
3/ (p(p))
02))
- -o: |oTF; 1
Begr pa
|||||| 20
1
Voz TF (Zoi 22)) l
(Lemma 1)
1 + L
SB R pre
|1rll00 7
T.E
(Lemma 3)
I22 + I.
-E
LE)
Let C, E, (Era(m))lll We have [(O S 1 1 + L Therefore, it holds that 2
L + L
lim J( lim ++00
uO + LEC,T" =
because the exponential term 7 dominates the polynomial terms of L. Since (0 is norm and thus always nonnegative, we have
J(O 0, L
lim
completing the proof.
B.6 Lemma 4
Proof. We note that given p, r, Sq, there is no randomness in TF, (Zo; 0). Therefore, we have
Ep, (G(,,) TF,, (Zo 07
Ep.r(G(s.)) (s4s)] TF (Zu: 0110)
v(sg) (dg, wtY1
1.
v(s)) 6 Σ &
(eq. (6))
k=
0 07 (I (TP)))(
-P))-1
0 o7 (I p'
=00 P'vv''
1> |
(Hölder's inequality)
sx | plollllllllo
(10011 1)
( = 1)
B.7
Corollary 1
Proof. (Corollary 1) J'
30

---

pa [(((s)) TFL Zo; 07. VoTFL
Ep.r [G((s) TF (Zo: (Zo: (o
= Bagage p
Bpr [G(s) - 7F (Z:: 02))
V [T; TFF (2)
Bpr (G(s)) T (Zo: 0
)1
Ea Baguue
B [G(g)) TFL (Zo: 0% ]]]]]]]]]
V&TFL (Zo: 11]
< spp
[1111100
V&TFL (Zo: 02))11 [h1]
(Lemma 4)
L2 +
1101|x07
+ LL
(Lemma 3)
L
= -
v LE)
2
Let C. R (p we have ) S 12+1 vCot 1 that
+ LEC.tt. Therefore, it holds
+
lim J'( lim L-00
2 vC. IgC,,==0
L-x
because the exponential term 7 dominates the polynomial terms of L. Since J((,). is a norm and thus always nonnegative, have
lim ) 0,
1-000
completing the proof.
C Experiment Details
C.1 Boyan's Chain
We include an illustration of Boyan's 's chain and its generation method here.
2
S
Figure 3: Boyan's 's chain with S states. Arrows indicate non-zero transition probabilities.
C.2 Compute Resources
We experiments in parallel single node of a CPU cluster. The node has 150 CPU cores and 150 GB memory. The wall clock time it takes to finish running the experiments is about 50 minutes for multi-task TD and 15 hours for multi-task MC.
C.3 Hyperparameters and More Implementation Details
We use NumPy [Harris et al., 2020] for data processing implementing MRPs. We use Py Torch [Ansel et al., 2024] to create and train our models. For data visualization, we use Mat- plotlib [Hunter, 2007] create the plots. Table lists the hyperparameters we used for experimenta- tion.
31

---

Algorithm 3: Boyan Chain MRP Generation (Adapted from Algorithm 2 of Wang et al.
[2025])
1: Input: state space size = |
2: Uniform [) // reward function
3: f- 0mxm // transition function
4: io I,...,, n do
5: 6 Uniform [0,1)]
6: 1 c-
7: p(i,i+ -
8: end for
9: ((1-1,n) 11
10: z Uniform 1)"]
11: 2 2
12: p(n,1:n)+ f-
13: Po aeaaiionrr_distribbueee
//initial distribution
14: Output: MRP (pv.p,r)
optimizer learning rate
Adam [Kingma and Ba, 2015]
0.001
weight decay
0.0
batch size
64
# of attention layers
30
of Boyan's chain states
discount factor
0.9
# of Monte Carlo rollout steps
200
# of random seeds
20
# of Boyan's 's chain tasks for training
20,000
# of validation instances validation context lengths
10
5, ......... 100
Table I: Hyperparameters and more training details.
D Hyperparameter Study
D.1 Transformer Depth
We conducted additional experiments to investigate the robustness of the emergence of ICTD under different numbers of attention layers. Keeping all other configurations unchanged, we pretrained 5-layer, 10-layer, and 60-layer Transformer for 20,000 tasks. Fach experiment was repeated for 20 times. Figure 4 plots the mean converged weights of the Transformers. We observe that the pattern of @TD emerged in all the tested depths for both multi-task TD and MC. We further tested the converged Transformers on unseen MRPs and report the mean and standard error of the MSVE as function of the context length, displayed in Figure We witness steady decrease of the MSVE as the context length grows, demonstrating clear in-context RL capability of the pretrained models.
32

---

(a) Multi-task TD with 5-layer Transformer
(b) Multi-task MC with a 5-layer Transformer
Multi-task TD with 10-layer Transformer
(d) Multi-task MC with 10-layer Transformer
(c) Multi-task TD with 60-layer Transformer
(f) Multi-task MC with a 60-layer Transformer
Figure 4: Mean parameters of 5-, 10-, and 60-layer Transformers after pretraining. Parameters are averaged over 20 trials and normalized to lie within the range [-1,1]
5 layers
layers
50
TD 60 layers
leyers
MC 10 layers
40
MC 60 layers
msve
20
10
20
100
context length
Figure 5: Mean and standard error of the averaged MSVEs against context lengths for 5-, 10-, and 60-layer Transformers. The curves are averaged over 20 random trials. The shaded areas represent the standard errors.
33

---

D.2 Number of States
We also studied the emergence of ICTD under different numbers of states of the MRPs for training. Fixing all other configurations, additionally pretrained the Transformer on Boyan's 's chains with and 10 states, respectively. Notably, since we employ state enumeration for training, the context length consistent with the number of states of the MRP during training. Hence, this experiment also ablates different training context sizes. Each experiment was repeated for 20 independent trials. Figure plots the mean converged weights of the Transformers. All them have clear pattern of gTb, Figure 7 illustrates the decreasing MSVEs vs, the increasing context length for the converged models the Again, all the pretrained models exhibit in-context policy evaluation behaviors.
6,
Multi-task TD with 3-state MRP
Multi-task MC with 3-state MRP
(c) Multi-task TD with a 10-state MRP
(d) Multi Must MC with 10-state MRP
Figure 6: Mean parameters the Transformers after pretraining on 3- and 10-state MRPs. Parameters averaged over 20 trials and normalized to lie within the range [-1,
3 states states
states MC 10 states
60
msve
20
20
40
60
80
100
context length
Figure 7: Mean and standard error of the averaged MSVEs against context lengths for Boyan's chains with and 10 states, respectively. The curves are averaged over 20 random trials. The shaded areas represent the standard errors.
34

---

D.3 MRP
Since we hypothesize that training with multi-task TD and MC induces ICTD capabiliities in the Transformers, and since TD is not environment-speen specific, then the Transformers pretrained on one class of MRPs should be capable of predicting the values of another class of MRPs. In our last ablation study, we pretrained the Transformers on the Loop environment used by Wang et al. [2025] and evaluated them on the data generated by the Boyan's 's chain MRPs. Algorithm sketches the generation procedure of the Loop MRPs, and Figure 8 is an illustrative example of 6-state Loop environment. Figure 9 shows the mean converged parameters of the Transformers, supporting the emergence of (1 for Transformers pretrained on the Loop MRPs. In addition, Figure 10 plots the consistent decrease the MS VEs against the context length on the tested Boyan's 's chain environments, supporting our hypothesis that the Transformer not overfitting one particular class of MRPs but learning a general prediction algorithm in its forward pass via pretraining with multi-task TD or MC.
Algorithm 4: Loop MRP Generation
1: Input: state space size n = |S]; threshold E (0, 1)
2: Uniform (-1.
// reward function
3: c Uniform [0, 1) R) // connectivity
4: c [[e ) // apply indicator function element-wise
5: for i = 1,...
do
6: c(i, 0 // no self-loop
7: +) +1 // can transition to next state
8: end for
9: c(n,n) +
10: 1
11: pm Uniform [0., 1) // transition function
12: p p // element- wise multiplication
13: for 1. do
14: = E--i((i,j) // normalizing constant
15: p(i,1:n + (nn))
end for
17: Po + staationary.disszibbttooo))) //initial distribution
Output: MRP (Po.p.
5
3
Figure 8: Loop with states. Arrows indicate non-zero transition probabilities.
35

---

(a) Multi-tast TD with Loop
(b) Multi task MC with Loop
Figure 9: Mean parameters of the Transformers after pretraining on the Loop MRPs, Parameters are averaged over 20 trials and normalized lie within range [-1,
TD Loop
200
msve
100
50
20
40
60
80
100
context length
Figure 10: Mean and standard error the averaged MSVEs against context lengths for Boyan's 's chains trained on the Loop MRPs. The curves are averaged over 20 random trials. The shaded areas represent the standard errors.
36


Reviews:
[{"review_id": "g05aEKPhPJ", "reviewer": "Reviewer_De38", "summary": "This 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.", "questions": "First, 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.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "*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", "was_revised": true, "final_justification": "- The main reason to accept is the good theoretical development of the In-Context Reinforcement Learning.\n- However, the limited scope of experiments which is induced by the theoretical limitations is rather difficult to ignore to assess the paper as having a major impact on at least one sub-area of AI, therefore borderline accept.", "rebuttal": "We sincerely thank the reviewer for the comprehensive review and questions. We appreciate the reviewer's opinions on our theoretical rigour and the timeliness and importance of our work. We hope our response below addresses the reviewer's concerns and questions.\n\n> The key open question from Wang et al. whether standard pre-training converges to the invariant/minimizer set remains unproven.\n\nWe agree that a convergence proof would significantly strengthen this work, and we have been considering this for a long time. It might not seem too difficult at first glance, given the success of the convergence proof of in-context regression in [2] and [3]. However, we want to highlight the unique challenge here -- the lack of the Polyak-Lojasiewicz (PL) inequality. Since the Transformer is highly nonlinear, PL is probably the only effective tool to use. Indeed, it's used in both [2] and [3]. Yet, we do not see any application of PL inequality in analyzing value-based RL algorithms. They are fundamentally different because regression is gradient descent (GD), whereas TD is not. TD is only semi-gradient. So, some nice symmetric property for establishing PL in regression does not hold in TD. Consequently, given the tremendous difficulty, the convergence proof, if it exists in the first place, deserves a project of its own. For example, [4] is the first paper to demonstrate the global optimality, and [2] is the first to show convergence. Both are very lengthy, even though they only studied regression. Our TD setting will inevitably make it more complicated. \n\n> 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 ...\n\nWe appreciate the reviewer for providing us with actionable suggestions. Before we delve into the additional results, we would like to note that during evaluation, we assessed the Transformers trained on 5-state MRPs across a range of context lengths (5-100). We believe the scalability of performance against the context length is a key and distinct feature of in-context learning. Thus, the reviewer's advice \"...mismatch settings -- train with short context, evaluate with long\" is precisely what we did for evaluating the Transformers' in-context policy evaluation capability. \nIn light of the reviewer's suggestions, we did the following additional hyperparameter studies:\n- Depth: Keeping everything else unchanged, we additionally trained and evaluated 5-, 10-, and 60-layer Transformers for 20,000 tasks. We also repeated each experiment with 20 random seeds for statistical rigour. As we cannot upload figures during the rebuttal period, we can only verbally confirm that the pattern of $\\theta^{TD}$ emerged in all the depths we tested, with the 60-layer Transformers having slightly more pronounced patterns. Let m denote the context length. In terms of quantitative evaluations, we paste here the approximate mean MSVEs on the shortest context (m = 5) and the longest context (m = 100) tested across 5-, 10-, and 60-layer Transformers trained via multi-task TD and MC.\n    \n    Multi-task TD:\n\n    |   | 5-layer | 10-layer | 60-layer |\n    |---|---------|----------|----------|\n    | m = 5| 12.0 | 7.0 | 2.1 |\n    | m = 100 | 7.0 | 4.8 | 1.1|\n\n    Multi-task MC:\n    |   | 5-layer | 10-layer | 60-layer |\n    |---|---------|----------|----------|\n    | m = 5| 10.9| 7.8 | 3.0 |\n    | m = 100 | 6.3 | 5.3 | 1.9\n\n    We observe that in all cases, the mean MSVE dropped significantly as the context length increased, confirming the emergence of ICPE. Additionally, we witness that the deeper the Transformers, the better the prediction. Since under $\\theta^{TD}$, one loop over the attention layer is equivalent to one step of TD update, we conjecture that the deeper Transformer applies more steps of TD updates internally, thus exhibiting better convergence. Lastly, we observe that the magnitudes of the mean MSVEs are of the same order between multi-task TD and MC. This observation provides further evidence that the Transformers trained via multi-task TD and MC are implementing the same in-context policy evaluation algorithm, at least behaviorally. It aligns with our theoretical result that $\\theta^{TD}$ is a global minimizer of the NEU loss for both multi-task TD and MC.\n\n- Context length/State number/Feature dimension: Because of our construction, the number of states of the MRP, the context length, and the feature dimension are equal for training. For this study, we repeated our experiment on a 3-state Boyan's chain and a 10-state Boyan's chain for 20 random seeds. Limited by compute, we cannot make the MRPs too large due to the increased Transformer parameter size and Monte Carlo rollouts. We can confirm the emergence of $\\theta^{TD}$ for both MRPs. Below is a table of the mean MSVEs from the quantitative evaluation.\n\n    Multi-task TD:\n\n    |   | 3-state MRP | 10-state MRP | \n    |---|---------|----------|\n    | m = 5| 2.9| 3.3 |\n    | m = 100 | 2.3 | 1.9 |\n\n\n    Multi-task MC:\n\n    |   | 3-state MRP | 10-state MRP | \n    |---|---------|----------|\n    | m = 5| 6.0| 3.8 |\n    | m = 100 | 4.75 | 1.9 |\n        \n    The consistent decrease in MSVE confirms the in-context policy evaluation capability of the Transformer in both MRPs.\n\n- MRP: We further designed an MRP called Loop as follows: Given states $s_0, s_1, \\dots, s_{n-1}$, $s_i$ can transition to $s_{i+1}$ for $i = 0, 1, \\dots, n-2$. Furthermore, $s_{n-1}$ can transition to $s_0$, forming a closed loop. For each pair of $i, j$ where $j \\notin \\\\{ i, ((i+1) \\mod n)\\\\}$, we sample a number $k(i,j)$ between 0 and 1 uniformly. Given a threshold $\\tau \\in [0,1]$, $s_i$ can transition to $s_j$ if and only if $k(i, j) \\ge \\tau$. We then generate the transition probabilities for each state, retaining the connectivity. With this design, we guarantee that there exists only one communicating class due to the loop structure. Furthermore, we can choose $\\tau$ to control the \"hops\" between states. We will also add the pseudocode for generating the Loop MRP in the next version of the manuscript. In this experiment, we chose $\\tau = 0.5$ and ran the same training and validation procedures as in Boyan's chain experiments. We observed a clear $\\theta^{TD}$ pattern for both multi-task TD and MC. For validation, we have the following table.\n\n    | |Multi-task TD | Multi-task MC | \n    |---|---------|----------|\n    | m = 5| 2.2| 3.5 |\n    | m = 100 | 0.3 | 1.1 |\n\n    Therefore, the numerical results suggest the emergence of ICPE capability of the Transformers when trained on this Loop MRP.\n\nWe would also tone down our claim \"empirically verifying our theoretical insights\" to something more conservative, such as \"preliminary empirical studies demonstrate overall consistency with our theoretical results\".\n\n> 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\nWe will update our claims to be more prudent to address the reviewer's concerns about overclaiming. We also thank and welcome the reviewer to follow up with other possible overclaims in our manuscript. In our defence, we wish to highlight that Wang et al., 2025 [1] justified the emergence of $\\theta^{TD}$ by showing that it lies in an invariance set of the expected version of the multi-task TD update in the single-layer case only. Nevertheless, they did not connect $\\theta^{TD}$ to any well-defined loss as in our analysis.\n\n[1] Jiuqi Wang, Ethan Blaser, Hadi Daneshmand, & Shangtong Zhang (2025). Transformers Can Learn Temporal Difference Methods for In-Context Reinforcement Learning. In The Thirteenth International Conference on Learning Representations.\n\n[2] Ruiqi Zhang, Spencer Frei, & Peter L. Bartlett (2024). Trained Transformers Learn Linear Models In-Context. Journal of Machine Learning Research, 25(49), 1–55.\n\n[3] Khashayar Gatmiry, Nikunj Saunshi, Sashank J. Reddi, Stefanie Jegelka, & Sanjiv Kumar (2024). Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning?. In Forty-first International Conference on Machine Learning.\n\n[4] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, & Suvrit Sra (2023). Transformers learn to implement preconditioned gradient descent for in-context learning. In Thirty-seventh Conference on Neural Information Processing Systems."}, {"review_id": "gIXDytA3no", "reviewer": "Reviewer_zvwM", "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.", "questions": "see above", "limitations": "yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n1. 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.\n\n2. 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.\n\n3. 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).\n\nWeaknesses:\n\n1. 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.\n\n2. 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).", "was_revised": true, "final_justification": "After carefully reading the author's response, I have decided to keep my initial score.", "rebuttal": "We deeply thank the reviewer for reviewing our work and acknowledging the robustness of our core theoretical contribution. We hope the following can address the concerns of the reviewer.\n\n> 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.\n\nWe acknowledge that this assumption leaves a gap between the theory and the canonical form of the modern Transformer architecture. However, linear attention is one of the most common simplifying assumptions for establishing exact equivalence between the forward pass of the Transformer and some principled algorithm. Some notable works include [1], [2], [3], and [4]. Furthermore, as the majority of our work builds upon [5], where they prove linear Transformers can and do implement TD learning in-context, we very much need to inherit their setup because we need to work with in-context TD. \n\n> Experiments use tiny MRPs (Boyan’s chain with 13 states). Claims lack validation in complex environments.\n\nSince our primary goal is to empirically confirm our answer to why ICTD emerges, we need to mirror the setup of our theories. Thus, it simultaneously requires us to generate a large number of MRPs and have full access to their information. This requirement is not achievable with complex environments. On the other hand, having full control over the environment has several benefits. For instance, we can analytically solve for the ground truth value function and the stationary distribution of the MRP, allowing us to compute the MSVE against the ground truth. With a complex environment, one can only approximate it by unrolling the policy for many steps, which is neither efficient nor accurate. To compensate for the lack of complexity in our tasks, we ran additional hyperparameter studies to demonstrate that the emergence of ICTD is robust, as presented in the response to Reviewer De38.\n\n> And the experiment misses a comparison to non-ICRL baselines (e.g., fine-tuning/offline-to-online RL).\n\nWe would like to claim that our goal in this work is not to propose a state-of-the-art RL algorithm but to provide an explanation for the emergence of in-context TD in Transformers through reinforcement pretraining, such as multi-task TD and MC. Thus, we design our experiments to verify our theory. Since other potential baselines are irrelevant to the goal of explanation, we do not include them in our empirical study.\n\n[1] Sander, M., Giryes, R., Suzuki, T., Blondel, M., & Peyré, G. (2024). How do transformers perform in-context autoregressive learning?. In Proceedings of the 41st International Conference on Machine Learning.\n\n[2] Chenyu Zheng, Wei Huang, Rongzhen Wang, Guoqiang Wu, Jun Zhu, & Chongxuan Li (2024). On Mesa-Optimization in Autoregressively Trained Transformers: Emergence and Capability. In The Thirty-eighth Annual Conference on Neural Information Processing Systems.\n\n[3] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, & Suvrit Sra (2023). Transformers learn to implement preconditioned gradient descent for in-context learning. In Thirty-seventh Conference on Neural Information Processing Systems.\n\n[4] Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., & Vladymyrov, M. (2023). Transformers learn in-context by gradient descent. In Proceedings of the 40th International Conference on Machine Learning.\n\n[5] Jiuqi Wang, Ethan Blaser, Hadi Daneshmand, & Shangtong Zhang (2025). Transformers Can Learn Temporal Difference Methods for In-Context Reinforcement Learning. In The Thirteenth International Conference on Learning Representations."}, {"review_id": "uuTsL2pycY", "reviewer": "Reviewer_N5DC", "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.", "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.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "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.", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate the reviewer for the honest review and thoughtful question. We are pleased to learn the compliments on the soundness of our theory and the effectiveness of our experiments. We hope our response below addresses the reviewer's concerns.\n\n>  A weakness is the heavy reliance on linear attention and one-hot features; practical relevance to nonlinear Transformers or continuous domains is still open.\n\nWe acknowledge that linear attention is not widely adopted in modern Transformer architectures. Nevertheless, it is a standard simplifying assumption in establishing precise equivalence between the forward pass of the Transformer and some principled algorithm. Some notable theoretical works in in-context learning using linear attentions include [1], [2], [3], and [4]. Furthermore, as our work extends [5] to answer why ICTD emerges, we also need to follow their setup using linear attentions.\nAs our work represents one of the earliest attempts to explain why ICTD emerges, to the best of our knowledge, we wish to begin with the tabular setting. As in many fundamental theoretical works in RL, the tabular setting lays out the foundation for studying linear function approximation by allowing us to isolate the core problem (i.e., the pretraining dynamics). The one-hot encoding is mathematically equivalent to the tabular representation.\n\n> How does MSVE change if you vary the number of layers?\n\nWe conducted an additional hyperparameter study on Transformer depth in response to the reviewer's question. Keeping everything else unchanged, we additionally trained and evaluated 5-, 10-, and 60-layer Transformers for 20,000 tasks. We also repeated each experiment with 20 random seeds for statistical rigour. Let m denote the context length. Pasted below are the approximate mean MSVEs on the shortest context (m = 5) and the longest context (m = 100) tested across 5-, 10-, and 60-layer Transformers trained via multi-task TD and MC.\n    \nMulti-task TD:\n\n|   | 5-layer | 10-layer | 60-layer |\n|---|---------|----------|----------|\n| m = 5| 12.0 | 7.0 | 2.1 |\n| m = 100 | 7.0 | 4.8 | 1.1|\n\nMulti-task MC:\n|   | 5-layer | 10-layer | 60-layer |\n|---|---------|----------|----------|\n| m = 5| 10.9| 7.8 | 3.0 |\n| m = 100 | 6.3 | 5.3 | 1.9\n            \nWe observe that in all cases, the mean MSVE dropped significantly when the context length increased, confirming the emergence of in-context policy evaluation. Furthermore, we notice that the deeper the Transformer, the better the prediction. Since under $\\theta^{TD}$, one loop over the attention layer is equivalent to one step of TD update, we conjecture that the deeper Transformers apply more steps of TD updates internally, thus exhibiting better convergence. Lastly, we observe that the magnitudes of the mean MSVEs are of the same order between multi-task TD and MC. This observation provides further evidence that the Transformers trained via multi-task TD and MC are implementing the same in-context policy evaluation algorithm, at least behaviorally. It aligns with our theoretical result that $\\theta^{TD}$ is a global minimizer of the NEU loss for both multi-task TD and MC. Please also consult our response to Reviewer De38 for more hyperparameter study outcomes.\n\n[1] Sander, M., Giryes, R., Suzuki, T., Blondel, M., & Peyré, G. (2024). How do transformers perform in-context autoregressive learning?. In Proceedings of the 41st International Conference on Machine Learning.\n\n[2] Chenyu Zheng, Wei Huang, Rongzhen Wang, Guoqiang Wu, Jun Zhu, & Chongxuan Li (2024). On Mesa-Optimization in Autoregressively Trained Transformers: Emergence and Capability. In The Thirty-eighth Annual Conference on Neural Information Processing Systems.\n\n[3] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, & Suvrit Sra (2023). Transformers learn to implement preconditioned gradient descent for in-context learning. In Thirty-seventh Conference on Neural Information Processing Systems.\n\n[4] Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., & Vladymyrov, M. (2023). Transformers learn in-context by gradient descent. In Proceedings of the 40th International Conference on Machine Learning.\n\n[5] Jiuqi Wang, Ethan Blaser, Hadi Daneshmand, & Shangtong Zhang (2025). Transformers Can Learn Temporal Difference Methods for In-Context Reinforcement Learning. In The Thirteenth International Conference on Learning Representations."}, {"review_id": "CBTPLgdjVu", "reviewer": "Reviewer_xY5j", "summary": "This 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.", "questions": "- 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?", "limitations": "See weaknesses above.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths\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.", "was_revised": true, "final_justification": "My concerns about the paper's clarity and the structural assumptions on the transformer are addressed in the rebuttal. The concern about whether the setup is overly toy-like and misaligned with real-world RL or ICRL applications is partially addressed—while related literature exists, I am not aware of ICRL application papers that adopt this exact setting. Overall, I lean toward accepting the paper, though it remains a borderline case.", "rebuttal": "We sincerely thank the reviewer for the constructive feedback and detailed questions. We thank the reviewer for acknowledging the novelty of our theoretical investigation and the quality of our writing. We hope our response below can address the concerns and the questions.\n\n> 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\n\nWe agree that we worked with a simplified RL setting. Nevertheless, we still preserve the challenge of temporal credit assignment (i.e., backward propagation of the rewards) -- a core challenge of modern RL that distinguishes our work from in-context regression. So we argue that our results still constitute a big leap over SOTA results on understanding in-context learning and in-context RL. The MRP formulation helps us to isolate the challenge of temporal credit assignment. Empirically, it allows us to easily compute the ground truth value function and the stationary distribution, which enables us to compute the MSVE accurately. We would not have these benefits if we were to work with control.\n\n> and can be solved by relatively simple methods.\n\nWe agree that our simplified problem setting can be solved by other relatively simple methods more efficiently. However, our goal is not to create a new SOTA for solving this simplified problem. Instead, we aim to advance the theoretical understanding of the in-context TD phenomenon first established by [2] and provide a finer explanation for why reinforcement pretraining, like multi-task TD and MC, can induce in-context policy evaluation capabilities in Transformers from a training dynamics point of view. We argue that this simplified model does serve this purpose well and creates new insights.\n\n> The writing contains some inconsistencies and relies on setups or results from prior work without sufficient explanation.\n\nWe apologize for the confusion caused by our notation inconsistency and insufficient explanations. We will update our manuscript to refrain from reusing the same notations for different meanings across sections. In our setup, $n = |\\mathcal{S}|$ because the context consists of an enumeration of the states. However, we agree that we shouldn't have used $n$ to also denote the sampled trajectory length earlier. We will also add more explanations of the setup and an explicit definition of multi-task TD inherited from [2].\n\n> the Transformer architecture and the sparse parameterization considered in this work represent a highly simplified setting that lacks scalability to realistic applications.\n\nWe acknowledge that our assumptions about the Transformer leave a gap between our theory and practice. In our defence, we would like to note that such assumptions are common and standard in works that aim to whitebox the in-context learning algorithms and explain their emergence in Transformers. For instance, notable works like [3], [4], [5], [6], [7], and [8] all assumed linear attention in their analysis. Furthermore, [5], [7], and [8] enforced the sparse parameter constraint. Without such assumptions, the manual computation of the forward propagation and the gradients of a multi-layer Transformer becomes intractably convoluted. In regard to the scalability to realistic applications, we argue that our goal is not to propose a state-of-the-art RL algorithm but to advance the explanability of the emergence of in-context RL as observed in large Transformer-based models [9].\n\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?\n\nThe reviewer's interpretation of multi-tasking is correct. By multi-tasking, we meant training the Transformer on multiple MRPs, each having its unique transition dynamics and reward function. \n\n> If so, does this find application in other reinforcement learning studies?\n\nMulti-task training is an important technique in meta-RL, where the agent trained on multiple tasks achieves generalization and exhibits zero-shot/few-shot learning capabilities. A notable early work is RL$^2$ [10]. See [11] for a more comprehensive survey of how such multi-task training is useful in meta RL. One can consider in-context RL algorithms to be a subset of meta RL algorithms. A distinctive feature of ICRL algorithms is that the agent achieves generalization via encoding an RL algorithm in the forward pass of its model. \n\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.\n\nWe thank the reviewer for carefully examining our manuscript and raising this question. We planned to examine the in-context policy evaluation capabilities of the Transformer by showing a decreasing MSVE against an increasing context length, analogous to Figure 1 of [2]. However, in our theoretical analysis, the context length has to match the number of states due to enumeration. Hence, if we were to choose to mirror our theoretical setup for evaluation, we would have to generate a different MRP for each context length. For instance, the MSVE reported on context length 5 would come from a different MRP than the one used to generate the context of length 10. In this case, it would not be an apple-to-apple comparison between the two MSVEs since they are computed for two different MRPs, thus defeating the purpose. Consequently, we must give up satisfying the state enumeration assumption and sample the contexts of different lengths from the same MRP for each validation trial. Recall that we proved $\\theta^{TD}$ is a global minimum of the NEU loss for both multi-task TD and MC as the layer goes infinitely deep.  If the weights of the Transformer indeed converge to $\\theta^{TD}$, then the model learns to implement TD learning in its forward pass. Since Wang et al. [2] generated their Figure 1 by unrolling the MRP, we choose to approximate their setting by directly sampling from the stationary distribution, as the state distribution would converge to the stationary distribution beginning from any distribution as the context gets longer.\n\n>Specifically, I would expect that unless the stationary distribution is uniform over the state space, such an approximation could introduce bias. \n\nSince we are using one-hot features, the value functions are fully representable by the features, which is equivalent to the tabular setting. Hence, as long as the context covers all the states, there would not be any bias in value estimation. So, as long as the sampling distribution supports the state space, there would not be any bias when the context is sufficiently long. As our MRPs are ergodic, the stationary distribution indeed supports the state space, which should be feasible for our evaluation purpose.\n\n[1] Sutton, R., Barto, A., & others (1998). Reinforcement learning: An introduction. (Vol. 1) MIT press Cambridge.\n\n[2] Jiuqi Wang, Ethan Blaser, Hadi Daneshmand, & Shangtong Zhang (2025). Transformers Can Learn Temporal Difference Methods for In-Context Reinforcement Learning. In The Thirteenth International Conference on Learning Representations.\n\n[3] Sander, M., Giryes, R., Suzuki, T., Blondel, M., & Peyré, G. (2024). How do transformers perform in-context autoregressive learning?. In Proceedings of the 41st International Conference on Machine Learning.\n\n[4] Chenyu Zheng, Wei Huang, Rongzhen Wang, Guoqiang Wu, Jun Zhu, & Chongxuan Li (2024). On Mesa-Optimization in Autoregressively Trained Transformers: Emergence and Capability. In The Thirty-eighth Annual Conference on Neural Information Processing Systems.\n\n[5] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, & Suvrit Sra (2023). Transformers learn to implement preconditioned gradient descent for in-context learning. In Thirty-seventh Conference on Neural Information Processing Systems.\n\n[6] Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., & Vladymyrov, M. (2023). Transformers learn in-context by gradient descent. In Proceedings of the 40th International Conference on Machine Learning.\n\n[7] Khashayar Gatmiry, Nikunj Saunshi, Sashank J. Reddi, Stefanie Jegelka, & Sanjiv Kumar (2024). Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning?. In Forty-first International Conference on Machine Learning.\n\n[8] Zhang, R., Frei, S., & Bartlett, P. (2024). Trained transformers learn linear models in-context. J. Mach. Learn. Res., 25(1).\n\n[9] Chanwoo Park, Xiangyu Liu, Asuman E. Ozdaglar, & Kaiqing Zhang (2025). Do LLM Agents Have Regret? A Case Study in Online Learning and Games. In The Thirteenth International Conference on Learning Representations.\n\n[10] Yan Duan, John Schulman, Xi Chen, Peter L. Bartlett, Ilya Sutskever, & Pieter Abbeel. (2017). RL$^2$: Fast Reinforcement Learning via Slow Reinforcement Learning.\n\n[11] Jacob Beck, Risto Vuorio, Evan Zheran Liu, Zheng Xiong, Luisa Zintgraf, Chelsea Finn, & Shimon Whiteson. (2025). A Tutorial on Meta-Reinforcement Learning."}]

Author rebuttal:

Meta-review: The theoretical results are judged sound and interesting by all reviewers. However, the paper has also some weaknesses:
- Narrow empirical scope (Reviewer De38), limited practical applicability (Reviewer zvwM), "practical relevance (...) is still open" (Reviewer N5DC), "not representative of core challenges in modern RL" and "lacks scalability" (Reviewer xY5j).
- Some open questions are left for future work, for instance the "remaining convergence gap" (Reviewer De38).

The overall contribution is still appreciated despite the weaknesses. Given the consensus from the reviews with all scores being 4 (weak accept), the recommendation is to accept the paper.
Decision: Accept (poster)