Memo: Training Memory-Efficient Embodied Agents with Reinforcement Learning

AI summary

Overview

Research area: Reinforcement learning for embodied agents, specifically long-horizon memory in transformer-based policies.

Technical level: Advanced. The paper assumes familiarity with transformers, attention/KV caching, POMDPs, and on-policy vs. off-policy RL algorithms.

Scope: A transformer architecture and training recipe that compresses long observation histories into accumulated summary tokens so that embodied agents can reason over long horizons with far less memory.

What This Paper Is About

Transformer policies for embodied tasks must condition on long histories to act well, but raw visual observations quickly exceed context limits and make the key-value (KV) cache expensive to store at inference. The authors propose Memo, which trains a transformer to periodically summarize past experience into a small set of summary tokens that are stored in a memory buffer and fed back into the model, instead of attending to the full history. The goal is to match or beat full-context transformers on long-horizon, memory-intensive tasks while using much less compute and storage.

Key Contributions

1. Memo, a context summarization framework for RL. Summary tokens are generated at fixed intervals (segments) and accumulated, so the model attends to condensed representations of prior states rather than a full-context cache. Summarization is trained end-to-end through the RL objective, with gradients propagating through all summaries.
2. Evaluation on sequential decision-making benchmarks. Memo is tested on a grid-world meta-RL benchmark (Dark-Key-To-Door) and a multi-object navigation task in photo-realistic indoor scenes (ExtObjNav), and shown to outperform transformers that require full-context storage.
3. A general memory augmentation applicable to both on-policy and off-policy RL. Memo is integrated with ReLIC (an adaptation of DD-PPO) for on-policy RL and with AMAGO for off-policy RL.
4. Analysis of design choices. Ablations cover summary length (16, 32, 64), segment length randomization, gradient propagation depth (comparison to Autocompressors with truncated backpropagation), streaming inference with a truncated KV cache, and fine-tuning on longer contexts.

Main Findings

• Summarization beats full in-context access at lower cost. On ExtObjNav, Memo outperforms the full-context transformer (FCT) baseline by 7.5% higher success rate (SR) and 2.5% higher SPL on average, while using 8x fewer tokens. The full-context baseline requires an 8-10x larger cache. Memo was trained with segment length 256, 32 summarization tokens, and a trainable rollout length of 4096.
• In-context learning continues past the training context. Both Memo and FCT improve on ExtObjNav up to about 10k steps, roughly 2.5x the training context length, after which both degrade. FCT degrades more than Memo in SR, while both degrade similarly in SPL. Omitting inter-episode attention (no-IEA) performs significantly worse.
• Memory accumulation beats fixed recurrent memory. Comparing against a Recurrent Memory Transformer (RMT) variant, Memo converges faster, improving training speed by over 10M steps on Dark-Key-To-Door. On ExtObjNav, RMT with 1x and 2x Memo's summary size (32, 64) still lags Memo by about 5% success rate, and larger summaries (RMT-128) destabilize training.
• On Dark-Key-To-Door, Memo and RMT match FCT. Trained with AMAGO, both reach peak validation performance by 40M training steps. FCT shows a notable performance drop around the 35-40M step mark across all seeds, while Memo converges stably.
• Long-horizon gradient propagation matters. Under an Autocompressors-style setup, the truncated-backpropagation variant AC (TBTT) performs much worse. AC (all segments) slightly surpasses Memo in SR and matches it in SPL during the first 6-8k steps, but degrades and converges to AC (TBTT) after 16k steps.
• Streaming inference is robust for Memo. With the KV cache truncated (streaming starting after 6k eval steps; T = 6k for the streaming transformer and T = 1024 for streaming Memo), Streaming Memo maintains and slightly improves its in-context learning trend, whereas the Streaming Transformer suffers a sharp decline.
• Fine-tuning on longer contexts helps both methods, but Memo more. Checkpoints trained to 1B steps and fine-tuned on a 4x larger context of 16,384 tokens for 500M steps improve. Memo (16k) outperforms Memo (4k) in SPL, and beats FCT (16k) in maximum SR and SPL. Memo (16k) generalizes up to 1.5x its training context length before slight degradation, while FCT (16k) maintains performance up to at least 2x.
• Summary length is a sensitive hyperparameter. Varying summary length between 16, 32, and 64 (compression ratios of 16x, 8x, and 4x) shows 32 is best for generalization over longer trajectories, 16 next, and 64 worst. The main gap appears after 6k environment steps.
• Randomizing segment lengths is important. Training with segment lengths sampled between [0.8, 1.2] x 256, i.e. a range of [205, 307], performs significantly better than a fixed segment length of 256, both in training and evaluation.
• Performance saturates around 60% SR on ExtObjNav, which the authors investigate through the fine-tuning experiment.

Methodology in Plain English

The authors treat each task as a partially observable Markov decision process made fully observable by feeding history to a sequence encoder. The policy and value heads read from the encoder's hidden state.

The core mechanism splits a long input sequence into segments of length l_seg and, at the end of each segment, has the transformer emit l_sum learnable summary tokens. From then on, the model attends only to accumulated summary tokens plus the observations within the current segment. Attention masking creates a deliberate bottleneck: older raw observations are excluded, so all long-range information must flow through the summaries. Positional indices are assigned relative to position within a segment, and the summary tokens at the end of a segment follow the indices of the last observation in that segment. Segment lengths are randomized during training within ±20% of a fixed length, while a fixed length is used for data collection and evaluation.

Because on-policy RL updates change the model weights, cached key-value entries become stale; the authors recompute both the KV cache and the summary vectors after each policy update, following the approach used by ReLIC.

Experiments alternate two RL backbones. On the ExtObjNav navigation task, trained with ReLIC, the policy consumes a 256x256 RGB image (encoded by a frozen, fine-tuned ViT-B-sized VC-1 encoder with an MLP reducing to 256 dimensions), a 32-dimensional position embedding, the previous action embedding, and the goal category one-hot vector. On Dark-Key-To-Door, the off-policy AMAGO implementation and its training details are kept fixed.

Why This Matters

Impact on research. The paper argues that task-driven memory formation can be learned end-to-end inside the RL objective, rather than applied as a post-hoc compression or supervision. It shows that accumulating summaries (with gradients flowing through all of them) is better than fixed-size recurrent memory or truncated backpropagation, and that full-context transformers, while strong, are both more expensive and less stable on long-horizon RL.

Real-world applications:

• Home and service robots that must remember where objects were seen during earlier navigation to reach new goals in the same building.
• Warehouse or logistics robots that revisit the same space repeatedly under changing goals.
• Assistive or companion agents operating over long sessions where only a small on-device memory budget is available.
• Any deployment where inference hardware limits the KV cache, since streaming with truncated memory is the operating regime Memo targets.

Industry relevance. The efficiency argument is practical: an 8-10x smaller cache directly reduces inference memory and enables deployment on constrained hardware, and the streaming results indicate acceptable behavior when the cache must be bounded.

Future Directions

• Investigate the factors limiting Memo's long-context generalization; Memo (16k) generalizes to about 1.5x its training context length while FCT (16k) reaches at least 2x.
• Explore more flexible memory mechanisms such as memory consolidation, where past summaries are progressively compressed.
• Reduce reliance on the RL objective alone by adding self-supervised objectives such as future prediction for training memory representations.
• Evaluate semantic generalization to entirely new object categories, which the authors deliberately avoided because it would conflate memorization with object recognition capability.

Target Audience

Researchers and practitioners working on long-horizon reinforcement learning, transformer-based policies, and embodied AI who care about the trade-off between memory footprint and long-term reasoning. It is most useful to readers already comfortable with attention mechanisms, KV caching, and on-policy versus off-policy RL training loops, since the paper leans heavily on those concepts. Readers primarily interested in navigation benchmarks may also find the ExtObjNav setup and its evaluation protocol informative.

Authors’ 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.