Year: 2020

Abstract:
A common optimization tool used in deep reinforcement learning is momentum, which consists in accumulating and discounting past gradients, reapplying them at each iteration. We argue that, unlike in supervised learning, momentum in Temporal Difference (TD) learning accumulates gradients that become doubly stale: not only does the gradient of the loss change due to parameter updates, the loss itself changes due to bootstrapping. We first show that this phenomenon exists, and then propose a first-order correction term to momentum. We show that this correction term improves sample efficiency in policy evaluation by correcting target value drift. An important insight of this work is that deep RL methods are not always best served by directly importing techniques from the supervised setting.

heading: 1 INTRODUCTION
text: 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.

heading: 1.1 REINFORCEMENT LEARNING AND TEMPORAL DIFFERENCE LEARNING
text: 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.

heading: 1.2 BIAS AND STALENESS IN MOMENTUM
text: 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).

heading: 2 CORRECTING MOMENTUM

heading: 2.1 IDENTIFYING BIAS
text: 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).

heading: 2.2 CORRECTING BIAS
text: 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).

heading: 3 EXPERIMENTAL RESULTS
text: We evaluate our proposed correction as well as the oracle (the best our method could perform) on common RL benchmarks and supervised problems. We evaluate each method with a variety of
2This computation of the gradient of Vθ(s′) may remind the reader of the so-called full gradient or residual gradient (Baird, 1995), but its purpose here is very different: we care about learning using semi-gradients, TD(0) is a principled algorithm, but we also care about how this semi-gradient evolves as parameters change, and thus we need to compute ∇θVθ(s′).
hyperparameters and multiple seeds for each hyperparameter setting. The full range of hyperparameters used as well as architectural details can be found in appendix D.
We will use the following notation for the optimizer used to train models: µ is the usual momentum, i.e. (3)&(4), and serves as a baseline. µ∗ is our oracle, defined in (6). µ̂ is the proposed correction, with updates as in (16)-(19), using the outer product approximation of Z. µ̂diag is the proposed correction, but with a diagonal approximation of Z. Throughout our figures, shaded areas are bootstrapped 95% confidence intervals over hyperparameter settings (if applicable) and runs.

heading: 3.1 SUPERVISED LEARNING
text: We first test our hypothesis, that there is bias in momentum, in a simple regression task and SVHN. For regression, we task a 4-layer MLP to regress to a 1-d function from a 1-d input. For illustration we use a smooth but non-trivial mixture of sines of increasing frequency in the domain x ∈ [−1, 1]:
x y(x) = 0.5 sin(2.14(x+ 2)) + 0.82 sin(9x+ 0.4)+
0.38 sin(12x) + 0.32 sin(38x− 0.1) (20)
Note that this choice is purely illustrative and that our findings extend to similar simple functions. We train the model on a fixed sample of 10k uniformly sampled points. We measure the mean squared error (MSE) when training a small MLP with momentum SGD versus our oracle momentum µ∗ and the outer product correction µ̂. As shown in Figure 1, we find that while there is a significant difference between the oracle and the baseline (p ≈ 0.001 from Welsh’s t-test), the difference is fairly small, and is no longer significant when using the corrected µ̂ (p ≈ 0.1).
We compare training a small convolutional neural network on SVHN (Netzer et al., 2011) with momentum SGD versus our oracle momentum µ∗ and the diagonalized correction momentum µ̂diag. The results are shown in Figure 2. We do not find significant differences except between the training loss of µ∗ and µ, and find that the oracle performs worse than the normal and corrected momentum.
From these experiments we conclude that, in supervised learning, there exists a quantifiable optimization bias to momentum, but that correcting it does not appear to offer any benefit. It improves performance only marginally at best, and degrades it at worst. This is consistent with the insight that µ∗ approximates the large batch gradient, and that large batches are often associated with overfitting or poor convergence in supervised learning (Wilson & Martinez, 2003; Keskar et al., 2016).

heading: 3.2 TEMPORAL DIFFERENCE LEARNING
text: We now test our hypotheses, that there is optimisation bias and that we can correct it, on RL problems. First, we test policy evaluation of the optimal policy on the Mountain Car problem (Singh
& Sutton, 1996) with a small MLP. We also test Acrobot and Cartpole environments and find very similar results (see appendix C). We then test our method on Atari (Bellemare et al., 2013) with convolutional networks.
Figure 3 shows policy evaluation on Mountain Car using a replay buffer (on-policy state transitions are sampled i.i.d. in minibatches). We compare the loss distributions (across hyperparameters and seeds) at step 5k, and find that all methods are significantly different (p < 0.001) from one another. Figure 4 shows online policy evaluation, i.e. the transitions are generated and learned from once, in-order, and one at a time (minibatch size of 1). There we see that the oracle µ∗ and full corrected version µ̂ are significantly different from the baseline µ (p < 0.001) and diagonalized correction µ̂diag, as well as µ from µ̂diag, while µ∗ and µ̂ are not significantly different (p > 0.1).
This suggests that the ζ matrix carries useful off-diagonal temporal information about parameters which co-vary, especially when the data is not used uniformly during learning. We test another possible explanation, which is that performance is degraded in online learning because the batch size is 1 (rather than 16 or 32 as in Figure 3). We find that a batch size of 1 does degrade µ̂diag’s performance significantly, but does not fully explain its poor online performance (see appendix Figure 10).
A central motivation of this work is that staleness in momentum arises from the change in gradients and targets. In theory, this is especially problematic for function approximators which tend to have interference, such as DNNs (Fort et al., 2019; Bengio et al., 2020), i.e. where taking a gradient step using a single input affects the output for virtually every other input. More interference means that as we change θ, the accumulated gradients computed from past θs become stale faster. We test this hypothesis by (1) measuring the drift of the value function in different scenarios, and (2) computing the cosine similarity between the corrected gradients of (13), ĝti , and their true value g t i = ∇Ji(θt).
In Figure 5 we compare the value drift of MLPs with that of linear models with Radial Basis Function (RBF) features. We compute the value drift of the target on recently seen examples, i.e. we compute the average (Vθi(s ′ i)−Vθt(s′i))2 for the last h = 2nmb/(1−β) examples, where nmb is the minibatch size. We compute the RBF features for s ∈ R2 as exp(−‖s − uij‖2/σ2) for a regular grid of uij ∈ R2 in the input domain. We find that the methods we try (even the oracle) are all virtually identical when using RBFs (see Figure 12 in appendix). We also find, as shown in Figure 5, that RBFs have very little value drift (due to their sparsity) compared to MLPs. This is consistent with our hypothesis that the method we propose is only useful if there is value drift–otherwise, there is no optimization bias incurred by using momentum. We can test this hypothesis further by artificially increasing the width of the RBFs, σ2, such that they overlap. As predicted, we find that reducing sparsity (increasing interference) increases value drift and increases the gap between our method and the baseline (Figure 6). This drift correlates with performance when changing σ2 (Figure 7).
In Figure 8 we measure the cosine similarity of the Taylor approximations ĝti = gi + Z > i (θt − θi) with the true gradients gti = ∇Ji(θt) for the last h = 2nmb/(1 − β) examples. We find that the similarity is relatively high (close to 1) throughout training but that it gets lower as the model converges to the true value V π . This is also consistent with our hypothesis that there is change (staleness) in gradients, while also validating the approach of using Taylor approximations to achieve this correction mechanism.
It is also possible to correct momentum for a subset of parameters; we try this on a per-layer basis and find that, perhaps counter-intuitively3, it is much better to correct the bottom layers (close to x) than the top layers (close to V ), although correcting all layers is the best (see appendix Figure 13).
We now apply our method on the Atari MsPacman game. We first do policy evaluation on an expert agent (we use a pretrained Rainbow agent (Hessel et al., 2018)). Since the architecture required to train an Atari agent is too large (n ≈ 4.8M) to maintain n2 values, we only use the diagonal version of our method. We also experiment with smaller (n ≈ 12.8k) models and the full correction,
3Although one may expect that changes close to V should produce more drift (it is commonly thought that bootstrapping happens in the last layer on top of relatively stable representations), the opposite is consistent with µ̂ interacting with interference in the input space, which the bottom layers have to learn to disentangle.
with similar results (see D.8). To (considerably) speed up learning, since the model is large, we additionally use per-parameter learning rates as in Adam (Kingma & Ba, 2015), where an estimate of the second moment is used as denominator in the parameter update; we denote this combination µ̂diag/ √ ∇2 + . We see in Figure 9 that our method provides a significant (p < 0.01) advantage.
Note that our method does not use frozen targets (as is usually necessary for this environment). A simple way to avoid momentum staleness and/or target drift in TD(0) is the use of frozen targets, i.e. to keep a separate θ̄ to compute Vθ̄(s
′), updated (θ̄ ← θ) at large intervals. Such a method is central to DQN (Mnih et al., 2013), but its downside is that it requires more updates to bootstrap. We find that for policy evaluation, frozen targets are much slower (both in Atari and simple environments) than our baseline (see appendix Figures 16 and 17).
We finally apply our method to control, first in combination with a policy gradient method, PPO (Schulman et al., 2017), in its policy evaluation part, and second with a 5-step Sarsa method, but find no improvement (or marginal at best) in either setting. As in simpler environ-
ments we measure cosine similarity and value drift. We find low similarity (≈ 0.1) but a 2× to 3× decrease in drift using our method, suggesting that while our method corrects drift, its effect on policy improvement is minor. We suspect that in control, even with a better value function, other factors such as exploration or overestimation come into play which are not addressed by our method.

heading: 4 DISCUSSION
text: We propose a method which improves momentum, when applied to DNNs doing TD learning, by correcting gradients for their staleness via an approximate Taylor expansion. We show that correcting this staleness is particularly useful when learning online using a TD objective, but less so for supervised learning tasks. We show that the proposed method corrects for value drift in the bootstrapping target of TD, and that the proposed approximate Taylor expansion is a useful tool that aligns well enough with the true gradients.
Shortcomings of the method In its most principled form, the proposed method requires computing a second-order derivative, which is impractical in most deep learning settings. While we do find experimentally that ignoring its contribution has a negligible effect in toy settings, we are unable to verify this for larger neural networks. Compared to the usual momentum update, the proposed method also requires performing two additional backward passes, as well as storing 2n + n2 additional scalars. This can be improved with a diagonal approximation requiring only 3n scalars, but this approximation is not as precise, which is currently an obstacle for large architectures. While these extra computations improve sample complexity, i.e. we get more out of each sample, they do not necessarily improve convergence speed in computation time (although, our method is suited to GPU parallelism and has reasonable speed even for large n2s, including in Atari). Sample complexity may be particularly important in settings where samples are expensive to get (e.g. real world systems), but less so when they are almost free (e.g. simulations).
Incorporating TD in optimization One meta-result stands out from this work: something is lost when naively applying supervised learning optimization tools to RL algorithms. In particular here, by simply taking into account the non-stationarity of the TD objective, we successfully improve the sample complexity of a standard tool (momentum), and demonstrate the potential value of incorporating elements of RL into supervised learning tools. The difference explored here with momentum is only one of the many differences between RL and supervised learning, and we believe there remain plenty of opportunities to improve deep learning methods by understanding their interaction with the peculiarities of RL.

heading: A RELATED WORK
text: To the best of our knowledge, no prior work attempts to derive a corrected momentum-SGD update adapted to the Temporal Difference method. That being said, a wealth of papers are looking to accelerate TD and related methods.

heading: On momentum, traces, and gradient acceleration in TD
text: From an RL perspective, our work has some similarity to the so-called eligibility traces mechanism. In particular, in the True Online TD(λ) method of van Seijen & Sutton (2014), the authors derive a strict-online update (i.e. weights are updated at every MDP step, using only information from past steps, rather than future information as in the λ-return perspective) where the main mechanism of the derivation lies in finding an update by assuming (at least analytically) that one can “start over” and reuse all past data iteratively at each step of training, and then from this analytical assumption derive a recursive update (that doesn’t require iterating through all past data). The extra values that have to be kept to compute the recursive updates are then called traces. This is akin to how we conceptualize µ∗, (6), and derive µ̂.
The conceptual similarities of the work of van Seijen & Sutton (2015) with our work are also interesting. There, the authors analyse what “retraining from scratch” means (i.e., again, iteratively restarting from θ0 ∈ Rm) but with some ideal target θ∗ (e.g. the current parameters) by redoing sequentially all the TD(0) updates using θ∗ for all the n transitions in a replay buffer, costing O(nm). They derive an online update showing that one can continually learn at a cost of O(m2) rather than payingO(nm) at each step. The proposed update is also reminiscent of our method in that it aims to perform an approximate batch update without computing the entire batch gradient, and also maintains extra momentum-like vectors and matrices. We note that the derivation there only works in the linear TD case.
In a way, such an insight can be found in the original presentation of TD(λ) of Sutton (1988), where the TD(λ) parameter update is written as (equation (4) in the original paper, but with adapted notation):
∆θt = α[rt + γVθt(st+1)− Vθt(st)] t∑
k=1
λt−k∇θtVθt(sk)
Remark the use of θt in the sum; in the linear case since ∇θtVθt(sk) = φ(sk), the sum does not depend on θt and thus can be computed recursively. A posteriori, if one can find a way to cheaply compute ∇θtVθt(sk) ∀k, perhaps using the method we propose, it may be an interesting way to perform TD(λ) using a non-linear function approximator.
Our analysis is also conceptually related to the work of Schapire & Warmuth (1996), where a worstcase analysis of TD∗(λ) is performed using a best-case learner as the performance upper bound. This is similar to our momentum oracle; just as the momentum oracle is the ”optimal” approximation of the accumulation gradients coming from all past training examples, the best-case learner of Schapire & Warmuth (1996) is the set parameters that is optimal when one is allowed to look at all past training examples (in contrast to an online TD learner).
Before moving on from TD(λ), let us remark that eligibility traces and momentum, while similar, estimate different quantities. The usual (non-replacing) traces estimate the exponential moving average of the gradient of Vθ, while momentum does so for the objective J (itself a function of Vθ):
et = (1− λ) t∑ k λt−k∇θVθ, µt = (1− β) t∑ k βt−k∇θJ(Vθ)
Our method also has similarities with residual gradient methods (Baird, 1995). A recent example of this is the work of Zhang et al. (2019), who adapt the residual gradient for deep neural networks. Residual methods learn by taking the gradient of the TD loss with respect to both the current value and the next state value V (S′), but this comes at the cost of requiring two independent samples of S′ (except in deterministic environments).
Similarly, our work is related to the “Gradient TD” family of methods (Sutton et al., 2008; 2009). These methods attempt to maintain an expectation (over states) of the TD update, which allows to
directly optimize the Bellman objective. While the exact relationship between GTD and “momentum TD” is not known, they both attempt to maintain an “expected update” and adjust parameters according to it; the first approximates the one-step linear TD solution, while the latter approximates the one-step batch TD update. Note that linear GTD methods can also be accelerated with momentum-style updates (Meyer et al., 2014), low-rank approximations for part of the Hessian (Pan et al., 2016), and adaptive learning rates (Gupta et al., 2019).
More directly related to this work is that of Sun et al. (2020), who show convergence properties of a rescaled momentum for linear TD(0). While most (if not every) deep reinforcement learning method implicitly uses some form of momentum and/or adaptive learning rate as part of the deep learning toolkit, Sun et al. (2020) properly analyse the use of momentum in a (linear) TD context. Gupta (2020) also analyses momentum in the context of a linear TD(0) and TD(λ), with surprising negative results suggesting naively applying momentum may hurt stability and convergence in minimal MDPs.
Another TD-aware adaptive method is that of Romoff et al. (2020), who derive per-parameter adaptive learning rates, reminiscent of RMSProp (Hinton et al., 2012), by considering a (diagonal) Jacobi preconditioning that takes into account the bootstrap term in TD.
Finally, we note that, as far as we know, recent deep RL works all use some form of adaptive gradient method, Adam (Kingma & Ba, 2015) being an optimizer of choice, closely followed by RMSProp (Hinton et al., 2012); notable examples of such works include those of Mnih et al. (2013), Schulman et al. (2017), Hessel et al. (2018), and Kapturowski et al. (2019). We also note the work of Sarigül & Avci (2018), comparing various SGD variants on the game of Othello, showing significant differences based on the choice of optimizer.
On Taylor approximations Balduzzi et al. (2017) note that while theory suggests that Taylor expansions around parameters should not be useful because of the ”non-convexity” of ReLU neural networks, there nonetheless exists local regions in parameter space where the Taylor expansion is consistent. Much earlier work by Engelbrecht (2000) also suggests that Taylor expansions of small sigmoid neural networks are easier to optimize. Using Taylor approximations around parameters to find how to prune neural networks also appears to be an effective approach with a long history (LeCun et al., 1990; Hassibi & Stork, 1993; Engelbrecht, 2001; Molchanov et al., 2016).
On policy-gradient methods and others While not discussed in this paper, another class of methods used to solve RL problems are PG methods. They consist in taking gradients of the objective wrt a directly parameterized policy (rather than inducing policies from value functions). We note in particular the work of Baxter & Bartlett (2001), who analyse the bias of momentum-like cumulated policy gradients (referred to as traces therein), showing that β the momentum parameter should be chosen such that 1/(1− β) exceeds the mixing time of the MDP. Let us also note the method of Vieillard et al. (2020), Momentum Value Iteration, which uses the concept of an exponential moving average objective for a decoupled (with its own parameters) action-value function from which the greedy policy being evaluated is induced. This moving average is therein referred to as momentum; even though it is not properly speaking the optimizational acceleration of Polyak (1964), its form is similar.

heading: B COMPLETE DERIVATIONS

heading: B.1 DERIVATION OF MOMENTUM CORRECTION
text: In momentum, µt can be rewritten as:
µt = (1− β) t∑ i=1 βt−i∇θiJi(θi) (21)
We argue that an ideal unbiased momentum µ∗t would approximate the large batch case by only discounting past minibatches and using current parameters θt rather than past parameters θi:
µ∗t def = (1− β) t∑ i=1 βt−i∇θtJi(θt) (22)
Note that the only difference between (21) and (22) is the use of θi vs θt. The only way to compute µ∗t exactly is to recompute the entire sum after every parameter update. Alternatively, we could somehow approximate this sum. Below we will define gti = ∇θtJi(θt), which we will then approximate with ĝti . We will then show that this approximation has a recursive form which leads to an algorithm.
We want to correct accumulated past gradients gi to their “ideal” form gti , as above. We do so with their Taylor expansion around θ, and we write the correction of the gradient gi computed at learning time i corrected at time t as:
gti = gi(θi + ∆θ(t− 1; i)) = gi +∇θgTi ∆θ(t− 1; i) + o(‖∆θ‖22) (23) ≈ ĝti = gi + ZTi ∆θ (24)
where ∆θ(t; i) = θt − θi, gi = ∇θiJi, Zi = ∇θigi. Here we are agnostic of the particular form of Z, which will depend on the loss and learning algorithm, and is not necessarily the so-called Hessian. To see why this is the case, and for the derivation of Z for the squared loss, cross-entropy and TD(0), see section B.2.
Let’s now express ĝti in terms of gi and updates µt. At each learning step, the current momentum µt is multiplied with the learning rate α to update the parameters, which allows us to more precisely write ĝti :
θt = θt−1 − αµt = θ0 − α t∑ i=1 µi (25)
∆θ(t; i) = θt − θi = θ0 − α t∑
k=1
µk − θ0 + α i∑
k=1 µk = −α t∑ k=i µk (26)
ĝti = gi + Z T i ∆θ(t− 1; i) = gi − αZTi t−1∑ k=i µk (27)
We can now write µ̂t, the approximated µ∗t using ĝ t i :
µ̂t = (1− β)gt + (1− β) t−1∑ i=1 βt−iĝti (28)
= (1− β)gt + (1− β) t−1∑ i=1 βt−igi − (1− β) t−1∑ i=1 βt−iαZTi t−1∑ k=i µ̂k (29)
= µt − α(1− β) t−1∑ i=1 t−1∑ k=i βt−iZTi µ̂k extract µ (30)
= µt − α(1− β) t−1∑ k=1 k∑ i=1 βt−iZTi µ̂k change the sum indices for convenience (31) = µt − ηt extract ηt (32)
Let’s try to find a recursive form for ηt:
ηt − ηt−1 = α(1− β) t−1∑ k=1 k∑ i=1 βt−iZTi µ̂k − α(1− β) t−2∑ k=1 k∑ i=1 βt−1−iZTi µ̂k
= α(1− β) ( t−2∑ k=1 k∑ i=1 (βt−iZTi µ̂k − βt−1−iZTi µ̂k) + t−1∑ i=1 βt−iZTi µ̂t−1 )
= α(1− β) ( t−2∑ k=1 k∑ i=1 (β − 1)βt−1−iZTi µ̂k + t−1∑ i=1 βt−iZTi µ̂t−1 )
= (β − 1)ηt−1 + αβ(1− β) t−1∑ i=1 βt−1−iZTi µ̂t−1
Let ζt = (1− β) t∑ i=1 βt−iZi
= (β − 1)ηt−1 + αβζt−1µ̂t−1 ηt − ηt−1 = −(1− β)ηt−1 + αβζTt−1µ̂t−1
ηt = βηt−1 + αβζ T t−1µ̂t−1
We can now write the full update as:
µ̂t = µt − ηt ηt = βηt−1 + αβζ T t−1µ̂t−1
µt = (1− β) t∑ i=1 βt−igi = (1− β)gt + βµt−1
ζt = (1− β) t∑ i=1 βt−iZi = (1− β)Zt + βζt−1
with η0 = 0.
This corresponds to an algorithm where one maintains η, µ and ζ.

heading: B.2 DERIVATION OF TAYLOR EXPANSIONS AND Z
text: For least squares regression, the expansion around θ of g(θ) is simply the Hessian of the loss δ2:
δ2 = 1
2 (y − fθ(x))2 (33)
g(θ) = ∇θδ2 (34) g(θ + ∆θ) = g(θ) +HTJ ∆θ + o(‖∆θ‖22) (35)
This Hessian has the form HJ = ∇θf ⊗∇θf + δ2∇2θf (36)
δ2 being small when a neural network is trained and ∇2 being expensive to compute, a common approximation to Hj is to only rely on the outer product. Thus we can write:
Zreg = ∇θf ⊗∇θf (37)
For classification, or categorical crossentropy, Z has exactly the same form, but where f is the log-likelihood (i.e. output of a log-softmax) of the correct class.
For TD(0), the expansion is more subtle. Since TD is a semi-gradient method, when computing g(θ), gradient of the TD loss δ2, we ignore the bootstrap target’s derivative, i.e. we hold Vθ(s′) constant (unlike in GTD). On the other hand, when computing the Taylor expansion around θ, we do care about how Vθ(s′) changes, and so its gradient comes into play:
g(θ) = (Vθ(x)− γVθ(x′)− r)∇θVθ(x) g(θ + ∆θ)i = gi(θ) + (∇θVθ(x)− γ∇θVθ(x′))∇θiVθ(x) ·∆θ + δ∇θ∇θiVθ(x) ·∆θ g(θ + ∆θ) = g(θ) + ((∇θVθ(x)− γ∇θVθ(x′))⊗∇θVθ(x))T ∆θ + δ∇2θVθ(x)T∆θ
Similarly for TD we ignore the second order derivatives and write:
ZTD = (∇θVθ(x)− γ∇θVθ(x′))⊗∇θVθ(x)
For an n-step TD objective, the correction is very similar:
ZTD(n) = (∇θVθ(xt)− γn∇θVθ(xt+n))⊗∇θVθ(xt)
For a forward-view TD(λ) objective this is also similar, but more expensive:
ZTD(λ) = (∇θVθ(xt)− (1− λ)∇θ(γλVθ(xt+1) + γ2λ2Vθ(xt+2) + ...))⊗∇θVθ(xt)
Finally, to avoid maintaining n× n values for Z, it is possible to only maintain the diagonal of Z or some block-diagonal approximation, at some performance cost.

heading: B.3 THE LINEAR CASE
text: Much of RL analysis is done in the linear case, as it is easier to make precise statements about convergence there. We write down such a case below, but we were unfortunately unable to derive any interesting analyses from it.
Recall that the proposed method attempts to approximate µ∗:
µ∗t = (1− β) t∑ i=1 βt−i∇θtJi(θt)
which in the linear case Vθ(x) = θ>φ(x) is simply:
= (1− β) t∑ i=1 βt−iδiφi
which depends on θi through δi the TD error. As such we can write Z as:
ZTD = (∇θVθ(x)− γ∇θVθ(x′))⊗∇θVθ(x) = (φ− γφ′)φ>
which interestingly does not depend on θ. To the best of our knowledge, and linear algebra skills, this lack of dependence on θ does not allow for a simplification of the proposed correction mechanism (in contrast with linear eligibility traces) due to the pairwise multiplicative t, i dependencies that emerge.

heading: B.4 MINIBATCHING
text: It is possible to only do 3 (or 2 in the supervised case) “backward” calls (e.g. torch.autograd.grad in PyTorch) even when computing this correction for several examples. This is in contrast to other adaptive gradient methods (Romoff et al., 2020) which require per-input
gradients (these can be computed efficiently with methods such as Dangel et al. (2020), but remain more expensive).
Because the gradient we care to correct is the gradient of the minibatch, i.e. the sum of gradients, then we only have to correct that sum, and not individual gradients. This allows us to compute Z out of the gradients of the minibatch instead of the individual gradients. When the gradient is:
gt = ∇θJt(θ) = ∇θ m∑ i Jt,i(θ) = m∑ i ∇θJt,i(θ) (38)
then the first derivative used in its Taylor expansion is, by linearity of the derivative:
∂gt ∂θ = m∑ i ∇θfi ⊗∇θfi + δt,iHf (39)
= ∇θ( ∑ fi)⊗∇θ( ∑ fi) +∇2θJt,i (40)
Ignoring the Hessian we only have to use the gradients obtained for the sum (or the mean if J is a mean) of losses of the minibatch, which requires only one backpropagation regardless of batch size. The derivation is analogous for TD(0). Note that this would not work if we used the second order derivative due to the multiplicative δt,i factors.

heading: C ADDITIONAL FIGURES

heading: D HYPERPARAMETER SETTINGS AND COMMENTS
text: All experiments are implemented using PyTorch (Paszke et al., 2019). We use Leaky ReLUs throughout. All experimental code is available in supplementary materials.
On Leaky ReLUs: we did experiment with ReLU, tanh, ELU, and SELU activation units. The latter 3 units have more stable Taylor expansions for randomly initialized neural networks, but in terms of experimental results, Leaky ReLUs were always significantly better.

heading: D.1 FIGURE 1
text: We use an MLP with 4 layers of width nh.
We use the cross-product of nh ∈ {8, 16, 32}, β ∈ {0.9, 0.99}, α ∈ {0.005, 0.01}, nmb ∈ {4, 16, 64}.

heading: D.2 FIGURE 2
text: We use a convolutional model with the following sequence of layers, following PyTorch convention: Conv2d(3, nh, 3, 2, 1), Conv2d(nh, 2nh, 3, 2, 1), Conv2d(2nh, 2nh, 3, 2, 1), Conv2d(2nh, nh, 3, 1, 1), Flatten(), Linear(16nh, 4nh ), Linear(4nh, 4nh), Linear(4nh, 10), with LeakyReLUs between each layer.
We use the cross-product of nh ∈ {8, 16}, β ∈ {0.9, 0.99}, α ∈ {0.005, 0.01}, nmb ∈ {4, 16, 64}.

heading: D.3 FIGURE 3 AND 16
text: We use an MLP with 4 layers of width nh.
We use the cross-product of nh ∈ {8, 16, 32}, β ∈ {0.9, 0.99}, α ∈ {0.5, 0.1, 0.05}, nmb ∈ {4, 16, 64}.

heading: D.4 FIGURE 4
text: We use an MLP with 4 layers of width nh.
We use the cross-product of nh ∈ {16}, β ∈ {0.9, 0.99}, α ∈ {0.005, 0.001, 0.0005}, nmb ∈ {1}.

heading: D.5 FIGURE 5 AND 12
text: We use a linear layer on top of a grid RBF representation with each gaussian having a variance of σ2/ngrid.
For the RBF we use ngrid = 20, α = 0.1, nmb = 16, β = 0.99. For the MLP we use nh = 16, α = 0.1, nmb = 16, β = 0.99.

heading: D.6 FIGURE 6 AND 7
text: We use RBFs with σ2 ∈ {4, 3.5, 3, 2.5, 2, 1.5, 1.25, 1, 0.75, 0.5}, ngrid ∈ {10, 20}, α ∈ {0.1, 0.01}.

heading: D.7 FIGURE 8
text: We use an MLP with 4 layers of width nh = 16, α = 0.1, nmb = 16, β = 0.95.

heading: D.8 FIGURE 9 AND 17
text: We use the convolutional model of Mnih et al. (2013) with the same default hyperparameters, following PyTorch convention: Conv2d(4, nh, 8, stride=4, padding=4), Conv2d(nh, 2nh, 4, stride=2, padding=2), Conv2d(2nh, 2nh, 3, padding=1), Flatten(), Linear(2nh×12×12, 16nh), Linear(16nh, nacts).
We use nh = 64, nmb = 32, for Adam we use α = 5× 10−5 and β = 0.99, for our method we use α = 10−4 and β = 0.9 (these choices are the result of a minor hyperparameter search of which the best values were picked, equal amounts of compute went towards our method and the baseline so as to avoid “poor baseline cherry picking”). For the frozen target baseline we update the target every 2.5k steps.
We were able to find similar differences between Adam and our method with a much smaller model, but using the full correction instead of the diagonal one. Although the full correction outperforms Adam when both use this small model, using so few parameters is not as accurate as the original model described above, and we omit these results.
This small model is: Conv2d(4, nh, 3, 2, 1), Conv2d(nh, nh, 3, 2, 1), Conv2d(nh, nh, 3, 2, 1), Conv2d(nh, nh, 3, 2, 1), Conv2d(nh, nacts, 6). For nh = 16 (which we used for experiments), this model has a little less than 16k parameters, making Z about 250M scalars. While this is large, this easily fits on a modern GPU, and the extra computation time required for the full correction mostly still comes from computing two extra backward passes, rather than from computing Z and the correction.
This is beyond the scope of this paper, but is worth of note: Interestingly this small model still works quite well for control (both with Adam and our method). We have not tested this extensively, but, perhaps contrary to popular RL wisdom surrounding Deep Q-Learning, we were able to train decent MsPacman agents with (1) no replay buffer, but rather 4 or more parallel environments (the more the better) as in A2C (Mnih et al., 2016) (2) no frozen target network (3) a Q network with only 16k parameters rather than the commonplace 4.8M-parameter Atari DQN model. The only “trick” required is to use 5-step Sarsa instead of 1-step TD (as suggested by the results of Fedus et al. (2020), although in their case a replay buffer is used).

heading: D.9 FIGURE 10
text: We use the same configuration than Figure 3 with nmb ∈ {1, 4, 16, 64}.
D.10 FIGURE 11
We use the same configuration than Figure 3 with β ∈ {.8, .9, .95, .975, .99, .995, .999}.
D.11 FIGURE 13
We use an MLP with 6 layers of width nh = 16, α = 0.1, nmb = 16, β = 0.9. We vary which layers get used in the corrected momentum; see caption.
D.12 FIGURE 14 AND 15
We use the same settings as Figure 3, with the exception that we do 5-step TD in Figure 14, for Acrobot, and 3-step TD for Figure 15, for Cartpole. Using n > 1 appears necessary for convergence for both our method and the baseline.

Review summaries:
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. 

This paper proposes a modification to momentum. Extra terms are added to account for the drift in parameters through the gradient updates that contribute to the momentum. The effect of this modification is studied on supervised learning and TD learning and on different representations. The conclusion is that while the original momentum is good enough for supervised learning, this modification matters when there is bootstrapping, especially on representations with high interference.