Transformers Learn Temporal Difference Methods for In-Context Reinforcement Learning
in-context reinforcement learning;policy evaluation;temporal difference learning
Yes
No ethics review needed.
1. Would it be possible for the authors to change the paper's title to something that better reflects what is shown by the paper?
2. Could the authors better justify the study of In-Context RL, both for Linear Attention-based transformers and otherwise?
The paper argues with proofs and empirical evidences for its results and is able to remove any possible doubt of their veracity. Furthermore it's a novel exploration of the capabilities of linear transformers and further helps the understanding of In-Context Reinforcement Learning.
The paper shows that linear transformers can execute TD(0) over their context both by constructing the matrices that would do so, but also by showing a pre-training method that gives rise to such properties. Finally they show similar constructions for other Policy Evaluation algorithms like TD-lambda, Average-Reward TD etc. Besides some empirical evidences the paper presents many proofs of its statements.
The title of the paper is rather strong, not only is the paper only about Linear Attention, as opposed to the commonly used Softmax attention, but also they prove that transformers with a specific kind of training implement TD-Learning, so I believe "Transformers Can Learn Temporal Difference Methods for In-Context Reinforcement Learning" would be more appropriate.  Furthermore one must wonder how relevant is studying Reinforcement Learning implemented by linear attention over such short horizons such as 40 steps, as most RL problems involve orders of magnitude longer trajectories, whereas attention tends to become computationally expensive and lose precision as the context length grows.
Yes
No ethics review needed.
1. The authors explain that in-context TD learning occurs when the task distribution is sufficiently challenging. I wonder why such a phenomenon occurs. They trained the transformer across thousands of Markov Reward Processes. Also, I wonder if the same TD learning effect would emerge internally within the transformer if it were trained in a single-task setup rather than a multi-task one.

2. In most prior research, transformers are directly applied to perform control tasks. However, in this paper, the transformer is used to estimate the value function. While a generalized approach to value function estimation is indeed important in reinforcement learning, this study does not directly address the control problem in RL. I would be interested to hear the authors' perspective on how their approach could be applied to control tasks in future research.

3. When a transformer estimates the value function, its internal operations align with TD learning. Most transformers, however, serve as policies rather than value function estimators. It would be interesting to know the authors’ perspective on whether, if the transformer were to act as a policy instead of estimating a value function, the internal operations would still resemble TD learning, or if another reinforcement learning algorithm would take place internally.
1. The authors mathematically analyze that defining a Transformer as a value function estimator and training it in a multi-task environment makes the attention mechanism within the Transformer equivalent to TD(0) learning. The proof process appears to be concrete, and the resulting application seems straightforward. Particularly, while previous reinforcement learning studies using Transformers have merely mentioned its high performance, this paper analyzes why the Transformer structure works so effectively. In this regard, this paper seems to present a notable novelty. 

2. The authors mathematically demonstrated that the internal operations of the transformer can be extended not only to TD(0) learning but also to various RL algorithms, such as TD($\lambda$), average reward TD, and residual gradient. This suggests that these algorithms can be applied in diverse ways to approximate a generalized value function through a transformer. 

3. The proposed theorems and corollaries in the paper are meticulously detailed in the appendix, where each proof is presented with clarity and rigor. This comprehensive approach allows readers to follow the logical reasoning behind each statement. This thoroughness underpins the robustness of the paper and enhances its reproducibility, which will be beneficial for future researchers.
The authors performed a mathematical analysis of the transformer architecture commonly used in the field of in-context reinforcement learning (ICRL). They proved that when a transformer with a linear attention structure is trained to estimate value functions in a multi-task environment, the forward pass within the transformer becomes equivalent to temporal-difference (TD) learning in reinforcement learning. Even when out-of-distribution task contexts are provided as input to the transformer, this TD learning effect is induced internally, enabling it to estimate the correct value function. The authors experimentally demonstrate this in a simple Boyan's chain environment.
1. The experimental setup and evaluation metrics may not effectively demonstrate the algorithm's impact. The authors defined the value function estimation error for new states in out-of-distribution (OOD) tasks as the evaluation metric in Figure 1. However, since the Boyan's chain environment is relatively simple, it's possible that generalization effects could also be achieved using standard architectures like MLPs and RNNs, resulting in graphs similar to those in Figure 1. This raises the question of whether the results in Figure 1 reflect the effects of the transformer's in-context TD learning or if they could be similarly obtained with conventional networks. To address this, it seems the authors should consider conducting comparative experiments.

2. The comparison experiments seem insufficient. If there were experiments comparing value estimation between the authors' proposed value-estimating Transformer algorithm and existing reinforcement learning algorithms that utilize Transformers (other than behavior cloning methods), it would support the claim that the proposed algorithm is more effective. Although this paper focuses more on mathematical analysis, it seems to not have enough experiments on other algorithms or tasks. 

3. The paper appears to be written in a way that makes it difficult to read. It seems to borrow equations and results from previous studies, but the notation and equations are challenging to follow at first glance. For instance, it is hard to understand how equations like Equation (4) were derived. A brief explanation of the existing equations would be helpful.
Yes
No ethics review needed.
1. In Corollary 3, the authors show that the transformer can implement TD(λ) updates with a specific mask matrix, $M_{TD(\lambda)}$. The matrix structure requires strict alignment of eligibility traces, which might be difficult to achieve in real tasks where temporal dependencies and eligibility traces shift dynamically. 

2. Lines 523-524 in the Conclusion are unclear. Are you trying to emphasize that policy iteration or optimization forms the key foundation?

3. In Lines 531-532, the claim seems unsupported; working well in a small-scale setting doesn’t necessarily imply the same for larger, more complex tasks.
It is a new try and an interesting direction to investigate the equivalence between transformers and in-context TD.
This paper provides theoretical analysis on the equivalence between a forward pass and in-context TD in a simple setting - linear single-layer transformers.
1. Many of the proofs rely on strict initial conditions, like specific embeddings and matrix structures. The validity of these proofs depends heavily on these controlled setups, which may not be feasible or sustainable in more complex, real-world settings. As a result, some claims in the paper appear intuitive and unsubstantiated.
2. While the paper claims the emergence of in-context TD learning through multi-task TD pretraining, the empirical analysis is constrained to quite simple tasks.
@inproceedings{
anonymous2024transformers,
title={Transformers Learn Temporal Difference Methods for In-Context Reinforcement Learning},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=Pj06mxCXPl},
note={under review}
}
Traditionally, reinforcement learning (RL) agents learn to solve new tasks by updating their parameters through interactions with the task environment. However, recent works have demonstrated that transformer-based RL agents, after certain pretraining procedures, can learn to solve new out-of-distribution tasks without parameter updates, a phenomenon known as in-context reinforcement learning (ICRL). The empirical success of ICRL is widely attributed to the hypothesis that the forward pass of these models implements an RL algorithm. However, no prior works have demonstrated a precise equivalence between a forward pass and any specific RL algorithm, even in simplified settings like transformers with linear attention. In this paper, we present the first proof by construction demonstrating that transformers with linear attention can implement temporal difference (TD) learning in the forward pass — referred to as in-context TD. We also provide theoretical analysis and empirical evidence demonstrating the emergence of in-context TD after training the transformer with a multi-task TD algorithm, offering the first constructive explanation for transformers’ ability to perform in-context reinforcement learning.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
in-context reinforcement learning policy evaluation temporal difference learning
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/f48c67a3c07ad4977007033ad17b360b938a9f5f.pdf
reinforcement learning
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Transformers Learn Temporal Difference Methods for In-Context Reinforcement Learning
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission