PaperlayerView exact revision on arXiv

Source-linked AI summary

Understanding and Preventing Capacity Loss in Reinforcement Learning

Clare Lyle, Mark Rowland, Will Dabney

arXiv:2204.09560v2cs.LG

TL;DR

Deep RL agents must continually adapt to changing prediction targets, but sequential target fitting can cause capacity loss and impede learning progress. The paper analyzes this phenomenon and introduces Initial Feature Regularization to preserve feature capacity. Capacity loss occurs across agents and environments, is especially damaging in sparse-reward settings, and InFeR improves performance on hard-exploration tasks such as Montezuma’s Revenge.

  • Problem

    Deep RL agents must continually adapt to changing prediction targets, but sequential target fitting can reduce their ability to learn new functions quickly.

  • Method

    The paper empirically analyzes target-fitting capacity and feature rank, then introduces InFeR to preserve feature subspaces present at initialization.

  • Results

    Capacity loss appears across several RL agents and environments, is especially damaging in sparse-reward settings, and InFeR improves performance on hard-exploration tasks such as Montezuma’s Revenge.

  • Takeaways & Limitations

    Preventing capacity loss can help agents benefit from learning signals throughout training, including under naive exploration in sparse-reward environments.

  • Takeaways & Limitations

    High feature rank is not sufficient for learning progress because other RL components must also be suitable.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Deep RL agents face non-stationary prediction problems that can reduce their ability to adapt, especially in sparse-reward environments. The paper studies this capacity loss and proposes representation-focused mitigation.

  • Deep RL is brittle in sparse-reward environments, where different random seeds can produce dramatically different outcomes.
  • Agents solve a sequence of evolving prediction tasks as their value functions and policies improve, but sequential fitting can harm adaptation to future targets.
  • Capacity loss occurs when agents lose the ability to quickly fit new prediction tasks, and in extreme cases it prevents learning progress.
  • The paper finds declining target-fitting ability across Atari environments and non-stationary reward tasks, alongside diminishing feature separability in sparse-reward settings.
  • Initial Feature Regularization preserves feature subspaces by regressing auxiliary outputs toward their initialization values.
  • Suitable representation learning can improve hard-exploration performance without smart exploration algorithms, indicating that poor sparse-reward performance also involves representation learning.

2 BACKGROUND

The paper frames value-based deep RL as learning action values in a Markov Decision Process from sampled transitions. Neural networks are updated with stochastic gradient descent using replay data and slower target parameters, while the learned features are analyzed as representations.

  • A Markov Decision Process specifies states, actions, rewards, transitions, and a discount factor for the agent’s interaction problem.
  • Q-learning updates a predicted action-value function toward a bootstrap target.
  • Practical updates use sampled transitions containing the current state, action, reward, and next state.
  • Deep RL uses neural-network function approximation, minibatches from replay buffers, and gradient descent on sampled transitions.
  • A separate, more slowly updated parameter set computes target values for the learning loss.
  • The study examines how variations in the learning objective shape learning dynamics, particularly the dynamics of learned features.
  • The network’s features are defined as the outputs of the penultimate layer used for additional representation-learning objectives.

3 CAPACITY LOSS

Sequentially changing prediction targets can erode networks’ ability to fit new targets and distinguish states, producing capacity loss that is especially damaging in sparse-reward reinforcement learning. The paper measures this through target fitting and feature rank, finding links between reduced representation capacity, representation collapse, and impaired learning progress.

  • Target-fitting capacity: Sequentially fitting changing targets can harm a network’s ability to adapt to new targets as its value function and policy evolve.The paper frames this as a consequence of reinforcement learning’s long sequence of changing prediction problems.
  • Target-fitting capacity: Target-fitting capacity measures how well a network reaches new targets from its current parameters within a limited optimization budget.The definition depends on the optimizer, target distribution, and chosen optimization budget.
  • Target-fitting capacity: Networks trained on successive MNIST targets show increasing error on later targets, while sufficiently over-parameterized models can exhibit positive forward transfer.The decline is strongest in smaller networks; models with roughly one million parameters for one thousand data points instead show positive transfer.
  • Target-fitting capacity: Agent checkpoints from value-based RL training are tested by regressing on randomly generated targets to assess whether target-fitting ability declines during training.The procedure loads checkpoint parameters, samples replay-buffer inputs, and fits new random targets under a fixed evaluation process.
  • Representation capacity: Feature rank approximates rapid adaptation capacity by measuring how easily states can be distinguished when only the final layer is updated.The estimator uses the dimension of the feature subspace above a singular-value cutoff and is designed to capture representation collapse.
  • Representation capacity: Higher feature rank corresponds to higher Montezuma’s Revenge performance, whereas sparse-reward environments can show declining rank and representation collapse; high rank alone is not sufficient for progress.Other requirements include useful learning signals, accurate credit assignment, stable updates, suitable optimization, and effective exploration.

4 INFER: MITIGATING CAPACITY LOSS WITH FEATURE REGULARIZATION

InFeR preserves initial feature subspaces to mitigate capacity loss, improving performance especially in sparse-reward environments. Its effects depend on network capacity and regularize the broader learning dynamics.

  • 4 INFER: MITIGATING CAPACITY LOSS WITH FEATURE REGULARIZATION: InFeR adds auxiliary linear prediction heads whose current outputs regress toward their initialization outputs, preserving feature subspaces.The method uses k heads over the feature representation and optionally scales initialization targets by β.
  • 4 INFER: MITIGATING CAPACITY LOSS WITH FEATURE REGULARIZATION: InFeR improved Rainbow performance across Atari, with significant gains concentrated in games where agents struggle to reach human performance.The evaluations used k = 10 heads, β = 100, and α = 0.1.
  • 4.1 INFER: FEATURE-SPACE REGULARIZATION: In sparse-reward Montezuma’s Revenge, adding InFeR to DDQN enabled it to exceed Rainbow’s noisy-networks approach during the last 40 million training frames.The DDQN agent used naive ϵ-greedy exploration and otherwise obtained zero reward throughout training.
  • 4.2 UNDERSTANDING HOW INFER WORKS: InFeR can slow progress in some environments, including Asteroids and Jamesbond, while improving performance on average across Atari games.The authors investigate this differential effect through random-feature and network-width experiments.
  • 4.2 UNDERSTANDING HOW INFER WORKS: Random appended features did not improve performance over Rainbow, indicating that InFeR’s effect on earlier network layers is crucial.The random-feature agents performed comparably to vanilla Rainbow agents rather than to InFeR agents.
  • 4.2 UNDERSTANDING HOW INFER WORKS: Doubling the penultimate-layer width reduced, eliminated, or sometimes reversed InFeR’s performance gap in environments where the smaller architecture underperformed.The result supports regularization of the entire network’s learning dynamics as the principal mechanism.

5 RELATED WORK

The paper relates its representation analysis to prior work on auxiliary tasks, value-function geometry and stability, and interference between subtasks. It also situates InFeR among regularization and distillation approaches to catastrophic forgetting.

  • 5 RELATED WORK: Prior auxiliary-task research has improved performance and encouraged desirable learned-representation properties across diverse settings.Related analyses have also examined the geometry and stability of RL value functions.
  • 5 RELATED WORK: Prior work studies interference between subtasks in reinforcement and supervised learning, including catastrophic forgetting.Proposed responses include regularization and distillation-based training algorithms.

6 CONCLUSIONS

The paper concludes that deep RL agents can lose the capacity to distinguish states and represent new target functions, especially in sparse-reward settings. InFeR preserves capacity and improves performance in settings where agents struggle to match human performance.

  • 6 CONCLUSIONS: Deep RL agents can lose the capacity to distinguish states and represent new target functions during training.In sparse-reward settings, representation collapse can completely prevent learning progress.
  • 6 CONCLUSIONS: Larger networks trained on rich learning signals are more robust to capacity loss than smaller networks trained on sparse targets.The conclusion identifies network size and learning-signal richness as important moderators.
  • 6 CONCLUSIONS: InFeR preserves capacity and improves performance across settings where deep RL agents have historically struggled to match human performance.The authors interpret the regularizer as a form of function-space regularization.
  • 6 CONCLUSIONS: In Montezuma’s Revenge, effective representation learning enabled good policies under sparse rewards and naive exploration strategies.The result suggests that inadequate representation learning contributes alongside exploration difficulty to poor sparse-reward performance.

A THEORETICAL RESULTS

The theoretical results establish consistency of the feature-rank estimator. Convergence follows from matrix convergence and continuity of eigenvalues, except at thresholds equal to eigenvalues of the population matrix.

  • A THEORETICAL RESULTS: The expected feature-covariance matrix is approximated elementwise, and the resulting matrix converges to its population counterpart.The argument applies the strong law of large numbers to elements of M = E[φφ⊤].
  • A THEORETICAL RESULTS: Because eigenvalues are continuous functions of the matrix, the singular values of the feature matrix converge to those of the population matrix almost surely.The proof identifies singular values of Φ with eigenvalues of M.
  • A THEORETICAL RESULTS: The threshold estimator N(λ1, . . . , λk; ϵ) converges for almost every threshold ϵ, excluding eigenvalues of the population matrix.The estimator counts eigenvalues exceeding ϵ.

A.2 FEATURE DYNAMICS

The analysis studies feature dynamics under continuous-time TD learning and shows that sparse rewards can drive representations toward collapse, especially with independently initialized prediction heads. It also relates this mechanism to saturation risks in deep ReLU agents and introduces non-stationary MNIST experiments for studying capacity loss.

  • Feature dynamics: The analysis models value functions as linear predictions over learned features and translates discrete expected TD updates into continuous-time feature dynamics.The representation is denoted Φ_t, with value estimates parameterized as V_t = ⟨Φ_t(x), w_t⟩.
  • Feature dynamics: The framework considers ensemble prediction with multiple independently initialized linear heads trained on TD targets over a shared representation.The analysis restates prior limiting results for different learning-rate and initialization scalings.
  • Sparse-reward collapse: Under sparse rewards and sub-optimal policies, the representation converges to the zero vector for every state, causing feature rank to tend toward zero.This result is stated to hold independently of whether the learning rate or weight-initialization variance is scaled with ensemble size.
  • Sparse-reward collapse: For an ergodic transition process, the limiting stable rank is one in the analyzed setting.The result follows from exponential evolution along transition-matrix eigenvectors.
  • Connection to deep agents: The linear analysis does not exactly match deep neural-network training, but it suggests collapse along output-relevant dimensions and saturation of ReLU features in sparse-reward agents.The paper specifically connects these risks to QR-DQN agents trained on environments such as Montezuma’s Revenge.
  • Non-stationary target experiments: The accompanying MNIST setup uses changing target functions, small fully connected ReLU networks, and varied widths and depths to study capacity loss.Training proceeds across randomly sampled target functions, carrying parameters from one iteration into the next.

B.2 ADDITIONAL EVALUATIONS

The additional MNIST evaluations vary target structure and network size to test how architecture and task difficulty affect capacity loss. They find that larger or deeper networks can slow capacity loss, while InFeR significantly mitigates declining target-fitting capacity.

  • Experimental design: The experiments compare random, hash, and threshold MNIST tasks while varying network width and depth.They evaluate target-fitting loss, feature rank, and stable feature rank across sequential target functions.
  • Network size effects: Target-fitting error increases over iterations in smaller hash-MNIST networks, whereas greater depth or width slows capacity loss and enables positive transfer in the largest networks studied.The figure reports mean squared error at the end of training for each task iteration.
  • Target structures: Hash-MNIST removes the neural-network-induced target structure by applying a sine transformation, making the task equivalent to memorizing labels.Random-MNIST retains structure because similar inputs tend to produce similar outputs under the target-generating network.
  • InFeR evaluation: InFeR significantly mitigates the decline in target-fitting capacity on the non-stationary MNIST reward-prediction task.The comparison is against the corresponding agent without the regularizer.

C ATARI EVALUATIONS

The Atari evaluations extend the analysis across training checkpoints and MNIST-style prediction tasks, examining target-fitting error and the effect of InFeR. The reported figures compare performance under changing network or regularization conditions.

  • Evaluation procedure: The evaluation procedure uses a common training process and loads agent checkpoints to compute the reported quantities.The same checkpoint-based procedure is used throughout the evaluation section.
  • Random-MNIST comparison: On random-MNIST, mean squared error after 2e3 training steps is compared across network configurations, with the figure caption emphasizing capacity loss in under-parameterized networks.The caption states that increasing depth or width slows capacity loss and enables positive transfer in the largest network studied.
  • Threshold-MNIST comparison: On threshold-MNIST, mean squared error after 2e3 training steps is likewise used to examine target-fitting behavior across network configurations.The supplied figure passage identifies the task and evaluation quantity but does not state a specific numerical outcome.
  • InFeR comparison: On the non-stationary random-reward MNIST task, adding InFeR produces networks that consistently outperform standard regression networks and exhibit minimal capacity loss by comparison.The comparison uses the same network architecture and the same sequence of targets.
  • Atari evaluation: The Atari evaluation includes hyperparameter sweeps for DDQN+InFeR using average capped human-normalized score at the end of training.Each contour plot marginalizes over hyperparameters not shown on its axes.

C.1 HYPERPARAMETER SENSITIVITY OF INFER IN DEEP REINFORCEMENT LEARNING AGENTS

The hyperparameter-sensitivity study evaluates InFeR with DDQN and Rainbow agents across auxiliary-prediction and loss-scale settings. It summarizes robustness using capped human-normalized returns and examines feature rank alongside agent performance.

  • Hyperparameter sensitivity: The sweeps vary the number of auxiliary predictions, cumulant scale, and auxiliary-loss scale for DDQN+InFeR and Rainbow+InFeR.The tested auxiliary-head counts are 1, 5, 10, and 20; cumulant scales are 10, 100, and 200; loss scales are 0.01, 0.05, 0.1, and 0.2.
  • Evaluation metric: The sweep metric is capped human-normalized return averaged across Montezuma’s Revenge, Hero, James Bond, and Ms. Pac-Man.The contour plots marginalize over hyperparameters not displayed on their axes.
  • Rainbow configuration: Rainbow+InFeR uses the Rainbow architecture and adds InFeR with 10 heads, gradient weight 0.1, and scale 100.Training follows the referenced Rainbow implementation and runs for 200 million frames with periodic evaluation.
  • Representation-performance analysis: Feature rank and agent performance are evaluated together on demonstrative Atari environments.The supplied figure passage identifies the paired quantities but does not report a specific numerical comparison.
  • Checkpoint analysis: The evaluation saves agent parameters and replay buffers every 10M frames to estimate feature dimension and target-fitting capacity.Evaluation frames are interspersed every 1M training frames during the 200-million-frame run.

C.2 FEATURE RANK

Feature rank generally declines during training, indicating that agent representations lose dimension despite increasingly diverse observations. InFeR often mitigates this decline, though neither the collapse nor the regularizer’s benefit is uniform across games.

  • Evaluation scope: The evaluation includes feature-rank analyses across additional games and full Rainbow, double-width Rainbow, raw-score, and Double DQN performance curves.These evaluations compare Rainbow with Rainbow+InFeR and assess performance using multiple score presentations and architectures.
  • InFeR across Atari: Across all 57 Atari games, InFeR often mitigates representation-dimension decline, but its effect is not uniformly beneficial.Some games show neither consistent learning progress nor the representation collapse targeted by InFeR.
  • Measurement procedure: Feature rank is estimated from the singular values of a representation matrix built from replay-buffer states, counting values above 0.01.The matrix contains penultimate-layer features for sampled origin states.
  • Feature-rank evolution: In most games, feature rank declines after the first checkpoint at 10M frames despite the agent observing increasingly diverse states.This indicates strong pressure toward representational degeneracy in online RL.
Loading 2204.09560v2…