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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Ellis, Benjamin, Jackson, Matthew T., Lupu, Andrei, Goldie, Alexander D., Fellows, Mattie, Whiteson, Shimon, Foerster, Jakob
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866915075674079232
author Ellis, Benjamin
Jackson, Matthew T.
Lupu, Andrei
Goldie, Alexander D.
Fellows, Mattie
Whiteson, Shimon
Foerster, Jakob
author_facet Ellis, Benjamin
Jackson, Matthew T.
Lupu, Andrei
Goldie, Alexander D.
Fellows, Mattie
Whiteson, Shimon
Foerster, Jakob
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2412_17113
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Adam on Local Time: Addressing Nonstationarity in RL with Relative Adam Timesteps
Ellis, Benjamin
Jackson, Matthew T.
Lupu, Andrei
Goldie, Alexander D.
Fellows, Mattie
Whiteson, Shimon
Foerster, Jakob
Machine Learning
Optimization and Control
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.
title Adam on Local Time: Addressing Nonstationarity in RL with Relative Adam Timesteps
topic Machine Learning
Optimization and Control
url https://arxiv.org/abs/2412.17113