Pre-training as Batch Meta Reinforcement Learning with tiMe 

1 INTRODUCTION . Deep Reinforcement Learning algorithms still require millions of environment interactions to obtain reasonable performance , hindering their applications ( Mnih et al. , 2015 ; Lillicrap et al. , 2016 ; Vuong et al. , 2018 ; Fujimoto et al. , 2018b ; Jaderberg et al. , 2018 ; Arulkumaran et al. , 2019 ; Huang et al. , 2019 ) . This is due to the lack of good pre-training methods . In supervised learning , a pre-trained network significantly reduces sample complexity when learning new tasks ( Zeiler & Fergus , 2013 ; Devlin et al. , 2018 ; Yang et al. , 2019 ) . Meta Reinforcement Learning ( RL ) has been proposed as a framework for pre-training in RL ( Wang et al. , 2016a ; Duan et al. , 2016 ; Finn et al. , 2017 ) . However , such methods still require the collection of millions of interactions during meta-train , which means that they face the same sample complexity challenge as standard RL algorithms . In this work , we use the following definition of pre-training : the ability to use data from a set of tasks to improve performance on unseen , but related tasks . In supervised learning , a key reason why pre-training is incredibly successful is that the dataset used for pre-training can be collected from naturally occurring large-scale processes . This removes the need to manually collect data and allows for scalable data collection , resulting in massive datasets . For example , Mahajan et al . ( 2018 ) pre-trains using existing images and their corresponding hashtags from Instagram to obtain state-of-the-art performance on ImageNet ( Russakovsky et al. , 2014 ) . In this paper , we seek to formalize pre-training in RL in a way that allows for scalable data collection . The data used for pre-training should be purely observational and the policies that are being optimized for should not need to interact with the environment during pre-training . To this end , we propose Batch Meta Reinforcement Learning ( BMRL ) as a formalization of pre-training in RL from only existing and observational data . During training , the learning algorithms only have access to a batch of existing data collected a priori from a family of Markov Decision Process ( MDP ) . During testing , the trained policies should perform well on unseen MDPs sampled from the family . A related setting is Batch RL ( Antos et al. , 2007 ; Lazaric et al. , 2008 ; Lange et al. , 2012 ) , which we emphasize assumes the existing data comes from a single MDP . To enable scalable data collection , this assumption must be relaxed : the existing data should come from a family of related MDPs . Consider smart thermostats , whose goal is to maintain a specific temperature while minimizing electricity cost . Assuming Markovian dynamics , the interactions between a thermostat and its environment can be modelled as a MDP . Data generated by a thermostat operating in a single building can be used to train Batch RL algorithms . However , if we consider the data generated by the same thermostat operating in different buildings , much more data is available . While the interactions between the same thermostat model and different buildings correspond to different MDPs , these MDPs share regularities which can support generalization , such as the physics of heat diffusion . In section 6 , we further discuss the relations between BMRL and other existing formulations . The first challenge in BMRL is the accurate inference of the unseen MDP identity . We show that existing algorithms which sample mini-batches from the existing data to perform Q-learning style updates converge to a degenerate value function , a phenomena we term MDP mis-identification . The second challenge is the interpolation of knowledge about seen MDPs to perform well on unseen MDPs . While Meta RL algorithms can explicitly optimize for this objective thanks to the ability to interact with the environment , we must rely on the inherent generalizability of the trained networks . To mitigate these issues , we propose tiMe , which learns from existing data to distill multiple value functions and MDP embeddings . tiMe is a flexible and scalable pipeline with inductive biases to encourage accurate MDP identity inference and rich supervision to maximize generalization . The pipeline consists of two phases . In the first phase , Batch RL algorithm is used to extract MDPspecific networks from MDP-specific data . The second phase distills the MDP-specific networks . To summarize , our contributions are three folds : ( 1 ) Formulation of Meta RL from observational data as Batch Meta Reinforcement Learning ( BMRL ) ( 2 ) A simple stage-wise approach which works well on standard benchmarks in the BMRL setting ( 3 ) Most importantly , demonstration that Meta RL from only observational data is possible . We hope this work will direct the attention of the meta RL community towards this research direction . 2 PRELIMINARIES . 2.1 BATCH REINFORCEMENT LEARNING . We model the environment as a Markov Decision Process ( MDP ) , uniquely defined as a 5 element tuple Mi = ( S , A , Ti , Ri , γ ) with state space S , action space A , transition function Ti , reward function Ri and discount factor γ ( Puterman , 1994 ; Sutton & Barto , 1998 ) . At each discrete timestep , the agent is in a state s , pick an action a , and arrives at the next state s′ and receives a reward r. The goal of the agent π is to maximize the expected sum of discounted rewards J ( π ) = Eτ∼π , Mi [ ∑∞ t=0 γ tRi ( st , i , at , i , s ′ t , i ) ] where τ = ( s0 , i , a0 , i , r0 , i , s1 , i , a1 , i , r1 , i , . . . ) is a trajectory generated by using π to interact with Mi . We will consider a family of MDPs , defined formally in subsection 2.3 . We thus index each MDP in this family with i . In Batch RL , policies are trained from scratch to solve a single MDP Mi using existing batch of N transition tuples Bi = { ( st , i , at , i , rt , i , s′t , i ) |t = 1 , . . . , N } without any further interaction with Mi . At test time , we use the trained policies to interact with Mi to obtain an empirical estimate of its performance J . Batch RL optimizes for the same objective as standard RL algorithms . However , during training , the learning algorithm only has access to Bi and are not allowed to interact withMi . 2.2 BATCH-CONSTRAINED Q-LEARNING . Fujimoto et al . ( 2018a ) identifies extrapolation error and value function divergence as the modes of failure when modern Q-learning algorithms are applied to the Batch RL setting . Concretely , deep Qlearning algorithms approximate the expected sum of discounted reward starting from a state-action pair E [ ∑∞ t=0 γ tR ( st , at , s ′ t ) |s0 = s , a0 = a ] with a value estimate Q ( s , a ) . The estimate can be learned by sampling transition tuples from the batch and applying the temporal difference update : Q ( s , a ) ← ( 1− αt ) Q ( s , a ) + αt ( r + γQ ( s′ , π ( s′ ) ) ) π ( s′ ) ∈ arg max a∈A Q ( s′ , a ) ( 1 ) The value function diverges if Q fails to accurately estimate the value of π ( s′ ) . Fujimoto et al . ( 2018a ) introduces Batch-Constrained Q-Learning , constraining π to select actions that are similar to actions in the batch to prevent inaccurate values estimation . Concretely , given s′ , a generator G outputs multiple candidate actions { am } m. A perturbation model ξ takes each state-candidate action pair as input and generates small correction term ξ ( s′ , am ) for each candidate . The corrected candidate action am + ξ ( s′ , am ) with the highest value as estimated by a learnt Q is π ( s′ ) : π ( s′ ) = arg max am+ξ ( s′ , am ) Q ( s , am + ξ ( s ′ , am ) ) { am = G ( s′ , zm ) } m zm ∼ N ( 0 , 1 ) Estimation error has also been previously studied and mitigated in model-free RL algorithms ( Hasselt , 2010 ; van Hasselt et al. , 2015 ; Fujimoto et al. , 2018b ) . 2.3 META REINFORCEMENT LEARNING . Meta RL optimizes for average return on a family of MDPs and usually assume that the MDPs in this family share S , A , γ . Each MDP is uniquely defined by a tuple ( Ti , Ri ) . A distribution p ( Ti , Ri ) defines a distribution over MDPs . During meta-train , we train a policy by sampling MDPs from this distribution and sampling trajectories from each sampled MDP , referred to as the meta-train MDPs . During meta-test , unseen MDPs are sampled from p ( Ti , Ri ) , referred to as the meta-test MDPs . The trained policy is used to interact with the meta-test MDPs to obtain estimate of its performance . The choice of whether to update parameters ( Finn et al. , 2017 ) or to keep them fixed during meta-test ( Hochreiter et al. , 2001 ) is left to the learning algorithms , both having demonstrated prior successes . 2.4 MDP IDENTITY INFERENCE WITH SET NEURAL NETWORK . A Meta RL policy needs to infer the meta-test MDP identity to pick actions with high return . Rakelly et al . ( 2019 ) introduces PEARL , which uses a set neural network ( Qi et al. , 2016 ; Zaheer et al. , 2017 ) f as the MDP identity inference function . f takes as input a context set c = { ( sk , ak , rk , s′k ) } k and infers the identity of a MDP in the form of distributed representation in continuous space . The parameters of f is trained to minimize the error of the critic Q : ( Q ( s , a , f ( c ) ) − ( r + V̄ ( s′ , f ( c ) ) ) ) 2 ( 2 ) where V̄ is a learnt state value function . PEARL also adopts an amortized variational approach ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ; Alemi et al. , 2016 ; Kingma & Welling , 2019 ) to train a probabilistic f , which is interpreted as an approximation to the true posterior over the set of possible MDP identities given the context set . 3 BATCH META REINFORCEMENT LEARNING . Let K be the number of meta-train MDPs , N be the number of transition tuples available from each meta-train MDP , θ be the parameter of the policy , we can formulate Batch Meta Reinforcement Learning ( BMRL ) as an optimization problem : arg max θ J ( θ ) = EMi∼p ( Ti , Ri ) [ Eτ∼πθ , Mi [ ∞∑ t=0 γtRi ( st , i , at , i , st′ , i ) ] ] ( 3 ) where the learning algorithms only have access to the batch B during meta-train : B = ∪Ki=1Bi Bi = { ( st , i , at , i , rt , i , s′t , i ) |t = 1 , . . . , N } Mi ∼ p ( Ti , Ri ) We assume we know which MDP each transition in the batch was collected from . This assumption simplifies our setting and is used to devise the algorithms . To maintain the flexibility of the formalization , we do not impose restrictions on the controller that generates the batch . However , the performance of learning algorithms generally increases as the training data becomes more diverse . MDP identity inference challenge To obtain high return on the unseen meta-test MDPs , the trained policies need to accurately infer their identities ( Ritter et al. , 2018 ; Gupta et al. , 2018 ; Humplik et al. , 2019 ) . In BMRL , previously proposed solutions based on Q-learning style updates , where mini-batches are sampled from the batch to minimize TD error , converge to a degenerate solution . subsection 5.1 provides experimental result that demonstrates the phenomena . In finite MDP , this degenerate solution is the optimal value function of the MDP constructed by the relative frequencies of transitions contained in the batch . We can formalize this statement with the following proposition . Proposition 1 . Let N ( s , a , s′ ) be the number of times the triple ( s , a , s′ ) appears in B ( with any reward ) . Performing Q-learning on finite S and A with all the Q-values Q ( s , a ) initialized to 0 and update rule ( 1 ) where ( s , a , s′ , r ) is sampled uniformly at random from B at each step t , will lead the Q-values to converge to the optimal Q-value of the MDP ( S , A , T̂ , R̂ , γ ) almost surely as long as αt ≥ 0 , ∑∞ t=0 αt =∞ , ∑∞ t=0 α 2 t < ∞ , where T̂ ( s , a , s′ ) = 1s=s′ , if ∑ s′′∈S N ( s , a , s′′ ) = 0 N ( s , a , s′ ) ∑ s′′∈S N ( s , a , s′′ ) , otherwise . R̂ ( s , a , s′ ) = 0 , if N ( s , a , s′ ) = 0∑ r : ( s , a , s′ , r ) ∈B r N ( s , a , s′ ) , otherwise . Thus , performing Q-learning style update directly on data sampled from the batch B fails to find a good policy because the value function converges to the optimal value function of the wrong MDP . We refer this phenomena as MDP mis-identification . The proof is shown in subsection A.1 . Interpolation of seen MDPs to unseen MDPs challenge The trained policies need to generalize from the meta-train MDPs to unseen meta-test MDPs . Meta RL tackles this challenge by formulating an optimization problem that explicitly optimizes for the average return of the meta-trained policy after additional gradient steps in unseen MDPs ( Finn et al. , 2017 ; Rothfuss et al. , 2018 ; Nichol et al. , 2018 ) . This is possible thanks to the ability to interact with the environment during metatrain . However , in the meta-train phase of BMRL , the learning algorithms do not have access to the environment . We must rely on the inherent generalizability of the trained networks to perform well on the unseen meta-test MDPs . The key challenge is therefore finding the right inductive biases in the architecture and training procedure to encourage such generalization . The need to find the right inductive biases in RL was highlighted by Botvinick et al . ( 2019 ) ; Zambaldi et al . ( 2019 ) ; Hessel et al . ( 2019 ) . We note that previous works phrase the need to find inductive biases as a means to forgo generality for efficient learning . In our setting , these two goals need not be mutually exclusive .

Summary: The paper studies batch meta learning, i.e. the problem of using a fixed experience from past tasks to learn a policy which can quickly adapt to a new related task. The proposed method combines the techniques of Fujimoto et al. (2018) for stabilizing batch off-policy learning with ideas from Rakelly et al. (2019) for learning a set invariant task embedding using task-specific datasets of transitions. They learn task-specific Q-values which are then distilled into a new Q function which is conditioned on the task embedding instead of the task ID. The embedding is further shaped using a next-state prediction auxiliary loss. The algorithmic ideas feel a bit too incremental and the experimental evaluation could be stronger--I'd recommend trying the method on more complicated environments and including ablation studies.