Paper deep dive
Bayesian Flow Networks for Offline Trajectory Planning
Ludvig Killingberg, Helge Langseth
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 94%
Last extracted: 8/27/2026, 4:49:30 AM
Summary
The paper introduces BFN-RL, a unified generative modeling framework for offline reinforcement learning based on Bayesian Flow Networks (BFNs). Unlike diffusion models that denoise data instances, BFNs iteratively evolve distribution parameters, allowing native handling of both discrete and continuous trajectory spaces. The framework combines a categorical BFN state-sequence model with a learned inverse-dynamics model to generate effective trajectories in both discrete planning and continuous control tasks.
Entities (6)
Relation Signals (6)
BFN-RL → uses → Bayesian Flow Networks
confidence 98% · We present BFN-RL, a unified generative modeling framework for offline RL based on Bayesian Flow Networks (BFNs).
BFN-RL → appliedto → Offline Reinforcement Learning
confidence 97% · We present BFN-RL, a unified generative modeling framework for offline RL
BFN-RL → handles → discrete and continuous trajectory spaces
confidence 95% · BFN-RL natively models both discrete and continuous trajectory spaces within a single probabilistic formulation.
BFN-RL → consistsof → Categorical Planner
confidence 93% · The categorical planner generates future state sequences
Categorical Planner → uses → Inverse Dynamics Model
confidence 92% · a learned inverse-dynamics model converts consecutive generated states into actions.
Bayesian Flow Networks → comparedto → Denoising Diffusion Probabilistic Models
confidence 90% · Unlike diffusion models that iteratively denoise corrupted data instances, BFNs operate by updating the parameters of an input distribution
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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.
Tags
Links
- Source: https://arxiv.org/abs/2608.25163v1
- Canonical: https://arxiv.org/abs/2608.25163v1
Trouble viewing inline? Open PDF directly →
Full Text
54,269 characters extracted from source content.
Expand or collapse full text
Bayesian Flow Networks for Offline Trajectory Planning Ludvig Killingberg Affiliation: Norwegian University of Science and Technology Email: ludvig.killingberg@ntnu.no Helge Langseth Affiliation: 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. Preprint 1 Introduction Offline reinforcement learning (RL) (Sutton and Barto, 2018; Levine et al., 2020) 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; Chen et al., 2021) 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; Levine et al., 2020). 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; Ajay et al., 2023). Despite their empirical success, contemporary sequence-modeling approaches rely almost exclusively on Gaussian Denoising Diffusion Probabilistic Models (DDPMs) (Ho et al., 2020; Janner et al., 2022; Ajay et al., 2023). 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; Lou et al., 2023). 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). 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). 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). • 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. 2 Preliminaries 2.1 Reinforcement Learning Reinforcement learning (RL) is a framework for learning to make decisions in an environment (Sutton and Barto, 2018). The interactions with the environment are modeled as a Markov decision process (MDP), which is a tuple (,,,ℛ,γ)(S,A,P,R,γ), where S is the state space, A is the action space, P is the transition function, ℛR is the reward function, and γ is the discount factor. In an environment where the agent performs the action a∈a in state s∈s , the next state, s′∈s is sampled from (s,a)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. Let rtr_t denote the reward received at time t, and let Rt=∑i=0∞γirt+iR_t= _i=0^∞γ^ir_t+i be the discounted cumulative reward obtained starting from time t. Now, one goal of RL is to learn a policy π:→π:S that maximizes the expected return [Rt]E[R_t], where the expectation is taken over the uncertainty defined by both the transition function and the stochastic strategy π. 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.2 Offline 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). 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) or using a conservative value function (Kumar et al., 2020). 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; Janner et al., 2021). 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; Ajay et al., 2023). 2.3 Denoising Diffusion Probabilistic Models Since the current state of the art in this domain (Janner et al., 2022; Ajay et al., 2023) rely on diffusion models as the generative model, we will give a brief introduction to denoising diffusion probabilistic models (DDPMs) (Ho et al., 2020). This will also serve as a backdrop for our discussion of Bayesian flow networks, which follows in Section 2.5. DDPMs are a type of generative model inspired by non-equilibrium thermodynamics. The model is defined by a forward process that slowly adds Gaussian noise to data, and its reverse, 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; Luo and Hu, 2021). Given data 0∼q()x_0 q(x), we define the forward process to produce a sequence of noisy samples 1,…,Kx_1,…,x_K, q(k|k−1)=(k,1−βkk−1,βk).q(x_k|x_k-1)=N(x_k; 1- _kx_k-1, _k I). where βi∈(0,1)i=1K\ _i∈(0,1)\_i=1^K is a carefully chosen variance schedule. A nice property of the forward process is that we can directly sample kx_k at any step i, because the distribution q(k|0)q(x_k|x_0) can be derived using the property that a sum of uncorrelated normally distributed random variables is normally distributed. Let ak=1−βka_k=1- _k and a¯k=∏i=1kak a_k= _i=1^ka_k, then q(k|0)=(k,a¯k0,(1−a¯k)).q(x_k|x_0)=N (x_k; a_k\,x_0, (1- a_k ) I ). Note also that q(k−1|k,0)=(k−1,~(k,0),β~k),q(x_k-1|x_k,x_0)=N (x_k-1; μ(x_k,x_0), β_k I ), (1) where ~(k,0)=a¯k−1βk1−a¯k0+ak(1−a¯k−1)1−a¯kk,β~k=βk1−a¯k−11−a¯k. μ(x_k,x_0)= a_k-1 _k1- a_kx_0+ a_k(1- a_k-1)1- a_kx_k, β_k= _k 1- a_k-11- 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 from q(k−1|k)q(x_k-1|x_k). Let pθ(k−1|k)p_θ(x_k-1|x_k) be a parameterized approximation of q(k−1|k)q(x_k-1|x_k). This means that we define a neural network model with trainable parameters θ that outputs θ(k,k) μ_θ(x_k,k) and Σθ(k,k) _θ(x_k,k) so that pθ(k−1|k)=(k−1,θ(k,k),Σθ(k,k)).p_θ(x_k-1|x_k)=N (x_k-1; μ_θ(x_k,k), _θ(x_k,k) ). Ho et al. (2020) chose to fix the variance term θ(k,k) _θ(x_k,k) as a constant σk2=β~k _k^2= β_k, see Eq. (2). We therefore only look at how θ(k,k) μ_θ(x_k,k) is estimated. First, we consider the identity ~k(k,0)=1ak(k−1−ak1−a¯kϵk), μ_k(x_k,x_0)= 1 a_k (x_k- 1-a_k 1- a_k ε_k ), where ϵk∼(0,) ε_k (0, I); cf. Eqs. (1) and (2). Since kx_k is known during training, we can choose to predict ϵk ε_k, rather than ~k μ_k directly. Empirically, this has shown better results. Let us define ϵθ(,k) ε_θ(x,k) as a model that predicts the noise, ϵk ε_k. This means that we can define θ(k,k)=1ak(k−1−ak1−a¯kϵθ(k,k)) μ_θ(x_k,k)= 1 a_k (x_k- 1-a_k 1- a_k ε_θ(x_k,k) ). Ho et al. (2020) derive the following loss function to minimize the difference between θ μ_θ and ~ μ: L(θ)=k∼[1,K],0,ϵk[βk22σk2ak(1−a¯k)‖ϵk−ϵθ(k,k)‖2].L(θ)=E_k [1,K],x_0, ε_k [ _k^22σ^2_ka_k(1- a_k) ε_k- ε_θ(x_k,k) ^2 ]. 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(θ)=E_k [1,K],x_0, _k _k- _θ(x_k,k) ^2. 2.4 Guided Diffusion We will discuss three ways diffusion models can condition on variables. The first, classifier-guided diffusion (Dhariwal and Nichol, 2021), takes as its starting point that we have a trained probabilistic classifier that classifies objects kx_k during the reverse process. We want to use this to produce an object 0x_0 that is classified as belonging to a predefined class y. The approach uses the gradients of this classifier’s allocated log-likelihood to the class y wrt. kx_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 ϵ¯θ ε_θ, guided by a classifier h(y|k,k)h(y|x_k,k) meant to estimate the probability that the noisy datapoint kx_k belongs to class y, would assume the following form: ϵ¯θ(k,k,y)=ϵθ(k,k)−wσk∇klogh(y|k,k), ε_θ(x_k,k,y)= ε_θ(x_k,k)-w _k _x_k h(y|x_k,k), where w is a hyper-parameter controlling the strength of the guidance. Secondly, classifier-free guidance (Ho and Salimans, 2021), 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), ε(x_k,k,y)=(w+1) ε_θ(x_k,k,y)-w ε_θ(x_k,k), where we again use w to 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). Finally, we can also employ inpainting (Lugmayr et al., 2022) to condition on partial observations. In the context of image generation, this implies conditioning on some pixels within the image. Consider an image 0x_0 divided into known pixels 0knownx_0^known and unknown pixels 0unknownx_0^unknown, and a mask m defining which pixels are known. During the reverse process, we define k−1known=a¯k0+1−a¯k,∼(,),x_k-1^known= a_kx_0+ 1- a_k z, z ( 0,I), i.e., k−1knownx_k-1^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 step k, k−1unknownx_k-1^unknown, are computed in standard fashion: k−1unknown=1ak(k−βk1−a¯kϵθ(k,k))+σk,∼(,).x_k-1^unknown= 1 a_k (x_k- _k 1- a_k ε_θ(x_k,k) )+ _k z, z ( 0,I). Finally, we have: k−1=⊙k−1known+(1−)⊙k−1unknown,x_k-1= m _k-1^known+(1- m) _k-1^unknown, (3) where ⊙ is elementwise multiplication. 2.5 Bayesian Flow Networks While there are variations of diffusion models that model discrete data (Lou et al., 2023; Austin et al., 2021), these models are not considered state of the art when it comes to generating high-quality discrete data (Graves et al., 2023). In an attempt to remedy this shortcoming, Graves et al. (2023) introduced Bayesian 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 evolves distribution 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), and are therefore a more natural choice than diffusion models for planning in discrete state spaces. Figure 1 illustrates the idea of Bayesian flow networks. At each step i, the parameters of a distribution (i θ_i) are updated with noisy samples from the data, i y_i. The level of added noise decreases for each step and is at step i dictated by the accuracy αi _i. The parameters i θ_i are defined as the Bayesian update of the parameters at the previous step, i−1 θ_i-1, with observation noise parameterized by αi _i using a predetermined update rule h(⋅)h(·). 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=defh(i−1,i−1,αi−1) θ_i def=h( θ_i-1, y_i-1, _i-1)…0 θ_0i+1=defh(i,i,αi) θ_i+1 def=h( θ_i, y_i, _i)…N θ_N Figure 1: Generative process for Bayesian flow networks. Figure 2 represents the training process of Bayesian Flow Networks. The aim is to iteratively update the parameters of a distribution, θ, beginning with a prior distribution 0 θ_0, so that eventually, sampling once from this distribution mirrors sampling from the data it is trained on. In this process, accuracy αi _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. Generate θi θ_i: Given a datapoint x, the parameters 0 θ_0 can be updated i times with noisy data to create parameters i θ_i. We will later see that accuracies are additive and that i θ_i can be generated in a single step with accuracy ∑j=0i−1αj _j=0^i-1 _j. 2. Neural network transformation: The parameters i θ_i are passed through a neural network, with weights ω, which outputs the parameters of a new distribution. This is referred to as the output distribution pOp_O. 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. KL divergence minimization: The loss function is the KL divergence from the sender distribution pSp_S to the receiver distribution pRp_R. The neural network weights, ω, are updated by stochastic gradient descent to minimize this loss. This training procedure shows that if no noise is added to pOp_O, the network will learn to collapse pOp_O onto x. With noise added, however, the network learns to give a probability to all ′ x relative to how likely they were to produce pRp_R. The level of noise added is determined by an accuracy schedule, α0,…,αN−1\ _0,…, _N-1\. The accuracy starts low and increases over time. i=defh(i−1,i−1,αi−1) θ_i def=h( θ_i-1, y_i-1, _i-1)…0 θ_0pOp_OpRp_Ri+1=defh(i,i,αi) θ_i+1 def=h( θ_i, y_i, _i)…N θ_NpSp_S xNeural Net()( ω)++++Sample i y_iNoiseℒ()=KL[pS∥pR]L( θ)=KL [p_S\|p_R ] Figure 2: Training process for Bayesian flow networks. Figure 3 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. Initial parameterization: The process starts with parameters, 0 θ_0, of a prior distribution. For discrete data, these represent uniform probability. 2. Neural network transformation: Similar to the training process, the neural network transforms the parameters i θ_i to produce an output distribution pOp_O. 3. Noise injection: Noise is convolved with the output distribution to create the receiver distribution. 4. Bayesian update: A sample from the receiver distribution is used to update i θ_i using the Bayesian update function h. 5. Iterate: Step 2-4 is repeated N times, after which the parameters are fed into the neural network one last time to produce the final pOp_O, from which a sample is taken. Note that the noise added to pOp_O to produce pRp_R follows the same accuracy schedule as during training. i=defh(i−1,i−1,αi−1) θ_i def=h( θ_i-1, y_i-1, _i-1)…0 θ_0pOp_OpRp_Ri+1=defh(i,i,αi) θ_i+1 def=h( θ_i, y_i, _i)…N θ_N++Neural Net()( ω)++Sample i y_iNoise 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 θ, h, pOp_O, pRp_R, and pSp_S shown in Figure 2 look like for categorical distributions. Consider data represented as a D dimensional vector =(x(1),…,x(D))∈1,…,ADx= (x^(1),…,x^(D) )∈ \1,…,A \^D, where A is the maximum size of the state space over (d) x^(d), and 1,…,A \1,…,A \ is the set of integers from 1 and A. 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 and 3. For discrete data, this distribution is modeled as a factorized categorical distribution with parameters =((1),…,(D)) θ= ( θ^(1),…, θ^(D) ), where each (d) θ^(d) includes parameters for a categorical distribution over A categories corresponding to variable d. Specifically, θa(d) _a^(d) represents the probability of category a for the variable d: pI(∣)=∏d=1DpI((d)∣(d)).p_I(x θ)= _d=1^Dp_I(x^(d) θ^(d)). Initially, the input distribution is uniform, meaning 0=[1A,…,1A] θ_0= [ 1A,…, 1A ]. Output distribution Ψ(i,i) _ ω( θ_i,i) is a neural network model that takes as input a D-dimensional parameter vector i θ_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 x, model inputs i θ_i, step counter i, and resulting model outputs Ψ(i,i)=(Ψ(1)(i,i),…,Ψ(D)(i,i))∈ℝA×D _ ω ( θ_i,i )= ( _ ω^(1) ( θ_i,i ),…, _ ω^(D) ( θ_i,i ) ) ^A× D. The network inputs i θ_i represents the parameters of the factorized categorical distribution pI(∣i)p_I(x θ_i), while i serves 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(x θ_i,i)= _d=1^D ^(d)_ ω( θ_i,i). Here, Ψ(d)(i,i) _ ω^(d)( θ_i,i) denotes A components of the network output corresponding to the parameters (θ1(d),…,θA(d)) (θ^(d)_1,…,θ^(d)_A ) of the categorical distribution for the d-th observation. Sender distribution A sample from the sender distribution is used to update the parameters θ. For =(y(1),…,y(D))∈Dy= (y^(1),…,y^(D) ) ^D, the sender distribution is defined as pS(∣;α)=(∣α(A−),αA),p_S (y ;α )=N (y α (Ae_x- 1 ),α A I ), where e_x is a unit vector of length A and element x is 11, also known as a one-hot-encoding. The accuracy of these samples is controlled by an accuracy parameter α∈ℝ+α ^+. When α is low, the samples provide limited information about x. As α increases, the samples become increasingly informative about x. Note that the value of i y_i determines the amount of information, not the variance. The reason for this is that i y_i will play the role of the logits. The value of (i)E( y_i) increases with αi _i, which means that we get a higher focus on the values where =1e_x=1. Receiver distribution The receiver distribution is defined according to the output distribution pOp_O, and pSp_S, this takes the form pR(∣i;i,α)=pO(′∣i;i)[pS(|′;α)],p_R(y θ_i;i,α)=E_p_O (x θ_i;i ) [p_S (y|x ;α ) ], In essence, this integrates over all ′∈1,…,ADx ∈\1,…,A\^D, considering the contribution of each possible ′x as weighted by its likelihood under the output distribution pO(|i,i)p_O(x| θ_i,i), effectively combines all potential sender distributions into a single receiver distribution. The Bayesian update function for discrete data introduced in Figure 1 is given by h(i,i)=ei⊙i∑a=1Aei,a(i)a,h ( θ_i,y_i )= e^y_i θ_i _a=1^Ae^y_i,a ( θ_i )_a, where ⊙ refers to the Hadamard product. For a detailed derivation of the update function, refer to the original work by Graves et al. (2023). 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 in pSp_S approaches zero, the process is driven to the distribution that maximizes the likelihood of sampling data from the distribution pOp_O. Graves et al. (2023) show that the loss function for an n-step procedure at step i is: LN(;,i,i,i)=N[ln(i∣αi(A−1),αiAI)−ln(∑a=1ApO(a∣i,i)(i∣αi(Aa−1),αiAI))]. splitL^N( ω; x,y_i, θ_i,i)=N\,E [ (y_i _i(Ae_x-1), _iAI )\\ - ( _a=1^Ap_O (a θ_i,i )N (y_i _i(Ae_a-1), _iAI ) ) ]. split Furthermore, Graves et al. (2023) shows that if you let N go to infinity you get a continuous-time loss function. In that case, the generative process is not bound by N used during training. We omit the detailed derivation here, but Graves et al. (2023) show that the continuous-time loss works marginally better than the discrete-time loss when N is large. We will use the continuous-time loss in our experiments. Algorithm 1 Generating Discrete Samples With Bayesian Flow Networks 1: β(1)∈ℝ+β(1) ^+, number of steps n∈ℕn , number of categories K∈ℕK 2: ←() θ← ( 1 K ) 3: for i=1 to ni=1 to n do 4: ∼pO(⋅∣,i) k p_O(· θ,i) 5: α←β(1)(2i−1n2)α←β(1)( 2i-1n^2) 6: ∼(α(K−1),αK) y (α(K e_k-1),α KI) 7: ′←e⊙ θ ← e y θ 8: ←θ′∑kθk′ θ← θ _kθ _k 9: end for 10: ∼pO(⋅∣,1) k p_O(· θ,1) 3 Related Work Several recent works have considered offline RL as a sequence modeling problem. Here we will consider the two most prominent. 3.1 Diffuser Diffuser (Janner et al., 2022) 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.2 Decision Diffuser The Decision Diffuser Ajay et al. (2023) 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; Lou et al., 2023). 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. 4 Method 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), BFN-RL only models sequences of states. 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). 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 the return. 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.1 Condition on Return There are two obvious ways we can condition on return: Either as done by the Diffuser (Janner et al., 2022) or as done by the Decision Diffuser (Ajay et al., 2023). 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). 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). 4.2 Conditioning on Current State When Diffuser (Janner et al., 2022) and Decision Diffuser (Ajay et al., 2023) condition on the current state, they both apply an inpainting technique specific to diffusion models, see Eq. (3). 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 shows this method implemented for discrete data. The alterations to BFN sampling (Algorithm 9 in Graves et al. (2023)) are shown in Algorithm 2. For continuous data, the Bayesian update at each step is made similarly, see Algorithm 3 in the Appendix. The operations involving the mask and conditioning values are those added to the regular BFN algorithm. In Algorithm 2, the mask m selects the observed variables, and 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β(t)=β(1)t^2, but parameterize its endpoint by C=Kβ(1)C=Kβ(1), where K is the number of categories. This keeps the terminal correct-versus-incorrect message-mean separation from changing merely because a representation uses more categories. We use C=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, with 12 updates for Empty-Random and 16 for DoorKey and BlockedUnlockPickup. For categorical tasks, we optimize ℒBFN+0.5ℒCEL_BFN+0.5L_CE, where ℒCEL_CE is the cross-entropy for predicting the clean category from the BFN input parameters. Algorithm 2 Inpainting Conditioning for Discrete Random Variables 1: β(1)∈ℝ+β(1) ^+, number of steps n∈ℕn , number of categories K, mask m, condition c 2: ← θ← 1 K 3: for i=1 to ni=1 to n do 4: t←i−1nt← i-1n 5: ∼k discrete_output_distribution(,t θ,t) 6: α←β(1)(2i−1n2)α←β(1) ( 2i-1n^2 ) 7: ∼(α(K−),αK)y (α(Ke_k- 1),α K I ) 8: ←α(K−)y_c←α(Ke_c- 1) 9: ←⊙+(1−)⊙y← m _c+(1- m) 10: ′←e⊙ θ ← e^y θ 11: =′∑k′ θ= θ _k θ _k 12: end for 13: ∼k discrete_output_distribution(,1 θ,1) 5 Experiments 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) datasets with the Gym-Mujoco suite of environments. We compare our method to the Decision Diffuser (Ajay et al., 2023) 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.1 Discrete 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×8 navigation and two-box Sokoban-7×7 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. Task Generator Success (%) Mean return Empty-Random-6×6 Categorical diffusion 100.0 ± 0.0 0.941 ± 0.012 BFN-RL 99.7 ± 0.6 0.895 ± 0.004 DoorKey-8×8 Categorical diffusion 99.3 ± 0.6 0.916 ± 0.010 BFN-RL 100.0 ± 0.0 0.934 ± 0.002 BlockedUnlockPickup Categorical diffusion 37.0 ± 13.5 0.305 ± 0.113 BFN-RL 32.7 ± 8.3 0.279 ± 0.072 FrozenLake-8×8 Categorical diffusion 85.5 ± 1.0 BFN-RL 86.7 ± 1.3 Sokoban-7×7 Categorical diffusion 59.0 ± 2.6 BFN-RL 68.7 ± 2.5 Table 1: Matched discrete-generator comparison. Entries are mean ± sample 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.2 Continuous Control The original Gaussian formulations of Diffuser (Janner et al., 2022) and Decision Diffuser (Ajay et al., 2023) 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 shows the performance of BFN-RL compared to state-of-the-art algorithms. The table shows that BFN-RL is competitive on most datasets. Dataset Environment BC CQL IQL DT T MOReL Diffuser D BFN-RL Med-Expert HalfCheetah 55.2 91.6 86.7 86.8 95 53.3 79.8 90.6 97.2 ± 0.1 Med-Expert Hopper 52.5 105.4 91.5 107.6 110.0 108.7 107.2 111.8 110.3 ± 0.8 Med-Expert Walker2d 107.5 108.8 109.6 108.1 101.9 95.6 108.4 108.8 106.6 ± 4.9 Medium HalfCheetah 42.6 44.0 47.4 42.6 46.9 42.1 44.2 49.1 47.8 ± 0.1 Medium Hopper 52.9 58.5 66.3 67.6 61.1 95.4 58.5 79.3 72.2 ± 7.1 Medium Walker2d 75.3 72.5 78.3 74.0 79 77.8 79.7 82.5 66.9 ± 5.5 Med-Replay HalfCheetah 36.6 45.5 44.2 36.6 41.9 40.2 42.2 39.3 42.5 ± 0.5 Med-Replay Hopper 18.1 95 94.7 82.7 91.5 93.6 96.8 100 91.2 ± 6.3 Med-Replay Walker2d 26.0 77.2 73.9 66.6 82.6 49.8 61.2 75 58.5 ± 4.8 Average 51.9 77.6 77 74.7 78.9 72.9 75.3 81.8 77.0 Table 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). 6 Conclusion In this work, we introduced a novel approach to reinforcement learning that leverages Bayesian flow networks (Graves et al., 2023) 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). 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 A Hyperparameters Here we present hyperparameters used in the experiments. For the discrete experiments, we used: • Quadratic accuracy schedule β(t)=β(1)t2β(t)=β(1)t^2, with Kβ(1)=36Kβ(1)=36, 18, and 72 for Empty-Random, DoorKey, and BlockedUnlockPickup, respectively. The planning horizons are 16, 32, and 32. • Learning rate 3×10−43× 10^-4, batch size 128, Adam optimizer (Kingma and Ba, 2014), 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, and Kβ(1)=72Kβ(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. • ϵθ _θ and fϕf_φ are trained for 2×1062× 10^6 steps using the Adam optimiser (Kingma and Ba, 2014) with a batch size of 256, a learning rate of 5×10−55× 10^-5. • We use a planning horizon H of 20. • For testing we used an exponential moving average of the weights with decay α=0.9995α=0.9995 • First-order BFN solver of Xue et al. (2024) using 10 sampling steps. • Guidance weight 1.2. • Temperature 0.5. • σ1=0.01 _1=0.01. Appendix B Algorithms Algorithm 3 gives the continuous-variable counterpart of the inpainting procedure in Algorithm 2. Algorithm 3 Inpainting Conditioning for Continuous Random Variables 1: σ1∈ℝ+ _1 ^+, number of steps n∈ℕn , mask m, condition c 2: ← μ← 0 3: ρ←1ρ← 1 4: for i=1 to ni=1 to n do 5: t←i−1nt← i-1n 6: ^(,t)← x( θ,t)← cts_output_distribution(,t,1−σ12 μ,t,1- _1^2) 7: α←σ1−2i/n(1−σ12/n)α← _1^-2i/n (1- _1^2/n ) 8: ∼(^(,t),α−1)y ( x( θ,t),α^-1 I ) 9: ←(1−σ12t)y_c←(1- _1^2t) c 10: ←⊙+(−)⊙y← m _c+( 1- m) 11: ←ρ+αρ+α μ← ρ μ+ ρ+α 12: ρ←ρ+αρ←ρ+α 13: end for 14: ^(,1)← x( θ,1)← cts_output_distribution(,1,1−σ12 μ,1,1- _1^2) Here ρ is the precision of the continuous BFN input distribution; it starts at the unit-precision prior and accumulates the message precisions α. References Agarwal et al. (2020) R. Agarwal, D. Schuurmans, and M. Norouzi An optimistic perspective on offline reinforcement learning. In International Conference on Machine Learning (ICML), Vol. 119, p. 104–114. External Links: Link Cited by: §1. Ajay et al. (2023) A. Ajay, Y. Du, A. Gupta, J. B. Tenenbaum, T. S. Jaakkola, and P. Agrawal Is conditional generative modeling all you need for decision making?. In International Conference on Learning Representations (ICLR), External Links: Link Cited by: §1, §1, §2.2, §2.3, §3.2, §4.1, §4.1, §4.2, §4, §5.2, Table 2, §5. Austin et al. (2021) J. Austin, D. D. Johnson, J. Ho, D. Tarlow, and R. van den Berg Structured denoising diffusion models in discrete state-spaces. arXiv preprint arXiv:2107.03006. External Links: Link Cited by: §1, §2.5, §3.2. Chen et al. (2021) L. Chen, K. Lu, A. Rajeswaran, K. Lee, A. Grover, M. Laskin, P. Abbeel, A. Srinivas, and I. Mordatch Decision transformer: reinforcement learning via sequence modeling. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 34, p. 15084–15097. External Links: Link Cited by: §1. Dhariwal and Nichol (2021) P. Dhariwal and A. Nichol Diffusion models beat GANs on image synthesis. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 34, p. 8780–8794. External Links: Link Cited by: §2.4. Fu et al. (2020) J. Fu, A. Kumar, O. Nachum, G. Tucker, and S. Levine D4RL: datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219. External Links: Link Cited by: §5. Graves et al. (2023) A. Graves, R. K. Srivastava, T. Atkinson, and F. Gomez Bayesian flow networks. arXiv preprint arXiv:2308.07037. External Links: Link Cited by: 1st item, §1, §2.5, §2.5, §2.5, §2.5, §4.2, §6. Ho et al. (2020) J. Ho, A. Jain, and P. Abbeel Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 33, p. 6840–6851. External Links: Link Cited by: §1, §2.3, §2.3, §2.3. Ho et al. (2022) J. Ho, T. Salimans, A. Gritsenko, W. Chan, M. Norouzi, and D. J. Fleet Video diffusion models. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 35, p. 8633–8646. External Links: Link Cited by: §2.3. Ho and Salimans (2021) J. Ho and T. Salimans Classifier-free diffusion guidance. In NeurIPS Workshop on Deep Generative Models and Downstream Applications, External Links: Link Cited by: §2.4, §2.4. Janner et al. (2022) M. Janner, Y. Du, J. B. Tenenbaum, and S. Levine Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning (ICML), p. 9902–9915. External Links: Link Cited by: §1, §1, §2.2, §2.3, §3.1, §4.1, §4.2, §4, §5.2. Janner et al. (2021) M. Janner, Q. Li, and S. Levine Offline reinforcement learning as one big sequence modeling problem. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 34, p. 1273–1286. External Links: Link Cited by: §1, §2.2. Kingma and Ba (2014) D. P. Kingma and J. Ba Adam: a method for stochastic optimization. arXiv preprint arXiv:1412.6980. External Links: Link Cited by: 2nd item, 2nd item. Kumar et al. (2020) A. Kumar, A. Zhou, G. Tucker, and S. Levine Conservative Q-learning for offline reinforcement learning. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 33, p. 1179–1191. External Links: Link Cited by: §2.2. Levine et al. (2020) S. Levine, A. Kumar, G. Tucker, and J. Fu Offline reinforcement learning: tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643. External Links: Link Cited by: §1, §2.2. Li et al. (2024) K. Li, Z. Deng, and D. Wang Guided-BFNs: towards visualizing and understanding bayesian flow networks in the context of trajectory planning. Unpublished. External Links: Link Cited by: §6. Lou et al. (2023) A. Lou, C. Meng, and S. Ermon Discrete diffusion modeling by estimating the ratios of the data distribution. arXiv preprint arXiv:2310.16834. External Links: Link Cited by: §1, §2.5, §3.2. Lugmayr et al. (2022) A. Lugmayr, M. Danelljan, A. Romero, F. Yu, R. Timofte, and L. Van Gool RePaint: inpainting using denoising diffusion probabilistic models. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), p. 11461–11471. External Links: Link Cited by: §2.4. Luo and Hu (2021) S. Luo and W. Hu Diffusion probabilistic models for 3d point cloud generation. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), p. 2837–2845. External Links: Link Cited by: §2.3. Peters et al. (2010) J. Peters, K. Mulling, and Y. Altun Relative entropy policy search. In AAAI Conference on Artificial Intelligence, Vol. 24, p. 1607–1612. External Links: Link Cited by: §2.2. Sutton and Barto (2018) R. S. Sutton and A. G. Barto Reinforcement learning: an introduction. MIT Press. External Links: Link Cited by: §1, §2.1. Xue et al. (2024) K. Xue, Y. Zhou, S. Nie, X. Min, X. Zhang, J. Zhou, and C. Li Unifying bayesian flow networks and diffusion models through stochastic differential equations. In International Conference on Machine Learning (ICML), Vol. 235, p. 55656–55681. External Links: Link Cited by: 5th item.