Towards Provable Emergence of In-Context Reinforcement Learning

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.

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

We 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.

and can be solved by relatively simple methods.

We 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.

The writing contains some inconsistencies and relies on setups or results from prior work without sufficient explanation.

We 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].

the Transformer architecture and the sparse parameterization considered in this work represent a highly simplified setting that lacks scalability to realistic applications.

We 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].

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?

The 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.

If so, does this find application in other reinforcement learning studies?

Multi-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.

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.

We 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.

Specifically, I would expect that unless the stationary distribution is uniform over the state space, such an approximation could introduce bias.

Since 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.

[1] Sutton, R., Barto, A., & others (1998). Reinforcement learning: An introduction. (Vol. 1) MIT press Cambridge.

[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.

[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.

[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.

[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.

[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.

[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.

[8] Zhang, R., Frei, S., & Bartlett, P. (2024). Trained transformers learn linear models in-context. J. Mach. Learn. Res., 25(1).

[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.

[10] Yan Duan, John Schulman, Xi Chen, Peter L. Bartlett, Ilya Sutskever, & Pieter Abbeel. (2017). RL$^2$: Fast Reinforcement Learning via Slow Reinforcement Learning.

[11] Jacob Beck, Risto Vuorio, Evan Zheran Liu, Zheng Xiong, Luisa Zintgraf, Chelsea Finn, & Shimon Whiteson. (2025). A Tutorial on Meta-Reinforcement Learning.

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.

A weakness is the heavy reliance on linear attention and one-hot features; practical relevance to nonlinear Transformers or continuous domains is still open.

We 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. As 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.

How does MSVE change if you vary the number of layers?

We 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.

Multi-task TD:

Multi-task MC:

We 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.

[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.

[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.

[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.

[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.

[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.

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.

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.

We 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.

Experiments use tiny MRPs (Boyan’s chain with 13 states). Claims lack validation in complex environments.

Since 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.

And the experiment misses a comparison to non-ICRL baselines (e.g., fine-tuning/offline-to-online RL).

We 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.

[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.

[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.

[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.

[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.

[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.

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.

The key open question from Wang et al. whether standard pre-training converges to the invariant/minimizer set remains unproven.

We 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.

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 ...

We 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. In light of the reviewer's suggestions, we did the following additional hyperparameter studies:

• 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.

  Multi-task TD:

  	5-layer	10-layer	60-layer
  m = 5	12.0	7.0	2.1
  m = 100	7.0	4.8	1.1

  Multi-task MC:

  	5-layer	10-layer	60-layer
  m = 5	10.9	7.8	3.0
  m = 100	6.3	5.3	1.9

  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.

• 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.

  Multi-task TD:

  	3-state MRP	10-state MRP
  m = 5	2.9	3.3
  m = 100	2.3	1.9

  Multi-task MC:

  	3-state MRP	10-state MRP
  m = 5	6.0	3.8
  m = 100	4.75	1.9

  The consistent decrease in MSVE confirms the in-context policy evaluation capability of the Transformer in both MRPs.

• 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.

  	Multi-task TD	Multi-task MC
  m = 5	2.2	3.5
  m = 100	0.3	1.1

  Therefore, the numerical results suggest the emergence of ICPE capability of the Transformers when trained on this Loop MRP.

We 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".

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.

We 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.

[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.

[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.

[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.

[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.