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

TL;DR

This paper introduces Bellman Diffusion, a new generative modeling framework that preserves linearity in distribution space for MDPs, enabling faster and more accurate distributional RL applications.

Contribution

It proposes Bellman Diffusion, a linear operator-based DGM that maintains the linearity of the Bellman equation, with divergence-based training and a novel SDE for sampling.

Findings

Achieves accurate distribution field estimations.

Converges 1.5x faster than histogram-based methods.

Effectively generates images with high fidelity.

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 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.