Research map / Reinforcement learning / Visual and representation RL
Memo: Training Memory-Efficient Embodied Agents with Reinforcement Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Memo, a transformer training recipe that interleaves learned summary tokens every segment so RL policies keep a compact memory, and evaluates it on ExtObjNav multi-object navigation and Dark-Key-To-Door meta-RL.
Problem
Transformer policies for long-horizon, memory-intensive embodied RL tasks face quadratic attention cost, large KV caches, and poor extrapolation to long contexts, while recurrent models have fixed-size memory and struggle with long-horizon gradient propagation.
Approach
Periodically generate learnable summary tokens at the end of each segment, feed them back as accumulated memory, and train end-to-end through the RL objective with gradients flowing across all summaries, so the model learns task-driven compression while attending over ~8x fewer tokens than a full-context transformer.
Key results
- On ExtObjNav, Memo beats the full-context transformer by 7.5% SR and 2.5% SPL on average while using 8x fewer tokens in context.
- Memo uses 51.8 MB GPU memory vs 546.5 MB for the full-context transformer, with 4.2x fewer FLOPs and about 2x lower latency at 32k eval steps.
- Streaming Memo keeps performance over a 32k-step trial with a fixed 1024-token window, while the streaming full-context transformer drops sharply after 6k steps.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace full-context attention with periodic summary tokens and accumulated memory | Memo versus Full-Context Transformer · ExtObjNav validation | success rate: 7.5% higher SR | Helped |
| Replace full-context attention with periodic summary tokens and accumulated memory | Memo versus Full-Context Transformer · ExtObjNav validation | SPL: 2.5% higher SPL | Helped |
| Use accumulated summaries instead of fixed-size recurrent memory | Memo versus RMT · ExtObjNav | success rate: about 5% | Helped |
| Allow gradients to propagate across all segments instead of only the two preceding segments (ablation) | AC · ExtObjNav | success rate | Helped |
| Allow gradients to propagate across all segments instead of only the two preceding segments (ablation) | AC · ExtObjNav | SPL | Helped |
| Accumulate summaries across segments instead of passing forward only newly generated summaries (ablation) | Memo · Dark-Key-To-Door | training speed: over 10M steps | Helped |
| Use RMT-128 instead of Memo (ablation) | RMT · ExtObjNav | training stability: RMT-128 | Hurt |
| Use 64 summary tokens instead of 32 (ablation) | Memo · ExtObjNav | generalization: 16, 32, and 64 | Hurt |
| Use 32 summary tokens instead of 16 (ablation) | Memo · ExtObjNav | generalization: 16/32/64 | Helped |
| Disable segment-length randomization and use a fixed segment length of 256 (ablation) | Memo · ExtObjNav | data efficiency: 256 | Hurt |
| Remove the summary re-encoding step and construct the KV cache directly from summary embeddings (ablation) | Memo · ExtObjNav | performance beyond training context length | Hurt |
| Retain only the most recent 6k KV cache elements | Full-Context Transformer · ExtObjNav streaming evaluation | success rate: 6k | Hurt |
Surprising findings
- Re-encoding summary tokens through the transformer matters: directly using per-layer summary embeddings as the KV cache performs poorly and degrades beyond training context length.
- Fine-tuning Autocompressor-style models with truncated BPTT (TBTT) performs much worse than propagating gradients across all segments, and AC (all segments) degrades over time to the TBTT level.
Limitations
- Evaluated only on object navigation in unseen scenes with a fixed object set; no semantic generalization to new object categories.
- Does not explore memory consolidation where past summaries are progressively compressed.
- Length extrapolation beyond training context limits remains an open challenge; Memo (16k) degrades beyond ~1.5x its training context.
Experimental setup
Datasets: HSSD YCB ExtObjNav training placements (self-created)
Benchmarks: ExtObjNav Dark-Key-To-Door Memory Maze 9x9 T-Maze
Metrics: success rate SPL average return GPU memory FLOPs latency
Baselines: Full-context transformer (FCT) Transformer without inter-episode attention Recurrent Memory Transformer (RMT) Autocompressors (AC) RL^2 Streaming Transformer StreamingLLM
Builds on: Recurrent Memory Transformer Autocompressors ReLIC AMAGO DD-PPO
Closest papers by meaning
- Spatially-Aware Transformers for Embodied Agents (ICLR 2024 · similarity 0.89)
- Embodied Agents Meet Personalization: Investigating Challenges and Solutions Through the Lens of Memory Utilization (ICLR 2026 · similarity 0.89)
- Benchmarking Agent Memory in Interdependent Multi-Session Agentic Tasks (ICML 2026 · similarity 0.88)
- 3D-Mem: 3D Scene Memory for Embodied Exploration and Reasoning (CVPR 2025 · similarity 0.88)
- Stable Hadamard Memory: Revitalizing Memory-Augmented Agents for Reinforcement Learning (ICLR 2025 · similarity 0.88)
- Reasoning in Visual Navigation of End-to-end Trained Agents: A Dynamical Systems Approach (CVPR 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.