Posterior-Averaged Bellman Flow Matching for Distributional Critics

Zeqiao Li; Yijing Wang; Haoyu Wang; Zheng Li; Zhiqiang Zuo

cs.LGOct 13, 2026

Posterior-Averaged Bellman Flow Matching for Distributional Critics

Authors: Zeqiao Li, Yijing Wang, Haoyu Wang, Zheng Li, Zhiqiang Zuo

Organizations: Tianjin Key Laboratory of Intelligent Unmanned Swarm Technology and System, School of Electrical and Information Engineering, Tianjin University, Tianjin 300072, China · Key Laboratory of System Control and Information Processing, Ministry of Education of China, Shanghai 200240, China

Abstract

Distributional critics model the full return distribution, providing richer signals for exploration and risk-sensitive control. Recent work has shown that flow-based critics can model these distributions flexibly and provide effective supervision across flow time for temporal-difference learning. However, unlike standard flow matching with fixed data endpoints, a flow critic is trained on Bellman targets recursively generated by a learned target critic under the current policy. Existing objectives therefore either use noisy single-endpoint supervision or rely on the learned flow itself for additional guidance. Our key observation is that an intermediate bridge state may be compatible with several Bellman endpoints, so its velocity should average their directions under the corresponding posterior. We propose Posterior-Averaged Bellman Flow Matching (PA-BFM), which samples an endpoint pool, infers the closed-form posterior over the latent anchor, and regresses to the posterior-mean velocity. Uniform anchor selection makes this target an exact finite-pool Rao--Blackwellization of the single-endpoint target: for every $K$ , it preserves the population target, cannot increase conditional variance, and its conditional variance scales as $O(K^{-1})$ . The endpoint pool supports conservative updates, while the twin return flows provide critic-disagreement and unsupported-mass diagnostics. Controlled studies show that PA-BFM reduces fixed-query target mean-squared error by $79.7\%$ and target-induced gradient variance by up to $98.0\%$ , while improving return-distribution fidelity. Across 50 OGBench tasks and four offline-to-online scenarios, PA-BFM achieves higher aggregate performance than Value Flows.

Figures & tables

Domain	IQL	ReBRAC	IFQL	FQL	IQN	Value Flows	PCBF	PA-BFM
	Gaussian Policies		Flow Policies		Distributional Critics
cube-double-play (5 tasks)	$6\pm 2$	$12\pm 3$	$14\pm 5$	$29\pm 6$	$42\pm 8$	$69\pm 4$	$71\pm 5$	$\mathbf{79\pm 4}$
cube-triple-play (5 tasks)	$1\pm 1$	$0\pm 0$	$0\pm 0$	$4\pm 2$	$6\pm 0$	$\mathbf{14\pm 3}$	$4\pm 1$	$\mathbf{14\pm 2}$
puzzle-3x3-play (5 tasks)	$9\pm 3$	$22\pm 2$	$19\pm 1$	$30\pm 4$	$15\pm 1$	$87\pm 13$	$53\pm 3$	$\mathbf{93\pm 3}$
puzzle-4x4-play (5 tasks)	$7\pm 2$	$14\pm 3$	$25\pm 8$	$17\pm 5$	$27\pm 4$	$27\pm 4$	$30\pm 4$	$\mathbf{52\pm 6}$
scene-play (5 tasks)	$28\pm 3$	$41\pm 7$	$30\pm 4$	$56\pm 2$	$40\pm 1$	$59\pm 4$	$54\pm 4$	$\mathbf{78\pm 2}$

Table 1: Offline evaluation on OGBench. Each row averages five tasks; state-based and image-based results use eight and four seeds. Uncertainty is the arithmetic mean of the five task-level standard deviations. Matched flow critics share policy-side components while retaining their original critic supervision. Bold denotes the highest displayed mean in each domain.

Figure 1: Controlled posterior-target evaluation. (a) Return-law $W_{1}$ ; (b) conditional target MSE; (c) target-induced gradient trace-covariance ratio; and (d) normalized $\ESS(\mathbf{w})/K$ at $K=8$ . Sol./D-MC and VF denote Solitaire Dice/Discrete-MC and Value Flows; solid/dashed in (c) follow the same environment order, and dotted in (d) is $1/K$ . Error bars are 95% confidence intervals; (b)–(c) use exact bridge-posterior anchors at fixed $(\tau,t,x_{t})$ .

Figure 2: Offline-to-online evaluation. Mean success and one-standard-deviation bands over eight seeds on four OGBench scenarios. Gray marks one million offline optimizer updates, followed by one million online environment steps.

Figure 3: Estimator components and pool-size sensitivity. Top: BCFM ( $K=1$ ), uniform and PA-BCFM ( $K=8$ ), and Value Flows. Bottom: pool-size sensitivity on two state-based tasks.

Figure 4: Critic diagnostics and adaptive pessimism on puzzle-4x4-play-task4 . Left: Error, Twin, and Out denote absolute prediction error, twin disagreement, and unsupported mass. Right: None, Fixed, Adapt., and VF denote undistorted, fixed, adaptive, and Value Flows variants.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Meaning	Symbol	Meaning
$\tau=(r,s^{\prime},d)$	transition outcome for fixed $(s,a)$	$\mu_{\tau}$	conditional Bellman endpoint law
$Y$ , $\mathbf{Y}$	endpoint and pool $(Y_{1},\ldots,Y_{K})$	$I$	uniformly sampled anchor index
$X_{t},x_{t}$	bridge random variable and realization	$\epsilon,p_{0}$	$\epsilon\sim\mathcal{N}(0,1)$ and its density $p_{0}$
$\ell_{t}(x\mid y)$	bridge likelihood	$u_{t}(x,y)$	endpoint-conditioned bridge direction
$u_{j}$	candidate direction $u_{t}(x_{t},Y_{j})$	$w_{j}$	posterior anchor responsibility
$\widehat{u}_{K}$	posterior-mean training target	$v_{\theta}$	learned velocity field

Appendix

Table 2: Core notation. Uppercase variables denote random quantities and lowercase variables their realizations. Implementation- and proof-specific symbols are defined in the subsections where they are used.

Process	Configuration
Bernoulli return	One-state continuing MRP; $R\sim\operatorname{Bernoulli}(0.5)$ ; $\gamma=0.5$ ; exact return law $Z\sim\operatorname{Unif}(0,2)$ .
Solitaire Dice	A fair six-sided die is repeatedly rolled. Rolls $2$ – $6$ produce reward $1$ and continue; roll $1$ produces reward $0$ and terminates; $\gamma=0.9$ .
Discrete-MC	Twenty-state nearest-neighbor chain with absorbing terminal states $0$ and $19$ . States $1$ – $18$ use the fixed cosine–exponential nearest-neighbor transition weights of the controlled-process configuration. Transitions to nonterminal states produce reward $1$ , transitions into a terminal state produce reward $0$ , and $\gamma=0.95$ .

Appendix

Table 3: Controlled-process settings. The Bernoulli process is used to fix the controlled-study configuration; Solitaire Dice and Discrete-MC are used for the reported evaluation.

$K$	Signed bias [95% CI]	Cond. variance	MSE [95% CI]
1	$0.004\,[-0.006,0.014]$	5.0407	$5.0447\,[4.86,5.23]$
2	$-0.003\,[-0.011,0.005]$	3.0301	$3.0331\,[2.91,3.16]$
4	$0.002\,[-0.005,0.009]$	1.7790	$1.7810\,[1.70,1.86]$
8	$-0.001\,[-0.006,0.004]$	1.0209	$1.0219\,[0.97,1.08]$
16	$0.001\,[-0.003,0.005]$	0.5693	$0.5703\,[0.54,0.61]$

Appendix

Table 4: Fixed-query conditional target estimation. Bias is the mean signed error; variance is the mean within-query conditional variance. Brackets give 95% process-stratified query-cluster bootstrap confidence intervals. Lower MSE is better.

Method	Solitaire	Discrete-MC	Avg.
BCFM ( $K=1$ )	$\begin{subarray}{c}6.644\\ {[6.31,6.98]}\end{subarray}$	$\begin{subarray}{c}1.454\\ {[1.36,1.55]}\end{subarray}$	$\begin{subarray}{c}4.049\\ {[3.86,4.24]}\end{subarray}$
Uniform ( $K=8$ )	$\begin{subarray}{c}6.252\\ {[5.92,6.58]}\end{subarray}$	$\begin{subarray}{c}3.630\\ {[3.38,3.88]}\end{subarray}$	$\begin{subarray}{c}4.941\\ {[4.73,5.15]}\end{subarray}$
Value Flows	$\begin{subarray}{c}6.650\\ {[6.25,7.05]}\end{subarray}$	$\begin{subarray}{c}4.202\\ {[3.91,4.49]}\end{subarray}$	$\begin{subarray}{c}5.426\\ {[5.18,5.67]}\end{subarray}$
PCBF ( $\lambda_{\mathrm{cv}}=0.3$ )	$\begin{subarray}{c}6.020\\ {[5.76,6.28]}\end{subarray}$	$\begin{subarray}{c}1.227\\ {[1.15,1.30]}\end{subarray}$	$\begin{subarray}{c}3.623\\ {[3.48,3.77]}\end{subarray}$
PA-BCFM ( $K=8$ )	$\begin{subarray}{c}\mathbf{5.891}\\ {[5.65,6.13]}\end{subarray}$	$\begin{subarray}{c}\mathbf{1.185}\\ {[1.12,1.25]}\end{subarray}$	$\begin{subarray}{c}\mathbf{3.538}\\ {[3.41,3.67]}\end{subarray}$

Appendix

Table 6: Learned return-law fidelity. Entries are seed-level means with 95% confidence intervals. Avg. is the equal-weight average over Solitaire Dice and Discrete-MC. Lower

1

-Wasserstein distance is better.

cube-double-play
	Gaussian Policies			Flow Policies		Distributional Critics
Task	BC	IQL	ReBRAC	IFQL	FQL	C51	IQN	CODAC	Value Flows	PCBF	PA-BFM
task1	$8\pm 3$	$27\pm 5$	$45\pm 6$	$35\pm 9$	$61\pm 9$	$9\pm 0$	$70\pm 14$	$80\pm 11$	$97\pm 1$	$92\pm 3$	$\mathbf{99\pm 1}$
task2 *	$0\pm 0$	$1\pm 1$	$7\pm 3$	$9\pm 5$	$36\pm 6$	$0\pm 0$	$24\pm 9$	$63\pm 4$	$76\pm 7$	$74\pm 7$	$\mathbf{89\pm 2}$
task3	$0\pm 0$	$0\pm 0$	$4\pm 1$	$8\pm 5$	$22\pm 5$	$0\pm 0$	$25\pm 6$	$66\pm 9$	$73\pm 4$	$81\pm 8$	$\mathbf{91\pm 2}$
task4	$0\pm 0$	$0\pm 0$	$1\pm 1$	$1\pm 1$	$5\pm 2$	$0\pm 0$	$10\pm 1$	$13\pm 2$	$30\pm 5$	$22\pm 5$	$\mathbf{44\pm 7}$

Appendix

Table 7: Full offline evaluation on state-based OGBench tasks. Results cover 25 tasks and eight training seeds. An asterisk marks the domain-level tuning task. Bold denotes the highest displayed task mean; ties at the highest mean are all bold, while fully tied rows are left unbolded.

visual-antmaze-medium-navigate
	Gaussian Policies		Flow Policies		Distributional Critics
Task	IQL	ReBRAC	IFQL	FQL	IQN	Value Flows	PCBF	PA-BFM
task1 *	$78\pm 9$	$54\pm 15$	$\mathbf{81\pm 3}$	$32\pm 3$	$62\pm 7$	$77\pm 4$	$65\pm 3$	$78\pm 7$
task2	$90\pm 3$	$\mathbf{96\pm 1}$	$87\pm 1$	$60\pm 2$	$88\pm 2$	$75\pm 5$	$73\pm 1$	$92\pm 2$
task3	$80\pm 6$	$97\pm 1$	$92\pm 1$	$35\pm 8$	$64\pm 5$	$81\pm 7$	$86\pm 4$	$\mathbf{98\pm 2}$
task4	$89\pm 4$	$93\pm 2$	$84\pm 3$	$35\pm 2$	$71\pm 6$	$71\pm 6$	$78\pm 3$	$\mathbf{96\pm 3}$

Appendix

Table 8: Full offline evaluation on 25 image-based OGBench tasks over four seeds. An asterisk marks the tuning task; bold denotes the highest displayed task mean, with fully tied rows left unbolded.

Offline checkpoint at $1000\text{k}$ steps
Method	Antmaze	Cube-D	Cube-T	Scene
IQN	$55{\pm}5$	$29{\pm}4$	$25{\pm}5$	$0{\pm}0$
IFQL	$11{\pm}3$	$12{\pm}3$	$2{\pm}1$	$0{\pm}0$
FQL	$\mathbf{87{\pm}7}$	$40{\pm}11$	$4{\pm}1$	$\mathbf{82{\pm}11}$
RLPD	$0{\pm}0$	$0{\pm}0$	$0{\pm}0$	$0{\pm}0$
IQL	$41{\pm}5$	$0{\pm}0$	$3{\pm}2$	$15{\pm}4$

Appendix

Table 9: Offline-to-online evaluation on four OGBench scenarios. The two panels report success at online entry and after one million online environment steps. Entries are means

\pm

standard deviations over eight seeds; bold denotes means within

95\%

of the highest scenario mean.

Variant	Positive overestimation $\downarrow$	Final success $\uparrow$
No distortion	$0.380\pm 0.033$	$0.490\pm 0.100$
Fixed $\eta=0.1$	$0.300\pm 0.027$	$0.505\pm 0.090$
Fixed $\eta=0.3$	$0.220\pm 0.024$	$0.455\pm 0.110$
Fixed $\eta=0.6$	$\mathbf{0.105\pm 0.019}$	$0.365\pm 0.130$
Adaptive	$0.165\pm 0.023$	$\mathbf{0.515\pm 0.090}$
Value Flows	$0.335\pm 0.041$	$0.280\pm 0.050$

Appendix

Table 10: Conservatism–performance trade-off on puzzle-4x4-play-task4 . The undistorted PA-BFM and Value Flows scores correspond to the task-level results in the main benchmark table. Lower positive overestimation and higher final success are better.

Method	Target work	State	Visual
BCFM ( $K=1$ )	2 target solves	3.6	10.5
Uniform ( $K=8$ )	16 target solves	3.8	11.0
PA-BCFM ( $K=4$ )	8 solves + wt.	3.8	11.0
PA-BCFM ( $K=8$ )	16 solves + wt.	4.0	11.8
PA-BCFM ( $K=16$ )	32 solves + wt.	4.5	13.2
Value Flows	path + confidence	4.4	12.5

Appendix

Table 11: Wall-clock training cost. State and visual columns report configuration-level hours for

10^{6}

-update budgets on one NVIDIA GeForce RTX 4090 (24 GB) under the two representative timing settings. Target work counts both target critics, so

K

ensemble-mean endpoints require

2K

target-flow solves.

Hyperparameter	Value
Optimizer	Adam ( Kingma and Ba, 2015 )
Batch size	$256$
Critic learning rate	$3\times 10^{-4}$
Critic hidden-layer sizes	$(512,512,512,512)$
Activation function	GELU
Actor layer normalization	Yes

Appendix

Table 12: Common hyperparameters for PA-BFM and matched flow-critic baselines. Entries apply where the corresponding component is present.

Domain/scenario	Evaluation	$\gamma$
cube-double	Off.+O2O	.995
cube-triple	Off.+O2O	.995
puzzle-3x3	Offline	.99
puzzle-4x4	Offline	.99
scene	Off.+O2O	.99
visual-antmaze- medium	Offline	.99

Appendix

Table 13: Domain/scenario settings. PA-BFM uses an undistorted

K=8

critic, and all main experiments use mean-based candidate scoring.

Distributional RL

Distributional reinforcement learning models the full range of possible returns to improve value estimation, risk-sensitive decisions, and policy evaluation.

See all 49 papers

Sep 26, 2026

Retimed Bellman Flows: Escaping the Impossible Triangle of Velocity Bootstrapping

Flow critics learn return distributions by transporting Gaussian noise to Bellman endpoints via continuous velocity fields. While velocity bootstrapping stabilizes training by querying a successor teacher, existing methods face a structural dilemma: on straight paths, no residual-free same-time affine mapping can preserve Gaussian initial noise while maintaining an unbiased target. To overcome this limitation, we introduce Retimed Bellman Flows (ReBF). ReBF queries the teacher critic at a dynamically shifted earlier flow time, aligning intermediate student and teacher trajectories. By combining this retimed clock with fresh, decoupled noise generation, ReBF constructs a provably conditionally unbiased velocity target that preserves the Bellman fixed point and contracts under Wasserstein distances. Empirically, ReBF reduces

W_1

distance to ground-truth return distributions by up to

7.7\times

on synthetic MRPs and outperforms existing flow critics across 38 challenging OGBench and D4RL offline reinforcement learning tasks.

Boyang Xu, Shengzhe Chen, Hao Yan

Arizona State University

Reinforcement Learning Distributional RL Offline RL

Ask PDF

May 8, 2026

Quantile-Coupled Flow Matching for Distributional Reinforcement Learning

Unlike standard expected-return Reinforcement Learning (RL), Distributional RL (DRL) models the full return distribution, making it better-suited for uncertainty-aware and risk-sensitive decision-making. Conditional Flow Matching (CFM) critics have recently attracted attention for modelling continuous, multi-modal return distributions. Despite this interest, there remains a substantial metric mismatch: DRL theory relies on the distributional Bellman operator being contractive in the

p

-Wasserstein distance, yet existing CFM critics are trained with arbitrary source-target couplings, so their flow-matching losses are not Wasserstein-aligned surrogates for matching Bellman target return distributions. In this work, we address this mismatch by proposing FlowIQN, a CFM critic that sorts source and Bellman target samples within each mini-batch to approximate the monotone optimal transport coupling, replacing arbitrary pairings with quantile-aligned flow paths. We prove that the loss of our quantile-coupled CFM critic yields a Wasserstein-aligned approximate projection compatible with the foundations of DRL. To our knowledge, FlowIQN is the first flow-matching distributional critic with an explicit Wasserstein-aligned projection guarantee. We further extend FlowIQN with shortcut models for efficient inference. Empirical results show that FlowIQN improves Wasserstein return-distribution accuracy over other CFM critics. It also yields competitive performance on offline RL benchmarks across multiple policy extraction methods, providing a theoretically grounded CFM critic that is readily compatible with DRL pipelines. Code: https://github.com/ori-goals/flowIQN.

Michael Groom, Victor-Alexandru Darvariu, Lars Kunze +2

Oxford Robotics Institute, University of Oxford, Oxford, United Kingdom · Bristol Robotics Laboratory, University of the West of England, Bristol, United Kingdom · Dyson Institute of Engineering & Technology, Malmesbury, United Kingdom

Distributional RL Conditional Flow Matching Optimal Transport

Ask PDF

May 7, 2026

Path-Coupled Bellman Flows for Distributional Reinforcement Learning

Distributional reinforcement learning (DRL) models the full return distribution, but existing finite-support or quantile-based methods rely on projections, while recent flow-based approaches can suffer from boundary mismatch at the flow source or from high-variance bootstrapping when current and successor noises are independent. We propose Path-Coupled Bellman Flows (PCBF), a continuous-time DRL method that learns return distributions with flow matching using source-consistent Bellman-coupled paths: the current path starts from the required base prior at

t{=}0

, reaches the Bellman target at

t{=}1

, and maintains a pathwise affine relation to the successor flow at intermediate times (without requiring time-

t

marginals to satisfy a distributional Bellman fixed point for all

t

). PCBF couples current and successor return flows through shared base noise and uses a

λ

-parameterized control-variate target:

λ{=}0

recovers an unbiased sample Bellman target, while

λ{>}0

trades controlled bias for variance reduction. Experiments on analytically tractable MRPs, OGBench, and D4RL show improved distributional fidelity and training stability, and competitive offline RL performance.

Boyang Xu, Qing Zou, Siqin Yang +1

School of Computing and Augmented Intelligence, Arizona State University, Tempe, AZ, USA

Distributional RL Flow Matching

Ask PDF

Abstract

Distributional RL

Explore similar work

Retimed Bellman Flows: Escaping the Impossible Triangle of Velocity Bootstrapping

Quantile-Coupled Flow Matching for Distributional Reinforcement Learning

Path-Coupled Bellman Flows for Distributional Reinforcement Learning