# Towards mental time travel: a hierarchical memory for reinforcement learning agents

arXiv: https://arxiv.org/abs/2105.14039 · HTML: https://mlsanity.com/paper/2105.14039/

## Key facts

- Cited by 40 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 64 references that are in the map.
- First posted on arXiv 28 May 2021.

## Builds on (most cited first, 30 of 37)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Perceiver: General Perception with Iterative Attention](https://mlsanity.com/paper/2103.03206/) (2021)
- [Effective Approaches to Attention-based Neural Machine Translation](https://mlsanity.com/paper/1508.04025/) (2015)
- [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (2018)
- [Generalization through Memorization: Nearest Neighbor Language Models](https://mlsanity.com/paper/1911.00172/) (2019)
- [Compressive Transformers for Long-Range Sequence Modelling](https://mlsanity.com/paper/1911.05507/) (2019)
- [Long Range Arena: A Benchmark for Efficient Transformers](https://mlsanity.com/paper/2011.04006/) (2020)
- [Self-supervised Learning: Generative or Contrastive](https://mlsanity.com/paper/2006.08218/) (2020)
- [The NarrativeQA Reading Comprehension Challenge](https://mlsanity.com/paper/1712.07040/) (2017)
- [Reinforcement Learning with Unsupervised Auxiliary Tasks](https://mlsanity.com/paper/1611.05397/) (2016)
- [Stabilizing Transformers for Reinforcement Learning](https://mlsanity.com/paper/1910.06764/) (2019)
- [Online Meta-Learning](https://mlsanity.com/paper/1902.08438/) (2019)
- [Hierarchical Transformers for Multi-Document Summarization](https://mlsanity.com/paper/1905.13164/) (2019)
- [HIBERT: Document Level Pre-training of Hierarchical Bidirectional Transformers for Document Summarization](https://mlsanity.com/paper/1905.06566/) (2019)
- [Adaptive Semiparametric Language Models](https://mlsanity.com/paper/2102.02557/) (2021)
- [Out of Order: How Important Is The Sequential Order of Words in a Sentence in Natural Language Understanding Tasks?](https://mlsanity.com/paper/2012.15180/) (2020)
- [Differentiable Patch Selection for Image Recognition](https://mlsanity.com/paper/2104.03059/) (2021)
- [Optimizing Agent Behavior over Long Time Scales by Transporting Value](https://mlsanity.com/paper/1810.06721/) (2018)
- [Unsupervised Predictive Memory in a Goal-Directed Agent](https://mlsanity.com/paper/1803.10760/) (2018)
- [Grounded Language Learning Fast and Slow](https://mlsanity.com/paper/2009.01719/) (2020)
- [Task Agnostic Continual Learning via Meta Learning](https://mlsanity.com/paper/1906.05201/) (2019)
- [Environmental drivers of systematicity and generalization in a situated agent](https://mlsanity.com/paper/1910.00571/) (2019)
- [The StreetLearn Environment and Dataset](https://mlsanity.com/paper/1903.01292/) (2019)

## Built on by (most cited first, 30 of 40)

- [WebShop: Towards Scalable Real-World Web Interaction with Grounded Language Agents](https://mlsanity.com/paper/2207.01206/) (2022)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [When Do Transformers Shine in RL? Decoupling Memory from Credit Assignment](https://mlsanity.com/paper/2307.03864/) (2023)
- [Memory, Benchmark & Robots: A Benchmark for Solving Complex Tasks with Reinforcement Learning](https://mlsanity.com/paper/2502.10550/) (2025)
- [Facing Off World Model Backbones: RNNs, Transformers, and S4](https://mlsanity.com/paper/2307.02064/) (2023)
- [Large-Scale Retrieval for Reinforcement Learning](https://mlsanity.com/paper/2206.05314/) (2022)
- [Recurrent Action Transformer with Memory](https://mlsanity.com/paper/2306.09459/) (2023)
- [On the link between conscious function and general intelligence in humans and machines](https://mlsanity.com/paper/2204.05133/) (2022)
- [Large language models can segment narrative events similarly to humans](https://mlsanity.com/paper/2301.10297/) (2023)
- [Memory Gym: Towards Endless Tasks to Benchmark Memory Capabilities of Agents](https://mlsanity.com/paper/2309.17207/) (2023)
- [Latent learning: episodic memory complements parametric learning by enabling flexible reuse of experiences](https://mlsanity.com/paper/2509.16189/) (2025)
- [Zipfian environments for Reinforcement Learning](https://mlsanity.com/paper/2203.08222/) (2022)
- [MemoNav: Working Memory Model for Visual Navigation](https://mlsanity.com/paper/2402.19161/) (2024)
- [Episodic memory in AI agents poses risks that should be studied and mitigated](https://mlsanity.com/paper/2501.11739/) (2025)
- [Generalization, Mayhems and Limits in Recurrent Proximal Policy Optimization](https://mlsanity.com/paper/2205.11104/) (2022)
- [MrSteve: Instruction-Following Agents in Minecraft with What-Where-When Memory](https://mlsanity.com/paper/2411.06736/) (2024)
- [Planning Irregular Object Packing via Hierarchical Reinforcement Learning](https://mlsanity.com/paper/2211.09382/) (2022)
- [Graph schemas as abstractions for transfer learning, inference, and planning](https://mlsanity.com/paper/2302.07350/) (2023)
- [Knowledge Retention for Continual Model-Based Reinforcement Learning](https://mlsanity.com/paper/2503.04256/) (2025)
- [Temporal-Spatial Causal Interpretations for Vision-Based Reinforcement Learning](https://mlsanity.com/paper/2112.03020/) (2021)
- [ELMUR: External Layer Memory with Update/Rewrite for Long-Horizon RL Problems](https://mlsanity.com/paper/2510.07151/) (2025)
- [Spatially-Aware Transformer for Embodied Agents](https://mlsanity.com/paper/2402.15160/) (2024)
- [Unraveling the Complexity of Memory in RL Agents: an Approach for Classification and Evaluation](https://mlsanity.com/paper/2412.06531/) (2024)
- [Beyond Sliding Windows: Learning to Manage Memory in Non-Markovian Environments](https://mlsanity.com/paper/2512.19154/) (2025)
- [EMO: Episodic Memory Optimization for Few-Shot Meta-Learning](https://mlsanity.com/paper/2306.05189/) (2023)
- [Grid Cell-Inspired Fragmentation and Recall for Efficient Map Building](https://mlsanity.com/paper/2307.05793/) (2023)
- [Neuro-Inspired Fragmentation and Recall to Overcome Catastrophic Forgetting in Curiosity](https://mlsanity.com/paper/2310.17537/) (2023)
- [Off-Beat Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2205.13718/) (2022)
- [MemoNav: Selecting Informative Memories for Visual Navigation](https://mlsanity.com/paper/2208.09610/) (2022)
- [Contrastive Distillation Is a Sample-Efficient Self-Supervised Loss Policy for Transfer Learning](https://mlsanity.com/paper/2212.11353/) (2022)

## Abstract (arXiv)

Reinforcement learning agents often forget details of the past, especially after delays or distractor tasks. Agents with common memory architectures struggle to recall and integrate across multiple timesteps of a past event, or even to recall the details of a single timestep that is followed by distractor tasks. To address these limitations, we propose a Hierarchical Chunk Attention Memory (HCAM), which helps agents to remember the past in detail. HCAM stores memories by dividing the past into chunks, and recalls by first performing high-level attention over coarse summaries of the chunks, and then performing detailed attention within only the most relevant chunks. An agent with HCAM can therefore "mentally time-travel" -- remember past events in detail without attending to all intervening events. We show that agents with HCAM substantially outperform agents with other memory architectures at tasks requiring long-term recall, retention, or reasoning over memory. These include recalling where an object is hidden in a 3D environment, rapidly learning to navigate efficiently in a new neighborhood, and rapidly learning and retaining new object names. Agents with HCAM can extrapolate to task sequences much longer than they were trained on, and can even generalize zero-shot from a meta-learning setting to maintaining knowledge across episodes. HCAM improves agent sample efficiency, generalization, and generality (by solving tasks that previously required specialized architectures). Our work is a step towards agents that can learn, interact, and adapt in complex and temporally-extended environments.