Distributional Reinforcement Learning with Diffusion Bridge Critics

TL;DR

DBC is the first work to employ the diffusion bridge model as the critic and derives an analytic integral formula to address discretization errors in DBC, which is essential in value estimation.

Abstract

Recent advances in diffusion-based reinforcement learning (RL) methods have demonstrated promising results in a wide range of continuous control tasks. However, existing works in this field focus on the application of diffusion policies while leaving the diffusion critics unexplored. In fact, since policy optimization fundamentally relies on the critic, accurate value estimation is far more important than policy expressiveness. Furthermore, given the stochasticity of most reinforcement learning tasks, it has been confirmed that the critic is more appropriately depicted with a distributional model. Motivated by these points, we propose a novel distributional RL method with Diffusion Bridge Critics (DBC). DBC directly models the inverse cumulative distribution function (CDF) of the Q value. This allows us to accurately capture the value distribution and prevents it from collapsing into a trivial Gaussian distribution owing to the strong distribution-matching capability of the diffusion bridge. Moreover, we further derive an analytic integral formula to address discretization errors in DBC, which is essential in value estimation. To our knowledge, DBC is the first work to employ the diffusion bridge model as the critic. Notably, DBC is also a plug-and-play component and can be integrated into most existing RL frameworks. Experimental results on MuJoCo robot control benchmarks demonstrate the superiority of DBC compared with previous distributional critic models.

Figures (3)

• Figure 1: The training pipeline of Diffusion Bridge Critics. Compared with Vanilla Diffusion Critic, DBC explicitly models the inverse cumulative distribution function (CDF) of Q-values and resolves the Gaussian Degradation Problem. Besides, the design of the integral-consistent discretization technique is also developed for accurate value estimation and stable policy optimization.
• Figure 2: Top row: Value diffusion fits the iterative target $r+\gamma \cdot Z_{k-1}^{\mathrm{FM}}$. Bottom row: DBC fits the iterative target $r+\gamma \cdot Z_{k-1}^{\mathrm{DBC}}$. DBC introduces quantile conditioning $\tau$, which consequently preserves multimodality of the distribution under iterative Bellman drift $Z_k \leftarrow r+\gamma Z_{k-1}$. We takes $100$ inner training steps to fit $Z_k$ for each method in iteration and 10000 steps to fit the initial distribution in Iter 0.
• Figure 3: Learning curves of different algorithms on five MuJoCo benchmarks. The x-axis denotes training epochs, and the y-axis denotes episodic return. Curves are smoothed for improved visualization.

Theorems & Definitions (14)

• Theorem 4.1: Gaussian Degradation of Diffusion Critics
• Theorem 4.2: Endpoint Consistency via Integral-Consistent Discretization
• Corollary 4.3: Integral-consistent discretization for DBC
• Theorem 1.1
• proof
• proof
• Lemma 3.1: Sample quantile as a minimizer of the quantile loss
• proof
• Theorem 3.2: Consistency of the sample quantile
• proof