Trasformatore Decisionale Potenziato da Recupero: Memoria Esterna per RL Contestuale

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

October 9, 2024

Autori: Thomas Schmied, Fabian Paischer, Vihang Patil, Markus Hofmarcher, Razvan Pascanu, Sepp Hochreiter

cs.AI

Abstract

L'apprendimento in contesto (ICL) è la capacità di un modello di apprendere un nuovo compito osservando alcuni esemplari nel suo contesto. Sebbene diffuso nell'elaborazione del linguaggio naturale (NLP), questa capacità è stata recentemente osservata anche in ambienti di Apprendimento per Rinforzo (RL). Tuttavia, i metodi precedenti di RL in contesto richiedono interi episodi nel contesto dell'agente. Dato che ambienti complessi portano tipicamente a episodi lunghi con ricompense sparse, questi metodi sono limitati ad ambienti semplici con episodi brevi. Per affrontare queste sfide, presentiamo il Decision Transformer potenziato da Recupero (RA-DT). RA-DT utilizza un meccanismo di memoria esterna per memorizzare esperienze passate da cui recupera solo sotto-traiettorie rilevanti per la situazione attuale. Il componente di recupero in RA-DT non richiede addestramento e può essere completamente agnostico al dominio. Valutiamo le capacità di RA-DT su ambienti a griglia, simulazioni robotiche e videogiochi generati proceduralmente. Sulle griglie, RA-DT supera i modelli di base, utilizzando solo una frazione della loro lunghezza di contesto. Inoltre, mettiamo in luce i limiti dei metodi attuali di RL in contesto su ambienti complessi e discutiamo le direzioni future. Per agevolare la ricerca futura, rilasciamo set di dati per quattro degli ambienti considerati.

English

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.