Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Distributional value gradients for stochastic environments

Baptiste Debes, Tinne Tuytelaars

TL;DR

This work addresses the challenge of gradient-aware credit assignment in stochastic environments by extending distributional reinforcement learning to model not only returns but also their gradients. It introduces Distributional Sobolev Training and a Sobolev Bellman operator that jointly bootstraps return and gradient distributions, implemented via a differentiable world model (cVAE) and trained with Max–Sliced MMD (MSMMD). The authors prove contraction properties for the Sobolev operator under both Wasserstein and MSMMD metrics and demonstrate improved stability and robustness in a stochastic toy task and MuJoCo benchmarks, including ablations on world-model choices and overestimation bias controls. The proposed Distributional Sobolev Deterministic Policy Gradient (DSDPG) framework advances sample efficiency and gradient-informed policy improvement in uncertain environments, with broad implications for model-based, gradient-aware RL and related fields.

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 gradient-aware RL. To validate our method, we first showcase its effectiveness on a simple stochastic reinforcement learning toy problem, then benchmark its performance on several MuJoCo environments.

Distributional value gradients for stochastic environments

TL;DR

This work addresses the challenge of gradient-aware credit assignment in stochastic environments by extending distributional reinforcement learning to model not only returns but also their gradients. It introduces Distributional Sobolev Training and a Sobolev Bellman operator that jointly bootstraps return and gradient distributions, implemented via a differentiable world model (cVAE) and trained with Max–Sliced MMD (MSMMD). The authors prove contraction properties for the Sobolev operator under both Wasserstein and MSMMD metrics and demonstrate improved stability and robustness in a stochastic toy task and MuJoCo benchmarks, including ablations on world-model choices and overestimation bias controls. The proposed Distributional Sobolev Deterministic Policy Gradient (DSDPG) framework advances sample efficiency and gradient-informed policy improvement in uncertain environments, with broad implications for model-based, gradient-aware RL and related fields.

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 gradient-aware RL. To validate our method, we first showcase its effectiveness on a simple stochastic reinforcement learning toy problem, then benchmark its performance on several MuJoCo environments.
Paper Structure (72 sections, 30 theorems, 323 equations, 16 figures, 3 tables, 3 algorithms)

This paper contains 72 sections, 30 theorems, 323 equations, 16 figures, 3 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Paper contributions
  3. Background
  4. Notation and RL objective
  5. Temporal‐difference as an affine operator
  6. A new Bellman operator
  7. Learning a useful critic
  8. Distributional Sobolev training
  9. Complete Sobolev TD.
  10. Theoretical results
  11. Towards a tractable metric
  12. Max–Sliced MMD.
  13. Contraction under Max–Sliced MMD.
  14. Approach
  15. Results
  16. ...and 57 more sections

Key Result

Proposition 1

Let $\pi$ be an $L_\pi$-Lipschitz continuous policy, and let $G(s)=\mathrm{Law}[\nabla_a Z^\pi(s,a)\mid_{a=\pi(s)}]$ and $\hat{G}(s)=\mathrm{Law}[\nabla_a \hat{Z}(s,a)\mid_{a=\pi(s)}]$ denote the true and estimated distributions of the action‐gradients at $a=\pi(s)$, respectively. Define the $p$–Was Then, specializing to $p=1$, the error between the true policy gradient $\nabla_\theta J(\theta)$ a

Figures (16)

  • Figure 1: Left: block diagram of our DSDPG algorithm, where the critic $Z_\phi$ maps noise $\xi$ and $(s,a)$ to Sobolev‐return samples, a cVAE world model generates next‐state–reward samples $(\hat{s}',\hat{r})$, MSMMD or MMD compares predicted and bootstrapped Sobolev‐return distributions, and the policy is updated via the critic’s mean (gradient flows shown as dashed arrows; inspired by singh2020samplebased). Right: pseudocode for estimating the biased MSMMD between predicted and bootstrapped Sobolev returns.
  • Figure 2: (a) Illustration of the 2D point‐mass environment with $N$ possible bonus locations. (b) Evaluation curves (10 means the agent reached the bonus) for our Distributional Sobolev (MMD Sobolev and MSMMD Sobolev, in orange and pink resp.), deterministic Sobolev (MAGE doro2020how, in red), and other baselines as $N$ varies (median over 5 seeds and 25%-75% IQR).
  • Figure 3: Evaluation of DSDPG (MSMMD/MMD Sobolev), deterministic Sobolev/MAGE doro2020how, TD3-Huber fujimoto2018addressing, IQN dabney2018implicit, standard MMD DBLP:journals/corr/abs-2007-12354killingberg2023the and MSMMD on six MuJoCo tasks. Results are reported over 10 random seeds with median and 95% bootstrap confidence intervals. We compare three settings: normal environment, multiplicative observation noise, and Gaussian dynamics noise luo2021distributional. (\ref{['fig:mujoco_grid:a']}) Final evaluation performance. (\ref{['fig:mujoco_grid:b']}) Normalized AUC over the entire training curve. Our DSDPG variants (MMD Sobolev and MSMMD Sobolev, red and brown) are on par or better than competing methods, shining especially on harder tasks and under noisy environments.
  • Figure 4: (a) Samples of the marginals of the full Sobolev distribution $[f(x;a);\nabla_x f(x;a)]$: output (left) and gradient (right). Blue: samples from the true distribution; red: samples from the Distributional Sobolev model trained via MMD; green: samples from the deterministic Sobolev baseline czarnecki2017sobolev. (b) Biased MMD score (lower is better) on the joint variable. (c) $L_2$ error between predicted sample mean and true mean. NB. Predicted (red and green) and true (blue) samples are highly overlapping.
  • Figure 5: Toy supervised learning problem. Comparison between MMD-based or L2 based modeling. Left panel: training curve of L2 loss (logscale) on gradient between true conditional expectation with regression prediction and with empirical mean from MMD-based model. Sobolev (blue) used gradient information to train either using MMD (full line) or L2 regression (dashed line). Right panel: average over the input space of the second order derivative (logscale) of predicted gradient from deterministic model (blue), MMD / stochastic (yellow) and with gradient information / Sobolev (dashed). Metrics averaged over 5 seeds.
  • ...and 11 more figures

Theorems & Definitions (58)

  • Proposition 1
  • Theorem 1: Action‐gradient Sobolev contraction
  • Theorem 2: Action–gradient Sobolev contraction under $\mathbf{MS}\mathrm{MMD}$
  • Definition 1: $p$-Wasserstein distance villani2009optimal
  • Definition 2: Supremum-$p$–Wasserstein distance zhang2021distributional, eq. (12)
  • Lemma 1: Push‐forward law identity
  • proof
  • Lemma 2: Affine form of the Sobolev Bellman operator
  • proof
  • Lemma 3: Affine push‐forward contraction in a normed space
  • ...and 48 more