Thread by @AndrewLampinen

How can RL agents recall the past in detail, in order to behave appropriately in the present? In our new preprint "Towards mental time travel: A hierarchical memory for RL agents" (arxiv.org/abs/2105.14039) we propose a memory architecture that steps in this direction.

We draw inspiration from the idea that human memory is like "mental time travel"—we can recall a specific event in the past, and relive it in some sequential detail, with relatively little interference from other events. This ability is key to our goal-directed use of memory.

By contrast, RL agent memories generally lack either this sparsity (attending to one or a few events in the past) or this detail (replaying an event in sequence, rather than just recalling a single vector from each event). This limits their ability to learn from their memory.

We therefore propose a Hierarchical Transformer Memory (HTM). HTM stores chunks of the past, together with a summary for each chunk. It recalls hierarchically: first attending over summaries to identify relevance, then time-traveling to each relevant chunk to attend in detail.

We show that HTM allows RL agents to excel across 6 varied memory domains, such as recalling a ballet, maintaining object permanence, rapidly learning to navigate efficiently in a new neighborhood, and rapidly learning new nouns and recalling them after distractor tasks.

Agents with HTM can extrapolate to recalling knowledge after an order of magnitude(!) more distractor phases than they were trained on, and can even generalize zero-shot from training within single episodes to evaluation of knowledge from several episodes before.

They also achieve near-optimal performance on One-Shot StreetLearn (openreview.net/forum?id=F-mvp…), a challenging domain which requires recombining memories of prior paths in order to plan new routes. HTM matches the performance of a memory specifically designed for this setting.

We think these results highlight the value of hierarchy, sparsity, and attention in RL agent memories. Combining these allows agents to effectively "mental time travel" into relevant memories, which allows them to better learn from the past to achieve their present goals.

We also think our work highlights some generally interesting ideas, including 1) hierarchical attention where the top level determines what deserves detailed attention and 2) key-value memories where the values are not just vectors, but more general structures (e.g. sequences).

Interested to hear your thoughts! Thanks to my awesome coauthors @scychan_brains, Andrea Banino, and @FelixHill84, as well as all our colleagues who provided amazing suggestions and support while working on this paper!