What Does Flow Matching Bring To TD Learning?
This work provides a mechanistic understanding of flow matching's benefits for TD learning, which is important for researchers and practitioners in reinforcement learning to develop more robust and efficient algorithms.
This paper investigates why flow matching is effective for Q-value function estimation in reinforcement learning. It finds that flow matching improves TD learning through test-time recovery, which dampens errors in value estimates, and plastic feature learning, which helps represent non-stationary TD targets. Empirically, flow-matching critics substantially outperform monolithic critics (2x in final performance and 5x in sample efficiency) in high-UTD online RL problems.
Recent work shows that flow matching can be effective for scalar Q-value function estimation in reinforcement learning (RL), but it remains unclear why or how this approach differs from standard critics. Contrary to conventional belief, we show that their success is not explained by distributional RL, as explicitly modeling return distributions can reduce performance. Instead, we argue that the use of integration for reading out values and dense velocity supervision at each step of this integration process for training improves TD learning via two mechanisms. First, it enables robust value prediction through \emph{test-time recovery}, whereby iterative computation through integration dampens errors in early value estimates as more integration steps are performed. This recovery mechanism is absent in monolithic critics. Second, supervising the velocity field at multiple interpolant values induces more \emph{plastic} feature learning within the network, allowing critics to represent non-stationary TD targets without discarding previously learned features or overfitting to individual TD targets encountered during training. We formalize these effects and validate them empirically, showing that flow-matching critics substantially outperform monolithic critics (2$\times$ in final performance and around 5$\times$ in sample efficiency) in settings where loss of plasticity poses a challenge e.g., in high-UTD online RL problems, while remaining stable during learning.