cs.LGMar 4, 2026
SaveWhat Does Flow Matching Bring To TD Learning?
Organizations: Carnegie Mellon University · University of Warsaw
Abstract
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 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 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 in final performance and around 5 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.
Figures & tables
Figure 1: Flow-matching critics with non-stationary TD targets. Values are computed by integrating a learned velocity field over multiple steps. This iterative process enables test-time recovery (left), where errors in early integration steps are dampened by later steps. When TD targets change ( ), the integration dynamics can absorb part of the shift, allowing earlier features to remain largely unchanged, improving plasticity .
| Env. | Success (%) | ||
|---|---|---|---|
| hmmaze-large | 52 / 30 | ||
| antmaze-giant | 86 / 74 | ||
| cube-double | 72 / 72 | ||
| hmmaze-medium | 94 / 94 |
Table 1: Comparing expected-value (E) vs. distributional (D) flow-matching critics on representative OGBench tasks. Each entry reports E / D . While both variants learn similar expected Q-values, D produces higher-variance estimates (that more closely reflect variation in the return distribution) but does not outperform E .
| Env. (default task) | Success as stale fraction increases |
|---|---|
| antsoccer-arena | \hbox{\pagecolor{tmpShade}45}\rightarrow{\color[rgb]{0.22,0.45,0.7}\mathbf{50}}\rightarrow\hbox{\pagecolor{tmpShade}47}\rightarrow\hbox{\pagecolor{tmpShade}46}\rightarrow\hbox{\pagecolor{tmpShade}35} |
| hmmaze-medium | {\color[rgb]{0.22,0.45,0.7}\mathbf{98}}\rightarrow\hbox{\pagecolor{tmpShade}43}\rightarrow\hbox{\pagecolor{tmpShade}63}\rightarrow\hbox{\pagecolor{tmpShade}62}\rightarrow\hbox{\pagecolor{tmpShade}39} |
| cube-double | \hbox{\pagecolor{tmpShade}72}\rightarrow{\color[rgb]{0.22,0.45,0.7}\mathbf{84}}\rightarrow\hbox{\pagecolor{tmpShade}82}\rightarrow\hbox{\pagecolor{tmpShade}58}\rightarrow\hbox{\pagecolor{tmpShade}64} |
Table 2: Effect of injecting staleness into early integration steps of the flow-matching critic. Entries are shaded by performance degradation relative to the best value within each environment as increases from left to right. Note that a flow-matching critic can still succeed in many cases when the first 25% or even 50% of the integration is done with a stale velocity field. On two out of three environments, using a stale velocity field improves performance slightly (best performance in each row is in bold ).
Figure 2: Performance of flow-matching ( floq ) and monolithic (FQL) critics when trained with target noise . Observe that flow-matching critics are much more robust to noise in TD targets, while performance of FQL (monolithic critic) degrades substantially faster, even when they start at a similar point (antmaze/antsoccer).
Figure 3: Feature norms. (a) Learned feature norms and average Q-values for monolithic critics (FQL) and flow-matching critics ( floq ) in the penultimate and last hidden layers. While the last hidden layer adapts to the scale of Q-values for both methods, the penultimate hidden layer in floq exhibits a much more rapid decrease in feature norms compared to FQL. This indicates that floq learns more flexible and adaptive features in the penultimate hidden layer that are largely decoupled from the magnitude of Q-values. (b) Penultimate hidden layer feature norms for floq trained with TD, SARSA, and MC targets. floq with TD shows the fastest decrease in feature norms, whereas SARSA and MC trends resemble those of the monolithic FQL critic. This suggests that flow-matching critics, particularly under TD learning, develop more robust representations under non-stationary targets.
Figure 4: Measuring feature plasticity on four tasks, by freezing all layers except the final two at steps (gray shaded region denotes the pre-freeze phase). Solid curves correspond to the default (fully trained) runs, while dashed curves show performance after freezing the penultimate hidden features. Across all environments, FQL with a monolithic critic exhibits a sharp performance collapse once features are frozen (brown vs orange), indicating in inability to represent future Q-functions. In contrast, flow-matching critics remain stable and continue to improve after freezing, showing substantially greater plasticity.
Figure 5: (a) Frozen features with a monolithic ResNet critic. Although a monolithic ResNet admits a computation graph similar to a flow-matching critic, freezing its features leads to a collapse in performance during subsequent offline RL training. (b) Frozen features with a monolithic transformer-based critic. Despite performing well with the standard FQL algorithm, monolithic transformer critics still suffer a performance collapse when their layers are frozen during subsequent training, indicating that the plasticity issue persists across model architectures. (c) Frozen features with a single integration step. With only one integration step, a flow-matching critic is more stable than a monolithic network, but less stable than full flow matching with multiple integration steps (performance drop indicated by the red arrow), highlighting the essential role of integration in preserving feature plasticity.
Figure 6: Impact of predicting velocities vs. final TD targets at each integration step. Predicting the final TD-target substantially degrades the performance of a flow-matching critic, yielding performance comparable to monolithic critics. This underscores the importance of dense velocity supervision in flow-matching critics.
Figure 7: Comparing monolithic and flow-matching critics with RLPD on four hard OG-bench environments. Each column corresponds to a different environment (humanoidmaze-medium, humanoidmaze-large, antsoccer-arena, antmaze-giant). Row 1: Performance curves for floq at UTD . Row 2: Performance curves for FQL at UTD . Row 3: Final success rate at K environment steps versus UTD. Row 4: Sample efficiency measured as the number of environment steps required to reach of the best achieved final performance (capped at a maximum of 100,000). Shaded regions indicate standard deviation across random seeds. Observe that floq outperforms monolithic FQL on success rates at high UTD values (row 3) by 2 on many tasks, stability (rows 1 vs 2), and efficiency (row 4) by 5 on many tasks.
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: While the last hidden layer adapts to the Q-value scale for both floq and monolithic critics, the penultimate hidden layer layer in floq exhibits a qualitatively distinct, rapid decrease in feature norms. This indicates that floq learns more adaptive representations in the penultimate hidden layer that are largely decoupled from the Q-value scale.
Figure 9: Flow-matching critics trained with TD learning ( floq ) shows the fastest decrease in penultimate hidden layer feature norms, whereas floq (SARSA) and floq (MC) trends resemble those of the monolithic FQL critics. Thus flow-matching critics develop particularly better representations under non-stationary TD targets.
Figure 10: Flow-matching critics with one integration step are more stable than monolithic critics but less stable than full flow matching. This emphasizes that test-time recovery in the form of integration plays a crucial role in preserving feature plasticity.
Figure 11: Performance profiles across all update-to-data (UTD) ratios. Each column corresponds to an environment and each row corresponds to a UTD ratio. Solid lines denote floq and dashed lines denote FQL. Curves show mean success rate over environment steps and shaded regions indicate standard deviation across random seeds.
Figure 12: OGBench [ 47 ] domains . These tasks include high-dimensional state and action spaces, sparse rewards, stochasticity, as well as hierarchical structure.
| Monolithic critic | Flow-matching critic (Euler integration; slice-wise supervision) |
|---|---|
| Function | |
| Single linear predictor: | -step Euler recursion (integration index , training step ): Output . If , the mean predictor is |
| Parameters | |
| For each slice : often grouped as | |
| Training loss | |
| Squared loss against time-varying target : | Slice-wise quadratic loss: |
Table 3: Summary of differences between monolithic and flow-matching predictors in the linear gradient-flow model. Flow matching induces an additional feature reweighting channel (via , driven by gain dynamics ) that can move the predictor even when , whereas monolithic predictors require direct motion of to change .
More in this topic
Reinforcement Learning
Reinforcement learning studies how agents learn sequential decisions from rewards and interaction, balancing exploration, performance, safety, and adaptation.
See all 1,717 papers