Investigating the Edge of Stability Phenomenon in Reinforcement Learning

TL;DR

This paper investigates the edge of stability phenomenon in reinforcement learning, revealing its presence in off-policy deep RL and highlighting how different loss functions influence this dynamic, contrasting with supervised learning behaviors.

Contribution

It extends the understanding of the edge of stability phenomenon from supervised learning to reinforcement learning, showing its occurrence in off-policy algorithms and the impact of loss function choices.

Findings

Edge of stability observed in off-policy deep RL.

DQN with Huber loss shows strong edge of stability effects.

C51 with cross entropy loss does not exhibit the phenomenon.

Abstract

Recent progress has been made in understanding optimisation dynamics in neural networks trained with full-batch gradient descent with momentum with the uncovering of the edge of stability phenomenon in supervised learning. The edge of stability phenomenon occurs as the leading eigenvalue of the Hessian reaches the divergence threshold of the underlying optimisation algorithm for a quadratic loss, after which it starts oscillating around the threshold, and the loss starts to exhibit local instability but decreases over long time frames. In this work, we explore the edge of stability phenomenon in reinforcement learning (RL), specifically off-policy Q-learning algorithms across a variety of data regimes, from offline to online RL. Our experiments reveal that, despite significant differences to supervised learning, such as non-stationarity of the data distribution and the use of bootstrapping, the edge of stability phenomenon can be present in off-policy deep RL. Unlike supervised learning, however, we observe strong differences depending on the underlying loss, with DQN -- using a Huber loss -- showing a strong edge of stability effect that we do not observe with C51 -- using a cross entropy loss. Our results suggest that, while neural network structure can lead to optimisation dynamics that transfer between problem domains, certain aspects of deep RL optimisation can differentiate it from domains such as supervised learning.