arXiv preprint ·

Posterior-Averaged Bellman Flow Matching for Distributional Critics

Zeqiao Li, Yijing Wang, Haoyu Wang, Zheng Li and Zhiqiang Zuo

Published
Source
arXiv preprint, posted before peer review
Original
Read it at arXiv
Topics
Machine LearningArtificial Intelligence

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.

The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.

Read it on ReadSearch

Each opens after you create an account, and brings you back to this paper.

Already a reader? Sign in

All new papers in Machine Learning

Get the few papers worth your time in Machine Learning every morning.

ReadSearch reads the new papers in your field each day and brings you the few worth reading, each with its key finding in a sentence.

Papers belong to their authors and publishers; every page here links to the original. All research topics