Towards mental time travel: a hierarchical memory for reinforcement learning agents
Authors: Andrew Kyle Lampinen
Stephanie C.Y. Chan
Andrea Banino
Felix Hill
Venue: NeurIPS 2021 Poster (2021)
Keywords: Memory
Reinforcement Learning
Hierarchical
Attention
Episodic

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 Chunk Attention Memory (HCAM), that 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 words. 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 improve 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.

Towards mental time travel: a hierarchical memory for reinforcement learning agents
Andrew K. Lampinen DeepMind London, UK
Stephanie Y. Chan
Andrea Banino
DeepMind London, UK scychan@deepmind..comm
DeepMind London, UK
lampinenn@deepmidd.com
abanino@deeepindd.com
Felix Hill Deep.indd London, UK felixhhil@ddeepmnndd.o com
Abstract
Reinforcement learning agents often forget details of the past, especially after de- lays or distractor tasks. Agents with common memory architectures struggle to recall and integrate across multiple timesteps details of single timestep that is followed by distractor tasks. To address these limitations, we propose 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 mem- ory architectures at tasks requiring long-tent rnngtermm recall, retention, reasoning over memory. These include recalling where an object is hidden in a 3D environment, rapidly learning to navigate efficiently in new neighborhood, and rapidly learn- ing and retaining new object names. Agents with HCAM can extrapolate to task sequences order of magnitude longer than they were trained on, and 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.
past event, even to recall the
Introduction
Human learning and generalization relies on our detailed memory of the past 55 52
16
30
43
watch show, can generally recall its scenes in some detail afterward. If explore neighborhood, we can recall our paths through it order to plan new routes. If we are exposed to a new noun, 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 able transport ourselves into past event and re-live in sequential detail, without attending to everything that has happened since. That is, our recall is both sparse (we attend small chunk of the past, or a few chunks) and detailed (we recover large amount information
35th Conference on Neural Information Processing Systems (NcurIPS 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 I 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 Tr [ can serve as RL agent memories in short tasks 44 we show that they
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 [4] can maintain information over slightly longer timescales, but they struggle with tasks that require recalling a temporally- structured event, rather than simply recalling a single past state. That is, they are particularly poor at the type of event recall that fundamental 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 cannot effectively 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 Hierarchical Chunk Attention Memory (HCAM). The fundamental insight of HCAM is to divide the past into distinct chunks before storing in memory, and 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 long- term recall ability of prior episodic memories [6 with the short-term sequential and reasoning power of transformers [56 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 5x 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 [4 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, RI. agents require self-superva
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 L.STSs on sequential language tasks suggests that they might be effective agent memories, and recent work succesfully used transformers as RL agent memories
2

---

specifically, gated version of TransformerXL [9]. However, using TransformerXL memories on challenging memory tasks still requires auxiliary sulf-supern [1)] 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
No attention
(a) Standard attention.
(b) Hierarchical attention.
(Self supervised
Artn. within top chunks
Memory
ReS
LSTM
MLP
MLF
Decoders
Memory
Atin.
Weighted sum of results
liyer)
Mensory
RCAM Alm.
across
chunks
Attestion
Keys
Chunks
layer
Memory
HCCA Alt.
LayerNorm
Attention
ResNlet
LSTM
Encoders
Input
Skip connection
Output
Image
Language
(c) HCAM attention block.
(d) HCAM in agent.
Figure I: HCAM motivation and implementation. 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. Top- level attention is performed over all summaries (left), and then standard attention is performed within and summed (right), and then added a residual input to produce the block output. (d) Our agent encodes inputs, processes them through 4-layer memory with an HCAM block in each layer, and decodes outputs for self-supervised and RI. losses.
the top-k chunks. The results are relevance-weighhe
The HCAM uses a hierarchical strategy to attend over memory (Fig. D.. Specifically, it divides the past into chunks. It stores each chunk in detail-ass as a sequence-togetther with a single summary key for that chunk. It can then recall hierarchically. by first attending to chunk summaries 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 linear projection ("query") layer. Then the input is normalized by LayerNorm, and the chunk relevance (Rt is computed as:
R = softmax ( (normed input): -S)
Then the model attends 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-heat head attention:
memory query results-
Σ
R₁ N.-MHA( (normed input, C.)
E top-k from R
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 transformer layer, to supplement standard attention. (For an open-source HCAM implementation, see A.I.
HCAM can attend to any point in its memory, but at each step it only attends to a few chunks, which much more compute- and memory-effiiint 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 computation 10x compared to TrXL. Even with k > 1, HCAM often runs faster (App. D.10
in many cases (App- D.5). reducing
3

---

Our agents (Fig. ld] consist of input encoders (a CNN or ResNet for visual input, and an LSTM for language), 4-layer memory (one HCAM, Gated TransformerXL 44, or an LSTM 20). followed by output MIPs for policy, value, and auxiliary losses. Our agents are trained using IMPALA [1]. To use HCAM in the memory, we adapt the Gated TransformerXI. agent memory
Agents
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 our tasks (App. D.7p. so we removed it. See App. A.2 for further agent details.
44
We store the initial layer inputs in memory that is, the inputs to the local attention-wwhs 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-poo across the chunk to generate a summary. We reset the memory between episodes, except where cross-episode memory 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-differentiabiliy of top-k selection (sce App. D for discussion). It also means that 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 cannot update the encoders through the memory. Thus, as 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 at each step we trained the agents (of all memory types) 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.]
Experiments
We evaluated HCAM in wide variety of domains (Fig. 2) Our tasks test the ability of HCAM to recall event details (SIc. to maintain object permanence (Sec. 3.2); to rapidly learn and recall
new words, and extrapolate far beyond the training distribution (Sec. 3.3) and to reason multiple
memories to generalize, even in settings that previously required task-specific architectures (Sec. 3.4).
"This is dax"
(a) Ballet.
(b) Object permanence.
(c) Rapid word learning.
D.C.
1-0 0
t- 125
t- 526
(d) Visual match.
(e) Paired associative inference. (f) One-shot StreetLearn navigation.
Figure 2: We evaluated our model across wide variety of task domains, some of which are modified or used from other recent memory papers 19 21 These include a variety environments (2D. 3D, real images) and tasks (finding hidden objects, language learning, and navigation).
3.1 Remembering the ballet
Our first experiment tests the ability recall spatiotemporal detail, which is important in real-worid events. We place the agent in 2D world, surrounded by "dancers" with different shapes and colors (Fig. 2a One at time, each dancer performs 16-step solo dance (moves in distinctive pattern), with a 6- or 48-step delay between dances. After all dances, the agent is cued to e.g. "go to the dancer who danced a square," and is rewarded if it goes to the correct dancer. This task requires recalling dances in detail, because no single step of a dance is sufficient to distinguish it from other dances. The task difficulty depends on the number of dances the agent must watch in an episode (between and 8) and the delay length. We varied both factors across episodes during training- this

---

likely increases the memory challenges, since the agent must adapt in different episodes to recalling different amounts of information for varying periods. See App. B.2 details.
Fig. shows the results. Agents with an HCAM or a Gated TransformerXI. (TrXL.) memory perform well at the shortest task, with only 2 dances per episode (Fig. Bap. LSTT agents learn much more slowly, showing the benefits of attention-basen memories. With 8 dances per episode (Fig. (b)) HCAM continues to perform well. However, TrXL's performance degrades, even though the entire task is short enough to fit within a single TrXL attention window. (The dances are largely in the XL region, where gradients are stopped for TrXL, but HCAM has the same limitation.) LSTM performance also degrades substantially. The advantage of HCAM is even more apparent with longer delays between the dances (Fig. Bcb. HCAM can also generalize well (2 90%) from training on up to dances to testing on 8 (App. D.2). Furthermore, HCAM is robust to varying memory chunk sizes (App. D.4) even if the chunk segmentation is not task aligned. Finally, HCAM can perform well on this task even if it is only allowed to select single memory chunk rather than the top-k (App. D.5), thereby attending to 8x fewer time-points than TrXL. (while a sparse TrXL is not advantageous, see App. [D.9]. App. Thus HCAM can increase computational efficiency. In summary, both attention-basee memories recall events better than LSTMs, and HCAM robustly outperforms TrXL at all but the easiest tasks.
100
100
100
Accuracy (86)
75
Accuracy (%%)
75
Accuracy (%5)
75
50
Chance
50
HCAM
25
25
Chance
Chance
LSTM
0
0
0
Training steps
1e9
Training steps
100
Training steps 109
(a) dances.
(b) 8 dances.
(c) dances, long delays.
The ballet task HCAM (teal) outperforms TiXL (yellow) and LSTM (red) memories at difficult ballets with many dances, especially when there are long delays between. In the shortest ballets, TrXL performs nearly as well as HCAM, while LSTMs require much more training. (b) With 8 dances, HCAM substantially outperforms TrXL, even though the entire task fits within one TrXL attention window. LSTMs do not get far above chance at choosing among the dancers. Cc) With longer delays between dances, HCAM performs similarly but TrXL degrades further. (3 runs per condition, lines=means, light regions=range, SD.. Training steps are agent/environment steps, not the number of parameter updates. Chance denotes random choice among the dancers, not random actions.)
Figure 3:
3.2 Object permanence
Our next tasks test the ability to remember the identity and location of objects after they have been hidden within visually-identica boxes, despite delays and looking away. These tasks are inspired by work object permanence in developmental psychology [e.g. [1] We implemented these tasks in a 3D Unity environment. The agent is placed in a room and fixed a position where it can see three boxes (Fig. 4ap. One at time, object jumps out of each box few times, and then returns to concealment inside. There may be a delay between one object and the next. After all objects have appeared, the box lids close. Then, the agent is released, and must face backwards. Finally, it is asked to go to the box with a particular object, and rewarded if it succeeds.
HCAM solves these tasks more effectively than the baselines. HCAM learns the shortest tasks (with no delays) substantially faster than TrXL (Fig. 4b). and LSTMs are unable to learn at all. When trained on tasks of varying lengths, HCAM was also able to master tasks with 30 seconds of delay after each object revealed itself, while TrXL performed near chance even with twice as much training (Fig, 4b). Furthermore, HCAM can learn these long tasks even without the shorter tasks to scaffold it (Fig. 4d
3.3 Rapid word learning with distractor tasks generalization across episodes
The preceding tasks focused on passively absorbing events, and then recalling them after delays. We next turn to a set of rapid-wor worddwoor--lannngg tasks (Fig. 2c) that require actively acquiring knowledge, and maintaining it in more challenging settings, with intervening distractor tasks (Fig: 5a). These
5

---

"Find the box with duck""
3x object reveals delays
Look away
Recall location
(a) The object permanence task.
100
100
100
Accuracy (%6
75
Accuracy (%%
Accuracy (%)
50
50
50
Chance
Chaose
Chance
25
HCAM
25
25
TIXL
LSTM
9
0
0
2
Training steps
1eb
Training stops
1ee
Training stops
1c9
(b) No delays, varying training.
(c) Long delays, varying training.
(d) Long delays, long-only training.
The object permanence task-HCAM learns faster than TrXL, and succeeds at longer tasks. (a) The task structure. (b) The shortest tasks, with no delays. When trained with varying delay lengths, HCAM rapidly and consistently learns to track object locations. TrXL learns more slowly and less consistently. (b) Long tasks, with 30-second delays between each object being revealed. HCAM learns quite well, while TrXL barely achieves above-chance performance. (c) HCAM can learn the long tasks even without simultaneous training on the shorter tasks to scaffold its learning. (3 runs per condition. Chance denotes random choice among the boxes, not random actions.)
Figure 4:
tasks are based on the work of Hill al. [9] who showed that transformer- -based agents in a 3D environment can meta-learn to bind objects to their names after learning each name only once. That is, immediately after a single exposure to novel object names, their agent can successfully lift objects specified by those names. Their agent then discards its memory before proceeding to the next episode. However, in more complex environments tasks will not be so cleanly structurd. There will often be distracting events and delays or even learning of other information--eefe the agent is tested.
We therefore created more challenging versions of these tasks by inserting distractor tasks between the learning and test phases. In each distractor task we placed the agent in room with a fixed set of three objects (that were never used for the word-learning portion), and asked it lift one of those objects (using fixed names). By varying the number of distractor phases, we were able to manipulate the demands on the memory. We trained our agents on tasks with 1, or distractor phases, and evaluated their ability to generalize to tasks with more distractors. HCAM and TrXL could learn the training tasks equally well, and both showed some ability to generalize to longer tasks (App. D.I]. However, HCAM was able extrapolate better to tasks with 10 distractor phases (Fig. [d]. and even showed some extrapolation to 20 distractors, an order of magnitude more than trained (Fig. 7a
Maintaining knowledge across episodes zero-shot It is a fundamental limitation of most meta- learning paradigms that they discard the information they have learned after each episode. The agents used by Hill et al. [19] discarded their memories end of each episode. Children, by contrast, can rapidly learn new words and maintain that knowledge even after learning other new words. This leads to a question: if HCAM can effectively maintain a word-object mapping despite intervening distractor tasks, could it maintain this knowledge across entire episodes of learning new words?
To answer this question, we used the agents trained on singlc-episode tasks. We evaluated these agents on tasks where they were sometimes tested on a word from several episodes before, despite having learned and been tested on other words in the intervening period (Fig. (b).. In training, the agents were never required maintain a set of words once they encountered a new learning phase with new words. Nevertheless, our agent with HCAM is able to generalize well to recalling words four episodes later (Fig. 5ep. It can even extrapolate along multiple dimensions, recalling words two episodes later with more distractor phases per episode than it ever saw in training (Fig- 57 In App. D.3 we show how the attention patterns of the hierarchical memory support this generalization across episodes. Agents with HCAM are able to generalize to memory tasks that are far from the training distribution.

---

"This das?
"Liin mekket"
"Pind zipe"
100
Learn
Distiittos
Accuracy (N6)
75
HCAN TTAL
"This a dox "Lift rocket"
"Find a dax"
"This is zark"
"Läf train"
"Find des""
50
Chance
Lear new words
Distractor tasks (N
Test word recall
Leam
Distractos
Test older
Training steps
(a) Rapid word-learning with distractor tasks.
(b) Recall across multiple episodes (cvaluation only).
(c) Train task with distrac- tors.
100
100
100
75
Accurasy (80
75
Accuracy (%)
75
Accurary
50
25
0
Training steps
Training steps
169
Training steps
160
(d) Eval. with 10 distractors.
(c) Eval. across 4 eps., 0 dist.
(f) Eval. across 2 eps., 5 dist. each.
Figure 5: The rapid word-learning tasks. Training tasks consist of learning three words (nouns),
completing distractor tasks, and then finding object named using one of the three learned words. We evaluated generalization to larger numbers of distractors. (b) During training, the agent never had to recall words for more than a single episode. However, we also evaluated the agent on its ability to recall a word learned several episodes earlicr. This tests the ability of our memory to bridge from rapid meta-learning longer-tern retention. (c) Agents with HCAM or TrXL memories learn the hardest
training task (2 distractors) well, although the agent with HCAM is faster to learn. When trained only single episodes with 0-2 distractor phases: <ch HCAM outperforms TrXL at extrapolation to 10 distractor phases. (e I HCAM can generalize strongly out-of- -iitti--disiitttioon to recalling words that were learned several episodes earlier, despite having completed other episodes in the intervening time. (3 runs per condition. Chance denotes random performance on the final evaluation choice; this requires complex behaviours: learning all names and completing all distractors and intermediate episodes to reach the final choice. This is why LSTM performance is below chance in one seed.)
3.4 Comparisons with alternative memory systems
There has been an increasing interest in architectures for memory in RL, as well in deep learning more broadly. In order to help situate our approach within this literature, we benchmarked our model against tasks used in several recent papers exploring alternative forms of memory. HCAM is able to perform well across all the tasks-reaching near-optimal performance, comparable with memory architectures that were specifically engineered for each tndk-and outperforms strong baselines.
Passive Visual Match First, compare the Passive Visual Match task from Hung et al. (21). In this task, the agent must remember color it sees at the beginning of the episode, in order to choose a matching color the end (Fig. 2d.. HCAM achieves reliably accurate performance (Fig. 6a), comparable to RMA, the best model from Hung et al. [21]. It would be interesting to combine HCAM with the value transport idea Hung al. proposed, in order to solve their active tasks.
Paired Associative Inference We next considered the Paired Associative Inference task 2 In each episode, the agent receives set of pairings of stimuli, and it must chain together transitive inferences over these pairings in order to respond to a probe (Fig. 2cp. These tasks therefore require sequential reasoning over multiple memories. HCAM achieves comparable performance MEMO the length 3 version of this task, and substantially outperforms baselines like Universal Transformer [10]. This likely because HCAM, like MEMO, respects the chunk structure of experiencing pairs of images.
One-Shot StreetLeeam Finally, we evaluated HCAM on One-Shot Streetl Leann (Fig. 21 proposed by Ritter al. [49], based on Mirowski et al. 40). One-shot StreetLearn is challenging meta-

---

10
Transttive inference accurasy
1000
MERO
HCAM
Phose: 3 roward
75
Transt:
Univ
Evalustion reward
25
Tro
20
50
Chence
15
Meniin
Crance
10
8.0
1.5
2.5
Training steps
Training steps
Training steps
1e9
(a) Passive Visual Match.
(b) Paired Associative Inference.
(c) One-Shot StreetLLeamn
Figure 6: Evaluating HCAM on tasks from other memory papers. Our agent achieves performance
competitive with the specialized model proposed solve each task. Top results (gold) and selected
baselines (red) from each paper are indicated. (a) The Passive Visual Match task 21
(a) The Paired
optimal performance, comparable to MEMO and above Universal Transformers [T]] (c) The One-Shot Streetl Leamn environment 49] HCAM
Associative Inference task (2. HCAM achieves
achieves comparable performance to EPN this difficult navigation task, although it is slower to
learn. Both substantially outperform strong baselines. (Some prior results were estimated from figures.
HCAM results aggregate over: (al 6 seeds. (b) seeds, selected by validation accuracy. (ch 2 seeds.)
learning setting. The agent is placed in a neighborhood. and must navigate to as many goals as it can within a fixed time. The agent receives its current state and goal visual inputs (Google Street images) and must learn about the neighborhood in early tasks within an episode, in order to navigate efficiently later. Ritter al. designed specialized Episodic Planning Network (EPN) to solve these tasks. HCAM-despir despite its more general purpose architecture- can achieve comparable performance EPN (although HCAM somewhat slower reach this level of performance). Ritter et al. showed that EPN achieves near-optimal planning in later parts of each episode, and therefore HCAM must be planning close to optimally to match EPN's performance. Strong baselines perform much worse. Merlin [60] strong agent with a learned episodic key-value memory and auxiliary unsupervised training-only achieves around half the performance of HCAM and EPN. TrXL is also less effective. This highlights the value of more sophisticated memory architectures, and in particular emphasizes the value of sequences stored in memory, rather than single vectors: EPN stores transitions (sequences of length 2) and HCAM stores longer sequences, while Merlin has only single vector values for each key (as does TrXL).
4 Discussion
In this paper we have proposed the Hierarchical Chunk Attention Memory. HCAM allows agents to attend in detail to past events without attending to irrelevant intervening information. Thus, HCAM effectively implements form of "mental time-travel" (55). Our main insight is that this ability can be achieved using a hierarchically structured memory that sparsely distributes detailed attention. Specifically, recall uses attention over memory chunk summarics to allocate more detailed attention within only the most relevant chunks. We see this as a potentially important insight, because both mental time travel [52) and hierarchical structure [16] are essential to the power of human memory. Furthermore, recent work has suggested that human memory behaves more like a "quantum wave function" distribution over several past events (37], which matches our approach of distribution over the top-k relevant chunks. HCAM incorporates a number of important features of human memory.
Correspondingly. we showed that agents with HCAM succeed at wide range of environments and tasks that humans depend on memory to solve. HCAM allows agents to remember sequential events, such ballet, in detail. HCAM allows agents to rapidly Iearn to navigate near-optimally in new neighborhood by planning over previous paths; to perform transitive inferences; and to maintain object permanence, despite delays, occlusion and looking away. HCAM allows agents to effectively learn new words from a single exposure, and maintain that knowledge across distractor tasks. The agents can extrapolate far beyond the training distribution to recall words after subsequent learning episodes without ever being trained to do so better memory systems can help with the challenge of bridging from meta-learning continual learning, which is receiving increasing interest [13 17 49 41).
The tasks we considered are challenging because they rely on flexible use of memory. The agent did not know in advance which paths to remember in One-Shot StreetLearn, which dancer would
8

---

be chosen in Ballet. Furthermore, we trained the agent simultaneously on procedurally generated episodes where task features varied- e.g. the agent did not know how many distractor phases would appear in an episode, or the length of delays, so it could not rely on a fixed policy that recalls information after a fixed time. HCAM might be especially useful in these difficult, variable settings.
HCAM is robust to hyperparameters such as chunk size (App. D.4) and the number k of memories selected at each layer and step (App. D.5 is even able to solve many tasks with A: 1. It also outperforms TrXL models that are either twice as wide or twice as deep, and correspondingly have many more parameters and, in the deeper case, make many more attention computations (App. [D.). Our comparisons to other approaches (Sec. (..4) show that our approach is competitive with state of the art, problem-specific memory architectures, and outperforms strong baselines like Merlin and Universal Transformers. Finally. in hyperparameter sweeps suggested by our reviewers to improve TrXI.ss performance, we found that HCAM was consistently more robust to hyperparameter variation (App. D.I1]. These observations suggest that our results should be relatively generalizable.
Self-supervised learning As in prior memory work [60
[9] we found that self-supervised learning was necessary to train the agent to store all task-relevant information in memory. Training the agent to reconstruct its input observations as outputs was sufficient in the tasks we explored, but might be useful in other settings.
more sophisticated forms of auxiliary learning 23 3
Memory in RL vs. supervised settings Our work has focused on improving the memory of situated RL agent. Compared to supervised settings such as language modelling, where a vanilla transformer can achieve very impressive performance [8 4. RL poses unique challenges to memory architectures. First, sparse rewards present a more impoverished learning signal than settings that provide detailed errors for each output. The ability of HCAM to restore a memory in detail may help ameliorate this problem. Second, the multimodal (language + vision) stimuli experienced by our agent contains much more information per step than the word-pieces alone that a language model experiences, and therefore our setting may place more demands on the memory architecture. Finally, our tasks require access to detailed, structural aspects of past memories, while many existing NLP tasks do not depend significantly on structure for example, Pham et al. [45] show that BERT ignores word order when solving many language understanding benchmarks, and produces equivalent outputs for shuffled inputs. Detailed memory will be most beneficial in settings in which the sequential structure of the past is important.
Nonetheless, we do not rule out possible applications of HCAM in supervised settings, such as language processing 48 4 video understanding (53), or multimodal perception [24]. HCAM would likely be most beneficial in tasks that strongly rely on long-term context and structure, c.g. the full-length version of NarrativeQA 29). Because videos often contain hierarchically-strun events, video models might especially benefit from HCAM-style attention. More broadly, the greater efficiency of sparse, hierarchical attention might allow detailed memory even in settings with limited computational resources, such embedded systems.
Episodic memories Nematzadeh et al. [42) suggested endowing transformers with external episodic memories. Several recent papers have proposed systems that can be interpreted episodic memories from this perspective. specifically looking up related contexts from the training corpus using either nearest neighbors 266 or attention 622. However, these approaches have generally only used this type of external memory to predict outputs, rather than allowing the model to perform further computations over the memories recalls, as in HCAM. Thus, these memories would be inadequate for most RL tasks, which require planning or reasoning with memories.
Various works have proposed other types of external/episodic memory, both in RL specifically [e.g. 60 21 1.I and in deep learning more broadly [e.g. 550 [5]. These approaches have generally not stored memories with the hierarchical summary-chunk structure of HCAM. Ritter et al. [49] proposed an episodic memory for navigation tasks that stores state transitions; this architecture can be seen as a step towards our approach of storing sequences as chunks. Indeed, in the challenging city navigation domain that Ritter et al. explored, we showed that HCAM achieves comparable performance to their EPN architecture, despite HCAM being much more general. Furthermore, both HCAM and EPN substantially outperform Merlin [60), a strong baseline that learns to store vector memories, rather than the rich representations of sequences stored by HCAM (or transitions stored by EPN). This highlights the value of rich, sequential memories. We suggest that episodic memories could benefit from moving beyond the idea of keys and values as single vectors. It can be useful to store more general ttuctures-suct as a time-sequence of states-aa as a single "vaauee" in memory.
9

---

One benefit of episodic memory is that memory replay can support learning [30 38 in particular by ameliorating catastrophic interference. It would therefore be interesting to explore whether HCAM's memory could be used for training. For example, could memory summaries be used to locate similar or contrasting memories that should be interleaved together [39]] to improve continual learning?
Transformer memories Many recent works have attempted to improve the computational effi- ciency of transformer attention over long sequences [e.g.28] (ag.28|588 (58. However, these approaches often perform poorly at even moderately long tasks [54]. Similarly, we found that on tasks like Ballet or One-Shot Streetl Transformer-Xi can perform suboptimally even when the entire task fits within a single, relatively short attention window. However, there could potentially be complementary benefits to combining approaches to obtain even more effective attention with hierarchical memory. which should be investigated in future work. In particular, some recent work has proposed simple inductive bias which allows Transformers to benefit from evaluation lengths much longer than they were trained on [46] while this strategy would likely not help with recalling specific instances in detail, it might be complementary to an approach like ours.
Other works have used a hierarchy of transformers with different timescales for supervised tasks [4 64 61 (22, for example encoding sentences with one transformer, and then using these embeddings as higher-level tokens. This approach improves performance on tasks like document summarization, thus supporting the value of hierarchy. Luong et al. [66 also showed benefits of both local and global attention computation in L.STM-based models. However, we are not aware prior transformers in which the coarse computations are used to select chunks for more detailed computation, as in HCAM (although recent work has demonstrated top-k patch selection in CNNs (8) nor are we aware of prior works that have demonstrated their approaches in RL, or beyond a single task domain.
Memory and adaptation One major benefit of memory is that a model can flexibly use its memo- ries in goal- context-dependent way in order to adapt new tasks [55 52 30). Adult humans
use our recall of rich, contextual memories in order to generalize effectively [22 43 While our tasks required some forms of goal-dependent memory use use-e.g. combining old paths plan new ones-it would be interesting to evaluate more drastic forms of adaptation. Recent work shows that transforming prior task representations allows zero-shot adaptation to substantially altered tasks [I). HCAM could potentially learn to transform and combine memories of prior tasks in order to adapt to radically different settings. Exploring these possibilities offers an exciting future direction.
Limitations & future directions While we see our contribution as a step towards more effective mental time-travel for agent memory, many aspects could be further improved. Our implementation requires the ability to keep each previous step in (hardware) memory, even if the chunk containing that step is irrelevant. This approach would be challenging to scale to the entire lifetime of cpisodic memory that humans retain. Nonetheless, storing the past is feasible up tens of thousands of steps on current accelerators. This could be extended further by using efficient ANN implementations, which have recently been used to perform ANN lookup over an entire language dataset (26).
The challenge of scaling to longer-term memory would be helped by deciding what to store, as in some episodic memory models [15], soft forgetting of rarely retrieved memories, and similar mechanisms. However, such mechanisms are simultaneously limiting, in that the model may be unable to recall knowledge that was not obviously useful at the time of encoding. Over longer timescales, HCAM might also benefit from including more layers of grooupnnn memory chunks into higher-order chunks recursively to achieve logarithmic complexity for recall. Even over moderate timescales, more intelligent segmentation of memory into chunks would potentially be beneficial human encoding and recall depends upon complex strategies for event segmentation 633 51 6 35| 77 HCAM might also benefit from improved chunk summaries; mean- aocaann over the chunk is a relatively naive approach to summarization, learning compression mechanism might be beneficial [47.. These possibilities provide exciting directions for future work.
Conclusions We have proposed a Hierarchical Chunk Attention Memory for RL agents. This architecture allows agents to recall their most relevant memories in detail, and to reason over those memories to achieve new goals. This approach outperforms (or matches the optimal performance of) a wide variety of baselines, across wide variety of task domains. It allows agents to remember where objects were hidden, and to efficiently learn to navigate a new neighborhood by planning from memory. It allows agents to recall words they have learned despite distracting intervening tasks, and even across episodes. These abilities to learn, adapt, and maintain new knowledge are critical to intelligent behavior, especially as the field progresses towards more complex environments. We
10

---

hope that our work will motivate further exploration of hierarchically structured memories, in RL and beyond. Hierarchical memory may have many benefits for learning, reasoning. adaptation, and intermediate- or long-term recall.
Acknowledgments and Disclosure of Funding
We would like to acknowledge Adam Santoro, Emilio Parisotto, Jay McClelland, David Raposo, Wilka Carvalho, Tim Scholtes, Tamara von Glehn, Sébastien Racanière, and the anonymous reviewers for helpful comments and suggestions.
References
[1] Renee Baillargeon. Infants' reasoning about hidden objects: evidence for event-general and event-specific expectations. Developmental science, 7(4):391-414, 2004.
[2] Andrea Banino, Adrià Puigdomenech Badia, Raphael Köster, Martin Chadwick, Vinicius Zambaldi, Demis Hassabis, Caswell Barry, Matthew Botvinick, Dharshan Kumaran, and Charles Blundell. Memo: A deep network for flexible combination of episodic memories. In International Conference on Learning Representations, 2020.
[3] Andrea Banino, Adrià Puidomenech Badia, Jacob Walker, Tim Scholtes, Jovana Mitrovic, and Charles Blundell. Coberl: Contrastive bert for reinforcement learning. arXiv preprint arXiv:2107.0543// 2021.
[4] Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005. /4/65, 2020.
[5] Massimo Caccia, Pau Rodriguez, Oleksiy Ostapenko, Fabrice Normandin, Min Lin, Lucas Caccia, Issam Laradji, Irina Rish, Alexande Lacoste, David Vazquez, et al. Online fast adap- tation and knowledge accumulation: a new approach to continual learning. arXiv preprint WrXiv:2003.05850 2020.
[6] Stephanic CY Chan, Marissa C Applegate, Neal W Morton, Sean M Polyn, and Kenneth Norman. Lingering representations of stimuli influence recall organization. Neuropsychologia. 97:72-82, 82, 2017.
[7] Claire H. C. Chang, Christina Lazaridi, Yaara Yeshurun, Kenneth A. Norman, and Uri Hasson. Relating the Past with the Present: Information Integration and Segregation during Ongoing Narrative Processing. Journal of Cognitive Neuroscience, pages 1-23, 04 2021. ISSN 0898- 929X. doi: 10. 221166/jocn___017000 URL/https: /ttp::::////oo///////////// org/10 1162/ (1777
[8] Jean-Baptiste Cordonnier, Aravindh Mahendran, Alexey Dosovitskiy, Dirk Weissenborn, Jakob Uszkoreit, and Thomas Unterthiner. Differentiable patch selection for image recognition. arXiv preprint arXiv:2/04.03359, 2021.
[9] Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdinov; Transformer-xl: Attentive language models beyond a fixed-length context. arXiv preprint arXiv:1901.02860, 2019.
[10] Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Uni- versal transformers. arXiv preprint frXiv:1807.038/11 2018.
[11] Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. In International Conference on Machine Learning, pages 1407-1416. PMLR, 2018
[12] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. arXiv preprint arXiv:2101.03961, 2021.
11

---

[13] Chelsea Finn, Aravind Rajeswaran, Sham Kakade, and Sergey Levine. Online leta-ne In International Conference on Machine Learning, pages 1920-1930. PMLR, 2019.
[14] Meire Fortunato, Melissa Tan, Ryan Faulkner, Steven Hansen, Adrià Puigdomènech Badia, Gavin Buttimore, Charlie Deck, Joel Z Leibo, and Charles Blundell. Generalization of rein- forcement learners with working and episodic memory. arXiv preprint arXiv: 100X:::1900.3406, 2019.
[15] Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwińska, Sergio Gómez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, et al. Hybrid computing using neural network with dynamic external memory. Nature, 538 (7626):471-476, 2016.
[16] Uri Hasson, Janice Chen, and Christopher J Honey. Hierarchical process memory: memory as an integral component of information processing. Trends in cognitive sciences, 19(6):304-313, 2015.
[17] Xu He, Jakub Sygnowski, Alexandre Galashov, Andrei A Rusu, Yee Whye Teh, and Razvan Pascanu. Task agnostic continual learning via meta learning, arXiv preprint arXiv: XXi::1906.0500,,,, 2019.
[18] Felix Hill, Andrew Lampinen, Rosalia Schneider, Stephen Clark, Matthew Botvinick, James McClelland, and Adam Santoro. Environmental drivers of systematicity and generalization in situated agent. International Conference on Learning Representations, 2020.
[19] Felix Hill, Olivier Tieleman, Tamara von Glehn, Nathaniel Wong, Hamza Merzic, and Stephen Clark. Grounded language learning fast and slow. In International Conference on Learning Representations, 2021.
[20] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735-1780, 1997.
[21] Chia-Chun Hung, Timothy Lillicrap, Josh Abramson, Yan Wu, Mehdi Mirza, Federico Carnevale, Arun Abuja, and Greg Wayne. Optimizing agent behavior over long time scales by transporting value. Nature communications, 10(1):1-12, 2019.
[22] D. Hunter. Matplotlib: 2d graphics environment. Computing in Science & Engineering, (3):90995, 2007. doi: 10.1 109/MCSE.2007.55.
[23] Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. In International Conference on Learning Representations, 2016.
[24] Andrew Jaegle, Felix Gimeno, Andrew Brock, Andrew Zisserman, Oriol Vinyals, and Joao Car- reira. Perceiver: General perception with iterative attention. arXiv preprint arXiv:2 2021.
[25] Nan Rosemary Ke, Anirudh Goyal, Olexa Bilaniuk, Jonathan Binas, Michael Mozer, Chris Pal, and Yoshua Bengio. Sparse attentive backtracking: Temporal credit assignment through reminding, In Advance in Neural Information Processing Systems, 2018.
[26] Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. General- ization through memorization: Nearest neighbor language models. In International Conference on Learning Representations, 2020.
[27] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 14Xiv 2014.
[28] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001 .04451, 2020.
[29] Tomáš Kočiský, Jonathan Schwarz, Phil Blunsom, Chris Dyer, Karl Moritz Hermann, Gábor Melis, and Edward Grefenstette. The narrativeqa reading comprehension challenge. Transac- tions of the Association for Computational Linguistics, 6:317-328, 2018.
12

---

[30] Dharshan Kumaran, Demis Hassabis, and James L McClelland. What learning systems do intelligent agents need? complementary learning systems theory updated. Trends cognitive sciences, 20(7):112-334 2016.
[31] Andrew K Lampinen and James L McClelland. Transforming task representations to perform novel tasks. Proceedings of the National Academy of Sciences, 117652)339770-32911 2020.
[32] Wenda Li, Lei Yu, Yuhuai Wu, and Lawrence C Paulson. Isarstep: benchmark for high-level mathematical reasoning. In International Conference on Learning Representations, 2020.
[33] Xiao Liu, Fanjin Zhang, Zhenyu Hou, Zhaoyu Wang, Li Mian, Jing Zhang, and Jie Tang. Self- supervised learning: Generative or contrastive. arXiv preprint arXiv:2006.082 1(2), 2020.
[34] Yang Liu and Mirella Lapata. Hicrarchical transformers for multi-document summarization. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 1070-5081, Florence, Italy, July 2019. Association for Computational Linguistics. doi:
10.18653/v1/199-15555 1500. URL https: //::://wwww /www. aclweb. .rg.oo...r/ooo.oooo.o.. httt///nhoooogg11911.
[35] Qihong Lu, Uri Hasson, and Kenneth A Norman. Learning to use episodic memory for event prediction. bioRxiv, 2020.
[36] Minh-Thang Luong, Hieu Pham, and Christopher Manning. Effective approaches to attention- based neural machine translation. arXiv preprint arXiv:1508. 04005, 2015.
[37] Jeremy R Manning. Episodic memory: Mental time travel quantum "memory wave" function? Psychological Review, 2021.
[38] James L McClelland, Bruce L McNaughton, and Randall C C O'Rellly. Why there are complemen- tary learning systems in the hippocampus and neocortex: insights from the successes and fail- ures of connectionist models of learning and memory. Psychological review, 102(3):419, 1995.
[39] James I. McClelland, Bruce L. McNaughton, and Andrew K Lampinen. Integration of new information in memory: new insights from a complementary learning systems perspective. Philosophical Transactions of the Royal Society B, (75(1799):201 20119977, 2020.
[40] Piotr Mirowski, Andras Banki-Horvath, Keith Anderson, Denis Teplyashin, Karl Moritz Her- mann, Mateusz Malinowski, Matthew Koichi Grimes, Karen Simonyan, Koray Kavukcuoglu, Andrew Zisserman, et al. The streetlcam environment and dataset. arXiv preprint arXiv:1903.01292, 2019.
[41] Arseny Moskvichev and James A Liu. Updater-extractor architecture for inductive world state representations. arXiv preprint arXiv:2104.05500 2021.
[42] Aida Nematzadeh, Sebastian Ruder, and Dani Yogatama. On memory in human and artificial language processing systems. In Bridging Al and Cognitive Science Workshop at ICLR 2020, 2020.
[43] Chi T. Ngo, Susan L. Benear, Haroon Popal, Ingrid R. Olson, and Nora S. Newcombe, Con- tingency of semantic generalization on episodic specificity varies across development. Cur- rent Biology, 2021. ISSN 0960-9822, doi: httpa////aaaaaaeeeee//////////// attps://wee www.ciiitwneeett.coooooooooooo/e wttp:tticcccccccccceeceeeeee http:scccccccccceeeeee/ecceccces
088.. URL
[44] Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayaku- mar, Max Jaderberg, Raphael Lopez Kaufmmnn, Aidan Clark, Seb Noury, et al. Stabilizing trans- formers for reinforcement learning. In International Conference on Machine Learning, pages 7487-7498. PMLR, 2020.
[45] Thang M. Pham, Trung Bui, Long Mai, and Anh Nguyen. Out of order: How important is the sequential order of words in sentence in natural language understanding tasks? arXiv preprint 102 180, 2020.
[46] Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108. 12409, 2021.
13

---

[47] Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. In International Conference on Learning Representations, 2020.
[48] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 21:1-67, 2020.
[49] Samuel Ritter, Ryan Faulkner, Laurent Sartran, Adam Santoro, Matt M Botvinick, and David Raposo. Rapid task-solving in novel environments. In International Conference on Learning Representations, 2021.
[50] Adam Santoro, Sergey Bartunov, Matthew Botvinick, Daan Wierstra, and Timothy Lillicrap. Meta-leaming with memory-augmentted neural networks. In International conference on machine learming, pages 1842-1850. PMLR, 2016.
[51] Ignasi Sols, Sarah DuBrow, Lila Davachi, and Lluís Fuentemilla. Event boundaries trigger rapid memory reinstatement of the prior events to promote their representation in long-term memory. Current Biology, 27(22):349--3504 2017.
[52] Thomas Suddendorf, Donna Rose Addis, and Michael C Corballis. Mental time travel and the shaping of the human mind. Philosophical Transactions of the Royal Society B: Biological Sciences, 364((1511::311713222 2009.
[53] Chen Sun, Fabien Baradel, Kevin Murphy, and Cordelia Schmid. Learning video representations using contrastive bidirectional transformer. arXiv preprint
2019.
[54] Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena : A benchmark for efficient transformers. In International Conference on Learning Representations, 2021. URL wwwpeeewweeeffooeee iidgore grC2k
[55] Endel Tulving. Memory and consciousness. Canadian Psychology/Psycchooge canadienne, 26 (1):1, 1985.
[56] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. arXiv preprint 1eXXiv11706.337 2017.
[57] Jane X Wang. Meta-learning in natural and artificial intelligence. Current Opinion in Behavioral Sciences, 38:90-95, 2021.
[58] Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint Stt:::20000047768 2020.
[59] Michael L. Waskom. seaborn: statistical data visualization. Journal of Open Source Software 6(60):3021, 2021. doi: 10.21105/joss. 03021. URL https:://doi.oor//00.21100 joss 03021
[60] Greg Wayne, Chia-Chun Hung. David Amos, Mehdi Mirza, Arun Ahuja, Agnieszka Grabska- Barwinska, Jack Rae, Piotr Mirowski, Joel Z Leibo, Adam Santoro, al. Unsupervised predictive memory in goal-directed agent. arXiv preprint 1rXiv:1803.
2018.
[61] Linyi Yang, Tin Lok James Ng, Barry Smyth, and Riuhai Dong. Html: Hierarchical transformer- based multi-task learning for volatility prediction. In Proceedings of The Web Conference 2020, pages 441-451, 2020.
[62] Dani Yogatama, Cyprien de Masson d' Autume, and Lingpeng Kong. Adaptive semiparametric language models. arXiv preprint arXiv:2102.02557 2021.
[63] Jeffrey M Zacks and Khena M Swallow. Event segmentation. Current direction in psychological science, 16(2):80-84, 2007.
[64] Xingxing Zhang, Furu Wei, and Ming Zhou. Hibert: Document level pre-training of hierarchical bidirectional transformers for document summarization. arXiv preprint 1905000000 2019.
14


Reviews:
[{"review_id": "bf_9TRW2Eqc", "reviewer": "Reviewer_5ryt", "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "\n1 Strengths:\n\nThe main strength of the paper is the sheer number of different environments that were used to benchmark the method. It really is impressive and shows that the method can potentially handle different scenarios. Taking it a step further and comparing the architecture to state-of-the-art results in other domains is also a big plus.\n\nAnother strength is that the approach is fairly intuitive - attending over chunks instead of over individual states makes intuitive sense. \n\nMany ablations are provided in the appendix.\n\n2 Weaknesses:\n\nTuning the chunk size for each environment feels a little hacky. As mentioned by the authors – learning the chunk size would be an interesting avenue of future work.\n\n3 Correctness:\n\nFrom scanning the appendix – Figure 10 shows that the self-supervised loss seems to be crucial to obtain decent performance in the Ballet and Rapid word learning tasks. The model performs close to chance without it. I don’t think this is emphasized nearly enough in the main text. The method doesn’t work without it…\n\nI find it hard to believe that the transformer XL performs so poorly on some of the simpler tasks. I can’t help but wonder if the baseline was properly tuned for these tasks. I scanned the appendix, and it seems like the same sweep over hyper parameters was performed for the main method as well as the baselines but I wonder if something as simple as tuning the learning rate for the baseline (which to my understanding was not done) could have improved its performance.\n\n4 clarity:\n\nFor the most part the paper is well written and easy to follow. I think the paper could have benefited from a background section where Transformers and RL could have been introduced. As is, without any background, section 2 feels very dense. In particular, lines 95-109 are extremely dense and should be handled with more care. \n\nI also think that Figure 11 from the appendix should somehow make its way into figure 5 of the main paper. It really feels strange to just show that the baselines are failing when in reality they do learn properly on the training task.\n\n5 Relation to prior work:\n\nThe relation to prior work is handled nicely in the discussion section.\n\nUPDATE: Thanks for the replies to my questions. I leave my score as it stands. (6)", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No code and minimal plans to release – this will make it extremely hard for the community to iterate on the approach.\n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "ToxqkQVsFr", "reviewer": "Reviewer_3pZF", "summary": "The authors proposed a memory network that breaks events into chunks, and retrieve memory hierarchically. They tested their model on a wide range of qualitatively different tasks and showed (sometime dramatic) improvement of performance against many baseline models", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality\nAt the high-level, the idea of storing memories in a tree structure may not be mind-blowing. It’s still important to show that it can work in practice. And the range of tasks tested is rather impressive.\n\nQuality\nThe quality of the work is high. The network architecture is intuitive and well-motivated. The empirical results are thorough and convincing. \n\nClarity\nThe paper is clearly written.\n\nSignificance\nThis work provides a solid and significant improvement in the use of memory network for RL agents.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1"}, "extra_text": {}}, {"review_id": "LXxBkakzVvQ", "reviewer": "Reviewer_BE11", "summary": "The paper proposes a new RL agent architecture called Hierarchical Transformer Memory (HTM), which adds a history-chunking mechanism to the base GTrXL/TrXL model. HTM is found to perform well on six memory tasks.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- Originality\n\nThe proposed architecture is novel. But the acronym HTM is too easily confused with the HTM network (Hierarchical Temporal Memory) introduced by Jeff Hawkins in 2004, and still in development today.\n\n\n- Quality\n\nThe paper’s central claim is that HTM’s history-chunking mechanism improves its ability to retrieve information as needed. The most relevant experimental baseline is TrXL (essentially HTM without the chunking). HTM outperforms TrXL on four of six RL domains, but no TrXL results are reported for the other two tasks (Passive Visual Match and Paired Associative Inference). Nevertheless, HTM performs on par with SOTA models on those two tasks.\n\nThe case for HTM is undermined by the scarcity of hyperparameter tuning as reported. Reliable experimental practice generally requires that hyperparameters for baseline models be tuned with as much care as those of the new model. Table 1 (in Appendix A) lists the settings of 28 hyperparameters (HPs) for each of the 6 experiments. Four of those HPs pertain only to HTM, one pertains only to TrXL, and the other 23 apply to both models. The paper discusses the tuning of only 6 HPs (in Table 2 and Appendix D.3). How were the values of the other 22 HPs obtained? Most of them vary by task, and so were presumably found through limited testing of HTM on short runs instead of full sweeps. In no case are separate values reported, one for HTM and one for TrXL. Since TrXL is essentially an ablation of HTM, one would be tempted to tune many or most of the shared HPs on HTM, then apply the best HP configuration to TrXL as well. But this would inadvertently introduce a strong bias in favor of HTM over TrXL. Attention is the central transformer operation, and history chunking alters it in a major way, so there is little reason to believe that a given HP setting will prove equally good for HTM and TrXL.\n\nAll of the reported experiments were run for billions of training steps, demonstrating that sufficient compute was available to tune HPs for HTM and TrXL independently, if only on shorter runs. In the absence of separate tuning, it’s hard to assess the benefit of history chunking.\n\n\n- Clarity\n\nIn general, the paper is very clear and well-written. Here are a few points for improvement:\n\nContrary to the answer on checklist item 3.b, section A.1 does not explain how most hyperparameters were chosen.\n\nIn Figures 3.a and 3.b, why does the HTM curve start at 100%?\n\nIn Figure 5, how could the LSTM do worse than random chance?\n\n\n- Significance\n\nAs noted in “Limitations & future directions”, HTM requires the permanent storage of data from each step. The severity of this limitation strikes me as fundamental, and I don’t expect the potential workarounds to make it a practical approach for most applications of interest in the future. But this is a personal view, and I could be wrong.\n\n\n- Conclusion\n\nThe open questions regarding the choices of hyperparameters lead me to doubt the attribution of HTM’s performance gains to its chunking mechanism. If I have overlooked or misunderstood something, I look forward to being corrected.\n\n**POST-RESPONSE UPDATE**\n\nI commend the authors for their helpful responses. They have clarified several issues related to hyperparameters, and have performed supplementary hyperparameter sweeps that improve the work.\n\nRegarding the prior use of the term HTM, I find no problem with any of the suggested alternatives. HCAM seems the clearest to my mind.\n\nI recommend replacing “chance level performance” with a clearer term, as it can too easily be interpreted to mean the performance of a trivial agent that simply takes random actions. \n\nI have raised my score by one point.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "12"}, "extra_text": {}}]

Author rebuttal:

Meta-review: - The proposed method is tackling an important problem, reasonable, and the experiments are well designed and executed.
- The major concerns from the reviewers (e.g., hyper parameter tuning, chunk size, TxXL results, etc.) are addressed well enough by the rebuttal. 
- For reproducibility and accessibility, I suggest to open source the code and make the environment available at your earliest convenience. 
- I suggest accepting the paper. 
Decision: Accept (Poster)