PaperlayerView exact revision on arXiv

Source-linked AI summary

Posterior-Averaged Bellman Flow Matching for Distributional Critics

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

arXiv:2610.17788v1cs.LGcs.AI

TL;DR

A bridge state can be compatible with multiple Bellman endpoints, making a single-endpoint target an uncertain supervision signal. PA-BFM averages endpoint directions using their posterior, and experiments report improved estimation, return-law fidelity, and aggregate performance over Value Flows.

  • Problem

    A bridge state does not identify which compatible Bellman endpoint was sampled, so using one continuation introduces avoidable label variance.

    Source: §3.1 · p. 5

  • Method

    PA-BFM samples a finite endpoint pool and regresses to the closed-form posterior-weighted velocity, an exact finite-pool Rao–Blackwellization of the single-endpoint target.

    Source: §1 · p. 2

  • Results

    Controlled and OGBench experiments show improved estimation, return-law fidelity, and aggregate offline and offline-to-online performance over Value Flows.

    Source: §5 · p. 13

  • Takeaways & Limitations

    The endpoint pool supports conservative posterior reweighting, while twin return flows provide disagreement and unsupported-mass diagnostics.

    Source: §1 · p. 2

    The exact finite-pool Rao–Blackwell identity applies to the undistorted target, but generally not to a distorted target.

    Source: §3.3 · p. 8

Abstract

from arXiv · show

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.

1 Introduction

Distributional flow critics face a trade-off between noisy supervision from a single Bellman endpoint and guidance from a learned successor field. PA-BFM instead averages compatible endpoint directions under their posterior, with finite-pool guarantees and reported empirical gains.

1 Introduction

  • At an intermediate bridge state, several Bellman endpoints may be compatible, so the population velocity averages their directions rather than identifying one continuation.Existing objectives use one sampled continuation or guidance from a learned successor field.Source: §1 · pp. 1–2
  • PA-BFM samples a finite Bellman endpoint pool and uses closed-form bridge posteriors to estimate the posterior-mean velocity without a learned successor field.Uniform anchor selection makes the target the conditional expectation of the single-endpoint BCFM target.Source: §1 · p. 2
  • For every finite K, uniform anchor selection preserves the population bridge velocity, cannot increase conditional target variance, and yields O(K−1) conditional-variance scaling.The guarantees also connect velocity error to gradient variability, terminal transport error, and Bellman tracking.Source: §1 · p. 2
  • PA-BFM shows improved estimation, return-law fidelity, and downstream performance across controlled studies, 50 OGBench tasks, and four offline-to-online scenarios.The conclusion reports higher aggregate performance than Value Flows.Source: §1 · p. 2
  • The endpoint pool supports conservative posterior tilting, while twin return flows provide critic-disagreement and unsupported-mass diagnostics.The main experiments use the undistorted target and matched flow-critic policy extraction.Source: §1 · p. 2

2 Preliminaries

The preliminaries establish distributional returns and flow matching before defining the Bellman endpoint law and the source-consistent bridge used by the critic.

2.1 MDPs and Distributional Returns

  • Distributional RL models the law of returns rather than only its expectation; the distributional Bellman operator is a γ-contraction under Wasserstein metrics with Zπ as its unique fixed point.The setup distinguishes environment time h from flow time t.Source: §2.1 · p. 3

2.2 Flow Matching

  • Flow matching learns a time-dependent vector field transporting Gaussian noise to a target distribution, while conditional flow matching trains it on tractable paths whose population velocity averages compatible path directions.The conditional paths interpolate between source noise ϵ and endpoint y.Source: §2.2 · p. 3

2.3 Bellman Endpoint Laws and Source-Consistent Bridges

  • The Bellman endpoint law is conditioned on a transition, but does not uniquely specify an intermediate path; the Bellman-affine path has source N(r, γ²), unlike the critic's common source N(0, 1).The paper therefore uses a source-consistent bridge.Source: §2.3 · p. 4
  • An endpoint pool averages successor-action and target-flow randomness within one transition, while replay sampling supplies the outer average over transition randomness.The endpoint pool estimates velocity conditional on one transition.Source: §2.3 · p. 4

3 Method

PA-BFM treats Bellman flow supervision as posterior estimation over compatible endpoints, averaging their directions while retaining the population bridge velocity. The method also supports conservative reweighting and twin-flow diagnostics, with policy-side components matched across flow critics.

3.1 Bellman Objectives and Posterior Velocity

  • Because observing an intermediate bridge state does not identify its sampled continuation, single-endpoint BCFM has avoidable label variance, while learned-field supervision is self-referential.PA-BFM changes how the Bellman direction is formed while leaving the target policy, return-flow parameterization, and policy-extraction protocol unchanged.Source: §3.1 · p. 5

3.2 Posterior-Averaged Bellman Flow Matching

  • PA-BCFM samples an endpoint pool, infers anchor responsibilities from bridge likelihoods, and averages candidate directions; every finite pool preserves the population target without increasing conditional target variance.Under the stated moment conditions, the remaining large-pool variance scales as O(K^-1).Source: §3.2 · pp. 5–7

3 Method

  • PA-BFM uses Bellman endpoint candidates from the transition-conditional law and analytic linear-Gaussian bridge likelihoods, avoiding retrieval or an auxiliary proposal distribution.The paper distinguishes PA-BFM, the complete algorithm, from PA-BCFM, its posterior-averaged critic-training objective.Source: §3 · p. 4, §3.2 · p. 6

3.3 Conservative Learning and Critic Diagnostics

  • The weighted endpoint pool supports conservative reweighting, but the exact finite-pool Rao–Blackwell identity applies to the undistorted target and generally not to distorted targets.The paper also defines separate state–action diagnostics from twin return flows, including upper-tail unsupported-mass evaluation.Source: §3.3 · pp. 8–9

3.4 Distributional Policy Improvement

  • PA-BFM changes only critic learning: matched flow-critic methods share policy-side procedures, while Figure 1 reports return-law W1, conditional target MSE, gradient trace-covariance ratio, and normalized ESS(w)/K.The figure’s fixed-query target MSE and gradient-variance panels use exact bridge-posterior anchors.Source: §3.4 · p. 9, Figure 1 · p. 10

4 Experiments

Experiments evaluate posterior-target estimation, return-law fidelity, benchmark performance, estimator components, pool-size sensitivity, and critic diagnostics. Across these settings, PA-BFM improves on the reported comparisons, with gains qualified by pool-size saturation and diagnostic-specific trade-offs.

4 Experiments

  • The controlled suite uses Bernoulli-return, Solitaire Dice, and Discrete-MC processes; downstream evaluation covers 50 OGBench tasks and four offline-to-online scenarios.Matched code paths and shared policy-side and optimization components isolate critic learning among the flow-critic methods.Source: §4 · p. 10

4.1 Controlled Posterior-Target Estimation

  • At K = 8, target MSE falls by 79.7% relative to K = 1, while PA-BCFM has the lowest overall mean W1, 3.538 [3.41, 3.67].At K = 8, target-induced gradient trace-covariance ratios are 0.081 [0.073, 0.090] on Solitaire Dice and 0.020 [0.017, 0.023] on Discrete-MC.Source: §4.1 · pp. 10–11

4.2 Offline and Online Performance

  • PA-BFM’s ten-domain average is 51.6 versus 42.3 for Value Flows, and after offline training its final aggregate success is 94.75 versus 86.25 for Value Flows.The offline comparison records 35 wins, 7 ties, and 8 losses; offline-to-online evaluation follows one million offline optimizer updates with one million online environment steps.Source: §4.2 · p. 11, Figure 2 · p. 11

4.3 Estimator and Pool-Size Sensitivity

  • Figure 3’s matched K = 8 comparison shows PA-BCFM improves final success over uniform aggregation by 19.5 points on cube-double task4 and 21.2 points on puzzle-4x4 task2.Increasing K from 1 to 8 raises final success from 13.8 to 44 and from 17 to 51; K = 16 yields 43.4 and 50.1, respectively.Source: §4.3 · p. 11, Figure 3 · p. 12

4.4 Critic Diagnostics and Adaptive Pessimism

  • On puzzle-4x4-play-task4, twin disagreement and unsupported mass correlate with absolute prediction error at ρ = 0.943 and ρ = 0.886, respectively.Adaptive tilting reduces positive overestimation from 0.380 to 0.165 while raising final success from 0.490 to 0.515.Source: Figure 4 · p. 12, §4.4 · p. 13

5 Conclusion

The conclusion presents PA-BFM as an exact finite-pool averaging of single-endpoint BCFM directions, with variance and downstream performance benefits. The accompanying appendix passages describe its Bellman endpoint setting, return-flow properties, and theoretical scope.

5 Conclusion

  • PA-BFM’s default target is an exact finite-pool Rao–Blackwellization of the single-endpoint BCFM direction under the latent-anchor posterior.The conclusion also reports K^-1 variance scaling, bounds under stated assumptions, and improved estimation, return-law fidelity, and aggregate performance over Value Flows.Source: §5 · p. 13

A.1 Core Notation and Bellman Endpoint Law

  • For a fixed state-action pair and transition outcome, the Bellman endpoint law is conditional on that transition; the estimator’s twin-critic candidates use shared-noise ensemble means.The passage identifies the transition realization as τ = (r, s′, d).Source: §A.1 · p. 15

A.2 Distributional Bellman Operators

  • Under standard finite-moment assumptions, the distributional Bellman operator is a γ-contraction in maximal Wasserstein distance; for d = 1, the conditional endpoint law is the point mass at reward.Source: §A.2 · p. 15

A.3 One-Dimensional Return Flows

  • In one dimension, uniqueness of the return-flow ODE makes its terminal map monotone and invertible on its image, enabling the inverse-flow CDF.Source: §A.3 · p. 16

B Theoretical Analysis

  • The theoretical analysis establishes finite-pool Rao–Blackwell and variance results, plus optimization, transport, and Bellman-tracking implications under fixed transition, policy, and critic.The analysis uses flow times t ∈ [0, 1 − δ] for δ > 0.Source: B Theoretical Analysis · p. 16

B.1 Path Geometry and Population Bridge Target

  • Endpoint pooling estimates the transition-conditional velocity, while replay sampling supplies the outer Bellman mixture; intermediate paths need not share the same marginals or velocity fields.The residual for practical DCFM can vanish for a shared constant field independently of reward and transition, so it does not by itself identify the Bellman transform.Source: §B.1 · p. 17

B.2 Finite-Pool Rao–Blackwellization

PA-BFM averages candidate endpoint velocities using the posterior over which endpoint generated the observed bridge state. Uniform anchor selection makes this a finite-pool Rao–Blackwellization with variance and gradient-covariance guarantees.

B.2 Finite-Pool Rao–Blackwellization

  • Uniform anchor selection makes every finite-pool target preserve the population bridge velocity and population minimizer.The anchor-in-pool target is the conditional expectation of the single-endpoint target.Source: §B.2 · p. 18
  • PA-BCFM averages the single-endpoint direction YI − ϵ using posterior probabilities over the latent anchor index.At K = 1, its target reduces exactly to the single-endpoint BCFM target.Source: §B.2 · p. 18
  • Conditional target variance cannot increase, and decreases strictly when posterior-supported endpoints imply different velocities.This follows from conditional total variance.Source: §B.2 · p. 18

B.3 Asymptotic Variance, ESS, and Gradient Covariance

  • The asymptotic ESS bound describes posterior-weight concentration, while PA-BCFM gradient covariance is no greater than BCFM’s without requiring conditionally independent candidate gradients.ESS is not an unconditional exact variance identity.Source: §B.3 · p. 20

B.4 Optimization, Transport, and Bellman Tracking

The appendix relates endpoint-pool estimation to optimization, transport, and Bellman tracking, then describes conservative reweighting, diagnostics, critic updates, and policy extraction. The exact guarantees apply under stated conditions; the adaptive schedule does not receive an operator-level contraction claim.

B.4 Optimization, Transport, and Bellman Tracking

  • The endpoint-pool contribution is O(K^-1) in the optimization floor and O(K^-1/2) in the terminal Wasserstein bound.The stochastic gradient is unbiased for the population-risk gradient under exact finite-pool unbiasedness and the stated uniform-integrability condition.Source: §B.4 · pp. 20–21
  • For changing policies, the Bellman tracking bound includes consecutive return-law change and is not a global actor–critic convergence claim.The fixed-policy bound uses one-step terminal-distribution error.Source: §B.4 · p. 21

B.5 Conservative Endpoint Reweighting

  • A fixed nonincreasing distortion yields a pessimistically ordered fixed point when its Bellman operator is monotone and both iterations converge; contraction requires γLg < 1.For nonconstant distortion, the finite-pool ratio is generally biased by O(K^-1); exact finite-pool unbiasedness is specific to the undistorted objective.Source: §B.5 · p. 22

B.6 Adaptive Distortion and Reliability Diagnostics

  • The adaptive pool-dependent distortion has no operator-level contraction guarantee, and unsupported mass is a model-implied tail probability rather than an OOD certificate.The contraction result applies to a fixed population distortion.Source: §B.6 · pp. 22–23

C.1 PA-BCFM Critic Update

  • The twin target critics share successor actions and return-noise samples, average their endpoint outputs, and use a posterior target for both critic losses.The endpoint pool, posterior construction, and target quantities are stop-gradient.Source: §C.1 · pp. 23–24

C.2 Conservative Reweighting and Diagnostics

  • Hard upper-tail truncation is not used in the reported experiments; the stop-gradient adaptive diagnostic is not covered by the fixed-distortion contraction result.Twin-flow diagnostics use shared Gaussian samples to isolate disagreement, and unsupported mass is evaluated using their ensemble average.Source: §C.2 · pp. 24–26

C.3 Policy Extraction and Shared Policy Training

  • Candidate actions are ranked by predicted mean return, and optional one-step policy training uses a behavior-regularized objective against the frozen critic.The policy-side procedure is matched across PA-BFM, Value Flows, and PCBF.Source: §C.3 · p. 26

C.4 Implementation and Complexity

The implementation vectorizes target-flow solves and reuses successor-state representations across endpoints. Endpoint construction scales linearly with minibatch size, endpoint count, and flow-solver steps.

C.4 Implementation and Complexity

  • The Bernoulli process fixes the controlled-study configuration, while Solitaire Dice and Discrete-MC are the reported evaluation processes.All three processes admit exact return laws.Source: §C.4 · p. 26, Table 3 · p. 27
  • Endpoint construction costs O(BmbKTflow) network work and O(BmbK) posterior arithmetic and storage.Policy extraction requires O(NMactTflow) return-flow evaluations per state.Source: §C.4 · p. 27

D.1 Controlled Estimator Studies

The controlled studies use exact return laws to assess estimator targets, gradient variability, and learned return-law fidelity. The supplied passages describe the evaluation design and point to broader result tables, but do not report their comparative outcomes.

D.1 Controlled Estimator Studies

  • The controlled processes provide exact return laws, allowing reference return distributions and posterior Bellman velocities to be computed without a Monte Carlo reference pool.Bernoulli is used for configuration selection; Solitaire Dice and Discrete-MC are the evaluation processes.Source: §D.1 · p. 27
  • Fixed-query target estimation compares pool sizes K ∈ {1, 2, 4, 8, 16} using conditional MSE, with bias and within-query conditional variance also reported.The passages specify the estimator protocol and table metrics but provide no comparative values.Source: §D.1 · pp. 27–28, Table 4 · p. 28
  • Learned return-law fidelity is evaluated across trained configurations using 50,000 samples per evaluated state–action pair and eight seeds.The supplied passage identifies Table 6's metric as 1-Wasserstein distance but gives no results.Source: §D.1 · pp. 28–29, Table 6 · p. 29

D.2 Full Offline OGBench Results

  • Tables 7 and 8 report the complete task-level results for the full offline OGBench evaluation.The supplied passage does not state comparative outcomes.Source: §D.2 · p. 29

D.3 Offline-to-Online Results

  • Table 9 reports scenario-level performance at online entry and at the final online checkpoint.The supplied passage does not state comparative outcomes.Source: §D.3 · p. 29

D.4 Critic Diagnostics and Adaptive Pessimism

Twin-critic disagreement and unsupported predicted mass both track absolute return-prediction error across action-shift levels. Adaptive pessimism reduces positive overestimation while modestly increasing final success.

D.4 Critic Diagnostics and Adaptive Pessimism

  • The diagnostic evaluation uses eight seeds and 128 held-out state–action pairs per seed, with Monte Carlo references based on 64 rollouts.The checkpoint is evaluated on puzzle-4x4-play-task4; the reference forces a shared random action first, then follows the checkpoint policy.Source: §D.4 · p. 29
  • Across six action-shift levels, twin-critic disagreement correlates with absolute return-prediction error at ρ=0.943, and unsupported mass at ρ=0.886.These are descriptive Spearman correlations.Source: §D.4 · p. 29
  • In the reported results, the adaptive variant reduces normalized positive overestimation from 0.380 to 0.165 while increasing final success from 0.490 to 0.515.Source: §D.4 · p. 32

D.5 Runtime Analysis

The runtime analysis compares representative state-based and visual configurations under a shared timing setup, and identifies K=8 as a cost-conscious default. Experimental details specify the environments and matched-method comparisons used across the broader evaluation.

D.5 Runtime Analysis

  • K=8 remains below Value Flows runtime in both measured settings, while K=16 adds cost with little gain.The runtime measurements are configuration-level values for 106-update budgets, not averages over benchmark domains.Source: §D.5 · pp. 32–33, Table 11 · p. 33

E.2 Methods for Comparison

  • Matched flow-critic comparisons share policy-side components and differ in critic supervision: endpoint matching, weighted consistency, a path-coupled control variate, or endpoint-pool posterior averaging.The matched group is BCFM, Value Flows, PCBF, and PA-BFM; isolation claims apply only within this group.Source: E Experimental Details · p. 33, §E.2 · p. 34

E.1 Environments

  • Controlled evaluation uses exact return laws for Solitaire Dice and Discrete-MC; the one-state Bernoulli-return process is reserved for configuration selection.The exact laws provide reference return distributions and posterior Bellman velocities without a Monte Carlo reference pool.Source: §E.1 · p. 33
  • The study covers 50 offline OGBench state and image tasks and four offline-to-online scenarios, each fine-tuned for one million online steps after one million offline updates.Image tasks use three stacked RGB frames, an IMPALA encoder, and random-crop augmentation.Source: §E.1 · p. 34

E.3 Evaluation Protocol

The evaluation protocol defines checkpoint aggregation, tuning, and uncertainty reporting for offline and offline-to-online studies. Architecture and hyperparameter details establish the shared settings, while reproducibility records cover versions, seeds, commands, and evaluation procedures.

E.3 Evaluation Protocol

  • Offline results use final checkpoints, average five task means within each domain, and report descriptive task-level win/tie/loss counts rather than significance tests.The tuning task remains in the five-task aggregate after configuration selection; offline-to-online curves report means and one-standard-deviation bands over eight training seeds.Source: §E.3 · p. 34

E.4 Architectures and Hyperparameters

  • Matched flow-critic methods follow Value Flows’ network and optimization settings, while reported runtime is end-to-end on a single RTX 4090 with batch size 256.Runtime values are representative configuration-level measurements, not domain averages; performance comparisons use matched optimizer-update counts.Source: §E.4 · p. 35, §E.5 · p. 35
  • PA-BFM’s default uses K=8 endpoints without pessimistic distortion, with bridge time sampled uniformly below 1−10^-5; this setting is used downstream.Pool sizes from 1 to 16 are varied only in controlled estimator and pool-sensitivity studies.Source: §E.4 · p. 35

E.5 Runtime and Reproducibility

  • The study records software and hardware versions, dataset revisions, seeds, commands, configuration overrides, checkpoint rules, and evaluation commands.Tables and figures use the same plotting and aggregation pipeline across methods.Source: §E.5 · p. 36

Related papers

Loading 2610.17788v1…