Retimed bellman flows improve learning of return distributions in reinforcement learning
Retimed Bellman Flows: Escaping the Impossible Triangle of Velocity Bootstrapping
Machine Learning
Summary
When teaching computers to learn from experience, it helps to understand the possible future rewards as distributions rather than single numbers. But current methods have trouble making this learning process stable and accurate, especially when trying to follow smooth paths of randomness. The authors introduce Retimed Bellman Flows (ReBF), a way to adjust the timing of learning steps and add fresh noise to keep learning reliable and unbiased. Their approach better matches the true rewards, improving performance on many tested tasks.
What this means in practice
- •For offline reinforcement learning engineers: Build more accurate models of future rewards by using retimed velocity fields improving training stability on offline datasets.
- •For robotics developers: Enhance robot decision systems by better estimating return distributions from past data to improve policy reliability without new exploration.
Authors
Boyang Xu, Shengzhe Chen, Hao Yan
Abstract
Flow critics learn return distributions by transporting Gaussian noise to Bellman endpoints via continuous velocity fields. While velocity bootstrapping stabilizes training by querying a successor teacher, existing methods face a structural dilemma: on straight paths, no residual-free same-time affine mapping can preserve Gaussian initial noise while maintaining an unbiased target. To overcome this limitation, we introduce Retimed Bellman Flows (ReBF). ReBF queries the teacher critic at a dynamically shifted earlier flow time, aligning intermediate student and teacher trajectories. By combining this retimed clock with fresh, decoupled noise generation, ReBF constructs a provably conditionally unbiased velocity target that preserves the Bellman fixed point and contracts under Wasserstein distances. Empirically, ReBF reduces $W_1$ distance to ground-truth return distributions by up to $7.7\times$ on synthetic MRPs and outperforms existing flow critics across 38 challenging OGBench and D4RL offline reinforcement learning tasks.