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

1 Introduction . Human learning and generalization relies on our detailed memory of the past [ 55 , 52 , 16 , 30 , 43 ] . If we watch a show , we can generally recall its scenes in some detail afterward . If we explore a new neighborhood , we can recall our paths through it in order to plan new routes . If we are exposed to a new noun , we can find that object by name later . We experience relatively little interference from delays or intervening events . Indeed , human episodic memory has been compared to “ mental time travel ” [ 55 , 52 , 37 ] —we are able to transport ourselves into a past event and re-live it in sequential detail , without attending to everything that has happened since . That is , our recall is both sparse ( we attend to a small chunk of the past , or a few chunks ) and detailed ( we recover a large amount of information 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . from each chunk ) . This combination gives humans the flexibility necessary to recover information from memory that is specifically relevant to the task at hand , even after long periods of time . If we want Reinforcement Learning ( RL ) agents to meaningfully interact with complex environments over time , our agents will need to achieve this type of memory . They will need to remember event details . They will need to retain information they have acquired , despite unrelated intervening tasks . They will need to rapidly learn many pieces of new knowledge without discarding previous ones . They will need to reason over their memories to generalize beyond their training experience . However , current models struggle to achieve rapid encoding combined with detailed recall , even over relatively short timescales . Meta-learning approaches [ 57 ] can slowly learn global task knowledge in order to rapidly learn new information , but they generally discard the details of that new information immediately to solve the next task . While LSTMs [ 20 ] , Transformers [ 56 ] , and variants like the TransformerXL [ 9 ] can serve as RL agent memories in short tasks [ 44 ] , we show that they are ineffective at detailed recall even after a few minutes . Transformer attention can be ineffective at long ranges even in supervised tasks [ 54 ] , and this problem is likely exacerbated by the sparse learning signals in RL . By contrast , prior episodic memory architectures [ 21 , 14 ] can maintain information over slightly longer timescales , but they struggle with tasks that require recalling a temporallystructured event , rather than simply recalling a single past state . That is , they are particularly poor at the type of event recall that is fundamental to human memory . We suggest that this is because prior methods lack 1 ) memory chunks longer than a single timestep and 2 ) sparsity of attention . This means that these models can not effectively “ time-travel ” to an event , and relive that specific event in detail , without interference from all other events . This also limits their ability to reason over their memories to adapt to new tasks—for example , planning a new path by combining pieces of previous ones [ 49 ] . Here , we propose a new type of memory that begins to address these challenges , by leveraging sparsity , chunking , and hierarchical attention . We refer to our architecture as a Hierarchical Chunk Attention Memory ( HCAM ) . The fundamental insight of HCAM is to divide the past into distinct chunks before storing it in memory , and to recall hierarchically . To recall , HCAM first uses coarse chunk summaries to identify relevant chunks , and then mentally travels back to attend to each relevant chunk in further detail . This approach combines benefits of the sparse and relatively longterm recall ability of prior episodic memories [ 60 , 49 ] with the short-term sequential and reasoning power of transformers [ 56 , 9 , 44 ] in order to achieve better recall and reasoning over memory than either prior approach . While we do not address the problem of optimally chunking the past here , instead employing arbitrary chunks of fixed length , our results show that even fixed chunking can substantially improve memory . ( Relatedly , Ke et al . [ 25 ] have shown that sparse , top-k retrieval of memory chunks can improve credit assignment in LSTMs with attention . ) We show that HCAM allows RL agents to recall events over longer intervals and with greater sample efficiency than prior memories . Agents with HCAM can learn new words and then maintain them over distractor phases . They can extrapolate far beyond the training distribution , to maintain and recall memories after 5× more distractors than during training , or after multiple episodes when trained only on single ones . Agents with HCAM can reason over their memories to plan paths near-optimally as they explore a new neighborhood , comparable to an agent architecture designed specifically for that task . HCAM is robust to hyperparameters , and agents with HCAM consistently outperform agents with Gated TransformerXL memories [ 44 ] that are twice as wide or twice as deep . Furthermore , HCAM is more robust to varying hyperparameters than other architectures ( App . D.11 ) . Thus HCAM provides a substantial improvement in the memory capabilities of RL agents . 1.1 Background . Memory in RL Learning signals are sparse in RL , which can make it challenging to learn tasks requiring long memory . Thus , a variety of sophisticated memory architectures have been proposed for RL agents . Most memories are based on LSTMs [ 20 ] , often with additional key-value episodic memories [ 60 , 14 ] . Often , RL agents require self-supervised auxiliary training because sparse task rewards do not provide enough signal for the agent to learn what to store in memory [ 60 , 14 , 19 ] . Transformers Transformers [ 56 ] are sequence models that use attention rather than recurrence . These models—and their variants [ 9 , 28 , 58 , 46 ] —have come to dominate natural language processing . This success of transformers over LSTMs on sequential language tasks suggests that they might be effective agent memories , and a recent work succesfully used transformers as RL agent memories [ 44 ] —specifically , a gated version of TransformerXL [ 9 ] . However , using TransformerXL memories on challenging memory tasks still requires auxiliary self-supervision [ 19 ] , and might even benefit from new unsupervised learning mechanisms [ 3 ] . Furthermore , even in supervised tasks Transformers can struggle to recall details from long sequences [ 54 ] . In the next section , we propose a new hierarchical attention memory architecture for RL agents that can help overcome these challenges . 2 Hierarchical Chunk Attention Memory . . . ( a ) Standard attention. . . . No attention ( b ) Hierarchical attention . Memory Input LayerNorm ·Attn.across chunks Keys ... Chunks ... Attn . within top-k chunks + Weighted sum of results OutputSkip connection ( c ) HCAM attention block . Encoders Image Language ResNet LSTM Memory ( 4 layers ) Mem . layer Attention HCAM Attn.Memory MLP ... Attention HCAM Attn.Memory MLP DecodersResN LSTM MLP MLP Image Lang . Self-supervised V -trace π V ( d ) HCAM in an agent . Figure 1 : HCAM motivation and implementation . ( a ) Standard transformer attention attends to every time step in a fixed window . ( b ) HCAM divides the past into chunks . It first attends to summaries of each chunk , and only attends in greater detail within relevant chunks . ( c ) HCAM is implemented using a key-value memory , where the keys are chunk summaries , and the values are chunk sequences . Toplevel attention is performed over all summaries ( left ) , and then standard attention is performed within the top-k chunks . The results are relevance-weighted and summed ( right ) , and then added to a residual input to produce the block output . ( d ) Our agent encodes inputs , processes them through a 4-layer memory with an HCAM block in each layer , and decodes outputs for self-supervised and RL losses . The HCAM uses a hierarchical strategy to attend over memory ( Fig . 1 ) . Specifically , it divides the past into chunks . It stores each chunk in detail—as a sequence—together with a single summary key for that chunk . It can then recall hierarchically , by first attending to chunk summaries to choose the top-k relevant chunks , and then attending within those most relevant chunks in further detail . In particular , let S denote the memory summary keys and Q be a linear projection ( “ query ” ) layer . Then the input is normalized by a LayerNorm , and the chunk relevance ( R ) is computed as : R = softmax ( Q ( normed input ) · S ) Then the model attends in detail within the chunks that have the top-k relevance scores . If the memory chunks are denoted by C and MHA ( query inputs , key/value inputs ) denotes multi-head attention : memory query results = ∑ i∈ top-k from R Ri ·MHA ( normed input , Ci ) That is , HCAM selects the most relevant memory chunks , then time travels back into each of those chunks to query it in further detail , and finally aggregates the results weighted by relevance . This result is added to the input to produce the output of the block . An HCAM block can be added to a transformer layer , to supplement standard attention . ( For an open-source HCAM implementation , see App . A.1 ) HCAM can attend to any point in its memory , but at each step it only attends to a few chunks , which is much more compute- and memory-efficient than full attention . With N chunks in memory , each of size C , using HCAM over the top-k chunks requires O ( N + kC ) attention computations , whereas full attention requires O ( NC ) . HCAM succeeds even with k = 1 in many cases ( App . D.5 ) , reducing computation 10× compared to TrXL . Even with k ≥ 1 , HCAM often runs faster ( App . D.10 ) . Agents Our agents ( Fig . 1d ) consist of input encoders ( a CNN or ResNet for visual input , and an LSTM for language ) , a 4-layer memory ( one of HCAM , a Gated TransformerXL [ 44 ] , or an LSTM [ 20 ] ) , followed by output MLPs for policy , value , and auxiliary losses . Our agents are trained using IMPALA [ 11 ] . To use HCAM in the memory , we adapt the Gated TransformerXL agent memory [ 44 ] . We replace the XL memory with an HCAM block between the local attention ( which provides short-term memory ) and feed-forward blocks of each layer . We find that gating each layer [ 44 ] is not beneficial to HCAM on our tasks ( App . D.7 ) , so we removed it . See App . A.2 for further agent details . We store the initial layer inputs in memory—that is , the inputs to the local attention—which performs better than storing the output of the local attention , and matches TransformerXL ( TrXL ) more closely . We accumulate these inputs into fixed-length memory chunks . When enough steps have accumulated to fill the current chunk , we add the current chunk to the memory . We average-pool across the chunk to generate a summary . We reset the memory between episodes , except where cross-episode memory is tested ( Section 3.3 ) . We stop gradients from flowing into the memory , just as TrXL [ 9 ] stops gradients flowing before the current window . This reduces challenges with the non-differentiability of top-k selection ( see App . D.5 for discussion ) . It also means that we only need to store past activations in memory , and not the preceding computations , so we can efficiently store a relatively long time period . However , the model must therefore learn to encode over shorter timescales , since it can not update the encoders through the memory . Thus , as in prior work referenced above , we show that self-supervised learning [ 33 ] , is necessary for successful learning on our tasks ( App . D.6 ) . Specifically , we use an auxiliary reconstruction loss—at each step we trained the agents ( of all memory types ) to reconstruct the input image and language as outputs [ 19 ] . This forces the agents to propagate these task-relevant features through memory , which ensures that they be encoded [ cf . 21 , 14 ] .

Summary: The authors propose an architecture to augment an RL agent’s memory using a novel hierarchical transformer. The main novelties of the approach are 1) that the transformer attends over chunks and 2) that it’s used for long horizon RL. They test their approach in many interesting environments that each test a different component of memory (delayed recall, object permanence, distractors…). The results show that across the different environments the proposed approach outperforms the tested baselines. Moreover, the approach competes well with specially tuned state-of-the-art methods in other environments.