Retrieval-Augmented Decision Transformer: External Memory for In-context RL

We thank the reviewer for their feedback. **Novelty and Technical Challenges:** It is true that we are not the first to leverage retrieval-augmentation. However, there are a number of fundamental differences to existing works. The main challenge we address in this work is to improve in-context RL via retrieval-augmentation. To achieve this, we address the following technical challenges: - **Reliance on Long Context:** Existing ICRL works rely on keeping multiple episodes in their context, which is challenging as episodes in RL can last for thousands of timesteps. To address this, we leverage retrieval-augmentation to alleviate the need for long contexts that span multiple episodes. Importantly, the retrieval-augmentation in RA-DT enables access to the entirety of previously collected experiences at inference time with small additional cost (see below). This would be very computationally expensive for existing ICRL methods. - **Challenges of Retrieval in RL:** Targeting in-context RL with retrieval requires a number of careful technical considerations, on which we conduct a number of ablation studies (see Appendix E). Therefore, we leverage a reweighting mechanism (similar to [1]) to provide useful context trajectories at training and to produce high reward behavior at inference time. For this, we come up with measures for utility that are adequate for the RL setting, which we study in our ablation section (see Section 4.6, Appendix E.2). In addition, we found it useful to introduce retrieval regularization strategies (Appendix E.3). - **Domain-agnostic Retrieval:** To use retrieval-augmentation in RL, an embedding model is required. While embedding models are readily available for language, they are rarely available in RL. Therefore, we proposed to leverage the FrozenHopfield [2] mechanism to map RL “tokens” to the language space. This allows us to leverage frozen LMs as embedding models for RL trajectories, which enables retrieval without domain-specific pre-training. In contrast, existing works that focus on retrieval in RL (online RL [3] and the Go [4]) relied on domain-specific embedding models. Furthermore, our work aims to provide an **improved understanding of in-context RL**. In addition to the proposed method, we also provide new insights into how in-context learning behaves in offline RL. Previous works primarily considered grid-world environments for showcasing in-context RL. In contrast, we also evaluate more complex environments and find that ICL does not transfer to those domains. Consequently, we illuminate the challenges that current in-context RL methods face (see Section 5). **Quality of Data in Retrieval Buffer:** We want to clarify that we do not make use of any demonstrations during online evaluation in this work. During inference, the retrieval buffer is empty and entirely populated by the agent’s generated experience. Consequently, on grid worlds, the agents are able to benefit from (initially) suboptimal data generated in previous trials. Also, see Figure 17 for an example.

Thank you for your helpful feedback. We are glad that you find retrieval-augmentation to be a compelling idea and the domain-agnostic embedding interesting. **Poor performance on non-gridworlds:** While RA-DT does outperform all baselines on grid-worlds, the reviewer is correct in stating that RA-DT does not achieve better in-context learning abilities on more complex environments such as Meta-World. RA-DT does outperform baselines in terms of final average performance on hold-out tasks in Meta-World and DMControl. However, note that, to the best of our knowledge, no other method is capable of solving these environments by improving in-context learning behaviour without leveraging expert demonstrations. This is the reason we added a longer discussion section on the shortcomings of existing ICRL methods, and directions we believe are critical for future ICRL research to study. To answer your question (Q3): we believe that RA-DT and other ICRL methods may suffer from all the listed shortcomings, and future algorithms could greatly benefit from improvements along the axes highlighted in the paper. We argue this detailed analysis is a contribution of the paper, as we hope will allow future research to focus on the outlined shortcomings that we identified. **Domain-agnostic embedding model:** It is true that we primarily compared the domain-agnostic embedding model on the grid world environments. However, across these environments, we observe similar performance trends for the domain-agnostic embedding model as for the domain-specific one. Considering that the domain-agnostic embedding model is pre-trained purely on text, we believe this is remarkable. Furthermore, the domain-agnostic embedding model reduces training cost, as the pre-trained LM is readily available. We did not conduct the comparison on more complex environments with the domain-agnostic embedding model in our initial version, because we observed that neither RA-DT nor any other baseline exhibits ICL behavior. Furthermore, the results by [1] strongly indicated that domain-agnostic LMs indeed transfer to more complex environments (such as Procgen). However, to validate this, we run domain-agnostic RA-DT on DMControl. We find that similar performance trends hold (see Figures 27 and 28). For our final version, we aim to add the domain-agnostic RA-DT curves for Meta-World and Procgen. **Discussion on Melo, ICML22:** Thank you for bringing up the paper by Melo [2]. We agree with the reviewer that this paper, as well as the follow-up [3], are important related works, and we added them to our discussion (Section 5) in our manuscript. However, there are considerable differences between our work (and related ICRL methods) and theirs. The most severe difference is that we consider an offline setting, in which the policy is learned from pre-collected trajectories via a supervised learning objective. In contrast, their work learns the policy via online interaction and proper RL objectives, which is more similar to online ICRL methods like Amago [4]. However, using RL objectives in the offline setting is not straightforward and requires to incorporate constraints (e.g., pessimism [5]). Consequently, we believe that an additional comparison against existing online methods is out-of-scope for this work. Nevertheless, we agree that online interaction and RL objectives (as used by [2,3]) can both be important components for ICL to emerge, which is a point we also raised in Section 5.

Thank you for your constructive feedback. We appreciate your positive assessment and are glad that you consider our method effective and the datasets we release a valuable resource for future research. **Differences to prior works:** As the reviewer correctly identified, our work does build on other works [1,2,3] to enable retrieval for in-context RL. However, there are a number of fundamental differences to these works: - [1] introduces the FrozenHopfield mechanism and uses it to leverage pre-trained LMs as a history compression mechanism for online RL. In contrast, we consider the offline in-context RL setting and propose to use the FrozenHopfield mechanism for retrieval-augmentation. While in NLP pre-trained models are readily available to embed sequences, pre-trained models are rarely available in RL. Therefore, we leverage the FrozenHopfield mechanism to map RL tokens to the language space. This allows us to repurpose frozen LMs as embedding models for RL trajectories, and enables retrieval without domain-specific pre-training. Furthermore, our focus is on offline in-context RL, whereas [1] focuses on online RL. - [2] leverages retrieval to enhance the generation abilities of multi-modal models for images and text. In contrast, we leverage retrieval-augmentation as an external memory mechanism to enhance the ICL abilities of RL agents. In addition to techniques from [2], we leverage deduplication of the retrieval index to regularize the retrieval process. - [3] conducts a large-scale experiment of simulated agents represented by language models that interact with one another. Furthermore, all inputs/outputs are defined via language. They leverage a retrieval mechanism to search for relevant experiences, which they reweight according to relevance and utility. Reweighting in the context of RL requires re-defining quantities such as utility due to the differences between the domains. Therefore, we use different reweighting terms to provide useful context trajectories at training time and to produce high reward behavior at inference time. For this, we come up with measures for utility that are adequate for the RL setting which we study in our ablation section (see Section 4.6, Appendix E.2) Note that the domains of these methods differ strongly from ours. Consequently, all the components we introduce in our setting require careful considerations that we investigate in our ablation studies (see Section 4.6, Appendix E). **Computational Overhead:** The reviewer is correct in stating that the process of searching for and reweighting experiences introduces additional computational overhead compared to methods that do not use retrieval. In contrast to other methods, however, RA-DT does not require keeping entire episodes in the context, which decouples the computational cost from the episode length. In our updated manuscript, we report the training times for grid worlds with varying grid sizes to highlight this effect. We find that RA-DT is up to almost 7 times faster at training time (for the largest grid, see Appendix E.7). At inference-time, RA-DT is slightly slower compared to baselines when retrieving on every step, but exhibits similar inference speeds when retrieving less frequently (see Appendix E.8). In Appendix E.6, we validate that less frequent retrieval does not severely affect performance. Importantly, the retrieval mechanism in RA-DT enables access to the entirety of the experiences collected across all ICL trials with negligible additional cost. The ability of RA-DT to access a broader range of experiences may be a reason for its enhanced down-stream performance. Our manuscript improved considerably by addressing and incorporating your comments: thank you. [1] “History compression via language models in reinforcement learning”, ICML 2022 \ [2] “Retrieval-augmented multimodal language modeling”, ICML 2023 \ [3] “Generative agents: Interactive simulacra of human behavior”, ArXiv 2023