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

Abstract

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 Transformer Memory (HTM), which helps agents to remember the past in detail. HTM 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 HTM can therefore "mentally time-travel"-remember past events in detail without attending to all intervening events. We show that agents with HTM 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 HTM can extrapolate to task sequences an order of magnitude longer than they were trained on, and can even generalize zero-shot from a meta-learning setting to maintaining knowledge across episodes. HTM 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.Preprint. Under review.

On the link between conscious function and general intelligence in humans and machines

“…Whilst the majority of deep RL algorithms use experience to update the policy and/or value functions offline, some RL methods have taken direct inspiration from episodic memory to create policies that are directly conditioned on past experiences (Blundell et al, 2016;Pritzel et al, 2017). Sufficiently large buffers of experiences, which can be queried as needed during the learning of a novel task, provide powerful retrospective learning (Lampinen, Chan, Banino, & Hill, 2021), but limited future-oriented imagination.…”

“…Initial attempts to imitate hippocampal episodic control (Lengyel & Dayan, 2007) in AI agents used non-parametric (Blundell et al, 2016) and semi-parametric (Pritzel et al, 2017) memories (which store state information from every timestep), enabling rapid learning when compared to standard deep RL algorithms. Scaling such memory is an open challenge, with possible solutions including more sophisticated storage/forgetting mechanisms, compression (Agostinelli, Arulkumaran, Sarrico, Richemond, & Bharath, 2019), and hierarchy (Lampinen et al, 2021). Still, these algorithms only correspond to replay, while combining episodic memories with generative models could lead to further abilities, such as planning to find outcomes that are of specific relevance to the agent (Zakharov et al, 2021).…”