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

arXiv: https://arxiv.org/abs/2410.01796 · HTML: https://mlsanity.com/paper/2410.01796/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion and flow policies for robotics, one of 14 directions in Diffusion models.
- Builds on 32 of its 45 references that are in the map.
- First posted on arXiv 2 Oct 2024.

## Placed in

- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)

## Builds on (most cited first, 30 of 32)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Generative Adversarial Networks](https://mlsanity.com/paper/2203.00667/) (2022)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [OpenAI Gym](https://mlsanity.com/paper/1606.01540/) (2016)
- [Rainbow: Combining Improvements in Deep Reinforcement Learning](https://mlsanity.com/paper/1710.02298/) (2017)
- [A Distributional Perspective on Reinforcement Learning](https://mlsanity.com/paper/1707.06887/) (2017)
- [Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning](https://mlsanity.com/paper/2208.06193/) (2022)
- [Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion](https://mlsanity.com/paper/2310.02279/) (2023)
- [Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions](https://mlsanity.com/paper/2209.11215/) (2022)
- [Sliced Score Matching: A Scalable Approach to Density and Score Estimation](https://mlsanity.com/paper/1905.07088/) (2019)
- [Implicit Quantile Networks for Distributional Reinforcement Learning](https://mlsanity.com/paper/1806.06923/) (2018)
- [Your Diffusion Model is Secretly a Zero-Shot Classifier](https://mlsanity.com/paper/2303.16203/) (2023)
- [StarCraft II: A New Challenge for Reinforcement Learning](https://mlsanity.com/paper/1708.04782/) (2017)
- [Convergence of denoising diffusion models under the manifold hypothesis](https://mlsanity.com/paper/2208.05314/) (2022)
- [Rapid Convergence of the Unadjusted Langevin Algorithm: Isoperimetry Suffices](https://mlsanity.com/paper/1903.08568/) (2019)
- [Generalized Sliced Wasserstein Distances](https://mlsanity.com/paper/1902.00434/) (2019)
- [Convergence of score-based generative modeling for general data distributions](https://mlsanity.com/paper/2209.12381/) (2022)
- [Training generative neural networks via Maximum Mean Discrepancy optimization](https://mlsanity.com/paper/1505.03906/) (2015)
- [Diffusion Models for Black-Box Optimization](https://mlsanity.com/paper/2306.07180/) (2023)
- [Counter-Strike Deathmatch with Large-Scale Behavioural Cloning](https://mlsanity.com/paper/2104.04258/) (2021)
- [A Comparative Analysis of Expected and Distributional Reinforcement Learning](https://mlsanity.com/paper/1901.11084/) (2019)
- [FP-Diffusion: Improving Score-based Diffusion Models by Enforcing the Underlying Score Fokker-Planck Equation](https://mlsanity.com/paper/2210.04296/) (2022)
- [PaGoDA: Progressive Growing of a One-Step Generator from a Low-Resolution Diffusion Teacher](https://mlsanity.com/paper/2405.14822/) (2024)

## Built on by (most cited first, 4 of 4)

- [Path-Coupled Bellman Flows for Distributional Reinforcement Learning](https://mlsanity.com/paper/2605.08253/) (2026)
- [Dual-Flow Reinforcement Learning with State-Aware Exploration](https://mlsanity.com/paper/2606.29820/) (2026)
- [Distributional Inverse Reinforcement Learning](https://mlsanity.com/paper/2510.03013/) (2025)
- [Retimed Bellman Flows: Escaping the Impossible Triangle of Velocity Bootstrapping](https://mlsanity.com/paper/2609.32828/) (2026)

## Abstract (arXiv)

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 conventional 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, 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.5x 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.
