Meta-Q-Learning

M E TA - Q - L E A R N I N G Rasool Fakoor1 , Pratik Chaudhari2∗ , Stefano Soatto1 , Alexander Smola1 1 Amazon Web Services 2 University of Pennsylvania Email : { fakoor , soattos , smola } @ amazon.com , pratikac @ seas.upenn.edu A B S T R A C T This paper introduces Meta-Q-Learning ( MQL ) , a new off-policy algorithm for meta-Reinforcement Learning ( meta-RL ) . MQL builds upon three simple ideas . First , we show that Q-learning is competitive with state-of-the-art meta-RL algorithms if given access to a context variable that is a representation of the past trajectory . Second , a multi-task objective to maximize the average reward across the training tasks is an effective method to meta-train RL policies . Third , past data from the meta-training replay buffer can be recycled to adapt the policy on a new task using off-policy updates . MQL draws upon ideas in propensity estimation to do so and thereby amplifies the amount of available data for adaptation . Experiments on standard continuous-control benchmarks suggest that MQL compares favorably with the state of the art in meta-RL . 1 I N T R O D U C T I O N Reinforcement Learning ( RL ) algorithms have demonstrated good performance on simulated data . There are however two main challenges in translating this performance to real robots : ( i ) robots are complex and fragile which precludes extensive data collection , and ( ii ) a real robot may face an environment that is different than the simulated environment it was trained in . This has fueled research into MetaReinforcement Learning ( meta-RL ) which develops algorithms that “ meta-train ” on a large number of different environments , e.g. , simulated ones , and aim to adapt to a new environment with few data . How well does meta-RL work today ? Fig . 1 shows the performance of two prototypical meta-RL algorithms on four standard continuous-control benchmarks.1 We compared them to the following simple baseline : an off-policy RL algorithm ( TD3 by Fujimoto et al . ( 2018b ) ) and which was trained to maximize the average reward over all training tasks and modified to use a “ context variable ” that represents the trajectory . All algorithms in this figure use the same evaluation protocol . It is surprising that this simple non-meta-learning-based method is competitive with state-of-the-art meta-RL algorithms . This is the first contribution of our paper : we demonstrate that it is not necessary to meta-train policies to do well on existing benchmarks . Our second contribution is an off-policy meta-RL algorithm named Meta-Q-Learning ( MQL ) that builds upon the above result . MQL uses a simple meta-training procedure : it maximizes the average ∗Work done while at Amazon Web Services 1We obtained the numbers for MAML and PEARL from training logs published by Rakelly et al . ( 2019 ) . rewards across all meta-training tasks using off-policy updates to obtain θ̂meta = arg max θ 1 n n∑ k=1 E τ∼Dk [ ` k ( θ ) ] ( 1 ) where ` k ( θ ) is the objective evaluated on the transition τ obtained from the task Dk ( θ ) , e.g. , 1-step temporal-difference ( TD ) error would set ` k ( θ ) = TD2 ( θ ; τ ) . This objective , which we call the multi-task objective , is the simplest form of meta-training . For adapting the policy to a new task , MQL samples transitions from the meta-training replay buffer that are similar to those from the new task . This amplifies the amount of data available for adaptation but it is difficult to do because of the large potential bias . We use techniques from the propensity estimation literature for performing this adaptation and the off-policy updates of MQL are crucial to doing so . The adaptation phase of MQL solves arg max θ { E τ∼Dnew [ ` new ( θ ) ] + E τ∼Dmeta [ β ( τ ; Dnew , Dmeta ) ` new ( θ ) ] − ( 1− ÊSS ) ‖θ − θ̂meta‖22 } ( 2 ) whereDmeta is the meta-training replay buffer , the propensity score β ( τ ; Dnew , Dmeta ) is the odds of a transition τ belonging to Dnew versusDmeta , and ÊSS is the Effective Sample Size between Dnew and Dmeta that is a measure of the similarly of the new task with the meta-training tasks . The first term computes off-policy updates on the new task , the second term performs β ( · ) -weighted off-policy updates on old data , while the third term is an automatically adapting proximal term that prevents degradation of the policy during adaptation . We perform extensive experiments in Sec . 4.2 including ablation studies using standard meta-RL benchmarks that demonstrate that MQL policies obtain higher average returns on new tasks even if they are meta-trained for fewer time-steps than state-of-the-art algorithms . 2 B A C K G R O U N D This section introduces notation and formalizes the meta-RL problem . We discuss techniques for estimating the importance ratio between two probability distributions in Sec . 2.2 . Consider a Markov Decision Processes ( MDP ) denoted by xt+1 = f k ( xt , ut , ξt ) x0 ∼ pk0 , ( 3 ) where xt ∈ X ⊂ Rd are the states and ut ∈ U ⊂ Rp are the actions . The dynamics fk is parameterized by k ∈ { 1 , . . . , n } where each k corresponds to a different task . The domain of all these tasks , X for the states and U for the actions , is the same . The distribution pk0 denotes the initial state distribution and ξt is the noise in the dynamics . Given a deterministic policy uθ ( xt ) , the actionvalue function for γ-discounted future rewards rkt : = r k ( xt , uθ ( xt ) ) over an infinite time-horizon is qk ( x , u ) = E ξ ( · ) [ ∞∑ t=0 γt rkt |x0 = x , u0 = u , ut = uθ ( xt ) ] . ( 4 ) Note that we have assumed that different tasks have the same state and action space and may only differ in their dynamics fk and reward function rk . Given one task k ∈ { 1 , . . . , n } , the standard Reinforcement Learning ( RL ) formalism solves for θ̂k = arg max θ ` k ( θ ) where ` k ( θ ) = E x∼p0 [ qk ( x , uθ ( x ) ) ] . ( 5 ) Let us denote the dataset of all states , actions and rewards pertaining to a task k and policy uθ ( x ) by Dk ( θ ) = { xt , uθ ( xt ) , r k , xt+1 = f k ( xt , uθ ( xt ) , ξt ) } t≥0 , x ( 0 ) ∼pk0 , ξ ( · ) ; we will often refer toDk as the “ task ” itself . The Deterministic Policy Gradient ( DPG ) algorithm ( Silver et al. , 2014 ) for solving ( 5 ) learns a ϕ-parameterized approximation qϕ to the optimal value func- tion qk by minimizing the Bellman error and the optimal policy uθ that maximizes this approximation by solving the coupled optimization problem ϕ̂k = arg min ϕ E τ∼Dk [ ( qϕ ( x , u ) − rk − γ qϕ ( x′ , uθ̂k ( x ′ ) ) ) 2 ] , θ̂k = arg max θ E τ∼Dk [ q ϕ̂k ( x , uθ ( x ) ) ] . ( 6 ) The 1-step temporal difference error ( TD error ) is defined as TD2 ( θ ) = ( qϕ ( x , u ) − rk − γ qϕ ( x′ , uθ ( x′ ) ) ) 2 ( 7 ) where we keep the dependence of TD ( · ) on ϕ implicit . DPG , or its deep network-based variant DDPG ( Lillicrap et al. , 2015 ) , is an off-policy algorithm . This means that the expectations in ( 6 ) are computed using data that need not be generated by the policy being optimized ( uθ ) , this data can come from some other policy . In the sequel , we will focus on the parameters θ parameterizing the policy . The parameters ϕ of the value function are always updated to minimize the TD-error and are omitted for clarity . 2 . 1 M E TA - R E I N F O R C E M E N T L E A R N I N G ( M E TA - R L ) Meta-RL is a technique to learn an inductive bias that accelerates the learning of a new task by training on a large of number of training tasks . Formally , meta-training on tasks from the meta-training set Dmeta = { Dk } k=1 , ... , n involves learning a policy θ̂meta = arg max θ 1 n n∑ k=1 ` kmeta ( θ ) ( 8 ) where ` kmeta ( θ ) is a meta-training loss that depends on the particular method . Gradient-based meta-RL , let us take MAML by Finn et al . ( 2017 ) as a concrete example , sets ` kmeta ( θ ) = ` k ( θ + α∇θ ` k ( θ ) ) ( 9 ) for a step-size α > 0 ; ` k ( θ ) is the objective of non-meta-RL ( 5 ) . In this case ` kmeta is the objective obtained on the task Dk after one ( or in general , more ) updates of the policy on the task . The idea behind this is that even if the policy θ̂meta does not perform well on all tasks in Dmeta it may be updated quickly on a new task Dnew to obtain a well-performing policy . This can either be done using the same procedure as that of meta-training time , i.e. , by maximizing ` newmeta ( θ ) with the policy θ̂meta as the initialization , or by some other adaptation procedure . The meta-training method and the adaptation method in meta-RL , and meta-learning in general , can be different from each other . 2 . 2 L O G I S T I C R E G R E S S I O N F O R E S T I M AT I N G T H E P R O P E N S I T Y S C O R E Consider standard supervised learning : given two distributions q ( x ) ( say , train ) and p ( x ) ( say , test ) , we would like to estimate how a model ’ s predictions ŷ ( x ) change across them . This is formally done using importance sampling : E x∼p ( x ) E y|x [ ` ( y , ŷ ( x ) ) ] = E x∼q ( x ) E y|x [ β ( x ) ` ( y , ŷ ( x ) ) ] ; ( 10 ) where y|x are the true labels of data , the predictions of the model are ŷ ( x ) and ` ( y , ŷ ( x ) ) is the loss for each datum ( x , y ) . The importance ratio β ( x ) = dpdq ( x ) , also known as the propensity score , is the Radon-Nikodym derivative ( Resnick , 2013 ) of the two data densities and measures the odds of a sample x coming from the distribution p versus the distribution q . In practice , we do not know the densities q ( x ) and p ( x ) and therefore need to estimate β ( x ) using some finite data Xq = { x1 , . . . , xm } drawn from q and Xp = { x′1 , . . . , x′m } drawn from p. As Agarwal et al . ( 2011 ) show , this is easy to do using logistic regression . Set zk = 1 to be the labels for the data in Xq and zk = −1 to be the labels of the data in Xp for k ≤ m and fit a logistic classifier on the combined 2m samples by solving w∗ = min w 1 2m ∑ ( x , z ) log ( 1 + e−zw > x ) + c ‖w‖2 . ( 11 ) This gives β ( x ) = P ( z = −1|x ) P ( z = 1|x ) = e−w ∗ > x . ( 12 ) Normalized Effective Sample Size ( ÊSS ) : A related quantity to β ( x ) is the normalized Effective Sample Size ( ÊSS ) which we define as the relative number of samples from the target distribution p ( x ) required to obtain an estimator with performance ( say , variance ) equal to that of the importance sampling estimator ( 10 ) . It is not possible to compute the ÊSS without knowing both densities q ( x ) and p ( x ) but there are many heuristics for estimating it . A popular one in the Monte Carlo literature ( Kong , 1992 ; Smith , 2013 ; Elvira et al. , 2018 ) is ÊSS = 1 m ( ∑m k=1 β ( xk ) ) 2∑m k=1 β ( xk ) 2 ∈ [ 0 , 1 ] ( 13 ) where X = { x1 , . . . , xm } is some finite batch of data . Observe that if two distributions q and p are close then the ÊSS is close to one ; if they are far apart the ÊSS is close to zero . 3 M Q L This section describes the MQL algorithm . We begin by describing the meta-training procedure of MQL including a discussion of multi-task training in Sec . 3.1 . The adaptation procedure is described in Sec . 3.2 . 3 . 1 M E TA - T R A I N I N G MQL performs meta-training using the multi-task objective . Note that if one sets ` kmeta ( θ ) , ` k ( θ ) = E x∼pk0 [ qk ( x , uθ ( x ) ) ] ( 14 ) in ( 8 ) then the parameters θ̂meta are such that they maximize the average returns over all tasks from the meta-training set . We use an off-policy algorithm named TD3 ( Fujimoto et al. , 2018b ) as the building block and solve for θ̂meta = arg min θ 1 n n∑ k=1 E τ∼Dk [ TD2 ( θ ) ] ; ( 15 ) where TD ( · ) is defined in ( 7 ) . As is standard in TD3 , we use two action-value functions parameterized by ϕ1 and ϕ2 and take their minimum to compute the target in ( 7 ) . This trick known as “ doubleQ-learning ” reduces the over-estimation bias . Let us emphasize that ( 14 ) is a special case of the procedure outlined in ( 8 ) . The following remark explains why MQL uses the multi-task objective as opposed to the meta-training objective used , for instance , in existing gradient-based meta-RL algorithms . Remark 1 . Let us compare the critical points of the m-step MAML objective ( 9 ) to those of the multi-task objective which uses ( 14 ) . As is done by the authors in Nichol et al . ( 2018 ) , we can perform a Taylor series expansion around the parameters θ to obtain ∇ ` kmeta ( θ ) = ∇ ` k ( θ ) + 2α ( m− 1 ) ( ∇2 ` k ( θ ) ) ∇ ` k ( θ ) +O ( α2 ) . ( 16 ) Further , note that∇ ` kmeta in ( 16 ) is also the gradient of the loss ` k ( θ ) + α ( m− 1 ) ‖∇ ` k ( θ ) ‖22 ( 17 ) up to first order . This lends a new interpretation that MAML is attracted towards regions in the loss landscape that under-fit on individual tasks : parameters with large ‖∇ ` k‖2 will be far from the local maxima of ` k ( θ ) . The parameters α and m control this under-fitting . Larger the number of gradient steps , larger the under-fitting effect . This remark suggests that the adaptation speed of gradient-based meta-learning comes at the cost of under-fitting on the tasks . 3 . 1 . 1 D E S I G N I N G C O N T E X T As discussed in Sec . 1 and 4.4 , the identity of the task in meta-RL can be thought of as the hidden variable of an underlying partially-observable MDP . The optimal policy on the entire trajectory of the states , actions and the rewards . We therefore design a recurrent context variable zt that depends on { ( xi , ui , ri ) } i≤t . We set zt to the hidden state at time t of a Gated Recurrent Unit ( GRU by Cho et al . ( 2014 ) ) model . All the policies uθ ( x ) and value functions qϕ ( x , u ) in MQL are conditioned on the context and implemented as uθ ( x , z ) and qϕ ( x , u , z ) . Any other recurrent model can be used to design the context ; we used a GRU because it offers a good trade-off between a rich representation and computational complexity . Remark 2 ( MQL uses a deterministic context that is not permutation invariant ) . We have aimed for simplicity while designing the context . The context in MQL is built using an off-the-shelf model like GRU and is not permutation invariant . Indeed , the direction of time affords crucial information about the dynamics of a task to the agent , e.g. , a Half-Cheetah running forward versus backward has arguably the same state trajectory but in a different order . Further , the context in MQL is a deterministic function of the trajectory . Both these aspects are different than the context used by Rakelly et al . ( 2019 ) who design an inference network and sample a probabilistic context conditioned on a moving window . RL algorithms are quite complex and challenging to reproduce . Current meta-RL techniques which build upon them further exacerbate this complexity . Our demonstration that a simple context variable is enough is an important contribution .

Summary: The authors propose meta Q-learning, an algorithm for off-policy meta RL. The idea is to meta-train a context-dependent policy to maximize the expected return averaged over all training tasks, and then adapt this policy to any new task by leveraging both novel and past experience using importance sampling corrections. The proposed approach is evaluated on standard Mujoco benchmarks and compared to other relevant meta-rl algorithms.