검색 보강 결정 트랜스포머: 맥락 내 외부 메모리를 위한 강화 학습

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

October 9, 2024

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

cs.AI

초록

컨텍스트 내 학습 (In-context learning, ICL)은 모델이 그 컨텍스트 내에서 몇 가지 예시를 관찰함으로써 새로운 작업을 학습하는 능력입니다. 자연어 처리(NLP)에서 흔히 볼 수 있지만, 이 능력은 최근 강화 학습(Reinforcement Learning, RL) 환경에서도 관찰되었습니다. 그러나 이전의 컨텍스트 내 강화 학습 방법은 에이전트의 컨텍스트에서 전체 에피소드를 요구했습니다. 복잡한 환경은 일반적으로 희소 보상을 동반하는 긴 에피소드로 이어지기 때문에, 이러한 방법은 짧은 에피소드를 갖는 간단한 환경으로 제한되었습니다. 이러한 도전에 대처하기 위해, 우리는 검색 보강 결정 트랜스포머(Retrieval-Augmented Decision Transformer, RA-DT)를 소개합니다. RA-DT는 외부 메모리 메커니즘을 사용하여 과거 경험을 저장하고, 현재 상황에 관련 있는 하위 경로만 검색합니다. RA-DT의 검색 구성 요소는 훈련이 필요 없으며 완전히 도메인에 중립적일 수 있습니다. 우리는 RA-DT의 성능을 그리드 월드 환경, 로봇 시뮬레이션, 그리고 절차적으로 생성된 비디오 게임에서 평가합니다. 그리드 월드에서 RA-DT는 기준선을 능가하며, 그들의 컨텍스트 길이의 일부만 사용합니다. 더 나아가, 우리는 복잡한 환경에서 현재의 컨텍스트 내 강화 학습 방법의 한계를 밝히고, 미래 방향에 대해 논의합니다. 미래 연구를 촉진하기 위해, 우리는 고려된 환경 중 네 가지에 대한 데이터셋을 공개합니다.

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.