Bellman Diffusion: Generative Modeling as Learning a Linear Operator in the Distribution Space

## Abstract

Abstract Deep Generative Models (DGMs), including Energy-Based Models (EBMs) and Score-based Generative Models (SGMs), have advanced high-fidelity data generation and complex continuous distribution approximation. However, their application in Markov Decision Processes (MDPs), particularly in distributional Reinforcement Learning (RL), remains underexplored, with the classical histogram-based methods dominating the field. This paper rigorously highlights that this application gap is caused by the nonlinearity of modern DGMs, which conflicts with the linearity required by the Bellman equation in MDPs. For instance, EBMs involve nonlinear operations such as exponentiating energy functions and normalizing constants. To address this problem, we introduce Bellman Diffusion , a novel DGM framework that maintains linearity in MDPs through gradient and scalar field modeling. With divergence-based training techniques to optimize neural network proxies and a new type of stochastic differential equation (SDE) for sampling, Bellman Diffusion is guaranteed to converge to the target distribution. Our empirical results show that Bellman Diffusion achieves accurate field estimations and is a capable image generator, converging 1.5 × 1.5\times 1.5 × faster than the traditional histogram-based baseline in distributional RL tasks. This work enables the effective integration of DGMs into MDP applications, unlocking new avenues for advanced decision-making frameworks.

### 1 Introduction

Even though Deep Generative Models (DGMs), such as Energy-Based Models (EBMs) (Teh et al., ), Generative Adversarial Networks (GANs) (Goodfellow et al., ), and the emerging Score-based Generative Models (SGMs) (Song et al., ; Sohl-Dickstein et al., ; Ho et al., ), have been highly developed to achieve high-fidelity generation (Li et al., ; Wang et al., ; Krishnamoorthy et al., ) and accurately approximate complex continuous distributions, their applications in Markov Decision Processes (MDPs), including Planning and Reinforcement Learning (RL), remain underexplored. Instead, a classical histogram-based methods (e.g., C51 (Bellemare et al., )) are still widely used in MDPs. We identify the main reason for this gap within the core of MDPs: the Bellman equation (Bellemare et al., ), which supports efficient model training but imposes a linear structure that conflicts with the strong nonlinearity of modern DGMs. In the following, we present a case study that shows this bottleneck using EBMs.

##### Motivation and problem.

The easiest way to model the return distribution with EBMs is to sample full state-action trajectories, so that one can directly use the returns computed from trajectories to train EBMs. However, this is not scalable since trajectory sampling is very costly in many RL environments. Some such examples are provided in Appendix D. An alternative is to follow the paradigm of Deep Q-Learning (Mnih et al., ), which applies the distributional Bellman equation to update the model with only partial trajectories. This equation is linear in form, which connects the return distribution $p_{z}(x)$ of a state $z$ with that of the next state $z^{\prime}$, expressed formally as:

$$ $p_{z}(x)=\sum_{z^{\prime},r}\alpha_{z,z^{\prime},r}p_{z^{\prime}}\Big{(}\frac{ x-r}{\gamma}\Big{)},$ (1) $$

where $r$ is the expected reward for transitioning between states, and $\alpha_{z,z^{\prime},r}$, $\gamma$ are constants determined by the RL environment. Although EBMs (or existing DGMs) are effective in distribution approximation (Lee et al., ), their application to the Bellman equation is constrained by the nonlinearity of how it is modeling
$\frac{e^{-E_{z}(x)}}{\mathcal{Z}_{z}}$, where $\mathcal{Z}_{z}$ is the normalization factor at each state $z$. To simply the case, suppose that the term $\mathcal{Z}_{z}$ is always $1$ for every state $z$, then we have:

$$ $\begin{aligned} E_{z}(x)&=-\log(\mathcal{Z}_{z}p_{z}(x))=-\log\Big{(}\sum_{z^{ \prime},r}\alpha_{z,z^{\prime},r}p_{z^{\prime}}\Big{(}\frac{x-r}{\gamma}\Big{) }\Big{)}\\ &\leq-\sum_{z^{\prime},r}\alpha_{z,z^{\prime},r}\log p_{z^{\prime}}\Big{(} \frac{x-r}{\gamma}\Big{)}=\sum_{z^{\prime},r}\alpha_{z,z^{\prime},r}E_{z}\Big{ (}\frac{x-r}{\gamma}\Big{)}\end{aligned},$ (2) $$

which holds due to Jensen’s inequality, indicating that the energy function $E_{z}(\mathbf{x})$ for a state $z$ is not linearly interpolated by the energy function for the next state $z^{\prime}$. As such, EBMs act as a nonlinear operator, disrupting the linearity of the distributional Bellman equation, thereby rendering EBMs inapplicable in this context. In Sec. 2, we analyze the modeling approaches of other modern DGMs and find that none can preserve the linearity of the Bellman update.

##### Our framework: Bellman Diffusion.

In this work, we re-examine the modeling approach of well-known DGMs to assess whether they meet the desired linearity property (Sec. 2), revealing that none of the existing DGMs satisfy this requirement. To address this bottleneck, we propose *Bellman Diffusion*, a novel DGM designed to overcome bottlenecks in applying DGMs to MDPs. The core idea is to model the gradient field $\nabla p_{z}(x)$ and scalar field $p_{z}(x)$ directly. Since these fields are linear operators, the linearity of Bellman equations is well preserved. For instance, after applying the gradient operator $\nabla$, the Bellman equation remains linear:

$$ $\nabla p_{z}(x)=\sum_{z^{\prime},r}\frac{\alpha_{z,z^{\prime},r}}{\gamma} \nabla p_{z^{\prime}}\Big{(}\frac{x-r}{\gamma}\Big{)}.$ (3) $$

We now use $p_{\mathrm{target}}$ to denote the target density of each state, replacing the previous notation $p_{z}$. Since $\nabla p_{\mathrm{target}}({\mathbf{x}})$ and $p_{\mathrm{target}}({\mathbf{x}})$ are generally inaccessible, we introduce field-based divergence measures (Definition 3.1) and transform them into training objectives (Eqs. (10), (11)): approximating fields $\nabla p_{\mathrm{target}}({\mathbf{x}})$ and $p_{\mathrm{target}}({\mathbf{x}})$ with neural network proxies ${\mathbf{g}}_{\bm{\phi}}$ and $s_{\bm{\varphi}}$.

Given these proxies, we introduce a new sampling method: Bellman Diffusion Dynamics, associated with the fields represented by the following stochastic differential equation (SDE):

$$ $\mathop{}\!\mathrm{d}{\mathbf{x}}(t)=\underbrace{\nabla p_{\mathrm{target}}({ \mathbf{x}}(t))}_{\approx{\mathbf{g}}_{\bm{\phi}}}\mathop{}\!\mathrm{d}t+ \underbrace{\sqrt{p_{\mathrm{target}}({\mathbf{x}}(t))}}_{\approx\sqrt{s_{\bm{ \varphi}}}}\mathop{}\!\mathrm{d}{\mathbf{w}}(t),\quad\text{starting from }{ \mathbf{x}}(0)\sim p_{0},$ (4) $$

where ${\mathbf{w}}(t)$ is a Brownian process and $p_{0}$ is any initial distribution. Once the fields are well approximated, we can replace the field terms in the above equation with learned proxies, resulting in a proxy SDE that can be solved forward in time to sample from $p_{\mathrm{target}}({\mathbf{x}})$.

##### Theoretical and empirical results.

Theoretically, we guarantee the convergence of our Bellman Diffusion Dynamics to the stationary distribution $p_{\mathrm{target}}$, regardless of the initial distribution (Theorem 4.1), and provide an error bound analysis accounting for neural network approximation errors (Theorem 4.2). Thus, Bellman Diffusion is a reliable standalone generative model.

Experimentally, we show the generative capabilities of Bellman Diffusion on real and synthetic datasets, confirming accurate field estimations, with promising results in image generation. We further apply Bellman Diffusion to classical distributional RL tasks, resulting in much more stable and $1.5\times$ faster convergence compared to the widely used histogram method (Bellemare et al., ). Notably, it can effectively learn and recover the target distributions with multiple unbalanced modes (Fig. 2 of Sec. 6), a challenge for score-based methods (Song & Ermon, ) due to the inherent nature of the score function.

In summary, Bellman Diffusion introduced in this paper stands as a novel and mathematically grounded generative modeling approach, paving the way for continuous density modeling in various applications within MDPs, such as Planning and distributional RL.

### 2 Linear Property for MDPs

In this section, we review modern DGMs and highlight the desired property to facilitate density estimation with Bellman updates, avoiding full trajectory updates.

#### 2.1 Modelings of Modern Deep Generative Models

DGMs aim to model the complex target distribution $p_{\mathrm{target}}(\mathbf{x})$ using a neural network-approximated *continuous* density, enabling new samples generation.
Below, we review well-known DGMs and offer high-level insights into how they define a *modeling operator* $\mathcal{M}$ that connects their own modeling functions to the desired density or its related statistics.

Energy-Based Models (EBMs) (Teh et al., ): These models define an energy function $E(\mathbf{x})$ and represent the probability as: $p_{\mathrm{target}}(\mathbf{x})\approx\frac{e^{-E(\mathbf{x})}}{Z}$, where $Z:=\int e^{-E(\mathbf{x})}\mathop{}\!\mathrm{d}{\mathbf{x}}$ is the partition function for normalizing probabilities. EBM defines a modeling operator $\mathcal{M}_{\mathrm{EBM}}\colon E(\cdot)\mapsto\frac{e^{-E(\cdot)}}{Z}$, linking the statistic $E(\cdot)$ to desired density.

Flow-Based Models (Rezende & Mohamed, ; Chen et al., ): These use a series of invertible transformations ${\mathbf{f}}({\mathbf{x}})$ to map data $\mathbf{x}$ to a latent space $\mathbf{z}$, with an exact likelihood: $p_{\mathrm{target}}(\mathbf{x})\approx\pi(\mathbf{z})\left|\det\frac{\partial{
\mathbf{f}}^{-1}({\mathbf{x}})}{\partial\mathbf{x}}\right|$. It determines a modeling operator $\mathcal{M}_{\mathrm{Flow}}\colon{\mathbf{f}}(\cdot)\mapsto\pi({\mathbf{f}}(
\cdot))\left|\det\frac{\partial{\mathbf{f}}^{-1}(\cdot)}{\partial\mathbf{x}}\right|$, connecting the transformation ${\mathbf{f}}(\cdot)$ to desired density.

Implicit Latent Variable Models: These models define a latent variable $\mathbf{z}$ and use a generative process $p(\mathbf{x}|\mathbf{z})$, where the latent space is sampled from a prior $\pi(\mathbf{z})$, usually taken as a standard normal distribution. Two popular models are VAE (Kingma, ) and GAN (Goodfellow et al., ). VAE maximizes a variational lower bound using an encoder network $q(\mathbf{z}|\mathbf{x})$ to approximate the posterior distribution, while GAN employs a discriminator to distinguish between real and generated data, with a generator learning to produce realistic samples but lacking an explicit likelihood.

Since VAEs and GANs are implicit models, they lack an explicit modeling operator like $\mathcal{M}_{\mathrm{EBM}}$ and $\mathcal{M}_{\mathrm{Flow}}$ that connects modeling functions to the desired density or its related statistics.

Score-Based Generative Models (SGMs) (Song et al., ): They involve a process that gradually adds noise to $p_{\mathrm{target}}$, resulting in a sequence of time-conditioned densities $\{p({\mathbf{x}}_{t},t)\}_{t\in[0,T]}$, where $t=0$ corresponds to $p_{\mathrm{target}}$ and $t=T$ corresponds to a simple prior distribution $\pi({\mathbf{z}})$. Then, SGMs reverse this diffusion process for sampling by employing the time-conditioned score ${\mathbf{s}}(\cdot,t):=\nabla\log p(\cdot,t)$ and solving the ordinary differential equation (Song et al., ) from $t=T$ to $t=0$ with $\bm{\phi}_{T}({\mathbf{x}}_{T})={\mathbf{x}}_{T}\sim\pi$:
$\mathop{}\!\mathrm{d}\bm{\Psi}_{t}({\mathbf{x}}_{T})=\left({\mathbf{f}}(\bm{
\Psi}_{t}({\mathbf{x}}_{T}),t)-\frac{1}{2}g^{2}(t){\mathbf{s}}(\bm{\Psi}_{t}({
\mathbf{x}}_{T}),t)\right)\mathop{}\!\mathrm{d}t$, where ${\mathbf{f}}$ and $g$ are pre-determined. This flow defines a pushforward map $\mathcal{V}^{T\rightarrow t}[{\mathbf{s}}]$ of the density as
$\mathcal{V}^{T\rightarrow t}[{\mathbf{s}}]\{\pi\}:=\pi\left(\bm{\Psi}_{t}^{-1}
(\cdot)\right)\left|\det\frac{\partial\bm{\Psi}_{t}^{-1}(\cdot)}{\partial
\mathbf{x}}\right|$. Thus, SGMs determine a modeling operator $\mathcal{M}_{\mathrm{SGM}}\colon{\mathbf{s}}\mapsto\mathcal{V}^{T\rightarrow 0
}[{\mathbf{s}}]\{\pi\}\approx p_{\mathrm{target}}$.

#### 2.2 Desired Linear Property in MDP

As the case of EBMs shown in Sec. 1, to leverage the strong capability of DGMs in density modeling with the Bellman update (Eq. (1)), the linearity of modeling operator $\mathcal{M}$ is crucial:
{mdframed}
Linear property of modeling. The modeling operator $\mathcal{M}$ defined by a DGM is linear: $\mathcal{M}\big{(}af+bg\big{)}=a\mathcal{M}\big{(}f\big{)}+b\mathcal{M}\big{(}
g\big{)}$, for any reals $a,b$ and functions $f,g$.
If $\mathcal{M}$ is linear, we can link future state densities or their statistics with the current state for efficient updates, as shown in the Bellman equation in Eq. (1):

$$ $\mathcal{M}(p_{z})(x)=\sum_{z^{\prime},r}\alpha_{z,z^{\prime},r}\mathcal{M}(p_ {z^{\prime}})\left(\frac{x-r}{\gamma}\right).$ $$

However, for the current well-established DGMs, their modeling operators are either not explicitly defined (e.g., VAE and GAN), thus lacking guaranteed linearity, or are nonlinear operators (e.g., $\mathcal{M}_{\mathrm{EBM}}$, $\mathcal{M}_{\mathrm{Flow}}$, and $\mathcal{M}_{\mathrm{SGM}}$ as shown in Ineq. (2)). Consequently, this restricts the application of these powerful DGMs to MDPs.

### 3 Method: Bellman Diffusion

In this section, we mainly provide an overview of Bellman Diffusion, presenting the usage, with its theoretical details later in Sec. 4. We defer all proofs to Appendix 3.

#### 3.1 Scalar and Vector Field Matching

##### Field matching.

Suppose we have a finite set of $D$-dimensional sample vectors $\mathcal{X}=\{\mathbf{x}_{i}\}_{1\leq i\leq N}$, where each point $\mathbf{x}_{i}$ is drawn from the distribution ${p_{\mathrm{target}}({\mathbf{x}})}$. Bellman Diffusion, as a generative model, aims to learn both the gradient field $\nabla{p_{\mathrm{target}}({\mathbf{x}})}$ and the scalar field ${p_{\mathrm{target}}({\mathbf{x}})}$. Like Fisher divergence (Antolín et al., ) for the score function $\nabla\log{p_{\mathrm{target}}({\mathbf{x}})}$, we introduce two divergence measures for $\nabla{p_{\mathrm{target}}({\mathbf{x}})}$ and ${p_{\mathrm{target}}({\mathbf{x}})}$.

###### Definition 3.1 (Field Divergences) .

Let $p(\cdot)$ and $q(\cdot)$ be continuous probability densities. The discrepancy between the two can be defined as

$$ $\mathcal{D}_{\mathrm{grad}}\big{(}p(\cdot),q(\cdot)\big{)}=\int p(\mathbf{x}) \left\lVert\nabla p(\mathbf{x})-\nabla q(\mathbf{x})\right\rVert^{2}\mathop{} \!\mathrm{d}{\mathbf{x}}$ (5) $$

using the gradient operator $\nabla$ in terms of ${\mathbf{x}}$, or as

$$ $\mathcal{D}_{\mathrm{id}}\big{(}p(\cdot),q(\cdot)\big{)}=\int p(\mathbf{x})(p( \mathbf{x})-q(\mathbf{x}))^{2}\mathop{}\!\mathrm{d}{\mathbf{x}}$ (6) $$

using the identity operator $\mathbb{I}$. Here, $\left\lVert\cdot\right\rVert$ denotes the $\ell_{2}$ norm.

As shown in Appendix B.2, the two measures above are valid statistical measures. These measures are used to empirically estimate the gradient field $\nabla{p_{\mathrm{target}}({\mathbf{x}})}$ and the scalar field ${p_{\mathrm{target}}({\mathbf{x}})}$ from real data $\mathcal{X}$. . Furthermore, our modeling defines a modeling operator given by $\mathcal{M}_{\mathrm{Bellman}}:=\begin{bmatrix}\nabla\\
\mathbb{I}\end{bmatrix}\colon p_{\mathrm{target}}(\cdot)\mapsto\begin{bmatrix}
\nabla p_{\mathrm{target}}(\cdot)\\
p_{\mathrm{target}}(\cdot)\end{bmatrix}$ which is linear in its input.

Similar to SGMs, we parameterize two neural networks, $\mathbf{g}_{\bm{\phi}}(\mathbf{x})$ and $s_{\bm{\varphi}}(\mathbf{x})\geq 0$, to approximate these fields, with learnable parameters $\bm{\phi}$ and $\bm{\varphi}$, using the following estimation loss functions:

$$ $\left\{\begin{aligned} \mathcal{L}_{\mathrm{grad}}(\bm{\phi})&:=\mathcal{D}_{ \mathrm{grad}}\big{(}{p_{\mathrm{target}}(\cdot)},\mathbf{g}_{\bm{\phi}}(\cdot )\big{)}=\mathbb{E}_{\mathbf{x}\sim{p_{\mathrm{target}}({\mathbf{x}})}}\Big{[} \|\nabla{p_{\mathrm{target}}({\mathbf{x}})}-\mathbf{g}_{\bm{\phi}}(\mathbf{x}) \|^{2}\Big{]}\\ \mathcal{L}_{\mathrm{id}}(\bm{\varphi})&:=\mathcal{D}_{\mathrm{id}}\big{(}{p_{ \mathrm{target}}(\cdot)},s_{\bm{\varphi}}(\cdot)\big{)}=\mathbb{E}_{\mathbf{x} \sim{p_{\mathrm{target}}({\mathbf{x}})}}\Big{[}({p_{\mathrm{target}}({\mathbf{ x}})}-s_{\bm{\varphi}}(\mathbf{x}))^{2}\Big{]}\end{aligned}\right.$ (7) $$

Since the terms $\nabla{p_{\mathrm{target}}({\mathbf{x}})}$ and ${p_{\mathrm{target}}({\mathbf{x}})}$ inside the expectation are generally inaccessible, these losses cannot be estimated via Monte Carlo sampling. The following proposition resolves this issue by deriving a feasible proxy for the loss functions

###### Proposition 3.1 (Equivalent Forms of Field Matching) .

The loss $\mathcal{L}_{\mathrm{grad}}(\bm{\phi})$ is given by

$$ $\mathcal{L}_{\mathrm{grad}}(\bm{\phi})=C_{\mathrm{grad}}+\lim_{\epsilon \rightarrow 0}\mathbb{E}_{\mathbf{x}_{1},\mathbf{x}_{2}\sim{p_{\mathrm{target} }({\mathbf{x}})}}\Big{[}\|\mathbf{g}_{\bm{\phi}}(\mathbf{x}_{1})\|^{2}+\mathrm {tr}(\nabla\mathbf{g}_{\bm{\phi}}(\mathbf{x}_{1}))\mathcal{N}(\mathbf{x}_{2}- \mathbf{x}_{1};\mathbf{0},\epsilon\mathbf{I}_{D})\Big{]},$ $$

and $\mathcal{L}_{\mathrm{id}}(\bm{\varphi})$ is expressed as

$$ $\mathcal{L}_{\mathrm{id}}(\bm{\varphi})=C_{\mathrm{id}}+\lim_{\epsilon \rightarrow 0}\mathbb{E}_{\mathbf{x}_{1},\mathbf{x}_{2}\sim{p_{\mathrm{target} }({\mathbf{x}})}}\Big{[}s_{\bm{\varphi}}(\mathbf{x}_{1})^{2}-2s_{\bm{\varphi}} (\mathbf{x}_{1})\mathcal{N}(\mathbf{x}_{2}-\mathbf{x}_{1};\mathbf{0},\epsilon \mathbf{I}_{D})\Big{]}.$ $$

Here, $\mathcal{N}({\mathbf{x}};\mathbf{0},\epsilon\mathbf{I}_{D})$ denotes a $D$-dimensional isotropic Gaussian density function with ${\mathbf{x}}$, and $C_{\mathrm{grad}}$ and $C_{\mathrm{id}}$ are constants independent of the model parameters $\bm{\phi}$ and $\bm{\varphi}$.

Building on the above proposition, we can obtain feasible approximations of the training losses. With $\epsilon$ fixed to be sufficiently small (see Sec. 6 for experimental setups), we have:

$$ $\displaystyle\left\{\begin{aligned} \widebar{\mathcal{L}}_{\mathrm{grad}}(\bm{ \phi};\epsilon)&:=C_{\mathrm{grad}}+\mathbb{E}_{\mathbf{x}_{1},\mathbf{x}_{2} \sim{p_{\mathrm{target}}({\mathbf{x}})}}\Big{[}\|\mathbf{g}_{\bm{\phi}}( \mathbf{x}_{1})\|^{2}+\mathrm{tr}(\nabla\mathbf{g}_{\bm{\phi}}(\mathbf{x}_{1}) )\mathcal{N}(\cdot,\mathbf{0},\epsilon\mathbf{I}_{D})\Big{]}\approx\mathcal{L} _{\mathrm{grad}}(\bm{\phi}),\\ \widebar{\mathcal{L}}_{\mathrm{id}}(\bm{\varphi};\epsilon)&:=C_{\mathrm{id}}+ \mathbb{E}_{\mathbf{x}_{1},\mathbf{x}_{2}\sim{p_{\mathrm{target}}({\mathbf{x}} )}}\Big{[}s_{\bm{\varphi}}(\mathbf{x}_{1})^{2}-2s_{\bm{\varphi}}(\mathbf{x}_{1 })\mathcal{N}(\mathbf{x}_{2}-\mathbf{x}_{1},\mathbf{0},\epsilon\mathbf{I}_{D}) \Big{]}\approx\mathcal{L}_{\mathrm{id}}(\bm{\varphi})\end{aligned}\right.$ (8) $$

We note that scalar and gradient fields can be modeled independently.
Moreover, as Bellman Diffusion directly matches these fields, it eliminates the need for the normalizing constant associated with costly spatial integrals in the density network required by EBMs.

###### Definition 3.1 (Field Divergences) .

###### Proposition 3.1 (Equivalent Forms of Field Matching) .

#### 3.2 Efficient Field Matching Losses

##### Slice trick for efficient training.

While the loss functions $\widebar{\mathcal{L}}_{\mathrm{grad}}(\bm{\phi};\epsilon)$ and $\widebar{\mathcal{L}}_{\mathrm{id}}(\bm{\varphi};\epsilon)$ support Monte Carlo estimation, the term $\mathrm{tr}(\nabla\mathbf{g}_{\bm{\phi}}(\mathbf{x}_{1}))$ in $\widebar{\mathcal{L}}_{\mathrm{grad}}(\bm{\phi};\epsilon)$ is computationally expensive, limiting the scalability in high dimensions. To address this problem, we apply the slice trick (Kolouri et al., ; Song et al., ) to estimate the trace term efficiently. The resulting objective is summarized in the following proposition.

###### Proposition 3.2 (Sliced Gradient Matching) .

We define the sliced version of $\mathcal{L}_{\mathrm{grad}}$ (i.e., Eq. (5)) as

$$ $\mathcal{L}_{\mathrm{grad}}^{\mathrm{slice}}(\bm{\phi})=\mathbb{E}_{\mathbf{v} \sim q(\mathbf{v}),\mathbf{x}\sim{p_{\mathrm{target}}({\mathbf{x}})}}\Big{[} \big{(}\mathbf{v}^{T}\nabla{p_{\mathrm{target}}({\mathbf{x}})}-\mathbf{v}^{T} \mathbf{g}_{\bm{\phi}}(\mathbf{x})\big{)}^{2}\Big{]},$ $$

where $\mathbf{v}$ represents the slice vector drawn from a continuous distribution $q(\mathbf{v})$. This sliced loss also has an equivalent form:

$$ $\mathcal{L}_{\mathrm{grad}}^{\mathrm{slice}}(\bm{\phi})=C_{\mathrm{grad}}^{ \prime}+\lim_{\epsilon\rightarrow 0}\mathbb{E}_{\begin{subarray}{c}\mathbf{v} \sim q(\mathbf{v});\\ \mathbf{x}_{1},\mathbf{x}_{2}\sim{p_{\mathrm{target}}({\mathbf{x}})}\end{ subarray}}\Big{[}(\mathbf{v}^{\top}\mathbf{g}_{\bm{\phi}}(\mathbf{x}_{1}))^{2} +(\mathbf{v}^{\top}\nabla_{\mathbf{x}_{1}}\mathbf{g}_{\bm{\phi}}(\mathbf{x}_{1 })\mathbf{v})\mathcal{N}(\mathbf{x}_{2}-\mathbf{x}_{1};\mathbf{0},\epsilon \mathbf{I}_{D})\Big{]},$ $$

where $C_{\mathrm{grad}}^{\prime}$ is another constant independent of the model parameters.

Similar to Eq. (8), we can define a proxy loss for $\mathcal{L}_{\mathrm{grad}}^{\mathrm{slice}}(\bm{\phi})$ as follows with a sufficiently small $\epsilon$:

$$ $\mathbb{E}_{\begin{subarray}{c}\mathbf{v}\sim q(\mathbf{v});\mathbf{x}_{1}, \mathbf{x}_{2}\sim{p_{\mathrm{target}}({\mathbf{x}})}\end{subarray}}\Big{[}( \mathbf{v}^{\top}\mathbf{g}_{\bm{\phi}}(\mathbf{x}_{1}))^{2}+(\mathbf{v}^{\top }\nabla_{\mathbf{x}_{1}}\mathbf{g}_{\bm{\phi}}(\mathbf{x}_{1})\mathbf{v}) \mathcal{N}(\mathbf{x}_{2}-\mathbf{x}_{1};\mathbf{0},\epsilon\mathbf{I}_{D}) \Big{]},$ (9) $$

which allows Monte Carlo estimation from samples $\mathcal{X}$. This proxy loss serves as a reasonable estimator (Lai et al., ) for $\mathcal{L}_{\mathrm{grad}}(\bm{\phi})$.

###### Proposition 3.2 (Sliced Gradient Matching) .

##### Slice trick for improving sample efficiency.

When the data dimension $D$ is large, the multiplier $\mathcal{N}(\mathbf{x}_{2}-\mathbf{x}_{1};\mathbf{0},\epsilon\mathbf{I}_{D})$ in the loss functions: $\widebar{\mathcal{L}}_{\mathrm{grad}}(\bm{\phi};\epsilon)$ and $\widebar{\mathcal{L}}_{\mathrm{id}}(\bm{\varphi};\epsilon)$, will become nearly zero due to the $(2\pi)^{-D/2}$ factor, requiring a very large batch size for accurate Monte Carlo estimation and leading to low data efficiency.

To resolve this issue, we apply an additional slice trick, projecting the $D$-dimensional Gaussian density $\mathcal{N}(\mathbf{x}_{2}-\mathbf{x}_{1};\mathbf{0},\epsilon\mathbf{I}_{D})$ into a 1-dimensional density $\mathcal{N}(\mathbf{w}^{T}\mathbf{x}_{2}-\mathbf{w}^{T}\mathbf{x}_{1},0,\epsilon)$ along a random direction $\mathbf{w}\sim q({\mathbf{w}})$, where $\mathbf{w}$ follows a slice vector distribution $q({\mathbf{w}})$. Combining with Eq. (9), this results in our ultimate gradient field matching loss:
{mdframed}

$$ $\widebar{\mathcal{L}}_{\mathrm{grad}}^{\mathrm{slice}}(\bm{\phi};\epsilon):= \mathbb{E}_{\begin{subarray}{c}\mathbf{w}\sim q(\mathbf{w}),\mathbf{v}\sim q( \mathbf{v});\\ \mathbf{x}_{1},\mathbf{x}_{2}\sim{p_{\mathrm{target}}({\mathbf{x}})}\end{ subarray}}\Big{[}(\mathbf{v}^{\top}\mathbf{g}_{\bm{\phi}}(\mathbf{x}_{1}))^{2} +(\mathbf{v}^{\top}\nabla_{\mathbf{x}_{1}}\mathbf{g}_{\bm{\phi}}(\mathbf{x}_{1 })\mathbf{v})\mathcal{N}(\mathbf{w}^{T}\mathbf{x}_{2}-\mathbf{w}^{T}\mathbf{x} _{1};0,\epsilon)\Big{]}.$ (10) $$

Similarly, we apply the same trick to $\widebar{\mathcal{L}}_{\mathrm{id}}(\bm{\varphi};\epsilon)$ for dimension projection and obtain:
{mdframed}

$$ $\widebar{\mathcal{L}}_{\mathrm{id}}^{\mathrm{slice}}(\bm{\varphi};\epsilon):= \mathbb{E}_{\begin{subarray}{c}\mathbf{w}\sim q(\mathbf{w});\\ \mathbf{x}_{1},\mathbf{x}_{2}\sim{p_{\mathrm{target}}({\mathbf{x}})}\end{ subarray}}\Big{[}s_{\bm{\varphi}}(\mathbf{x}_{1})^{2}-2s_{\bm{\varphi}}( \mathbf{x}_{1})\mathcal{N}(\mathbf{w}^{T}\mathbf{x}_{2}-\mathbf{w}^{T}\mathbf{ x}_{1};0,\epsilon)\Big{]}.$ (11) $$

We adopt $\widebar{\mathcal{L}}_{\mathrm{grad}}^{\mathrm{slice}}(\bm{\phi};\epsilon)$ and $\widebar{\mathcal{L}}_{\mathrm{id}}^{\mathrm{slice}}(\bm{\varphi};\epsilon)$ for vector and scalar field matching losses, as they offer more practical and efficient objectives than $\mathcal{L}_{\mathrm{grad}}(\bm{\phi})$ and $\mathcal{L}_{\mathrm{id}}(\bm{\varphi})$, respectively. Empirically, these adaptations significantly stabilize the model in experiments.

#### 3.3 Bellman Diffusion Dynamics

Suppose that neural networks $\mathbf{g}_{\bm{\phi}}(\mathbf{x}),s_{\bm{\varphi}}(\mathbf{x})$ accurately estimate the target fields $\nabla p_{\mathrm{target}}({\mathbf{x}})$ and $p_{\mathrm{target}}({\mathbf{x}})$, one can sample from $p_{\mathrm{target}}({\mathbf{x}})$ by approximating the score function as:

$$ $\nabla\log{p_{\mathrm{target}}({\mathbf{x}})}=\frac{\nabla{p_{\mathrm{target}} ({\mathbf{x}})}}{{p_{\mathrm{target}}({\mathbf{x}})}}\approx\frac{\mathbf{g}_{ \bm{\phi}}(\mathbf{x})}{s_{\bm{\varphi}}(\mathbf{x})}$ $$

and then applying Langevin dynamics (Bussi & Parrinello, ):

$$ $\mathop{}\!\mathrm{d}\mathbf{x}(t)=\nabla\log{p_{\mathrm{target}}({\mathbf{x}} )}\mathop{}\!\mathrm{d}t+\sqrt{2}\mathop{}\!\mathrm{d}\bm{\omega}(t)\approx \frac{\mathbf{g}_{\bm{\phi}}(\mathbf{x})}{s_{\bm{\varphi}}(\mathbf{x})}\mathop {}\!\mathrm{d}t+\sqrt{2}\mathop{}\!\mathrm{d}\bm{\omega}(t),$ $$

where $\bm{\omega}(t)$ is a standard Brownian motion. However, this approach can be numerically unstable due to the division(^1^11For example, if $s_{\bm{\varphi}}(\mathbf{x})$ is around 0.01, its inverse can magnify the estimation error of $\mathbf{g}_{\bm{\phi}}(\mathbf{x})$ by 100 times.). This issue is unavoidable as $p_{\mathrm{target}}({\mathbf{x}})$ vanishes when $\left\lVert\mathbf{x}\right\rVert\rightarrow\infty$. Additionally, it doesn’t support the distributional Bellman update for MDPs as mentioned in Sec. 1. To solve this, we propose a new SDE for sampling from $p_{\mathrm{target}}({\mathbf{x}})$, called *Bellman Diffusion Dynamics*:
{mdframed}

$$ $\mathop{}\!\mathrm{d}\mathbf{x}(t)=\nabla p_{\mathrm{target}}(\mathbf{x}(t)) \mathop{}\!\mathrm{d}t+\sqrt{p_{\mathrm{target}}(\mathbf{x}(t))}\mathop{}\! \mathrm{d}\bm{\omega}(t).$ (12) $$

We also provide the theoretical motivation and derivation of Eq. (12) in Appendix B.1.

In practice, once the neural network approximations $\mathbf{g}_{\bm{\phi}}(\mathbf{x})\approx\nabla p_{\mathrm{target}}(\mathbf{x})$ and $s_{\bm{\varphi}}(\mathbf{x})\approx p_{\mathrm{target}}(\mathbf{x})$ are both well-learned, we can derive the following *empirical Bellman Diffusion Dynamics*, a feasible proxy SDE for Eq. (12):

$$ $\mathop{}\!\mathrm{d}\mathbf{x}(t)=\mathbf{g}_{\bm{\phi}}(\mathbf{x})\mathop{} \!\mathrm{d}t+\sqrt{s_{\bm{\varphi}}(\mathbf{x})}\mathop{}\!\mathrm{d}\bm{ \omega}(t).$ (13) $$

Bellman Diffusion learns and samples using both the scalar and gradient fields, allowing it to better approximate low-density regions and unbalanced target weights—unlike SGMs, which relies solely on the score (i.e., gradient-log density) (Song & Ermon, ) and cannot recognize weighted modes in simple mixture data; see Theorem 4.2 and Sec. 6.1.

In Sec. 4, we also provide a steady-state analysis of Eq. (12) and an error analysis for Eq. (13), supporting the rationale behind our Bellman Diffusion Dynamics.

#### 3.4 Summary of Training and Sampling Algorithms

To summarize Bellman Diffusion as a deep generative model, we outline the training and sampling steps in Algorithm 1 and Algorithm 2. For training, we first sample real data $\mathbf{x}_{1},\mathbf{x}_{2}$ from dataset $\mathcal{X}$ (line 2) and slice vectors $\mathbf{v},\mathbf{w}$ from some predefined distributions $q(\mathbf{v}),q(\mathbf{w})$ (line 3)(^2^22Here, we follow the practice in Song et al. () by using a single slice vector to approximate the expectation over $q({\mathbf{v}})$ or $q({\mathbf{w}})$, trading variance for reduced computational cost.). Then, we estimate the loss functions $\widebar{\mathcal{L}}_{\mathrm{grad}}^{\mathrm{slice}}(\bm{\phi};\epsilon)$, $\widebar{\mathcal{L}}_{\mathrm{id}}^{\mathrm{slice}}(\bm{\varphi};\epsilon)$ using Monte Carlo sampling (lines 4-6). Finally, the model parameters $\bm{\phi}$ and $\bm{\varphi}$ are updated via gradient descent (lines 7-8).

For inference, we begin by sampling $\mathbf{x}(0)$ from an arbitrary distribution, such as standard normal (line 1). Then, after setting the number of steps $T$ and step size $\eta$, we iteratively update $\mathbf{x}(0)$ to $\mathbf{x}(\eta T)$ following Eq. (12) (lines 3-7).

Figure: Algorithm 1 Training

### 4 Main Theory

In this section, we provide theoretical support for Bellman Diffusion Dynamics, including steady-state and error analyses. We defer all proofs to Appendix C.

#### 4.1 Steady-State Analysis of Bellman Diffusion Dynamics

Let $p_{t}$ be the marginal density of Bellman Diffusion Dynamics given by Eq. (12), starting from any initial density $p_{0}$. The following theorem shows that, regardless of the initial distribution $p_{0}$, $p_{t}$ converges to the stationary distribution, which is exactly $p_{\text{target}}({\mathbf{x}})$, as $t\rightarrow\infty$, at an exponential rate.

###### Theorem 4.1 (Convergence to the Steady State) .

Let $p_{\mathrm{target}}$ be the target density satisfying Assumption C.1. Then, for any initial density $p_{0}$, we have the following KL and Wasserstein-2 bounds:

$$ $\displaystyle W_{2}^{2}\big{(}p_{t},p_{\mathrm{target}}\big{)}\lesssim\mathrm{ KL}\big{(}p_{t}\|p_{\mathrm{target}}\big{)}\lesssim e^{-2\alpha t}\mathrm{KL} \big{(}p_{0}\|p_{\mathrm{target}}\big{)}.$ $$

Here, $\alpha>0$ is some constant determined by $p_{\mathrm{target}}$, and $\lesssim$ hides multiplicative constants that depend only on $p_{\mathrm{target}}$.

This theorem implies that as $t\to\infty$, $p_{t}\to p_{\mathrm{target}}$ in both KL and Wasserstein-2 senses. Thus, it justifies that by using our sampling method, which involves solving the SDE in Eq. (12), we can ensure that samples will be obtained from the target distribution $p_{\mathrm{target}}$.

###### Theorem 4.1 (Convergence to the Steady State) .

#### 4.2 Error Analysis of Empirical Bellman Diffusion Dynamics

We let $p_{t;\bm{\phi},\bm{\varphi}}$ denote the marginal density from the empirical Bellman Diffusion Dynamics in Eq. (13), starting from any initial density $p_{0}$. The following theorem extends the result in Theorem 4.1 by providing an error analysis. It accounts for network approximation errors in $\mathbf{g}_{\bm{\phi}}(\mathbf{x})\approx\nabla p_{\mathrm{target}}(\mathbf{x})$ and $s_{\bm{\varphi}}(\mathbf{x})\approx p_{\mathrm{target}}(\mathbf{x})$, and gives an upper bound on the Wasserstein-2 discrepancy between $p_{t;\bm{\phi},\bm{\varphi}}$ and $p_{\mathrm{target}}$.

###### Theorem 4.2 (Error Analysis of Neural Network Approximations) .

Let $p_{\mathrm{target}}$ be the target distribution satisfying Assumptions C.1 and C.2. Suppose the dynamics in Eqs. (12) and (13) start from the same initial condition sampled from $p_{0}$. For any $\varepsilon>0$, if $T=\mathcal{O}(\log 1/\varepsilon^{2})$ and $\varepsilon_{\mathrm{est}}=\mathcal{O}\Big{(}\frac{\varepsilon}{\sqrt{T}e^{
\frac{1}{2}LT}}\Big{)}$, such that
$\left\lVert g_{\bm{\phi}}(\cdot)-\nabla p_{\mathrm{target}}(\cdot)\right\rVert
_{\infty}\leq\varepsilon_{\mathrm{est}}$ and $\left|s_{\bm{\varphi}}(\cdot)-p_{\mathrm{target}}(\cdot)\right|_{\infty}\leq
\varepsilon_{\mathrm{est}}$,
where $L>0$ is the Lipschitz constant associated with $p_{\mathrm{target}}$, then

$$ $\displaystyle W_{2}(p_{T;\bm{\phi},\bm{\varphi}},p_{\mathrm{target}})\leq\varepsilon.$ $$

From the above theorem, our dynamics can function as a standalone generative model, capable of learning the target distribution $p_{\mathrm{target}}$. Using advanced techniques such as Chen et al. (); De Bortoli (); Kim et al. (; ), a tighter bound between $p_{T;\bm{\phi},\bm{\varphi}}$ and $p_{\mathrm{target}}$ in $W_{2}$ or other divergences could be achieved. Moreover, discrete-time versions of both Theorems 4.1 and 4.2 can be derived with more advanced analysis. However, we defer this to future work, as the current focus is on establishing the core principles.

###### Theorem 4.2 (Error Analysis of Neural Network Approximations) .

### 5 Related Work

##### Deep generative modelings.

Bellman Diffusion stands as a new class of generative models. Sec. 2.1 revisits modern DGMs and re-examines whether their modeling satisfies the linear property. Additional details on related works concerning DGMs can be found in Appendix A.

##### Markov decision processes.

Limited by the linearity of the distributional Bellman equation, Previous works (Bellemare et al., ; Hessel et al., ; Dabney et al., ) in planning and distributional RL have relied on conventional generative models to represent state-level return distributions. For instance, the widely used C51 (Bellemare et al., ) is a histogram model, resulting in discrete approximation errors. In contrast, Bellman Diffusion is a new type of diffusion model that serves as an expressive distribution approximator without discretization errors. As shown in Sec. 6, our experiments demonstrate that Bellman Diffusion achieves significantly faster and much more stable convergence properties.

### 6 Experiments

Figure: Figure 1: Bellman Diffusion captures the uniform distribution supported on disjoint spans. The leftmost subfigure presents the training data histogram, while the next three show the estimated density, derivative functions, and samples generated by Bellman Diffusion.
Refer to caption: extracted/5896758/images/1.png

Figure: Figure 2: Bellman Diffusion learns the unbalanced Gaussian mixture, which is hard for score-based models. The subfigures, from left to right, display the training data, estimated scalar and gradient fields, and samples generated by our Bellman Diffusion.
Refer to caption: extracted/5896758/images/3.png

To verify the effectiveness of our method: Bellman Diffusion, we have conducted extensive experiments on multiple synthetic and real benchmarks across different tasks (e.g., generative modeling and MDPs). The aim of our experiments is to verify that Bellman Diffusion is both a capable generative model and an effective distributional RL model. We also place the experiment setup in Appendix E and other minor experiments in Appendix F.

#### 6.1 Synthetic Datasets

In this part, we aim to show that Bellman Diffusion can accurately estimate the scalar and gradient fields $\nabla p_{\mathrm{target}}({\mathbf{x}}),p_{\mathrm{target}}({\mathbf{x}})$ and the associated sampling dynamics can recover the data distribution in terms of the estimation models $\mathbf{g}_{\bm{\phi}}(\mathbf{x}),s_{\bm{\varphi}}(\mathbf{x})$. For visualization purpose, we will adopt low-dimensional synthetic data (i.e., $D=1$, or $2$) in the studies.

##### 1-dimensional uniform distribution.

As shown in the leftmost subfigure of Fig. 1, we apply Bellman Diffusion to a uniform distribution over three disjoint spans. A key challenge is approximating the discontinuous data distribution using continuous neural networks $g_{\bm{\phi}}({\mathbf{x}})$ and $s_{\bm{\varphi}}({\mathbf{x}})$.

Interestingly, the results in Fig. 1 indicate that the estimated field models $\mathbf{g}_{\bm{\phi}}(\mathbf{x})$ and $s_{\bm{\varphi}}(\mathbf{x})$ closely approximate the correct values on the support (e.g., $[-1.0,2.0]$) and behave reasonably in undefined regions. For instance, $g_{\bm{\phi}}({\mathbf{x}})$ resembles a negative sine curve on $[-4.5,-0.5]$, aligning with the definitions of one-sided derivative. Notably, our Bellman Diffusion Dynamics yield a generation distribution nearly matching the training samples, demonstrating effectiveness in learning discontinuous data distributions.

##### 2-Dimensional Mixture of Gaussian (MoG).

Bellman Diffusion effectively approximates the density and gradient fields for multimodal distributions, even with unbalanced weights, which are not recognizable for SGMs. We investigate this using a MoG distribution of three modes with differing weights ($0.45$, $0.45$, $0.1$), as shown in the leftmost subfigure of Fig. 2.

Figure: Figure 3: Bellman Diffusion learns unusually clustered data. The subfigures, from left to right, show the training data, estimated density field, gradient field, and generated samples.
Refer to caption: extracted/5896758/images/5.png

The right three subfigures of Fig. 2 present the experimental results, demonstrating accurate estimation of both the scalar and gradient fields for the target data distribution $p_{\mathrm{target}}({\mathbf{x}})$ and $\nabla p_{\mathrm{target}}({\mathbf{x}})$. The three clustering centers of the training data correspond to the three density peaks in the scalar field (leftmost subfigure) and the critical points in the gradient field (middle subfigure). Notably, Bellman Diffusion successfully recovers the unbalanced modes of the target distribution and accurately estimates the density and gradient fields, even in low-density regions—a challenge for SGMs (Song & Ermon, ) due to its score design.

##### 2-dimensional moon-shaped data.

To demonstrate the ability of Bellman Diffusion to learn distributions with disjoint supports, we test it on the two moon dataset, where samples cluster into two disjoint half-cycles, as shown in the leftmost subfigure of Fig. 3.

The right three parts of Fig. 3 shows that the estimated scalar and gradient fields $p_{\mathrm{target}}({\mathbf{x}})$ and $\nabla p_{\mathrm{target}}({\mathbf{x}})$ match the training samples, with correctly positioned density peaks (leftmost subfigure) and critical points (middle subfigure). Our diffusion sampling dynamics accurately recover the shape of the training data, even in low-density regions. Thus, we conclude that Bellman Diffusion is effective in learning from complex data.

#### 6.2 High-dimensional Data Generation

In this section, we follow the common practice (Song et al., ) to examine the scalability of our approach across multiple UCI tabular datasets (Asuncion et al., ), including Abalone, Telemonitoring, Mushroom, Parkinson’s, and Red Wine. We apply several preprocessing steps to these datasets, such as imputation and feature selection, resulting in data dimensions of $7$, $16$, $5$, $15$, and $10$, respectively. For evaluation metrics, we utilize the commonly used Wasserstein distance (Rüschendorf, ) and maximum mean discrepancy (MMD) (Dziugaite et al., ). The performance of a generative model is considered better when both metrics are lower.

The experimental results are presented in Table 1. We observe that, regardless of the dataset or metric, Bellman Diffusion performs competitively with DDPM (Ho et al., ; Song & Ermon, ), a diffusion model known for its scalability. For instance, although our model falls short of DSM by just $0.139$ points on the Telemonitoring dataset (as measured by the Wasserstein distance), it achieves an even better MMD score on the Abalone dataset.

We further demonstrate in Appendix F.1 that Bellman Diffusion is compatible with VAE (Kingma, ), allowing latent generative model training similar to latent diffusion models (Rombach et al., ) for higher-resolution image generation, highlighting the scalability of Bellman Diffusion. These results confirm that Bellman Diffusion is a scalable generative model.

**Table 1: Performances on multiple high-dimensional datasets, showing that Bellman Diffusion is a capable generative model in high dimensions.**
| Dataset | Denoising Diffusion Models | Our Model: Bellman Diffusion |  |  |
| --- | --- | --- | --- | --- |
| Wasserstein $\downarrow$ | MMD ($10^{-3}$) $\downarrow$ | Wasserstein $\downarrow$ | MMD ($10^{-3}$)$\downarrow$ |  |
| Abalone | $0.975$ | $5.72$ | $0.763$ | $5.15$ |
| Telemonitoring | $2.167$ | $10.15$ | $2.061$ | $9.76$ |
| Mushroom | $1.732$ | $4.29$ | $1.871$ | $5.12$ |
| Parkinsons | $0.862$ | $3.51$ | $0.995$ | $3.46$ |
| Red Wine | $1.151$ | $3.83$ | $1.096$ | $3.91$ |

#### 6.3 Applications to RL

Figure: Figure 4: The $2\times 3$ subfigures, arranged from left to right and top to bottom, show a full trajectory of Bellman Diffusion, interacting with a maze environment. Each subfigure consists of the state on the left, gradient field in the middle, and scalar field on the right.
Refer to caption: extracted/5896758/images/6-1.png

Figure: Figure 5: The left and right subfigures respectively show the initial and terminal states of Bellman Diffusion, interacting with an environment of balance control. Every subfigure is composed of the observation on the left, gradient field in the middle, and scalar field on the right.
Refer to caption: extracted/5896758/images/7-1.png

We demonstrate that Bellman Diffusion is effective for distributional RL, a classical task in MDPs. A method successful in RL can also handle simpler tasks (e.g., planning). Compared to previous baselines (e.g., C51 (Bellemare et al., )), our method learns a continuous return distribution at each state, avoiding errors from discrete approximations. We expect better convergence properties than the baseline, as our approach minimizes error propagation along state transitions.

Case studies.
We apply Bellman Diffusion to two OpenAI Gym environments (Brockman, ): Frozen Lake and Cart Pole. Concrete implementations are detailed in Appendix D. Frozen Lake is a maze where actions (e.g., moving up) may yield unexpected outcomes (e.g., moving left), while Cart Pole involves balancing a pole on a movable car.

Figure: Figure 6: The return dynamics on the Cart Pole, with the colored areas as the confidence interval and the maximum return as $500$.
Refer to caption: extracted/5896758/images/9.png

Results in Figs. 4 and 5 show that Bellman Diffusion accurately estimates state-level return distributions and their derivatives. For instance, as the agent approaches the goal in the maze, the expected return shifts from $0.5$ to $1$, reflecting that the agent receives no rewards until it reaches the goal.

Faster and more stable convergence. With the same model sizes and different random seeds, Bellman Diffusion and C51, a widely used distributional RL method, are both run on the environment of Cart Pole $10$ times. The results of return dynamics over training steps are shown in Fig. 6. We can see that, while both models have the potential to reach the maximum return in the end, Bellman Diffusion converges much faster than C51, with highly stable dynamics. These results confirm our previous conjecture about the superiority of Bellman Diffusion.

### 7 Conclusion

In this work, we highlight the limitations of modern DGMs, including emerging SGMs, in the context of MDPs and distributional RL. By identifying the critical necessity for linearity in these applications, we introduce Bellman Diffusion, a novel DGM that effectively maintains this linearity by directly modeling gradient and scalar fields. Through the development of divergence measures, we train neural network proxies for these fields, facilitating a new sampling dynamic: Bellman Diffusion Dynamics, via a SDE that ensures convergence to the target distribution. Our experiment results show that Bellman Diffusion not only achieves accurate field estimations and generates high-quality images, but also shows improved performance in classical distributional RL tasks. These findings suggest that Bellman Diffusion is a promising approach for advancing the integration of DGMs within the RL frameworks.

## Appendix

### Appendix A More Related Works on DGMs

Deep generative models (DGMs) have gained significant attention in recent years due to their ability to learn complex data distributions and generate high-fidelity samples. This literature review covers several prominent categories of DGMs, including Variational Autoencoders (VAEs), Generative Adversarial Networks (GANs), energy-based methods, flow-based methods, and diffusion models.

##### Variational Autoencoders (VAEs).

Variational Autoencoders (VAEs) are a class of generative models that leverage variational inference to approximate the posterior distribution of latent variables given the data. The VAE framework is based on the evidence lower bound (ELBO), which can be expressed as:

$$ $\mathcal{L}(\bm{\theta},\bm{\phi};{\mathbf{x}})=\mathbb{E}_{q_{\bm{\phi}}({ \mathbf{z}}|{\mathbf{x}})}[\log p_{\bm{\theta}}({\mathbf{x}}|{\mathbf{z}})]- \mathcal{D}_{\mathrm{KL}}(q_{\bm{\phi}}({\mathbf{z}}|{\mathbf{x}})||p({\mathbf {z}})),$ $$

where $q_{\bm{\phi}}({\mathbf{z}}|{\mathbf{x}})$ is the approximate posterior, $p_{\bm{\theta}}({\mathbf{x}}|{\mathbf{z}})$ is the likelihood, and $\mathcal{D}_{\mathrm{KL}}$ denotes the Kullback-Leibler divergence. VAEs have shown remarkable success in generating images and other complex data types Kingma ().

##### Generative Adversarial Networks (GANs).

Generative Adversarial Networks (GANs) consist of two neural networks, a generator $G$ and a discriminator $D$, that compete against each other. The generator aims to create realistic samples $G({\mathbf{z}})$ from random noise ${\mathbf{z}}$, while the discriminator attempts to distinguish between real samples ${\mathbf{x}}$ and generated samples $G({\mathbf{z}})$. The objective function for GANs can be formulated as:

$$ $\min_{G}\max_{D}\mathbb{E}_{{\mathbf{x}}\sim p_{\mathrm{target}}({\mathbf{x}}) }[\log D({\mathbf{x}})]+\mathbb{E}_{{\mathbf{z}}\sim p_{{\mathbf{z}}}({\mathbf {z}})}[\log(1-D(G({\mathbf{z}})))],$ $$

where $p_{\mathrm{data}}({\mathbf{x}})$ is the data distribution and $p_{{\mathbf{z}}}({\mathbf{z}})$ is the prior distribution on the noise. GANs have become popular for their ability to produce high-quality images and have been applied in various domains Goodfellow et al. ().

##### Energy-Based Models

Energy-based models (EBMs) define a probability distribution through an energy function $E({\mathbf{x}})$ that assigns lower energy to more probable data points. The probability of a data point is given by:

$$ $p({\mathbf{x}})=\frac{1}{Z}\exp(-E({\mathbf{x}})),$ $$

where $Z=\int\exp(-E({\mathbf{x}}))d{\mathbf{x}}$ is the partition function. Training EBMs typically involves minimizing the negative log-likelihood of the data (LeCun et al., ). They have been successfully applied in generative tasks, including image generation and modeling complex data distributions.

##### Flow-Based Methods

Flow-based methods, such as Normalizing Flows (NFs), learn a bijective mapping between a simple distribution ${\mathbf{z}}$ and a complex data distribution ${\mathbf{x}}$ through a series of invertible transformations. The probability density of the data can be expressed as:

$$ $p({\mathbf{x}})=p({\mathbf{z}})\left|\det\frac{\partial{\mathbf{f}}^{-1}}{ \partial{\mathbf{x}}}\right|,$ $$

where ${\mathbf{f}}$ is the invertible transformation from ${\mathbf{z}}$ to ${\mathbf{x}}$. Flow-based models allow for efficient exact likelihood estimation and have shown promise in generating high-quality samples (Rezende & Mohamed, ).

##### Diffusion Models

Diffusion models are a class of generative models that learn to generate data by reversing a gradual noising process. The generative process can be described using a stochastic differential equation (SDE):

$$ $\mathop{}\!\mathrm{d}{\mathbf{x}}_{t}={\mathbf{f}}({\mathbf{x}}_{t},t)\mathop{ }\!\mathrm{d}dt+g(t)\mathop{}\!\mathrm{d}d{\mathbf{w}}_{t},$ $$

where ${\mathbf{w}}_{t}$ is a Wiener process, and ${\mathbf{f}}({\mathbf{x}}_{t},t)$ and $g(t)$ are functions defining the drift and diffusion terms, respectively. The model learns to recover the data distribution from noise by training on the denoising score matching objective (Ho et al., ; Song et al., ). Diffusion models have recently gained attention for their impressive image synthesis capabilities.

### Appendix B Theoretical Results and Proofs for Sec. 3

#### B.1 Motivation of the Proposed Dynamics
in Eq. ( 12 )

##### 1 1 1 1 -dimensional Case.

Let us first consider the one-dimensional case:

$$ $\mathop{}\!\mathrm{d}x(t)=f(x(t))\mathop{}\!\mathrm{d}t+g(x(t))\mathop{}\! \mathrm{d}w(t).$ $$

Based on the Fokker–Planck equation Risken & Risken (), the probability distribution $p(x,t)$ of dynamics $x(t)$ satisfies

$$ $\displaystyle\frac{\partial p(x,t)}{\partial t}$ $\displaystyle=-\frac{\partial}{\partial x}\Big{(}f(x)p(x,t)\Big{)}+\frac{1}{2} \frac{\partial^{2}}{\partial x^{2}}\Big{(}g^{2}(x)p(x,t)\Big{)}$ $\displaystyle=\frac{\partial}{\partial x}\Big{(}-f(x)p(x,t)+\frac{1}{2}\frac{ \partial}{\partial x}\Big{(}g^{2}(x)p(x,t)\Big{)}\Big{)}.$ $$

Suppose that the density $p(x,t)$ converges as $t\rightarrow\infty$, then we have $\partial p(x,t)/\partial t\mid_{t\rightarrow\infty}=0$. As a result, the above equality indicates that

$$ $-f(x)p(x,\infty)+\frac{1}{2}\frac{\mathop{}\!\mathrm{d}}{\mathop{}\!\mathrm{d} x}\Big{(}g^{2}(x)p(x,\infty)\Big{)}=C.$ $$

Suppose the constant $C$ is 0 0, then we have

$$ $g^{2}(x)\frac{\mathop{}\!\mathrm{d}p(x,\infty)}{\mathop{}\!\mathrm{d}x}=\Big{( }2f(x)-\frac{\mathop{}\!\mathrm{d}g^{2}(x)}{\mathop{}\!\mathrm{d}x}\Big{)}p(x, \infty).$ $$

One way to make this equality hold is to have the below setup:

$$ $\left\{\begin{aligned} g^{2}(x)&=p(x,\infty)\\ f(x)&=\frac{1}{2}\Big{(}\frac{\mathop{}\!\mathrm{d}p(x,\infty)}{\mathop{}\! \mathrm{d}x}+\frac{\mathop{}\!\mathrm{d}g^{2}(x)}{\mathop{}\!\mathrm{d}x}\Big{ )}\Big{)}=\frac{\mathop{}\!\mathrm{d}p(x,\infty)}{\mathop{}\!\mathrm{d}x}.\end {aligned}\right.$ $$

Therefore, the following dynamics:

$$ $\mathop{}\!\mathrm{d}x(t)=\frac{\mathop{}\!\mathrm{d}p_{\mathrm{target}}(x(t)) }{\mathop{}\!\mathrm{d}x}\mathop{}\!\mathrm{d}t+\sqrt{p_{\mathrm{target}}(x(t) )}\mathop{}\!\mathrm{d}w(t),$ $$

will converge to distribution $p_{\mathrm{target}}(x)$ as $t\rightarrow\infty$, regardless of the initial distribution $p(x,0)$.

##### D 𝐷 D italic_D -dimensional Case.

For the general situation, the dynamics will be

$$ $\mathop{}\!\mathrm{d}\mathbf{x}(t)=\nabla_{\mathbf{x}(t)}p_{\mathrm{target}}({ \mathbf{x}}(t))\mathop{}\!\mathrm{d}t+\sqrt{p_{\mathrm{target}}({\mathbf{x}}(t ))}\mathop{}\!\mathrm{d}\bm{\omega}(t),$ $$

Let us check this expression. Firstly, the Fokker–Planck equation indicates that

$$ $\displaystyle\frac{\partial p(\mathbf{x},t)}{\partial t}$ $\displaystyle=-\nabla_{\mathbf{x}}\cdot\big{(}p(\mathbf{x},t)\nabla_{\mathbf{x }}p_{\mathrm{target}}({\mathbf{x}})\big{)}+\frac{1}{2}\Delta_{\mathbf{x}}\big{ (}p_{\mathrm{target}}({\mathbf{x}})p(\mathbf{x},t)\big{)}$ $\displaystyle=\nabla_{\mathbf{x}}\cdot\Big{(}-p(\mathbf{x},t)\nabla_{\mathbf{x }}p_{\mathrm{target}}({\mathbf{x}})+\frac{1}{2}\nabla_{\mathbf{x}}\big{(}p_{ \mathrm{target}}({\mathbf{x}})p(\mathbf{x},t)\big{)}\Big{)},$ $$

where $\Delta_{\mathbf{x}}=\nabla_{\mathbf{x}}\cdot\nabla_{\mathbf{x}}$ is the Laplace operator. With the Leibniz rule, we have

$$ $\displaystyle-p(\mathbf{x},t)\nabla_{\mathbf{x}}p_{\mathrm{target}}({\mathbf{x }})+\frac{1}{2}\nabla_{\mathbf{x}}\big{(}p_{\mathrm{target}}({\mathbf{x}})p( \mathbf{x},t)\big{)}$ $\displaystyle=$ $\displaystyle-p(\mathbf{x},t)\nabla_{\mathbf{x}}p_{\mathrm{target}}({\mathbf{x }})+\frac{1}{2}p(\mathbf{x},t)\nabla_{\mathbf{x}}p_{\mathrm{target}}({\mathbf{ x}})+\frac{1}{2}p_{\mathrm{target}}({\mathbf{x}})\nabla_{\mathbf{x}}p(\mathbf{ x},t)$ $\displaystyle=$ $\displaystyle\frac{1}{2}p_{\mathrm{target}}({\mathbf{x}})\nabla_{\mathbf{x}}p( \mathbf{x},t)-\frac{1}{2}p(\mathbf{x},t)\nabla_{\mathbf{x}}p_{\mathrm{target}} ({\mathbf{x}}).$ $$

Combining the above two equations, we get

$$ $\frac{\partial p(\mathbf{x},t)}{\partial t}=\frac{1}{2}\nabla_{\mathbf{x}} \cdot\Big{(}p_{\mathrm{target}}({\mathbf{x}})\nabla_{\mathbf{x}}p(\mathbf{x},t )-p(\mathbf{x},t)\nabla_{\mathbf{x}}p_{\mathrm{target}}({\mathbf{x}})\Big{)}.$ $$

By applying the Leibniz rule to divergence operators, we have

$$ $\left\{\begin{aligned} \nabla_{\mathbf{x}}\cdot\big{(}p_{\mathrm{target}}({ \mathbf{x}})\nabla_{\mathbf{x}}p(\mathbf{x},t)\big{)}&=\big{<}\nabla_{\mathbf{ x}}p_{\mathrm{target}}({\mathbf{x}}),\nabla_{\mathbf{x}}p(\mathbf{x},t)\big{>} +p_{\mathrm{target}}({\mathbf{x}})\Delta_{\mathbf{x}}p(\mathbf{x},t)\\ \nabla_{\mathbf{x}}\cdot\big{(}p(\mathbf{x},t)\nabla_{\mathbf{x}}p_{\mathrm{ target}}({\mathbf{x}})\big{)}&=\big{<}\nabla_{\mathbf{x}}p(\mathbf{x},t), \nabla_{\mathbf{x}}p_{\mathrm{target}}({\mathbf{x}})\big{>}+p(\mathbf{x},t) \Delta_{\mathbf{x}}p_{\mathrm{target}}({\mathbf{x}})\end{aligned}\right..$ $$

Therefore, the original partial differential equation (PDE) can be simplified as

$$ $\frac{\partial p(\mathbf{x},t)}{\partial t}=\frac{1}{2}\Big{(}p_{\mathrm{ target}}({\mathbf{x}})\Delta_{\mathbf{x}}p(\mathbf{x},t)-p(\mathbf{x},t)\Delta _{\mathbf{x}}p_{\mathrm{target}}({\mathbf{x}})\Big{)}.$ $$

Since the dynamics will converge, we set $p_{\mathrm{target}}({\mathbf{x}})=p(\mathbf{x},\infty)$. Then, we get

$$ $\frac{\partial p(\mathbf{x},t)}{\partial t}\Big{|}_{t\rightarrow\infty}=\frac{ 1}{2}\Big{(}p(\mathbf{x},\infty)\Delta_{\mathbf{x}}p(\mathbf{x},\infty)-p( \mathbf{x},\infty)\Delta_{\mathbf{x}}p(\mathbf{x},\infty)\Big{)}=0.$ $$

Therefore, the dynamics lead to sampling from a given distribution $p_{\mathrm{target}}({\mathbf{x}})$.

#### B.2 Validity of Field Divergences.

The first step is to check whether $\mathcal{D}_{\mathrm{grad}}$, $\mathcal{D}_{\mathrm{id}}$ are well defined divergence measures. To this end, we have the below conclusion.

###### Theorem B.1 (Well-defined Divergences) .

Suppose that $p(\cdot),q(\cdot)$ are probability densities that are second-order continuously differentiable (i.e., in $\mathcal{C}^{2}$) and that $p({\mathbf{x}})\neq 0$ for all ${\mathbf{x}}$. Then the divergence measure $\mathcal{D}_{\mathrm{grad}}({p(\cdot)},q(\cdot))$ defined by Eq. (5) and that $\mathcal{D}_{\mathrm{id}}({p(\cdot)},q(\cdot))$ formulated by Eq. (6) are both valid statistical divergence measures, satisfying the following three conditions:

- •
Non-negativity: $\mathcal{D}_{*}(p(\cdot),q(\cdot))$ is either zero or positive;
- •
Null condition: $\mathcal{D}_{*}(p(\cdot),q(\cdot))=0$ if and only if $p({\mathbf{x}})=q(\mathbf{x})$ for every point $\mathbf{x}$;
- •
Positive definiteness: $\mathcal{D}_{*}(p(\cdot),p(\cdot)+\delta p(\cdot))$ is a positive-definite quadratic form for any infinitesimal displacement $\delta p(\cdot)$ from $p(\cdot)$.

Here the subscript $*$ represents either $\mathrm{grad}$ or $\mathrm{id}$.

###### Proof.

Non-negativity condition obviously holds for both $D_{\mathrm{grad}}$ and $D_{\mathrm{id}}$, due to their definitions.

For the null condition,

$$ $D_{\mathrm{grad}}(p(\cdot),q(\cdot))=0\quad\text{implies}\quad p({\mathbf{x}}) \|\nabla p({\mathbf{x}})-\nabla q({\mathbf{x}})\|^{2}=0\text{ for all }{ \mathbf{x}}.$ $$

This implies $\nabla p({\mathbf{x}})=\nabla q({\mathbf{x}})$ for all $x$. Since the gradients are equal, $p({\mathbf{x}})$ and $q({\mathbf{x}})$ differ by at most a constant. For probability densities, this constant must be zero, so $p({\mathbf{x}})=q({\mathbf{x}})$. On the other hand, for $D_{\mathrm{id}}(p(\cdot),q(\cdot))=0$:

$$ $D_{\mathrm{id}}(p(\cdot),q(\cdot))=0\quad\text{implies}\quad(p({\mathbf{x}})-q ({\mathbf{x}}))^{2}=0\text{ for all }{\mathbf{x}}$ $$

This directly implies that $p({\mathbf{x}})=q({\mathbf{x}})$ for all ${\mathbf{x}}$. Hence, both $D_{\mathrm{grad}}$ and $D_{\mathrm{id}}$ satisfy the null condition: $D_{*}(p(\cdot),q(\cdot))=0$ if and only if $p({\mathbf{x}})=q({\mathbf{x}})$ for all ${\mathbf{x}}$.

At last, we prove that the two measurements satisfy the positive definiteness condition. For $D_{\mathrm{grad}}(p(\cdot),p(\cdot)+\delta p(\cdot))$:

$$ $D_{\mathrm{grad}}(p(\cdot),p(\cdot)+\delta p(\cdot))=\int p({\mathbf{x}})\| \nabla p({\mathbf{x}})-\nabla(p({\mathbf{x}})+\delta p({\mathbf{x}}))\|^{2}\, \mathop{}\!\mathrm{d}{\mathbf{x}}=\int p({\mathbf{x}})\|\nabla\delta p({ \mathbf{x}})\|^{2}\,\mathop{}\!\mathrm{d}{\mathbf{x}}.$ $$

This expression is quadratic in $\delta p({\mathbf{x}})$, and since norms are positive definite, $D_{\mathrm{grad}}$ is positive definite for any infinitesimal displacement $\delta p({\mathbf{x}})$. On the other hand, for $D_{\mathrm{id}}(p(\cdot),p(\cdot)+\delta p(\cdot))$:

$$ $D_{\mathrm{id}}(p(\cdot),p(\cdot)+\delta p(\cdot))=\int p({\mathbf{x}})(p({ \mathbf{x}})-(p({\mathbf{x}})+\delta p({\mathbf{x}})))^{2}\,\mathop{}\!\mathrm {d}{\mathbf{x}}=\int p({\mathbf{x}})(\delta p({\mathbf{x}}))^{2}\,\mathop{}\! \mathrm{d}{\mathbf{x}}.$ $$

Again, this is quadratic in $\delta p({\mathbf{x}})$, making $D_{\mathrm{id}}$ positive definite for any infinitesimal displacement $\delta p({\mathbf{x}})$. Thus, both $D_{\mathrm{grad}}$ and $D_{\mathrm{id}}$ are positive-definite quadratic forms for any infinitesimal displacement $\delta p({\mathbf{x}})$ from $p({\mathbf{x}})$. This concludes the proof.
∎

Since measures $\mathcal{D}_{\mathrm{grad}}$, $\mathcal{D}_{\mathrm{id}}$ are well defined, it is valid to derive the corresponding loss functions $\mathcal{L}_{\mathrm{grad}},\mathcal{L}_{\mathrm{id}}$ as formulated in Eq. (7).

###### Theorem B.1 (Well-defined Divergences) .

###### Proof.

#### B.3 Proof to Proposition 3.1 .

We prove for the case of

$$ $\mathcal{L}_{\mathrm{id}}(\bm{\varphi})=C_{\mathrm{id}}+\lim_{\epsilon \rightarrow 0}\mathbb{E}_{\mathbf{x}_{1},\mathbf{x}_{2}\sim{p_{\mathrm{target} }({\mathbf{x}})}}\Big{[}s_{\bm{\varphi}}(\mathbf{x}_{1})^{2}-2s_{\bm{\varphi}} (\mathbf{x}_{1})\mathcal{N}(\mathbf{x}_{2}-\mathbf{x}_{1};\mathbf{0},\epsilon \mathbf{I}_{D})\Big{]}.$ $$

Similar argument can be applied to the case for $\mathcal{L}_{\mathrm{grad}}(\bm{\phi})$.
Suppose that we have a divergence loss function:

$$ $\mathcal{L}_{\mathrm{id}}(\bm{\varphi})=\int p_{\mathrm{target}}(\mathbf{x}) \Big{(}p_{\mathrm{target}}(\mathbf{x})-s_{\bm{\varphi}}(\mathbf{x})\Big{)}^{2} \mathop{}\!\mathrm{d}\mathbf{x}.$ $$

Then, we can expand the term as

$$ $\mathcal{L}_{\mathrm{id}}(\bm{\varphi})=\int p_{\mathrm{target}}(\mathbf{x})^{ 3}\mathop{}\!\mathrm{d}\mathbf{x}-2\int p_{\mathrm{target}}(\mathbf{x})^{2}s_{ \bm{\varphi}}(\mathbf{x})\mathop{}\!\mathrm{d}\mathbf{x}+\int p_{\mathrm{ target}}(\mathbf{x})s_{\bm{\varphi}}(\mathbf{x})^{2}\mathop{}\!\mathrm{d} \mathbf{x}.$ $$

For the second integral, we apply the trick again:

$$ $\displaystyle\int p_{\mathrm{target}}(\mathbf{x})^{2}s_{\bm{\varphi}}(\mathbf{ x})\mathop{}\!\mathrm{d}\mathbf{x}=\int p_{\mathrm{target}}(\mathbf{x})p_{ \mathrm{target}}(\mathbf{y})s_{\bm{\varphi}}(\mathbf{x})\delta(\mathbf{y}- \mathbf{x})\mathop{}\!\mathrm{d}\mathbf{x}\mathop{}\!\mathrm{d}\mathbf{y}$ $\displaystyle=$ $\displaystyle\mathbb{E}_{\mathbf{x}\sim p_{\mathrm{target}}(\mathbf{x}), \mathbf{y}\sim p_{\mathrm{target}}(\mathbf{y})}\Big{[}s_{\bm{\varphi}}(\mathbf {x})\delta(\mathbf{y}-\mathbf{x})\Big{]}$ $\displaystyle=$ $\displaystyle\lim_{\epsilon\rightarrow 0}\mathbb{E}_{\mathbf{x}_{1},\mathbf{x} _{2}\sim p_{\mathrm{target}}(\mathbf{x})}\Big{[}s_{\bm{\varphi}}(\mathbf{x}_{1 })\mathcal{N}(\mathbf{x}_{2}-\mathbf{x}_{1};\mathbf{0},\epsilon\mathbf{I}_{D}) \Big{]}.$ $$

Here, we use that $\mathcal{N}(\cdot;0,\epsilon\mathbf{I}_{D})$ weakly converges to $\delta(\cdot)$ as $\epsilon\rightarrow 0$. Therefore, we simplify the loss function as

$$ $\displaystyle\mathcal{L}_{\mathrm{id}}(\bm{\varphi})$ $\displaystyle=C_{\mathrm{id}}+\mathbb{E}_{\mathbf{x}\sim p_{\mathrm{target}}( \mathbf{x})}\Big{[}s_{\bm{\varphi}}(\mathbf{x})^{2}\Big{]}-2\lim_{\epsilon \rightarrow 0}\mathbb{E}_{\mathbf{x}_{1},\mathbf{x}_{2}\sim p_{\mathrm{target} }(\mathbf{x})}\Big{[}s_{\bm{\varphi}}(\mathbf{x}_{1})\mathcal{N}(\mathbf{x}_{2 }-\mathbf{x}_{1};\mathbf{0},\epsilon\mathbf{I}_{D})\Big{]},$ $$

where $C_{\mathrm{id}}$ is a constant without learnable parameter $\bm{\varphi}$.

#### B.4 Proof to Proposition 3.2

###### Proof.

We recall the sliced version of $\mathcal{L}_{\mathrm{grad}}$ as:

$$ $\mathcal{L}_{\mathrm{grad}}^{\mathrm{slice}}(\bm{\phi})=\mathbb{E}_{\mathbf{v} \sim q(\mathbf{v}),\mathbf{x}\sim{p_{\mathrm{target}}({\mathbf{x}})}}\Big{[} \big{(}\mathbf{v}^{T}\nabla{p_{\mathrm{target}}({\mathbf{x}})}-\mathbf{v}^{T} \mathbf{g}_{\bm{\phi}}(\mathbf{x})\big{)}^{2}\Big{]}.$ $$

By expanding the quadratic term $(\cdot)^{2}$ inside the recursive expectations, we have

$$ $\displaystyle\mathcal{L}_{\mathrm{grad}}^{\mathrm{slice}}=\mathbb{E}_{\mathbf{ v}}\Big{[}\int p_{\mathrm{target}}(\mathbf{x})(\mathbf{v}^{\top}\nabla_{ \mathbf{x}}\mathbf{g}_{\bm{\phi}}(\mathbf{x}))^{2}\mathop{}\!\mathrm{d}{ \mathbf{x}}-2\int p_{\mathrm{target}}(\mathbf{x})(\mathbf{v}^{\top}\nabla_{ \mathbf{x}}p_{\mathrm{target}}(\mathbf{x}))(\mathbf{v}^{\top}\nabla_{\mathbf{x }}\mathbf{g}_{\bm{\phi}}(\mathbf{x}))\mathop{}\!\mathrm{d}\mathbf{x}+C_{ \mathrm{grad}}^{\prime}\Big{]},$ $$

where, $C_{\mathrm{grad}}^{\prime}:=\int p_{\mathrm{target}}(\mathbf{x})(\mathbf{v}^{
\top}\nabla_{\mathbf{x}}p_{\mathrm{target}}(\mathbf{x}))^{2}\mathop{}\!\mathrm
{d}\mathbf{x}$ is a constant independent of trainable parameter.
We will further simplify the second term came from the cross product as:

$$ $\displaystyle\int p_{\mathrm{target}}(\mathbf{x})(\mathbf{v}^{\top}\nabla_{ \mathbf{x}}p_{\mathrm{target}}(\mathbf{x}))(\mathbf{v}^{\top}\nabla_{\mathbf{x }}\mathbf{g}_{\bm{\phi}}(\mathbf{x}))\mathop{}\!\mathrm{d}{\mathbf{x}}$ $\displaystyle=$ $\displaystyle\frac{1}{2}\int(\mathbf{v}^{\top}\nabla_{\mathbf{x}}p_{\mathrm{ target}}(\mathbf{x})^{2})(\mathbf{v}^{\top}\nabla_{\mathbf{x}}\mathbf{g}_{\bm{ \phi}}(\mathbf{x}))\mathop{}\!\mathrm{d}{\mathbf{x}}$ $\displaystyle=$ $\displaystyle\frac{1}{2}\int\Big{(}\nabla_{\mathbf{x}}p_{\mathrm{target}}( \mathbf{x})^{2}\Big{)}^{\top}\Big{(}(\mathbf{v}^{\top}\nabla_{\mathbf{x}} \mathbf{g}_{\bm{\phi}}(\mathbf{x}))\mathbf{v}\Big{)}\mathop{}\!\mathrm{d}{ \mathbf{x}}.$ $$

Note that we can replace the gradient field $\nabla_{\mathbf{x}}\mathbf{g}_{\bm{\phi}}(\mathbf{x})$ with neural network $\mathbf{g}_{\bm{\phi}}(\mathbf{x})$. By applying the integration by parts, this equality can be expanded as

$$ $\frac{1}{2}\int\nabla_{\mathbf{x}}\cdot\Big{(}p_{\mathrm{target}}(\mathbf{x})^ {2}(\mathbf{v}^{\top}\mathbf{g}_{\bm{\phi}}(\mathbf{x}))\mathbf{v}\Big{)} \mathop{}\!\mathrm{d}{\mathbf{x}}-\frac{1}{2}\int p_{\mathrm{target}}(\mathbf{ x})^{2}\nabla_{\mathbf{x}}\cdot\Big{(}(\mathbf{v}^{\top}\mathbf{g}_{\bm{\phi}} (\mathbf{x}))\mathbf{v}\Big{)}\mathop{}\!\mathrm{d}{\mathbf{x}}.$ $$

Let us first handle the first term in the above equation. Applying Gauss’s divergence theorem to a ball $\mathbb{B}(R)$ centered at the origin with radius $R>0$, we get

$$ $\int_{\mathbb{B}(R)}\nabla_{\mathbf{x}}\cdot\Big{(}p_{\mathrm{target}}(\mathbf {x})^{2}(\mathbf{v}^{\top}\mathbf{g}_{\bm{\phi}}(\mathbf{x}))\mathbf{v}\Big{)} \mathop{}\!\mathrm{d}{\mathbf{x}}=\int_{\partial\mathbb{B}(R)}\mathbf{n}( \mathbf{x})^{\top}\Big{(}p_{\mathrm{target}}(\mathbf{x})^{2}(\mathbf{v}^{\top} \mathbf{g}_{\bm{\phi}}(\mathbf{x}))\mathbf{v}\Big{)}.$ (14) $$

where $\mathbf{n}(\mathbf{x})$ is the unit norm vector to the region boundary $\partial\mathbb{B}(R)$. Suppose that $p_{\mathrm{target}}(\mathbf{x})$ decays sufficiently fast as $\left\lVert\mathbf{x}\right\rVert_{2}\rightarrow\infty$,
for instance, $\lim_{\|\mathbf{x}\|_{2}\rightarrow\infty}p_{\mathrm{target}}(\mathbf{x})/
\left\lVert{\mathbf{x}}\right\rVert_{2}^{D}=0$ (see Assumption C.1 (iii)), then this term vanishes as $R\rightarrow\infty$.

For the second term in the expansion, we have

$$ $\displaystyle\nabla_{\mathbf{x}}\cdot\Big{(}(\mathbf{v}^{\top}\mathbf{g}_{\bm{ \phi}}(\mathbf{x}))\mathbf{v}\Big{)}=\sum_{1\leq i\leq D}\frac{\partial(( \mathbf{v}^{\top}\mathbf{g}_{\bm{\phi}}(\mathbf{x}))v_{i})}{\partial x_{i}}= \sum_{1\leq i\leq D}\sum_{1\leq j\leq D}\frac{v_{i}v_{j}{\mathbf{g}}_{\bm{\phi },j}(\mathbf{x})}{\partial x_{i}}=\mathbf{v}^{\top}\nabla_{\mathbf{x}}{\mathbf {g}}_{\bm{\phi}}(\mathbf{x})\mathbf{v}.$ $$

Here, we write ${\mathbf{v}}=(v_{i})_{1\leq i\leq D}$ and ${\mathbf{x}}=(x_{i})_{1\leq i\leq D}$.
Collecting the above derivations, we have

$$ $\int p_{\mathrm{target}}(\mathbf{x})(\mathbf{v}^{\top}\nabla_{\mathbf{x}}p_{ \mathrm{target}}(\mathbf{x}))(\mathbf{v}^{\top}\nabla_{\mathbf{x}}\mathbf{g}_{ \bm{\phi}}(\mathbf{x}))\mathop{}\!\mathrm{d}{\mathbf{x}}=-\frac{1}{2}\int p_{ \mathrm{target}}(\mathbf{x})^{2}(\mathbf{v}^{\top}\nabla_{\mathbf{x}}{\mathbf{ g}}_{\bm{\phi}}(\mathbf{x})\mathbf{v})\mathop{}\!\mathrm{d}{\mathbf{x}}.$ (15) $$

Therefore, the loss function can be converted into

$$ $\mathcal{L}_{\mathrm{grad}}^{\mathrm{slice}}=\mathbb{E}_{\mathbf{v}}\Big{[} \int p_{\mathrm{target}}(\mathbf{x})(\mathbf{v}^{\top}\nabla_{\mathbf{x}} \mathbf{g}_{\bm{\phi}}(\mathbf{x}))^{2}\mathop{}\!\mathrm{d}{\mathbf{x}}+\int p _{\mathrm{target}}(\mathbf{x})^{2}(\mathbf{v}^{\top}\nabla_{\mathbf{x}}g_{\bm{ \theta}}(\mathbf{x})\mathbf{v})\mathop{}\!\mathrm{d}{\mathbf{x}}\Big{]}+C_{ \mathrm{grad}}^{\prime}.$ (16) $$

We apply the same trick from the proof of Proposition 3.1—using Dirac expansion—to enable Monte Carlo estimation for the second inner term:

$$ $\displaystyle\int p_{\mathrm{target}}(\mathbf{x})^{2}\Big{(}\mathbf{v}^{\top} \nabla_{\mathbf{x}}g_{\bm{\theta}}(\mathbf{x})\mathbf{v}\Big{)}\mathop{}\! \mathrm{d}{\mathbf{x}}$ $\displaystyle=$ $\displaystyle\int p_{\mathrm{target}}(\mathbf{x})\Big{(}\int p_{\mathrm{target }}(\mathbf{y})\delta(\mathbf{y}-\mathbf{x})d\mathbf{y}\Big{)}\Big{(}\mathbf{v} ^{\top}\nabla_{\mathbf{x}}g_{\bm{\theta}}(\mathbf{x})\mathbf{v}\Big{)}\mathop{ }\!\mathrm{d}{\mathbf{x}}$ $\displaystyle=$ $\displaystyle\int p_{\mathrm{target}}(\mathbf{x}_{1})p_{\mathrm{target}}( \mathbf{x}_{2})(\mathbf{v}^{\top}\nabla_{\mathbf{x}_{1}}g_{\bm{\theta}}( \mathbf{x}_{1})\mathbf{v})\delta(\mathbf{x}_{2}-\mathbf{x}_{1})\mathop{}\! \mathrm{d}{\mathbf{x}}_{1}\mathop{}\!\mathrm{d}{\mathbf{x}}_{2}$ $\displaystyle=$ $\displaystyle\mathbb{E}_{\mathbf{x}_{1},\mathbf{x}_{2}\sim p_{\mathrm{target}} (\mathbf{x})}\Big{[}(\mathbf{v}^{\top}\nabla_{\mathbf{x}_{1}}g_{\bm{\theta}}( \mathbf{x}_{1})\mathbf{v})\delta(\mathbf{x}_{2}-\mathbf{x}_{1})\Big{]}.$ $$

Combining the above two identities, we have

$$ $\mathcal{L}_{\mathrm{grad}}^{\mathrm{slice}}=\mathbb{E}_{\mathbf{v}\sim p_{ \mathrm{slice}}(\mathbf{v}),\mathbf{x}_{1},\mathbf{x}_{2}\sim p_{\mathrm{ target}}(\mathbf{x})}\Big{[}(\mathbf{v}^{\top}g_{\bm{\theta}}(\mathbf{x}_{1})) ^{2}+(\mathbf{v}^{\top}\nabla_{\mathbf{x}_{1}}g_{\bm{\theta}}(\mathbf{x}_{1}) \mathbf{v})\delta(\mathbf{x}_{2}-\mathbf{x}_{1})\Big{]}+C_{\mathrm{grad}}^{ \prime},$ (17) $$

which completes the proof.
∎

###### Proof.

### Appendix C Proofs for Sec. 4

#### C.1 Prerequisites for Theoretical Analysis.

We introduce some notations and terminologies. We recall the definition of *KL divergence* between $p_{\mathrm{target}}$ and density $p$ as

$$ $\mathrm{KL}\big{(}p\|p_{\mathrm{target}}\big{)}:=\int_{\mathbb{R}^{D}}p({ \mathbf{x}})\log\frac{p({\mathbf{x}})}{p_{\mathrm{target}}({\mathbf{x}})} \mathop{}\!\mathrm{d}{\mathbf{x}}.$ $$

*Fisher divergence* between $p_{\mathrm{target}}$ and $p$ is defined as:

$$ $J_{p_{\mathrm{target}}}(p):=\int_{\mathbb{R}^{D}}p({\mathbf{x}})\left\lVert \nabla_{x}\log\frac{p({\mathbf{x}})}{p_{\mathrm{target}}({\mathbf{x}})}\right \rVert^{2}\mathop{}\!\mathrm{d}{\mathbf{x}}.$ $$

*Wasserstein-2 distance* ($W_{2}$) between $p_{\mathrm{target}}$ and $p$ is defined as:

$$ $W_{2}^{2}(p,p_{\mathrm{target}}):=\inf_{\gamma\sim\Gamma(\mu,\nu)}\mathbb{E}_{ ({\mathbf{x}},{\mathbf{y}})\sim\gamma}\left\lVert{\mathbf{x}}-{\mathbf{y}} \right\rVert_{2}^{2},$ $$

where $\Gamma(\mu,\nu)$ is the set of all couplings of $(\mu,\nu)$.

The following summarizes the two assumptions for our main theorems in Sec. 4.

###### Assumption C.1 .

Assume the target density $p_{\mathrm{target}}$ satisfies the following conditions:

- (i)
$p_{\mathrm{target}}(\cdot)\in\mathcal{C}^{2}$. That is, it is second-order continuously differentiable;
- (ii)
*Log-Sobolev inequality*: there is a constant $\alpha>0$ so that the following inequality holds for all continuously differentiable density $p$:
$\mathrm{KL}\big{(}p\|p_{\mathrm{target}}\big{)}\leq\frac{1}{2\alpha}J_{p_{
\mathrm{target}}}(p).$
(18)
- (iii)
$p_{\mathrm{target}}$ is either compactly supported with $M:=\left\lVert p_{\mathrm{target}}\right\rVert_{L^{\infty}}<\infty$, or it decays sufficiently fast as $\left\lVert\mathbf{x}\right\rVert_{2}\rightarrow\infty$:
$\displaystyle\lim_{\|\mathbf{x}\|_{2}\rightarrow\infty}\frac{p_{\mathrm{target
}}(\mathbf{x})}{\left\lVert{\mathbf{x}}\right\rVert_{2}^{D}}=0.$

###### Assumption C.2 .

Assume the target density $p_{\mathrm{target}}$ satisfies the following additional conditions:

- (i)
There is a $L>0$ so that for all ${\mathbf{x}},{\mathbf{y}}$
$\displaystyle\left\lVert p_{\mathrm{target}}({\mathbf{x}})-p_{\mathrm{target}}
({\mathbf{y}})\right\rVert_{2}^{2}\leq L\left\lVert{\mathbf{x}}-{\mathbf{y}}
\right\rVert_{2}^{2}\quad\text{and}\quad\left\lVert\nabla p_{\mathrm{target}}(
{\mathbf{x}})-\nabla p_{\mathrm{target}}({\mathbf{y}})\right\rVert_{2}^{2}\leq
L
\left\lVert{\mathbf{x}}-{\mathbf{y}}\right\rVert_{2}^{2}.$

###### Assumption C.1 .

###### Assumption C.2 .

#### C.2 Proofs of Theorem 4.1

###### Proof.

Recall our dynamics is

$$ $\displaystyle\mathop{}\!\mathrm{d}{\mathbf{x}}(t)=\nabla p_{\mathrm{target}}({ \mathbf{x}}(t))\mathop{}\!\mathrm{d}t+\sqrt{p_{\mathrm{target}}({\mathbf{x}}(t ))}\mathop{}\!\mathrm{d}{\mathbf{w}}(t).$ $$

The Fokker-Planck equation of our dynamics with density $p_{t}=p_{t}({\mathbf{x}}):=p({\mathbf{x}},t)$ is

$$ $\partial_{t}p({\mathbf{x}},t)=\frac{1}{2}\nabla\cdot\Big{(}p_{\mathrm{target}} ({\mathbf{x}})p({\mathbf{x}},t)\nabla\log\frac{p({\mathbf{x}},t)}{p_{\mathrm{ target}}({\mathbf{x}})}\Big{)}.$ (19) $$

This is due to the following derivation Risken & Risken (), where we demonstrated for the $D=1$ case. The probability distribution $p({\mathbf{x}},t)$ of dynamics ${\mathbf{x}}(t)$ at point ${\mathbf{x}}$ and time $t$ with $f(x)=\partial_{x}p_{\mathrm{target}}(x)$ and $g^{2}(x)=p_{\mathrm{target}}(x)$ is governed by:

$$ $\displaystyle\frac{\partial p(x,t)}{\partial t}$ $\displaystyle=-\frac{\partial}{\partial x}\Big{(}f(x)p(x,t)\Big{)}+\frac{1}{2} \frac{\partial^{2}}{\partial x^{2}}\Big{(}g^{2}(x)p(x,t)\Big{)}$ $\displaystyle=\frac{\partial}{\partial x}\Big{(}-\partial_{x}p_{\mathrm{target }}(x)p(x,t)+\frac{1}{2}\frac{\partial}{\partial x}\Big{(}p_{\mathrm{target}}(x )p(x,t)\Big{)}\Big{)}$ $\displaystyle=\frac{1}{2}\frac{\partial}{\partial x}\Big{(}p_{\mathrm{target}} (x)\partial_{x}p(x,t)-\partial_{x}p_{\mathrm{target}}(x)p(x,t)\Big{)}$ $\displaystyle=\frac{1}{2}\frac{\partial}{\partial x}\Big{(}p_{\mathrm{target}} (x)p(x,t)\partial_{x}\log\frac{p(x,t)}{p_{\mathrm{target}}(x)}\Big{)}.$ $$

Here, in the last equality we use the identity:

$$ $\displaystyle\partial_{\mathbf{x}}\log\frac{p(x,t)}{p_{\mathrm{target}}(x)}= \frac{p_{\mathrm{target}}(x)}{p(x,t)}\partial_{\mathbf{x}}\Big{(}\frac{p(x,t)} {p_{\mathrm{target}}(x)}\Big{)}^{2}=\frac{p_{\mathrm{target}}(x)\partial_{x}p( x,t)-\partial_{x}p_{\mathrm{target}}(x)p(x,t)}{p_{\mathrm{target}}(x)p(x,t)}.$ $$

For a general $D$, the same computation can be carried out to derive Eq. (19).

We now prove the KL bound of convergence using a similar argument motivated by Vempala & Wibisono ().

$$ $\displaystyle\frac{\mathop{}\!\mathrm{d}}{\mathop{}\!\mathrm{d}t}\mathrm{KL} \big{(}p_{t}\|p_{\mathrm{target}}\big{)}$ $\displaystyle=$ $\displaystyle~{}\frac{\mathop{}\!\mathrm{d}}{\mathop{}\!\mathrm{d}t}\int_{ \mathbb{R}^{D}}p_{t}\log\frac{p_{t}}{p_{\mathrm{target}}}\mathop{}\!\mathrm{d} {\mathbf{x}}$ $\displaystyle\stackrel{{\scriptstyle{(a)}}}{{=}}$ $\displaystyle\int_{\mathbb{R}^{D}}\frac{\partial}{\partial t}p_{t}\log\frac{p_ {t}}{p_{\mathrm{target}}}\mathop{}\!\mathrm{d}{\mathbf{x}}+\int_{\mathbb{R}^{D }}p_{t}\frac{\partial}{\partial t}\log\frac{p_{t}}{p_{\mathrm{target}}}\mathop {}\!\mathrm{d}{\mathbf{x}}$ $\displaystyle\stackrel{{\scriptstyle{(b)}}}{{=}}$ $\displaystyle\int_{\mathbb{R}^{D}}\frac{\partial}{\partial t}p_{t}\log\frac{p_ {t}}{p_{\mathrm{target}}}\mathop{}\!\mathrm{d}{\mathbf{x}}$ $\displaystyle\stackrel{{\scriptstyle{(c)}}}{{=}}$ $\displaystyle\frac{1}{2}\int_{\mathbb{R}^{D}}\Big{[}\nabla_{x}\cdot\big{(}p_{ \mathrm{target}}p_{t}\nabla_{x}\log\frac{p_{t}}{p_{\mathrm{target}}}\big{)} \Big{]}\log\frac{p_{t}}{p_{\mathrm{target}}}\mathop{}\!\mathrm{d}{\mathbf{x}}$ $\displaystyle\stackrel{{\scriptstyle{(d)}}}{{=}}$ $\displaystyle-\int_{\mathbb{R}^{D}}p_{\mathrm{target}}p_{t}\left\lVert\nabla_{ x}\log\frac{p_{t}}{p_{\mathrm{target}}}\right\rVert^{2}\mathop{}\!\mathrm{d}{ \mathbf{x}}$ $\displaystyle\stackrel{{\scriptstyle{(e)}}}{{\leq}}$ $\displaystyle-M\int_{\mathbb{R}^{D}}p_{t}\left\lVert\nabla_{x}\log\frac{p_{t}} {p_{\mathrm{target}}}\right\rVert^{2}\mathop{}\!\mathrm{d}{\mathbf{x}}$ $\displaystyle=$ $\displaystyle-MJ_{p_{\mathrm{target}}}(p_{t})$ $\displaystyle\leq$ $\displaystyle-\frac{M}{2\alpha}\mathrm{KL}\big{(}p_{t}\|p_{\mathrm{target}} \big{)}.$ $$

Here, (a) follows from the chain rule; (b) uses the identity $\int p_{t}\frac{\partial}{\partial t}\log\frac{p_{t}}{p_{\mathrm{target}}}
\mathop{}\!\mathrm{d}{\mathbf{x}}=\int\frac{\partial}{\partial t}p_{t}\mathop{
}\!\mathrm{d}{\mathbf{x}}=\frac{\mathop{}\!\mathrm{d}}{\mathop{}\!\mathrm{d}t}
\int p_{t}\mathop{}\!\mathrm{d}{\mathbf{x}}=0$; (c) follows from the Fokker-Planck Eq. (19); (d) is due to integration by parts and Assumption C.1 (iii); and (e) comes from Assumption C.1 (iii).

Thus, applying Grönwall’s inequality, we can get

$$ $\displaystyle\mathrm{KL}\big{(}p_{t}\|p_{\mathrm{target}}\big{)}\lesssim e^{-2 \alpha t}\mathrm{KL}\big{(}p_{0}\|p_{\mathrm{target}}\big{)}.$ $$

Since $p_{\mathrm{target}}$ satisfies the LSI, it also satisfies the Talagrand’s inequality Otto & Villani ():

$$ $\displaystyle\frac{\alpha}{2}W_{2}^{2}\big{(}p_{t},p_{\mathrm{target}}\big{)} \leq\mathrm{KL}\big{(}p_{t}\|p_{\mathrm{target}}\big{)}.$ $$

Therefore, we have

$$ $\displaystyle W_{2}^{2}\big{(}p_{t},p_{\mathrm{target}}\big{)}\leq\frac{2}{ \alpha}\mathrm{KL}\big{(}p_{t}\|p_{\mathrm{target}}\big{)}\lesssim\frac{2}{ \alpha}e^{-2\alpha t}\mathrm{KL}\big{(}p_{0}\|p_{\mathrm{target}}\big{)}.$ $$

This completes the proof. We notice that “Talagrand’s inequality implies concentration of measure of Gaussian type” allowing us to remove the compact support assumption on $p_{\mathrm{target}}$ while maintaining the validity of the theorem.
∎

###### Proof.

#### C.3 Proofs of Theorem 4.2

###### Proof.

In the proof we will extensively using a simple form of Cauchy-Schwarz (CS) inequality:

$$ $\displaystyle(u_{1}+u_{2}+\cdots+u_{n})^{2}\leq n(u_{1}^{2}+u_{2}^{2}+\cdots+u _{n}^{2}),$ $$

for $u_{i}\in\mathbb{R}$, $i=1,\cdots,n$. We aim at obtaining the following bound:

$$ $\displaystyle W_{2}^{2}(p_{T;\bm{\phi},\bm{\varphi}},p_{\mathrm{target}}) \lesssim\varepsilon_{\mathrm{est}}^{2}Te^{LT}+\frac{2}{\alpha}e^{-\alpha T} \mathrm{KL}\big{(}p_{0}\|p_{\mathrm{target}}\big{)}.$ (20) $$

To achieve it, we compare the random vector processes $\{\mathbf{x}(t)\}_{t\in[0,T]}$ and $\{\widehat{{\mathbf{x}}}(t)\}_{t\in[0,T]}$, governed by the following dynamics:

$$ $\displaystyle\mathop{}\!\mathrm{d}\mathbf{x}(t)$ $\displaystyle=\nabla p_{\mathrm{target}}(\mathbf{x}(t))\mathop{}\!\mathrm{d}t+ \sqrt{p_{\mathrm{target}}(\mathbf{x}(t))}\mathop{}\!\mathrm{d}\bm{\omega}(t)$ $\displaystyle\mathop{}\!\mathrm{d}\widehat{{\mathbf{x}}}(t)$ $\displaystyle=g_{\bm{\phi}}(\widehat{{\mathbf{x}}}(t))\mathop{}\!\mathrm{d}t+ \sqrt{s_{\bm{\varphi}}(\widehat{{\mathbf{x}}}(t))}\mathop{}\!\mathrm{d} \widehat{{\mathbf{w}}}(t).$ $$

Their strong solutions in the Itô sense are:

$$ $\displaystyle\mathbf{x}(t)$ $\displaystyle=\mathbf{x}(0)+\int_{0}^{T}\nabla p_{\mathrm{target}}(\mathbf{x}( t))\mathop{}\!\mathrm{d}t+\int_{0}^{T}\sqrt{p_{\mathrm{target}}(\mathbf{x}(t)) }\mathop{}\!\mathrm{d}\bm{\omega}(t)$ $\displaystyle\widehat{{\mathbf{x}}}(t)$ $\displaystyle=\widehat{{\mathbf{x}}}(0)+\int_{0}^{T}g_{\bm{\phi}}(\widehat{{ \mathbf{x}}}(t))\mathop{}\!\mathrm{d}t+\int_{0}^{T}\sqrt{s_{\bm{\varphi}}( \widehat{{\mathbf{x}}}(t))}\mathop{}\!\mathrm{d}\widehat{{\mathbf{w}}}(t).$ $$

Set random vectors ${\mathbf{a}}(t):=\nabla p_{\mathrm{target}}(\mathbf{x}(t))-g_{\bm{\phi}}(
\widehat{{\mathbf{x}}}(t))$ and ${\mathbf{b}}(t):=\sqrt{p_{\mathrm{target}}(\mathbf{x}(t))}-\sqrt{s_{\bm{
\varphi}}(\widehat{{\mathbf{x}}}(t))}$, we then have

$$ $\displaystyle\mathbb{E}\big{[}\left\lVert\mathbf{x}(T)-\widehat{{\mathbf{x}}}( T)\right\rVert_{2}^{2}\big{]}$ $\displaystyle\leq\mathbb{E}\Bigg{[}\Big{(}\mathbf{x}(0)-\widehat{{\mathbf{x}}} (0)+\int_{0}^{T}{\mathbf{a}}(t)\mathop{}\!\mathrm{d}t+\int_{0}^{T}{\mathbf{b}} (t)\mathop{}\!\mathrm{d}\bm{\omega}(t)\Big{)}^{2}\Bigg{]}$ $\displaystyle\leq 3\mathbb{E}\Big{[}\left\lVert\mathbf{x}(0)-\widehat{{\mathbf {x}}}(0)\right\rVert_{2}^{2}\Big{]}+3\mathbb{E}\Big{[}\big{(}\int_{0}^{T}{ \mathbf{a}}(t)\mathop{}\!\mathrm{d}t\big{)}^{2}\Big{]}+3\mathbb{E}\Big{[}\big{ (}\int_{0}^{T}{\mathbf{b}}(t)\mathop{}\!\mathrm{d}\bm{\omega}(t)\big{)}^{2} \Big{]}$ $\displaystyle\lesssim\mathbb{E}\Big{[}\left\lVert\mathbf{x}(0)-\widehat{{ \mathbf{x}}}(0)\right\rVert_{2}^{2}\Big{]}+T\mathbb{E}\Big{[}\int_{0}^{T}{ \mathbf{a}}^{2}(t)\mathop{}\!\mathrm{d}t\Big{]}+\mathbb{E}\Big{[}\int_{0}^{T}{ \mathbf{b}}^{2}(t)\mathop{}\!\mathrm{d}t\Big{]}$ $\displaystyle\lesssim\mathbb{E}\Big{[}\left\lVert\mathbf{x}(0)-\widehat{{ \mathbf{x}}}(0)\right\rVert_{2}^{2}\Big{]}+T\mathbb{E}\Big{[}\int_{0}^{T}\left \lVert\nabla p_{\mathrm{target}}(\mathbf{x}(t))-\nabla p_{\mathrm{target}}( \widehat{{\mathbf{x}}}(t))\right\rVert_{2}^{2}\mathop{}\!\mathrm{d}t\Big{]}$ $\displaystyle~{}~{}+T\mathbb{E}\Big{[}\int_{0}^{T}\left\lVert\nabla p_{\mathrm {target}}(\widehat{{\mathbf{x}}}(t))-g_{\bm{\phi}}(\widehat{{\mathbf{x}}}(t)) \right\rVert_{2}^{2}\mathop{}\!\mathrm{d}t\Big{]}$ $\displaystyle~{}~{}+\mathbb{E}\Big{[}\int_{0}^{T}\left|p_{\mathrm{target}}( \mathbf{x}(t)-p_{\mathrm{target}}(\widehat{{\mathbf{x}}}(t))\right|\mathop{}\! \mathrm{d}t\Big{]}+\mathbb{E}\Big{[}\int_{0}^{T}\left|p_{\mathrm{target}}( \widehat{{\mathbf{x}}}(t))-s_{\bm{\varphi}}(\widehat{{\mathbf{x}}}(t))\right| \mathop{}\!\mathrm{d}t\Big{]}$ $\displaystyle\lesssim\mathbb{E}\Big{[}\left\lVert\mathbf{x}(0)-\widehat{{ \mathbf{x}}}(0)\right\rVert_{2}^{2}\Big{]}+LT\int_{0}^{T}\mathbb{E}\big{[} \left\lVert\mathbf{x}(t)-\widehat{{\mathbf{x}}}(t)\right\rVert_{2}^{2}\big{]} \mathop{}\!\mathrm{d}t+\varepsilon_{\mathrm{est}}^{2}T.$ $$

Here, we apply the Cauchy-Schwarz (CS) inequality and the Itô isometry in the third inequality, the CS inequality and $(\sqrt{u}-\sqrt{v})^{2}\leq\left|u-v\right|$ ($u$, $v\geq 0$) in the fourth inequality, and the estimation error assumption in the last equality.

Since the dynamics in Eqs. (12) and (13) start from the same initial condition sampled from $p_{0}$, we have $\mathbb{E}\Big{[}\left\lVert\mathbf{x}(0)-\widehat{{\mathbf{x}}}(0)\right
\rVert_{2}^{2}\Big{]}=0$.
Applying the Grönwall’s inequality and the definition of the Wasserstein-2 distance, then we obtain

$$ $\displaystyle W_{2}^{2}(p_{T;\bm{\phi},\bm{\varphi}},p_{T})\lesssim\varepsilon _{\mathrm{est}}^{2}Te^{LT}.$ $$

Combining the above inequality and the result of Theorem 4.1 that

$$ $\displaystyle W_{2}^{2}\big{(}p_{T},p_{\mathrm{target}}\big{)}\lesssim\frac{2} {\alpha}e^{-\alpha T}\mathrm{KL}\big{(}p_{0}\|p_{\mathrm{target}}\big{)},$ $$

we finally derive the following inequality by applying CS inequality

$$ $\displaystyle W_{2}^{2}(p_{T;\bm{\phi},\bm{\varphi}},p_{\mathrm{target}}) \lesssim\varepsilon_{\mathrm{est}}^{2}Te^{LT}+\frac{2}{\alpha}e^{-\alpha T} \mathrm{KL}\big{(}p_{0}\|p_{\mathrm{target}}\big{)}.$ $$

∎

###### Proof.

### Appendix D Bellman Diffusion for MDPs

Figure: Figure 7: $15\times 15$ randomly sampled images from our latent Bellman Diffusion model that is trained on the MNIST dataset. We can see that most of the results are high-quality.
Refer to caption: extracted/5896758/images/8.png

While it is clear that Bellman Diffusion can be used as a generative model, we show that our framework applies to MDPs, which is not the case of its counterpart (e.g., SGMs).

#### D.1 Disfavored Full Trajectory Sampling

As mentioned in Sec. 2, a DGM that is qualified to be applied with the efficient Bellman update needs to satisfy some linearity condition, otherwise one can only sample full state-action trajectories to train the DGM, which is too costly for many RL environments. To understand this point, suppose that there is an $1$-dimensional maze with $N$ blocks, with a robot moving from the leftmost block to the rightmost block. If one directly trains the return model with the returns computed from full trajectories, then the robot has to try to move to the final block after each action, resulting in a time complexity at least as $\mathcal{O}(N\cdot N)=\mathcal{O}(N^{2})$ for every episode. In contrast, if the return model can be trained with partial trajectories (e.g., $1$ step) through the Bellman equation, then the time complexity would be significantly reduced (e.g., $\mathcal{O}(N^{2})$). There are many RL environments where the number $N$ can be very big. For example, StarCraft II (Vinyals et al., ) and Counter-Strike (Pearce & Zhu, ), where a full trajectory can contain over ten thousand steps.

#### D.2 Efficient Bellman Diffusion Update

Figure: Algorithm 3 Planning with Bellman Diffusion

Conventionally, the MDP is specified by a $5$-tuple $(\mathcal{Z},\mathcal{A},p_{\mathrm{tran}},p_{\mathrm{rwd}},\gamma)$, where $\mathcal{Z}$ is the state space, $\mathcal{A}$ is the action space, $p_{\mathrm{tran}}:\mathcal{Z}\times\mathcal{A}\times\mathcal{Z}\rightarrow
\mathbb{R}$ represents the transition probability, $p_{\mathrm{rwd}}:\mathcal{Z}\times\mathcal{A}\rightarrow\mathbb{R}$ denotes the reward model, and $\gamma$ is the discount factor. Given a policy $\pi$ that takes some action $a\in\mathcal{A}$ at every state $z\in\mathcal{Z}$, we aim to estimate the probability distribution of discounted return $X=\sum_{t\geq 1}\gamma^{t-1}R_{t}$ for every state $s$ or state-action pair $(s,a)$.

Let us take the state return $X_{z},z\in\mathcal{Z}$ as an example. To apply Bellman Diffusion to that case, we have to parameterize the gradient and scalar field models $g_{\bm{\phi}}(x),s_{\bm{\varphi}}(x)$ for every state $z$, which is memory consuming for a large state space $\mathcal{Z}$. We adopt the notation $g_{\bm{\phi}}(x)$, rather than $\mathbf{g}_{\bm{\phi}}(x,z)$, because the return distribution is typically one dimensional. To tackle this issue, we share the model parameters $\bm{\phi},\bm{\varphi}$ among different states, letting the models $g_{\bm{\phi}}(x,z),s_{\bm{\varphi}}(x,z)$ depend on state $z$. With this parameterization strategy, we show a planning algorithm in Algorithm 3.

In that procedure, we assume that the reward at the end state is a scalar and $\xi$ is a small value. This assumption applies to most of the scenarios. For example, one will win or fail in the end of a game. We also respectively set $\mathrm{x}_{\mathrm{min}},\mathrm{x}_{\mathrm{max}}$ as the minimum and maximum returns one can get. this algorithm is for planning, and it can be further applied to RL by adding a step of action selection: most of the time choosing the action that leads to the best expected return, and otherwise doing random exploration.

A baseline is C51 (Bellemare et al., ), which models the return distribution with a simple categorical distribution. Its training algorithm is as shown in Algorithm 1 of their paper. Obviously, diffusion model can model much more complex distributions than a categorical parameterization. Therefore, our model will perform better than C51 in the task of distributional RL.

### Appendix E Experiment Settings

Unless specified, we construct the gradient and scalar field models $\mathbf{g}_{\bm{\phi}}(\mathbf{x})$ and $s_{\bm{\varphi}}(\mathbf{x})$ using MLPs (Pinkus, ). We employ Adam algorithm (Kingma, ) for optimization, without weight decay or dropout. The parameter $\epsilon$ in the loss functions $\widebar{\mathcal{L}}_{\mathrm{grad}}^{\mathrm{slice}}(\bm{\phi};\epsilon)$ and $\widebar{\mathcal{L}}_{\mathrm{id}}^{\mathrm{slice}}(\bm{\varphi};\epsilon)$ ranges from $0.1$ to $1.0$, depending on the task. For the sampling dynamics defined in Eq. (12), we typically set $T=300$ and $\eta=0.1$. All models are trained on a single A100 GPU with 40GB memory, taking only a few tens of minutes to a few hours.

### Appendix F Additional Experiments

Due to the limited space, we put the minor experiments here in the appendix. The main experiments involving field estimation, generative modeling, and RL are placed in the main text.

#### F.1 Image Generation

While image generation is not the main focus of our paper, we show that Bellman Diffusion is also promising in that direction. We adopt a variant of the widely used architecture of latent diffusion (Rombach et al., ), with VAE to encode images into latent representations and Bellman Diffusion to learn the distribution of such representations. We run such a model on MNIST (Deng, ), a classical image dataset. The results are shown in Fig. 7. We can see that most generated images are high-quality. This experiment verify that Bellman Diffusion is applicable to high-dimensional data, including image generation.

#### F.2 Ablation Studies

**Table 2: The experiment results of our case studies, with Wasserstein distance as the metric.**
| Method | Abalone | Telemonitoring |
| --- | --- | --- |
| Bellman Diffusion w/ $\epsilon=0.5,n=1$ | $0.975$ | $2.167$ |
| Bellman Diffusion w/ $\epsilon=1.0,n=1$ | $1.113$ | $2.379$ |
| Bellman Diffusion w/ $\epsilon=0.1,n=1$ | $0.875$ | $2.075$ |
| Bellman Diffusion w/ $\epsilon=0.01,n=1$ | $1.567$ | $3.231$ |
| Bellman Diffusion w/ $\epsilon=0.5,n=2$ | $0.912$ | $2.073$ |
| Bellman Diffusion w/ $\epsilon=0.5,n=3$ | $0.895$ | $1.951$ |

There are some important hyper-parameters of Bellman Diffusion that need careful studies to determine their proper values for use. This part aims to achieve this goal. We adopt two tabular datasets: Abalone and Telemonitoring, with the Wasserstein distance as the metric.

##### The variance of Gaussian coefficients.

The loss functions $\widebar{\mathcal{L}}_{\mathrm{grad}}(\bm{\phi};\epsilon),\widebar{\mathcal{L}
}_{\mathrm{id}}(\bm{\varphi};\epsilon)$ of both gradient and scalar matching contain a term $\epsilon$, which is to relax their original limit forms for practical computation. As shown in the first $4$ rows of Table 2, either too big or too small value of term $\epsilon$ leads to worse performance of our Bellman Diffusion model. These experiment results also make sense because too big $\epsilon$ will significantly deviate the loss functions from their limit values, and too small $\epsilon$ will also cause numerical instability.

##### Number of slice vectors.

Intuitively, more slice vectors will make our loss estimation more accurate, leading to better model performance. The experiment results in the first and the last two rows of Table 2 confirm this intuition, but also indicate that such performance gains are not notable. Therefore, we adopt $n=1$ slice vectors in experiments to maintain high efficiency.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 1
questions: * Please, check weakness section.
* Instead of Eq12, why not use SVGD sampler and the derived distribution following:
Messaoud S, Mokeddem B, Xue Z, Pang L, An B, Chen H, Chawla S. S $^ 2$ AC: Energy-Based Reinforcement Learning with Stein Soft Actor Critic. ICLR., 2024 ?
The sampler doesn’t depend on the score of the updated particle itself, which makes it have the desired property of linearity.
soundness: 1
strengths: * The author propose a sampling technique that combines both the target density and its derivative and show that it theory it leads to convergence to the target.
summary: The authors propose a new genarative model, Bellman diffusion, that only depends on the density and its derivative. This choice enables learning return distributions in distributional RL as it's consistent with the linearity requirement of an MDP.
weaknesses: * Structure and clarity: The paper can be restructured in a better way.  For instance, the motivation would be clearer if it would have integrated some material from Sec. 21 can help convey in the intro. Also, in the intro, the authors refer to equations (83) from the approach section. This assumes that the paper has to be read twice?

* Notation issues:
** x is undefined in Eq1
** The prime is Eq2, line 2, right most term, should be applied to z
** Extra dot in line 185

* Formulation lacks preciseness: 
** Eq 1 in linear in p but not linear in x. The fact that the linearity is related to the operator can only be understood in Sec 2.2
** The explanation of eq2 can be clearer is it would have explicitly mentioned that the energy formulation transforms the equality in eq2 into an inequality. 

* Approach:
** What's the use of the isotropic Gaussan in Eq 199? There are no references to how this proxy is derived. It's similar to score matching with a new Gaussian term?
** The score matching loss is not scalable (even the sliced version is not). It's established that score matching is not equivalent to fisher divergence because the derivation requires integration by parts which assumes access to an infinite number of samples. This is not the case in practice.
** A sub-section or a figure/algorithm showcasing how this generative model can be used in distributional RL is missing.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. How strong are the assumptions (Assumptions C.1 and C.2) required for proving Theorems 4.1 and 4.2, compared with those typically used in related theoretical analysis? It would be helpful to see more justification regarding their validity. For example, could you provide examples of other works in the field that rely on similar assumptions?
soundness: 4
strengths: 1. The proposed model addresses a key gap in using deep generative models for Markov decision processes. It is well motivated by maintaining the linearity of the Bellman equation.
2. Following the motivation, the authors build a solid theoretical framework around the proposed model. Some important theorems are stated and proved, including the steady-state convergence theorem (Theorem 4.1) and error bounds for neural network approximation (Theorem 4.2). Note that I do not fully follow the proof in the appendix, and cannot guarantee its correctness.
3. The experiments performed are in a variety of domains, including synthetic point distributions, images, and RL enviroments. I especially like the experiments on OpenAI Gym, which demonstrates the fast and stable convergence of the proposed model over the conventional baseline.
summary: This paper proposes a new generative framework, Bellman Diffusion, tailored for applications in Markov Decision Processes (MDPs) and distributional reinforcement learning. Because of the inherent non-linearity in the modeling operator, traditional generative models such as energy-based models (EBMs) and score-based generative models (SGMs) cannot be applied to RL contexts; they cannot preserve the linearity of the Bellman update. The proposed new model addresses this challenge by modeling both gradient and scalar fields directly in the distribution space, thereby maintaining linearity and enabling effective generative modeling in MDPs.

The paper introduces divergence-based training methods to optimize neural networks to approximate both fields and defines a specialized SDE for sampling. On the theoretical side, the paper proves that the proposed model will converge to a target distribution (in terms of both KL divergence and Wasserstein distance) regardless of the initial distribution. Experimentally, Bellman Diffusion demonstrates superior performance in estimating and generating target distributions. Also, the proposed model performs well on two OpenAI Gym environments and converges faster than the baseline model.
weaknesses: While this paper is primarily focused on methodology and theoretical contributions, I believe there is room for improvement on the experimental side:

1. The abstract claims that the proposed model is a "capable image generator." However, the only image generation results provided are on MNIST (Figure 7, Appendix), and these are purely qualitative. This claim would be better supported with quantitative experiments on real-world image datasets.

2. Are there specific reasons preventing the application of Bellman Diffusion to larger or more complex RL environments? Currently, only two simple examples are shown. Additionally, if possible at all, it would be valuable to compare the performance of Bellman Diffusion with that of denoising diffusion models on these RL tasks, to observe how the non-linearity of traditional diffusion models impacts their performance in this context.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: One concurrent work [1] also seems to optimize the diffusion model in a temporal difference manner, and they used existing diffusion models, rather than proposing a new SDE with linear operators as in Bellman Diffusion. Although I understand that [1] deals with the successor measure instead of value distribution, could the authors briefly discuss about the difference or relationship between [1] and Bellman Diffusion? 

[1] Liam Schramm and Abdeslam Boularias. Bellman Diffusion Models.
soundness: 3
strengths: The motivation behind this paper is both clear and compelling. In reinforcement learning, many distributions exhibit inherent structures that can be effectively described through Bellman iteration. Beyond the return or value distribution, other examples include the successor state distribution, which describes the probability of transitioning to a state s' at any time after taking an action a in a state s. Learning these distributions—even with powerful generative models—can be challenging if the data’s inherent structure isn’t properly respected. In this regard, Bellman Diffusion offers a practical solution, enabling the modeling of such structured distributions in a way that aligns with their underlying characteristics.
summary: This paper introduces Bellman Diffusion, a novel generative model designed to approximate both the gradient and scalar field of the data’s probability distribution. Bellman Diffusion is distinguished by its linear modeling property, meaning that if we aim to model a combined distribution f+g, we can simply sum the fields of f and g. This property is particularly advantageous for distributional reinforcement learning (RL). By first modeling the distribution of returns in terminal states, we can subsequently perform Bellman-like updates on the approximated fields, ultimately deriving the distribution for all states in the Markov Decision Process (MDP). In essence, Bellman Diffusion’s linearity makes it well-suited for distribution modeling tasks where data exhibits a linear structure. The authors validate Bellman Diffusion’s effectiveness through evaluations on low-dimensional toy tasks, high-dimensional benchmarks, and applications in distributional RL.
weaknesses: The presentation of the paper is poor. Since the authors motivate their approach through its application in distributional reinforcement learning, it would be more effective to explain how Bellman Diffusion can be integrated into and facilitate distributional RL (currently in Appendix D) directly within the main method section.

The evaluation of the proposed method is quite limited in both comparisons with existing methods and also the explanations about the performance. 
+ In Section 6.1, the authors primarily demonstrate that Bellman Diffusion can model synthetic data distributions (though there are apparent modeling inconsistencies). Including comparisons with established baseline methods, such as DDPM, would provide a clearer picture of the advantages Bellman Diffusion offers.
+ In Section 6.2, given that the tasks lack internal structures, it raises the question of why Bellman Diffusion outperforms DDPM in certain tasks. Further explanation here would clarify the observed performance benefits and also certain claims made in line 431. Regarding Section 6.3, the benchmark environments (FrozenLake and CartPole) are too simple to truly assess the method’s effectiveness. These environments have relatively small state and action spaces, leading to straightforward return distributions. In contrast, existing methods like C51 are validated on more complex tasks, such as those in Atari environments, which exhibit greater structural complexity. In light of this, I think it would be necessary to extend the evaluations to more complex environments to truly reflect the effectiveness of the proposed method.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
questions: In addition to the aforementioned weakness comments, I request the authors to provide clarifications for the following questions.

1.	Are there any limitations or failure cases where Bellman Diffusion does not perform well compared to other generative modeling approaches?

2.	How computationally expensive is the training process for Bellman Diffusion compared to other generative modeling approaches? Are there any tricks used to improve training efficiency?

3.	The paper shows experimental anecdotes on how Bellman Diffusion could effectively learn multi-modal distributions. How does it compare to other methods like normalizing flows or mixture density networks for multi-modal density estimation?

4.	The slice trick is used to improve sample efficiency. How sensitive is the method to the choice of slice vector distributions q(v) and q(w)? Are there guidelines for selecting these?

5.	Can the authors please share details of the implementation codebase that was used for the experimental evaluations ?
soundness: 2
strengths: 1.	The paper put forth an innovative solution to a well-defined and very interesting problem in applying generative models to MDPs.

2.	The authors provide a rigorous theoretical basis for their algorithmic innovations, including convergence guarantees and error bound analysis.

3.	The paper offers detailed algorithms for training and sampling, making the method easier to follow for implementation purposes.

4.	The authors demonstrate the effectiveness of Bellman Diffusion on both synthetic and real datasets, and also show improvements towards distributional RL tasks.
summary: This paper introduces Bellman Diffusion, a deep generative modeling framework designed to overcome limitations in applying existing generative models to Markov Decision Processes (MDPs), especially when the downstream task is distributional Reinforcement Learning (RL). The authors identify a key issue: the nonlinearity of modern deep generative models conflicts with the linearity required by the Bellman equation in MDPs. To address this, they propose modeling the gradient field and scalar field of the target distribution directly, which preserves linearity.
weaknesses: 1.	The method introduces additional complexity compared to traditional histogram-based approaches, which may limit its adoption in some practical scenarios.

2.	While the paper shows improvements over histogram-based methods, comparisons with other advanced generative modeling techniques in RL contexts are not extensively explored.

3.	The paper does not thoroughly address the scalability of the method to very high-dimensional problems or large-scale RL environments.

4.	The authors must improve quality and readability of their figures and other visualizations. For instance, Figure 4 is especially hard to follow.

5. There are multiple typos/grammatical errors across the manuscript which should be fixed.

Meta-review: This paper introduces Bellman Diffusion, which exploits advanced generative models for distributional reinforcement learning with linearity in MDPs. The method is well-motivated to fill the gap in distributional RL, and derived with the practical objectives. 

However, the major issues raised by the authors lies in the following aspects: 

1, The presentation of the methodology is difficult to follow. 

2, The empirical evaluation is mainly focusing on simple RL environments, with limited baselines. Therefore, it is difficult to justify the advantages in practice. 

In sum, the paper proposed an interesting method for distributional RL, however, it is not ready enough to be published yet.

Decision: Reject