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