Decoupling Exploration and Exploitation for Meta-Reinforcement Learning without Sacrifices

1 INTRODUCTION . A general-purpose agent should be able to perform multiple related tasks across multiple related environments . Our goal is to develop agents that can perform a variety of tasks in novel environments , based on previous experience and only a small amount of experience in the new environment . For example , we may want a robot to cook a meal ( a new task ) in a new kitchen ( the environment ) after it has learned to cook other meals in other kitchens . To adapt to a new kitchen , the robot must both explore to find the ingredients , and use this information to cook . Existing meta-reinforcement learning ( meta-RL ) methods can adapt to new tasks and environments , but , as we identify in this work , struggle when adaptation requires complex exploration . In the meta-RL setting , the agent is presented with a set of meta-training problems , each in an environment ( e.g. , a kitchen ) with some task ( e.g. , make pizza ) ; at meta-test time , the agent is given a new , but related environment and task . It is allowed to gather information in a few initial ( exploration ) episodes , and its goal is to then maximize returns on all subsequent ( exploitation ) episodes , using this information . A common meta-RL approach is to learn to explore and exploit end-to-end by training a policy and updating exploration behavior based on how well the policy later exploits using the information discovered from exploration ( Duan et al. , 2016 ; Wang et al. , 2016a ; Stadie et al. , 2018 ; Zintgraf et al. , 2019 ; Humplik et al. , 2019 ) . With enough model capacity , such approaches can express optimal exploration and exploitation , but they create a chicken-and-egg problem that leads to bad local optima and poor sample efficiency : Learning to explore requires good exploitation to gauge the exploration ’ s utility , but learning to exploit requires information gathered via exploration ; therefore , with only final performance as signal , one can not be learned without already having learned the other . For example , a robot chef is only incentivized to explore and find the ingredients if it already knows how to cook , but the robot can only learn to cook if it can already find the ingredients by exploration . To avoid the chicken-and-egg problem , we propose to optimize separate objectives for exploration and exploitation by leveraging the problem ID—an easy-to-provide unique one-hot for each training meta- 1Project web page : https : //anonymouspapersubmission.github.io/dream/ training task and environment . Such a problem ID can be realistically available in real-world meta-RL tasks : e.g. , in a robot chef factory , each training kitchen ( problem ) can be easily assigned a unique ID , and in a recommendation system that provides tailored recommendations to each user , each user ( problem ) is typically identified by a unique username . Some prior works ( Humplik et al. , 2019 ; Kamienny et al. , 2020 ) also use these problem IDs , but not in a way that avoids the chicken-and-egg problem . Others ( Rakelly et al. , 2019 ; Zhou et al. , 2019b ; Gupta et al. , 2018 ; Gurumurthy et al. , 2019 ; Zhang et al. , 2020 ) also optimize separate objectives , but their exploration objectives learn suboptimal policies that gather task-irrelevant information ( e.g. , the color of the walls ) . Instead , we propose an exploitation objective that automatically identifies task-relevant information , and an exploration objective to recover only this information . We learn an exploitation policy without the need for exploration , by conditioning on a learned representation of the problem ID , which provides taskrelevant information ( e.g. , by memorizing the locations of the ingredients for each ID / kitchen ) . We also apply an information bottleneck to this representation to encourage discarding of any information not required by the exploitation policy ( i.e. , task-irrelevant information ) . Then , we learn an exploration policy to only discover task-relevant information by training it to produce trajectories containing the same information as the learned ID representation ( Section 4 ) . Crucially , unlike prior work , we prove that our separate objectives are consistent : optimizing them yields optimal exploration and exploitation , assuming expressive-enough policy classes and enough meta-training data ( Section 5.1 ) . Overall , we present two core contributions : ( i ) we articulate and formalize a chicken-and-egg coupling problem between optimizing exploration and exploitation in meta-RL ( Section 4.1 ) ; and ( ii ) we overcome this with a consistent decoupled approach , called DREAM : Decoupling exploRation and ExploitAtion in Meta-RL ( Section 4.2 ) . Theoretically , in a simple tabular example , we show that addressing the coupling problem with DREAM provably improves sample complexity over existing end-to-end approaches by a factor exponential in the horizon ( Section 5 ) . Empirically , we stress test DREAM ’ s ability to learn sophisticated exploration strategies on 3 challenging , didactic benchmarks and a sparse-reward 3D visual navigation benchmark . On these , DREAM learns to optimally explore and exploit , achieving 90 % higher returns than existing state-of-the-art approaches ( PEARL , E-RL2 , IMPORT , VARIBAD ) , which struggle to learn an effective exploration strategy ( Section 6 ) . 2 RELATED WORK . We draw on a long line of work on learning to adapt to related tasks ( Schmidhuber , 1987 ; Thrun & Pratt , 2012 ; Naik & Mammone , 1992 ; Bengio et al. , 1991 ; 1992 ; Hochreiter et al. , 2001 ; Andrychowicz et al. , 2016 ; Santoro et al. , 2016 ) . Many meta-RL works focus on adapting efficiently to a new task from few samples without optimizing the sample collection process , via updating the policy parameters ( Finn et al. , 2017 ; Agarwal et al. , 2019 ; Yang et al. , 2019 ; Houthooft et al. , 2018 ; Mendonca et al. , 2019 ) , learning a model ( Nagabandi et al. , 2018 ; Sæmundsson et al. , 2018 ; Hiraoka et al. , 2020 ) , multi-task learning ( Fakoor et al. , 2019 ) , or leveraging demonstrations ( Zhou et al. , 2019a ) . In contrast , we focus on problems where targeted exploration is critical for few-shot adaptation . Approaches that specifically explore to obtain the most informative samples fall into two main categories : end-to-end and decoupled approaches . End-to-end approaches optimize exploration and exploitation end-to-end by updating exploration behavior from returns achieved by exploitation ( Duan et al. , 2016 ; Wang et al. , 2016a ; Mishra et al. , 2017 ; Rothfuss et al. , 2018 ; Stadie et al. , 2018 ; Zintgraf et al. , 2019 ; Humplik et al. , 2019 ; Kamienny et al. , 2020 ; Dorfman & Tamar , 2020 ) . These approaches can represent the optimal policy ( Kaelbling et al. , 1998 ) , but they struggle to escape local optima due to a chicken-and-egg problem between learning to explore and learning to exploit ( Section 4.1 ) . Several of these approaches ( Humplik et al. , 2019 ; Kamienny et al. , 2020 ) also leverage the problem ID during meta-training , but they still learn end-to-end , so the chicken-and-egg problem remains . Decoupled approaches instead optimize separate exploration and exploitation objectives , via , e.g. , Thompson-sampling ( TS ) ( Thompson , 1933 ; Rakelly et al. , 2019 ) , obtaining exploration trajectories predictive of dynamics or rewards ( Zhou et al. , 2019b ; Gurumurthy et al. , 2019 ; Zhang et al. , 2020 ) , or exploration noise ( Gupta et al. , 2018 ) . While these works do not identify the chicken-and-egg problem , decoupled approaches coincidentally avoid it . However , existing decoupled approaches , including those ( Rakelly et al. , 2019 ; Zhang et al. , 2020 ) that leverage the problem ID , do not learn optimal exploration : TS ( Rakelly et al. , 2019 ) explores by guessing the task and executing a policy for that task , and hence can not represent exploration behaviors that are different from exploitation ( Russo et al. , 2017 ) . Predicting the dynamics ( Zhou et al. , 2019b ; Gurumurthy et al. , 2019 ; Zhang et al. , 2020 ) is inefficient when only a small subset of the dynamics are relevant to solving the task . In contrast , we propose a separate mutual information objective for exploration , which both avoids the chicken-and-egg problem and yields optimal exploration when optimized ( Section 5 ) . Past work ( Gregor et al. , 2016 ; Houthooft et al. , 2016 ; Eysenbach et al. , 2018 ; Warde-Farley et al. , 2018 ) also optimize mutual information objectives , but not for meta-RL . Exploration in general RL . The general RL setting ( i.e. , learning from scratch ) also requires targeted exploration to gather informative samples that enables learning a policy to solve the problem . In contrast to exploration algorithms for general RL ( Bellemare et al. , 2016 ; Pathak et al. , 2017 ; Burda et al. , 2018 ; Leibfried et al. , 2019 ) , which must visit many novel states to find regions with high reward , exploration in meta-RL can be even more targeted by leveraging prior experience from different problems during meta-training . As a result , DREAM can learn new tasks in just two episodes ( Section 6 ) , while learning from scratch can require thousands or even millions of episodes . 3 PRELIMINARIES . Meta-reinforcement learning . The meta-RL setting considers a family of Markov decision processes ( MDPs ) 〈S , A , Rµ , Tµ〉 with states S , actions A , rewardsRµ , and dynamics Tµ , indexed by a onehot problem ID µ ∈M , drawn from a distribution p ( µ ) . Colloquially , we refer to the dynamics as the environment , the rewards as the task , and the entire MDP as the problem . Borrowing terminology from Duan et al . ( 2016 ) , meta-training and meta-testing both consist of repeatedly running trials . Each trial consists of sampling a problem ID µ ∼ p ( µ ) and running N + 1 episodes on the corresponding problem . Following prior evaluation settings ( Finn et al. , 2017 ; Rakelly et al. , 2019 ; Rothfuss et al. , 2018 ; Fakoor et al. , 2019 ) , we designate the first episode in a trial as an exploration episode consisting of T steps for gathering information , and define the goal as maximizing the returns in the subsequent N exploitation episodes ( Figure 1 ) . Following Rakelly et al . ( 2019 ) ; Humplik et al . ( 2019 ) ; Kamienny et al . ( 2020 ) , the easy-to-provide problem ID is available for meta-training , but not meta-testing trials . We formally express the goal in terms of an exploration policy πexp used in the exploration episode and an exploitation policy πtask used in exploitation episodes , but these policies may be the same or share parameters . Rolling out πexp in the exploration episode produces an exploration trajectory τ exp = ( s0 , a0 , r0 , . . . , sT ) , which contains information discovered via exploration . The exploitation policy πtask may then condition on τ exp and optionally , its history across all exploitation episodes in a trial , to maximize exploitation episode returns . The goal is therefore to maximize : J ( πexp , πtask ) = Eµ∼p ( µ ) , τ exp∼πexp [ V task ( τ exp ; µ ) ] , ( 1 ) where V task ( τ exp ; µ ) is the expected returns of πtask conditioned on τ exp , summed over the N exploitation episodes in a trial with problem ID µ. End-to-end meta-RL . A common meta-RL approach ( Wang et al. , 2016a ; Duan et al. , 2016 ; Rothfuss et al. , 2018 ; Zintgraf et al. , 2019 ; Kamienny et al. , 2020 ; Humplik et al. , 2019 ) is to learn to explore and exploit end-to-end by directly optimizing J in ( 1 ) , updating both from rewards achieved during exploitation . These approaches typically learn a single recurrent policy π ( at | st , τ : t ) for both exploration and exploitation ( i.e. , πtask = πexp = π ) , which takes action at given state st and history of experiences spanning all episodes in a trial τ : t = ( s0 , a0 , r0 , . . . , st−1 , at−1 , rt−1 ) . Intuitively , this policy is learned by rolling out a trial , producing an exploration trajectory τ exp and , conditioned on τ exp and the exploitation experiences so far , yielding some exploitation episode returns . Then , credit is assigned to both exploration ( producing τ exp ) and exploitation by backpropagating the exploitation returns through the recurrent policy . Critically , estimates of the expected exploitation returns in ( 1 ) ( e.g. , from a single roll-out or value-function approximation ) form the learning signal for exploration . Directly optimizing the objective J this way can learn optimal exploration and exploitation strategies , but optimization is challenging , which we show in Section 4.1 .

Summary: This paper introduces DREAM, a meta-RL approach that decouples exploration from exploitation. An exploitation policy learns to maximize rewards that are conditioned on an encoder that learns task relevant information. Then an exploration policy learns to collect data that maximizes the mutual information between the encoder and explored states. The work is compared against multiple baselines in simple tasks.