Correcting Momentum in Temporal Difference Learning

1 INTRODUCTION . Temporal Difference ( TD ) learning ( Sutton , 1988 ) is a fundamental method of Reinforcement Learning ( RL ) , which works by using estimates of future predictions to learn to make predictions about the present . To scale to large problems , TD has been used in conjunction with deep neural networks ( DNNs ) to achieve impressive performance ( e.g . Mnih et al. , 2013 ; Schulman et al. , 2017 ; Hessel et al. , 2018 ) . Unfortunately , the naive application of TD to DNNs has been shown to be brittle ( Machado et al. , 2018 ; Farebrother et al. , 2018 ; Packer et al. , 2018 ; Witty et al. , 2018 ) , with extensions such as the n-step TD update ( Fedus et al. , 2020 ) or the TD ( λ ) update ( Molchanov et al. , 2016 ) only marginally improving performance and generalization capabilities when coupled with DNNs . Part of the success of DNNs , including when applied to TD learning , is the use of adaptive or accelerated optimization methods ( Hinton et al. , 2012 ; Sutskever et al. , 2013 ; Kingma & Ba , 2015 ) to find good parameters . In this work we investigate and extend the momentum algorithm ( Polyak , 1964 ) as applied to TD learning in DNNs . While accelerated TD methods have received some attention in the literature , this is typically done in the context of linear function approximators ( Baxter & Bartlett , 2001 ; Meyer et al. , 2014 ; Pan et al. , 2016 ; Gupta et al. , 2019 ; Gupta , 2020 ; Sun et al. , 2020 ) , and while some studies have considered the mix of DNNs and TD ( Zhang et al. , 2019 ; Romoff et al. , 2020 ) , many are limited to a high-level analysis of hyperparameter choices for existing optimization methods ( Sarigül & Avci , 2018 ; Andrychowicz et al. , 2020 ) ; or indeed the latter are simply applied as-is to train RL agents ( Mnih et al. , 2013 ; Hessel et al. , 2018 ) . For an extended discussion of related work , we refer the reader to appendix A . As a first step in going beyond the naive use of supervised learning tools in RL , we examine momentum . We argue that momentum , especially as it is used in conjunction with TD and DNNs , adds an additional form of bias which can be understood as the staleness of accumulated information . We quantify this bias , and propose a corrected momentum algorithm that reduces this staleness and is capable of improving performance . 1.1 REINFORCEMENT LEARNING AND TEMPORAL DIFFERENCE LEARNING . A Markov Decision Process ( MDP ) ( Bellman , 1957 ; Sutton & Barto , 2018 ) M = 〈S , A , R , P , γ〉 consists of a state space S , an action space A , a reward function R : S → R and a transition function P ( s′|s , a ) . RL agents usually aim to optimize the expectation of the long-term return , G ( St ) = ∑∞ k=t γ k−tR ( Sk ) where γ ∈ [ 0 , 1 ) is called the discount factor . Policies π ( a|s ) map states to action distributions . Value functions V π and Qπ map states and states-action pairs to expected returns , and can be written recursively : V π ( St ) = Eπ [ G ( St ) ] = Eπ [ R ( St , At ) + γV ( St+1 ) |At ∼ π ( St ) ] Qπ ( St , At ) = Eπ [ R ( St , At ) + γ ∑ aπ ( a|St+1 ) Q ( St+1 , a ) ] We approximate V π with Vθ . We can train Vθ via regression to observed values of G , but these recursive equations also give rise to the Temporal Difference ( TD ) update rules for policy evaluation , relying on current estimates of V to bootstrap , which for example in the tabular case is written as : V ( St ) ← V ( St ) − α ( V ( St ) − ( R ( St ) + γV ( St+1 ) ) ) , ( 1 ) where α ∈ [ 0 , 1 ) is the step size . Alternatively , estimates of this update can be performed by a so-called semi-gradient ( Sutton & Barto , 2018 ) algorithm where the “ TD ( 0 ) loss ” is minimized : θt+1 = θt − α∇θt ( Vθt ( St ) − ( R ( St ) + γV̄θt ( St+1 ) ) ) 2 , ( 2 ) with V̄ meaning we consider V constant for the purpose of gradient computation . 1.2 BIAS AND STALENESS IN MOMENTUM . The usual form of momentum ( Polyak , 1964 ; Sutskever et al. , 2013 ) in stochastic gradient descent ( SGD ) maintains an exponential moving average with factor β of gradients w.r.t . to some objective J , changing parameters θt ∈ Rn ( t is here SGD time rather than MDP time ) with learning rate α : µt = βµt−1 + ( 1− β ) ∇θt−1Jt ( θt−1 ) ( 3 ) θt = θt−1 − αµt ( 4 ) We assume here that the objective Jt is time-dependent , as is the case for example in minibatch training or online learning . Note that other similar forms of this update exist , notably Nesterov ’ s accelerated gradient method ( Nesterov , 1983 ) , as well as undampened methods that omit ( 1− β ) in ( 3 ) or replace ( 1− β ) with α , found in popular deep learning packages ( Paszke et al. , 2019 ) . We make the observation that , at time t , the gradients accumulated in µ are stale . They were computed using past parameters rather than θt , and in general we ’ d expect ∇θtJt ( θt ) 6= ∇θkJt ( θk ) , k < t. As such , the update in ( 4 ) is a biased update . In supervised learning where one learns a mapping from x to y , this staleness only has one source : θ changes but the target y stays constant . We argue that in TD learning , momentum becomes doubly stale : not only does the value network change , but the target ( the equivalent of y ) itself changes1 with every parameter update . Consider the TD objective in ( 2 ) , when θ changes , not only does V ( s ) change , but V ( s′ ) as well . The objective itself changes , making past gradients stale and less aligned with recent gradients ( even more so when there is gradient interference ( Liu et al. , 2019 ; Achiam et al. , 2019 ; Bengio et al. , 2020 ) , constructive or destructive ) . Note that several sources of bias already exist in TD learning , notably the traditional parametric bias ( of the bias-variance tradeoff when selecting capacity ) , as well as the bootstrapping bias ( of the error in V ( s′ ) when using it as a target ; using a frozen target prevents this bias from compounding ) . We argue that the staleness in momentum we describe is an additional form of bias , slowing down or preventing convergence . This has been hinted at before , e.g . Gupta ( 2020 ) suggests that momentum hinders learning in linear TD ( 0 ) . We wish to understand and possibly correct this staleness in momentum . In this work we propose answers to the following questions : • Is this bias significant in supervised learning ? No , the bias exists but has a minimal effect at best when comparing to an unbiased oracle . • Is this bias significant in TD learning ? Yes , we can quantify the bias , and comparisons to an unbiased oracle reveal significant differences . • Can we correct µt to remove this bias ? Yes , we derive an online update that approximately corrects µt using only first order derivatives . • Does the correction help in TD learning ? Yes , using a staleness-corrected momentum improves sample complexity , in policy evaluation , especially in an online setting . 1Interestingly , even in most recent value-based control works ( Hessel et al. , 2018 ) a ( usually frozen ) copy is used for stability , meaning that the target only changes when the copy is updated . This is considered a “ trick ” which it would be compelling to get rid of , since it slows down learning , and since most recent policy-gradient methods ( which still use a value function ) do not make use of such copies ( Schulman et al. , 2017 ) . 2 CORRECTING MOMENTUM . 2.1 IDENTIFYING BIAS . Here we propose an experimental protocol to quantify the bias of momentum . We assume we are minimizing some objective in the minibatch or online setting . In such a stochastic setting , momentum is usually understood as a variance reduction method , or as approximating the large-batch setting ( Botev et al. , 2017 ) , but as discussed above , momentum also induces bias in the optimization . We note that µt , the momentum at time t , can be rewritten as : µt = ( 1− β ) t∑ i=1 βt−i∇θiJi ( θi ) , ( 5 ) and argue that an ideal unbiased momentum µ∗t would approximate the large batch case by only discounting past minibatches and using current parameters rather than past parameters : µ∗t def = ( 1− β ) t∑ i=1 βt−i∇θtJi ( θt ) . ( 6 ) Note that the only difference between ( 5 ) and ( 6 ) is the use of θi vs θt . The only way to exactly compute µ∗t is to recompute the entire sum after every parameter update . We will consider this our unbiased oracle . To compute µ∗t empirically , since beyond a certain power βk becomes small , we will use an effective horizon of h = 2/ ( 1− β ) steps ( i.e . start at i = t− h rather than at i = 1 ) . We call the difference between µt and µ∗t the bias , but note that in the overparameterized stochastic gradient case , minor differences in gradients can quickly send parameters in different regions . This makes the direct measure of ‖µt − µ∗t ‖ uninformative . Instead , we measure the optimization bias by simply comparing the loss of a model trained with µt against that of a model trained with µ∗t . Finally , we note that , in RL , momentum approximating the batch case is related to ( approximately ) replaying an entire buffer at once ( instead of sampling transitions ) . This has been shown to also have interesting forms in the linear case ( van Seijen & Sutton , 2015 ) , reminiscent of the correction derived below . We also note that , while the mathematical expression of momentum and eligibility traces ( Sutton , 1988 ) look fairly similar , they estimate a very different quantity ( see appendix A ) . 2.2 CORRECTING BIAS . Here we propose a way to approximate µ∗ , and so derive an approximate correction to the bias in momentum for supervised learning , as well as for Temporal Difference ( TD ) learning . We consider here the simple regression and TD ( 0 ) cases , for the online ( minibatch size 1 ) case . We show the full derivations , and results for any minibatch size , TD ( λ ) and n-step TD , in appendix B . In least squares regression with loss δ2 we can write the gradient gt as : gt ( θt ) = ( yt − fθt ( xt ) ) ∇θtfθt ( xt ) = ∇θtδ2t /2 . ( 7 ) We would like to “ correct ” this gradient as we move away from θt . A simple way to do so is to compute the Taylor expansion around θt of gt : gt ( θt + ∆θ ) = gt ( θt ) +∇θtgt ( θt ) > ∆θ + o ( ‖∆θ‖22 ) , ( 8 ) ≈ gt ( θt ) + ( δ∇2θtfθt ( xt ) −∇θtfθt ( xt ) ⊗∇θtfθt ( xt ) ) > ∆θ , ( 9 ) where ⊗ is the outer product , ∇2 the second derivative . We note that the term multiplying ∆θ is commonly known as “ the Hessian ” of the loss . We also note that Taylor expansions around parameters have a rich history in deep learning ( LeCun et al. , 1990 ; Molchanov et al. , 2016 ) , and that in spite of its non-linearity , the parameter space of a deep network is filled with locally consistent regions ( Balduzzi et al. , 2017 ) in which Taylor expansions are accurate . The same correction computation can be made for TD ( 0 ) with TD loss δ2 . Remember that we write t as the learning time ; we denote an MDP transition as ( st , at , rt , s′t ) , making no assumption on the distribution of transitions : gt ( θt ) = ( Vθt ( st ) − rt − γVθt ( s′t ) ) ∇θtVθt ( st ) = ∇θtδ2t /2 , ( 10 ) gt ( θt + ∆θ ) ≈ gt ( θt ) + ( ∇θt ( Vθt ( st ) − γVθt ( s′t ) ) ⊗∇θtVθt ( st ) + δ∇2θtVθt ( st ) ) > ∆θ . ( 11 ) Here , because we are using a semi-gradient method , the term multiplying ∆θ is not exactly the Hessian : when computing ∇θδ2 , we hold V ( s′ ) constant , but when computing ∇θg , we need to consider Vθ ( s′ ) a function of θ as it affects g , and so compute its gradient.2 Without loss of generality , let us write equations like ( 9 ) and ( 11 ) using the matrix Zt ∈ Rn×n : gt ( θt + ∆θ ) ≈ gt ( θt ) + Z > t ∆θ , ( 12 ) where the form of Zt , the “ Taylor term ” , depends on the loss ( e.g . the Hessian in ( 9 ) ) . Recall that in ( 6 ) we define µ∗ as the discounted sum of gradients using θt for losses Ji , i ≤ t. We call those gradients corrected gradients , gti : = gi ( θi + ( θt − θi ) ) . At timestep t we update θt with αµt , thus we substitute ∆θ = θt − θi = −α ∑t k=i µk in ( 12 ) and get : ĝti def = gi ( θi ) − αZ > i t−1∑ k=i µk ≈ gti . ( 13 ) We can now approximate the unbiased momentum µ∗t using ( 13 ) , which we denote µ̂ : µ̂t def = ( 1− β ) t∑ i=1 βt−iĝti ( 14 ) = µt − α ( 1− β ) t−1∑ k=1 k∑ i=1 βt−iZ > i µ̂k . ( 15 ) Noting that µt can be computed as usual and that the second term of ( 15 ) has a recursive form , which we denote ηt ( see appendix B ) , we rewrite µ̂ as follows : µ̂t = µt − ηt ( 16 ) ηt = βηt−1 + αβζ > t−1µ̂t−1 ( correction term ) ( 17 ) µt = ( 1− β ) t∑ i=1 βt−igi = ( 1− β ) gt + βµt−1 ( normal momentum ) ( 18 ) ζt = ( 1− β ) t∑ i=1 βt−iZi = ( 1− β ) Zt + βζt−1 . ( “ momentum ” of Taylor terms ) ( 19 ) Algorithmically , this requires maintaining 2 vectors µ and η of size n , the number of parameters , and one matrix ζ of size n× n. In practice , we find that only maintaining the diagonal of ζ can also work and can avoid the quadratic growth in n. The computation of Z also calls on computing second order derivatives ∇2f ( e.g . in ( 9 ) and ( 11 ) ) , which is impractical for large architectures . In this work , as is commonly done due to the usually small magnitude of ∇2f ( Bishop , 2006 , section 5.4 ) , we ignore them and only rely on the outer product of gradients . Ignoring the second derivatives , computing Z for TD ( 0 ) requires 3 backward passes , for g = ∇δ2 , ∇γV ( s′ ) , and ∇V ( s ) . In the online case g = δ∇V ( s ) , requiring only 2 backward passes , but in the more general minibatch of size m > 2 case , it is more efficient with modern automatic differentiation packages to do 3 backward passes than 2m passes ( see appendix B.4 ) . Finally , we note that this method easily extends to forward-view TD ( λ ) and n-step TD methods , all that is needed is to compute Z appropriately ( see appendix B.2 ) .

Summary: This paper extends the idea of momentum, commonly used in optimization literature, to Temporal Difference (TD) learning which is a widely used algorithm for policy evaluation in Reinforcement learning literature. The main challenge in this work is to account for the 'optimization bias' introduced by using momentum based updates. The authors propose to do this via a Taylor approximation to the TD update and empirically show the merits of this idea on some toy data sets as well as on an Atari game.