PaperlayerView exact revision on arXiv

Source-linked AI summary

Bayesian Flow Networks for Offline Trajectory Planning

Ludvig Killingberg, Helge Langseth

arXiv:2608.25163v1cs.LGcs.AI

TL;DR

Offline trajectory planning needs a generative framework that handles discrete states alongside continuous domains. BFN-RL uses Bayesian Flow Networks for unified sequence generation and is competitive across both settings, with consistent gains on Sokoban.

  • Problem

    Recent trajectory-modeling approaches rely mainly on Gaussian diffusion, leaving discrete planning in need of a native categorical formulation.

  • Method

    BFN-RL updates distribution parameters through Bayesian flows to generate state sequences across categorical and continuous spaces, then infers actions with inverse dynamics.

  • Results

    Across discrete tasks, BFN-RL is competitive with categorical diffusion and consistently stronger on Sokoban, while remaining competitive in continuous control.

  • Takeaways & Limitations

    These results support Bayesian Flow Networks as a common planning framework for categorical and continuous domains.

Abstract

from arXiv · show

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.

1 Introduction

BFN-RL addresses the limitations of Gaussian diffusion for discrete offline RL by using Bayesian Flow Networks to unify categorical and continuous trajectory planning. It combines categorical state-sequence generation with learned inverse dynamics and evaluates the formulation across discrete and continuous control tasks.

  • Offline RL learns decision policies from static datasets without real-time environment interaction, making it suitable for high-risk domains such as autonomous driving and medical decision-making.
  • Contemporary sequence-modeling approaches rely mainly on Gaussian DDPMs, whereas discrete states and actions require a categorical formulation.
  • BFN-RL updates input-distribution parameters through Bayesian inference and continuous-time parameter flows instead of iteratively denoising corrupted data instances.
  • BFN-RL provides a unified parameter-flow framework for offline sequence modeling across discrete and continuous state-action spaces.
  • A categorical BFN state-sequence model combined with learned inverse dynamics produces a discrete planner.
  • The same BFN planning formulation is evaluated as viable for both discrete and continuous control problems.

2 Preliminaries

This section introduces offline reinforcement learning, diffusion-model preliminaries, and Bayesian flow networks as a generative framework for continuous and discrete data. It emphasizes that BFNs evolve distribution parameters through noisy observations rather than directly modeling noisy data instances.

  • Offline reinforcement learning: Offline reinforcement learning learns effective policies from a fixed dataset without online environment interactions.Unlike online learning, exploration is absent, so policy optimization focuses on effectively using the available dataset.
  • Denoising diffusion probabilistic models: Diffusion models generate samples by adding noise in a forward process and iteratively reconstructing them through a learned reverse process.The reverse model approximates q(x_k−1|x_k) using neural-network predictions of the reverse-process parameters.
  • Bayesian flow networks: BFNs are generative models designed to handle continuous, discrete, and discretized data, addressing limitations of diffusion models for high-quality discrete generation.The passage positions BFNs as a novel alternative because existing discrete diffusion variants are not considered state of the art for high-quality discrete data.
  • Bayesian flow networks: BFNs iteratively update distribution parameters θ_i using noisy data samples, with decreasing noise controlled by an accuracy parameter α_i.Each update applies a predetermined rule h(·) to the previous parameters and observation noise, producing an evolved distribution distinct from the data distribution.
  • Bayesian flow networks: For categorical data, increasing α makes noisy samples increasingly informative about the underlying one-hot encoded value.The sample value y_i determines the information conveyed and serves the role of logits, while α controls sample accuracy.

3 Related Work

Related work frames offline RL as sequence modeling, including unconditional diffusion for state-action trajectories and state-only diffusion paired with inverse dynamics. For discrete generation, categorical diffusion offers an alternative baseline to the categorical BFN.

  • Diffusion-based offline RL: Diffuser models state-action trajectories with unconditional diffusion, using inpainting for current or terminal states and reward-model gradients to guide high-return sampling.Inpainting supports goal-conditioned planning by imposing desired end states.
  • Diffusion-based offline RL: Decision Diffuser models state sequences and uses inverse dynamics to predict the action connecting consecutive states.This lets the diffusion model focus on state sequences rather than state-action sequences.
  • Categorical generation: Native categorical diffusion models provide an alternative approach for discrete generation.The discrete experiments compare categorical BFN against uniform-categorical diffusion that predicts clean trajectories using the exact reverse posterior.

4 Method

BFN-RL is a Bayesian-flow sequence generator for offline RL that plans in discrete and continuous domains by modeling state sequences conditioned on return. It fixes the current state through inpainting during sampling and uses a separate inverse-dynamics network to derive actions.

  • Sequence generation: BFN-RL generates state sequences for offline reinforcement learning in both discrete and continuous planning domains.The method models sequences of states rather than actions, following Decision Diffuser’s state-sequence formulation.
  • Action inference: A second inverse-dynamics network converts consecutive generated states into actions.BFN-RL therefore generates states first and derives the corresponding action sequence afterward.
  • Return conditioning: The BFN is trained on sequences with varying returns, and a desired return biases inference toward trajectories with corresponding cumulative reward.Return is defined as the sum of discounted future rewards and measures sequence quality.
  • Return conditioning: BFN-RL conditions directly on return in a classifier-free manner, enabling integration with existing sequence-modeling architectures.This choice follows the Decision Diffuser conditioning framework, which the passage describes as easier to implement and better-performing than the alternative.
  • State conditioning: During sampling, inpainting fixes the current state by replacing observed variables with their clean categorical values throughout the BFN process.The mask selects observed variables, while the condition supplies their clean values.

5 Experiments

Experiments evaluate BFN-RL as a unified planner across discrete grid-world tasks and continuous D4RL MuJoCo control. BFN-RL is competitive with categorical diffusion and state-of-the-art continuous-control methods, with a consistent advantage on Sokoban but no uniform discrete-task superiority.

  • Experimental settings: Experiments cover discrete grid-world planning and continuous D4RL MuJoCo control, comparing BFN-RL with categorical diffusion and offline RL baselines.The discrete experiments test categorical trajectory generation, while the continuous experiments assess whether the same framework remains competitive in control.
  • Discrete planning: BFN-RL uses categorical state-sequence generation with learned inverse dynamics across MiniGrid, FrozenLake, and Sokoban tasks.MiniGrid generation is return-conditioned, whereas FrozenLake and Sokoban use goal conditioning; consecutive generated states are mapped to native categorical actions.
  • Discrete planning: BFN-RL is competitive with categorical diffusion but does not uniformly outperform it: diffusion is faster in Empty-Random, while BFN-RL has higher DoorKey return.On BlockedUnlockPickup, variation across training seeds exceeds the difference in mean success.
  • Discrete planning: 9.7 percentage points: BFN-RL improves mean Sokoban success over categorical diffusion and has higher success for all three training seeds.The matched protocol controls training data, architecture, conditioning, inpainting, inverse dynamics, planning, execution, and evaluation seeds.
  • Continuous control: BFN-RL is competitive on most D4RL MuJoCo datasets against state-of-the-art continuous-control algorithms.Table 2 reports mean and standard error over three random seeds, with non-BFN-RL numbers taken from Ajay et al. (2023).

6 Conclusion

BFN-RL provides a shared Bayesian-flow framework for sequence planning in discrete and continuous domains. Results show competitiveness across both settings and motivate extending the framework to joint continuous-categorical environments.

  • Core contribution: BFN-RL enables sequence-generation-based reinforcement learning in both discrete and continuous planning domains.The framework could naturally assign continuous or categorical BFN distributions to trajectory variables within one joint model.
  • Empirical findings: Across discrete tasks, BFN-RL matches categorical diffusion competitively and is consistently stronger on Sokoban, while remaining competitive with established offline trajectory models in continuous control.These findings support BFNs as a common planning framework across categorical and continuous domains.
  • Related work: Guided-BFNs study continuous state–action trajectories, whereas BFN-RL’s categorical experiments plan with discrete states and actions.Guided-BFNs supplement conditional guidance with gradients from a learned reward model during sampling.

Appendix A Hyperparameters

The appendix specifies task-specific hyperparameters for discrete categorical planning and shared training, sampling, and guidance settings for continuous experiments. Discrete configurations vary schedules, horizons, architectures, and replanning, while continuous experiments use fixed optimization and solver settings.

  • Discrete experiments: Discrete tasks use a quadratic accuracy schedule with Kβ(1) = 36, 18, and 72 for Empty-Random, DoorKey, and BlockedUnlockPickup, with horizons 16, 32, and 32.The schedule is β(t) = β(1)t2.
  • Discrete experiments: Discrete models use task-dependent temporal U-Net widths, 12 or 16 categorical BFN sampling updates, temperatures 0.7 or 0.9, and one generated candidate.The U-Net base width is 32 for Empty-Random, 48 for DoorKey, and 48 or 96 in the BlockedUnlockPickup sweep.
  • Discrete experiments: FrozenLake and Sokoban use horizons 64 and 32, base widths 64 and 96, batch sizes 256 and 64, Kβ(1) = 72, temperature 0.9, 16 sampling updates, and one-action replanning.Both tasks also use cross-entropy weight 0.5.
  • Continuous experiments: Continuous experiments train ϵθ and fϕ for 2 × 106 steps with Adam, batch size 256, learning rate 5 × 10−5, horizon H of 20, and 10 BFN sampling steps.Testing uses exponential-moving-average weights with decay α = 0.9995 and guidance weight 1.2.

Appendix B Algorithms

Appendix B presents Algorithm 3, the continuous-variable counterpart to Algorithm 2’s inpainting procedure. It iteratively updates the continuous BFN input distribution from masked conditioning information and produces the final output distribution.

  • Algorithm 3 specifies inpainting conditioning for continuous random variables as the continuous-variable counterpart of Algorithm 2.
  • The procedure requires σ1 ∈R+, a step count n ∈N, a mask m, and a condition c, then initializes µ ←0 and ρ ←1.
  • At each step, it forms masked observations y ←m ⊙yc + (1 −m) ⊙y and updates the mean using µ ←ρµ+αy.
  • The algorithm accumulates precision via ρ ←ρ + α and returns ˆx(θ, 1) from CTS_OUTPUT_DISTRIBUTION(µ, 1, 1 −σ2 1).ρ starts at the unit-precision prior and accumulates the message precisions α.
Loading 2608.25163v1…