Bellman Diffusion Models

The successor state measure is a central object of study in reinforcement learning (RL). A common statement of the objective is to find the policy that induces the state occupancy measure with the highest expected reward [4, 3, 6, 5, 7]. The state occupancy measure (SOM) has also received considerable attention in the RL theory community, as a number of provably efficient exploration schemes revolve around regularizing the state occupancy measure [1, 2, 8]. We explore a closely related concept, the state successor measure (SSM), which is the probability distribution over future states, given that the agent is currently at state s and takes action a. Despite their utility, the problem of learning the successor measure or state occupancy measure has received relatively little attention in the empirical RL community. While the full reasons for this are difficult to pin down, we argue that it is in large part due to the lack of an expressive and learnable representation that can be easily normalized. We argue that diffusion models can address this deficiency.