Understanding and Preventing Capacity Loss in Reinforcement Learning

TL;DR

Capacity loss emerges as a non-stationarity-induced bottleneck in deep reinforcement learning, where networks gradually lose the ability to quickly adapt to new targets and distinguish states, a problem that is acute in sparse-reward environments. The authors introduce Initial Feature Regularization (InFeR), a function-space regularizer that regresses auxiliary outputs toward their initialization to preserve capacity, and demonstrate robust gains across Rainbow and DDQN agents, most notably in Montezuma's Revenge and various Atari games. By linking representation dynamics to learning progress, the work argues for preserving plasticity as a core design objective alongside exploration. The findings suggest broad implications for non-stationary prediction problems in RL and point to future avenues in representation learning and regularization techniques that maintain adaptability.

Abstract

The reinforcement learning (RL) problem is rife with sources of non-stationarity, making it a notoriously difficult problem domain for the application of neural networks. We identify a mechanism by which non-stationary prediction targets can prevent learning progress in deep RL agents: \textit{capacity loss}, whereby networks trained on a sequence of target values lose their ability to quickly update their predictions over time. We demonstrate that capacity loss occurs in a range of RL agents and environments, and is particularly damaging to performance in sparse-reward tasks. We then present a simple regularizer, Initial Feature Regularization (InFeR), that mitigates this phenomenon by regressing a subspace of features towards its value at initialization, leading to significant performance improvements in sparse-reward environments such as Montezuma's Revenge. We conclude that preventing capacity loss is crucial to enable agents to maximally benefit from the learning signals they obtain throughout the entire training trajectory.

Figures (20)

• Figure 1: Networks trained to fit a sequence of different targets on MNIST data see increasing error on new target functions as the number of tasks trained on increases.
• Figure 2: Networks see reduced ability to fit new targets over the course of training in two demonstrative Atari environments.
• Figure 3: Feature rank and performance over the course of training for Montezuma's Revenge (left) and Pong (right). We observe that feature rank is higher for environments and auxiliary tasks which provide denser reward signals than for sparse reward problems.
• Figure 4: (a): Agent capacity vs human-normalized score in games where Rainbow does not achieve superhuman performance. While $\text{feature rank}\xspace$ does not appear to solely determine agent performance, there is a positive correlation between $\text{feature rank}\xspace$ and human-normalized score. Bottom row contains Rainbow agents trained with the regularizer presented in Equation \ref{['eq:infer']}. (b) An 'unlucky' seed from our evaluations on the sparsified version of Pong, where learning progress occurs only after the agent recovers from representation collapse.
• Figure 5: (a) Visualization of InFeR. (b) Analysis of the effect of InFeR on capacity loss. (c) Effect of InFeR on performance in Montezuma's Revenge with respect to Rainbow and Double DQN baselines. (d) Performance of InFeR relative to Rainbow on all 57 Atari games.

Theorems & Definitions (5)

• Definition 1: Target-fitting capacity
• Definition 2: Feature rank
• Theorem 1: Lyle et al., 2021
• Corollary 1
• proof