Rapid Task-Solving in Novel Environments

1 INTRODUCTION . 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 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 it 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. , 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 : 1 . 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 . ∗Equal contribution . 3 . Design a 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 problems than those seen in training . 5 . Demonstrate that EPNs learn a value iteration-like planning algorithm that iteratively propagates information about state-state connectivity outward from the goal state . 2 PROBLEM FORMULATION . 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 define such a distribution , we first formalize the notion of an environment e as a 4-tuple ( S , A , Pa , R ) consisting of states , actions , a state-action transition function , and a distribution over reward functions . We then define the notion of a task t in environment e as a Markov decision process ( MDP ) ( S , A , Pa , r ) that results from sampling a reward function r from R. We can now define 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 , Pa , r ) sampled from a task distributionD . In the rapid task-solving in novel environments ( RTS ) paradigm , we instead sample problems by first 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 : Ee∼E [ Er∼Re [ Je , r ( θ ) ] ] , 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 metaRL objective Ee∼E [ Je ( θ ) ] . 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 ) efficiently 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. , plan1 ) . 3 THE LIMITATIONS OF PRIOR AGENTS . 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 Game2 , 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 ( 4×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 5 possible actions : move left , move right , move up , move down , and collect . If the agent chooses the ” collect ” action when its current location symbol matches the goal symbol , a reward of 1 is received . Otherwise , the agent receives a reward of 0 . At the beginning of each episode , a new set of symbols is sampled , effectively inducing a new transition function . The agent 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 to completing a task – , a new goal is sampled in and the transition function stays fixed . 1Following Sutton & Barto ( 1998 ) , we refer to the ability to choose actions based on predictions of not-yetexperienced events as “ planning ” . A successful agent will ( 1 ) efficiently 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 , Pa ) and must solve the current task while finding 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 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 optimally 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 architectures ’ mechanisms for processing retrieved memories . We designed Episodic Planning Networks to overcome this limitation . 4 EPISODIC PLANNING NETWORKS . 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 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 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 ) 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 ) a policy network . 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 as input to the planner . We then add a new 2In a striking parallel , the Memory Game played a key role in the development of memory abilities in early episodic memory deep-RL agents ( Wayne et al. , 2018 ) . 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-based update function is applied over the memories to produce a processed 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 selfattention , each of which is then passed through a shared MLP . The resulting vectors are aggregated by a feature-wise max operation . Note that the self-attention ( SA ) does not 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 a 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 5 , Figure 4 evidence that the final SA state comes in fact to resemble an iteratively improved value map . The specific self-attention update function we used was : Bi+1 = f ( Bi + φ ( Bi ) ) φ ( x ) = MHA ( LayerNorm ( x ) ) where MHA is the multi-head dot-product attention described by Vaswani et al . ( 2017 ) . Layer normalization was applied to the inputs of each iteration , and residual connections were added around the multi-head attention operation.3 In our experiments , f was a ReLU followed by a 2-layer MLP shared across rows . Policy network – A 2-layer multi-layer perceptron ( MLP ) was applied to the output of the planner to produce policy logits and a baseline . When using pixel inputs , embeddings of the observations were produced by a 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 in the Memory & Planning Game . It matches the shortest-path optimum after only a few tasks in a novel environment ( 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 . 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 ) .

Summary: 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.