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

TL;DR

This paper introduces Adam-Rel, a modified optimizer for reinforcement learning that resets the timestep after target network updates, effectively addressing nonstationarity and improving performance across various RL benchmarks.

Contribution

The paper proposes Adam-Rel, an adaptation of Adam that resets its timestep within epochs to better handle nonstationarity in RL, outperforming standard Adam.

Findings

Adam-Rel reduces large updates caused by nonstationary gradients.

Improved performance of Adam-Rel on Atari and Craftax benchmarks.

Gradient norm increases are observed in RL, validating the theoretical model.

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.