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: This paper proposes Meta Q-Learning (MQL), an algorithm for efficient off-policy meta-learning. The method relies on a simple multi-task objective which provides initial parameter values for the adaptation phase. Adaptation is performed by gradient descent, minimizing TD-error on the new validation task (regularizing towards initial parameter values). To make adaptation data efficient, the method makes heavy use of off-policy data generated during meta-training, by minimizing its importance weighted TD-error. Importance weights are estimated via a likelihood ratio estimator, and are also used to derive the effective sample size of the meta-training batch, which is used to adaptively weight the regularization term. Intuitively, this has the effect of turning off regularization when meta-training trajectories are “close” to validation trajectories. One important but somewhat orthogonal contribution of the paper is to highlight the importance of context in meta-learning and fast adaptation. Concretely, the authors show that a simple actor-critic algorithm (TD3), whose policy and value are conditioned on a context variable derived from a recurrent network performs surprisingly well in comparison to SoTA meta-learning algorithms like PEARL. MQL is evaluated on benchmark meta-RL environments from continuous control tasks and is shown to perform competitively with PEARL.