Stable Deep Reinforcement Learning via Isotropic Gaussian Representations
Abstract
Deep reinforcement learning systems often suffer from unstable training dynamics due to non-stationarity, where learning objectives and data distributions evolve over time. We show that under non-stationary targets, isotropic Gaussian embeddings are provably advantageous. In particular, they induce stable tracking of time-varying targets for linear readouts, achieve maximal entropy under a fixed variance budget, and encourage a balanced use of all representational dimensions--all of which enable agents to be more adaptive and stable. Building on this insight, we propose the use of Sketched Isotropic Gaussian Regularization for shaping representations toward an isotropic Gaussian distribution during training. We demonstrate empirically, over a variety of domains, that this simple and computationally inexpensive method improves performance under non-stationarity while reducing representation collapse, neuron dormancy, and training instability.
Lay Summary
Deep reinforcement learning systems learn by continuously interacting with their environment. As they learn, both the data they observe and the targets they optimize change over time. This non-stationarity often causes training to become unstable, leading to issues such as representation collapse, dormant neurons, and reduced learning performance. While prior work has largely addressed these problems through new optimization algorithms or architectural changes, we investigate a different question: what properties should the learned representations themselves have in order to remain stable under changing learning conditions? We show theoretically and empirically that representations with an isotropic Gaussian structure are particularly well suited for learning under non-stationarity. Such representations distribute information evenly across dimensions, avoid over-reliance on a small number of features, and are less sensitive to changes in training targets. Building on this insight, we apply a lightweight regularization method that encourages learned representations to follow an isotropic Gaussian distribution. Across a range of reinforcement learning benchmarks, this simple modification improves training stability, reduces representation collapse and neuron dormancy, and leads to better performance. Our results suggest that representation geometry is a fundamental factor in stable reinforcement learning and that shaping representations may provide a simple alternative to more complex optimization-based solutions.