Rapid Task-Solving in Novel Environments
Authors: Samuel Ritter
Ryan Faulkner
Laurent Sartran
Adam Santoro
Matthew Botvinick
David Raposo
Venue: ICLR 2021 Poster (2021)
Keywords: deep reinforcement learning
meta learning
deep learning
exploration
planning

Abstract: We propose the challenge of rapid task-solving in novel environments (RTS), wherein an agent must solve a series of tasks as rapidly as possible in an unfamiliar environment. An effective RTS agent must balance between exploring the unfamiliar environment and solving its current task, all while building a model of the new environment over which it can plan when faced with later tasks. While modern deep RL agents exhibit some of these abilities in isolation, none are suitable for the full RTS challenge. To enable progress toward RTS, we introduce two challenge domains: (1) a minimal RTS challenge called the Memory&Planning Game and (2) One-Shot StreetLearn Navigation, which introduces scale and complexity from real-world data. We demonstrate that state-of-the-art deep RL agents fail at RTS in both domains, and that this failure is due to an inability to plan over gathered knowledge. We develop Episodic Planning Networks (EPNs) and show that deep-RL agents with EPNs excel at RTS, outperforming the nearest baseline by factors of 2-3 and learning to navigate held-out StreetLearn maps within a single episode. We show that EPNs learn to execute a value iteration-like planning algorithm and that they generalize to situations beyond their training experience.

Published as a conference paper at ICLR 2021
RAPID TASK-SOLVING IN NOVEL ENVIRONMENTS
Samuel Ritter, Ryan Faulkner,' Laurent Sartran, Adam Santoro, Matthew Botvinick, David Raposo
DeepMind London, UK
fritters, rfaulk, 1sartran, adamsantoro, botvinick, (opposo)@ooggee...o
ABSTRACT
We propose the challenge of rapid task-solving in novel environments (RTS), wherein an agent must solve series of tasks rapidly as possible in unfa- miliar environment. An effective RTS agent must balance between exploring the unfamiliar environment and solving its current task, all while building a model of the new environment over which it can plan when faced with later tasks. While modern deep RL agents exhibit some these abilities in isolation, none are suit- able for the full RTS challenge. To enable progress toward RTS, we introduce two challenge domains: (1) a minimal RTS challenge called the Memory & Planning Game and (2) One-Shot StreetLearn Navigation, which introduces scale and com- plexity from real-world data. We demonstrate that state-of-the-an deep RL agents fail at RTS in both domains, and that this failure is due to inability to plan over gathered knowledge. We develop Episodic Planning Networks (EPNs) and show that deep-RI. agents with EPNs excel RTS, outperforming the nearest baseline by factors of 2-3 and learning to navigate held-out Streetl earnn maps within sin- gle episode. We show that EPNs learn to execute value iteration-like planning algorithm and that they generalize to situations beyond their training experience.
INTRODUCTION
An ideal AI system would be useful immediately upon deployment a new environment, and would become more useful as it gained experience there. Consider for example household robot deployed a new home. Ideally, the new owner could turn the robot on and ask it to get started, say, by cleaning the bathroom. The robot would general knowledge about household layouts to find the bathroom and cleaning supplies. As it carried out this task, it would gather information for use in later tasks, noting for example where the clothes hampers are in the rooms passes. When faced with its next task, say, doing the laundry, it would use its newfound knowledge of the hamper locations to efficiently collect the laundry. Humans make this kind of rapid task-solving in novel environments (RTS) look easy (Lake et al., (017)) but yet it remains aspiration for Al.
Prominent deep RL systems display some of the key abilities required, namely, exploration and planning. But, they need many episodes over which to explore (Ecoffet et al.) 2019:; Badia et al.) 2020) and to learn models for planning (Schrittwieser et al.) 2019). This in part because they treat each new environment in isolation, relying generic exploration and planning algorithms. We propose to overcome this limitation by treating RTS as a meta-rriinforcement learning (RL) problem, where agents learn exploration policies and planning algorithms from distribution over RTS challenges. Our contributions are to:
1. Develop two domains for studying letanle RTS: the minimal and interpretable Mem- ory ory& &lannnng Game and the scaled-up One-Shot StreetLeamn
2. Show that previous meta-RI. agents fail at RTS because of limitations in their ability to plan using recently gathered information.
"Equal contribution.

---

Published as a conference paper at ICLR 2021
a)
sentolee enviooment
sumple envrormen:
State
Goal
b)
Prey, 5939
Episodio
Itrate
Pres ardinn
Stared MLP
Qualit
Stade
Embeddings
Poel WAN
Appaeed
max
elements
Willd momory
Append
Pakicy
Chassen
Bum
Figure 1: (a) Rapid Task Solving in Novel Environments (RTS) setup. A new environment is sampled in every episode. Each episode consists of a sequence tasks which are defined by sampling new goal state and new initial state. The agent has fixed number of steps per episode complete as many tasks as possible. (b) Episodic Planning Network (EPN) architecture. The EPN uses multiple iterations single shared self-attention function over memories retrieved from an episodic storage.
3. Design new architecture the Episodic Planning Network (EPN) that overcomes this limitation, widely outperforming prior agents across both domains.
4. Show that EPNs learn exploration and planning algorithms that generalize to larger prob- lems than those seen training.
5. Demonstrate that EPNs learn value iteration-lik like planning algorithm that iteratively prop- agates information about state-state connectivity outward from the goal state.
2 PROBLEM FORMULATION
Our objective is build agents that maximize reward over sequence of tasks in novel en- vironment, Our basic approach is to have agents learn to do this through exposure to distributions over multi-task environments. To define such a distribution, ronment e as 4-tuple (S, A, Pa: R) consisting of states, actions, state-action transition function, and distribution over reward functions. We then define the notion of task in environment e as Markov decision process (MDP) (S, A, Pa. r) that results from sampling reward function from R.
first formalize the notion of envi-
We can now define framework for learning to solve tasks in novel environments simple gener- alization of the popular meta-RI. framework (Wang al. 2017 Duan et al., 2016). In meta-RI., the agent is trained on MDPs (S, A, Pa,r) r) sampled from a task distribution D. In the rapid task-solving novel environments (RTS) paradigm, we instead sample problems by first sampling an environ- ment from environment distribution E, then sequentially sampling tasks, ie. MDPs, from that environment's 's reward function distribution (see Figure 1).
An agent can be trained for RTS by maximizing the following objective:
E EE [Je, [J,,
where Je,r is the expected reward in environment e with reward function r. When there is only one reward function per environment, the inner expectation disappears and we recover the usual meta- RL objective (J).). RTS can be viewed as meta-RL. RL with the added complication that the reward function changes within the inner learning loop.
2

---

Published as a conference paper at ICLR 2021
Bambled envirommnn
Cracks
X
X
@
Ctseersiion
renard 58
70
Foudoon of cracle
0.8
8
0.-
Opriex replostion
0.-
06
DPM
LETM
169
CPN
LSTM
Marin
MRA
Optirall incloraiin
case
Number *
optiral planning
artic rirciie
Figure 2: Memory&PPlanning Game. (a) Example 4 environment (not observable by the agent) and sttt-ggoal observation. (b) Training curves. Performance measured by the average reward per episode, which corresponds to the average number of tasks completed within a 100-step episode (showing the best runs from large hyper-parameter sweep for each model). (c) Performance mea- sured in the last third of the episodes (post-training), relative to an oracle with perfect information that takes the shortest path the goal. (d) Example trajectory of a trained EPN agent in the first three tasks of an episode. In the first task, the agent explores optimally without repeating states. In the subsequent tasks, the agent takes the shortest path to the goal. (e) Number of steps taken by an agent to reach the nh goal in an episode.
While meta-RL formalizes the problem of learning learn, RTS formalizes both the problems of learning to learn and learning to plan. To maximize reward while the reward function is constantly changing non-trivial novel environments, agents must learn to (1) efficiently explore and effec- tively encode the information discovered during exploration (i.e., learn) and (2) use that encoded knowledge to select actions by predicting trajectories has never experienced (i.e.,plar) plan
3 THE LIMITATIONS OF PRIOR AGENTS
To test whether past deep-RL agents can explore and plan in novel environments, introduce minimal problem that isolates the challenges of (1) exploring and remembering the dynamics of novel environments and (2) planning over those memories. The problem we propose is a simple variation of the well-known Memory Game wherein players must remember the locations of cards in grid. The variation we propose, which we call the Memory Sammory&Planningg Game, extends the chal- lenge require planning as well as remembering (see Figure 2).
In the Memory&PPanning Game, the agent occupies an environment consisting of a grid of symbols (4 x 4). The observation consists of two symbols one which corresponds to the agent's current location, and another that corresponds to the "goal" the agent is tasked with navigating to. The agent can not see its relative location with respect to other symbols in the grid. At each step the agent selects one of possible actions: move left, move right, move up, move down, and collect. the agent chooses the "collect" action when its current location symbol matches the goal symbol, a reward of is received. Otherwise, the agent receives a reward of At the beginning of each transition function. The agent
episode, new set of symbols is sampled, effectively inducing is allowed a fixed number of steps (100) per episode to "collect" as many goals as possible. Each time the agent collects a goal which corresponds completing a task new goal sampled in and the transition function stays fixed.
Following Sutton & Barto 1998 we refer to the ahility to choose actions based predictions of not-yet- experienced events "planning".

---

Published as a conference paper at ICLR 2021
Titanwekre
Ditanwekro
9)
30
EPH
d)
1.0
Onie
c)
46
BIN LSTM
LETM
Number steps to gool
Avinigg roward
-30
0.6
16-
<0
0.-
16
0.
0.0
0.5
1.0
1.5
2.0
2.5
3.0
3.5
4.0
EPW
MRA
Merin
L974
10 15 N task eplaodo
20
25
Training steps
169
Figure 3: One-Shot StreetLeam. (a) Four example states from two randomly sampled neighbor- hoods. (b) Example connectivity graphs. (c) Evaluation performance measured on neighborhoods of held-out city throughout the course of training (showing the best run from large hyper-put parameetee sweep for each model). (d) Performance in last third of episode relative to an oracle with perfect information. (e) Number of steps taken by an agent reach the th goal in an episode.
A successful agent will (1) efficiently explore the grid to discover the current set of symbols and their connectivity and (2) plan shortest paths goal symbol if it has (or can infer) all of the transitions it needs to connect its current location and the current goal location. This setup supplies minimal case of the RTS problem: at the beginning of each episode the agent faces new transition structure (S, A, Pa) and must solve the current task while finding information that will be useful for solving future tasks on that same transition structure. subsequent tasks, the agent must use its stored knowledge of the current grid's connectivity to plan shortest paths.
In this symbolic domain, we find evidence that previous agents learn to explore but not to plan. Specifically, they match the within-trial planning optimum; that is, a strategy that explores opti- mally within each task, but forgets everything about the current environment when the task ends (Figure 2). We hypothesize that this failure is the result of the limited expressiveness of past archi- tectures mechanisms for processing retrieved memories. We designed Episodic Planning Networks to overcome this limitation.
EPISODIC PLANNING NETWORKS
Past episodic memory agents used their memory stores by querying the memories, summing the retrieved slots, then projecting the result through multi-layered perceptron (MLP) (Fortunato al.)
2019 Wayne et al. 2018). We hypothesize that these agents fail to plan because the weighted sum of retrieved slots not sufficiently representationally expressive, and the MLP is not, on its own, sufficiently computationally expressive to support planning in non-trivial environments. To test this hypothesis, we replace the weighted sum and MIP with an an iterative eelf-attention-based architecture designed to support implicit planning.
[T)) consists of three main components: (1) an episodic memory which reflects the agent's experience in the episode so far; (2) an iterative self-attention module that operates over memories, followed by a pooling operation (we will refer to this component as the planner); and (3) policy network.
The architecture, which we call the Episodic Planning Network (EPN, see Figure
Episodic memory At each timestep, we start by appending the embedding of the current goal to each slot in the episodic memory. The result is passed input to the planner. We then add new
In striking parallel, the Memory Game played key role in the development of memory abilities in early episodic memory Reep-RI. agents (Wayne et al. 2018.
4

---

Published as a conference paper at ICLR 2021
memory to the episodic memory. which represents that timestep's transition: it is the concatenation of embeddings of the current observation, the previous action, and the previous observation.
Planner A self-attention-basee based update function is applied over the memories to produce pro- cessed representation of the agent's past experience in the episode and reflecting the agent's belief about the environment. This self-attention operation is iterated some number of times, sharing parameters across iterations. The current state is appended to each element resulting from the self- attention, each of which then passed through shared MLP. The resulting vectors are aggregated by a feature-wise max operation.
Note that the self-attention (SA) does have access to the current state until the very end it does all of its updates with access only to the goal state and episodic memories. This design reflects the intuition that the SA function might learn to compute something like value map, which represents the distance to goal from all states in memory. If the final SA state does in fact come to represent such a value map, then a simple MLP should be sufficient to compute the action that, from the current state, leads to the nearest/highest value state with respect to the goal. We find in Section [ Figure eeviedence that the final SA state comes in fact to resemble an iteratively improved value map.
The specific self-attention update function we used was:
B,+1 fB, + ((B)) <(x) = MHA (1 (layyerooo)))()))
where MHA is the multi-head dot-product attention described by Vaswani al. (2017). Layer nor- malization was applied to the inputs of each iteration, and residual connections were added around the multi-head attention operation In our experiments, was ReLU followed by 2-layer MLP shared across rows.
Policy network- 2-layer multi-layer perceptron (MLP) was applied the output of the planner to produce policy logits and a baseline.
When using pixel inputs, embeddings of the observations were produced by 3-layer convolutional neural network. We trained all agents to optimize an actor-critic objective function using IMPALA, a framework for distributed RL training (Espeholt et al.) 2018). For further details, please refer to the Supplemental Material.
Our results show that the EPN succeeds in both exploration and planning the Memory Planning Game. It matches the shortest-path optimum after only few tasks in novel environment (Figure 2b). and its exploration and planning performance closely matches the performance of a hand-coded agent which combines a strong exploration strategy avoiding visiting the same state twice with an optimal planning policy (Figure 2i.e). For these results and the ones we will present in the following section, all of the episodic memories attend to all of the others. This update function scales quadratically with the number of memories (and the number of steps per episode). In the Supplemental Material present results showing that we recover 92% of the performance using variant of the update function that scales linearly with the size of the memory (see Figure 5 in the Suppl. Material).
STREETLEARN
We now test whether agent extend its success in the Memory&Planning Game domain with high-dimensional pixel-based state observations, varied real-world state-connectivity; longer planning depths, and larger time scales. We introduce the One-Shot StreetLearn domain (see Figure Th), wherein environments are sampled as neighborhoods from the StreetLearn dataset of Google SireetView images and their connectivity (Mirowski al., 2019). Tasks are then sampled by select- ing a position and orientation that the agent must navigate to from its current location.
In past work with StreetLeeam Learn, agents were faced with a single large map (e.g. a city), in which they could leam to navigate over billions of experiences (Mirowski et al. 2018). Here, we partition the
Applyng the layer normalization to the output of the multi-head attention, removing the layer normal- ization altogether, did produce good results in our experiments.
5

---

Published as a conference paper at ICLR 2021
50
Sumbbr of 10 130
Dradde Ail
1.0
Number of
Karodona
Frection of arecle
0.8
Number of teratons
0.6
0.6
16-
Class
0.4
10-
0.2
0.2
gareralooaoon
0.0
Training stops
te0
Nr. intereections
Distamce
4)
-04
-02
laraaioo
Itraaions
Figure 4: Iteration analysis and generalization. (a) Performance of an EPN agent on held-out neigh- borhoods with intersections using planners with 1, and 4 self-attention iterations (showing runs for each condition). (b) Performance on new neighborhoods that were larger (7 and intersections) than the ones used during training intersections). (c) Distance-to- to-ggal to-goal decoding accuracy from the output of the planner after I to iterations. (d) The ability of EPN activations to predict distance to goal expands out from the goal state (blue arrow) as the number of self-attention iterations increases. See section 6.2 for details.
StreetLearn data into many different maps (e.g. neighborhoods of different cities), so that agents can be trained to rapidly solve navigation tasks in new maps. This approach analogous to the partitioning of the ImageNet data into many one-shot classification tasks that spurred advances in one-shot learning (Vinyals et al. 2016). In this case, rather than learning to classify in one shot, agents should learn plan in one shot, that is, plan after one or fewer observations of each transition in the environment.
In One-Shot StreetLearn, the agent's observations consist an image representing the current state and an image representing the goal state (Figure [ip). The agent receives no other information from the environment. The available actions are turn right, which orients the agent clockwise toward the next available direction of motion from its current node; turn left, which does the same in the other direction; and move forward, which moves the agent along the direction it is facing to the next available node. In each episode, we sample new neighborhood with 5 intersections from one of 12 cities. To reduce the exploration problem while keeping the planning difficulty constant, we removed all the locations between intersections that corresponded degree-2 nodes. This resulted in graphs with a median of 12 nodes (see example connectivity graph in Figure 3b. leftmost panel). Every time the agent reaches goal, a new starting- and goal-state pair is sampled, initiating new task, until the fixed episode step limit reached (200 steps). Neighborhoods with 5 intersections allow us sample from approximately 380 possible tasks (i.e. starting- and goal-state pairs) ensuring small chance of repeating any single task. A step that takes the agent to the goal state results in reward of 1. Any other step results in a reward of 0.
One-Shot StreetLeaan has several important characteristics for developing novel-environment plan- ning agents. First, we can generate many highly varied environments for training by sampling neighborhoods from potentially any place in the world. This allows us to generate virtually unlim- ited maps with real-world connectivity structure for agents learn take advantage of. Second, the agent's observations are rich, digh-dimensional visual inputs that simulate person's 's experience when navigating. These inputs can principle be used by an agent infer some of the structure of a new environment. Third, can scale the planning depth arbitrarily by changing the size of the sampled neighborhoods a useful feature for iteratively developing increasing capable planning systems. Fourth, the time-scale of planning be varied independently of the planning depth by varying the number of nodes between intersections. This a useful feature for developing tempo- rally abstract (i.e. jumpy) planning abilities (see Section 7 for discussion).

---

Published as a conference paper at ICLR 2021
EPNs Outperform All Baselines- We trained EPN- agents and previous deep-RL agents on One-Shot StreetLeam neighborhoods with 5 intersections. Our agent significantly outperformed the baselines, successfully reaching 28.7 goals per episode (averaged over 100 consecutive episodes, Figure ( It is important to note that this performance was measured on neighborhoods of a held-out city, which the agent was visiting for the first time. Baseline agents with an episodic memory system (Merlin and MRA) did not exceed 14.5 goals per episode, performing better than ones without (LSTM) that were only able to reach 10.0.
EPNs optimally navigate new maps after 10-15 tasks - Plotting the average number of steps taken by the agent to complete the nhh task of an episode reveals that the agents needs fewer and fewer steps to complete a new task, matching the minimum number of steps required to reach a new goal after 10-15 tasks (Figure 3d. the last third of every episode, EPNs achieve on average 88% 2% of the performance of an oracle with perfect information that systematically takes the shortest-path to the goals (Figure We chose the last third of the episodes for this comparison as a way to exclude steps used for exploration and thus have better estimate of the agent's planning capabilities.
EPNs Generalize to Larger Maps - We trained EPN agents on neighborhoods with intersections and evaluated them (with gradient learning turned off) on larger neighhoroods with 7 and 9 intersec- tions. We found that the trained agents achieved strong performance relative to the oracle, suffering only small marginal drops in performance even maps nearly double the size of those seen during training (Figure 4b. in blue).
More "thinking time" improves performance A desireable property for a planning algorithm would be for performance to increase as "thinking time" increases, allowing the user to obtain better performance simply by allowing more computation time. We test whether EPNs have this property by evaluating the effect of the number of self-attention iterations (i.e., "thinking time") on evalua- tion performance. We observed a systematic performance boost with each additional iteration used in training" (see Figure 4a). When evaluated on neighborhoods with 5 intersections (same neigh- borhood size as the ones used during training) the EPN agent's performance approaches saturation with 2 or more iterations, as indicated by the fraction-of- -oforill performance measure in Figure 4p (leftmost points). When we increase the size of the neighborhoods (7 and 9 intersections) the performance difference between 2 and 4 iterations becomes more pronounced, as indicated by the divergence between the blue and green dashed lines in Figure 4p.
Trained EPNs recursively refine a value-like map - To demonstrate this we froze the weights of a trained EPN agent and replaced its policy network with an MLP decoder of the same size. We then trained the decoder (with a stop gradient to the planner) output the distance from a random state to a random goal, while providing the planner with memories containing all the transitions of a randomly sampled environment. We repeated this experiment with six decoders that received inputs from the planner after a fixed number of iterations, from to 6
The training loss of the different decoders revealed steady increase in the ability to compute dis- tance to goal as we increase the number of self-attention iterations from 1 to 6 (see Figure 6p in the Suppl. Material). This gain is direct consequence of improved decoding for longer distances, as made evident by the gradual rightward shift in the drop-off of the classification accuracy plotted against distance to goal (Figure 4b) This holds true even when we increase the number of iterations beyond the number of iterations used during the training of the planner. Altogether, this suggests that the planner is able to "look" further ahead when given more iterations.
We can visualize this result by selecting a single evaluation environment (in this case, larger neigh- borhood with 12 intersections) with fixed goal location, and measuring the likelihood resulting from all possible states in that environment. This manipulation reveals a spatial pattern in the ability to compute distance to goal, spreading radially from the goal location as we increase the number of self-attention iterations (see Figure 40 and Figure 6p in the Suppl. Material).
Ablations - Besides providing a strong baseline, MERLIN/MRA also represent a key ablation for understanding the causes of EPNs' performance: they have everything the EPN agent has except the planning module. LSTMs represent the next logical ablation: they lack both episodic memory and
"The experiments described here were restricted to hot-hot inputs for quicker tumaround.
Note that while the decoders were trained in supervised setting using I to 6 iterations, the weights of the planner were trained once in the RL. setting using 4 iterations.
7

---

Published as a conference paper at ICLR 2021
the planning module. The performance across these three types of architectures shows that episodic memory alone helps very slightly, but is not enough to produce planning behavior; for that, the self-attention module is necessary.
Another important ablation study is the progressive removal of the later planning iterations (Figure 4a). From this experiment learned that one step of self-attention isn't enough: iterating the planning module is necessary for full performance.
6 RELATED WORK
Deriving repurposable environment knowledge solve arbitrary tasks has been long-standing goal in the field of RI. (Foster & Dayan 2002 Dayan 1993) toward which progress has recently been made (Sutton al., 2011 Schaul et al. 2015; Borsa al., 2018). While quick to accommodate new tasks, these approaches require large amounts of time and experience to learn about new envi- ronments. This can be attributed to their reliance on gradient descent for encoding knowledge and their inability to bring prior knowledge to new environments (Botvinick et al. 2019). Our approach overcomes these problems by recording information about new environments activations instead of weights and learning prior knowledge in the form of exploration and planning policies.
The possibility of learning to explore with recurrent neural networks was demonstrated by Wang al. (2017) and Duan et al. (2016) in the context of bandit problems. Wang et al. (2017) further demonstrated that an LSTM could learn to display the behavioral hallmarks of model-based control in other words, learn to plan in a minimal task widely used for studying planning in humans
Daw et al. (011T). The effectiveness of learned, implicit planning was demonstrated further in fully- observed spatial environments using grid-structured memory and convolutional update functions (Tamar et al. 2016 Lee et al. [2088b Guez et al. 2019). Gupta et al. (2017) extended this approach partially observed environments with ground-trute truth ego-motion by training agents to emulate an expert that computes shortest paths. Our present work can be seen as extending learned planning by (1) solving partially-observel environments by learning to gather the information needed for planning, and (2) providing an architecture appropriate for broader range of (e.g. non-spatial) planning problems.
Model- based reinforcement learning (MBRL) has long been concerned with building agents capable of constructing and planning over environment models (for review see Moerland al., (2020). However, we are not aware of any MBRL agents that learn a deployable model within single episode of experience, as required for the RTS setting.
Learning to plan can be seen as specific case of learning to solve combinatorial optimization problems, a notion that has been taken up by recent work (for review see Bengio et al. 2018). Es- pecially relevant to our work is Kool et al (2018), who show that transformers can learn to solve large combinatorial optimization problems, comparing favorably with more resource-intensse dustrial solvers. This result suggests that in higher-depth novel-enviromment planning problems than we consider in our current experiments, the transformer-based architecture may continue to be effective.
in-
Savinov et al. (2018) develop an alternative to end-to-end learned planning that learns a distance metric over observations for use with a classical planning algorithm. Instead of learning to explore, this system relies on expert trajectories and random walks. It's worth noting that hand-written plan- ning algorithms do not provide the benefits of domain-adapted planning demonstrated by Kool al. (2018), and that design modifications would be needed to extend this approach to tasks requiring abstract planning e.g. jumpy planning and planning over belief states whereas the end-to-end learning approach can be applied to such problems out-of-the-box.
Recent work has shown episodic memory to be effective in extending the capabilities of deep-RL. agents to memory intensive tasks (Oh et al.) 2016) Wayne et al. 2018 Ritter et al.) 2018b) Fortu- nato et al. 2019). We chose episodic memory because of the following desirable properties. First, because the episodic store is non non-parametric, it can grow arbitrarily with the complexity of the envi- ronment, and approximate k-nearest neighbors method make it possible to scale to massive episodic memories in practice, as in Pritzel al. (2017). This means that memory fidelity need not decay with time. Second, episodic memory imposes no particular assumptions about the environment's stucture, making it potentially appropriate choice for variety of non-spatial applications such as
8

---

Published as a conference paper at ICLR 2021
chemical synthesis Segler al. (2018) and web navigation Gur et al. (2018) well as abstract
planning problems of interest in research, such as Gregor et al. ' (2019) Voxel environment tasks.
Our approach can be seen implementing episodic model-baset control (EMBC), a concept
cently developed in cognitive neuroscience (Vikbladh et al. 2017: Ritter 2019). While episodic
control (Gershman & Daw 2017 Lengyel & Dayan 2008 B lundell et al. (0166) produces value
estimates using memories of individual experiences in a model-free manner, EMBC uses episodic
(2018a) showed that a deep-RL agent with episodic memory and a minimal learned planning module (an MLP) could learn to produce behavior consistent with EMBC (Vikbladh et al. (017). Our current work can be seen as using iterated self-attention to scale EMBC to much larger implicit models than the MLPs of past work could support.
memories to inform model that predicts the outcomes of actions. |Ritter
7 FUTURE WORK
We showed that deep-RL agents with EPNs can meta-learn to explore, build models on-the-fly, and plan over those models, enabling them to rapidly solve sequences of tasks in unfamiliar environ- ments. This demonstration paves the way for important future work.
Temporally abstract planning (or, "jumpy" planning, Akilesh et al. 2019) may be essential for agents to succeed in temporally extended environments like the real world. our method makes no assumptions about the timescale of planning, unlike other prominent approaches to learned planning (Schrittwieser et al., 2019; Racanière et al. 2017). Training and evaluating EPNs on problems that benefit from jumpy planning, such One-Shot Streetleear with all intermediate nodes, may be enough to obtain strong temporal abstraction performance.
Planning over belief states may be essential for agents succeed in dynamic partially-observed environments. Planning over belief states might be accomplished simply by storing belief states such those developed by Gregor al. (2019) and training on problem distribution that requires planning over information stored in those states.
Humans are able to solve seemingly open-ended variety of tasks, as highlighted for example by the Frostbite Challenge (Lake et al. 2017). The EPN architecture is in principle suitable for any task class wherein the reward function can be represented by an input vector, so future work may test EPNs in RTS problems with broader task distributions than those addressed in this work, e.g., by using generative grammars over tasks or by taking human input (Fu al. 2019).
REFERENCES
Akilesh, Suriya Singh, Anirudh Goyal, Alexander Neitz, and Aaron Courville. Toward jumpy planning. In International Conference Machine Learning, 2019.
Adrià Puigdomênech Badia, Bilal Piot, Steven Kapturowski, Pablo Sprechmann, Alex Vitvitskyi, Daniel Guo, and Charles Blundell. Agent57: Outperforming the atari human benchmark. arXiv preprint arXiv:2003. 2020.
Yoshua Bengio, Andrea Lodi, and Antoine Prouvost. Machine learning for combinatorial optimiza- tion: a methodological tour d'horizon. arXiv preprint /8/ 2018.
Charles Blundell, Benigno Uria, Alexander Pritzel, Yazhe Li, Avraham Ruderman, Joel Z Leibo, Jack Rae, Daan Wierstra, and Demis Hassabis. Model-free episodic control. arXiv preprint arXiv:1606.044 044600 2016.
Diana Borsa, André Barreto, John Quan, Daniel Mankowitz, Rémi Munos, Hado Hasselt, David Silver, and Tom Schaul. Universal successor features approximators. arXiv preprint arXiv:1812.07626, 2018
Mathew Botvinick, Sam Ritter, Jane X Wang, Zeb Kurth-Nelson, Charles Blundell, and Demis Hassabis. Reinforcement learning, fast and slow. Trends in cognitive sciences, 2019.

---

Published as a conference paper at ICLR 2021
Nathaniel D Daw, Samuel Gershman, Ben Seymour, Peter Dayan, and Raymond Dolan. Model- based influences on humans' choices and striatal prediction errors. Neuron, 69(6):1204-12155, 2011.
Peter Dayan. Improving generalization for temporal difference learning: The successor representa- tion. Neural Computation, 5(4):613-624, 1993.
Yan Duan, John Schulman, Xi Chen, Peter L. Bartlett, Ilya Sutskever, and Pieter Abbeel. 212: Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv: 1.02779, 2016.
Adrien Ecoffet, Joost Huizinga, Joel Lehman, Kenneth O Stanley, and Jeff Clune. Go-explore: a approach for hard-exploration problems. arXiv preprint arXiv: 10995, 2019.
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Volodymir Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, al. Impala: Scalable distributed deep-rl with importance weighted actor- learner architectures. arXiv preprint arXiv: 1802.01660,. 2018.
Meire Fortunato, Melissa Tan, Ryan Faulkner, Steven Hansen, Adrià Puigdomènech Badia, Gavin Buttimore, Charles Deck, Joel Z. Leibo, and Charles Blundell. Generalization of reinforcement learners with working and episodic memory. Advances in Neural Information Processing Systems, pp. 12448-12457 2019.
David Foster and Peter Dayan. Structure in the space of value functions. Machine Learning, 49 2--3):325-346, 2002.
Justin Fu, Anoop Korattikara, Sergey Levine, and Sergio Guadarrama. From language to goals: Inverse reinforcement learning for vision-based instruction following. arXiv preprint 1rXiv:190..07742, 2019.
Samuel Gershman and Nathaniel D Daw. Reinforcement learning and episodic memory in humans and animals: an integrative framework. Annual review of psychology, 68:101-128, 2017.
Karol Gregor, Danilo Jimenez Rezende, Frederic Besse, Yan Wu, Hamza Merzic, and Aaron van den Oord. Shaping belief states with generative environment models for rl. In Advances in Neural Information Processing Systems, pp. 13475-13487, 2019.
Arthur Guez, Mehdi Mirza, Karol Gregor, Rishabh Kabra, Sébastien Racanière, Théophane Weber, David Raposo, Adam Santoro, Laurent Orseau, Tom Eccles, et al. An investigation of model-free planning. arXiv preprint arXiv: 190110359,9 2019.
Saurabh Gupta, James Davidson, Sergey Levine, Rahul Sukthankar, and Jitendra Malik. Cognitive mapping and planning for visual navigation. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 2-255, 2017.
Izzeddin Gur, UIrich Rückert, Aleksandra Faust, and Dilek Hakkani-Türr Learning to navigate the 2018. URL
web. CoRR,
09195
Wouter Kool, Herke Van Hoof, and Max Welling, Attention, learn to solve routing problems! arXiv preprint arXiv:1800 08475, 2018.
Brenden M Lake, Tomer D Ullman, Joshua B Tenenbaum, and Samuel Gershman. Building machines that learn and think like people. Behavioral and bruin sciences, 40, 2017.
Juho Lee, Yoonho Lee, Jungtaek Kim, Adam R Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention-based permutation invariant neural networks. arXiv preprint arXiv: 2018a.
Lisa Lee, Emilio Parisotto, Devendra Singh Chaplot, Eric Xing, and Ruslan Salakhutdinov. Gated path planning networks. arXiv preprint arXiv: 1o0:.0040600 2018b.
Máté Lengyel and Peter Dayan. Hippocampal contributions to control: the third way. In Advances in neural information processing systems, pp. 889-896, 2008
10

---

Published as a conference paper at ICLR 2021
Piotr Mirowski, Matt Grimes, Mateusz Malinowski, Karl Moritz Hermann, Keith Anderson, Denis Teplyashin, Karen Simonyan, Andrew Zisserman, Raia Hadsell, et al. Learning to navigate in cities without map. In Advances in Neural Information Processing Systems, pp. 2419-2430. 2018.
Piotr Mirowski, Andras Banki-Horvath, Keith Anderson, Denis Teplyashin, Karl Moritz Hermann, Mateusz Malinowski, Matthew Koichi Grimes, Karen Simonyan, Koray Kavukcuoglu, Andrew Zisserman, et al. The streetlearn environment and dataset. arXiv preprint arXiv: WrXiv:9933.0129 2019.
Thomas M Moerland, Joost Broekens, and Catholijn M Jonker. Model-based reinforcement learn- ing: A survey. arXiv preprint 2xiiv:0006.661222 2020.
Junhyuk Oh, Valliappa Chockalingam, Satinder Singh, and Honglak Lee. Control of memory, active perception, and action in minecraft. arXiv preprint arXiv:1605.09/28, 2016.
Alexander Pritzel, Benigno Uria, Sriram Srinivasan, Adria Puigdomenech Badia, Oriol Vinyals, Demis Hassabis, Daan Wierstra, and Charles Blundell. Neural episodic control. In Proceedings of the 34th International Conference on Machine Learning-Volle Volume 70, pp. 2827-2836. JMLR. org, 2017.
Sébastien Racanière, Théophane Weber, David Reichert, Lars Buesing, Arthur Guez, Danilo Jimenez Rezende, Adria Puigdomenech Badia, Oriol Vinyals, Nicolas Heess, Yujia Li, et al. mmagination-augmente agents for deep reinforcement learning. In Advances in neural information processing systems, pp. 5690-5701, 2017.
Samuel Ritter. Mein-Reiifforcement Learning with Episodic Recall: An Integrative Theory of Reward-Driven Learning. PhD thesis, Princeton University, 2019.
Samuel Ritter, Jane Wang, Zeb Kurth-Nelson, and Matthew M Botvinick. Episodic control as meta-reinforrement learning. bioRxiv, pp. 360537, 2018a.
Samuel Ritter, Jane X Wang, Zeb Kurth-Nelson. Siddhant M Jayakumar, Charles Blundell, Razvan Pascanu, and Matthew Botvinick. Been there, done that: Meta-learning with episodic recall. arXiv preprint arXiv: arXiv:1805.0969) 2018b.
Nikolay Savinov, Alexey Dosovitskiy, and Vladlen Koltun. Semi-parametric topological memory for navigation. arXiv preprint
2018
Tom Schaul, Daniel Horgan, Karol Gregor, and David Silver. Universal value function approxima- tors. In International Conference on Machine Learning, pp. 1312-1320, 2015.
Julian Schrittwesser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with learned model. arXiv preprint arXiv: 1.08265, 2019.
Marwin HS Segler, Mike Preuss, and Mark Waller. Planning chemical syntheses with deep neural networks and symbolic ai. Nature, 555(7698):604, 2018.
Richard S Sutton and A Barto. Introduction to reinforcement learning, volume 2. MIT press Cam- bridge, 1998.
Richard S Sutton, Joseph Modayil, Michael Delp. Thomas Degris, Patrick M Pilarski, Adam White, and Doina Precup. Horde: A scalable real-time architecture for learning knowledge from unsuper- vised sensorimotor interaction. In The 10th International Conference on Autonomous Agents and Multiagent Systems- Volume 2, pp. 761-768. International Foundation for Autonomous Agents and Multiagent Systems, 2011.
Aviv Tamar, Yi Wu, Garrett Thomas, Sergey Levine, and Picter Abbeel. Value iteration networks. In Advances in Neural Information Processing Systems, pp. 2154-2162, 2016.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
11

---

Published as a conference paper at ICLR 2021
Oliver Vikbladh, Daphna Shohamy, and Nathaniel Daw. Episodic contributions model-basee reinforcement learning. In Annual conference on cognitive computational neuroscience, CCN, 2017.
Oriol Vinyals, Charles Blundell, Timothy Lillicrap, Daan Wierstra, et al. Matching networks for one shot learning. In Advances in neural information processing systems, pp. 3630-3638, 2016.
JX Wang, Z Kurth-Nelson, D Tirumala, H Soyer, JZ. Leibo, R Munos, C Blundell, D Kumaran, and M Botivnick. Learning reinforcement learn. arxiv 161 1611.05763, 2017.
Greg Wayne, Chia-Chun Hung, David Amos, Mchdi Mirza, Arun Ahuja, Agnieszka Grabska- Barwinska, Jack Rae, Piotr Mirowski, Joel Z Leibo, Adam Santoro, al. Unsupervised predictive memory in a goal-directed agent, arXiv preprint arXiv: 10760, 2018.
12

---

Published as a conference paper at ICLR 2021
A APPENDIX
A.1 ONE-SHOT STREETLEARN
The StreetLeean dataset Mirowski et al. (019) provides a connectivity graph between nodes, and panoramic images (panos) for each node. As movement in the environment is limited to turning left, turning right, or moving forward to the next node, we only need from the panoramic images the frames corresponding when oriented towards next possible location. This defines a mapping from oriented edge in the full StreetLearn graph to single frame. We compute ahead of time and store it in data structure enabling efficient random accesses.
Dataset preparation
what an observer would in a given location,
Observation associated to an edge As our agent operates on reduced graph containing only 2 nodes, we define the observation to an oriented edge A-B in the reduced graph to be the image associated to the first edge in the shortest path from A to on the full graph.
Movement- When the agent moves forward, its new location becomes the neighbor it was oriented towards, and reoriented towards the neighbor most compatible with its current orientation.
Distributed learning setup-
Each actor is assigned to a random city, out of 12 cities/regions in Europe: Amsterdam, Brussels, Dublin, Lisbon, Madrid, Moscow, Rome, Vienna, Warsaw, Paris North-East, Paris North-Was and Paris South-West. A special evaluator actor, which does not send send trajectories to the learner, is assigned to withheld region: Paris South-Fast.
Sampling neighborhoods In the beginning of new episode, we start by sampling a random location in the region assigned to the actor i.e. random node of the region's connectivity graph. This node will become the center-node of the new neighborhood and is added as the first node of the sampled graph. We then proceed to traverse the connectivity graph, breadth first, adding visited nodes to the sampled graph, until it contains the number of intersections required. An intersection node with degree greater than 2; the number of intersections is a parameter of environment that we can set for each experiment and it determines the difficulty (depth) of the planning problem. In our experiments we decided to remove all d nodes of the sampled graph. This manipulation allowed us to simplify the exploration problem substantially without reducing the planning difficulty.
A.2 RL SETUP
We used an actor-critic setup for all RI. experiments reported in this paper, following the distributed training and V-Trace algorithm implementations described by Espeholt et al. (2018). The distributed agent consisted of 1000 actors that produced trajectories of experience on CPU, and single learner running on a Tensor Processing Unit (TPU), which learned a policy π and a baseline using mini- batches of actors' experiences provided via queue. The length of the actors' trajectories is set to the unroll length of the learner. Training was done using the RMSprop optimization algorithm. Please the table below for values of fixed hyperparameters and intervals used for hyperparameter tuning.
Hyperparameter
Values
Agent
Mini-batch size
[32, 128]
Unroll length
[10, 40]
Entropy cost
[1c-3, 1c-2]
Discount
[0.9, 0.95]
RMSprop
Learning rate
[le- 5, 4e-4] le-4
Epsilon
Momentum
0
Decay
0.99
Table 1: Hyperparameter values and tuning intervals used in RL experiments.
13

---

Published as a conference paper at ICLR 2021
A.3 ARCHITECTURE DETAILS
Vision - The visual system of our agents, which produced embeddings for state and goal from the raw pixel inputs, was identical to the one described in Espeholt et al. (2018), except ours was smaller. comprised residual-convolution
layers, each one with a single residual block, instead of two. The number of output channels each layer was 16, 32 and 32. We used smaller final linear layer with 64 units.
Planner For the multi-head attention, queries, keys and values were produced with an embedding size of 64, using to 4 attention heads. In our experiments, we did not observe a significant benefit from using more than a single attention head. The feedforward block of each attention step was 2-layer MLP with 64 units per layer (shared row-wise). Both the self-attention block and the feedforward block were shared across iterations. After appending the state the output of the final attention iteration, used another MLP (shared row-wise) consisting of 2-layers with 64 units. The output of this MLP applied to each row was then aggregated using a max pooling, teature-wise, operation.
Policy network- The input to the policy network was the result of concatenating the output of the planner and the stategoaa embedding pair (which be seen skip connection). We used 2-layer MLP with 64 units per layer followed by a ReLU and two separate linear layers to produced the policy logits (π) and the baseline (V).
A.4 SCALABILITY
We experimented with two variants of the update function described in Section In the first, all of the episodic memories attend to all of the others to produce a tensor with the same dimensions as the episodic memory. This variant, which we refer to as all-to-all (A2A), scales quadratically with the number of memories (and the number of steps per episode). The second version, which scales more favorably, takes the & most recent memories from the episodic memory (M) the initial belief state (B). On each self-attention iteration, each slot B attends to each slot in M, producing a tensor with the same dimensions B (i.e. i vectors), which then self-attend to another to produce the next belief state. The idea behind this design is that the querying from the vectors to the full memory can learn select and summarize information be composed via self-attention. The self-attention among the k vectors may then learn to compute relationships, e.g. state-state paths, using this condensed representation. The benefit that this architecture scales only linearly with the number of memories (N) in the episodic memory, because the quadratic time self-attention need only be applied over the i vectors. This approach is similar to the inducing points of Lee et al. 2018a). Because this scales with N * k rather than 2 we call this variant N-by-k (abbreviated, Nxk).
Figure 5 compares the performance obtained the Seemory&Plamiin Game with the two archi- tecture variants, A2A and Nxk. With set to one half of the original memory capacity (k= the Nxk agent recovers 92% of the performance of the A2A. This result indicates that we can indeed overcome the quadratic complexity and opens up the possibility of applying EPNs to problems with longer timescales.
14

---

Published as a conference paper at ICLR 2021
30
1.0
0.8-
Average teward
Fraidion of oracle
0.8
15-
0.4-
exploration
10-
02-
0.0
(A2A)
EPN (8o5)
0.0
1.0
1.5
2.5
3.0
Treining steps
Figure 5: Comparison between architecture variants. The Nxk architecture variant, which scales linearly with the total number of memories, recovers 92% of the performance of the A2A variant, which scales quadratically.
2.0
Dislens----goll lose
1.0
1.5-
-08
-08
1.0
-04
Herations
filtaratoons
horations
Optmizaton
Figure 6: (a) Distance-to-gun decoding loss from the output the planner (with frozen weights) after 1 to iterations. Each additional iteration improves distance decoding. (b) The ability of EPNs infer distance to goal expands with the number self-attention iterations, even when that number goes beyond the number of iterations used during training of the EPN.
15


Reviews:
[{"review_id": "lDwOWeDVmAd", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 5, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper is presenting a Transformer-based architecture for episodic memory-based sequential decision making.\nThe authors are interested in solving RTS and perform experiments that are focused on the case of mapless navigation.\nThe paper is insisting on the capacity to seamlessly coupling the exploration and planning parts of the overall task.\nThe authors are using a very pertinent dataset of google street maps for mapless navigation but seem to basically claim the superiority of Transformer against LSTM based policy in this particular setting which does not seem surprising.\nA large quantity of work has been recently done on this task [1] [2] [3]\nOne remark is that the planning phase of the algorithm is completely embedded into the model which can be problematic depending on the nature of the task, especially regarding the possible need for decision explainability.\nSo, beyond parametric variation of tasks, as presented here, a good improvement could be to evaluate this approach on non-parametric variations of tasks as proposed in the manipulation framework of the Stanford meta world [4].\n\nRefs\n[1] Gated Path Planning Networks, 2018, Salakhutdinov and al\n[2] LEARNING TO EXPLORE USING ACTIVE NEURAL SLAM, 2019, Salakhutdinov and al\n[3] Semantic Curiosity for Active Visual Learning, 2020, Chaplot and al\n[4] Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning, 2019, Levine and al\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "jv82-mHFV8-", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors propose a non-parametric memory based on the transformer/self-attention architecture to learn over tasks that require planning from previously experienced tasks. They describe this style of learning as a form of meta reinforcement learning where individual episodes are a collection of tasks in the same environment and individual environments are sampled from a 'meta' environment. This is very similar in spirit to the RL-squared framework of Duan et al. 2017.\n\nAs is typical of meta-learning environments, there is an outer loop and an inner loop. The outer loop is composed of environments sampled from a 'meta' environment and within each sampled environment are tasks sampled with different initial states and goals within that environment. \n\nUnlike the baselines in this paper, the proposed method utilizes a non-parametric memory that stores all (action, prev. state, observation) tuples that were experienced within an episode. This allows a model to leverage experiences from previous tasks to accomplish the current goal. The unique contribution presented in this work lies in using self-attention to integrate these past experiences.\n\nFurthermore, results are demonstrated on a toy environment and a street view environment to demonstrate scalability to harder tasks.\n\nThe authors convincingly show that the proposed episodic memory architecture:\n\n- Learns from experiences from past tasks by improving as time progresses within an episode\n- Learning is improved when the self-attention memory is allowed more recurrent iterations\n- Inner-loop learning is accomplished in the absence of gradient updates\n\nFurther improvements:\n\n-Are multiple episodes present within each minibatch or are parameter updates computed from minibatches extracted from single episodes? Is this method sensitive to how the episodes/tasks are presented during training? In many settings, it is not be possible to have all environments available to sample from at any given training iteration due to limited computation or data resources.\n\n-There is little to no discussion about the computational resources necessary as compared to the baseline methods described in this paper. The authors do provide some information in the supplementary, but these tradeoffs are likely significant enough to warrant some discussion in the main text. There is clearly a computational complexity difference between the self-attention mechanism and baseline models like the LSTM.\n\n-How important is tuning the timescale of the environment in this work? How does the computational complexity and model performance scale with course and finer grains of timescale?\n\nDuan, Yan, et al. \"Rl $^ 2$: Fast reinforcement learning via slow reinforcement learning.\" arXiv preprint arXiv:1611.02779 (2016).", "extra_scores": {}, "extra_text": {}}, {"review_id": "jpd-SctUbl9", "reviewer": "AnonReviewer1", "summary": "", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposes 2 new benchmarks for Rapid Task Solving (RTS), that evaluate RL agents on the ability to memorize past experiences and learn to plan to solve new tasks in different environments rapidly. The paper also proposes Episodic Planning Networks (EPN), an RL method that replaces a weighted sum and multi-layer peceptrons in memory networks with self-attention.  The proposed EPNs proved to significantly outperform baseline methods with memory and an LSTM method without memory. The paper formulates RTS as an extension of meta-reinforcement learning framework, where in addition to optimizing over a distribution of tasks, the objective is to also optimize over a distribution over environments. \n\nI think the work proposed in the paper is an important step towards developing general agents that learn to adapt quickly. To that end, it would be beneficial for the authors to release their code to benchmark new and older methods.\n\nI have a few clarification questions.\n\nHow is the reward defined for each task in the 2 benchmarks? Is it 1 for reaching the goal and 0 otherwise?\nGiven that the 2 domains are defined in terms of relations between entities (connections between symbols in the Memory&Planning game and neighborhoods in the other) how would relational and symbolic RL methods perfom compared to the baselines and the proposed EPN? I believe these would provide a fairer baseline performance compared to the LSTM used.  ", "extra_scores": {}, "extra_text": {}}, {"review_id": "52LtkNAEqQG", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary\n\nThe authors present the challenge of rapid task-solving (RTS), where the goal is to solve a series of tasks as quickly as possible in a new, shared environment. This challenge requires both memory and planning capabilities from an agent, and the authors demonstrate that current SOTA modern deep RL agents with memory components still fail at this challenge. The authors instead propose “episodic planning networks” (EPNs), which contain architectural changes that allow for better long-term planning, and demonstrate their success on two challenges: a memory & planning game, and a navigation task from Street View images. Finally, the authors demonstrate that their proposed EPNs intuitively learn a planning algorithm similar to BFS, hinting at their generalization ability.\n\n---\n\n### Strengths\n\n- The authors present a strong motivating example and an interesting challenge of rapid task-solving that seems both understudied and aspirational for current AI systems.\n- The authors develop challenging but useful benchmark problems to compare their proposed method against, including the simpler planning game, and the more realistic and practical navigation task.\n- Along with these benchmarks, the authors develop an interesting set of metrics, including the average training reward and the fraction of the oracle’s performance near the end of an episode, although the details of how these metrics were chosen and not others was less well-explained.\n- The ablation experiment on how performance changes when replacing the iterative attention with a single pass is extremely thorough and well reasoned, along with showing strong empirical evidence for the authors’ claims.\n- The authors clearly motivate why EPNs are necessary by demonstrating that current algorithms fail at the planning component, and that EPNs are able to optimally navigate and generalize to larger maps.\n- The authors’ insight that EPNs recursively refine a reachability map, similar to value iteration, is an exciting insight into how compute can be scaled at test-time with the same training setting for better generalization.\n- The qualitative visualization in Figure 4d of how the recursive refinement works is very convincing and easy to interpret.\n\n---\n\n### Weaknesses\n\n- Some claims the authors make are only partially or weakly supported. For instance, the authors hypothesize that the two main reasons prior techniques fail is because the weighted sum of retrieved slots is not representationally expressive, and that the MLP is not computationally expressive to support planning. Are there ways to test these two hypotheses, ideally in isolation, potentially with ablative experiments?\n- The description of the architecture in section 4 is hard to follow, although Figure 1b is very helpful. Can section 4 be formalized in terms of the specific modules and how they’re parameterized? \n- I have some confusion around the metrics used to evaluate, specifically how these ones are chosen over other potential choices. For instance, the choice of 37% as the optimal exploration rate within each task, and the definition of oracle performance for the last third of an episode are not clearly explained.\n- The authors note that the self attention network does not have access to the current state, only the goal state and episodic memories. This isn’t clear, as in diagram 1b it looks like the episodic storage has previous states as well as the current observed state, with the goal appended, and that the self-attention is happening across this entire vector. Some clarification around which keys and the features of those keys the self-attention is attending to would be helpful.\n- Along the same note, the authors’ motivation for the self attention network not having access to the current state is that it would learn to compute a value map. How does performance change with a different design choice here, either with the goal state only available at the end and current state throughout, or have both the goal and current states available throughout the iterative self-attention?\n- The authors also note that if the self-attention network did represent a reachability map, then a simple MLP on this map and the current state would be sufficient. Would a simple MLP on the map alone, not with the current state, also be sufficient? Specifically, if the map has information about all states, and the map states that the goal is not reachable from anywhere else in the map, would that not imply that the goal is reachable from the current location, even if the current location’s state is unknown?\n- The authors claim that the specific self-attention function used was crucial. Ideally the authors would present ablations on this design choice, as the one they used looks similar to that used in transformers, with multi-headed attention and layernorm. \n- The performance measure reported by the EPNs seems a bit suspect. Specifically, the baselines get to 5-10 tasks, but EPN gets to 28.7, where the median neighborhood has 12 locations. Does this mean that the EPN performance is with repeated starting positions or repeated goals? How novel are the latter set of tasks presented?\n- The authors note in figure 4a that performance is roughly correlated with the number of iterations on self-attention. However, the authors do not explain why performance with 2 iterations converges to roughly the same final reward as with 4 iterations, and that 2 iterations converges more quickly. Further, it is a bit confusing that the performance is correlated with number of self-attention iterations for generalization (4b), and the distance to goal metric (4c), but that none of these are reflected in the average reward experiment (4a). \n- A reference is made to figure 5a at the end of section 5 on the progressive removal of planning iterations -  is this meant to be figure 4a, or a different figure?\n- In my mind, the comparison to Merlin, LSTM, and MRA serve as ablation studies. What would other baselines look like? Are there other techniques in meta-learning that could be applied here without an explicit architecture design of episodic memory or a planner? I would expect there to be several methods that could be applied to this setting - I don’t have a clear sense of what is SOTA in episodic model-based control, but would want to see those comparisons. As an example, “Explore then Execute: Adapting without Rewards via Factorized Meta-Reinforcement Learning” (DREAM) may have similar properties that make it a good comparison.\n\n### Recommendation\n\nOverall, I vote for accepting this paper. It demonstrates strong empirical results compared to previous SOTA in episodic model-based control. I do have some additional clarification questions on the experiments, and I wish the architecture section was more formalized. My main concern about this work is that it is hard to place in context with other potential baselines or areas of research, for example meta learning. Since I am not an expert on episodic model-based control, maybe this comparison is not realistic however.\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes the challenge of rapid task-solving in unfamiliar environments and presents an approach to achieve this called Episodic Planning Networks -- a non-parametric memory based on the transformer architecture to learn tasks that require planning from previously experienced tasks, following a form of meta-RL.  The problem and approach are compelling, with strong empirical results.  The paper is well-written and is an exciting contribution.  This is a clear accept.

In response to the initial reviews, the authors updated their paper to improve the formalization and address other concerns in the reviews, which were viewed favorably by the reviewers as a good improvement. Based on the reviewer discussions, the work could still be placed better in context with respect to other literature.
Decision: Accept (Poster)