Path-Coupled Bellman Flows for Distributional Reinforcement Learning

Conference: ICML 2026
arXiv: 2605.08253
Code: None
Area: Reinforcement Learning / Distributional RL / Flow Matching
Keywords: Distributional RL, Flow Matching, Bellman Equation, Control Variates, Offline RL

TL;DR

The geometric "affine transport" nature of the distributional Bellman equation is explicitly woven into the flow matching path: the same base noise simultaneously drives the paths of both the current and successor states, while a \(\lambda\) control variate shifts between bias and variance. This results in a distributional critic that is source-consistent, Bellman-endpoint-consistent, and stable.

Background & Motivation

Background: Distributional Reinforcement Learning (DRL) models the return as a complete distribution \(Z^\pi(s,a)\) rather than just its expectation, better capturing uncertainty. Leading approaches include categorical projection (C51) or quantile regression (QR-DQN / IQN), while a recent trend attempts to replace discrete projections with continuous probability transport models like diffusion or flow matching.

Limitations of Prior Work: Existing distributional methods face two major issues. First, discrete supports and projections (categorical, quantile) introduce heuristic projection biases, limiting distributional expressivity. Second, recent flow-based methods (e.g., the DCFM term in Value Flows, Bellman Diffusion) attempt to force the Bellman affine mapping \(Z \stackrel{d}{=} R + \gamma Z'\) onto every intermediate timestep of the flow path. Consequently, at \(t=0\), the path start becomes \(R + \gamma U\) instead of the required Gaussian prior \(U\)—directly conflicting with the hard constraint of "fixed source distribution" in flow matching. Third, even if endpoints match, sampling independent noise for current and successor states means Bellman consistency is only enforced at endpoints, leading to high-variance per-sample training targets and unstable critic learning.

Key Challenge: Flow matching requires the path to originate from a designated prior, whereas the Bellman operator naturally translates the distribution. Forcing intermediate steps to satisfy the Bellman fixed point destroys the source boundary. Meanwhile, independent noise, though simple to sample, causes the two paths to "drift" at intermediate times, preventing trajectory-level variance control.

Goal: Retain the endpoint geometry of flow matching (\(t=0\) Gaussian, \(t=1\) Bellman target) while re-injecting Bellman geometry into the path and providing an adjustable bias-variance balance mechanism.

Key Insight: The author observes that the Bellman equation is essentially affine transport. Rather than forcing the marginals to satisfy the Bellman fixed point at every intermediate step, it is better to strictly match endpoints and use "shared base noise" to couple the current and successor paths into geometrically correlated line segments. Thus, intermediate marginals are no longer required to be equal, but their velocity fields satisfy a Bellman-shaped algebraic relationship that can be explicitly exploited.

Core Idea: Replace original pointwise Bellman paths with source-consistent Bellman interpolation paths, and have both paths share the same base noise \(X_0\). Based on this, the BCFM objective is rewritten as a control-variate form: "sampling target + \(\lambda\) × (successor velocity prediction − sampled velocity)". \(\lambda=0\) reduces to unbiased BCFM, while \(\lambda>0\) trades controlled bias for variance reduction.

Method

Overall Architecture

PCBF is a two-stage framework consisting of a "flow-based distributional critic + offline policy extraction." The core component is a time-dependent velocity field \(v_\theta(t, Z_t \mid s, a)\) that solves the ODE \(dZ_t/dt = v_\theta(t, Z_t)\), transporting Gaussian noise \(X_0\) at \(t=0\) to return samples at \(t=1\). During training, a Polyak-updated target network \(v_{\theta^-}\) is maintained. For each minibatch transition \((s,a,r,s',a')\): (1) sample shared base noise \(X_0 \sim \mathcal{N}(0,I)\) and time \(t \sim \text{Unif}[0,1]\); (2) use the target network to integrate \(X_0\) at \((s',a')\) to the successor terminal \(X' = \psi_{\theta^-}^1(X_0 \mid s', a')\); (3) use the same \(X_0\) to construct the time-synchronized successor interpolation \(Z_t^{s'} = (1-t)X_0 + tX'\) and current interpolation \(Z_t^s = (1-t)X_0 + t(R+\gamma X')\); (4) calculate the \(\lambda\) control-variate target \(u_t^\lambda\) and minimize \(\|v_\theta(t, Z_t^s \mid s, a) - u_t^\lambda\|_2^2\). During inference, explicit Euler integration transports \(X_0\) to \(t=1\) to obtain return samples, and candidate actions are ranked by sample means for offline policy extraction.

Key Designs

1. Source-Consistent Bellman-Coupled Path:

   • Function: Maintains the flow matching source boundary \(Z_0=X_0\) while ensuring \(t=1\) strictly lands on the Bellman endpoint \(R+\gamma X'\).
   • Mechanism: Applying the pointwise Bellman path \(Z_t^D = R + \gamma Z_t'\) directly leads to \(Z_0^D = R+\gamma U \neq U\), violating the source boundary. The author rewrites the current path as \(Z_t^s = (1-t)X_0 + t(R+\gamma X')\), which is equivalent to \(Z_t^s = tR + \gamma Z_t^{s'} + (1-t)(1-\gamma)X_0\). The last term, the "residual anchor \((1-t)(1-\gamma)X_0\)", acts as a patch—it ensures that at \(t=0\), the \(\gamma X_0\) term is exactly filled back to \(X_0\), and it automatically vanishes at \(t=1\), satisfying both boundaries while maintaining a one-to-one correspondence with Bellman geometry.
   • Design Motivation: To completely decouple the geometric constraints of flow matching (source = noise) from Bellman-guided stochasticity. This allows the critic to use standard flow matching objectives while preserving the Bellman operator intact.

2. Shared-Noise Path Coupling:

   • Function: Drives the flow paths for both \((s,a)\) and \((s',a')\) with the same \(X_0\) to maintain trajectory-level alignment at every \(t\), rather than just matching at the \(t=1\) endpoint.
   • Mechanism: Traditional approaches sample noise independently for current and successor states, causing the per-sample target \(Y = R + \gamma X' - X_0\) variance to explode. With shared noise, \(X' = \psi_{\theta^-}^1(X_0 \mid s', a')\) originates from the same source as the current path. The author proves this is a "latent variable synchronous coupling" that maintains a contraction rate of \(\gamma\) and exhibits an additional \(t\gamma\) contraction for PCBF interpolations: \(\sup_{s,a} (\mathbb{E}|X_t^G - X_t^H|^p)^{1/p} \le t\gamma D_p(G,H)\). This implies that the difference between the two trajectories approaches 0 as \(t \to 0\) and grows slowly over flow time.
   • Design Motivation: Converts "distribution-level Bellman comparison" into "trajectory-level pointwise comparison," naturally reducing critic training variance and improving the robustness of Euler discretization under low NFE.

3. \(\lambda\)-Parameterized Control-Variate Target:

   • Function: Builds upon shared-noise coupling to unify the unbiased high-variance BCFM target and the low-variance biased target (using model-predicted successor velocity) into a family of adjustable targets.
   • Mechanism: Defines the control variate \(C_t = v_{\theta^-}(t, Z_t^{s'} \mid s', a') - (X' - X_0)\). For linear interpolation, the true successor path velocity is constant at \(X'-X_0\), so \(C_t\) measures the error between the target network's velocity prediction and the sampled velocity. The training target is \(u_t^\lambda := (R + \gamma X' - X_0) + \lambda [v_{\theta^-}(t, Z_t^{s'}) - (X' - X_0)]\). \(\lambda=0\) is unbiased BCFM; at \(\lambda=\gamma\), the \(X'\) term in the target is effectively replaced by the velocity prediction, which is particularly effective when the target network is not yet stable. The author also provides a closed-form bias \(\kappa(t,\gamma,\sigma,\rho)\) for the linear Gaussian case, proving that under shared noise (\(\rho=1\)), the bias decays as \(\mathcal{O}((1-\gamma)(1-t))\), and derives the variance-minimizing \(\lambda^\star(t) = \gamma(1-t) + \rho t\).
   • Design Motivation: Completely separates "Bellman endpoint correctness" from "variance control." \(\lambda\) is merely the strength of the control variate and does not geometrically affect the endpoint or source distribution; when the target network is sufficiently accurate, it serves as an efficient baseline estimator.

Loss & Training

The final training objective is \(\mathcal{L}(\theta) = \mathbb{E}_{(s,a,r,s',a'),X_0,t}[\|v_\theta(t, Z_t^s \mid s, a) - u_t^\lambda\|_2^2]\), accompanied by Polyak averaging for \(v_{\theta^-}\) and target flow integration for \(X' = \psi_{\theta^-}^1(X_0 \mid s', a')\). Inference utilizes Euler steps \(N \in \{4, 8, 16, 32\}\). The paper identifies the continuity equation for Bellman interpolation marginals and the population optimal velocity field \(v^\star_{s,a}(x,t) = \mathbb{E}[(R+\gamma X_1') - U \mid X_t = x]\), proving that the PCBF target is a Monte Carlo regression estimate of this conditional expectation. Total bias has an \(L_2\) bound independent of Gaussian assumptions: \(|\mathcal{B}_{s,a}[C_{\bar v}](x,t)| \le \|\bar v - \bar v^\star\|_{L_2(\mu_{x,t})} + \sigma_x\).

Key Experimental Results

Main Results

Evaluation covers 38 offline RL tasks (20 state-based OGBench manipulation, 10 pixel-based OGBench, 8 D4RL Adroit) plus three analytically solvable MRP toys (Solitaire Dice, Bernoulli MRP, Discrete MC Chain). Baselines include quantile critics (IQN, CODAC), scalar flow critics (FloQ), scalar value methods (IQL, FQL), and the most closely related flow-based distributional RL method, Value Flows.

Dataset (Total Tasks)	Metric	PCBF	Strongest Baseline	Remarks
OGBench cube-double-play (5)	mean ± std	\(\mathbf{71\pm5}\)	Value Flows \(69\pm4\)	PCBF leads; best-case scenario.
OGBench puzzle-4x4-play (5)	mean ± std	\(\mathbf{30\pm4}\)	IQN \(27\pm4\) / VF \(27\pm4\)	Long-horizon combinatorial; PCBF slightly better.
OGBench scene-play (5)	mean ± std	\(54\pm4\)	FloQ/VF \(\sim 58\pm4\)	Comparable range; no win.
D4RL Adroit (8)	mean ± std	\(\mathbf{69\pm2}\)	FQL \(\mathbf{71\pm4}\)	Tied for best within 95% interval.
visual-antmaze-teleport (5)	mean ± std	\(\mathbf{14\pm4}\)	Value Flows \(13\pm4\)	Pixel-based slightly ahead.
OGBench cube-triple-play (5)	mean ± std	\(4\pm1\)	Value Flows \(\mathbf{14\pm3}\)	Failure case; long-horizon sparse rewards.

PCBF shows the greatest advantage in tasks where "distribution tails / multimodal returns significantly impact action ranking." In tasks dominated by visual representation bottlenecks or ultra-long-horizon sparse rewards, it loses to Value Flows, suggesting that critic improvements translate to policy quality only in ranking-sensitive scenarios.

Ablation Study

Configuration	Key Metrics	Description
Shared-noise PCBF (full)	Min \(r_{corr}(t,N)\)	Lowest error across all \(t \in [0,1]\) and NFE \(\in \{4,8,16,32\}\).
Independent-noise ablation	Significantly > shared-noise	Only sampling method changed; same Bellman geometry.
Value Flows (dcfm=0)	CDF close to ground truth	Capable of tracking on toy tasks.
Value Flows (dcfm=0.5/1)	Systematic CDF variance underestimation	Particularly severe in long-horizon multimodal tasks; consistent with source boundary conflict.
PCBF on toy Solitaire/Discrete MC	Matches ground-truth CDF	Stable across all \(\lambda\) values.

Key Findings

• Shared-noise coupling is the primary source of gain: On Solitaire Dice, measuring the "corrected Bellman residual" \(r_{corr}(t,N) = \mathbb{E}[|\hat Z_t^s - (tR + \tilde\gamma \hat Z_t^{s'} + (1-t)(1-\tilde\gamma)U)|]\) with the same solver budget (NFE = 4/8/16/32) shows the shared-noise version is lower than the independent-noise version across all \((t,N)\), proving that coupling significantly reduces discretization error.
• Decoupling path geometry from variance control brings stability: Increasing the DCFM coefficient in Value Flows increasingly damages distributional accuracy (corroborating theoretical analysis of source boundary conflict), whereas PCBF's \(\lambda\) shows extremely low hyperparameter sensitivity, converging stably across nearly the entire \(\lambda \in [0, \gamma]\) interval.
• Failure modes suggest future directions: PCBF does not excel in long-horizon sparse (cube-triple-play) and high-resolution visual (visual-cube-double-play) tasks. This is primarily because policy extraction protocols, \(\lambda\) scheduling, and visual encoders have not yet been optimized for PCBF, suggesting that critic improvements require corresponding actor/representation upgrades to be fully realized.

Highlights & Insights

• The approach of using a "residual anchor \((1-t)(1-\gamma)X_0\)" to patch endpoint conflicts is elegant: adding a single term simultaneously satisfies the seemingly contradictory constraints of "must start at \(X_0\)" and "must end at \(R+\gamma X'\)" with geometric simplicity and zero engineering overhead.
• Analyzing "shared noise" as a latent variable version of synchronous coupling yields a beautiful \(t\gamma\) contraction property at intermediate times—the first explicit trajectory-level contraction provided in flow-based critic literature.
• The connection between the \(\lambda\)-control variate and the velocity network is ingenious: for linear interpolation, \(X'-X_0\) is exactly the true successor path velocity, making the velocity network prediction \(v_{\theta^-}\) a natural baseline without requiring auxiliary networks or reparameterization.
• This paradigm of "incorporating Bellman geometry via flow matching" can be transferred to Q-function distillation, reward shaping, or even actor-side policy flows—wherever there is an "operator-defined endpoint + probability transport" setting.

Limitations & Future Work

• Evaluated only on offline RL; does not address online exploration. Whether shared noise remains stable during active exploration needs verification.
• \(\lambda\) is currently a task-dependent fixed hyperparameter (the paper suggests \(\lambda \approx \gamma\) in early stages), lacking an automatic scheduling strategy. The theoretical optimal \(\lambda^\star(t) = \gamma(1-t)+\rho t\) assumes linear Gaussianity, and its applicability to general tasks is unknown.
• Significant performance drops in long-horizon sparse and pixel-level tasks; attributed to lack of synergetic optimization with policy extraction and visual encoders, indicating a lower ceiling for standalone distributional critic improvements.
• Training costs are comparable to Value Flows but significantly higher than scalar critics (multiple target flow integrations per step), which may be unfriendly for large-scale online RL deployment.

Related Work & Insights

• vs. Value Flows (DCFM): Value Flows attempts to force Bellman relations at all intermediate steps, resulting in conflicts with the Gaussian source boundary and high sensitivity to the DCFM hyperparameter. PCBF avoids this by retaining only endpoint constraints and using \(\lambda\) to control variance.
• vs. Bellman Diffusion / DFC: these methods use independent noise for endpoint matching and lack path coupling, leading to velocity field drift and high variance. PCBF's shared-noise synchronous coupling directly addresses this.
• vs. FloQ: FloQ uses flow matching to parameterize scalar Q, essentially accelerating numerical integration. PCBF learns the entire return distribution, exploiting tail and multimodal structures.
• vs. IQN / CODAC: Quantile methods avoid projection bias via quantile regression, but expressivity remains limited by the number of quantiles. PCBF is a truly continuous distribution, and its training objective includes a native \(\lambda\) bias-variance knob.

Rating

• Novelty: ⭐⭐⭐⭐ Proposing source-consistent path correction alongside shared-noise coupling is the cleanest solution for flow-based distributional RL to date.
• Experimental Thoroughness: ⭐⭐⭐⭐ 38 tasks + toy MRP CDF calibration + discretization diagnostics, missing only online RL experiments.
• Writing Quality: ⭐⭐⭐⭐ Theory and empirical evidence correspond closely; the logical chain from endpoint conflict to patching, coupling, and control variates is clear.
• Value: ⭐⭐⭐⭐ Establishes a standard paradigm for "embedding operators into flow paths" in flow-based critics, with high long-term reference value.

Related Papers

• [ICML 2026] Offline Multi-agent Reinforcement Learning via Sequential Score Decomposition
• [ICML 2026] CoCoEdit: Content-Consistent Image Editing via Region Regularized Reinforcement Learning
• [ICLR 2026] Offline Reinforcement Learning with Generative Trajectory Policies
• [ICML 2026] The Coupling Within: Flow Matching via Distilled Normalizing Flows
• [ICML 2026] Adapting Noise to Data: Generative Flows from Learned 1D Processes