Retrieval-Augmented Decision Transformer: External Memory for In-context RL (2410.07071v2)

Abstract: In-context learning (ICL) is the ability of a model to learn a new task by observing a few exemplars in its context. While prevalent in NLP, this capability has recently also been observed in Reinforcement Learning (RL) settings. Prior in-context RL methods, however, require entire episodes in the agent's context. Given that complex environments typically lead to long episodes with sparse rewards, these methods are constrained to simple environments with short episodes. To address these challenges, we introduce Retrieval-Augmented Decision Transformer (RA-DT). RA-DT employs an external memory mechanism to store past experiences from which it retrieves only sub-trajectories relevant for the current situation. The retrieval component in RA-DT does not require training and can be entirely domain-agnostic. We evaluate the capabilities of RA-DT on grid-world environments, robotics simulations, and procedurally-generated video games. On grid-worlds, RA-DT outperforms baselines, while using only a fraction of their context length. Furthermore, we illuminate the limitations of current in-context RL methods on complex environments and discuss future directions. To facilitate future research, we release datasets for four of the considered environments.

• The paper presents the integration of external memory into the Decision Transformer, enabling the retrieval of relevant sub-trajectories for efficient RL decision-making.
• The methodology employs a domain-agnostic vector index and dual relevance-utility reweighting to enhance performance across grid-world, robotics, and video game environments.
• Experimental results show that RA-DT outperforms baselines with shorter context lengths, paving the way for improved long-horizon planning in sparse reward settings.

Overview of Retrieval-Augmented Decision Transformer in Reinforcement Learning

The paper presents the Retrieval-Augmented Decision Transformer (RA-DT), a novel method designed to enhance In-context Learning (ICL) in Reinforcement Learning (RL). Reinforcement Learning traditionally encounters challenges in complex environments, characterized by long episodes and sparse rewards. The RA-DT aims to address these using an innovative approach that incorporates an external memory mechanism, allowing RL models to efficiently retrieve past experiences and leverage sub-trajectories relevant to current tasks.

Key Contributions

RA-DT introduces an external memory component to the Decision Transformer framework, enabling the retrieval of informative sub-trajectories from past interactions. This is achieved through a vector index populated with embedded memories of past sub-trajectories. RA-DT adopts retrieval techniques akin to Retrieval-Augmented Generation in NLP to enhance the in-context decision-making capabilities of RL agents.

Evaluation and Results

The evaluation showcases RA-DT's superior performance on grid-world environments, robotics simulations, and video games. Notably, on grid-world tasks, RA-DT achieves higher performance than baseline methods while utilizing shorter context lengths. The method leverages a domain-agnostic model for trajectory embedding, demonstrating performance comparable to domain-specific models.

Technical Insights

• External Memory and Retrieval: RA-DT builds a vector index using an embedding model, enabling the retrieval of sub-trajectories similar to the current query. The retrieval process employs maximum inner product search to select relevant experiences efficiently.
• Relevance and Utility Reweighting: The paper highlights a dual-dimensional approach for assessing the retrieved sub-trajectories based on relevance and utility, refining the retrieval results for improved performance on RL tasks.
• Domain-Agnostic Retrieval Model: Experimentation with embedding techniques, such as FrozenHopfield with pre-trained LMs like BERT, allows RA-DT to perform effectively without the need for domain-specific pre-training.

Implications and Future Directions

The integration of retrieval mechanisms in RL presents a promising advancement for dealing with tasks requiring long-horizon planning and sparse feedback environments. The research implies that retrieval-based models in RL can significantly shrink computational requirements while maintaining or enhancing performance.

Looking forward, the paper acknowledges the potential benefits of combining RA-DT with learned interactive simulations or expert demonstrations to facilitate complex task adaptation. The domain-agnostic retrieval approach could be further refined for broader application in diverse RL environments, contributing to more generalized autonomous systems.

Conclusion

The RA-DT represents a substantial step forward in RL by facilitating more efficient and effective learning in context. While demonstrating robust performance in grid-based challenges and video game environments, the framework opens pathways for more seamless adaptation in real-world applications. The method highlights a promising direction for leveraging retrieval mechanisms, offering insights into the future of scalable and adaptable RL systems.