Bayesian Flow Networks for Offline Trajectory Planning

arXiv cs.LG Papers

Summary

The article presents BFN-RL, a unified generative modeling framework for offline reinforcement learning based on Bayesian Flow Networks, capable of generating effective trajectories across discrete and continuous state spaces.

arXiv:2608.25163v1 Announce Type: new Abstract: Offline reinforcement learning (RL) leverages static datasets to learn decision policies without real-time environment interaction. While recent sequence-modeling approaches rely on continuous diffusion models for trajectory synthesis, applying these methods to discrete planning tasks requires a categorical formulation rather than the standard Gaussian construction. We present BFN-RL, a unified generative modeling framework for offline RL based on Bayesian Flow Networks (BFNs). By iteratively evolving distribution parameters rather than noisy data instances, BFN-RL natively models both discrete and continuous trajectory spaces within a single probabilistic formulation. The categorical planner generates future state sequences, and a learned inverse-dynamics model converts consecutive generated states into actions. Evaluations in discrete planning and continuous control show that BFN-RL can generate effective trajectories across both categorical and continuous state spaces. Our results establish BFNs as a versatile generative foundation for offline trajectory planning across data modalities.
Original Article
View Cached Full Text

Cached at: 08/27/26, 09:34 AM

# Bayesian Flow Networks for Offline Trajectory Planning
Source: [https://arxiv.org/html/2608.25163](https://arxiv.org/html/2608.25163)
Ludvig KillingbergHelge LangsethAffiliation:Norwegian University of Science and Technology

###### Abstract

Offline reinforcement learning \(RL\) leverages static datasets to learn decision policies without real\-time environment interaction\. While recent sequence\-modeling approaches rely on continuous diffusion models for trajectory synthesis, applying these methods to discrete planning tasks requires a categorical formulation rather than the standard Gaussian construction\. We present BFN\-RL, a unified generative modeling framework for offline RL based on Bayesian Flow Networks \(BFNs\)\. By iteratively evolving distribution parameters rather than noisy data instances, BFN\-RL natively models both discrete and continuous trajectory spaces within a single probabilistic formulation\. The categorical planner generates future state sequences, and a learned inverse\-dynamics model converts consecutive generated states into actions\. Evaluations in discrete planning and continuous control show that BFN\-RL can generate effective trajectories across both categorical and continuous state spaces\. Our results establish BFNs as a versatile generative foundation for offline trajectory planning across data modalities\.

## 1Introduction

Offline reinforcement learning \(RL\)\([Sutton and Barto, 2018](https://arxiv.org/html/2608.25163#bib.bib18);[Levine et al\., 2020](https://arxiv.org/html/2608.25163#bib.bib17)\)is a powerful paradigm that leverages static, previously collected datasets to learn effective decision policies without requiring real\-time environment interaction\. By eliminating the safety hazards associated with online exploration, offline RL is particularly suited for high\-risk domains such as autonomous driving and medical decision\-making\. In recent years, framing offline RL as a conditional sequence\-modeling task\([Janner et al\., 2021](https://arxiv.org/html/2608.25163#bib.bib12);[Chen et al\., 2021](https://arxiv.org/html/2608.25163#bib.bib15)\)has emerged as a compelling alternative to traditional value\-based methods, which frequently suffer from value overestimation on out\-of\-distribution state\-action pairs\([Agarwal et al\., 2020](https://arxiv.org/html/2608.25163#bib.bib16);[Levine et al\., 2020](https://arxiv.org/html/2608.25163#bib.bib17)\)\. By viewing trajectory generation through the lens of conditional generative modeling, sequence\-based agents synthesize high\-return trajectories by capturing complex temporal dependencies across long horizons\([Janner et al\., 2022](https://arxiv.org/html/2608.25163#bib.bib3);[Ajay et al\., 2023](https://arxiv.org/html/2608.25163#bib.bib2)\)\.

Despite their empirical success, contemporary sequence\-modeling approaches rely almost exclusively on Gaussian Denoising Diffusion Probabilistic Models \(DDPMs\)\([Ho et al\., 2020](https://arxiv.org/html/2608.25163#bib.bib4);[Janner et al\., 2022](https://arxiv.org/html/2608.25163#bib.bib3);[Ajay et al\., 2023](https://arxiv.org/html/2608.25163#bib.bib2)\)\. While diffusion models excel in continuous control domains, discrete states and actions require a categorical formulation rather than the standard Gaussian construction\([Austin et al\., 2021](https://arxiv.org/html/2608.25163#bib.bib20);[Lou et al\., 2023](https://arxiv.org/html/2608.25163#bib.bib19)\)\. This motivates studying a generative framework that handles both data types natively\.

To address this limitation, we present*BFN\-RL*, an offline reinforcement learning framework grounded in Bayesian Flow Networks \(BFNs\)\([Graves et al\., 2023](https://arxiv.org/html/2608.25163#bib.bib11)\)\. Unlike diffusion models that iteratively denoise corrupted data instances, BFNs operate by updating the parameters of an input distribution via Bayesian inference driven by continuous\-time parameter flows\([Graves et al\., 2023](https://arxiv.org/html/2608.25163#bib.bib11)\)\. This formulation yields a unified generative paradigm that natively handles categorical, continuous, and discretized variables within the same probabilistic framework\.

Our main contributions are summarized as follows:

- •Unified Generative Planning Framework:We introduce BFN\-RL, establishing a parameter\-flow paradigm for offline sequence modeling that operates seamlessly across both discrete and continuous state\-action spaces\([Graves et al\., 2023](https://arxiv.org/html/2608.25163#bib.bib11)\)\.
- •Categorical Trajectory Planning:We combine a categorical BFN state\-sequence model with learned inverse dynamics to obtain a discrete planner\.
- •Cross\-Domain Evaluation:We demonstrate that the same BFN planning formulation is viable in both discrete and continuous control problems\.

## 2Preliminaries

### 2\.1Reinforcement Learning

Reinforcement learning \(RL\) is a framework for learning to make decisions in an environment\([Sutton and Barto, 2018](https://arxiv.org/html/2608.25163#bib.bib18)\)\. The interactions with the environment are modeled as a Markov decision process \(MDP\), which is a tuple\(𝒮,𝒜,𝒫,ℛ,γ\)\(\\mathcal\{S\},\\mathcal\{A\},\\mathcal\{P\},\\mathcal\{R\},\\gamma\), where𝒮\\mathcal\{S\}is the state space,𝒜\\mathcal\{A\}is the action space,𝒫\\mathcal\{P\}is the transition function,ℛ\\mathcal\{R\}is the reward function, andγ\\gammais the discount factor\. In an environment where the agent performs the actiona∈𝒜a\\in\\mathcal\{A\}in states∈𝒮s\\in\\mathcal\{S\}, the next state,s′∈𝒮s^\{\\prime\}\\in\\mathcal\{S\}is sampled from𝒫⁡\(s,a\)\\mathcal\{P\}\(s,a\), i\.e\., is only dependent on the current state and action, not the history of previous states and actions\. In other words, the domain adheres to the Markov property\. Letrtr\_\{t\}denote the reward received at timett, and letRt=∑i=0∞γi​rt\+iR\_\{t\}=\\sum\_\{i=0\}^\{\\infty\}\\gamma^\{i\}r\_\{t\+i\}be the discounted cumulative reward obtained starting from timett\. Now, one goal of RL is to learn a policyπ:𝒮→𝒜\\pi:\\mathcal\{S\}\\rightarrow\\mathcal\{A\}that maximizes the expected return𝔼⁡\[Rt\]\\mathbb\{E\}\[R\_\{t\}\], where the expectation is taken over the uncertainty defined by both the transition function and the stochastic strategyπ\\pi\.

The exploration\-exploitation trade\-off is a fundamental challenge in reinforcement learning, typically associated with online learning scenarios where agents iteratively interact with an environment to learn optimal policies\. Exploration involves sampling actions to gather information about the environment, potentially leading to the discovery of better strategies, while exploitation entails leveraging known information in an attempt to maximize the expected returns\. Much of the research in online RL is dedicated to striking a balance between exploration and exploitation, devising algorithms that effectively navigate this trade\-off to converge to optimal or near\-optimal policies\.

### 2\.2Offline RL

In the realm of offline reinforcement learning, the primary objective is to learn effective policies from a fixed dataset, without the need for online interactions\([Levine et al\., 2020](https://arxiv.org/html/2608.25163#bib.bib17)\)\. In this context, where the exploration aspect is inherently absent, the focus shifts towards effectively utilizing the available dataset to optimize policies\. Traditionally, RL has been concerned with estimating stationary policies or single\-step models, leveraging the Markov property to factorize problems over time\. However, applying standard RL methods to offline settings is challenging because methods relying on an estimated value function often suffer from overestimating the value of out\-of\-distribution state\-action pairs\. Various methods have been proposed to address this issue, including constraining the policy to be close to the data distribution\([Peters et al\., 2010](https://arxiv.org/html/2608.25163#bib.bib10)\)or using a conservative value function\([Kumar et al\., 2020](https://arxiv.org/html/2608.25163#bib.bib9)\)\. Our solution, on the other hand, is to produce a sequence of steps that will be generated conditionally on the objective of the RL agent\. An intriguing perspective emerges when we view RL through the lens of sequence modeling\. Instead of treating it as a specialized domain, we can consider RL as a generic sequence modeling problem\. The crux of this viewpoint lies in producing a sequence of actions that leads to a sequence of high rewards\. Earlier work has solved this by conditioning the model on returns such that trajectories with high returns can be generated in online settings\([Ajay et al\., 2023](https://arxiv.org/html/2608.25163#bib.bib2);[Janner et al\., 2021](https://arxiv.org/html/2608.25163#bib.bib12)\)\. By adopting this perspective, we can simplify design decisions and dispense with many components commonly found in offline RL algorithms\. This approach not only demonstrates flexibility across various tasks such as long\-horizon dynamics prediction, imitation learning, goal\-conditioned RL, and offline RL but also yields state\-of\-the\-art planners in sparse\-reward, long\-horizon scenarios\([Janner et al\., 2022](https://arxiv.org/html/2608.25163#bib.bib3);[Ajay et al\., 2023](https://arxiv.org/html/2608.25163#bib.bib2)\)\.

### 2\.3Denoising Diffusion Probabilistic Models

Since the current state of the art in this domain\([Janner et al\., 2022](https://arxiv.org/html/2608.25163#bib.bib3);[Ajay et al\., 2023](https://arxiv.org/html/2608.25163#bib.bib2)\)rely on diffusion models as the generative model, we will give a brief introduction to denoising diffusion probabilistic models \(DDPMs\)\([Ho et al\., 2020](https://arxiv.org/html/2608.25163#bib.bib4)\)\. This will also serve as a backdrop for our discussion of Bayesian flow networks, which follows in Section[2\.5](https://arxiv.org/html/2608.25163#S2.SS5)\. DDPMs are a type of generative model inspired by non\-equilibrium thermodynamics\. The model is defined by aforward processthat slowly adds Gaussian noise to data, and itsreverse, that amounts to learning to iteratively denoise the noisy data\. Diffusion models have primarily been used for image generation, but have also shown state\-of\-the\-art performance in other domains, like video generation and 3D model\([Ho et al\., 2022](https://arxiv.org/html/2608.25163#bib.bib5);[Luo and Hu, 2021](https://arxiv.org/html/2608.25163#bib.bib6)\)\.

Given data𝐱0∼q⁡\(𝐱\)\\mathbf\{x\}\_\{0\}\\sim q\(\\mathbf\{x\}\), we define the*forward process*to produce a sequence of noisy samples𝐱1,…,𝐱K\\mathbf\{x\}\_\{1\},\\dots,\\mathbf\{x\}\_\{K\},

q⁡\(𝐱k\|𝐱k−1\)=𝒩⁡\(𝐱k,1−βk​𝐱k−1,βk​𝑰\)\.q\(\\mathbf\{x\}\_\{k\}\|\\mathbf\{x\}\_\{k\-1\}\)=\\mathcal\{N\}\(\\mathbf\{x\}\_\{k\};\\sqrt\{1\-\\beta\_\{k\}\}\\mathbf\{x\}\_\{k\-1\},\\beta\_\{k\}\\bm\{I\}\)\.where\{βi∈\(0,1\)\}i=1K\\\{\\beta\_\{i\}\\in\(0,1\)\\\}\_\{i=1\}^\{K\}is a carefully chosen variance schedule\. A nice property of the forward process is that we can directly sample𝐱k\\mathbf\{x\}\_\{k\}at any stepii, because the distributionq⁡\(𝐱k\|𝐱0\)q\(\\mathbf\{x\}\_\{k\}\|\\mathbf\{x\}\_\{0\}\)can be derived using the property that a sum of uncorrelated normally distributed random variables is normally distributed\. Letak=1−βka\_\{k\}=1\-\\beta\_\{k\}anda¯k=∏i=1kak\\overline\{a\}\_\{k\}=\\prod\_\{i=1\}^\{k\}a\_\{k\}, then

q⁡\(𝐱k\|𝐱0\)=𝒩⁡\(𝐱k,a¯k​𝐱0,\(1−a¯k\)​𝑰\)\.q\(\\mathbf\{x\}\_\{k\}\|\\mathbf\{x\}\_\{0\}\)=\\mathcal\{N\}\\left\(\\mathbf\{x\}\_\{k\};\\sqrt\{\\overline\{a\}\_\{k\}\}\\,\\mathbf\{x\}\_\{0\},\\left\(1\-\\overline\{a\}\_\{k\}\\right\)\\bm\{I\}\\right\)\.Note also that

q⁡\(𝐱k−1\|𝐱k,𝐱0\)=𝒩⁡\(𝐱k−1,𝝁~​\(𝐱k,𝐱0\),β~k​𝑰\),q\(\\mathbf\{x\}\_\{k\-1\}\|\\mathbf\{x\}\_\{k\},\\mathbf\{x\}\_\{0\}\)=\\mathcal\{N\}\\left\(\\mathbf\{x\}\_\{k\-1\};\\tilde\{\\bm\{\\mu\}\}\(\\mathbf\{x\}\_\{k\},\\mathbf\{x\}\_\{0\}\),\\tilde\{\\beta\}\_\{k\}\\bm\{I\}\\right\),\(1\)where

𝝁~​\(𝐱k,𝐱0\)=a¯k−1​βk1−a¯k​𝐱0\+ak​\(1−a¯k−1\)1−a¯k​𝐱k,β~k=βk​1−a¯k−11−a¯k\.\\tilde\{\\bm\{\\mu\}\}\(\\mathbf\{x\}\_\{k\},\\mathbf\{x\}\_\{0\}\)=\\frac\{\\sqrt\{\\overline\{a\}\_\{k\-1\}\}\\beta\_\{k\}\}\{1\-\\overline\{a\}\_\{k\}\}\\mathbf\{x\}\_\{0\}\+\\frac\{\\sqrt\{a\_\{k\}\}\(1\-\\overline\{a\}\_\{k\-1\}\)\}\{1\-\\overline\{a\}\_\{k\}\}\\mathbf\{x\}\_\{k\},\\qquad\\tilde\{\\beta\}\_\{k\}=\\beta\_\{k\}\\frac\{1\-\\overline\{a\}\_\{k\-1\}\}\{1\-\\overline\{a\}\_\{k\}\}\.\(2\)
While the forward process creates a noisy representation of data, the*reverse process*aims to iteratively recreate samples from noise by modeling and then sampling fromq⁡\(𝐱k−1\|𝐱k\)q\(\\mathbf\{x\}\_\{k\-1\}\|\\mathbf\{x\}\_\{k\}\)\. Letpθ​\(𝐱k−1\|𝐱k\)p\_\{\\theta\}\(\\mathbf\{x\}\_\{k\-1\}\|\\mathbf\{x\}\_\{k\}\)be a parameterized approximation ofq⁡\(𝐱k−1\|𝐱k\)q\(\\mathbf\{x\}\_\{k\-1\}\|\\mathbf\{x\}\_\{k\}\)\. This means that we define a neural network model with trainable parameters𝜽\\bm\{\\theta\}that outputs𝝁θ​\(𝐱k,k\)\\bm\{\\mu\}\_\{\\theta\}\(\\mathbf\{x\}\_\{k\},k\)andΣθ​\(𝐱k,k\)\\Sigma\_\{\\theta\}\(\\mathbf\{x\}\_\{k\},k\)so that

pθ​\(𝐱k−1\|𝐱k\)=𝒩⁡\(𝐱k−1,𝝁θ​\(𝐱k,k\),Σθ​\(𝐱k,k\)\)\.p\_\{\\theta\}\(\\mathbf\{x\}\_\{k\-1\}\|\\mathbf\{x\}\_\{k\}\)=\\mathcal\{N\}\\left\(\\mathbf\{x\}\_\{k\-1\};\\bm\{\\mu\}\_\{\\theta\}\(\\mathbf\{x\}\_\{k\},k\),\\Sigma\_\{\\theta\}\(\\mathbf\{x\}\_\{k\},k\)\\right\)\.[Ho et al\. \(2020\)](https://arxiv.org/html/2608.25163#bib.bib4)chose to fix the variance term𝚺θ​\(𝐱k,k\)\\bm\{\\Sigma\}\_\{\\theta\}\(\\mathbf\{x\}\_\{k\},k\)as a constantσk2=β~k\\sigma\_\{k\}^\{2\}=\\tilde\{\\beta\}\_\{k\}, see Eq\. \([2](https://arxiv.org/html/2608.25163#S2.E2)\)\. We therefore only look at how𝝁θ​\(𝐱k,k\)\\bm\{\\mu\}\_\{\\theta\}\(\\mathbf\{x\}\_\{k\},k\)is estimated\. First, we consider the identity

𝝁~k​\(𝐱k,𝐱0\)=1ak​\(𝐱k−1−ak1−a¯k​ϵk\),\\bm\{\\tilde\{\\mu\}\}\_\{k\}\(\\mathbf\{x\}\_\{k\},\\mathbf\{x\}\_\{0\}\)=\\frac\{1\}\{\\sqrt\{a\_\{k\}\}\}\\left\(\\mathbf\{x\}\_\{k\}\-\\frac\{1\-a\_\{k\}\}\{\\sqrt\{1\-\\overline\{a\}\_\{k\}\}\}\\bm\{\\epsilon\}\_\{k\}\\right\),whereϵk∼𝒩⁡\(0,𝑰\)\\bm\{\\epsilon\}\_\{k\}\\sim\\mathcal\{N\}\(0,\\bm\{I\}\); cf\. Eqs\. \([1](https://arxiv.org/html/2608.25163#S2.E1)\) and \([2](https://arxiv.org/html/2608.25163#S2.E2)\)\. Since𝐱k\\mathbf\{x\}\_\{k\}is known during training, we can choose to predictϵk\\bm\{\\epsilon\}\_\{k\}, rather than𝝁~k\\bm\{\\tilde\{\\mu\}\}\_\{k\}directly\. Empirically, this has shown better results\. Let us defineϵθ​\(𝐱,k\)\\bm\{\\epsilon\}\_\{\\theta\}\(\\mathbf\{x\},k\)as a model that predicts the noise,ϵk\\bm\{\\epsilon\}\_\{k\}\. This means that we can define𝝁θ​\(𝐱k,k\)=1ak​\(𝐱k−1−ak1−a¯k​ϵθ​\(𝐱k,k\)\)\\bm\{\\mu\}\_\{\\theta\}\(\\mathbf\{x\}\_\{k\},k\)=\\frac\{1\}\{\\sqrt\{a\_\{k\}\}\}\\left\(\\mathbf\{x\}\_\{k\}\-\\frac\{1\-a\_\{k\}\}\{\\sqrt\{1\-\\overline\{a\}\_\{k\}\}\}\\bm\{\\epsilon\}\_\{\\theta\}\(\\mathbf\{x\}\_\{k\},k\)\\right\)\.[Ho et al\. \(2020\)](https://arxiv.org/html/2608.25163#bib.bib4)derive the following loss function to minimize the difference between𝝁θ\\bm\{\\mu\}\_\{\\theta\}and𝝁~\\bm\{\\tilde\{\\mu\}\}:

L⁡\(θ\)=𝔼k∼\[1,K\],𝐱0,ϵk​\[βk22​σk2​ak​\(1−a¯k\)​‖ϵk−ϵθ​\(𝐱k,k\)‖2\]\.L\(\\theta\)=\\mathbb\{E\}\_\{k\\sim\[1,K\],\\mathbf\{x\}\_\{0\},\\bm\{\\epsilon\}\_\{k\}\}\\left\[\\frac\{\\beta\_\{k\}^\{2\}\}\{2\\sigma^\{2\}\_\{k\}a\_\{k\}\(1\-\\overline\{a\}\_\{k\}\)\}\\left\\lVert\\bm\{\\epsilon\}\_\{k\}\-\\bm\{\\epsilon\}\_\{\\theta\}\(\\mathbf\{x\}\_\{k\},k\)\\right\\rVert^\{2\}\\right\]\.
They also present the following simplified loss function that turns out to give better empirical results:

L⁡\(θ\)=𝔼k∼\[1,K\],𝐱0,ϵk​‖ϵk−ϵθ​\(𝐱k,k\)‖2\.L\(\\theta\)=\\mathbb\{E\}\_\{k\\sim\[1,K\],\\mathbf\{x\}\_\{0\},\\epsilon\_\{k\}\}\\left\\lVert\\epsilon\_\{k\}\-\\epsilon\_\{\\theta\}\(\\mathbf\{x\}\_\{k\},k\)\\right\\rVert^\{2\}\.

### 2\.4Guided Diffusion

We will discuss three ways diffusion models can condition on variables\. The first,classifier\-guideddiffusion\([Dhariwal and Nichol, 2021](https://arxiv.org/html/2608.25163#bib.bib7)\), takes as its starting point that we have a trained probabilistic classifier that classifies objects𝐱k\\mathbf\{x\}\_\{k\}during the reverse process\. We want to use this to produce an object𝐱0\\mathbf\{x\}\_\{0\}that is classified as belonging to a predefined classyy\. The approach uses the gradients of this classifier’s allocated log\-likelihood to the classyywrt\.𝐱k\\mathbf\{x\}\_\{k\}to “push” the reverse process towards objects that are aligned with the conditioning information\. This method has the advantage that the diffusion model does not have to be trained with conditioning variables, the guidance is only related to the reverse and only needs the classifier to be trained on conditioning information\. A model predictorϵ¯θ\\overline\{\\bm\{\\epsilon\}\}\_\{\\theta\}, guided by a classifierh⁡\(y\|𝐱k,k\)h\(y\|\\mathbf\{x\}\_\{k\},k\)meant to estimate the probability that the noisy datapoint𝐱k\\mathbf\{x\}\_\{k\}belongs to classyy, would assume the following form:

ϵ¯θ​\(𝐱k,k,y\)=ϵθ​\(𝐱k,k\)−w​σk​∇𝐱k​log⁡h⁡\(y\|𝐱k,k\),\\overline\{\\bm\{\\epsilon\}\}\_\{\\theta\}\(\\mathbf\{x\}\_\{k\},k,y\)=\\bm\{\\epsilon\}\_\{\\theta\}\(\\mathbf\{x\}\_\{k\},k\)\-w\\sigma\_\{k\}\\nabla\_\{\\mathbf\{x\}\_\{k\}\}\\log\{h\(y\|\\mathbf\{x\}\_\{k\},k\)\},wherewwis a hyper\-parameter controlling the strength of the guidance\.

Secondly,classifier\-free guidance\([Ho and Salimans, 2021](https://arxiv.org/html/2608.25163#bib.bib8)\), plugs the conditioning variable directly into the denoising network as an auxiliary input variable during training\. At test time, the auxiliary variable can be set to the conditioning value\. In this setting, the model predictor takes the following form:

ϵ~​\(𝐱k,k,y\)=\(w\+1\)​ϵθ​\(𝐱k,k,y\)−w​ϵθ​\(𝐱k,k\),\\bm\{\\tilde\{\\epsilon\}\}\(\\mathbf\{x\}\_\{k\},k,y\)=\(w\+1\)\\bm\{\\epsilon\}\_\{\\theta\}\(\\mathbf\{x\}\_\{k\},k,y\)\-w\\bm\{\\epsilon\}\_\{\\theta\}\(\\mathbf\{x\}\_\{k\},k\),where we again usewwto denote the hyper\-parameter that controls the strength of the guidance\. Classifier\-free guidance has shown better practical performance than classifier\-guided diffusion\([Ho and Salimans, 2021](https://arxiv.org/html/2608.25163#bib.bib8)\)\.

Finally, we can also employ inpainting\([Lugmayr et al\., 2022](https://arxiv.org/html/2608.25163#bib.bib13)\)to condition on partial observations\. In the context of image generation, this implies conditioning on some pixels within the image\. Consider an image𝐱0\\mathbf\{x\}\_\{0\}divided into known pixels𝐱0known\\mathbf\{x\}\_\{0\}^\{\\text\{known\}\}and unknown pixels𝐱0unknown\\mathbf\{x\}\_\{0\}^\{\\text\{unknown\}\}, and a mask𝒎\\bm\{m\}defining which pixels are known\. During the reverse process, we define

𝐱k−1known=a¯k​𝐱0\+1−a¯k​𝒛,𝒛∼𝒩⁡\(𝟎,𝐈\),\\mathbf\{x\}\_\{k\-1\}^\{\\text\{known\}\}=\\sqrt\{\\overline\{a\}\_\{k\}\}\\mathbf\{x\}\_\{0\}\+\\sqrt\{1\-\\overline\{a\}\_\{k\}\}\\bm\{z\},\\qquad\\bm\{z\}\\sim\\mathcal\{N\}\(\\bm\{0\},\\mathbf\{I\}\),i\.e\.,𝐱k−1known\\mathbf\{x\}\_\{k\-1\}^\{\\text\{known\}\}is chosen equal to what the forward process would have produced had it started from the known parts of the image\. The unknown pixels at stepkk,𝐱k−1unknown\\mathbf\{x\}\_\{k\-1\}^\{\\text\{unknown\}\}, are computed in standard fashion:

𝐱k−1unknown=1ak​\(𝐱k−βk1−a¯k​ϵθ​\(𝐱k,k\)\)\+σk​𝒛,𝒛∼𝒩⁡\(𝟎,𝐈\)\.\\mathbf\{x\}\_\{k\-1\}^\{\\text\{unknown\}\}=\\frac\{1\}\{\\sqrt\{a\_\{k\}\}\}\\left\(\\mathbf\{x\}\_\{k\}\-\\frac\{\\beta\_\{k\}\}\{\\sqrt\{1\-\\overline\{a\}\_\{k\}\}\}\\bm\{\\epsilon\}\_\{\\theta\}\(\\mathbf\{x\}\_\{k\},k\)\\right\)\+\\sigma\_\{k\}\\bm\{z\},\\qquad\\bm\{z\}\\sim\\mathcal\{N\}\(\\bm\{0\},\\mathbf\{I\}\)\.Finally, we have:

𝐱k−1=𝒎⊙𝐱k−1known\+\(1−𝒎\)⊙𝐱k−1unknown,\\mathbf\{x\}\_\{k\-1\}=\\bm\{m\}\\odot\\mathbf\{x\}\_\{k\-1\}^\{\\text\{known\}\}\+\(1\-\\bm\{m\}\)\\odot\\mathbf\{x\}\_\{k\-1\}^\{\\text\{unknown\}\},\(3\)where⊙\\odotis elementwise multiplication\.

### 2\.5Bayesian Flow Networks

While there are variations of diffusion models that model discrete data\([Lou et al\., 2023](https://arxiv.org/html/2608.25163#bib.bib19);[Austin et al\., 2021](https://arxiv.org/html/2608.25163#bib.bib20)\), these models are not considered state of the art when it comes to generating high\-quality discrete data\([Graves et al\., 2023](https://arxiv.org/html/2608.25163#bib.bib11)\)\. In an attempt to remedy this shortcoming,[Graves et al\. \(2023\)](https://arxiv.org/html/2608.25163#bib.bib11)introducedBayesian flow networks\(BFNs\), a novel generative model capable of generating continuous, discrete, and discretized data\. BFNs resemble diffusion models in that they generate data in an iterative process\. Unlike diffusion models, however, the BFN analogy to the diffusion models’ reverse process iteratively evolvesdistribution parameters, not noised versions of data\. The high\-level idea is to start from a prior distribution and iteratively update the distribution conditioned on a data point sampled from a noisy version of the previous distribution\. BFNs have been shown to perform well on discrete data\([Graves et al\., 2023](https://arxiv.org/html/2608.25163#bib.bib11)\), and are therefore a more natural choice than diffusion models for planning in discrete state spaces\.

Figure[1](https://arxiv.org/html/2608.25163#S2.F1)illustrates the idea of Bayesian flow networks\. At each stepii, the parameters of a distribution \(𝜽i\\bm\{\\theta\}\_\{i\}\) are updated with noisy samples from the data,𝒚i\\bm\{y\}\_\{i\}\. The level of added noise decreases for each step and is at stepiidictated by the accuracyαi\\alpha\_\{i\}\. The parameters𝜽i\\bm\{\\theta\}\_\{i\}are defined as the Bayesian update of the parameters at the previous step,𝜽i−1\\bm\{\\theta\}\_\{i\-1\}, with observation noise parameterized byαi\\alpha\_\{i\}using a predetermined update ruleh⁡\(⋅\)h\(\\cdot\)\. While each evolved distribution is distinct from the data distribution, the idea is that the compound of all generated distributions should approximate the data distribution\. To sample from the data distribution, an initial uninformative distribution is iteratively updated, gradually concentrating around a single data point\. Once the distribution has evolved sufficiently, a single sample from this nearly degenerate distribution will closely approximate a sample from the data distribution\.

𝜽i​=def​h​\(𝜽i−1,𝒚i−1,αi−1\)\\bm\{\\theta\}\_\{i\}\\overset\{\\text\{def\}\}\{=\}h\(\\bm\{\\theta\}\_\{i\-1\},\\bm\{y\}\_\{i\-1\},\\alpha\_\{i\-1\}\)…\\dots𝜽0\\bm\{\\theta\}\_\{0\}𝜽i\+1​=def​h​\(𝜽i,𝒚i,αi\)\\bm\{\\theta\}\_\{i\+1\}\\overset\{\\text\{def\}\}\{=\}h\(\\bm\{\\theta\}\_\{i\},\\bm\{y\}\_\{i\},\\alpha\_\{i\}\)…\\dots𝜽N\\bm\{\\theta\}\_\{N\}

Figure 1:Generative process for Bayesian flow networks\.Figure[2](https://arxiv.org/html/2608.25163#S2.F2)represents the training process of Bayesian Flow Networks\. The aim is to iteratively update the parameters of a distribution,𝜽\\bm\{\\theta\}, beginning with a prior distribution𝜽0\\bm\{\\theta\}\_\{0\}, so that eventually, sampling once from this distribution mirrors sampling from the data it is trained on\. In this process, accuracyαi\\alpha\_\{i\}refers to how well the updated parameters reflect the true data after observing information\. The accuracy quantifies the expected quality of each update, meaning how much closer, in expectation, the updated distribution is to the true data distribution\.

This training process involves the following key steps:

1. 1\.Generateθi\\bm\{\\theta\}\_\{i\}:Given a datapoint𝒙\\bm\{x\}, the parameters𝜽0\\bm\{\\theta\}\_\{0\}can be updatediitimes with noisy data to create parameters𝜽i\\bm\{\\theta\}\_\{i\}\. We will later see that accuracies are additive and that𝜽i\\bm\{\\theta\}\_\{i\}can be generated in a single step with accuracy∑j=0i−1αj\\sum\_\{j=0\}^\{i\-1\}\\alpha\_\{j\}\.
2. 2\.Neural network transformation:The parameters𝜽i\\bm\{\\theta\}\_\{i\}are passed through a neural network, with weights𝝎\\bm\{\\omega\}, which outputs the parameters of a new distribution\. This is referred to as the*output distribution*pOp\_\{O\}\.
3. 3\.Addition of noise:A new*sender distribution*pSp\_\{S\}, is created by adding noise to the data according to a predefined schedule\. Meanwhile, a*receiver distribution*pRp\_\{R\}, is created by convolving the output distribution with the same noise\.
4. 4\.KL divergence minimization:The loss function is the KL divergence from the sender distributionpSp\_\{S\}to the receiver distributionpRp\_\{R\}\. The neural network weights,𝝎\\bm\{\\omega\}, are updated by stochastic gradient descent to minimize this loss\.

This training procedure shows that if no noise is added topOp\_\{O\}, the network will learn to collapsepOp\_\{O\}onto𝒙\\bm\{x\}\. With noise added, however, the network learns to give a probability to all𝒙′\\bm\{x\}^\{\\prime\}relative to how likely they were to producepRp\_\{R\}\. The level of noise added is determined by an accuracy schedule,\{α0,…,αN−1\}\\\{\\alpha\_\{0\},\\dots,\\alpha\_\{N\-1\}\\\}\. The accuracy starts low and increases over time\.

𝜽i​=def​h​\(𝜽i−1,𝒚i−1,αi−1\)\\bm\{\\theta\}\_\{i\}\\overset\{\\text\{def\}\}\{=\}h\(\\bm\{\\theta\}\_\{i\-1\},\\bm\{y\}\_\{i\-1\},\\alpha\_\{i\-1\}\)…\\dots𝜽0\\bm\{\\theta\}\_\{0\}pOp\_\{O\}pRp\_\{R\}𝜽i\+1​=def​h​\(𝜽i,𝒚i,αi\)\\bm\{\\theta\}\_\{i\+1\}\\overset\{\\text\{def\}\}\{=\}h\(\\bm\{\\theta\}\_\{i\},\\bm\{y\}\_\{i\},\\alpha\_\{i\}\)…\\dots𝜽N\\bm\{\\theta\}\_\{N\}pSp\_\{S\}𝒙\\bm\{x\}Neural Net\(𝝎\)\(\\bm\{\\omega\}\)\+\+\+\+Sample𝒚i\\bm\{y\}\_\{i\}Noiseℒ\(𝜽\)=KL\[pS∥pR\]\\mathcal\{L\}\(\\bm\{\\theta\}\)=KL\\left\[p\_\{S\}\\\|p\_\{R\}\\right\]

Figure 2:Training process for Bayesian flow networks\.Figure[3](https://arxiv.org/html/2608.25163#S2.F3)represents the generative process of Bayesian flow networks\. The key difference from the training process is that the parameters are updated based on samples from the receiver distribution, rather than the sender distribution\. It unfolds as follows:

1. 1\.Initial parameterization:The process starts with parameters,𝜽0\\bm\{\\theta\}\_\{0\}, of a prior distribution\. For discrete data, these represent uniform probability\.
2. 2\.Neural network transformation:Similar to the training process, the neural network transforms the parameters𝜽i\\bm\{\\theta\}\_\{i\}to produce an output distributionpOp\_\{O\}\.
3. 3\.Noise injection:Noise is convolved with the output distribution to create the receiver distribution\.
4. 4\.Bayesian update:A sample from the receiver distribution is used to update𝜽i\\bm\{\\theta\}\_\{i\}using the Bayesian update functionhh\.
5. 5\.Iterate:Step 2\-4 is repeatedNNtimes, after which the parameters are fed into the neural network one last time to produce the finalpOp\_\{O\}, from which a sample is taken\. Note that the noise added topOp\_\{O\}to producepRp\_\{R\}follows the same accuracy schedule as during training\.

𝜽i​=def​h​\(𝜽i−1,𝒚i−1,αi−1\)\\bm\{\\theta\}\_\{i\}\\overset\{\\text\{def\}\}\{=\}h\(\\bm\{\\theta\}\_\{i\-1\},\\bm\{y\}\_\{i\-1\},\\alpha\_\{i\-1\}\)…\\dots𝜽0\\bm\{\\theta\}\_\{0\}pOp\_\{O\}pRp\_\{R\}𝜽i\+1​=def​h​\(𝜽i,𝒚i,αi\)\\bm\{\\theta\}\_\{i\+1\}\\overset\{\\text\{def\}\}\{=\}h\(\\bm\{\\theta\}\_\{i\},\\bm\{y\}\_\{i\},\\alpha\_\{i\}\)…\\dots𝜽N\\bm\{\\theta\}\_\{N\}\+\+Neural Net\(𝝎\)\(\\bm\{\\omega\}\)\+\+Sample𝒚i\\bm\{y\}\_\{i\}Noise

Figure 3:Generative process for Bayesian flow networks\.A comprehensive description of Bayesian Flow Networks is beyond the scope of this paper, but we aim to give the reader a clear understanding of how they differ from diffusion models\. We will now look at what𝜽\\bm\{\\theta\},hh,pOp\_\{O\},pRp\_\{R\}, andpSp\_\{S\}shown in Figure[2](https://arxiv.org/html/2608.25163#S2.F2)look like for categorical distributions\.

Consider data represented as aDDdimensional vector𝐱=\(x\(1\),…,x\(D\)\)∈\{1,…,A\}D\\mathbf\{x\}=\\left\(x^\{\(1\)\},\\dots,x^\{\(D\)\}\\right\)\\in\\left\\\{1,\\dots,A\\right\\\}^\{D\}, whereAAis the maximum size of the state space over𝒙\(d\)\\bm\{x\}^\{\(d\)\}, and\{1,…,A\}\\left\\\{1,\\dots,A\\right\\\}is the set of integers from 1 andAA\. We will model this as a categorical distribution\.

#### Input distribution

The input distribution defines the probability of the data given the parameters fed into the neural network, as shown in Figures[2](https://arxiv.org/html/2608.25163#S2.F2)and[3](https://arxiv.org/html/2608.25163#S2.F3)\. For discrete data, this distribution is modeled as a factorized categorical distribution with parameters𝜽=\(𝜽\(1\),…,𝜽\(D\)\)\\bm\{\\theta\}=\\left\(\\bm\{\\theta\}^\{\(1\)\},\\dots,\\bm\{\\theta\}^\{\(D\)\}\\right\), where each𝜽\(d\)\\bm\{\\theta\}^\{\(d\)\}includes parameters for a categorical distribution overAAcategories corresponding to variabledd\. Specifically,θa\(d\)\\theta\_\{a\}^\{\(d\)\}represents the probability of categoryaafor the variabledd:

pI​\(𝐱∣𝜽\)=∏d=1DpI​\(𝐱\(d\)∣𝜽\(d\)\)\.p\_\{I\}\(\\mathbf\{x\}\\mid\\bm\{\\theta\}\)=\\prod\_\{d=1\}^\{D\}p\_\{I\}\(\\mathbf\{x\}^\{\(d\)\}\\mid\\bm\{\\theta\}^\{\(d\)\}\)\.
Initially, the input distribution is uniform, meaning𝜽0=\[1A,…,1A\]\\bm\{\\theta\}\_\{0\}=\\left\[\\frac\{1\}\{A\},\\dots,\\frac\{1\}\{A\}\\right\]\.

#### Output distribution

Ψ𝝎​\(𝜽i,i\)\\Psi\_\{\\bm\{\\omega\}\}\(\\bm\{\\theta\}\_\{i\},i\)is a neural network model that takes as input aDD\-dimensional parameter vector𝜽i\\bm\{\\theta\}\_\{i\}, where each element are parameters of a categorical distribution\. The output is of the same type\. The output distribution for discrete data is defined based on the data𝐱\\mathbf\{x\}, model inputs𝜽i\\bm\{\\theta\}\_\{i\}, step counterii, and resulting model outputsΨ𝝎​\(𝜽i,i\)=\(Ψ𝝎\(1\)​\(𝜽i,i\),…,Ψ𝝎\(D\)​\(𝜽i,i\)\)∈ℝA×D\\Psi\_\{\\bm\{\\omega\}\}\\left\(\\bm\{\\theta\}\_\{i\},i\\right\)=\\left\(\\Psi\_\{\\bm\{\\omega\}\}^\{\(1\)\}\\left\(\\bm\{\\theta\}\_\{i\},i\\right\),\\dots,\\Psi\_\{\\bm\{\\omega\}\}^\{\(D\)\}\\left\(\\bm\{\\theta\}\_\{i\},i\\right\)\\right\)\\in\\mathbb\{R\}^\{A\\times D\}\. The network inputs𝜽i\\bm\{\\theta\}\_\{i\}represents the parameters of the factorized categorical distributionpI​\(𝐱∣𝜽i\)p\_\{I\}\(\\mathbf\{x\}\\mid\\bm\{\\theta\}\_\{i\}\), whileiiserves as an additional input that represents the process time\. The output distribution is thus defined as

pO​\(𝐱∣𝜽i,i\)=∏d=1DΨ𝝎\(d\)​\(𝜽i,i\)\.p\_\{O\}\(\\mathbf\{x\}\\mid\\bm\{\\theta\}\_\{i\},i\)=\\prod\_\{d=1\}^\{D\}\\Psi^\{\(d\)\}\_\{\\bm\{\\omega\}\}\(\\bm\{\\theta\}\_\{i\},i\)\.
Here,Ψ𝝎\(d\)​\(𝜽i,i\)\\Psi\_\{\\bm\{\\omega\}\}^\{\(d\)\}\(\\bm\{\\theta\}\_\{i\},i\)denotesAAcomponents of the network output corresponding to the parameters\(θ1\(d\),…,θA\(d\)\)\\left\(\\theta^\{\(d\)\}\_\{1\},\\dots,\\theta^\{\(d\)\}\_\{A\}\\right\)of the categorical distribution for thedd\-th observation\.

#### Sender distribution

A sample from the sender distribution is used to update the parameters𝜽\\bm\{\\theta\}\. For𝐲=\(y\(1\),…,y\(D\)\)∈𝒴D\\mathbf\{y\}=\\left\(y^\{\(1\)\},\\dots,y^\{\(D\)\}\\right\)\\in\\mathcal\{Y\}^\{D\}, the sender distribution is defined as

pS​\(𝐲∣𝐱;α\)=𝒩⁡\(𝐲∣α⁡\(A​𝐞𝐱−𝟏\),α​A​𝑰\),p\_\{S\}\\left\(\\mathbf\{y\}\\mid\\mathbf\{x\};\\alpha\\right\)=\\mathcal\{N\}\\left\(\\mathbf\{y\}\\mid\\alpha\\left\(A\\mathbf\{e\_\{x\}\}\-\\bm\{1\}\\right\),\\alpha A\\bm\{I\}\\right\),where𝐞𝐱\\mathbf\{e\_\{x\}\}is a unit vector of lengthAAand element𝐱\\mathbf\{x\}is11, also known as a one\-hot\-encoding\. The accuracy of these samples is controlled by an accuracy parameterα∈ℝ\+\\alpha\\in\\mathbb\{R\}^\{\+\}\. Whenα\\alphais low, the samples provide limited information about𝐱\\mathbf\{x\}\. Asα\\alphaincreases, the samples become increasingly informative about𝐱\\mathbf\{x\}\. Note that the value of𝒚i\\bm\{y\}\_\{i\}determines the amount of information, not the variance\. The reason for this is that𝒚i\\bm\{y\}\_\{i\}will play the role of the logits\. The value of𝔼⁡\(𝒚i\)\\mathbb\{E\}\(\\bm\{y\}\_\{i\}\)increases withαi\\alpha\_\{i\}, which means that we get a higher focus on the values where𝐞𝐱=1\\mathbf\{e\_\{x\}\}=1\.

#### Receiver distribution

The receiver distribution is defined according to the output distributionpOp\_\{O\}, andpSp\_\{S\}, this takes the form

pR​\(𝐲∣𝜽i;i,α\)=𝔼pO​\(𝐱′∣𝜽i;i\)​\[pS​\(𝐲\|𝐱′;α\)\],p\_\{R\}\(\\mathbf\{y\}\\mid\\bm\{\\theta\}\_\{i\};i,\\alpha\)=\\mathbb\{E\}\_\{p\_\{O\}\\left\(\\mathbf\{x\}^\{\\prime\}\\mid\\bm\{\\theta\}\_\{i\};i\\right\)\}\\left\[p\_\{S\}\\left\(\\mathbf\{y\}\|\\mathbf\{x\}^\{\\prime\};\\alpha\\right\)\\right\],
In essence, this integrates over all𝐱′∈\{1,…,A\}D\\mathbf\{x\}^\{\\prime\}\\in\\\{1,\\dots,A\\\}^\{D\}, considering the contribution of each possible𝐱′\\mathbf\{x\}^\{\\prime\}as weighted by its likelihood under the output distributionpO​\(𝐱\|𝜽i,i\)p\_\{O\}\(\\mathbf\{x\}\|\\bm\{\\theta\}\_\{i\},i\), effectively combines all potential sender distributions into a single receiver distribution\.

The Bayesian update function for discrete data introduced in Figure[1](https://arxiv.org/html/2608.25163#S2.F1)is given by

h⁡\(𝜽i,𝐲i\)=e𝐲i⊙𝜽i∑a=1Ae𝐲i,a​\(𝜽i\)a,h\\left\(\\bm\{\\theta\}\_\{i\},\\mathbf\{y\}\_\{i\}\\right\)=\\frac\{e^\{\\mathbf\{y\}\_\{i\}\}\\odot\\bm\{\\theta\}\_\{i\}\}\{\\sum\_\{a=1\}^\{A\}\{e^\{\\mathbf\{y\}\_\{i,a\}\}\\left\(\\bm\{\\theta\}\_\{i\}\\right\)\_\{a\}\}\},where⊙\\odotrefers to the Hadamard product\. For a detailed derivation of the update function, refer to the original work by[Graves et al\. \(2023\)](https://arxiv.org/html/2608.25163#bib.bib11)\.

At each step, the objective is to minimize the KL divergence from the sender distribution to the receiver distribution over all variables\. As the noise inpSp\_\{S\}approaches zero, the process is driven to the distribution that maximizes the likelihood of sampling data from the distributionpOp\_\{O\}\.[Graves et al\. \(2023\)](https://arxiv.org/html/2608.25163#bib.bib11)show that the loss function for annn\-step procedure at stepiiis:

LN\(𝝎;𝐱,𝐲i,𝜽i,i\)=N𝔼\[ln𝒩\(𝐲i∣αi\(A𝐞𝐱−1\),αiAI\)−ln\(∑a=1ApO\(a∣𝜽i,i\)𝒩\(𝐲i∣αi\(A𝐞a−1\),αiAI\)\)\]\.\\begin\{split\}L^\{N\}\(\\bm\{\\omega\};\\mathbf\{\\mathbf\{x\}\},\\mathbf\{y\}\_\{i\},\\bm\{\\theta\}\_\{i\},i\)=N\\,\\mathbb\{E\}\\Bigg\[\\ln\\mathcal\{N\}\\left\(\\mathbf\{y\}\_\{i\}\\mid\\alpha\_\{i\}\(A\\mathbf\{e\_\{x\}\}\-1\),\\alpha\_\{i\}AI\\right\)\\\\ \-\\ln\\left\(\\sum\_\{a=1\}^\{A\}p\_\{O\}\\left\(a\\mid\\bm\{\\theta\}\_\{i\},i\\right\)\\mathcal\{N\}\\left\(\\mathbf\{y\}\_\{i\}\\mid\\alpha\_\{i\}\(A\\mathbf\{e\}\_\{a\}\-1\),\\alpha\_\{i\}AI\\right\)\\right\)\\Bigg\]\.\\end\{split\}
Furthermore,[Graves et al\. \(2023\)](https://arxiv.org/html/2608.25163#bib.bib11)shows that if you letNNgo to infinity you get a continuous\-time loss function\. In that case, the generative process is not bound byNNused during training\. We omit the detailed derivation here, but[Graves et al\. \(2023\)](https://arxiv.org/html/2608.25163#bib.bib11)show that the continuous\-time loss works marginally better than the discrete\-time loss whenNNis large\. We will use the continuous\-time loss in our experiments\.

Algorithm 1Generating Discrete Samples With Bayesian Flow Networks1:

β⁡\(1\)∈ℝ\+\\beta\(1\)\\in\\mathbb\{R\}^\{\+\}, number of steps

n∈ℕn\\in\\mathbb\{N\}, number of categories

K∈ℕK\\in\\mathbb\{N\}
2:

𝜽←\(𝟏𝑲\)\\bm\{\\theta\}\\leftarrow\\left\(\\frac\{\\bm\{1\}\}\{\\bm\{K\}\}\\right\)
3:for

i=1​to​ni=1\\text\{ to \}ndo

4:

𝒌∼pO\(⋅∣𝜽,i\)\\bm\{k\}\\sim p\_\{O\}\(\\cdot\\mid\\bm\{\\theta\},i\)
5:

α←β​\(1\)​\(2​i−1n2\)\\alpha\\leftarrow\\beta\(1\)\(\\frac\{2i\-1\}\{n^\{2\}\}\)
6:

𝒚∼𝒩⁡\(α⁡\(K​𝒆𝒌−1\),α​K​𝐈\)\\bm\{y\}\\sim\\mathcal\{N\}\(\\alpha\(K\\bm\{e\_\{k\}\}\-1\),\\alpha K\\mathbf\{I\}\)
7:

𝜽′←e𝒚⊙𝜽\\bm\{\\theta\}^\{\\prime\}\\leftarrow e^\{\\bm\{y\}\}\\odot\\bm\{\\theta\}
8:

𝜽←θ′∑kθk′\\bm\{\\theta\}\\leftarrow\\frac\{\\theta^\{\\prime\}\}\{\\sum\_\{k\}\\theta^\{\\prime\}\_\{k\}\}
9:endfor

10:

𝒌∼pO\(⋅∣𝜽,1\)\\bm\{k\}\\sim p\_\{O\}\(\\cdot\\mid\\bm\{\\theta\},1\)

## 3Related Work

Several recent works have considered offline RL as a sequence modeling problem\. Here we will consider the two most prominent\.

### 3\.1Diffuser

Diffuser\([Janner et al\., 2022](https://arxiv.org/html/2608.25163#bib.bib3)\)models state\-action trajectories using an unconditional diffusion model\. At test time, the current state is imposed through inpainting, while gradients from a differentiable reward model guide sampling toward high\-return trajectories\. Terminal states can likewise be imposed through inpainting to obtain goal\-conditioned plans\. The authors further show that the inpainting technique can be used to condition desired end states, effectively allowing the model to solve planning problems it has not been specifically trained for\.

### 3\.2Decision Diffuser

The Decision Diffuser[Ajay et al\. \(2023\)](https://arxiv.org/html/2608.25163#bib.bib2)differs from Diffuser in two main ways: how it models actions, and how it conditions on rewards\. First, Decision Diffuser leverages an inverse dynamics model to capture the relationship between states and actions\. This inverse dynamics model estimates actions conditioned on states, effectively predicting the action that brought the environment from one state to the next\. This lets the diffusion model focus only on state sequences, rather than on sequences containing both state and action\. The authors show empirically that using an inverse dynamics model is advantageous in deterministic environments, but that the performance reduces to the same level as the Diffuser as more stochasticity is introduced into the environment\. Second, the Decision Diffuser conditions on return\-to\-go in a classifier\-free manner\. This means that the required return is fed into the model during training so that the model learns which sequences to associate with that return\. The desired return can again be fed into the model at test time when generating a sequence of future states\.

Native categorical diffusion models provide an alternative route to discrete generation\([Austin et al\., 2021](https://arxiv.org/html/2608.25163#bib.bib20);[Lou et al\., 2023](https://arxiv.org/html/2608.25163#bib.bib19)\)\. In the discrete experiments below, we therefore compare the categorical BFN with a uniform\-categorical diffusion model that predicts the clean trajectory and uses the exact reverse posterior\.

## 4Method

We propose a sequence\-generating approach to reinforcement learning based on Bayesian flow networks capable of planning in both discrete and continuous domains\. From now on we will refer to our method as BFN\-RL\. Like the Decision Diffuser\([Ajay et al\., 2023](https://arxiv.org/html/2608.25163#bib.bib2)\), BFN\-RL only models sequences ofstates\. Return is supplied as a network condition, whereas the current state is fixed by inpainting during sampling\. We then utilize a second inverse\-dynamics network to model the actions conditioned on the states\. In the context of diffusion models, this approach has been shown to provide superior performance compared to modeling state\-action pairs from a single model\([Janner et al\., 2022](https://arxiv.org/html/2608.25163#bib.bib3)\)\. We expect the same benefits when using BFNs as the generative model, but consider further examination of this hypothesis as future work\.

In our method, the BFN is trained to generate sequences conditioned on thereturn\. The return is the sum of all discounted future rewards and is therefore a measure of the quality of a sequence\. The network learns to model the distribution over sequences with both high and low returns\. At test time, the desired return biases generation toward trajectories with the corresponding cumulative reward\.

### 4\.1Condition on Return

There are two obvious ways we can condition on return: Either as done by the Diffuser\([Janner et al\., 2022](https://arxiv.org/html/2608.25163#bib.bib3)\)or as done by the Decision Diffuser\([Ajay et al\., 2023](https://arxiv.org/html/2608.25163#bib.bib2)\)\. Considering that the latter option is significantly easier to implement, showed better performance, and does not involve training an extra model, we opted to condition directly on return in a classifier\-free manner as was done in Decision Diffuser\([Ajay et al\., 2023](https://arxiv.org/html/2608.25163#bib.bib2)\)\. At each BFN step, the neural network receives the current parameters of the state distribution, factorized over trajectory time, together with the return and BFN time\. The observed current state is not supplied through a separate direct\-conditioning pathway\.

By conditioning directly on return during the generation process, the model learns to generate trajectories associated with different levels of cumulative reward\. At inference time, this allows the generation process to be biased toward higher\-return trajectories by conditioning on a desired return value\. Furthermore, adopting the Decision Diffuser\-style conditioning framework enables straightforward integration with existing sequence\-modeling architectures, including the temporal U\-net architecture used by[Ajay et al\. \(2023\)](https://arxiv.org/html/2608.25163#bib.bib2)\.

### 4\.2Conditioning on Current State

When Diffuser\([Janner et al\., 2022](https://arxiv.org/html/2608.25163#bib.bib3)\)and Decision Diffuser\([Ajay et al\., 2023](https://arxiv.org/html/2608.25163#bib.bib2)\)condition on the current state, they both apply an inpainting technique specific to diffusion models, see Eq\. \([3](https://arxiv.org/html/2608.25163#S2.E3)\)\. During the reverse diffusion process, the part of the sequence that is known \(the first state\), is replaced by the true value diffused to the appropriate amount for that step\. We have adopted a similar technique for BFNs\. For discrete data, we set the probabilities of the categorical distribution of the known parts of the object to the appropriate probability for that BFN step\. Algorithm[2](https://arxiv.org/html/2608.25163#alg2)shows this method implemented for discrete data\. The alterations to BFN sampling \(Algorithm 9 in[Graves et al\. \(2023\)](https://arxiv.org/html/2608.25163#bib.bib11)\) are shown in Algorithm[2](https://arxiv.org/html/2608.25163#alg2)\. For continuous data, the Bayesian update at each step is made similarly, see Algorithm[3](https://arxiv.org/html/2608.25163#alg3)in the Appendix\.

The operations involving the mask and conditioning values are those added to the regular BFN algorithm\. In Algorithm[2](https://arxiv.org/html/2608.25163#alg2), the mask𝒎\\bm\{m\}selects the observed variables, and𝒄\\bm\{c\}supplies their clean categorical values\. All reported experiments use this inpainting mechanism to fix the current state throughout sampling\.

For the categorical experiments, we retain the quadratic BFN accuracy scheduleβ⁡\(t\)=β⁡\(1\)​t2\\beta\(t\)=\\beta\(1\)t^\{2\}, but parameterize its endpoint byC=K​β​\(1\)C=K\\beta\(1\), whereKKis the number of categories\. This keeps the terminal correct\-versus\-incorrect message\-mean separation from changing merely because a representation uses more categories\. We useC=36C=36, 18, and 72 for Empty\-Random, DoorKey, and BlockedUnlockPickup, respectively\. The continuous\-time BFN objective is augmented by clean\-data cross entropy with weight 0\.5\. Reported BFN results use the categorical sampling procedure in Algorithm[2](https://arxiv.org/html/2608.25163#alg2), with 12 updates for Empty\-Random and 16 for DoorKey and BlockedUnlockPickup\. For categorical tasks, we optimizeℒBFN\+0\.5​ℒCE\\mathcal\{L\}\_\{\\mathrm\{BFN\}\}\+0\.5\\mathcal\{L\}\_\{\\mathrm\{CE\}\}, whereℒCE\\mathcal\{L\}\_\{\\mathrm\{CE\}\}is the cross\-entropy for predicting the clean category from the BFN input parameters\.

Algorithm 2Inpainting Conditioning for Discrete Random Variables1:

β⁡\(1\)∈ℝ\+\\beta\(1\)\\in\\mathbb\{R\}^\{\+\}, number of steps

n∈ℕn\\in\\mathbb\{N\}, number of categories

KK,mask𝒎\\bm\{m\}, condition𝒄\\bm\{c\}

2:

𝜽←𝟏𝑲\\bm\{\\theta\}\\leftarrow\\frac\{\\bm\{1\}\}\{\\bm\{K\}\}
3:for

i=1​to​ni=1\\text\{ to \}ndo

4:

t←i−1nt\\leftarrow\\frac\{i\-1\}\{n\}
5:

𝐤∼\\mathbf\{k\}\\simdiscrete\_output\_distribution\(

𝜽,t\\bm\{\\theta\},t\)

6:

α←β​\(1\)​\(2​i−1n2\)\\alpha\\leftarrow\\beta\(1\)\\left\(\\frac\{2i\-1\}\{n^\{2\}\}\\right\)
7:

𝐲∼𝒩⁡\(α⁡\(K​𝐞𝐤−𝟏\),α​K​𝑰\)\\mathbf\{y\}\\sim\\mathcal\{N\}\\left\(\\alpha\(K\\mathbf\{e\_\{k\}\}\-\\bm\{1\}\),\\alpha K\\bm\{I\}\\right\)
8:𝐲𝐜←α⁡\(K​𝐞𝐜−𝟏\)\\mathbf\{y\_\{c\}\}\\leftarrow\\alpha\(K\\mathbf\{e\_\{c\}\}\-\\bm\{1\}\)

9:𝐲←𝒎⊙𝐲𝐜\+\(1−𝒎\)⊙𝐲\\mathbf\{y\}\\leftarrow\\bm\{m\}\\odot\\mathbf\{y\_\{c\}\}\+\(1\-\\bm\{m\}\)\\odot\\mathbf\{y\}

10:

𝜽′←e𝐲⊙𝜽\\bm\{\\theta\}^\{\\prime\}\\leftarrow e^\{\\mathbf\{y\}\}\\odot\\bm\{\\theta\}
11:

𝜽=𝜽′∑k𝜽k′\\bm\{\\theta\}=\\frac\{\\bm\{\\theta\}^\{\\prime\}\}\{\\sum\_\{k\}\\bm\{\\theta\}^\{\\prime\}\_\{k\}\}
12:endfor

13:

𝐤∼\\mathbf\{k\}\\simdiscrete\_output\_distribution\(

𝜽,1\\bm\{\\theta\},1\)

## 5Experiments

We evaluate our method on two sets of tasks, one with discrete action and state space, and one with continuous action and state space\. In the discrete case, we use a grid world environment\. For the continuous case, we use the D4RL\([Fu et al\., 2020](https://arxiv.org/html/2608.25163#bib.bib1)\)datasets with the Gym\-Mujoco suite of environments\. We compare our method to the Decision Diffuser\([Ajay et al\., 2023](https://arxiv.org/html/2608.25163#bib.bib2)\)and other state of the art offline RL methods\. The discrete experiments test whether a categorical BFN can generate executable state plans without hand\-coded transition rules\. The continuous experiments test whether the same planning formulation remains competitive on standard offline RL benchmarks\.

### 5\.1Discrete Experiments

We evaluate BFN\-RL on three MiniGrid tasks, as well as FrozenLake and Sokoban\. In each environment, a categorical BFN generates state sequences and a learned inverse\-dynamics model maps consecutive states to actions from the environment’s native categorical action space\. The MiniGrid environments have seven possible actions\. We use the same temporal U\-Net family across all discrete tasks, with task\-specific widths, planning horizons, state representations, and conditioning inputs\.

Empty\-Random requires navigation to a randomly located goal\. DoorKey additionally requires the agent to collect a key and open a door, while BlockedUnlockPickup requires moving an obstructing object before collecting a target box\. Their training sets contain 2,000, 5,000, and 3,000 successful trajectories, respectively\. The DoorKey and BlockedUnlockPickup datasets include random perturbations followed by expert recovery\. We additionally evaluate slippery FrozenLake\-8×\\times8 navigation and two\-box Sokoban\-7×\\times7 using 30,000 and 10,000 successful trajectories, respectively\.

At each replanning step, the observed state is fixed through inpainting and one trajectory is sampled\. Empty\-Random uses a planning horizon of 16 and replans after every action\. DoorKey and BlockedUnlockPickup use a horizon of 32 and execute at most 16 actions from each sampled trajectory, replanning earlier if the observed state departs from the generated plan\. FrozenLake and Sokoban use horizons of 64 and 32, respectively, and replan after every action\. No hand\-coded transition projection or analytic inverse dynamics is used\.

The MiniGrid experiments condition generation on return, whereas FrozenLake and Sokoban use goal conditioning\. To isolate the effect of the generative model, the categorical\-diffusion baseline is matched to BFN\-RL in training data, temporal U\-Net architecture, state representation, conditioning signal, current\-state inpainting, inverse\-dynamics model, planning horizon, action\-execution protocol, and evaluation seeds\. The baseline uses uniform categorical corruption with a cosine signal\-survival schedule\.

Table 1:Matched discrete\-generator comparison\. Entries are mean±\\pmsample standard deviation across three training seeds\. Evaluation uses 100 episodes per seed except for FrozenLake, which uses 200\. Environment and planner seeds are matched between generators\.The diffusion baseline reaches the goal more quickly in Empty\-Random, while BFN\-RL has the higher return in DoorKey\. On BlockedUnlockPickup, variation between training seeds is larger than the difference in mean success\. Overall, BFN\-RL is a competitive categorical planner, although it does not uniformly outperform categorical diffusion\.

BFN\-RL and categorical diffusion perform similarly on FrozenLake\. On Sokoban, BFN\-RL has higher success for each of the three training seeds and improves the mean success rate by 9\.7 percentage points\. This identifies a harder discrete task on which the BFN generator provides a consistent advantage under the matched protocol\.

### 5\.2Continuous Control

The original Gaussian formulations of Diffuser\([Janner et al\., 2022](https://arxiv.org/html/2608.25163#bib.bib3)\)and Decision Diffuser\([Ajay et al\., 2023](https://arxiv.org/html/2608.25163#bib.bib2)\)are not natively categorical\. We evaluate BFN\-RL on D4RL MuJoCo to test whether the same planning framework remains competitive in continuous control\.

Table[2](https://arxiv.org/html/2608.25163#S5.T2)shows the performance of BFN\-RL compared to state\-of\-the\-art algorithms\. The table shows that BFN\-RL is competitive on most datasets\.

DatasetEnvironmentBCCQLIQLDTTTMOReLDiffuserDDBFN\-RLMed\-ExpertHalfCheetah55\.291\.686\.786\.89553\.379\.890\.697\.2±\\pm0\.1Med\-ExpertHopper52\.5105\.491\.5107\.6110\.0108\.7107\.2111\.8110\.3±\\pm0\.8Med\-ExpertWalker2d107\.5108\.8109\.6108\.1101\.995\.6108\.4108\.8106\.6±\\pm4\.9MediumHalfCheetah42\.644\.047\.442\.646\.942\.144\.249\.147\.8±\\pm0\.1MediumHopper52\.958\.566\.367\.661\.195\.458\.579\.372\.2±\\pm7\.1MediumWalker2d75\.372\.578\.374\.07977\.879\.782\.566\.9±\\pm5\.5Med\-ReplayHalfCheetah36\.645\.544\.236\.641\.940\.242\.239\.342\.5±\\pm0\.5Med\-ReplayHopper18\.19594\.782\.791\.593\.696\.810091\.2±\\pm6\.3Med\-ReplayWalker2d26\.077\.273\.966\.682\.649\.861\.27558\.5±\\pm4\.8Average51\.977\.67774\.778\.972\.975\.381\.877\.0Table 2:The table summarizes the test performance of BFN\-RL and various other methods for continuous control\. The results indicate that BFN\-RL can perform at a level comparable to the state of the art\. We report mean and standard error over 3 random seeds\. All numbers except for BFN\-RL are from[Ajay et al\. \(2023\)](https://arxiv.org/html/2608.25163#bib.bib2)\.

## 6Conclusion

In this work, we introduced a novel approach to reinforcement learning that leverages Bayesian flow networks\([Graves et al\., 2023](https://arxiv.org/html/2608.25163#bib.bib11)\)for sequence generation\. Our method is capable of planning in both discrete and continuous domains\. This shared planning framework also suggests a natural extension to joint continuous\-categorical environments, in which each trajectory variable is assigned the corresponding continuous or categorical BFN distribution within one model\. The current state is fixed by inpainting throughout sampling, while return conditioning avoids the need for an additional return classifier\.

Across the discrete tasks, BFN\-RL is competitive with matched categorical diffusion and is consistently stronger on Sokoban\. In continuous control, it remains competitive with established offline trajectory models\. Together, these results support BFNs as a common planning framework for categorical and continuous domains\.

Independent work on Guided\-BFNs reports strong results in continuous trajectory\-planning environments by supplementing conditional guidance with gradients from a learned reward model during sampling\([Li et al\., 2024](https://arxiv.org/html/2608.25163#bib.bib22)\)\. That work considers continuous state–action trajectories, whereas our categorical experiments address planning with discrete states and actions\.

## Declarations

This research was funded by internal funding from the Norwegian University of Science and Technology\.

## Appendix AHyperparameters

Here we present hyperparameters used in the experiments\. For the discrete experiments, we used:

- •Quadratic accuracy scheduleβ⁡\(t\)=β⁡\(1\)​t2\\beta\(t\)=\\beta\(1\)t^\{2\}, withK​β​\(1\)=36K\\beta\(1\)=36, 18, and 72 for Empty\-Random, DoorKey, and BlockedUnlockPickup, respectively\. The planning horizons are 16, 32, and 32\.
- •Learning rate3×10−43\\times 10^\{\-4\}, batch size 128, Adam optimizer\([Kingma and Ba, 2014](https://arxiv.org/html/2608.25163#bib.bib14)\), linear warmup for 200 updates on Empty\-Random and 300 updates otherwise, followed by cosine decay\.
- •A generic temporal U\-Net with base width 32 for Empty\-Random, 48 for DoorKey, and 48 or 96 in the BlockedUnlockPickup sweep\.
- •A learned inverse\-dynamics MLP with two hidden layers of 256 units\.
- •The continuous\-time categorical BFN objective is augmented by a cross\-entropy term with weight 0\.5\.
- •Categorical BFN sampling uses 12 updates for Empty\-Random and 16 for DoorKey and BlockedUnlockPickup, temperature 0\.7 for Empty\-Random and 0\.9 otherwise, and one generated candidate\.
- •Empty\-Random replans after each action\. For DoorKey and BlockedUnlockPickup, at most 16 actions are executed from a plan, with immediate replanning after a state disagreement\.
- •The matched categorical\-diffusion baseline uses cosine uniform corruption, clean\-state cross entropy, and 12 or 16 deterministic reverse updates, matching the BFN update count for each task\.
- •FrozenLake and Sokoban use horizons of 64 and 32, base widths of 64 and 96, batch sizes of 256 and 64, andK​β​\(1\)=72K\\beta\(1\)=72\. Both use cross\-entropy weight 0\.5, temperature 0\.9, 16 sampling updates, and one\-action replanning\.

Most hyperparameters and model architectures for continuous experiments that are not specific to Bayesian Flow Networks are similar to those used in the official Decision Diffuser implementation\.

Hyperparameters used for the continuous experiments:

- •The Inverse dynamics model is an MLP with two layers with 256 units and ReLU activations\.
- •ϵθ\\epsilon\_\{\\theta\}andfϕf\_\{\\phi\}are trained for2×1062\\times 10^\{6\}steps using the Adam optimiser\([Kingma and Ba, 2014](https://arxiv.org/html/2608.25163#bib.bib14)\)with a batch size of 256, a learning rate of5×10−55\\times 10^\{\-5\}\.
- •We use a planning horizonHHof 20\.
- •For testing we used an exponential moving average of the weights with decayα=0\.9995\\alpha=0\.9995
- •First\-order BFN solver of[Xue et al\. \(2024\)](https://arxiv.org/html/2608.25163#bib.bib21)using 10 sampling steps\.
- •Guidance weight 1\.2\.
- •Temperature 0\.5\.
- •σ1=0\.01\\sigma\_\{1\}=0\.01\.

## Appendix BAlgorithms

Algorithm[3](https://arxiv.org/html/2608.25163#alg3)gives the continuous\-variable counterpart of the inpainting procedure in Algorithm[2](https://arxiv.org/html/2608.25163#alg2)\.

Algorithm 3Inpainting Conditioning for Continuous Random Variables1:

σ1∈ℝ\+\\sigma\_\{1\}\\in\\mathbb\{R\}^\{\+\}, number of steps

n∈ℕn\\in\\mathbb\{N\},mask𝒎\\bm\{m\}, condition𝒄\\bm\{c\}

2:

𝝁←𝟎\\bm\{\\mu\}\\leftarrow\\bm\{0\}
3:

ρ←1\\rho\\leftarrow 1
4:for

i=1​to​ni=1\\text\{ to \}ndo

5:

t←i−1nt\\leftarrow\\frac\{i\-1\}\{n\}
6:

𝐱^​\(𝜽,t\)←\\mathbf\{\\hat\{x\}\}\(\\bm\{\\theta\},t\)\\leftarrowcts\_output\_distribution\(

𝝁,t,1−σ12\\bm\{\\mu\},t,1\-\\sigma\_\{1\}^\{2\}\)

7:

α←σ1−2i/n\(1−σ12/n\)\\alpha\\leftarrow\\sigma\_\{1\}^\{\-2i/n\}\\left\(1\-\\sigma\_\{1\}^\{2/n\}\\right\)
8:

𝐲∼𝒩⁡\(𝐱^​\(𝜽,t\),α−1​𝑰\)\\mathbf\{y\}\\sim\\mathcal\{N\}\\left\(\\mathbf\{\\hat\{x\}\}\(\\bm\{\\theta\},t\),\\alpha^\{\-1\}\\bm\{I\}\\right\)
9:𝐲𝐜←\(1−σ12​t\)​𝒄\\mathbf\{y\_\{c\}\}\\leftarrow\(1\-\\sigma\_\{1\}^\{2t\}\)\\bm\{c\}

10:𝐲←𝒎⊙𝐲𝐜\+\(𝟏−𝒎\)⊙𝐲\\mathbf\{y\}\\leftarrow\\bm\{m\}\\odot\\mathbf\{y\_\{c\}\}\+\(\\bm\{1\}\-\\bm\{m\}\)\\odot\\mathbf\{y\}

11:

𝝁←ρ​𝝁\+α​𝐲ρ\+α\\bm\{\\mu\}\\leftarrow\\frac\{\\rho\\bm\{\\mu\}\+\\alpha\\mathbf\{y\}\}\{\\rho\+\\alpha\}
12:

ρ←ρ\+α\\rho\\leftarrow\\rho\+\\alpha
13:endfor

14:

𝒙^​\(𝜽,1\)←\\bm\{\\hat\{x\}\}\(\\bm\{\\theta\},1\)\\leftarrowcts\_output\_distribution\(

𝝁,1,1−σ12\\bm\{\\mu\},1,1\-\\sigma\_\{1\}^\{2\}\)

Hereρ\\rhois the precision of the continuous BFN input distribution; it starts at the unit\-precision prior and accumulates the message precisionsα\\alpha\.

## References

- Agarwalet al\.\(2020\)R\. Agarwal, D\. Schuurmans, and M\. NorouziAn optimistic perspective on offline reinforcement learning\.InInternational Conference on Machine Learning \(ICML\),Vol\.119,pp\. 104–114\.External Links:[Link](https://proceedings.mlr.press/v119/agarwal20c.html)Cited by:[§1](https://arxiv.org/html/2608.25163#S1.p1.1)\.
- Ajayet al\.\(2023\)A\. Ajay, Y\. Du, A\. Gupta, J\. B\. Tenenbaum, T\. S\. Jaakkola, and P\. AgrawalIs conditional generative modeling all you need for decision making?\.InInternational Conference on Learning Representations \(ICLR\),External Links:[Link](https://openreview.net/forum?id=sP1fo2K9DFG)Cited by:[§1](https://arxiv.org/html/2608.25163#S1.p1.1),[§1](https://arxiv.org/html/2608.25163#S1.p2.1),[§2\.2](https://arxiv.org/html/2608.25163#S2.SS2.p1.1),[§2\.3](https://arxiv.org/html/2608.25163#S2.SS3.p1.1),[§3\.2](https://arxiv.org/html/2608.25163#S3.SS2.p1.1),[§4\.1](https://arxiv.org/html/2608.25163#S4.SS1.p1.1),[§4\.1](https://arxiv.org/html/2608.25163#S4.SS1.p2.1),[§4\.2](https://arxiv.org/html/2608.25163#S4.SS2.p1.1),[§4](https://arxiv.org/html/2608.25163#S4.p1.1),[§5\.2](https://arxiv.org/html/2608.25163#S5.SS2.p1.1),[Table 2](https://arxiv.org/html/2608.25163#S5.T2),[§5](https://arxiv.org/html/2608.25163#S5.p1.1)\.
- Austinet al\.\(2021\)J\. Austin, D\. D\. Johnson, J\. Ho, D\. Tarlow, and R\. van den BergStructured denoising diffusion models in discrete state\-spaces\.arXiv preprint arXiv:2107\.03006\.External Links:[Link](https://arxiv.org/abs/2107.03006)Cited by:[§1](https://arxiv.org/html/2608.25163#S1.p2.1),[§2\.5](https://arxiv.org/html/2608.25163#S2.SS5.p1.1),[§3\.2](https://arxiv.org/html/2608.25163#S3.SS2.p2.1)\.
- Chenet al\.\(2021\)L\. Chen, K\. Lu, A\. Rajeswaran, K\. Lee, A\. Grover, M\. Laskin, P\. Abbeel, A\. Srinivas, and I\. MordatchDecision transformer: reinforcement learning via sequence modeling\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Vol\.34,pp\. 15084–15097\.External Links:[Link](https://proceedings.neurips.cc/paper/2021/hash/7f489f642a0ddb10272b5c31057f0663-Abstract.html)Cited by:[§1](https://arxiv.org/html/2608.25163#S1.p1.1)\.
- Dhariwal and Nichol \(2021\)P\. Dhariwal and A\. NicholDiffusion models beat GANs on image synthesis\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Vol\.34,pp\. 8780–8794\.External Links:[Link](https://proceedings.neurips.cc/paper_files/paper/2021/hash/49ad23d1ec9fa4bd8d77d02681df5cfa-Abstract.html)Cited by:[§2\.4](https://arxiv.org/html/2608.25163#S2.SS4.p1.1)\.
- Fuet al\.\(2020\)J\. Fu, A\. Kumar, O\. Nachum, G\. Tucker, and S\. LevineD4RL: datasets for deep data\-driven reinforcement learning\.arXiv preprint arXiv:2004\.07219\.External Links:[Link](https://arxiv.org/abs/2004.07219)Cited by:[§5](https://arxiv.org/html/2608.25163#S5.p1.1)\.
- Graveset al\.\(2023\)A\. Graves, R\. K\. Srivastava, T\. Atkinson, and F\. GomezBayesian flow networks\.arXiv preprint arXiv:2308\.07037\.External Links:[Link](https://arxiv.org/abs/2308.07037)Cited by:[1st item](https://arxiv.org/html/2608.25163#S1.I1.i1.p1.1),[§1](https://arxiv.org/html/2608.25163#S1.p3.1),[§2\.5](https://arxiv.org/html/2608.25163#S2.SS5.SSS0.Px4.p3.2),[§2\.5](https://arxiv.org/html/2608.25163#S2.SS5.SSS0.Px4.p4.1),[§2\.5](https://arxiv.org/html/2608.25163#S2.SS5.SSS0.Px4.p6.1),[§2\.5](https://arxiv.org/html/2608.25163#S2.SS5.p1.1),[§4\.2](https://arxiv.org/html/2608.25163#S4.SS2.p1.1),[§6](https://arxiv.org/html/2608.25163#S6.p1.1)\.
- Hoet al\.\(2020\)J\. Ho, A\. Jain, and P\. AbbeelDenoising diffusion probabilistic models\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Vol\.33,pp\. 6840–6851\.External Links:[Link](https://proceedings.neurips.cc/paper/2020/hash/4c5bcfec8584af0d967f1ab10179ca4b-Abstract.html)Cited by:[§1](https://arxiv.org/html/2608.25163#S1.p2.1),[§2\.3](https://arxiv.org/html/2608.25163#S2.SS3.p1.1),[§2\.3](https://arxiv.org/html/2608.25163#S2.SS3.p3.2),[§2\.3](https://arxiv.org/html/2608.25163#S2.SS3.p3.3)\.
- Hoet al\.\(2022\)J\. Ho, T\. Salimans, A\. Gritsenko, W\. Chan, M\. Norouzi, and D\. J\. FleetVideo diffusion models\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Vol\.35,pp\. 8633–8646\.External Links:[Link](https://proceedings.neurips.cc/paper_files/paper/2022/hash/39235c56aef13fb05a6adc95eb9d8d66-Abstract-Conference.html)Cited by:[§2\.3](https://arxiv.org/html/2608.25163#S2.SS3.p1.1)\.
- Ho and Salimans \(2021\)J\. Ho and T\. SalimansClassifier\-free diffusion guidance\.InNeurIPS Workshop on Deep Generative Models and Downstream Applications,External Links:[Link](https://openreview.net/forum?id=qw8AKxfYbI)Cited by:[§2\.4](https://arxiv.org/html/2608.25163#S2.SS4.p2.1),[§2\.4](https://arxiv.org/html/2608.25163#S2.SS4.p2.2)\.
- Janneret al\.\(2022\)M\. Janner, Y\. Du, J\. B\. Tenenbaum, and S\. LevinePlanning with diffusion for flexible behavior synthesis\.InInternational Conference on Machine Learning \(ICML\),pp\. 9902–9915\.External Links:[Link](https://proceedings.mlr.press/v162/janner22a.html)Cited by:[§1](https://arxiv.org/html/2608.25163#S1.p1.1),[§1](https://arxiv.org/html/2608.25163#S1.p2.1),[§2\.2](https://arxiv.org/html/2608.25163#S2.SS2.p1.1),[§2\.3](https://arxiv.org/html/2608.25163#S2.SS3.p1.1),[§3\.1](https://arxiv.org/html/2608.25163#S3.SS1.p1.1),[§4\.1](https://arxiv.org/html/2608.25163#S4.SS1.p1.1),[§4\.2](https://arxiv.org/html/2608.25163#S4.SS2.p1.1),[§4](https://arxiv.org/html/2608.25163#S4.p1.1),[§5\.2](https://arxiv.org/html/2608.25163#S5.SS2.p1.1)\.
- Janneret al\.\(2021\)M\. Janner, Q\. Li, and S\. LevineOffline reinforcement learning as one big sequence modeling problem\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Vol\.34,pp\. 1273–1286\.External Links:[Link](https://proceedings.neurips.cc/paper_files/paper/2021/hash/099fe6b0b444c23836c4a5d07346082b-Abstract.html)Cited by:[§1](https://arxiv.org/html/2608.25163#S1.p1.1),[§2\.2](https://arxiv.org/html/2608.25163#S2.SS2.p1.1)\.
- Kingma and Ba \(2014\)D\. P\. Kingma and J\. BaAdam: a method for stochastic optimization\.arXiv preprint arXiv:1412\.6980\.External Links:[Link](https://arxiv.org/abs/1412.6980)Cited by:[2nd item](https://arxiv.org/html/2608.25163#A1.I1.i2.p1.1),[2nd item](https://arxiv.org/html/2608.25163#A1.I2.i2.p1.1)\.
- Kumaret al\.\(2020\)A\. Kumar, A\. Zhou, G\. Tucker, and S\. LevineConservative Q\-learning for offline reinforcement learning\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Vol\.33,pp\. 1179–1191\.External Links:[Link](https://proceedings.neurips.cc/paper/2020/hash/0d2b2061826a5df3221116a5085a6052-Abstract.html)Cited by:[§2\.2](https://arxiv.org/html/2608.25163#S2.SS2.p1.1)\.
- Levineet al\.\(2020\)S\. Levine, A\. Kumar, G\. Tucker, and J\. FuOffline reinforcement learning: tutorial, review, and perspectives on open problems\.arXiv preprint arXiv:2005\.01643\.External Links:[Link](https://arxiv.org/abs/2005.01643)Cited by:[§1](https://arxiv.org/html/2608.25163#S1.p1.1),[§2\.2](https://arxiv.org/html/2608.25163#S2.SS2.p1.1)\.
- Liet al\.\(2024\)K\. Li, Z\. Deng, and D\. WangGuided\-BFNs: towards visualizing and understanding bayesian flow networks in the context of trajectory planning\.Unpublished\.External Links:[Link](https://openreview.net/forum?id=JRXcvEg3OB)Cited by:[§6](https://arxiv.org/html/2608.25163#S6.p3.1)\.
- Louet al\.\(2023\)A\. Lou, C\. Meng, and S\. ErmonDiscrete diffusion modeling by estimating the ratios of the data distribution\.arXiv preprint arXiv:2310\.16834\.External Links:[Link](https://arxiv.org/abs/2310.16834)Cited by:[§1](https://arxiv.org/html/2608.25163#S1.p2.1),[§2\.5](https://arxiv.org/html/2608.25163#S2.SS5.p1.1),[§3\.2](https://arxiv.org/html/2608.25163#S3.SS2.p2.1)\.
- Lugmayret al\.\(2022\)A\. Lugmayr, M\. Danelljan, A\. Romero, F\. Yu, R\. Timofte, and L\. Van GoolRePaint: inpainting using denoising diffusion probabilistic models\.InIEEE/CVF Conference on Computer Vision and Pattern Recognition \(CVPR\),pp\. 11461–11471\.External Links:[Link](https://openaccess.thecvf.com/content/CVPR2022/html/Lugmayr_RePaint_Inpainting_Using_Denoising_Diffusion_Probabilistic_Models_CVPR_2022_paper.html)Cited by:[§2\.4](https://arxiv.org/html/2608.25163#S2.SS4.p3.1)\.
- Luo and Hu \(2021\)S\. Luo and W\. HuDiffusion probabilistic models for 3d point cloud generation\.InIEEE/CVF Conference on Computer Vision and Pattern Recognition \(CVPR\),pp\. 2837–2845\.External Links:[Link](https://openaccess.thecvf.com/content/CVPR2021/html/Luo_Diffusion_Probabilistic_Models_for_3D_Point_Cloud_Generation_CVPR_2021_paper.html)Cited by:[§2\.3](https://arxiv.org/html/2608.25163#S2.SS3.p1.1)\.
- Peterset al\.\(2010\)J\. Peters, K\. Mulling, and Y\. AltunRelative entropy policy search\.InAAAI Conference on Artificial Intelligence,Vol\.24,pp\. 1607–1612\.External Links:[Link](https://ojs.aaai.org/index.php/AAAI/article/view/7727)Cited by:[§2\.2](https://arxiv.org/html/2608.25163#S2.SS2.p1.1)\.
- Sutton and Barto \(2018\)R\. S\. Sutton and A\. G\. BartoReinforcement learning: an introduction\.MIT Press\.External Links:[Link](http://incompleteideas.net/book/the-book-2nd.html)Cited by:[§1](https://arxiv.org/html/2608.25163#S1.p1.1),[§2\.1](https://arxiv.org/html/2608.25163#S2.SS1.p1.1)\.
- Xueet al\.\(2024\)K\. Xue, Y\. Zhou, S\. Nie, X\. Min, X\. Zhang, J\. Zhou, and C\. LiUnifying bayesian flow networks and diffusion models through stochastic differential equations\.InInternational Conference on Machine Learning \(ICML\),Vol\.235,pp\. 55656–55681\.External Links:[Link](https://proceedings.mlr.press/v235/xue24d.html)Cited by:[5th item](https://arxiv.org/html/2608.25163#A1.I2.i5.p1.1)\.

Similar Articles

Path-Coupled Bellman Flows for Distributional Reinforcement Learning

arXiv cs.LG

This paper introduces Path-Coupled Bellman Flows (PCBF), a continuous-time distributional reinforcement learning method that uses flow matching to model return distributions without heuristic projections. It addresses boundary mismatch and high-variance issues in previous flow-based approaches by coupling current and successor return flows through shared base noise.