Adam on Local Time: Addressing Nonstationarity in RL with Relative Adam Timesteps

Authors: benjamin  ellis
matthew t jackson
andrei  lupu
alexander d goldie
mattie  fellows
shimon  whiteson
jakob n foerster
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 11:18:56

Abstract: In reinforcement learning (RL), it is common to apply techniques used broadly in 
machine learning such as neural network function approximators and momentum-based optimizers. However, such tools were largely developed for supervised learning rather than nonstationary RL, leading practitioners to adopt target networks, clipped policy updates, and other RL-specific implementation tricks to combat this mismatch, rather than directly adapting this toolchain for use in RL. 
In this paper, we take a different approach and instead address the effect of nonstationarity by adapting the widely used Adam optimiser. 
We first analyse the impact of nonstationary gradient magnitude --- such as that caused by a change in target network --- on Adam's update size, demonstrating that such a change can lead to large updates and hence sub-optimal performance.
To address this, we introduce Adam-Rel.
Rather than using the global timestep in the Adam update, Adam-Rel uses the *local* timestep within an epoch, essentially resetting Adam's timestep to 0 after target changes.
We demonstrate that this avoids large updates and reduces to learning rate annealing in the absence of such increases in gradient magnitude. Evaluating Adam-Rel in both on-policy and off-policy RL, we demonstrate improved performance in both Atari and Craftax.
We then show that increases in gradient norm occur in RL in practice, and examine the differences between our 
theoretical model and the observed data.

Main paper: 
 introduction: 
 Reinforcement Learning (RL) aims to learn robust policies from an agent's experience. This has the potential for large scale real-world impact in areas such as autonomous driving or improving logistic chains. Over the last decade, a number of breakthroughs in supervised learning-such as convolutional neural networks and the Adam optimizer-have expanded the deep learning toolchain and been transferred to RL, enabling it to begin fulfilling this potential. However, since RL agents are continuously learning from new data they collect under their changing policy, the optimisation objective is fundamentally nonstationary. Furthermore, temporal difference (TD) approaches bootstrap the agent's update from its own value predictions, exacerbating the nonstationarity in the objective function. This is in stark contrast to the stationary supervised learning setting for which the deep learning toolchain was originally developed. Therefore, to apply these tools successfully, researchers have developed a variety of implementation tricks on top of this base to stabilise training [8, 6, 5] . This has resulted in a proliferation of little-documented design choices that are vital for performance, contributing to the reproducibility crisis in RL [9] .We believe that in the long term, a more robust approach is to augment this toolchain for RL, rather than building on top of it. To this end, in this paper we examine the interaction between nonstationarity and the Adam optimizer [7] . Adam's update rule, where equations are applied element-wise (i.e. per parameter), is defined bym t = β 1 m t-1 + (1 -β 1 )g t , mt = m t (1 -β 1 t ) , v t = β 2 v t-1 + (1 -β 2 )g t 2 , vt = v t (1 -β 2 t ) , u t = mt √ vt + ϵ , θ t = θ t-1 -αu t .Here, g t is the gradient, θ t a parameter to be optimized, and α the learning rate. The resulting update is the ratio of two different momentum terms: one for the first moment, m t , and one for second moment, v t , of the gradient. These terms use different exponential decay coefficients, β 1 and β 2 . Under stationary gradients, the (1 -β i ) weighting ensures that, in the limit, the overall magnitude of the two momenta is independent of the value chosen for each of the coefficients. However, since both momentum estimates are initialised to 0, they must be renormalised for a given (finite) timestep t, to account for the "missing parts" of the geometric series [7] , resulting in vt and mt .Crucially, t counts the update steps since the beginning of training and thus bakes in the assumption of stationarity that is common in supervised learning. In particular, this renormalisation breaks down if the loss is nonstationary. Consider a task change late in training, which results in gradients orders of magnitudes higher than those of the prior (near convergence) task. Clearly, this is analogous to the situation at the beginning of training where all momentum estimates are 0. However, the t parameter, and therefore the renormalisation, does not account for this.In this paper, we demonstrate that changes in the gradient scale can lead to large updates that persist over a long horizon. Previous work [10, 11] has suggested that old momentum estimates can contaminate an agent's update and propose resetting the entire optimizer state when the target changes as a solution. However, by discarding previous momentum estimates, we hypothesise that this approach needlessly sacrifices valuable information for optimization. Instead, we propose retaining momentum estimates and only resetting t, which we refer to as Adam-Rel. In the limit of gradient sparseness, we show that the Adam-Rel update size remains bounded, converging to 1 in the limit of a large gradient, unlike Adam. Furthermore, if such gradient magnitude increases do not occur, Adam-Rel reduces to learning rate annealing, a common method for stabilising optimization.When evaluated against the original Adam and Adam with total resets, we demonstrate that our method improves PPO's performance in Craftax-Classic [12] and the Atari-57 challenge from the Arcade Learning Environment [13] . Additionally, we demonstrate improved performance in the off-policy setting by evaluating DQN on the Atari-10 suite of tasks [14] . We then examine the gradients in practice and show that there are significant increases in gradient magnitude following changes in the objective. Finally, we examine the discrepancies between our theoretical model and observed gradients to better understand the effectiveness of Adam-Rel.2 Background 
reinforcement learning: 
 Definition Reinforcement learning agents learn a policy π in a Markov Decision Process [15, MDP] , a tuple M = ⟨S, A, T , R, γ⟩ where S is the set of states, A is the set of actions, T : S × A → P(S) is the transition function, R : S × A → R is the reward function and γ is the discount factor. At each timestep t, the agent observes a state s t ∈ S and takes an action a t drawn from π(•|s t ) before transitioning to a new state s t+1 ∈ S and receiving reward r t drawn from R(s t , a t ). The goal of the agent is to maximise the expected discounted returnE π,T [ ∞ t=0 γ t r t ].Nonstationarity in RL In contrast with supervised learning, where a single stationary objective is typically optimised, reinforcement learning is inherently nonstationary. Updates to the policy induce changes not only in the distribution of observations seen at a given timestep, but also the return distribution, and hence value function being optimised. This arises regardless of how these updates are performed. However, one particular reason for nonstationarity in RL is the use of bootstrapped value estimates via TD learning [15] , which optimises the below objectiveL(θ) = [sg {r t + γV π θ (s t+1 )} -V π θ (s t )] 2 ,where sg is the stop-gradient operator. In this update, the target r t + γV π θ (s t+1 ) depends on the parameters θ and therefore changes as these are updated.These target changes can either be more gradual, as in the case of continuous updates to the value function in TD learning, or more abrupt, as in the case of the use of target networks in DQN.In this work, we focus on abrupt objective changes; changes of objectives that do not involve a smoothing method such as Polyak averaging [1] , and the resulting sudden change of supervised learning problem. More explicitly, we consider optimising a stationary loss function L(θ, ϕ), where θ are the parameters to be optimised and ϕ is the other parameters of the loss function (such as the parameters of a value network), which are not updated throughout optimisation, but does not include the training data.We consider a setting where at a certain timestep t in our training, we transition from optimising L(θ t , ϕ 1 ) to optimising L(θ t+1 , ϕ 2 ) for some ϕ 1 , ϕ 2 . Such individual objectives are still nonstationary. For example, significant changes in the policy would induce changes in the data distribution, which would then affect the underlying loss landscape, but we do not consider such non-stationarity in this work. This setting is very common throughout RL. Bootstrapped value estimates are the most common cause of this, but it also occurs in PPO's actor update, where each new rollout induces a different supervised learning problem due to the actor and critic updates. This is optimised for a fixed number of updates before collecting new data.We refer to these sequences of supervised learning problems as sequentially-optimised stationary objectives. In this work, we use this framing to propose an approach that is consistent throughout each stationary period of optimization and applies corrections to make optimization techniques valid when nonstationarity is introduced via objective changes. Bengio et al. [11] propose the gradient contamination hypothesis, which states that current optimizer momentum estimates can point in the opposite direction to the gradient following a change in objective, thereby hindering optimization. A previous approach to this problem is that of Asadi et al. [10] , where they propose resetting Adam's momentum estimates and timestep to 0 throughout training. We refer to this method as Adam-MR. Finally, Dohare et al. [16] propose setting the Adam hyperparameters to equal values, such that β 1 = β 2 , suggesting that this can help avoid performance collapse.Proximal Policy Optimization Proximal Policy Optimization [4, PPO] is a policy optimisation based RL method. It uses a learned critic V π ϕ trained by a TD loss to estimate the value function, and a clipped actor update of the formEQUATIONwhere the policy ratio r (θ,t) = πθ (at|st) π(at|st) is the ratio of the stochastic policy to optimise πθ and π, the previous policy. A π is the advantage, which is typically estimated using generalised advantage estimation [17] . Clipping the policy ratio aims to avoid performance collapse by preventing policy updates larger than ϵ.Optimisation of the PPO objective proceeds by first rolling out the policy to collect data, and then iterating over this data in a sequence of epochs. Each of these epochs splits the collected data into a sequence of mini-batches, over which the above update is calculated.Momentum [1, 2] is a method for enhancing stochastic gradient descent by accumulating gradients in the direction of repeated improvement. The typical formulation of momentum for each element i ism i t = βm i t-1 + g i t , θ i t = θ i t-1 -αm i t ,where β is the momentum coefficient, g t ∈ R n is the gradient at the current step, m t ∈ R n is the gradient incorporating momentum, α is the scalar learning rate and θ ∈ R n are the parameters to be optimised. With momentum, update directions with low curvature have their contribution to the gradient amplified, considerably reducing the number of steps required for convergence.In the introduction, we described the update equations for Adam [7] , the most popular optimizer that uses momentum. Adam's update is designed to keep its updates within a trust region, which depends on a learning rate α.Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: We discuss limitations, along with suggestions for future work, in Section 7.We also examine how our theoretical assumptions match practice in Section 5.4. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
nonstationary optimization with adam: 
 We now investigate the effect of nonstationarity on Adam by analysing its update rule after a sudden change in gradient. As a simplified model of gradient instability, we assume optimization with Adam starts at timestep t = -t ′ with a constant gradient g i -t ′ = g, 0 < g < ∞ until timestep 0. Following t = 0, we model instability by increasing the gradient by a factor of k, as might occur in a nonstationary optimization setting. This givesEQUATIONFor larger values of t ′ , the short term effects of Adam's initialisation on the momentum terms dissipate and mt and vt converge to stable values. By taking the limit of t ′ → ∞, we investigate the effect of a sudden change in gradient g i t on the update size u i t after a long period of training. This allows for any effects from the initialisation of momentum terms m-t ′ ,t and v-t ′ ,t to dissipate: Theorem 3.1. Assume that ϵ = 0. Let g i t be defined as in Equation (2) and mi -t ′ ,t and vi -t ′ ,t be the momentum terms at timestep t given Adam starts at timestep -t ′ . It follows that:EQUATIONProof. See Appendix A.For large k, Theorem 3.1 proves that the element-wise momentum term after the change in gradient at t = 0 is approximately 1-β1 √ 1-β2 . For the most commonly used values of β 1 = 0.9 and β 2 = 0.999, this is √ 10, which is much larger than the intended unit update which Adam is designed to maintain. The top plot in Figure 1 , which shows the Adam update size against t for different values of k, demonstrates that the update peaks significantly higher than the desired 1 before slowly converging back to 1.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: We provide anonymised links to our code in the Appendix, and only run on open-source environments, allowing for our experiments to be reproduced.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
adam with relative timesteps: 
 To fix the problems analysed in the previous section, we introduce Adam-Rel. At the start of each new supervised learning problem, Adam-Rel resets Adam's t parameter to 0, rather than incrementing it from its previous value. This one-line change is illustrated for PPO in Algorithm 1.At the start of training, both momentum terms in Adam are 0. Therefore, at the first timestep, when the first gradient is encountered, the magnitude of the gradient is infinite relative to the current momentum estimate. As explained in Section 3, this induces a large update. However, dividing the momentum estimates by (1 -β 1 t ) and (1 -β 2 t ) fixes this issue by correcting for this sparsity. Therefore, by resetting t to 0, Adam handles changes in gradient magnitude resulting from the change of supervised learning problem.If we examine the same update as in the previous section adjusted by Adam-Rel, assuming that we reset Adam's t just before the gradient scales to kg, we find it comes to Assumes that optimization starts at time -t ′ , which is large, and that the gradients up until time 0 are g and then there is an increase in the gradient to kg.EQUATION)Algorithm 1 Pseudocode for PPO with Adam, Adam-Rel, and Adam-MR. m 0 = 0, v 0 = 0, t = 0 ▷ Initialise Adam state j = 0 ▷ Initialise number of updates to 0for k = 1 to K do Rollout policy π θ k to collect batch D t = 0 t = 0, m j = 0, v j = 0 for epoch = 1 to E do for mini-batch B in D do g j = ∇ θj-1 L PPO (θ j-1 ) + L TD (θ j-1 ) j = j + 1 t = t + 1 m j = β 1 m j-1 + (1 -β 1 )g j v j = β 2 v j-1 + (1 -β 2 )g j 2 mj = mj (1-β1 t ) vj = vj (1-β2 t ) θ j = θ j-1 -α mj √ vj +ϵend for end for end for As k → ∞, this tends to 1. This means that Adam-Rel ensures approximately unit update size in the case of a large increase in magnitude in the gradient, at the expense of a potentially smaller update at the point t is reset. Figure 1 shows the update size of Adam-Rel as t -t ′ increases. The update size is smaller at the start, but never reaches significantly above 1.However, the above analysis does not show how Adam and Adam-Rel differ in practice, where large changes in gradient magnitude may not occur. Examining the bottom of Figure 1 , we can see that for lower values of k, Adam-Rel rapidly decays the update size before increasing it. Functionally, this behaves like a learning rate schedule. Over a short horizon (e.g., 16 steps is common in PPO), this effect is similar to learning rate annealing, whilst over a longer horizon (e.g., approximately 1000 steps in DQN) it is akin to learning rate warmup, both of which are popular techniques in optimising stationary objectives. Therefore, the benefits of Adam-Rel are twofold: first, it guards against large increases in gradient magnitude by capping the size of potential updates, and secondly, if such large gradient increases do not occur, it reduces to a form of learning rate annealing, which is commonly employed in optimising stationary objectives. 
experimental setup: 
 To evaluate Adam-Rel, we explore its impact on DQN and PPO, two of the most popular algorithms in off-policy and on-policy RL respectively.To do so, we first train DQN [18, 19] agents with Adam-Rel on the Atari-10 benchmark for 40M frames, evaluating performance against agents trained with Adam and Adam-MR. We then extensively evaluate our method's impact on PPO [4, 19, 20] , training agents on Craftax-Classic-1B [12] -a JAX-based reimplementation of Crafter [21] where the agent is allocated 1 billion environment interactions-and the Atari-572 suite [13] for 40 million frames. In doing so, our benchmarks respectively evaluate the performance of Adam-Rel on exceedingly long training horizons and its ) for PPO and Adam, Adam-MR and Adam-Rel for DQN on Atari-57 and Atari-10 respectively. Atari-10 uses a subset of Atari tasks to estimate median performance across the whole suite. Details can be found in [14] .Error bars are 95% stratified bootstrapped confidence intervals. Results are across 10 seeds except for Adam (β 1 = β 2 ), which is 3 seeds.robustness when applied to a diverse range of environments. We then analyse the differences between Adam-Rel and Adam's updates. We compare 8 seeds on the Craftax-Classic environment for this purpose, recording the update norm, maximum update, and gradient norm of every update.Figure 2 shows the performance of DQN agents trained with Adam-Rel against those trained with Adam-MR and Adam on the Atari-10 benchmark [14] . We tune the learning rate of each method, keeping all other hyperparameters fixed at values tuned for Adam in CleanRL [19] . Adam-Rel outperforms Adam, achieving 65.7% vs. 28.8% human-normalized performance. Furthermore, the stark performance difference between Adam-Rel and Adam-MR (23.5%) demonstrates the advantage of retaining momentum information across target changes (so long as appropriate corrections are applied), thereby contradicting the gradient contamination hypothesis discussed in Bengio et al. [11] and Asadi et al. [10] .More surprisingly, Adam-MR performs substantially worse than Adam, contrasting with the findings of Asadi et al. [10] . We evaluate on a different set of Atari games and tune both Adam and Adam-MR separately, which may account for the differences. However, these results suggest that preventing any gradient information from crossing over target changes is an excessive correction and can even harm performance. We additionally evaluate on the set of games used by Asadi et al. [10] , the results of which can be found in Appendix B. We find that Adam-Rel outperforms the Adam baseline in IQM.We also find that, although our implementation of Adam-MR again significantly under-performs relative to the Adam baseline, we approximately match the returns listed in their work.We also evaluate Adam-Rel when soft target changes are used, by comparing Adam and Adam-Rel on Atari-10 when using DQN with Polyak averaging. We find that Adam-Rel also outperforms Adam in this setting. These results, along with a more detailed discussion, can be found in Appendix C.Craftax Figure 3 shows the performance of PPO agents trained on Craftax-1B over 8 seeds. Most strikingly, Adam-MR, which resets the optimizer completely when PPO samples a new batch, achieves dramatically poorer performance across all metrics. This deficit is unsurprising when compared to its performance on DQN, where the optimizer has many more updates between resets and so can achieve a superior momentum estimate, and demonstrates the impact of not retaining any momentum information after resets in on-policy RL. Similarly, Adam with β 1 = β 2 [16] achieves poorer performance than Adam-Rel on all metrics and has no significant different against Adam with default hyperparameters.Furthermore, Adam-Rel outperforms Adam on all metrics. Whilst the performance on the number of achievements is similar, we follow the evaluation procedure recommended in Hafner [21] and report score, calculated as the geometric mean of success rates for all achievements. This metric applies logarithmic scaling to the success rate of each achievement, thereby giving additional weight to those that are hardest to accomplish. We see that Adam-Rel clearly outperforms Adam in score, as well as on the two hardest achievements (collecting diamonds and eating a plant). These behaviours require [16] . Bars show the 95% stratified bootstrap confidence interval, with mean marked, over 8 seeds [22] .a strong policy to discover so are learned late in training, suggesting that Adam-Rel improves the plasticity of PPO. Atari-57 Figure 2 shows the performance of PPO agents on Atari-57. As before, entirely resetting the optimizer significantly harms performance when compared to resetting only the count. Across all environments, Adam-Rel also improves over Adam, outperforming it in 33 out of the 55 games tested and IQM across games. Adam with β 1 = β 2 also fails to improve over the baseline.To further analyse the impact of Adam-Rel over Atari-57, we plot the performance profile of human-normalized score (Figure 4 ). Whilst the performance of the two methods is similar over the bottom half of the profile, we see a major increase in performance in the top half. Namely, at the 75th percentile of scores Adam-Rel achieves a human-normalized performance of 338% vs. 220% achieved by Adam. This demonstrates the ability of Adam-Rel to improve policy performance on tasks where Adam is successful but suboptimal, without sacrificing performance on harder tasks.In this section we connect our theoretical exposition in Section 3 to our experimental results. Specifically, we first examine whether gradients increase in magnitude due to nonstationarity, to what extent predictions from our model match the resulting updates, and how Adam's update differs from Adam-Rel's in practice.To this end, we collect gradient (i.e., before passing through the optimizer) and update (i.e., the final change applied to the network) information from PPO on Craftax-Classic. We follow the experimental setup in Section 5 but truncate the Craftax-Classic runs to 250M steps to reduce the data processing required. The results are shown in Figure 5 .Comparing Theory and Practice In Figure 5 , both Adam and Adam-Rel face a significant increase in gradient norm immediately after starting optimisation on a new objective resulting from a new batch of trajectories collected under an updated policy and value function. While this matches the assumptions we make in our work, the magnitude of the increase is much less than some of the values explored in Section 3.For Adam, this is approximately 29% and for Adam-Rel it is around 45%. The grad norm profiles look similar in each case, with the norm peaking early before decreasing below its initial average value. This decrease and the initial ramp both deviate from the step function we assume in our model. It is obvious that our theoretical model of gradients, which requires an increase in the gradient magnitude on each abrupt change in the objective, cannot hold throughout training in its entirety because this would require the gradient norm to increase without bound. Figure 5 : Adam and Adam-Rel compared to the theoretical model. To make this plot, we divided all the updates in the PPO run into chunks, each of which was optimising a stationary objective. We then averaged over all the chunks. The red dashed lines show the different epochs for each batch of data. The assumption about the gradient under the model is shown in the grad norm plot. Note that the update norm plot for Adam and Adam-Rel has separate y-axes. The shading represents standard error.However, we find that despite this discrepancy, for Adam-Rel the update predicted by our model fairly closely matches the shape of the true update norm, i.e., a fast drop at the beginning followed by flattening (the scaling is not comparable between observed and predicted values).For Adam, our model explains the initial overshoot of the update norm but then fails to predict the rapid decrease, which results from the fast drop in the true gradient norm. Given the simplicity of our modeling assumptions, we find these results overall encouraging.On Spherical Cows Under the assumption of a step increase in gradients of an infinite relative magnitude Adam-Rel results in a flat update, while Adam would drastically overshoot. Clearly, this assumption does not hold in practice, as we have shown above. However, we believe that this mismatch between reality and assumption is encouraging, since our experimental results show that Adam-Rel is still effective in this regime. Our hypothesis is that there are two benefits to designing Adam-Rel under these assumptions. First of all, it avoids overshoots even under large gradient steps and secondly, when there are less drastic gradient steps it undershoots, which might have similar effects to a fast learning rate annealing. These kind of annealing schedules (over longer horizons) are popular when optimising stationary losses [23, 24] . 
related work: 
 Optimization in Reinforcement Learning Plasticity loss [25] [26] [27] refers to the loss in ability of models to fit new objectives as they are trained. This is particularly relevant in nonstationary settings such as RL and continual learning, where the model is continuously fitting changing objectives. Many solutions have been proposed, including resetting network layers [28] [29] [30] [31] [32] , policy distillation [25] , LayerNorm [33, 34] , regressing outputs to their initial values [26] , resetting dead units [35] and adding output heads during training [36] . These solutions, in particular resetting layers during training [28, 32] , have contributed towards state-of-the-art performance on Atari 100k [30] . However, of these works, only Lyle et al. [33] investigate the relationship between the optimizer and nonstationarity, demonstrating that by reducing the momentum coefficient of the second-moment gradient estimate in Adam, the fraction of dead units no longer increases. However, these works focus on plasticity loss, which is a symptom of nonstationarity, and only analyse off-policy RL. In contrast, we address nonstationarity directly and evaluate both on-policy and off-policy RL.Meta-reinforcement learning [37] [38] [39] provides an alternative approach to designing optimizers for reinforcement learning. Rather than manually identifying problems and handcrafting solutions for RL optimization, this line of work seeks to automatically discover these solutions by meta-learning components of the optimization process. Often these methods parameterize the agent's loss function with a neural network, allowing it to be optimized through meta-gradients [40] [41] [42] or zeroth-order methods [43, 20, 44] . Recently, Lan et al. [45] proposed meta-learning a black-box optimizer directly, demonstrating competitive performance with Adam on a range of RL tasks. However, these works are limited by the distribution of tasks they were trained on, and using handcrafted optimizers in RL is still far more popular.Adam Extensions Cyclical update schedules [46] have previously been applied in supervised learning as a mechanism for simplifying hyperparameter tuning and improving performance, and Loshchilov and Hutter [47] propose the use of warm learning rate restarts with cosine decay for improving the training of convolutional nets. Liu et al. [48] examine the combination of Adam and learning rate warmup, proposing RAdam to stabilise training. However, all of these methods focus on supervised learning and therefore assume stationarity.There has also been some investigation of the interaction between deep RL and momentum-based optimization. Henderson et al. [49] investigate the effects of different optimizer settings and recommend sensible parameters, but do not investigate resetting the optimizer. Bengio et al. [11] identify the problem of contamination of momentum estimates and propose a solution based on a Taylor expansion. Dohare et al. [16] investigate policy collapse in RL when training for longer than methods were tuned for and propose setting β 1 = β 2 to address this. By contrast, we investigate training for a standard number of steps and focus on improved overall empirical performance, rather than avoiding policy collapse. Asadi et al. [10] , which is perhaps the most similar to our work, also aim to tackle contamination, but do so differently, by simply resetting the Adam momentum states to 0 whenever the target network changes in the value-based methods DQN and Rainbow. However, they do not consider resetting of Adam's timestep parameter, and explain their improved results by suggesting that old, bad, momentum estimates contaminate the gradients when training on a new objective. By contrast, we demonstrate that resetting only the timestep suffices for better performance on a range of tasks and therefore that the contamination hypothesis does not explain the better performance of resetting the optimizer. We also demonstrate that retaining momentum estimates can be essential for performance, particularly in on-policy RL.Adam in RL To adapt Adam for use in RL, prior work has commonly applied a number of modifications compared to its use in supervised learning [8] . The first is to set the parameter ϵ to 10 -5 , which is a higher value than the 10 -8 typically used in supervised learning. Additionally many reinforcement learning algorithms use gradient clipping before passing the gradients to Adam. Typically gradient vectors are clipped by their L 2 norm.A higher value of ϵ reduces the sensitivity of the optimizer to sudden large gradients. If an objective has been effectively optimized and hence the gradients are very small, then a sudden target change may lead to large gradients. v typically updates much more slowly than m and therefore this causes the update size to increase significantly, potentially causing performance collapse. However, this implementation detail is not mentioned in the PPO paper [4] , and subsequent investigations omit it [6, 5] . Clipping the gradient by the norm also aims at preventing performance collapse. Andrychowicz et al. [6] find this to increase performance slightly when set to 0.5.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: We provide details of our hyperparameter settings in Appendix F, as well as detailing our experimental setup in Section 5.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
limitations and future work: 
 In this work we have mostly examined abrupt nonstationarity, where there are distinct changes of target, as it is in that setting where our method can be most cleanly applied. However, a range of RL methods face continuous nonstationarity, such as when applying Polyak averaging [1] to smoothly update target networks after every optimization step. We have demonstrated improved performance in this algorithm in Appendix C. However, further investigation into how to apply resetting in this setting would be beneficial.There are also many promising avenues for future work. First, while we have focused on RL, it would be interesting to apply Adam-Rel to other domains that feature nonstationarity such as RLHF, training on synthetic data, or continual learning. Additionally, we also note that while our results are promising, it was not possible to investigate all RL settings and environments in this work, and we therefore encourage future work in settings such as continuous control. Secondly, Adam-Rel is designed with the principle that large updates can harm learning, but it is not clear in general what properties of update sizes are desirable in nonstationary settings. Understanding this more clearly may help produce meaningful improvements in optimisation. Relatedly, it would be beneficial to better understand the nature of gradients in RL tasks, in particular how they change throughout training for different methods and what effect this has on performance. Finally, re-examining other aspects of the RL toolchain that are borrowed from supervised learning could produce further advancements by designing architectures, optimisers and methods specifically suited for problems in RL.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes] Justification: In reporting our results, we follow the recommendations of Agarwal et al. [22] . We provide details of the error bars in the figure captions for each plot.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. Justification: This is foundational machine learning research and as such has no direct path to negative societal consequences separate from advancement in machine learning. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
conclusion: 
 We presented a simple, theoretically-motivated method for handling nonstationarity via the Adam optimizer. By analysing the impact of large changes in gradient size, we demonstrated how directly applying Adam to nonstationary problems can lead to unstable update sizes, before demonstrating how timestep resetting corrects for this instability. Following this, we performed an extensive evaluation of Adam-Rel against Adam and Adam-MR in both on-policy and off-policy settings, demonstrating significant empirical gains. We then demonstrated that increases in gradient magnitude after abrupt objective changes occur in practice and compared the predictions of our simple theoretical model with the observed data in a complex environment. Adam-Rel can be implemented as a simple, single-line extension to any Adam-based algorithm with discrete nonstationarity (e.g. target network updates), leading to major improvements in performance across environments and algorithm classes. We hope that the ease of implementation and effectiveness of Adam-Rel will encourage researchers to use it as a de facto component of future RL algorithms, providing a step towards robust and performant RL. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The paper contains no such risky models or data.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Starting from the definition of the momentum term in Adam's update rule:m i t = (1 -β 1 ) t j=-t ′ β 1 t-j g i j , = (1 -β 1 )   g -1 j=-t ′ β 1 t-j + kg t j=0 β 1 t-j   , = (1 -β 1 )β 1 t g   -1 j=-t ′ β 1 -j + k t j=0 β 1 -j   , = (1 -β 1 )β 1 t g   β 1 t ′ -1 j=0 β 1 j + k t j=0 β -1 1 j   .From the solution to the sum of a geometric series:m i t = (1 -β 1 )β 1 t g β 1 1 -β 1 t ′ 1 -β 1 + k 1 -β 1 -(t+1) 1 -β 1 -1 , = (1 -β 1 )β 1 t g β 1 1 -β 1 t ′ 1 -β 1 + k β 1 -t -β 1 1 -β 1 , = g β 1 t+1 (1 -β 1 t ′ ) + k(1 -β 1 t+1 ) .Similarly for v i t , it follows:v t = g 2 β 2 t+1 (1 -β 2 t ′ ) + k 2 (1 -β 2 t+1) .Substituting v i t and m i t into the Adam momentum updates with ϵ = 0 yields:mi -t ′ ,t vi -t ′ ,t = 1 -β 2 t ′ +t+1 1 -β 1 t ′ +t+1 • g β 1 t+1 (1 -β 1 t ′ ) + k(1 -β 1 t+1 ) g 2 β 2 t+1 (1 -β 2 t ′ ) + k 2 (1 -β 2 t+1 ).Taking the limit t ′ → ∞ with β 1 , β 2 ∈ [0, 1) yields our desired result:lim t ′ →∞ mi -t ′ ,t vi -t ′ ,t = β 1 t+1 + k(1 -β 1 t+1 ) β 2 t+1 + k 2 (1 -β 2 t+1 ).Asadi et al. [10] find in their paper that their method, when applied to DQN, gives roughly comparable performance to their Adam baseline. However, in our paper we find that Adam-MR performs significantly worse than the Adam baseline, even when compared on the same games as in Figure 6 . There Adam-Rel performs better than Adam on the inter-quartile mean, but worse on the median. However, given this is a selection of just 12 games of very different difficulties, the median is often likely in this case to reduce to a single game for most algorithms.To investigate this disparity, we compare our results for Adam-MR to theirs in Table 1 . We estimated their scores in each game from the appropriate figures in their paper. Overall we see that our Table 1 : Comparison with the results from Asadi et al. [10] . The scores are estimated by taking the performance at 40M frames from the figures in their paper. We compare to both K = 1000, which is our default hyperparameter, and K = 8000, which is their default hyperparameter.Environment Adam-MR (K=1000) [10] Adam-MR(K=8000) [10] Adam-Rel (K=1000) Adam-MR (K=1000) [10] Adam-MR (K=8000) [10] implementation, which uses K = 1000, performs significantly better than their implementation with K = 1000. It is also better in mean but worse in median and inter-quartile mean than their K = 8000 implementation. In short, our results broadly match theirs reported after a similar amount of training, but our Adam baseline performs significantly better than theirs. However, there are a number of differences in our evaluation. First we run for 10M steps (40M frames) whereas they run for 30M steps (120M frames). Secondly, they use the Dopamine [50] settings for Atari, whereas we use the more standard ones used by DQN [18] . We kept these settings throughout our paper to avoid significant hyperparameter tuning by evaluating in as standard settings as possible. We believe these results demonstrate the correctness of our implementation of their work and that our method still performs favourably.We also run experiments on DQN with Polyak averaging to examine the effect of Adam-Rel in cases where there are soft-target updates.We set τ = 0.02. This was chosen so that after 250 steps, the previous target update frequency, the original target parameters would contribute just 0.5% to the new target. We found the best learning rate to be lower, at 5 × 10 -5 . The results are shown in Figure 7 . As shown in that figure, Adam-Rel outperforms Adam in this setting as well, achieving a higher median value, although still retaining a long tail of negative results. We also note that although Adam achieves better performance than the baseline without Polyak averaging, Adam-Rel performs worse than when Polyak averaging is not used. This may be due to resetting t being less effective when soft-target changes are used, or that more extensive tuning may improve its performance.For the Atari experiments (both DQN and PPO), we based our implementation on CleanRL [19] . This code is available here. For the Craftax experiments, we based our implementation on PureJaxRL [20] . This code is available here. For our DQN experiments, we swept over learning rates for Adam-MR, Adam-Rel and Adam. For PPO experiments, we swept over learning rate, max gradient norm and GAE λ values, as we found these to differ from the PPO defaults. Experiments were performed on an internal cluster of NVIDIA V100 GPUs. Experiments were scheduled using slurm, with 10 CPU cores per GPU.The Atari PPO experiments required around 10000 GPU hours to complete, including hyperparameter tuning. The DQN experiments, because of the computational inefficiency of DQN, take much longer to run (approximately 2 days per experiment), and hence used a total of 14000 GPU hours, despite there being many fewer parallel runs. The Craftax-Classic experiments took around 300 GPU hours to complete. Justification: As claimed in the introduction, we provide an analysis of the Adam update rule under nonstationary gradients in Section 3, introduce and analyse Adam-Rel in Section 4, then evaluate Adam, Adam-Rel, and Adam-MR on Atari and Craftax in Section 5.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.We exclude 2 out of the 57 games, Montezuma's Revenge and Venture, after observing that all algorithms achieve a human-normalized score of 0.This is the inter-quartile mean over environments as opposed to the more usual over environments and seeds. This is because Asadi et al.[10] do not provide individual seed data. 


Decision: This paper proposes Adam-Rel, a small modification to Adam, that is motivated by the goal of reducing the impact of non-stationarity on policy optimization due to a sudden increase in update norms.

Reviews for this paper are mixed. The authors had concerns regarding the clarity/quality of initial reviews, but the reviewers clarified their opinions in discussion, and also engaged further in closed discussion about the strengths and weaknesses of this paper. The final discussion focused on two key issues: whether large update norms really are the problem, and whether the experimental evidence sufficient to say that the proposed approach is a good general solution.  I think there are reasonable reasons for doubting this paper on both issues. 

The authors demonstrate that at least on the problems they tested, update norms do increase at expected times with Adam, and Adam-Rel alleviates this. But this doesn't necessarily imply that fixing the update norms will always improve optimization. It is entirely possible that since the update norms do not increase without bounds in practice, Adam with suitable hyperparameters can adapt well to them. Similarly, it is possible to come up with scenarios where momentum reset (Adam-MR, a competing approach) might be more beneficial than Adam-Rel --- a reviewer suggests a scenario where the direction of the gradient changes. Could it be the case that Adam-MR is better on some problems (that were not tested) while Adam-Rel is better on other problems? Finally, famously tricky-to-get-right implementation details of PPO make it difficult to trust results when they are on a single benchmark (Crafter). 

Clearly, we can't be sure without much more experimentation. Moreover, there is an issue in deep RL research that algorithms from many papers can't be trusted due to insufficient rigor in designing experiments, and they turn out to not be useful in carefully controlled larger scale studies. This is the reason reviewers are being extra cautious.

Nevertheless, after careful consideration, I believe this paper should be accepted as it meets the bar for the results to be interesting enough to warrant further investigation by the broader community. My decision is based on the following:
- The authors have made effort to adopt a cautious tone in the paper (e.g. results "suggest" instead of "show"), instead of overly confident claims.
- The limitations of idealized theoretical analysis are made clear and then examined experimentally, again in the spirit of presenting and analyzing evidence.
- Many limitations of proposed approach are clearly discussed, and I expect that the authors would not object to discussing more.

I believe the only thing missing from this paper is an explicit admission that although the paper has spent considerable resources on experiments, RL environments and algorithms are so varied that it is quite possible that the results don't generalize broadly, and further investigation is warranted. I strongly encourage the authors to add a caveat like this to complete the limitations section, in line with their cautious approach throughout the paper.

Reviews: review: paper_summary: paper_summary: This paper studies the effect of non-stationarity on the Adam optimizer. It shows that the standard use of the Adam optimizer can lead to large updates, which can cause sub-optimal performance. To address this issue, Adam-Rel is introduced, which resets Adam's timestep parameter to zero when the target network changes. Finally, experiments show that Adam-Rel provides performance improvements over Adam.
strength_weakness: strength_weakness: The same problem of large updates by Adam in non-stationary problems has been studied before (Lyle et al., 2023; Dohare et al., 2023). They both use the solution of setting $\beta_1 = \beta_2$. The authors seem aware of this proposed solution because the work by Dohare et al. (2023) is discussed in the paper. However, Adam-rel is not compared with Adam with $\beta_1 = \beta_2$. A comparison with Adam with $\beta_1 = \beta_2$ is needed to accept this paper. Even if Adam with $\beta_1 = \beta_2$ performs better, this paper can be accepted, as it studies the problem in much more detail than any prior work.

I'm willing to change my score if this comparison is added to the paper.


Lyle et al. Understanding plasticity in neural networks. ICML, 2023.
Dohare et la. (2023). Overcoming policy collapse in Deep RL. EWRL 2023.
questions: questions: How does Adam with $\beta_1 = \beta_2$ perform in your experiments?
limitations: limitations: The authors adequately discuss the limitations.

-----------------
UPDATE: I increased my score as the authors provide comparison with Dohare et al. (2023).
paper_summary: paper_summary: This paper introduces a simple approach to address the issue of large updates commonly encountered with Adam optimizers in deep learning applications. The authors focus on a specific scenario prevalent in deep reinforcement learning: the updating of target networks. Instead of the conventional approach of resetting both the timestep and the momentum variables in the Adam optimizer, the authors propose only resetting the timestep. This modification leads to enhanced learning performance, as demonstrated through evaluations conducted in both the Atari and Craftax environments, utilizing a range of on-policy and off-policy algorithms.
strength_weakness: strength_weakness: The research question addressed by the authors, focusing solely on large gradient updates due to changes in the target network, seems somewhat narrow. This focus may overlook the broader applicability and efficacy of the proposed solution across various reinforcement learning scenarios. Notably, many algorithms have achieved better stability using softer target updates, like Polyak averaging, which the authors acknowledge [1,2,3]. It would be beneficial to see how the proposed optimizer modification compares in environments where such techniques are traditionally employed.

Moreover, the approach of resetting parameters in an optimizer, though straightforward, could be considered a drastic and potentially inelegant solution. This method might introduce other issues, such as affecting the convergence behavior of the optimizer. While I currently lack an alternative suggestion, exploring methods that adjust the scale of parameter updates dynamically, rather than resetting them, might offer a more refined solution.

[1] Lillicrap, Timothy P., et al. "Continuous control with deep reinforcement learning." arXiv preprint arXiv:1509.02971(2015).

[2] Kaplanis, Christos, Murray Shanahan, and Claudia Clopath. "Continual reinforcement learning with complex synapses." International Conference on Machine Learning. PMLR, 2018.

[3] Schwarzer, Max, et al. "Bigger, better, faster: Human-level atari with human-level efficiency." International Conference on Machine Learning. PMLR, 2023.
questions: questions: 1. Have the authors considered extending the setting towards changes in data distribution rather than only due to a change in target network? 
2. The paper includes a comprehensive study using both off-policy and on-policy algorithms. Could the authors clarify the rationale behind this choice? Specifically, which type of algorithm is more affected by drastic changes in data distribution, and why?
3. Figure 4 and Figure 5 should be re-labelled such that Figure 4 should be Figure 5 and Figure 5 should be Figure 4 since in the current version, Figure 5 appears before Figure 4. 
4. I find Figure 5 hard to interpret. Could the authors provide more detail on what the curves represent? Additionally, enhancing the clarity of the caption might help. 
5. Line 86: The authors mention a change of objective. However, it seems the objective function remains consistent throughout. Could the authors clarify if they are actually referring to changes in data distribution rather than the objective itself?
6. Line 259: The authors mention 'matching the trend.' Do you mean matching the overall trend rather than the specific shape or pattern?
7. Line 261:  Could the authors clarify what you mean by 'overshooting' and 'undershooting'? The current usage is confusing.
8. Line 283: The use of 'However' at the beginning of Line 283 immediately follows another sentence starting with 'However' in Line 280. Could these be rephrased to improve the flow of the text?
limitations: limitations: None. The authors have presented the concerned limitations of their work in the paper.
paper_summary: paper_summary: One of the main challenges of reinforcement learning is its inherent nonstationary nature. Such non-stationarity can cause learning difficulties. The tools currently available for deep reinforcement learning are largely borrowed from deep learning, such as the Adam optimizer, which this paper focuses on. The authors show that Adam, under nonstationarity, can have large updates leading to learning difficulties. Thus, they propose a simple modification to Adam in which they reset the time parameter at every epoch in PPO or every target update in DQN. This slight modification seems to significantly improve performance, agreeing with the observations by several works that proposed similar modifications to Adam [1,2,3] and showed gains in performance.
\
\
[1] Emmanuel Bengio, Joelle Pineau, and Doina Precup. Correcting momentum in temporal difference learning. arXiv preprint arXiv:2106.03955, 2021.
\
[2] Kavosh Asadi, Rasool Fakoor, and Shoham Sabach. Resetting the optimizer in deep rl: An empirical study. arXiv preprint arXiv:2306.17833, 2023.
\
[3] Shibhansh Dohare, Qingfeng Lan, and A Rupam Mahmood. Overcoming policy collapse in deep reinforcement learning. In Sixteenth European Workshop on Reinforcement Learning, 2023.
strength_weakness: strength_weakness: - The main weakness of the paper is its limited comparison to other methods, particularly the methods that are closely related to Adam-Rel (e.g. [1,3]). Additionally, the only competitor is Adam-MR [2] which almost always gives worse performance, contradicting previous works. There is no similar theoretical analysis on Adam-MR explaining the poor performance or some empirical evidence convincing the reader why it fails.
- Writing can be improved, especially when explaining parts of Adam. This can be done using more proper mathematical terms such as “estimators” or “correcting the biasedness of an estimator,” etc.
- The title is written in such a way that suggests your method is applicable to a wide range of RL methods, but it seems that this is not the case, and you only focus on PPO and DQN. I suggest a change that reflects your contributions without overstating them. 
- No pseudocode is given for Adam-Rel with DQN. It’s inferred from the context that a time parameter reset is done before each target function update, but the pseudocode needs to be there to confirm this for the reader.
- Dohare et al. (2023) showed that policy collapse in PPO with MuJoCo is more pronounced than DQN in Atari. It might be because the observations are bounded in Atari but not in MuJoCo, which increases the non-stationarity. The authors didn’t consider any environments other than pixel-based ones with bounded observation ranges. I think it’s necessary also to show some results on MuJoCo environemnts (or similar environments) to confirm that the results are qualitatively similar in those settings.
\
\
\
**Minor issues:**
- “It uses a learned critic trained by a TD loss to estimate the value function, and a clipped actor update of the form.” -> This is not accurate. PPO doesn’t use a learned critic. The critic is learned as well.
- “over which the above update is calculated” -> You wrote an objective in Eq. 1, not an update
- “Adam, the most popular optimizer that  uses momentum” -> It doesn’t use the momentum you defined in the background. You might want to clear this up for the reader.
\
\
[1] Emmanuel Bengio, Joelle Pineau, and Doina Precup. Correcting momentum in temporal difference learning. arXiv preprint arXiv:2106.03955, 2021.
\
[2] Kavosh Asadi, Rasool Fakoor, and Shoham Sabach. Resetting the optimizer in deep rl: An empirical study. arXiv preprint arXiv:2306.17833, 2023.
\
[3] Shibhansh Dohare, Qingfeng Lan, and A Rupam Mahmood. Overcoming policy collapse in deep reinforcement learning. In Sixteenth European Workshop on Reinforcement Learning, 2023.
questions: questions: - One of the main assumptions in the theoretical work is that the gradients become near zero at convergence, and then we get sudden large updates once the gradients start to be large. This might be the case for PPO, but it doesn’t seem to be the case for DQN. Dohare et al. (2023) showed that setting beta1 to equal beta improves PPO performance but has a smaller impact on DQN. Adam-Rel shows improvement in both PPO and DQN, but DQN's motivation is not clear. Can the authors explain why they expect Adam-Rel to improve DQN?
- I’m puzzled by Adam-MR's poor performance. Your results seem to contradict (as you mentioned) two previous works. Would it be possible that you used an incorrect implementation for Adam-MR?
- Do the authors have any intuition as to why Adam-MR performs poorly? In particular, do they have a similar theoretical analysis for Adam-MR showing that its updates can be large or even larger than Adam?
limitations: limitations: The authors adequately discussed the limitations of their work.
paper_summary: paper_summary: Due to problems with momentum and a non-stationary target, the authors propose resetting the value of ‘t’ in Adam, which is used to determine bias correction on the momentum terms, when updating the target. This approach is validated on PPO and DQN in Atari and Craftax.
strength_weakness: strength_weakness: - In terms of new insight, this paper offers very little that hasn’t been said in prior work (Bengio et al., Asadi et al.). Similarly, analysis in the paper is based around the update size. However, simply resetting would also bound the update size, so I feel like there is some lacking analysis or evidence to explain why maintaining momentum estimates is a reasonable approach. 
- Since the idea is not analyzed in depth, I would say there is limited empirical evidence (in terms of evaluations over environments and algorithms), especially compared to prior work.
questions: questions: It would be interesting to see more analysis/insight either empirically or theoretically with regards to what happens to the gradient. For example, do momentum estimates correlate with the error terms over new iterations? While resetting seems like a brute force approach to mitigating problems with momentum and non-stationarity, it’s unclear that maintaining the momentum estimate is necessarily better.
limitations: limitations: Satisfactory.

Comments: title: Thank you for your response
Discussion period is closing
Request for engagement
Thank you for your response
Thank you for the Response
Rebuttal Follow up
comment: I thank the authors for their response and providing additional empirical evidence. I believe my concerns have been sufficiently addressed so I will increase my score.
Thank you for taking the time to respond to our rebuttal. There appear to be two concerns with our work, which we address below.

Firstly, you mentioned a concern with the experiments conducted, saying that you do not believe that they are “thorough enough” and that “2 algorithms and 2 domains does not convince me this is a general-purpose RL technique.” In this respect, we believe that precedent from past submissions provides the fairest guide for what merits acceptance. Previous similar empirical research [1,3,4,5,6] from top conferences has focussed on only one algorithmic setting and evaluated on a single environment, typically Atari. Furthermore, both of the works that you cite [1,2] do not have as thorough empirical evaluations, either not evaluating in multiple domains or in on-policy and off-policy settings. Whilst you are entitled to hold higher standards for empirical rigor, we exceed the standard for empirical evaluation set by past top conferences, and hope you will reconsider your decision on these terms. 

Secondly, you mention that “one might expect other approaches which also bound the update size to provide a similar benefit. Since this appears not to be the case, it suggests that you have not correctly, or fully, identified the problem…” We entirely agree with your point that other methods also bound update size and that update size is not the full extent of the optimization problem. Where we disagree is that optimization in RL is multifaceted and there are a number of issues to consider, some of which alternative methods sacrifice. As a reductive example, an optimizer which never updated the agent would bound update size, but would not provide similar benefit to our method. Instead, treating update size as one of many well-established considerations in optimization [7, 8], we theoretically demonstrate our method’s impact in this aspect, then empirically validate its impact on overall RL performance. As there is no single issue and no theoretical or empirical result can entirely encapsulate the challenges in optimization, we hope you agree that this is not a fair standard for a paper.

Finally we would like to draw your attention to some new results we have obtained, evaluating our setting when Polyak averaging is used. 

We ran additional experiments in the Atari-10 setting, evaluating DQN with Adam and Polyak averaging against DQN with Adam-Rel and Polyak averaging. We cannot upload the results figure here, but we show the results below, where the values in brackets are the lower and upper 95% bootstrapped confidence intervals respectively.

| Method                         | Regressed Median    |
|--------------------------------|---------------------|
| Adam (w/ Polyak Averaging)     | 0.096 (0.089, 0.13) |
| Adam-Rel (w/ Polyak Averaging) | 0.42 (0.26, 0.50)   |

This also demonstrates that our method extends to the case of Polyak averaging.

These results got a slightly lower average score -- Polyak averaging is not typically used in Atari with DQN. We chose the coefficient, $\tau$, to result in a similar length of optimisation to the original hyperparameters. We therefore set $\tau = 0.98$ so that after 250 updates, the original parameters contribute $0.6\%$ to the target parameters. Although the scores attained are lower, this change seems to affect Adam and Adam-Rel similarly.

We hope that our response, in addition to these extra experimental results, encourages you to raise your score to vote for acceptance. 

[1] Asadi, Kavosh, Rasool Fakoor, and Shoham Sabach. "Resetting the optimizer in deep rl: An empirical study." Advances in Neural Information Processing Systems 36 (2024).

[2] Bengio, Emmanuel, Joelle Pineau, and Doina Precup. "Correcting momentum in temporal difference learning." arXiv preprint arXiv:2106.03955 (2021).

[3] Lyle, Clare, et al. "Understanding plasticity in neural networks." International Conference on Machine Learning. PMLR, 2023.

[4] Schwarzer, Max, et al. "Bigger, better, faster: Human-level atari with human-level efficiency." International Conference on Machine Learning. PMLR, 2023.

[5] Ceron, Johan Samir Obando, Marc G. Bellemare, and Pablo Samuel Castro. "Small batch deep reinforcement learning." Thirty-seventh Conference on Neural Information Processing Systems. 2023.

[6] Lyle, Clare, Mark Rowland, and Will Dabney. "Understanding and Preventing Capacity Loss in Reinforcement Learning." International Conference on Learning Representations.

[7] Goyal, Priya, et al. "Accurate, large minibatch sgd: Training imagenet in 1 hour." arXiv preprint arXiv:1706.02677 (2017).

[8] Gotmare, Akhilesh, et al. "A closer look at deep learning heuristics: Learning rate restarts, warmup and distillation." arXiv preprint arXiv:1810.13243 (2018).
Thank you for your engagement with our previous response. The discussion period is closing soon and we would like to remind you to respond to our previous comment, as well as inform you about some new results we have obtained.

To check your hypothesis about polyak averaging, we ran additional experiments in the Atari-10 setting, evaluating DQN with Adam and polyak averaging against DQN with Adam-Rel and polyak averaging. We cannot upload the results figure here, but we show the results below, where the values in brackets are the lower and upper 95% bootstrapped confidence intervals respectively.

| Method                         | Regressed Median    |
|--------------------------------|---------------------|
| Adam (w/ Polyak Averaging)     | 0.096 (0.089, 0.13) |
| Adam-Rel (w/ Polyak Averaging) | 0.42 (0.26, 0.50)   |

This also demonstrates that our method extends to the case of Polyak averaging.

These results got a slightly lower average score -- polyak averaging is not typically used in Atari with DQN. We chose the coefficient, $\tau$, to result in a similar length of optimisation to the original hyperparameters. We therefore set $\tau = 0.98$ so that after 250 updates, the original parameters contribute $0.6\%$ to the target parameters. Although the scores attained are lower, this change seems to affect Adam and Adam-Rel similarly.

We hope that these additional results, along with our rebuttal and other positive qualities of our paper, convince you to raise your score.
Thank you for the response. 

Prior work has shown possible limitations with Adam in the RL setting. Your work builds on this insight. Further insight might include new settings (e.g., EMA), or deeper analysis of the problem. 

If one of your key arguments for the performance boost is that the proposed approach bounds the update size, then one might expect other approaches which also bound the update size to provide a similar benefit. Since this appears not to be the case, it suggests that you have not correctly, or fully, identified the problem that your method corrects. 

The analysis is not in-depth because only one element of the algorithm is considered is detail (update size). However, analysis on update size is only sufficiently valuable if we are certain that update size is the problem. I don't see why this is true. 

Regarding empirical evidence, while I appreciate the current set of experiments, I do not believe that they are thorough enough given the limited insight. If the other reviewers feel otherwise, then I simply disagree with the other reviewers. 2 algorithms and 2 domains does not convince me this is a general-purpose RL technique.  

This paper either needs more convincing experiments, empirical insights, or stronger theoretical results to pass the bar for acceptance (one of three, not necessarily all three). In its current state, I view this as a minor heuristic change to prior work with insufficient justification. As such, I do not believe the paper is above the bar for acceptance.
Thank you for your response and for providing a comparison with Adam $\beta_1 = \beta_2$. The new results address my concerns, and I've raised my score to reflect that.
Thank you to reviewers **37TD** and **uFqs** for their engagement with our rebuttal!

Now the discussion period is almost over, we ask that reviewers **JkaJ** and **SmBt** please also read and respond to our rebuttal. We have provided an extensive response to each of your reviews and are eager for you to update your scores or provide further clarification.
I appreciate the extensive response to my concerns. I thank the authors for adding the comparison with Dohare et al. (2023). It is clear that Adam-Rel performs better than setting $\beta_1=\beta_2$. Given the new result with Dohare et al. (2023) and the authors' willingness to clarify the missing details, including adding the DQN-Adam-Rel pseudocode and fixing other writing issues, I have raised my score.
Thank you very much for the clarification of your review and extensive response to our comments. We address your concerns below. 

## The setting

Your objection to our work, as we understand it, is that you are concerned that polyak averaging, which we left out discussion of, would address precisely the problem that we are trying to solve.

However, this is not the point. There may be many solutions to the phenomenon that we describe. Our goal here is not to present **the only solution** but instead to **investigate current common practice and propose a simple and practical solution**. The methods we investigate are popular and widely used and therefore relevant for our paper. Our solution of resetting the timestep has a number of advantages:
1. It can be applied to both on-policy and off-policy methods
2. It is very simple to implement 
3. It has strong empirical performance in a range of environments

Polyak averaging, although it satisfies the latter two points, cannot be applied to PPO, or related algorithms, as there is no target in a similar way. Therefore, even if 
Polyak averaging did indeed address this problem, it's still not clear it is a better solution than what we have proposed.

Thank you for clarifying your remarks about our method being 
"drastic" and "potentially inelegant". On our method being "drastic", you state:

> Resetting is drastic because it generally removes previously learned information, such as momentum, and affects adaptive learning rates, which can result in longer convergence times.

We are not sure what you mean by this. Our method **does not remove momentum information**. Adam-Rel retains momentum information. It is clear that this is essential to the performance of our method given the poor performance of Adam-MR on the same tasks. Throwing away momentum information seems to significantly harm performance. Secondly, although our method clearly does affect adaptive update sizes, we **used exactly the same learning rate schedules as the base implementation**. Therefore, although this could potentially cause problems with our method, empirically, using identical learning rate schedules as with Adam achieved similarly good results. 

You also state

> Resetting disrupts the continuous learning process, causing the model to lose valuable accumulated optimization information.

**This is not true.** No information is 'lost' by resetting. We merely reset the $t$ parameter to $0$, which is only used to scale the momentum terms. The momentum estimates are exactly the same after as they were before. Instead, resetting $t$ rescales the update, but still retains all information that was there previously.

## Data Distribution

The only way to change the data distribution is to update the policy you roll out to collect data. In this sense, 
a change in target, either by updating $\pi$ in PPO or by updating the target network, do not have any effect on this.

Therefore the only consideration here is how much individual policy updates change the encountered data distribution. In this sense, the *mechanism* that causes this problem is the same. However, this is not a well-studied phenomenon and thus a definitive answer is not possible. If I were to guess, I would think that PPO is less-affected by this problem than DQN because of its clipped objective. However, this is entirely conjecture.

## Parameters

By $\theta$ we just mean the parameters of any neural networks that are optimised. This includes the critic, if it exists, and the policy. Thank you for pointing out this source of confusion -- we will amend this wording in the camera ready copy.


Thank you again for engaging with our rebuttal! We hope that the above clarifications encourage you to raise your review score.
Thank you for your rebuttal. Upon reading your comments, I first want to apologize for any confusion my review may have caused. It was never my intention to undermine the authors’ work. More importantly, I recognize that I should provide more detailed insights and reasoning to support the claims I made. Let me unpack my claims further, and I would be glad to clarify if needed.

# 1. Narrow Setting
The scenario referred to as narrow pertains to the situation the authors address: a (hard) change in the target network. The authors claimed (lines 10-11) and demonstrated that changes in the target network impact learning performance, as shown in Figure 2.

The reason for my claim that the study is narrow is as follows: Do learning difficulties persist if we use a soft target network update, such as Polyak averaging, instead of a hard update? 

Do we still observe large updates? If Polyak averaging can overcome learning inefficiencies and the issue of large updates, what additional benefits does the proposed method, Adam-Rel, provide, given that Polyak averaging offers a simpler solution without resetting?

# 2. Why is Resetting Drastic?
Resetting is drastic because it generally removes previously learned information, such as momentum, and affects adaptive learning rates, which can result in longer convergence times.

# 3. Why is Resetting potentially inelegant?
Resetting disrupts the continuous learning process, causing the model to lose valuable accumulated optimization information.

# 4. What would be a more refined solution? 
A potentially more refined solution would be to dynamically adapt Adam's timestep, such as using the difference in the gradient update norm. This approach would be less drastic and more elegant.

# Q2: Impact of Data Distribution Change
I refer to the change in data distribution after a hard update of the network generated by a different policy. How does this affect off-policy algorithms like DQN and on-policy algorithms like PPO? Are they both impacted in the same way? If not, which class of algorithms experiences more drastic effects and why?

# Q4: Suggestions for Figure 5
Label each subfigure in Figure 5, resulting in 5a, 5b, 5c, and 5d for easier reference.
Why is the trend in the update norm of Adam so different compared to the theoretical model?
The gradient norm per parameter update for both Adam and Adam-Rel exhibits the same trend, which is vastly different from the theoretical model. Why is this so?

To the best of my knowledge, these effects were not discussed in the main paper.

# Q5: Clarification on Parameters
From lines 87-90, you state: “More explicitly, we consider optimizing a stationary loss function $L(\theta, \phi)$, where $\theta$ are the parameters to be optimized and $\phi$ is the other parameters of the loss function (such as the parameters of a value network), which are not updated throughout optimization, but does not include the training data.”
What do the parameters $\theta$ consist of? Are they the parameters of the Adam optimizer or the parameters of the ANN pr both? Why is the value network not being updated during optimization?

The phrase “which are not updated throughout optimization, but does not include training data” is confusing. Does this imply that the training data is updated throughout optimization?

# Conclusion

My main reason for rejecting this submission is due to my concerns about the limited and narrow setting the authors have studied. I have reviewed the feedback from other reviewers, and I am still not convinced that this work warrants acceptance. However, I will re-evaluate based on the progress of the discussion between the authors and the other reviewers, as well as our ongoing discussion.