Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining

We sincerely appreciate the reviewer's detailed review and insightful feedback on our paper. In the following, we hope to address some points raised in the review.

“I am unsure if the proposed framework can offer insights on improving pretraining approaches for ICRL. I recommend that the authors address this in their manuscript.”

Thanks for the suggestion. In Appendix A.1 in our revision, we have included a discussion of some practical implications of our results, as well as potential ways to further improve ICRL practice.

As examples, our results highlight that the training labels (expert actions) in the offline dataset matter a lot. When it is itself a proper algorithm that only depends on the past observations, we can learn ${\sf Alg}_E$ (cf. Theorem 9). By contrast, when ${\sf Alg}_E$ is the ground truth optimal action $a^\star$ (which involves knowledge of the underlying MDP), we can learn the posterior average of this algorithm given past observations, which corresponds to the Thompson sampling algorithm as in Decision-Pretrained Transformers (cf. Theorem 11).

“Can you provide further explanation on the importance of calculating loss on the entire trajectory in Equation 3?”

The reason we calculate the loss over the entire trajectory is that we want the learned algorithm to imitate the expert algorithm at every time step. Unlike in-context learning for supervised learning problems (Garg et al. 2022) in which the loss only needs to be calculated on a new test data point since an i.i.d. assumption on the in-context data is assumed, in ICRL the actions chosen by the expert algorithm may depend on the time $t$ and previous actions and rewards. Therefore, to fully imitate the expert at every time step, the loss needs to be calculated on the entire trajectory.

We are grateful to the reviewer for their thorough review and valuable feedback on our paper. In the following, we hope to address each point raised in the review in the order presented.

The paper currently lacks a strong discussion of limitations… mismatch between offline and expert algorithm… non-vacuous practical take-aways?

Thanks for the suggestion. We have added a discussion of the limitations as well as practical implications in our revision, which is currently in Appendix A.1 (due to the space limit) with a pointer in the Conclusion section in the main text. As one example of a non-vacuous practical take-away, our results highlight that the training labels (expert actions) in the offline dataset matter a lot. When it is itself a proper algorithm that only depends on the past observations, we can learn ${\sf Alg}_E$ (cf. Theorem 9). By contrast, when ${\sf Alg}_E$ is the ground truth optimal action $a^\star$ (which involves knowledge of the underlying MDP), we can learn the posterior average of this algorithm given past observations, which corresponds to the Thompson sampling algorithm as in Decision-Pretrained Transformers (cf. Theorem 11). A discussion about this can also be found in “Three special cases of expert algorithms” (Page 3) where the “reduced algorithm” $\overline{\sf Alg}_E$ denotes this posterior average.

… partial observability. If the expert algorithm depends on unobservables, I believe it would be possible to create settings that lead to a failure to imitate the expert; essentially leading to “Self-Delusions” - see [1]. Taking a wild guess here, but this would probably imply additional assumptions around identifiability?

We appreciate the reviewer raising the question of partial observability. It would be helpful if the reviewer could clarify what they mean by "partial observability" in this context.

We agree that when the expert algorithm depends on variables that are unobserved by the learner, the transformer will fail to perfectly imitate the expert. Our Theorem 6 formalizes this notion: we show that the transformer will imitate the expectation of the expert algorithm conditional on the observables, rather than the exact expert algorithm itself. In this setting, the full environment (the decision process instance) is unobservable to the learner.

However, the "self-delusions" discussed in [1] seem to refer to a different type of "partial observability" than what we are considering above. It would be helpful if the reviewer could clarify what they mean by partial observability in this context. Specifically, is the reviewer referring to the formalism of “partially observable Markov decision processes”? We want to ensure we fully address the reviewer's concern, so additional details would allow us to refine our discussion accordingly.

[1] Shaking the foundations: delusions in sequence models for interaction and control, Ortega et al. 2021.

… amortized algorithm vs expert algorithm …

As posited by the reviewer, our statistical theory only ensures that the pretrained transformer learns an “amortized algorithm” that is only guaranteed to match the expert algorithm in-distribution (same training/test distribution for the bandit/MDP instances), even though our algorithmic construction gives the existence of a transformer that actually approximates the expert algorithm in the whole input space.

In some of our early experiments (not included in the paper), we have indeed found that the learned transformers do not generalize that well to out-of-distribution instances such as shift of reward distribution. Similar phenomena have been observed in other types of in-context learning problems, e.g. Garg et al. (2022), among many others. We have added a brief discussion about this in Appendix A.1 in our revision.

(How) Could ICRL be used to surpass the performance of the expert (e.g. going from human chess as the expert to superhuman chess)? … Or, alternatively, is it possible make a theoretical statement showing that ICRL performance is capped at exactly the level of the expert (even if we used ICRL to generate a new dataset of supervised trajectories where it slightly surpasses expert performance)?

In our setting, we can indeed only expect the learned ICRL algorithm to at most match the expert (i.e., $\overline{\sf Alg}_E$) but not surpass it. However, as the reviewer pointed out, that does not rule out surpassing the expert, e.g. by online RL style training where the transformer can interact with the environment in an online fashion. We would like to leave that as future work.

After Eq (2): “all the algorithm outputs [...] along the trajectory can be computed in a single forward propagation.” Does that also imply having masked (causal) attention?

Yes, we use the masked (causal) attention in this work (see Definition 1).

We are grateful to the reviewer for their thorough review and valuable feedback on our paper. In the following, we hope to address each point raised in the review in the order presented.

I am not sure the conclusion is a little bit narrow. The conclusion is limited in a supervised paradigm, using pretraining, and in an ICRL setting. It seems that these conditions restrict the conclusions to be general.

Regarding ICRL: We believe the ICRL paradigm (especially with transformers) has received lots of recent attention, and many works keep improving it including Algorithm Distillation (AD; Laskin et al. 2022), Decision-Pretrained Transformers (DPT; Lee et al. 2023), and more recent ones such as AMAGO [1]. In general, we believe ICRL is a promising way of using transformers to do reinforcement learning, and offers an alternative to other paradigms such as Decision Transformers. Regarding supervised pretraining: we believe this is a natural way for transformers to learn to do ICRL, and is actually an advantage (easier and more stable to train) compared with alternatives such as reward maximization using RL training.

[1] AMAGO: Scalable In-Context Reinforcement Learning for Adaptive Agents, Grigsby et al., 2023.

Would different pre-trained method lead to different results?

Right, other pretraining methods may lead to different results. While we use a supervised MLE loss, there exist other training methods such as cumulative reward maximization, e.g. in Duan et al., (2016). However, that method does not give similar guarantees in our setting. We have added a discussion about this in Appendix A.1 in our revision.

Can learning-from-scratch DT figure out similar conclusion? What’s the difference?

By “DT” does the reviewer mean Decision Transformers, and “learning-from-scratch DT” as some adaptation of DT to the ICRL setting? (Since the original paradigm of DT only learns a goal-conditioned policy for a fixed environment and cannot do ICRL in new environments). In that case, the “Learning-from-scratch DT” in our opinion is very different from the original DT and not very different from our setup other than perhaps goal-conditioning.

Will your conclusions help to inspire the empirical study? Such as figuring out better DT-based offline RL algorithms when using D4RL datasets?

We believe our conclusions may provide several intuitions for empirical study on ICRL:

• The training labels (expert actions) in the offline dataset matter a lot. For example, AD (using actions generated from some expert algorithm) learns the nominal algorithm, DPT (using the optimal actions $a_i^*$) learns Thompson sampling algorithm instead.
• The distribution ratio between the offline algorithm and the expert algorithm may affect the generalization properties of the learned algorithm. As shown in both our theories (Theorem 6) and simulations (Figure 4), it is important to have a small distribution ratio between the offline algorithm $A_0$ and the expert algorithm $A_E$ in order to achieve good online performance, otherwise the performance may substantially degrade. This suggests that it is good to always incorporate trajectories generated purely from the expert algorithm (“on-policy ICRL”) into the offline dataset, when possible.

We would like to thank the reviewer for taking the time to review our paper and provide helpful feedback! We truly appreciate your comments and suggestions, and believe they can make our work better. In the following, we hope to address each of the points made in the review, in the order as they came out.

… Two main sets of results are not as dependent/relevant to Transformers… any supervised pretraining models that satisfy the approximate realizability assumption should suffice Theorem 6. Is this true?

Theorem 6 indeed does not depend on transformers, and it holds for any model class that satisfies the approximate realizability assumption. We believe this is not a concern but rather a benefit, and this is expected as the result only involves statistical analysis that does not rely on the structure of the model class.

For the theorems in Section 4, the paper mostly constructs a specific Transformer structure (e.g., 3-head attention + MLP with a specified QVK) and says that this certain Transformer can mimic the gradient update in the no-regret RL/bandit algorithms. My understanding is that the Transformer is an overkill, and this can generally apply to NNs that have similar structures. Is this true?

We agree with the overall statement that the constructions in Section 4 can be done on other NN architectures that have the similar structures. However, we emphasize that in order for the construction to be efficient (with a small number of layers and parameters), our proof heavily uses the specific structure of attention, specifically its ability for token $t$ to efficiently retrieve information from token $1$ to $t-1$, as well as the inner product (between query and key). It is unclear to us whether other sequential architectures such as RNNs/LSTMs can do these as efficiently.

The idea of proofs of Theorems in Section 4 seems essentially similar to that of Bai et al. 2023, though with nuances.

For our proof ideas for Section 4, we believe the only core similarity with Bai et al. (2023) is that we also use some form of gradient descent as an iterative algorithm to approximate more sophisticated procedures such as matrix inverse. Other than that, our constructions involve several new ingredients, such as accelerated gradient descent (Theorem 8), approximate matrix square root (Theorem 10), and dynamic programming (Theorem 12), as well as lower-level operations such as token copying. In addition to transformer constructions, Section 4 also involves RL arguments for propagating the approximation error from the single-step log-probabilities to the regret, and for bounding the regret of the soft UCBVI algorithm (Theorem 13).

If I understand correctly, the distribution ratio R in Equation (6) can be arbitrarily large. This significantly limits the usefulness of the bound. In fact, in many cases (except the Algorithm distillation case), this ratio should be sufficiently large as Alg_E and Alg_0 are usually different, and the difference is accumulated over the entire trajectory.

We agree that the distribution ratio R can be large when Alg_E and Alg_0 are different. We believe how to get rid of the dependence on R, e.g. under additional structural conditions for the RL problem or the algorithms is an important question, which we would like to leave as future work.

Minor issues: $M$ is used as both an MDP environment and the number of heads… $n$ in Equation (5) and onwards has not been formally introduced. Is this the number of trajectories/samples?

Thanks for noticing these. We have fixed the $M$ in our updated version. $n$ is the number of pretraining trajectories ($nT$ is the number of samples), which we introduced in the first paragraph of Section 3.

Thank you again for your review! As the author-reviewer discussion period ends today, please let us know if our response has adequately addressed your questions and concerns, and we'd be happy to answer if you have additional questions.

Reading all other reviews and the authors' response I personally believe that all issues raised, including all issues raised by ycZf, have been well addressed by the authors. I have therefore raised my confidence accordingly.

I agree with the authors that the independence of Theorem 6 on the transformer architecture is an advantage, and I disagree with ycZf that this makes the title a mismatch or implies that the results are only weakly related to transformers - the overwhelming fraction of modern work on in-context RL is done via transformers, which makes it fair, in my opinion, to focus the paper on transformers, even though some of the main results hold for other neural network architectures. Similarly, the results in Section 4 benefit from a focus on architectures that have attention (aka transformers) - while weak (and as the authors state, inefficient) versions of the results could be obtained for other neural network architectures, transformers are the important special case that give Section 4 (practical) relevance.

Issue 3 raised by ycZf, regarding the practicality of the bound in Eq. (6), is somewhat warranted, I think: of course it would be nice to have a tighter bound - but given the nature of in-context RL, I would be surprised to see any theory that could come up with a much tighter bound. I would argue that Eq. (6) is a proper formalization of a very naturally and intuitively expected condition. I therefore consider Eq. (6) highly valuable as a formal expression of a fundamental condition for in-context RL to work. Adding additional structural assumptions (as pointed out by the reviewers) is interesting and of high practical relevance, but will not invalidate the more general Eq. (6). In addition, finding other practical ways to keep the bound low are important (but again: this seems a very natural thing to do in in-context RL; analogously in offline-RL there is only so much that can be done about distributional shift; under arbitrary shifts performance bounds are expected to become arbitrarily large). I would be happy to discuss this further in the reviewer discussion, and could see how one might make a point that Eq. (6) is perhaps a contribution of lower significance w.r.t. immediately coming up with an improved ICRL method - I personally would argue that that is not the aim of the paper.