English

Path-Coupled Bellman Flows for Distributional Reinforcement Learning

Machine Learning 2026-05-12 v1 Artificial Intelligence

Abstract

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 \emph{boundary mismatch} at the flow source or from \emph{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 \textbf{source-consistent Bellman-coupled paths}: the current path starts from the required base prior at t=0t{=}0, reaches the Bellman target at t=1t{=}1, and maintains a pathwise affine relation to the successor flow at intermediate times (without requiring time-tt marginals to satisfy a distributional Bellman fixed point for all tt). PCBF couples current and successor return flows through shared base noise and uses a λ\lambda-parameterized control-variate target: λ=0\lambda{=}0 recovers an unbiased sample Bellman target, while λ>0\lambda{>}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.

Keywords

Cite

@article{arxiv.2605.08253,
  title  = {Path-Coupled Bellman Flows for Distributional Reinforcement Learning},
  author = {Boyang Xu and Qing Zou and Siqin Yang and Hao Yan},
  journal= {arXiv preprint arXiv:2605.08253},
  year   = {2026}
}

Comments

Accepted to the 43rd International Conference on Machine Learning (ICML 2026)

R2 v1 2026-07-01T12:58:37.291Z