What Does Flow Matching Bring To TD Learning?

Bhavya Agrawalla; Michal Nauman; Aviral Kumar

cs.LGMar 4, 2026

What Does Flow Matching Bring To TD Learning?

Authors: Bhavya Agrawalla, Michal Nauman, Aviral Kumar

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 $\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.

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 (

\theta\rightarrow\theta^{\prime}

), the integration dynamics can absorb part of the shift, allowing earlier features to remain largely unchanged, improving plasticity .

Env.	Success (%)	$Q_{\theta}(\mathbf{s},\mathbf{a})$	$\mathrm{Var}_{\mathbf{z}}(Q)$
hmmaze-large	52 / 30	$-180/-170$	$0.2/\mathbf{4.5}$
antmaze-giant	86 / 74	$-190/-200$	$0.1/\mathbf{0.7}$
cube-double	72 / 72	$-130/-130$	$1.1/\mathbf{6.3}$
hmmaze-medium	94 / 94	$-170/-170$	$0.3/\mathbf{2.3}$

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 $\kappa$ 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

\kappa=\{0,25,50,75,100\}\%

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

T=0.5M

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

\in\{32,64,128\}

. Row 2: Performance curves for FQL at UTD

\in\{32,64,128\}

. Row 3: Final success rate at

200

K environment steps versus UTD. Row 4: Sample efficiency measured as the number of environment steps required to reach

75\%

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

\times

on many tasks, stability (rows 1 vs 2), and efficiency (row 4) by 5

\times

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: $f_{\mathrm{mono}}(\mathbf{x};m)=w(m)^{\top}\mathbf{x}.$	$T$ -step Euler recursion (integration index $t$ , training step $m$ ): $\!\!\!\!\!\!s_{i+1}(\mathbf{x},\mathbf{z};m)=(1+h\,v_{i}(m))\,s_{i}(\mathbf{x},\mathbf{z};m)+h\,u_{i}(m)^{\top}\mathbf{x},\penalty\ \penalty\ \penalty\ s_{1}=\mathbf{z}.$ Output $\hat{y}(\mathbf{x},\mathbf{z};m)\coloneqq s_{T}(\mathbf{x},\mathbf{z};m)$ . If $\mathbb{E}[\mathbf{z}]=0$ , the mean predictor is $\!\!\!\!\!f_{\mathrm{FM}}(\mathbf{x};m)=\sum_{\ell=1}^{T-1}\beta_{\ell}(m)\,u_{\ell}(m)^{\top}\mathbf{x};\penalty\ \penalty\ \beta_{\ell}(m)=h\prod_{j=\ell+1}^{T-1}\bigl(1+h\,v_{j}(m)\bigr).$
Parameters
$w(m)\in\mathbb{R}^{d}.$	For each slice $i\in\{1,\dots,T-1\}$ : $u_{i}(m)\in\mathbb{R}^{d},\hskip 9.24994ptv_{i}(m)\in\mathbb{R},$ often grouped as $w_{i}(m)\coloneqq\begin{bmatrix}u_{i}(m)\\ v_{i}(m)\end{bmatrix}\in\mathbb{R}^{d+1}.$
Training loss
Squared loss against time-varying target $Y(m)$ : $\mathcal{L}_{\mathrm{mono}}(m,w)=\mathbb{E}\big[(w^{\top}\mathbf{x}-y(m))^{2}\big].$	Slice-wise quadratic loss: $\mathcal{L}_{i}(m,w_{i})=\mathbb{E}\Big[\big(w_{i}^{\top}\tilde{\mathbf{x}}_{i}(m)-(y(m)-\mathbf{z})\big)^{2}\Big],$ $\tilde{\mathbf{x}}_{i}(m)\coloneqq\begin{bmatrix}\mathbf{x}\\ S_{i}(m)\end{bmatrix},\penalty\ \penalty\ S_{i}(m)=\alpha_{i}\mathbf{z}+(1-\alpha_{i})y(m),\penalty\ \penalty\ \alpha_{i}=\frac{T-i}{T-1}.$

Appendix

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

\dot{\beta}_{\ell}

, driven by gain dynamics

\dot{v}_{k}

) that can move the predictor even when

\dot{u}_{\ell}=0

, whereas monolithic predictors require direct motion of

w(m)

to change

f_{\mathrm{mono}}(\cdot;m)

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

Oct 26, 2025

FlowCritic: Bridging Value Estimation with Flow Matching in Reinforcement Learning

Reliable value estimation serves as the cornerstone of reinforcement learning (RL) by evaluating long-term returns and guiding policy improvement, significantly influencing the convergence speed and final performance. Existing works improve the reliability of value function estimation via multi-critic ensembles and distributional RL, yet the former merely combines multi point estimation without capturing distributional information, whereas the latter relies on discretization or quantile regression, limiting the expressiveness of complex value distributions. Inspired by flow matching's success in generative modeling, we propose a generative paradigm for value estimation, named FlowCritic. Departing from conventional regression for deterministic value prediction, FlowCritic leverages flow matching to model value distributions and generate samples for value estimation.

Shan Zhong, Shutong Ding, He Diao +3

School of Mechanical and Electrical Engineering, University of Electronic Science and Technology of China, Chengdu 611731, China · School of Information Science and Technology, ShanghaiTech University · School of Electrical and Electronic Engineering, Nanyang Technological University, Singapore +1

Value Function Estimation Distributional RL

Ask PDF

Oct 13, 2026

Posterior-Averaged Bellman Flow Matching for Distributional Critics

Distributional critics model the full return distribution, providing richer signals for exploration and risk-sensitive control. Recent work has shown that flow-based critics can model these distributions flexibly and provide effective supervision across flow time for temporal-difference learning. However, unlike standard flow matching with fixed data endpoints, a flow critic is trained on Bellman targets recursively generated by a learned target critic under the current policy. Existing objectives therefore either use noisy single-endpoint supervision or rely on the learned flow itself for additional guidance. Our key observation is that an intermediate bridge state may be compatible with several Bellman endpoints, so its velocity should average their directions under the corresponding posterior. We propose Posterior-Averaged Bellman Flow Matching (PA-BFM), which samples an endpoint pool, infers the closed-form posterior over the latent anchor, and regresses to the posterior-mean velocity. Uniform anchor selection makes this target an exact finite-pool Rao--Blackwellization of the single-endpoint target: for every

K

, it preserves the population target, cannot increase conditional variance, and its conditional variance scales as

O(K^{-1})

. The endpoint pool supports conservative updates, while the twin return flows provide critic-disagreement and unsupported-mass diagnostics. Controlled studies show that PA-BFM reduces fixed-query target mean-squared error by

79.7\%

and target-induced gradient variance by up to

98.0\%

, while improving return-distribution fidelity. Across 50 OGBench tasks and four offline-to-online scenarios, PA-BFM achieves higher aggregate performance than Value Flows.

Zeqiao Li, Yijing Wang, Haoyu Wang +2

Tianjin Key Laboratory of Intelligent Unmanned Swarm Technology and System, School of Electrical and Information Engineering, Tianjin University, Tianjin 300072, China · Key Laboratory of System Control and Information Processing, Ministry of Education of China, Shanghai 200240, China

Distributional RL Flow Matching Flow-Based Generative Modeling Temporal-Difference Learning

Ask PDF

Jun 6, 2026

Q-VGM: Q-Guided Value-Gradient Matching for Offline-to-Online RL of Flow-Matching VLA Policies

We propose Q-Guided Value-Gradient Matching (Q-VGM), an offline-to-online RL method for fine-tuning flow-matching vision-language-action policies with a learned critic. Directly applying critic gradients to flow policies requires backpropagation through the multi-step denoising process (BPTT), which is costly and unstable at VLA scale. Motivated by an optimal-control view of denoising, we derive a local approximation connecting clean-action value gradients to local velocity corrections at late denoising steps. Q-VGM uses this connection to construct critic-guided targets for fine-tuning the pretrained action expert through velocity matching. Gradients pass only through local velocity predictions, avoiding backpropagation through the denoising trajectory. An action-sensitive chunk critic is trained with IQL offline and TD learning online. On LIBERO, Q-VGM improves the few-shot-SFT four-suite average from 84.6% to 90.7% offline using 150 policy rollouts per suite, and reaches 98.4% after online training. Compared with on-policy RL, Q-VGM achieves 5.3x higher sample efficiency on average across four LIBERO suites, measured by the episode budget for 95% rollout success rate. On three real-world bimanual tasks, offline Q-VGM improves average success from 66.7% to 98.3%.

Ziqian Wang, Rui Zhang, Yitian Liu +3

Shanghai Jiao Tong University · University of Michigan, Ann Arbor

Flow Matching Visuomotor Policy Learning Vision-Language-Action Models Offline-To-Online RL

Ask PDF

Abstract

Reinforcement Learning

Explore similar work

FlowCritic: Bridging Value Estimation with Flow Matching in Reinforcement Learning

Posterior-Averaged Bellman Flow Matching for Distributional Critics

Q-VGM: Q-Guided Value-Gradient Matching for Offline-to-Online RL of Flow-Matching VLA Policies