Moments Matter:Stabilizing Policy Optimization using Return Distributions

• URL: http://arxiv.org/abs/2601.01803v1
• Date: Mon, 05 Jan 2026 05:27:11 GMT
• Title: Moments Matter:Stabilizing Policy Optimization using Return Distributions
• Authors: Dennis Jabs, Aditya Mohan, Marius Lindauer,
• Abstract summary: In continuous control tasks, even small parameter shifts can produce unstable gaits.<n>We propose an alternative that takes advantage of environmentality to update-induced variability.
• Score: 9.430246534202857
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Provably Efficient Algorithms for S- and Non-Rectangular Robust MDPs with General Parameterization [85.91302339486673]
  We study robust Markov decision processes (RMDPs) with general policy parameterization under s-rectangular and non-rectangular uncertainty sets.<n>We prove novel Lipschitz and Lipschitz-smoothness properties for general policy parameterizations that extends to infinite state spaces.<n>We design a projected gradient descent algorithm for s-rectangular uncertainty and a Frank-Wolfe algorithm for non-rectangular uncertainty.
  arXiv  Detail & Related papers  (2026-02-11T21:44:20Z)
• Not All Rollouts are Useful: Down-Sampling Rollouts in LLM Reinforcement Learning [55.15106182268834]
  Reinforcement learning with verifiable rewards (RLVR) has emerged as the leading approach for enhancing reasoning capabilities in large language models.<n>It faces a fundamental compute and memory asymmetry: rollout generation is embarrassingly parallel and memory-light, whereas policy updates are communication-heavy and memory-intensive.<n>We introduce PODS (Policy Optimization with Down-Sampling), which decouples rollout generation from policy updates by training only on a strategically selected subset of rollouts.
  arXiv  Detail & Related papers  (2025-04-18T17:49:55Z)
• Distributionally Robust Policy Learning under Concept Drifts [33.44768994272614]
  This paper studies a more nuanced problem -- robust policy learning under the concept drift.<n>We first provide a doubly-robust estimator for evaluating the worst-case average reward of a given policy.<n>We then propose a learning algorithm that outputs the policy maximizing the estimated policy value within a given policy class.
  arXiv  Detail & Related papers  (2024-12-18T19:53:56Z)
• Near-Optimal Dynamic Regret for Adversarial Linear Mixture MDPs [63.47351876442425]
  We study episodic linear mixture MDPs with the unknown transition and adversarial rewards under full-information feedback. We propose a novel algorithm that combines the benefits of two popular methods: occupancy-measure-based and policy-based. Our algorithm enjoys an $widetildemathcalO(d sqrtH3 K + sqrtHK(H + barP_K$)$ dynamic regret, where $d$ is the feature dimension.
  arXiv  Detail & Related papers  (2024-11-05T13:55:52Z)
• Policy Gradient with Active Importance Sampling [55.112959067035916]
  Policy gradient (PG) methods significantly benefit from IS, enabling the effective reuse of previously collected samples. However, IS is employed in RL as a passive tool for re-weighting historical samples. We look for the best behavioral policy from which to collect samples to reduce the policy gradient variance.
  arXiv  Detail & Related papers  (2024-05-09T09:08:09Z)
• The Role of Baselines in Policy Gradient Optimization [83.42050606055822]
  We show that the emphstate value baseline allows on-policy. emphnatural policy gradient (NPG) to converge to a globally optimal. policy at an $O (1/t) rate gradient. We find that the primary effect of the value baseline is to textbfreduce the aggressiveness of the updates rather than their variance.
  arXiv  Detail & Related papers  (2023-01-16T06:28:00Z)
• Autoregressive Bandits [58.46584210388307]
  We propose a novel online learning setting, Autoregressive Bandits, in which the observed reward is governed by an autoregressive process of order $k$. We show that, under mild assumptions on the reward process, the optimal policy can be conveniently computed. We then devise a new optimistic regret minimization algorithm, namely, AutoRegressive Upper Confidence Bound (AR-UCB), that suffers sublinear regret of order $widetildemathcalO left( frac(k+1)3/2sqrtnT (1-G
  arXiv  Detail & Related papers  (2022-12-12T21:37:36Z)
• Global Optimality and Finite Sample Analysis of Softmax Off-Policy Actor Critic under State Distribution Mismatch [26.23228392999382]
  We establish the global optimality and convergence rate of an off-policy actor critic algorithm.<n>Our work goes beyond existing works on the optimality of policy gradient methods.
  arXiv  Detail & Related papers  (2021-11-04T16:48:45Z)
• Variance Penalized On-Policy and Off-Policy Actor-Critic [60.06593931848165]
  We propose on-policy and off-policy actor-critic algorithms that optimize a performance criterion involving both mean and variance in the return. Our approach not only performs on par with actor-critic and prior variance-penalization baselines in terms of expected return, but also generates trajectories which have lower variance in the return.
  arXiv  Detail & Related papers  (2021-02-03T10:06:16Z)
• Minimax-Optimal Off-Policy Evaluation with Linear Function Approximation [49.502277468627035]
  This paper studies the statistical theory of batch data reinforcement learning with function approximation. Consider the off-policy evaluation problem, which is to estimate the cumulative value of a new target policy from logged history.
  arXiv  Detail & Related papers  (2020-02-21T19:20:57Z)
• Distributional Soft Actor-Critic: Off-Policy Reinforcement Learning for Addressing Value Estimation Errors [13.534873779043478]
  We present a distributional soft actor-critic (DSAC) algorithm to improve the policy performance by mitigating Q-value overestimations. We evaluate DSAC on the suite of MuJoCo continuous control tasks, achieving the state-of-the-art performance.
  arXiv  Detail & Related papers  (2020-01-09T02:27:18Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.