[2601.20071] Distributional value gradients for stochastic environments

Computer Science > Machine Learning arXiv:2601.20071 (cs) [Submitted on 27 Jan 2026 (v1), last revised 2 Mar 2026 (this version, v3)] Title:Distributional value gradients for stochastic environments Authors:Baptiste Debes, Tinne Tuytelaars View a PDF of the paper titled Distributional value gradients for stochastic environments, by Baptiste Debes and 1 other authors View PDF HTML (experimental) Abstract:Gradient-regularized value learning methods improve sample efficiency by leveraging learned models of transition dynamics and rewards to estimate return gradients. However, existing approaches, such as MAGE, struggle in stochastic or noisy environments, limiting their applicability. In this work, we address these limitations by extending distributional reinforcement learning on continuous state-action spaces to model not only the distribution over scalar state-action value functions but also over their gradients. We refer to this approach as Distributional Sobolev Training. Inspired by Stochastic Value Gradients (SVG), our method utilizes a one-step world model of reward and transition distributions implemented via a conditional Variational Autoencoder (cVAE). The proposed framework is sample-based and employs Max-sliced Maximum Mean Discrepancy (MSMMD) to instantiate the distributional Bellman operator. We prove that the Sobolev-augmented Bellman operator is a contraction with a unique fixed point, and highlight a fundamental smoothness trade-off underlying contraction in ...