R APID T ASK -S OLVING IN N OVEL E NVIRONMENTS
Anonymous authors
Paper under double-blind review
A BSTRACT
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.
I NTRODUCTION
An ideal AI system would be useful immediately upon deployment in a new environment, and would become more useful as it gained experience there.
Consider for example a household robot deployed in 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 use general knowledge about household layouts to ﬁnd 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 it passes.
When faced with its next task, say, doing the laundry, it would use its newfound knowledge of the hamper locations to efﬁciently collect the laundry.
Humans make this kind of rapid task-solving in novel environments (RTS) look easy (Lake et al., 2017), but as yet it remains an aspiration for AI.
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 is in part because they treat each new environment in isolation, relying on generic exploration and planning algorithms.
We propose to overcome this limitation by treating RTS as a meta-reinforcement learning (RL) problem, where agents learn exploration policies and planning algorithms from a distribution over
RTS challenges.
Our contributions are to:
Develop two domains for studying meta-learned RTS: the minimal and interpretable Memory&Planning Game and the scaled-up One-Shot StreetLearn. 2.
Show that previous meta-RL agents fail at RTS because of limitations in their ability to plan using recently gathered information.
Design a new architecture – the Episodic Planning Network (EPN) – that overcomes this limitation, widely outperforming prior agents across both domains.
Show that EPNs learn exploration and planning algorithms that generalize to larger problems than those seen in training.
Demonstrate that EPNs learn a value iteration-like planning algorithm that iteratively propagates information about state-state connectivity outward from the goal state.
[Figure]
P ROBLEM F ORMULATION
Our objective is to build agents that can maximize reward over a sequence of tasks in a novel environment.
Our basic approach is to have agents learn to do this through exposure to distributions over multi-task environments.
To deﬁne such a distribution, we ﬁrst formalize the notion of an environment e as a 4-tuple ( S, A, P a , R ) consisting of states, actions, a state-action transition function, and a distribution over reward functions.
We then deﬁne the notion of a task t in environment e as a
Markov decision process (MDP) ( S, A, P a , r ) that results from sampling a reward function r from
R .
We can now deﬁne a framework for learning to solve tasks in novel environments as a simple generalization of the popular meta-RL framework (Wang et al., 2017; Duan et al., 2016).
In meta-RL, the agent is trained on MDPs ( S, A, P a , r ) sampled from a task distribution D .
In the rapid task-solving in novel environments (RTS) paradigm, we instead sample problems by ﬁrst sampling an environment e from an environment distribution E , then sequentially sampling tasks, i.e. MDPs, from that environment’s reward function distribution (see Figure 1).
An agent can be trained for RTS by maximizing the following objective:
[Equation]
where J e,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 metaRL objective E e ∼ E
[ J e ( θ )] .
RTS can be viewed as meta-RL with the added complication that the reward function changes within the inner learning loop.
While meta-RL formalizes the problem of learning to learn, RTS formalizes both the problems of learning to learn and learning to plan.
To maximize reward while the reward function is constantly changing in non-trivial novel environments, agents must learn to (1) efﬁciently explore and effectively encode the information discovered during exploration (i.e., learn ) and (2) use that encoded knowledge to select actions by predicting trajectories it has never experienced (i.e., plan 1 ). [Figure]
T HE L IMITATIONS OF P RIOR A GENTS
To test whether past deep-RL agents can explore and plan in novel environments we introduce a 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 2 , wherein players must remember the locations of cards in a grid.
The variation we propose, which we call the Memory&Planning Game, extends the challenge to require planning as well as remembering (see Figure 2).
In the Memory&Planning Game, the agent occupies an environment consisting of a grid of symbols.
The agent can see one symbol which corresponds to the agent’s current location, and it sees a “goal” symbol, which it is tasked with navigating to.
The agent can not see its relative location with respect to other symbols in the grid.
The agent’s actions are: move left, move right, move up, move down, and choose.
The agent receives reward when it takes the “choose” action while its current location matches the goal location.
At the beginning of each episode, a new set of symbols is sampled, effectively inducing a new transition function.
Each time the agent ﬁnds a goal – which corresponds to completing a task –, a new goal is sampled in and the transition function stays ﬁxed.
A successful agent will (1) efﬁciently explore the grid to discover the current set of symbols and their connectivity and (2) plan shortest paths to goal symbol if it has seen (or can infer) all of the transitions it needs to connect its current location and the current goal location.
This setup supplies a minimal case of the RTS problem: at the beginning of each episode the agent faces a new transition structure ( S, A, P a ) and must solve the current task while ﬁnding information that will be useful for solving future tasks on that same transition structure.
In subsequent tasks, the agent must use its stored knowledge of the current grid’s connectivity to plan shortest paths.
In this symbolic domain, we ﬁnd evidence that previous agents learn to explore but not to plan.
Speciﬁcally, they match the within-trial planning optimum; that is, a strategy that explores optimally within each task, but forgets everything about the current environment the task ends (Figure [Figure] A v e r age r e w a r d Merlin
EPN
Oracle LSTM MRA Merlin EPN LSTM
MRA
[Figure]
intersections
[Figure]
intersections [Figure] ) c) [Figure]
[Figure]
We hypothesize that this failure is the result of the limited expressiveness of past architectures’ mechanisms for processing retrieved memories.
We designed Episodic Planning Networks to overcome this limitation.
E PISODIC P LANNING N ETWORKS
Past episodic memory agents used their memory stores by querying the memories, summing the retrieved slots, then projecting the result through a multi-layered perceptron (MLP) (Fortunato et al., 2019; Wayne et al., 2018).
We hypothesize that these agents fail to plan because the weighted sum of retrieved slots is not sufﬁciently representationally expressive, and the MLP is not sufﬁciently computationally expressive to support planning in non-trivial environments.
To test this hypothesis, we replace the weighted sum and MLP with an an iterative self-attention-based architecture designed to support implicit planning.
The architecture, which we call the Episodic Planning Network (EPN, see Figure 1b), begins with a set M of episodic memories which reﬂect the agent’s experience in the episode so far.
In the present experiments, the agent stores a memory in a new slot on each timestep.
The memory represents that timestep’s transition: it is the concatenation of embeddings of the current observation, the previous action, and the previous observation.
The EPN appends an embedding of the current goal g to each slot in M , producing M g .
Then, a self-attention-based update function S is applied to produce a processed representation reﬂecting the agent’s belief about the environment B g 1 .
S is iterated some number of times u (sharing weights) to produce B gu .
Finally, the current state c is appended to each slot in B gu to produce B gcu , and then each resulting slot is passed through a shared MLP.
The resulting vectors are aggregated by a feature-wise max operation, then fed to the policy network to inﬂuence behavior.
Note that the self-attention network does not have access to the current state until the very end – it does all of the self-attention updates with access only to the goal state and episodic memories.
This design reﬂects the intuition that the function S might learn to compute something like a value map, which represents the distance to goal from all states in memory.
If the ﬁnal self-attention state
B gu does in fact come to represent such a reachability map, then the simple MLP over B gcu should be sufﬁcient to compute the action from c that leads to the nearest/highest value state with respect g .
We ﬁnd in Section 5, Figure 4 evidence that B g does in fact come to resemble an iteratively improved reachability map.
The speciﬁc self-attention function we used was crucial for full performance.
In all references above to self-attention, the speciﬁc update function we used was:
[Equation]
where MHA is the multi-head dot-product attention described by Vaswani et al.
In our experiments, f was a ReLU followed by a 2-layer MLP shared across rows.
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 in the Memory&Planning
Game.
It matches the shortest-path optimum after a few tasks (Figure 2c), 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 2d,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 to produce B g with same dimensionality as M . This update function scales quadratically with the number of memories (and the number of steps per episode).
In the Supplemental Material we present results showing that we recover 92% of the performance using a variant of the update function that scales linearly with the size of the memory (see Figure 5 in the Suppl. Material).
O NE -S HOT S TREET L EARN
We now test whether our agent can extend its success in the Memory&Planning Game to a 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 1a), wherein environments are sampled as neighborhoods from the StreetLearn dataset of Google
StreetView images and their connectivity (Mirowski et al., 2019).
Tasks are then sampled by selecting a position and orientation that the agent must navigate to from its current location.
In past work with StreetLearn, agents were faced with a single large map (e.g. a city), in which they could learn to navigate over billions of experiences (Mirowski et al., 2018).
StreetLearn city maps are fully connected graphs where nodes are generated from point-wise samples every few meters or so in the given city.
For each node and neighbour pair there is a 64 × 64 RGB pixel input unique to that point-wise sample and direction of gaze - this pixel representation is the observation state of the environment.
We partition the StreetLearn data into a 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 is analogous to the partitioning of the ImageNet data into many one-shot classiﬁcation 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 to plan in one shot, that is, to plan after one or fewer observations of each transition in the environment.
In One-Shot StreetLearn, the agent’s observations consist of an image representing the current state and an image representing the goal state (Figure 1).
The agent receives no other information from the environment.
The agent’s actions are “turn left”, which orients the agent toward the next available direction of motion from its current node; “turn right”, which does the same in the other direction; and “move forward”, which moves the agent along the direction it’s facing to the next available node.
One-Shot StreetLearn has several important characteristics for developing novel-environment planning 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 unlimited maps with real-world connectivity structure for agents to learn to take advantage of.
Second, the agent’s observations are rich, high-dimensional visual inputs that simulate a person’s experience when navigating.
These inputs can in principle be used by an agent to infer some of the structure of a new environment.
Third, we 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 can be varied independently of the planning depth by varying the number of nodes between intersections.
This is a useful feature for developing temporally abstract (i.e. jumpy) planning abilities (see Section 7 for discussion).
[Figure]
In the current experiments, to reduce the exploration problem while keeping the planning difﬁculty constant, we removed all locations between intersections that corresponded to degree-2 nodes.
This resulted in neighborhoods with a median of 12 nodes (see example connectivity graphs in
Figure 3b).
EPNs Outperform All Baselines: We trained EPN-based agents and previous deep-RL agents on
One-Shot StreetLearn neighborhoods with 5 intersections.
Our agent signiﬁcantly outperformed the baselines, successfully reaching 28.7 goals per episode (averaged over 100 consecutive episodes, Figure 3c).
It is important to note that this performance was measured on neighborhoods of a held-out city, which the agent was visiting for the ﬁrst 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 n th 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 collect each goal after 10–15 tasks (Figure 3e). In the last third of every episode, EPNs achieve on average ( 88% ± 2% ) of the performance of oracle with perfect information (Figure 3d).
EPNs Generalize to Larger Maps: We trained EPN agents on neighborhoods with 5 intersections and evaluated them (with gradient learning turned off) on neigbhoroods with 7 and 9 intersections.
We found that the trained agents achieved strong performance relative to the oracle, suffering only small marginal drops in performance even on 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 evaluation performance.
We observed a systematic performance boost with each additional iteration used in training 3 (see Figure 4a).
Trained EPNs recursively reﬁne a value iteration-like reachability 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) to 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 ﬁxed number of iterations, from 1 to 6.
The training loss of the different decoders revealed a steady increase in the ability to compute distance to goal as we increase the number of self-attention iterations from 1 to 6 (see Figure 6a in the Suppl.
Material).
This gain is a direct consequence of improved decoding for longer distances, as made evident by the gradual rightward shift in the drop-off of the classiﬁcation accuracy plotted against distance to goal (Figure 4c).
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, a neighborhood with 12 intersections) with a ﬁxed 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 4d and Figure 6b in the Suppl. Material).
Ablations: MERLIN/MRA 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 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 5a).
From this experiment we learned that one step of self-attention isn’t enough: iterating the planning module is necessary for full performance.  R ELATED W ORK Deriving repurposable environment knowledge to solve arbitrary tasks has been a long-standing goal in the ﬁeld of RL (Foster & Dayan, 2002; Dayan, 1993) toward which progress has recently been made (Sutton et al., 2011; Schaul et al., 2015; Borsa et al., 2018).
While quick to accommodate new tasks, these approaches require large amounts of time and experience to learn about new environments.
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 in 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 et al. (2017) and Duan et al. 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. The effectiveness of learned, implicit planning was demonstrated further in fullyobserved spatial environments using grid-structured memory and convolutional update functions (Tamar et al., 2016; Lee et al., 2018b; Guez et al., 2019).
Gupta et al.
(2017) extended this approach to partially observed environments with ground-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-observed environments by learning to gather the information needed for planning, and (2) providing an architecture appropriate for a 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 et al.
However, we are not aware of any MBRL agents that learn a deployable model within a single episode of experience , as is required for the RTS setting.
a supervised setting using 1 to 6 iterations, the weights of the planner were trained once in the RL setting using 4 iterations.
Learning to plan can be seen as a speciﬁc 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).
Especially 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-intensive industrial solvers.
This result suggests that in higher-depth novel-environment planning problems than we consider in our current experiments, the transformer-based architecture may continue to be effective.
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 planning algorithms do not provide the beneﬁts of domain-adapted planning demonstrated by Kool et al.
(2018), and that design modiﬁcations 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; Fortunato et al., 2019).
We chose episodic memory because of the following desirable properties.
First, because the episodic store is non-parametric, it can grow arbitrarily with the complexity of the environment, and approximate k-nearest neighbors method make it possible to scale to massive episodic memories in practice, as in Pritzel et al.
This means that memory ﬁdelity need not decay with time.
Second, episodic memory imposes no particular assumptions about the environment’s stucture, making it a potentially appropriate choice for a variety of non-spatial applications such as chemical synthesis Segler et al.
and web navigation Gur et al.
(2018), as well as abstract planning problems of interest in AI research, such as Gregor et al.’s Voxel environment tasks. Our approach can be seen as implementing episodic model-based control (EMBC), a concept recently developed in cognitive neuroscience (Vikbladh et al., 2017; Ritter, 2019). While episodic control (Gershman & Daw, 2017; Lengyel & Dayan, 2008; Blundell et al., 2016) produces value estimates using memories of individual experiences in a model-free manner, EMBC uses episodic memories to inform a model that predicts the outcomes of actions.
Ritter et al.
(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., 2017). 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.
We showed that deep-RL agents with EPNs can meta-learn to explore, build models on-the-ﬂy, and plan over those models, enabling them to rapidly solve sequences of tasks in unfamiliar environments.
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`ere et al., 2017).
Training and evaluating EPNs on problems that beneﬁt from jumpy planning, such as One-Shot StreetLearn with all intermediate nodes, may be enough to obtain strong temporal abstraction performance.
Planning over belief states may be essential for agents to succeed in dynamic partially-observed environments. Planning over belief states might be accomplished simply by storing belief states such as those developed by Gregor et al. and training on a problem distribution that requires planning over information stored in those states.
Humans are able to solve a 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 et al., 2019).
EARN
Dataset preparation – The StreetLearn dataset Mirowski et al.
(2019) 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 to what an observer would see in a given location, when oriented towards a next possible location.
This deﬁnes a mapping from oriented edge in the full StreetLearn graph to a single frame.
We compute it ahead of time and store it in a data structure enabling efﬁcient random accesses.
Observation associated to an edge – As our agent operates on a reduced graph containing only degree-2 nodes, we deﬁne the observation to an oriented edge A–B in the reduced graph to be the image associated to the ﬁrst edge in the shortest path from A to B on the full graph.
Movement – When the agent moves forward, its new location becomes the neighbor it was oriented towards, and is 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-West and
Paris South-West.
A special evaluator actor, which does not send send trajectories to the learner, is assigned to a withheld region: Paris South-East.
Sampling neighborhoods – In the beginning of a new episode, we start by sampling a random location in the region assigned to the actor – i.e. a random node of the region’s connectivity graph.
This node will become the center-node of the new neighborhood and is added as the ﬁrst node of the sampled graph.
We then proceed to traverse the connectivity graph, breadth ﬁrst, adding visited nodes to the sampled graph, until it contains the number of intersections required.
An intersection is a 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 difﬁculty (depth) of the planning problem.
In our experiments we decided to remove all degree-2 nodes of the sampled graph.
This manipulation allowed us to simplify the exploration problem substantially without reducing the planning difﬁculty.
A. RL SETUP
We used an actor-critic setup for all RL experiments reported in this paper, following the distributed training and V-Trace algorithm implementations described by Espeholt et al.
The distributed agent consisted of 1000 actors that produced trajectories of experience on CPU, and a single learner running on a Tensor Processing Unit (TPU), which learned a policy π and a baseline V using minibatches of actors’ experiences provided via a 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 see the table below for values of ﬁxed hyperparameters and intervals used for hyperparameter tuning.
[Table]
A. A RCHITECTURE 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.
It comprised 3 residual-convolutional layers, each one with a single residual block, instead of two.
The number of output channels in each layer was 16, 32 and 32.
We used a smaller ﬁnal linear layer with 64 units.
Planner – For the multi-head attention, queries, keys and values were produced with an embedding size of 64, using 1 to 4 attention heads.
In our experiments, we did not observe a signiﬁcant beneﬁt from using more than a single attention head.
The feedforward block of each attention step was a 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 to the output of the ﬁnal attention iteration, we 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, feature-wise, operation.
Policy network – The input to the policy network was the result of concatenating the output of the planner and the state-goal embedding pair (which can be seen as a skip connection).
We used a 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 ).
We experimented with two variants of the update function S described in Section 4. In the ﬁrst, all of the episodic memories attend to all of the others to produce B g with same dimensionality as M . 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 k most recent memories as the initial belief state B g 0 . On each iteration of S , each slot in B gi attends to each slot in M g , producing | B g | vectors (i.e. k vectors), which then self-attend to one another to produce the next belief state B gi +1 . The idea behind this design is that the query from the k vectors to the full memory can learn to select and summarize information to be composed via self-attention. The self-attention among the k vectors may then learn to compute relationships, e.g. state-state reachability, using this condensed representation.
The beneﬁt is that this architecture scales only linearly with the size of the full memory N, because the quadratic time self-attention need only be applied over the k vectors.
This approach is similar to the inducing points of Lee et al.
(2018a). Because this scales with N ∗ k rather than N 2 , we call this variant N-by-k (abbreviated, Nxk). Figure 5 compares the performance obtained on the Memory&Planning Game with the two architecture variants, A2A and Nxk.
With k set to one half of the original memory capacity ( k =50 ), 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.
[Figure]
[Figure]
e
[Figure]
t
[Figure]