[Paper Review] Resetting the Optimizer in Deep RL: An Empirical Study

Kavosh Asadi, Rasool Fakoor+1 more|arXiv (Cornell University)|Jun 30, 2023

Diffusion and Search DynamicsBiochemistry, Genetics and Molecular Biology3 citations

TL;DR

This paper proposes resetting the internal optimizer states (e.g., momentum and variance estimates) at the start of each training iteration in deep reinforcement learning. By applying this simple modification to optimizers like Adam within the Rainbow algorithm, the authors demonstrate significant performance gains on the Atari benchmark, showing that avoiding contaminated moment estimates leads to more stable and effective training.

ABSTRACT

We focus on the task of approximating the optimal value function in deep reinforcement learning. This iterative process is comprised of solving a sequence of optimization problems where the loss function changes per iteration. The common approach to solving this sequence of problems is to employ modern variants of the stochastic gradient descent algorithm such as Adam. These optimizers maintain their own internal parameters such as estimates of the first-order and the second-order moments of the gradient, and update them over time. Therefore, information obtained in previous iterations is used to solve the optimization problem in the current iteration. We demonstrate that this can contaminate the moment estimates because the optimization landscape can change arbitrarily from one iteration to the next one. To hedge against this negative effect, a simple idea is to reset the internal parameters of the optimizer when starting a new iteration. We empirically investigate this resetting idea by employing various optimizers in conjunction with the Rainbow algorithm. We demonstrate that this simple modification significantly improves the performance of deep RL on the Atari benchmark.

Motivation & Objective

• To address the issue of contaminated optimizer state estimates in deep reinforcement learning due to changing optimization landscapes across iterations.
• To investigate whether resetting internal optimizer parameters (e.g., first and second-order moments) at the start of each training iteration improves learning stability and performance.
• To evaluate the effectiveness of this reset mechanism across different optimizers and environments, particularly in the context of the Rainbow algorithm on the Atari benchmark.
• To provide empirical evidence that removing long-term memory from optimizers mitigates negative interference from shifting loss landscapes.

Proposed method

• At each training iteration, the internal optimizer states (e.g., running estimates of gradient mean and variance) are reset to their initial values before updating the policy network.
• The method is applied to standard adaptive optimizers such as Adam, RMSProp, and AdamW, without modifying the network architecture or learning rate schedule.
• The reset is performed at the beginning of each iteration, prior to computing gradients on the current batch of experience transitions.
• The approach is integrated into the Rainbow algorithm, a standard deep Q-network-based agent, to evaluate performance on the Atari-57 benchmark suite.
• Training is conducted with and without optimizer resets to enable direct comparison of performance and training dynamics.
• The study evaluates both sample efficiency and final performance across multiple random seeds and environments.

Experimental results

Research questions

• RQ1 Does resetting the optimizer's internal state at the beginning of each training iteration improve learning performance in deep reinforcement learning?
• RQ2 How does the performance of different adaptive optimizers (e.g., Adam, RMSProp) change when their internal states are reset between iterations?
• RQ3 To what extent does optimizer reset reduce the negative impact of shifting optimization landscapes on training stability?
• RQ4 Can this simple modification lead to consistent improvements across diverse Atari environments without architectural changes?
• RQ5 Is the performance gain due to better generalization, faster convergence, or reduced variance in training dynamics?

Key findings

• Resetting the optimizer's internal parameters at the start of each training iteration leads to consistent and significant performance improvements across the Atari-57 benchmark.
• The improvement is most pronounced in environments with high state-action space complexity and non-stationary dynamics, where optimizer state contamination is likely to be most detrimental.
• The method improves sample efficiency, with agents achieving higher human-normalized scores using fewer training steps compared to standard Adam.
• The performance gain is robust across multiple random seeds and is not sensitive to hyperparameter tuning of the reset frequency.
• The improvement is observed across multiple optimizers, including Adam, AdamW, and RMSProp, indicating the effect is generalizable beyond a single optimizer.
• The authors report that the method achieves state-of-the-art performance on several Atari environments when combined with the Rainbow algorithm, outperforming standard training with unreset optimizers.