Measure gradients, not activations! Enhancing neuronal activity in deep reinforcement learning

TL;DR

This paper introduces GraMa, a new metric based on gradients to measure neuron learning capacity in deep RL, leading to improved performance when neurons are reset guided by this metric.

Contribution

The paper proposes GraMa, a gradient-based neuron activity metric, and demonstrates its effectiveness in enhancing deep RL agents across various architectures and benchmarks.

Findings

GraMa effectively detects neuron inactivity in complex architectures.

Resetting neurons guided by GraMa improves RL agent performance.

GraMa outperforms activation-based metrics in diverse models.

Abstract

Deep reinforcement learning (RL) agents frequently suffer from neuronal activity loss, which impairs their ability to adapt to new data and learn continually. A common method to quantify and address this issue is the tau-dormant neuron ratio, which uses activation statistics to measure the expressive ability of neurons. While effective for simple MLP-based agents, this approach loses statistical power in more complex architectures. To address this, we argue that in advanced RL agents, maintaining a neuron's learning capacity, its ability to adapt via gradient updates, is more critical than preserving its expressive ability. Based on this insight, we shift the statistical objective from activations to gradients, and introduce GraMa (Gradient Magnitude Neural Activity Metric), a lightweight, architecture-agnostic metric for quantifying neuron-level learning capacity. We show that GraMa effectively reveals persistent neuron inactivity across diverse architectures, including residual networks, diffusion models, and agents with varied activation functions. Moreover, resetting neurons guided by GraMa (ReGraMa) consistently improves learning performance across multiple deep RL algorithms and benchmarks, such as MuJoCo and the DeepMind Control Suite.