Paper: Learning to Solve New sequential decision-making Tasks with In-Context Learning (ICLR-2024)

Reviewer comment (Experiments):
The paper highlights the ability to perform well on unseen tasks but this actually relies on having a lot of demos from related tasks. Can the authors better clarify the relationship between the data they train on and the unseen tasks they evaluate on?

Author response:
1. **Relying on a lot of expert demonstrations** - Note that we use only 7 expert demonstrations for each level at test time which we’d argue is not a large number for learning (as echoed by reviewer qktf) entirely new tasks with unseen states, actions, dynamics, and rewards. 
2. **Relationship between the train data and test data** - We ask the reviewer to look at the common response where we clearly explained the difference between the train and test tasks.

Reviewer comment (Novelty):
Novelty-wise, the method is very similar to works like Prompt-DT, except actually requires stronger data assumptions (full expert demos).

Author response:
3. **Comparison with Prompt DT** -  Please refer to our common response where we explain in greater detail about this.

Reviewer comment (Novelty):
A lot of the insights in the empirical study are not that interesting, e.g. the paper highlights results like showing that in-context learning improves with trajectory burstiness, but it is not surprising that having demos similar to the query inside the context improves the performance. Can the authors give more clarity on what the most surprising, interesting takeaways are from the study?

Author response:
4. **Surprising aspect of our work** - Our paper is **first to demonstrate few-shot learning of new Procgen and MiniHack tasks from only a handful of demonstrations, without any online feedback, after training on a relatively small number of other tasks from the same domain (11 in the case of procgen and 12 in the case on MH)**. 
We believe this result is by no means trivial (and it surprised us) given that prior work was either only able to generalize to minor task variations (slightly different reward functions but same states, actions and dynamics, as in PromptDT) or required orders of magnitude more compute and training environments to generalize to more challenging tasks (as in Ada which trains for 100B timesteps on a task pool size of 25B, however at that point it’s not clear test tasks are OOD from train tasks, whereas in our case they are clearly OOD). Note that the train and test tasks are vastly different having different states, actions, dynamics, and reward functions. While we agree the approach is simple we believe this is in fact a strength of our work as it should be easy to adapt, apply to other settings, scale, and build upon by the community. 
Our paper performs an extensive analysis of different design choices and implementation details demonstrating their importance in getting the method to work. We believe these insights will be of interest to researchers and practitioners alike. In particular, we demonstrate a minimum model size, burstiness, data size and diversity are needed to obtain strong results. 
In addition, section 5.2 discusses in depth the limitations of this approach and the fact that it doesn’t work as well in certain settings with high stochasticity or high-stakes scenarios where a single wrong action can be fatal. Again, this supports the point that it is by no means “a given” that this approach works. In fact we believe there is a need for more research in this direction to develop methods that are robust to highly stochastic and unforgiving environments. It’s possible that offline learning and transformers are not even the best solution to these settings. All of this is to say that the results are much more nuanced and we try to convey this in our paper.

Reviewer comment (Novelty):
The proposed approach seems as a small adaptation of pre-existing approaches, i.e. causal transformer with multi-trajectory training, to new benchmarks (MiniHack and Procgen) in the offline setting.

Author response:
1. **Novelty** - We wrote a detailed comment about the novelty aspect of our work and how it is different from previous works. Could you please point us to the works, apart from the previous works considered in the common response, which are seemingly similar to the problem setting we consider?

Reviewer comment (Evaluation):
It is mentioned in the paper that all the results are produced using 3 seeds. In my opinion, for such noisy benchmarks evaluating on only 3 seeds is not enough to reliably estimate the mean and variance.

Author response:
3. **Number of seeds** - We also ran 5 seed experiments for procgen experiments and we didn’t notice any big differences in the conclusions. We will update the appendix with 5 seed results.

Reviewer comment (Evaluation):
The results are not clear to me - for example, in Figure 4 the episodic return is very low compared to the score reported by [2].

Author response:
4. **Comparing with Procgen[2]**  -  The results in the original Procgen paper are for a completely different setting than ours so this is not a fair comparison. **In the Procgen paper, agents are trained with PPO on multiple levels of a single game (e.g. Ninja) and tested on unseen levels from the same game (i.e. also Ninja), thus probing IID generalization In our paper, agents are trained offline on data from multiple tasks (e.g. Ninja) and tested on entirely new tasks at test time (e.g. Jumper) thus probing OOD generalization which is more challenging. ** Thus, the lower performance in our paper is expected.

Reviewer comment (Experiments):
Are the above few trajectories (1-7) sampled from the same level as the query level?

Author response:
6. **Few-shot trajectories sampled from the same level?** Yes, the trajectories are from the same level we are testing on. Note that, this is a challenging setting (look at common response for more details) where test tasks are OOD. Moreover, the environments are stochastic which adds more challenges. Hence the model might still struggle to generalize to OOD tasks just from a few demonstrations.  We will clarify this in the paper.

Reviewer comment (Evaluation):
Why is the return in Figure 4 so low compared to the return reported in [2]?

Author response:
4. **Comparing with Procgen[2]**  -  The results in the original Procgen paper are for a completely different setting than ours so this is not a fair comparison. **In the Procgen paper, agents are trained with PPO on multiple levels of a single game (e.g. Ninja) and tested on unseen levels from the same game (i.e. also Ninja), thus probing IID generalization In our paper, agents are trained offline on data from multiple tasks (e.g. Ninja) and tested on entirely new tasks at test time (e.g. Jumper) thus probing OOD generalization which is more challenging. ** Thus, the lower performance in our paper is expected.

Reviewer comment (Experiments):
Is the Procgen dataset evaluated on the easy or hard difficulty mode?

Author response:
7. **Procgen task difficulty** - We consider “easy” difficulty mode in our experiments.

Reviewer comment (Presentation):
Are the results in Figure 3 normalized?

Author response:
8. **Normalized scores** - No, we do not normalize anything and we report the raw scores for both MiniHack and Procgen.

Reviewer comment (Novelty):
Existing papers such as Prompt-DT (Xu et al, 2022) and AdA (Team et al, 2023) have shown similar results (in some cases, with even less presumptive data than full demonstrations) though in different domains. So it's perhaps not too surprising that we see this type of generalization in Procgen as a result. The main result that models trained with demonstrations in the context can perform better than models without the demo context is also expected.

Author response:
1. **PrompDT and AdA**: We would like the reviewer to refer to the common response where we contrast our method with prior work. We differ from PromptDT paper in terms of many aspects including multiple environment training, scale of our models, OOD evaluations etc. With regards to AdA, we would like to highlight that the setting is different from our work. AdA considers online RL setting with transformers in contrast with offline few-shot imitation learning setting, and requires orders of magnitude more compute and training environments to generalize to more challenging tasks (100B timesteps on a task pool size of 25B). Regarding surprising aspects of the results, we want to emphasize again that the setting we consider in this paper is very challenging and to the best of our knowledge this setting hasn’t been explored in prior work.

Reviewer comment (Experiments):
There are a couple of claims that do not seem sufficiently supported: (1) meta-RL methods "tend to be difficult to use in practice and require more than a handful of demonstrations or extensive fine-tuning," (2) "in sequential decision-making it is crucial for the context to contain full trajectories (or sequences of predictions) to cover the potentially wide range of states the agent may find itself in at deployment" (see Questions).

Author response:
2. **Meta RL** - The problem setting considered in our paper differs from that of Meta-RL, where training is online or offline and, at test time, the agent learns from online interactions and feedback (i.e., rewards) within a new environment. In contrast, our paper focuses on the few-shot imitation learning setting, where training occurs offline based on expert demonstrations from a set of tasks, and at test time, the agent learns a new task from only a handful of offline demonstrations without interactions or feedback in the new environment. Therefore, while Meta-RL is loosely related – a fact we acknowledge in the related work section – these approaches are not directly comparable, as they operate under vastly different assumptions. However, following your suggestion, we will rephrase that sentence to ensure our claims are well supported by our experiments.

Reviewer comment (Novelty):
The concepts of burstiness from Chan et al and trajectory burstiness have a pretty weak relation. In the case of this paper, it seems pretty clear from the get-go that demonstration contexts from the same level as the query would be more relevant than from any other levels.

Author response:
6. **Comparison to burstiness in Chan et al.** We think the burstiness in chan et al and trajectory burstiness are closely connected. Chan et al defines a bursty sequence as a sequence where there are multiple occurrences of the same classes which are similar to query (Figure 1 in [1]). Similarly, in our work, a trajectory bursty sequence is a sequence where there are multiple trajectories which are from the same level as the query.

Reviewer comment (Writing):
What does trajectory burstiness mean for the zero-shot model in Fig. 6(a)?

Author response:
5. **What does burstiness 0.0 mean?** Burstiness 0.0 means that in the entire dataset, there is a zero probability of finding the sequences which are bursty.

Reviewer comment (Experiments):
"[Meta-RL methods] tend to be difficult to use in practice and require more than a handful of demonstrations or extensive fine-tuning" --> Including some of these comparisons in the experiments would be help support this statement.

Author response:
2. **Meta RL** - The problem setting considered in our paper differs from that of Meta-RL, where training is online or offline and, at test time, the agent learns from online interactions and feedback (i.e., rewards) within a new environment. In contrast, our paper focuses on the few-shot imitation learning setting, where training occurs offline based on expert demonstrations from a set of tasks, and at test time, the agent learns a new task from only a handful of offline demonstrations without interactions or feedback in the new environment. Therefore, while Meta-RL is loosely related – a fact we acknowledge in the related work section – these approaches are not directly comparable, as they operate under vastly different assumptions. However, following your suggestion, we will rephrase that sentence to ensure our claims are well supported by our experiments.

Reviewer comment (Experiments):
"Our key finding is that in contrast to (self-)supervised learning where the context can simply contain a few different examples (or predictions), in sequential decision-making it is crucial for the context to contain full trajectories (or sequences of predictions) to cover the potentially wide range of states the agent may find itself in at deployment." --> Full trajectories as opposed to what? The experiments only show comparisons between full demos vs no demos, and didn't study other potential contexts, such as partial demos or non-expert trajectories. I think a study of different potential contexts and what is required for in-context learning would be interesting.

Author response:
4. **Partial Demonstrations** - We want to highlight that we consider partial demonstrations in our procgen experiments because of the long episode lengths in some of the procgen tasks.

Reviewer comment (Theory):
"This means that the agent manages to perform well on the new task even without copying actions from its context. This suggests the model is leveraging information stored in its weights during training, also referred to as in-weights learning" --> Could you elaborate on how not copying the context actions suggests in-weights learning as opposed to ICL? This conclusion seems to equate ICL with the ability to copy context actions.

Author response:
3. **Copying Actions** - Thank you for pointing this out. Recent works in ICL, such as 'Incontext Learning and Induction Heads'[1] by Anthropic, suggest that 'induction heads' are a key feature in transformer models for ICL. Induction head is  a circuit whose function is to look back over the sequence for previous instances of the current token (call it A ), find the token that came after it last time (call it B ), and then predict that the same completion will occur again (e.g. forming the sequence [A][B] … [A] → [B] ). This is what we mean by copying in our context. We understand that the term 'copying' is used loosely in the paper which caused this confusion. We will rephrase our sentence structure to reflect this understanding, including relevant citations.

Reviewer comment (Experiments):
During testing, the agent is given a handful of expert demonstrations. Are these all demonstrations on the same task and same level? If not, how does this work for the baselines hashmap if they are using demos from different levels? If so, why is the transformer trained with several sequences of demos from different levels? Why not just train with demonstrations from the same level and task always?

Author response:
* **Few-shot evaluation** - During few-shot evaluation, the model is given access to a limited number of expert demonstrations from the task on which it is being evaluated. For instance, when assessing performance on a Ninja task from Procgen, we condition our model on a few expert demonstrations of the Ninja task, specifically from the same level on which it will be tested. We then unroll the policy of the transformer in the online environment and measure the obtained reward. This itself is a challenging evaluation because the model never saw the “Ninja” task during the training, hence OOD, and learning from just a handful of demonstrations in this scenario could be very challenging.  
* **Hashmap evaluation** - Hashmap baseline evaluation is similar to that of the model. Again, taking the example of the “Ninja” task, we hash a few expert demonstrations coming from the same level on which it will be tested. This baseline suffers if the environment is stochastic but could act as an oracle in case the environment is deterministic. Finally, we use the same demonstrations for all the baselines and our method for fair comparison.  
* **Training** - When the trajectory burstiness probability is 1.0 (which is the case for figure 3 and 4), in each sequence the demonstrations come from the same level of a given task. So each sequence in the training set looks like this [ demonstration 1 from task A level A, demonstration 2 from task A level A, demonstration 3 from task A level A, demonstration 4 from task A level A ….]. In this work, we have 10k levels per  task and 11 training tasks in case of procgen resulting in 11 * 10k = 110K training sequences.

Reviewer comment (Experiments):
Related to this, how should one interpret the statement “we collect offline data from 11 Procgen tasks and train a transformer on Procgen sequences comprising of five episodes from the same level,” which suggests that all demos in the context come from the same level? What does burstiness even mean here if the levels are never varied?

Author response:
* **Training** - When the trajectory burstiness probability is 1.0 (which is the case for figure 3 and 4), in each sequence the demonstrations come from the same level of a given task. So each sequence in the training set looks like this [ demonstration 1 from task A level A, demonstration 2 from task A level A, demonstration 3 from task A level A, demonstration 4 from task A level A ….]. In this work, we have 10k levels per  task and 11 training tasks in case of procgen resulting in 11 * 10k = 110K training sequences.  
* **Sequence construction** - If all or at least two demonstrations are coming from the same level in a given sequence, then the sequence is called a bursty sequence with trajectory burstiness probability 1.0. This means that the relevant information required to solve the task is entirely in the sequence hence this design aids in-context learning. However, note that the trajectories, despite being from the same level, may vary because of the inherent environment stochasticity which means that the model still needs to generalize to some extent  to handle these cases.

Reviewer comment (Experiments):
What does it mean for BC-1 to condition on ‘one demonstration’? Does this mean it is given a full demonstration in the same task and same level? In other words, does the context look like this: [expert demo, history observed so far]? How would this be different from the proposed method if it were just limited to training on two sequences?

Author response:
**Behavioral Cloning (BC-1) baseline** - Yes, BC-1 means during the *evaluation time* (and *not* during the training), we condition the model with one expert demonstration. This is different from our method because during training BC-1 and BC-0 are trained with a single trajectory sequence and they only differ just in the way they are evaluated. The sequence construction looks exactly like you rightly mentioned. The aim of BC-1 baseline is to check if the model is ever making use of the expert demonstration info during the test time. However, this evaluation could be somewhat OOD for the model to generalize which is reflected in the results.

Reviewer comment (Evaluation):
What is the maximal achievable reward in each of the environments? This could be helpful to better understand the final results.

Author response:
**Upper bound of performance** - We will add these numbers in the appendix. However, we do want to highlight that, to our knowledge, there are not so many works which consider the setting we consider in our paper so it is hard to benchmark our current results. Ideally we want to move our numbers close to that of an oracle which is obtained by doing online RL from scratch in the test environment and these are already reported in [1].

Reviewer comment (Novelty):
The paper could better distinguish its work from prior methods. A more thorough comparison would help readers better understand the present problem setting and method.

Author response:
**Additional discussion on Related works** - We agree that additional discussion about recent works would be beneficial. We will add the works you mentioned in your review and extend our related works section in the appendix.

Reviewer comment (Novelty):
The method appears to be very similar to Prompt-DT [1], perhaps without the return conditioning. There is already a short discussion in the related work section, but this ought to be carefully dissected.

Author response:
**Similarities with PromptDT** - We wrote a detailed response to this concern in the “common response” section above where we contrast our work with  all the papers mentioned in your review. We encourage you to read the response and let us know if you have any additional questions.

Reviewer comment (Novelty):
The method is also very similar to DPT [2], which considered training and conditioning on expert demonstrations to solve new tasks. If there is a difference, both of these papers seem like highly relevant baselines.

Author response:
**Additional discussion on Related works** - We agree that additional discussion about recent works would be beneficial. We will add the works you mentioned in your review and extend our related works section in the appendix.

Reviewer comment (Novelty):
It is also likely worth distinguishing the method from other in-context RL works like [3] and [4].

Author response:
**Additional discussion on Related works** - We agree that additional discussion about recent works would be beneficial. We will add the works you mentioned in your review and extend our related works section in the appendix.

Reviewer comment (Novelty):
Beyond transformers, there is additional work on meta/few-shot imitation learning that could be helpful to discuss.

Author response:
**Additional discussion on Related works** - We agree that additional discussion about recent works would be beneficial. We will add the works you mentioned in your review and extend our related works section in the appendix.

Reviewer comment (Experiments):
How long are each of these sequences? I.e., what is T?

Author response:
**Value of T** -  In case of Minihack, the average episode length  (T) is about 100 steps. Incase of procgen, some tasks have very large episodes (> 1000 steps). So, to accommodate multiple trajectories, we truncate the episodes to have at maximum of 200 steps.  For the procgen evaluations, we do transformer rollouts for a maximum of 200 steps for that reason.

Reviewer comment (Theory):
In what settings would you expect this method to work (or these analyses be useful) under the current assumptions, beyond gameplaying?

Author response:
**Applicability beyond game playing** - We would like to highlight that our approach doesn’t use any domain-specific information or assumptions, making it broadly applicable for any other domains beyond games. We demonstrate that our approach works for two different domains without any specialized adjustments (Procgen and MiniHack), further supporting the generality of our method. Note that these domains have vastly different observation spaces, one consisting of discrete grids and the other high-dimensional images.  
One limitation of our approach, which we discuss in section 5.2, is the finite context length so if the trajectories are longer than the context our method suffers. Also, we expect performance to be worse in very stochastic or "unforgiving" environments where a single wrong action is fatal. Please see Section 5.2 where we talk extensively about these failure modes and show how they affect performance.