ReLIC: A Recipe for 64k Steps of In-Context Reinforcement Learning for Embodied AI

TL;DR

ReLIC introduces a novel in-context reinforcement learning approach enabling embodied agents to adapt rapidly to new environments using extensive experience, novel policy updates, and long observation histories, outperforming existing meta-RL methods.

Contribution

The paper presents ReLIC, a new method combining large-scale RL, partial updates, and Sink-KV mechanisms for effective in-context learning in embodied AI.

Findings

ReLIC outperforms meta-RL baselines in unseen house navigation tasks.

ReLIC can perform few-shot imitation learning without expert demonstrations.

The combination of large-scale RL, partial updates, and Sink-KV is crucial for success.

Abstract

Intelligent embodied agents need to quickly adapt to new scenarios by integrating long histories of experience into decision-making. For instance, a robot in an unfamiliar house initially wouldn't know the locations of objects needed for tasks and might perform inefficiently. However, as it gathers more experience, it should learn the layout of its environment and remember where objects are, allowing it to complete new tasks more efficiently. To enable such rapid adaptation to new tasks, we present ReLIC, a new approach for in-context reinforcement learning (RL) for embodied agents. With ReLIC, agents are capable of adapting to new environments using 64,000 steps of in-context experience with full attention while being trained through self-generated experience via RL. We achieve this by proposing a novel policy update scheme for on-policy RL called "partial updates'' as well as a Sink-KV mechanism that enables effective utilization of a long observation history for embodied agents. Our method outperforms a variety of meta-RL baselines in adapting to unseen houses in an embodied multi-object navigation task. In addition, we find that ReLIC is capable of few-shot imitation learning despite never being trained with expert demonstrations. We also provide a comprehensive analysis of ReLIC, highlighting that the combination of large-scale RL training, the proposed partial updates scheme, and the Sink-KV are essential for effective in-context learning. The code for ReLIC and all our experiments is at https://github.com/aielawady/relic