Normalization and effective learning rates in reinforcement learning

TL;DR

This paper introduces Normalize-and-Project (NaP), a re-parameterization technique that maintains a constant effective learning rate in normalization-based deep reinforcement learning, improving robustness and performance.

Contribution

The paper proposes NaP, a simple re-parameterization that couples normalization with weight projection to stabilize effective learning rates in reinforcement learning models.

Findings

NaP maintains constant effective learning rates during training.

NaP improves robustness to nonstationarity in benchmarks.

NaP can be applied to architectures like ResNets and transformers.

Abstract

Normalization layers have recently experienced a renaissance in the deep reinforcement learning and continual learning literature, with several works highlighting diverse benefits such as improving loss landscape conditioning and combatting overestimation bias. However, normalization brings with it a subtle but important side effect: an equivalence between growth in the norm of the network parameters and decay in the effective learning rate. This becomes problematic in continual learning settings, where the resulting effective learning rate schedule may decay to near zero too quickly relative to the timescale of the learning problem. We propose to make the learning rate schedule explicit with a simple re-parameterization which we call Normalize-and-Project (NaP), which couples the insertion of normalization layers with weight projection, ensuring that the effective learning rate remains constant throughout training. This technique reveals itself as a powerful analytical tool to better understand learning rate schedules in deep reinforcement learning, and as a means of improving robustness to nonstationarity in synthetic plasticity loss benchmarks along with both the single-task and sequential variants of the Arcade Learning Environment. We also show that our approach can be easily applied to popular architectures such as ResNets and transformers while recovering and in some cases even slightly improving the performance of the base model in common stationary benchmarks.