Paper: Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining (ICLR-2024)

Reviewer comment (Novelty):
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.

Author response:
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).

Reviewer comment (Theory):
Can you provide further explanation on the importance of calculating loss on the entire trajectory in Equation 3?

Author response:
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.

Reviewer comment (Theory):
Another issue that is missing from the discussion (and perhaps even the analysis) is 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? This point becomes even more important when considering humans as “expert policies”.

Author response:
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.

Reviewer comment (Experiments):
This is perhaps more a question than a weakness with actionable improvement (so consider this point addressed if it can easily be clarified/answered). I am wondering about the relationship between the transformer learning an **amortized algorithm** (i.e. the trained net behaves as if it were running the algorithm under the hood, but only on-distribution because it is not actually running the algorithm but has learned a function that produces the same outputs) **vs. actually implementing something close to the expert algorithm**. It seems to me that the paper suggests the latter (which is why it is important that Transformers can implement accelerated gradient descent and matrix square roots algorithmically). Naively I would have guessed that the transformer learns an amortized solution rather than the actual algorithm. On-distribution the two are indistinguishable, but typically the amortized solution generalizes less well. There are some experiments that could tease apart the two possibilities (but I’d also be interested in a general comment by the authors; most importantly whether anything in the paper would break with an amortized solution). Sketches for experiments (I do not necessarily expect to see them performed, but they might help clarify the discussion):  
- Let the offline algorithm not fully cover the state space. At test time input these parts of the state-space (e.g. a coordinate input in a gridworld that is much larger than anything the Transformer has ever seen; but same dynamics, same reward function, etc). An amortized solution should struggle in this case (and potentially break down catastrophically), whereas learning the correct algorithm should not.  
- Change the reward function at test-time (e.g. picking up bananas in a gridworld has positive reward and apples have negative reward during pre training, but at test-time the rewards are reversed). The amortized solution should struggle in this case.

Author response:
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.

Reviewer comment (Experiments):
Would different pre-trained method lead to different results?

Author response:
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.

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

Author response:
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.

Reviewer comment (Theory):
My main concern is that the two main sets of results are not as dependent/relevant to Transformers. For Theorem 6 (supervised pretraining can imitate the expert algorithm) — this seems to be a general result that is not specific to Transformers. In fact, any supervised pretraining models that satisfy the approximate realizability assumption should suffice Theorem 6. Is this true?

Author response:
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.

Reviewer comment (Theory):
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?

Author response:
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.

Reviewer comment (Novelty):
The idea of proofs of Theorems in Section 4 seems essentially similar to that of Bai et al. 2023, though with nuances.

Author response:
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).

Reviewer comment (Theory):
If I understand correctly, the distribution ratio $\mathcal{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.

Author response:
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.

Reviewer comment (Writing):
$M$ is used as both an MDP environment and the number of heads.

Author response:
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.

Reviewer comment (Writing):
Not sure if I missed this — $n$ in Equation (5) and onwards has not been formally introduced. Is this the number of trajectories/samples?

Author response:
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.