Parrot: Data-Driven Behavioral Priors for Reinforcement Learning

1 INTRODUCTION . Reinforcement Learning ( RL ) is an attractive paradigm for robotic learning because of its flexibility in being able to learn a diverse range of skills and its capacity to continuously improve . However , RL algorithms typically require a large amount of data to solve each individual task , including simple ones . Since an RL agent is generally initialized without any prior knowledge , it must try many largely unproductive behaviors before it discovers a high-reward outcome . In contrast , humans rarely attempt to solve new tasks in this way : they draw on their prior experience of what is useful when they attempt a new task , which substantially shrinks the task search space . For example , faced with a new task involving objects on a table , a person might grasp an object , stack multiple objects , or explore other object rearrangements , rather than re-learning how to move their arms and fingers . Can we endow RL agents with a similar sort of behavioral prior from past experience ? In other fields of machine learning , the use of large prior datasets to bootstrap acquisition of new capabilities has been studied extensively to good effect . For example , language models trained on large , diverse datasets offer representations that drastically improve the efficiency of learning downstream tasks ( Devlin et al. , 2019 ) . What would be the analogue of this kind of pre-training in robotics and RL ? One way we can approach this problem is to leverage successful trials from a wide range of previously seen tasks to improve learning for new tasks . The data could come from previously learned policies , from human demonstrations , or even unstructured teleoperation of robots ( Lynch et al. , 2019 ) . In this paper , we show that behavioral priors can be obtained through representation learning , and the representation in question must not only be a representation of inputs , but actually a representation of input-output relationships – a space of possible and likely mappings from states to actions among which the learning process can interpolate when confronted with a new task . What makes for a good representation for RL ? Given a new task , a good representation must ( a ) provide an effective exploration strategy , ( b ) simplify the policy learning problem for the RL algorithm , and ( c ) allow the RL agent to retain full control over the environment . In this paper , we address ∗Equal contribution . Correspondence to Avi Singh ( avisingh @ berkeley.edu ) . Trial 1 Trial 2 Trial n-1 Trial n Transfer Behavioral Prior Learning New Tasks all of these challenges through learning an invertible function that maps noise vectors to complex , high-dimensional environment actions . Building on prior work in normalizing flows ( Dinh et al. , 2017 ) , we train this mapping to maximize the ( conditional ) log-likelihood of actions observed in successful trials from past tasks . When dropped into a new MDP , the RL agent can now sample from a unit Gaussian , and use the learned mapping ( which we refer to as the behavioral prior ) to generate likely environment actions , conditional on the current observation . This learned mapping essentially transforms the original MDP into a simpler one for the RL agent , as long as the original MDP shares ( partial ) structure with previously seen MDPs ( see Section 3 ) . Furthermore , since this mapping is invertible , the RL agent still retains full control over the original MDP : for every possible environment action , there exists a point within the support of the Gaussian distribution that maps to that action . This allows the RL agent to still try out new behaviors that are distinct from what was previously observed . Our main contribution is a framework for pre-training in RL from a diverse multi-task dataset , which produces a behavioral prior that accelerates acquisition of new skills . We present an instantiation of this framework in robotic manipulation , where we utilize manipulation data from a diverse range of prior tasks to train our behavioral prior , and then use it to bootstrap exploration for new tasks . By making it possible to pre-train action representations on large prior datasets for robotics and RL , we hope that our method provides a path toward leveraging large datasets in the RL and robotics settings , much like language models can leverage large text corpora in NLP and unsupervised pretraining can leverage large image datasets in computer vision . Our method , which we call Prior AcceleRated ReinfOrcemenT ( PARROT ) , is able to quickly learn tasks that involve manipulating previously unseen objects , from image observations and sparse rewards , in settings where RL from scratch fails to learn a policy at all . We also compare against prior works that incorporate prior data for RL , and show that PARROT substantially outperforms these prior works . 2 RELATED WORK . Combining RL with demonstrations . Our work is related to methods for learning from demonstrations ( Pomerleau , 1989 ; Schaal et al. , 2003 ; Ratliff et al. , 2007 ; Pastor et al. , 2009 ; Ho & Ermon , 2016 ; Finn et al. , 2017b ; Giusti et al. , 2016 ; Sun et al. , 2017 ; Zhang et al. , 2017 ; Lynch et al. , 2019 ) . While demonstrations can also be used to speed up RL ( Schaal , 1996 ; Peters & Schaal , 2006 ; Kormushev et al. , 2010 ; Hester et al. , 2017 ; Vecerı́k et al. , 2017 ; Nair et al. , 2018 ; Rajeswaran et al. , 2018 ; Silver et al. , 2018 ; Peng et al. , 2018 ; Johannink et al. , 2019 ; Gupta et al. , 2019 ) , this usually requires collecting demonstrations for the specific task that is being learned . In contrast , we use data from a wide range of other prior tasks to speed up RL for a new task . As we show in our experiments , PARROT is better suited to this problem setting when compared to prior methods that combine imitation and RL for the same task . Generative modeling and RL . Several prior works model multi-modal action distributions using expert trajectories from different tasks . IntentionGAN ( Hausman et al . ) and InfoGAIL ( Li et al. , 2017 ) learn multi-modal policies via interaction with an environment using an adversarial imitation approach ( Ho & Ermon , 2016 ) , but we learn these distributions only from data . Other works learn these distributions from data ( Xie et al. , 2019 ; Rhinehart et al. , 2020 ) and utilize them for planning at test time to optimize a user-provided cost function . In contrast , we use the behavioral prior to augment model-free RL of a new task . This allows us to learn policies for new tasks that may be substantially different from prior tasks , since we can collect data specific to the new task at hand , and we do not explicitly need to model the environment , which can be complicated for high-dimensional state and action spaces , such as when performing continuous control from images observations . Another line of work ( Ghadirzadeh et al. , 2017 ; Hämäläinen et al. , 2019 ; Ghadirzadeh et al. , 2020 ) explores using generative models for RL , using a variational autoencoder ( Kingma & Welling , 2014 ) to model entire trajectories in an observation-independent manner , and then learning an open-loop , single-step policy using RL to solve the downstream task . Our approach differs in several key aspects : ( 1 ) our model is observation-conditioned , allowing it to prioritize actions that are relevant to the current scene or environment , ( 2 ) our model allows for closed-loop feedback control , and ( 3 ) our model is invertible , allowing the high-level policy to retain full control over the action space . Our experiments demonstrate these aspects are crucial for solving harder tasks . Hierarchical learning . Our method can be interpreted as training a hierarchical model : the lowlevel policy is the behavioral prior trained on prior data , while the high-level policy is trained using RL and controls the low-level policy . This structure is similar to prior work in hierarchical RL ( Dayan & Hinton , 1992 ; Parr & Russell , 1997 ; Dietterich , 1998 ; Sutton et al. , 1999 ; Kulkarni et al. , 2016 ) . We divide prior work in hierarchical learning into two categories : methods that seek to learn both the low-level and high-level policies through active interaction with an environment ( Kupcsik et al. , 2013 ; Heess et al. , 2016 ; Bacon et al. , 2017 ; Florensa et al. , 2017 ; Haarnoja et al. , 2018a ; Nachum et al. , 2018 ; Chandak et al. , 2019 ; Peng et al. , 2019 ) , and methods that learn temporally extended actions , also known as options , from demonstrations , and then recompose them to perform long-horizon tasks through RL or planning ( Fox et al. , 2017 ; Krishnan et al. , 2017 ; Kipf et al. , 2019 ; Shankar et al. , 2020 ; Shankar & Gupta , 2020 ) . Our work shares similarities with the data-driven approach of the latter methods , but work on options focuses on modeling the temporal structure in demonstrations for a small number of long-horizon tasks , while our behavioral prior is not concerned with temporally-extended abstractions , but rather with transforming the original MDP into one where potentially useful behaviors are more likely , and useless behaviors are less likely . Meta-learning . Our goal in this paper is to utilize data from previously seen tasks to speed up RL for new tasks . Meta-RL ( Duan et al. , 2016 ; Wang et al. , 2016 ; Finn et al. , 2017a ; Mishra et al. , 2017 ; Rakelly et al. , 2019 ; Mendonca et al. , 2019 ; Zintgraf et al. , 2020 ; Fakoor et al. , 2020 ) and meta-imitation methods ( Duan et al. , 2017 ; Finn et al. , 2017c ; Huang et al. , 2018 ; James et al. , 2018 ; Paine et al. , 2018 ; Yu et al. , 2018 ; Huang et al. , 2019 ; Zhou et al. , 2020 ) also seek to speed up learning for new tasks by leveraging experience from previously seen tasks . While meta-learning provides an appealing and principled framework to accelerate acquisition of future tasks , we focus on a more lightweight approach with relaxed assumptions that make our method more practically applicable , and we discuss these assumptions in detail in the next section . 3 PROBLEM SETUP . Our goal is to improve an agent ’ s ability to learn new tasks by incorporating a behavioral prior , which it can acquire from previously seen tasks . Each task can be considered a Markov decision process ( MDP ) , which is defined by a tuple ( S , A , T , r , γ ) , where S and A represent state and action spaces , T ( s′|s , a ) and r ( s , a ) represent the dynamics and reward functions , and γ ∈ ( 0 , 1 ) represents the discount factor . Let p ( M ) denote a distribution over such MDPs , with the constraint that the state and action spaces are fixed . In our experiments , we treat high-dimensional images as s , which means that this constraint is not very restrictive in practice . In order for the behavioral prior to be able to accelerate the acquisition of new skills , we assume the behavioral prior is trained on data that structurally resembles potential optimal policies for all or part of the new task . For example , if the new task requires placing a bottle in a tray , the prior data might include some behaviors that involve picking up objects . There are many ways to formalize this assumption . One way to state this formally is to assume that prior data consists of executions of near-optimal policies for MDPs f ( z ; s ) pi ( z|s ) a s Z0 Z1 Z2 Z3=a drawn according to M ∼ p ( M ) , and the new task M ? is likewise drawn from p ( M ) . In this case , the generative process for the prior data can be expressed as : M ∼ p ( M ) , πM ( τ ) = argmax π Eπ , M [ RM ] , τM ∼ πM ( τ ) , ( 1 ) where τM = ( s1 , a1 , s2 , a2 , . . . , sT , aT ) is a sequence of state and actions , πM ( τ ) denotes a nearoptimal policy ( Kearns & Singh , 2002 ) for MDPM andRM = ∑∞ t=0 γ trt . When incorporating the behavioral prior for learning a new task M ? , our goal is the same as standard RL : to find a policy π that maximizes the expected return argmaxπ Eπ , M ? [ RM ? ] . Our assumption on tasks being drawn from a distribution p ( M ) shares similarities with the meta-RL problem ( Wang et al. , 2016 ; Duan et al. , 2016 ) , but our setup is different : it does not require accessing any task in p ( M ) except the new task we are learning , M ? . Meta-RL methods need to interact with the tasks in p ( M ) during meta-training , with access to rewards and additional samples , whereas we learn our behavioral prior simply from data , without even requiring this data to be labeled with rewards . This is of particular importance for real-world problem settings such as robotics : it is much easier to store data from prior tasks ( e.g. , different environments ) than to have a robot physically revisit those prior settings and retry those tasks , and not requiring known rewards makes it possible to use data from a variety of sources , including human-provided demonstrations . In our setting , RL is performed in only one environment , while the prior data can come from many environments . Our setting is related to meta-imitation learning ( Duan et al. , 2017 ; Finn et al. , 2017c ) , as we speed up learning new tasks using data collected from past tasks . However , meta-imitation learning methods require at least one demonstration for each new task , whereas our method can learn new tasks without any demonstrations . Further , our data requirements are less stringent : meta-imitation learning methods require all demonstrations to be optimal , require all trajectories in the dataset to have a task label , and requires “ paired demonstrations ” , i.e . at least two demonstrations for each task ( since meta-imitation methods maximize the likelihood of actions from one demonstration after conditioning the policy on another demonstration from the same task ) . Relaxing these requirements increases the scalability of our method : we can incorporate data from a wider range of sources , and we do not need to explicitly organize it into specific tasks .

Summary: This paper introduces PARROT, a novel approach for pretraining a reinforcement learning agent on near-optimal trajectories by learning a behavioral prior. Essentially, the authors learn a word2vec style embedding of actions for a simple virtual single-arm environment. This embedding will naturally place more common examples from its training data towards the center of the gaussian, making sampling them during training time more likely. The authors demonstrate that this approach outperforms existing pretraining methods in this domain.