Moments Matter:Stabilizing Policy Optimization using Return Distributions

TL;DR

This work addresses instability in policy optimization under environmental and algorithmic stochasticity by linking the post-update returns to a distributional critic. Instead of attempting to precisely estimate the costly post-update return distribution $\mathcal{R}(\theta)$, the authors model the full return distribution with a distributional critic $Z_\phi(s,a)$ using a $51$-quantile support and regularize the PPO advantage with higher-order moments (skewness and kurtosis) to penalize asymmetric and heavy-tailed updates. Empirical results on Brax continuous-control tasks show that moment-based regularization narrows $\mathcal{R}(\theta)$ by up to 75% in challenging setups with weak critic–return alignment (e.g., Walker2D), while preserving comparable evaluation performance; in settings with strong alignment (e.g., HalfCheetah), vanilla PPO already performs well. The approach offers a practical route to more stable policy updates in stochastic environments, with potential benefits for real-world transfer and reliability of RL algorithms.

Abstract

Deep Reinforcement Learning (RL) agents often learn policies that achieve the same episodic return yet behave very differently, due to a combination of environmental (random transitions, initial conditions, reward noise) and algorithmic (minibatch selection, exploration noise) factors. In continuous control tasks, even small parameter shifts can produce unstable gaits, complicating both algorithm comparison and real-world transfer. Previous work has shown that such instability arises when policy updates traverse noisy neighborhoods and that the spread of post-update return distribution $R(θ)$, obtained by repeatedly sampling minibatches, updating $θ$, and measuring final returns, is a useful indicator of this noise. Although explicitly constraining the policy to maintain a narrow $R(θ)$ can improve stability, directly estimating $R(θ)$ is computationally expensive in high-dimensional settings. We propose an alternative that takes advantage of environmental stochasticity to mitigate update-induced variability. Specifically, we model state-action return distribution through a distributional critic and then bias the advantage function of PPO using higher-order moments (skewness and kurtosis) of this distribution. By penalizing extreme tail behaviors, our method discourages policies from entering parameter regimes prone to instability. We hypothesize that in environments where post-update critic values align poorly with post-update returns, standard PPO struggles to produce a narrow $R(θ)$. In such cases, our moment-based correction narrows $R(θ)$, improving stability by up to 75% in Walker2D, while preserving comparable evaluation returns.

Figures (4)

• Figure 1: Overview of the architecture of PPO. We modify the critic to a distributional critic that uses a $51$-quantile support distribution as outlined by schneider-icra24a to approximate the return distribution. We then utilize the higher-order moments of this distribution to regularize the advantage estimation.
• Figure 2: Relationship between returns and critic values sampled from $\mathcal{R}(\theta)$ and $\mathcal{V}(\theta)$ of PPO using the final checkpoint across Brax environments.
• Figure 2: Standard deviation $\sigma$ of $\mathcal{R}(\theta)$, averaged over $10$ seeds, of policies produced by different algorithms across continuous control tasks. The lowest $\sigma$ in each environment is highlighted.
• Figure 3: Final evaluation returns across $10$ seeds for Brax environments.