Memo: Training Memory-Efficient Embodied Agents with Reinforcement Learning

TL;DR

The paper tackles the challenge of long-horizon decision making in embodied RL where full-context transformers become intractable. It introduces Memo, a memory-augmented transformer that learns to compress past experiences into summary tokens at segment boundaries, forming a compact memory buffer that can be attended to in future steps. Through end-to-end RL training, Memo demonstrates superior efficiency and robust extrapolation compared to full-context baselines, achieving better in-context learning and navigation performance on ExtObjNav and related tasks while using substantially less memory. The work offers a general, memory-efficient approach applicable to both on-policy and off-policy RL and shows promising robustness in streaming inference, with ablations highlighting the importance of gradient propagation through summaries and accumulation of memory tokens. Overall, Memo advances scalable long-horizon RL by coupling learned memory compression with transformer-based decision models, enabling practical deployment under inference budgets.

Abstract

To enable embodied agents to operate effectively over extended timeframes, it is crucial to develop models that form and access memories to stay contextualized in their environment. In the current paradigm of training transformer-based policies for embodied sequential decision-making tasks, visual inputs often overwhelm the context limits of transformers, while humans can maintain and utilize a lifetime of experience compressed as memories. Significant compression is possible in principle, as much of the input is irrelevant and can be abstracted. However, existing approaches predominantly focus on either recurrent models with fixed-size memory or transformers with full-context reliance. In this work, we propose Memo, a transformer-based architecture and training recipe for reinforcement learning (RL) on memory-intensive, long-horizon tasks. Memo incorporates the creation and retrieval of memory by interleaving periodic summarization tokens with the inputs of a model during training. We demonstrate Memo's effectiveness on a gridworld meta-RL benchmark and a multi-object navigation task in photo-realistic indoor settings. Memo outperforms naive long-context transformer baselines while being more compute and storage efficient. Additionally, Memo generalizes better to longer contexts at inference time and remains robust in streaming settings, where historical context must be truncated to fit inference constraints. Our code is available at: https://github.com/gunshi/memo.

Figures (11)

• Figure 1: Architecture Diagram of Memo. The frames $O_{1-3l_{seg}}$ depict the input sensor observations to the agent at timesteps $1-3l_{seg}$. The figure depicts the information flow between three consecutive segments of a much longer context of inputs provided to a transformer during training. The summary tokens bottleneck the information passed on from one chunk of inputs to the next.
• Figure 2: (Left) Overhead view of a training scene in Habitat simulator. The ExtObjNav task involves placing multiple objects around the house, and then sampling an object category as the goal each time the previous goal is reached. The figure shows some sample objects that may be placed around the house, such as Banana, Cracker Box, etc. (Right) Val success rate and SPL curves over 32k in-context learning steps in novel scenes, for different methods trained with ReLIC. We compare Memo to the the full-context transformer (FCT), the FCT variant which does not attend over previous episodes (no IEA), the recurrent memory transformer (RMT), and the Autocompressors (AC) variant.
• Figure 3: (Left) We plot the average return over evaluation trials of 500 steps, consisting of 10 or more episodes, plotted against the training progress. We compare Memo, RMT and FCT, when training with the AMAGO RL algorithm over 3 random seeds. (Right) On ExtObjNav, we find that restricting gradient propagation to a few consecutive windows (AC (TBTT)) performs much worse than allowing all past segments to remain trainable (AC (all segments)). While AC (all segments) initially performs slightly better than Memo, it suffers from severe performance degradation over time, converging to AC (TBTT) after 16k steps.
• Figure 4: (Left) We show a comparison between the accumulating context (default setting) and the streaming version of Memo and the full-context transformer, where streaming starts after 6k eval steps. This comparison highlights the robustness of Streaming Memo, which maintains and even slightly improves performance, whereas Streaming Transformer suffers a sharp decline. (Right) We see that finetuning over a longer number of steps during training serves to partially fix the degradation due to context length extrapolation.
• Figure 5: Memo ablations.(Left) Effect of varying the number of memory tokens (16/32/64), where 32 outperforms 16, and 16 outperforms 64. (Right) Performance of Memo on ExtObjNav with and without randomization of the segment lengths at the end of which summarization is done, showing the latter is significantly less data-efficient.