Bellman Diffusion Models

TL;DR

This paper investigates using diffusion models to represent the successor state measure in policies, demonstrating that Bellman flow constraints simplify the diffusion-based Bellman update for reinforcement learning.

Contribution

It introduces a novel approach of applying diffusion models to the successor state measure, leveraging Bellman flow constraints for simplified policy updates.

Findings

Bellman flow constraints lead to a simple Bellman update on diffusion step distribution.

Diffusion models effectively represent policies in offline reinforcement learning.

The approach improves policy modeling by integrating diffusion with Bellman equations.

Abstract

Diffusion models have seen tremendous success as generative architectures. Recently, they have been shown to be effective at modelling policies for offline reinforcement learning and imitation learning. We explore using diffusion as a model class for the successor state measure (SSM) of a policy. We find that enforcing the Bellman flow constraints leads to a simple Bellman update on the diffusion step distribution.