Path-Coupled Bellman Flows for Distributional Reinforcement Learning

Distributional reinforcement learning (DRL) models the full return distribution, but typically relies on finite-dimensional categorical or quantile approximations, often involving projection or quantile-regression approximations to the Bellman target, together with independently sampled bootstrap targets that obscure transport structure and add variance. We present Path-Coupled Bellman Flows (PCBF), a continuous-time DRL method that encodes Bellman endpoint consistency and pathwise Bellman-coupled geometry within generative flow trajectories. PCBF represents return distributions via flow matching and couples the paths of consecutive states through shared base noise, yielding a geometric Bellman flow relation between velocity fields. This structure enables a $\lambda$-parameterized control-variate target that reduces training variance while preserving the source and Bellman endpoint geometry. Experiments on analytically tractable MRPs, OGBench, and D4RL show improved distributional fidelity, training stability, and competitive offline RL performance.

Many AI systems learn to make decisions by estimating what rewards may happen in the future. In difficult tasks, looking only at the average reward can hide important risks: the same action may sometimes lead to success and sometimes to failure. Existing methods often represent these possible outcomes in ways that are either too coarse or not well aligned across consecutive decision steps, which can make training noisy and unstable. We introduce Path-Coupled Bellman Flows, a method that learns the full range of possible future rewards using a continuous model. The key idea is to connect the model’s prediction for the current situation with its prediction for the next situation by making them start from the same simple random seed and move along coordinated paths toward the correct learning target. This shared-path design helps the model stay consistent as it learns from one decision step to the next. We also add a tunable correction that can reduce noisy updates during training. In experiments on challenging decision-making benchmarks, our method better captured return distributions and achieved competitive control performance. This can help AI systems make more reliable decisions in long-horizon control and reinforcement-learning tasks.