Learning to Solve New sequential decision-making Tasks with In-Context Learning

We thank the reviewer for the comments. In this response we will respond to all the comments and questions the reviewer raised.

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.
3. Comparison with Prompt DT - Please refer to our common response where we explain in greater detail about this.
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.

We thank the reviewer for their time. Here are the responses to the questions -

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?
2. Comparison with CQL - We refer the reviewer to look at the common response where we do comparison of our method with all the related previous works including CQL.
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.
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.
5. Clarity about test demonstrations - The expert demonstrations used at test time come from an expert policy trained on that particular environment. For example, when evaluating our transformer model on Ninja task, we use an expert PPO policy trained only on Ninja to collect expert trajectories and condition on those demonstrations.
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.
7. Procgen task difficulty - We consider “easy” difficulty mode in our experiments.
8. Normalized scores - No, we do not normalize anything and we report the raw scores for both MiniHack and Procgen.

We thank the reviewer for thoughtful comments. We appreciate the reviewers comments on strengths of our work. In this response, we will address the questions you raised in your review.

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

We thank the reviewer for their thoughtful comments and feedback. We appreciate the reviewer acknowledging the challenges posed by our setting, our thorough experimental analysis and discussion on limitations of our work. In this response, we hope to successfully answer your remaining concerns.

1. 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.
2. 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.
3. Regarding Evaluation
   • 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.
4. 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.
5. 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.
6. 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].

Comparison with PromptDT -

Here, we point out the core differences between PromptDT and our work.

Our work has some parallels with PromptDT in how we construct sequences. Specifically, the sequences in the PromptDT paper are structured as [(prompt_trajectories), query trajectory], with the transformer being optimized using MSE loss solely on the query action predictions. Early in our research, we conducted experiments where the loss was exclusively focused on the query predictions and observed that the models had difficulties in generalizing at scale. We discovered that the design choices we emphasize in our paper like trajectory burstiness, data and model scale and environment stochasticity are essential for enabling models to few-shot learn new sequential decision-making tasks in entirely new tasks, including unseen states, actions, dynamics, and rewards.

Comment about Partial Trajectories in PromptDT - Finally we want to highlight that PromptDT evaluated with partial demonstrations (not expert) during the test time and they do show promising results. However, the caveat here is that 1) their models are conditioned on return-to-go tokens and 2) their experiments are continuous control tasks. For these continuous control tasks, the task can be defined just with the help of achievable return. This means that they do not need full demonstrations to perform well on new tasks. This is vastly different from our setting and we are not sure if promptDT would be successful in the setting we consider.

[1] Xu M et al. Prompting decision transformers for few-shot policy generalization. International conference on machine learning 2022.

[2] Kumar, Aviral, et al. "Conservative q-learning for offline reinforcement learning." Advances in Neural Information Processing Systems (2020): 1179-1191.

[3] Laskin M et al. In-context reinforcement learning with algorithm distillation. arXiv preprint arXiv:2210.14215. 2022.

[4] Lu C et al. Structured state space models for in-context reinforcement learning. arXiv preprint arXiv:2303.03982. 2023.

[5] Kirk R et al. “A Survey of Zero-shot Generalisation in Deep Reinforcement Learning” Journal of Artificial Intelligence Research 76 (2023) 201-264

[6] Lee JN et al. Supervised Pretraining Can Learn In-Context Reinforcement Learning. arXiv preprint arXiv:2306.14892. 2023.

[7] Adaptive Agents Team, “Human-Timescale Adaptation in an Open-Ended Task Space”.