Path-Coupled Bellman Flows for Distributional Reinforcement Learning

TL;DR

The paper introduces Path-Coupled Bellman Flows (PCBF), a novel continuous-time distributional RL method that improves distributional fidelity and training stability by coupling return flows through shared noise and a control-variate target.

Contribution

PCBF is the first to use source-consistent Bellman-coupled paths with flow matching, addressing boundary mismatch and high-variance issues in existing flow-based distributional RL methods.

Findings

PCBF achieves better distributional fidelity in analytical MRPs.

PCBF demonstrates improved training stability on benchmark datasets.

PCBF attains competitive offline RL performance.

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{=}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 $\lambda$-parameterized control-variate target: $\lambda{=}0$ recovers an unbiased sample Bellman target, while $\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.