Reward-Weighted Flow Matching
- Reward-Weighted Flow Matching is a method that replaces uniform flow supervision with reward-conditioned signals, weighting trajectories to favor high-reward outcomes.
- It unifies approaches such as exponential sample reweighting, reward-corrected regression, and advantage-based updates to align velocity fields in applications like robotics and text-to-image generation.
- Empirical studies indicate improved convergence rates and performance, though effective regularization is essential to prevent model collapse and ensure stable training.
Reward-Weighted Flow Matching (RWFM) denotes a family of alignment methods for flow-matching and closely related diffusion models in which the pretraining regression signal is modulated by rewards, advantages, or reward-corrected targets so that the learned velocity field preferentially transports probability mass toward high-reward outcomes. In the most literal usage, RWFM is sample-weighted flow matching, with each training example multiplied by a scalar such as ; in later work, the same label also covers analytically derived reward-corrected regression targets, advantage-weighted rectification, and per-timestep reward weighting in online RL for generative flows (Pfrommer et al., 20 Jul 2025, Fan et al., 9 Feb 2025, Bergmeister et al., 11 May 2026, Fu et al., 29 Jun 2026). The resulting design space spans robotics, text-to-image generation, text-to-speech, and continuous-control policies, and it intersects with KL-regularized RL, optimal control, preference alignment, and inference-time reward guidance (Wang et al., 22 Jun 2026, Gubernatorov et al., 16 Mar 2026, Holderrieth et al., 5 Feb 2026).
1. Definition and conceptual scope
The common core of RWFM is the replacement of uniform flow-matching supervision by reward-conditioned supervision. In the direct formulation introduced for flow-matching policies, the standard conditional flow-matching loss is multiplied by a reward weight,
so that high-reward trajectories contribute more strongly to regression updates (Pfrommer et al., 20 Jul 2025). In the continuous-flow fine-tuning literature, the same idea appears as reward-weighted conditional flow matching and its online version, where the training distribution is itself sampled from the current model and the reward weight is typically chosen in Boltzmann form, (Fan et al., 9 Feb 2025).
Subsequent papers broaden the term. The RAM derivation states explicitly that, if RWFM means post-training by regressing a flow-matching target corrected by reward, then RWFM is equivalent to RAM’s reward-corrected flow-matching regression (Bergmeister et al., 11 May 2026). FlowAWR likewise defines RWFM as aligning flow fields by weighting the flow-matching objective with exponentiated rewards derived from a KL-regularized optimal policy, but then argues that a centered, magnitude-aware, advantage-weighted rectification is the more stable realization (Fu et al., 29 Jun 2026). DenseGRPO, ReFPO, FlowTTS-GRPO, and related methods do not always use the name as a formal method label, yet they are repeatedly interpreted as RWFM because rewards or advantages weight the per-sample or per-timestep flow update (Deng et al., 28 Jan 2026, Wang et al., 19 Jun 2026, Wang et al., 22 Jun 2026).
This suggests that RWFM is no longer a single loss but an umbrella concept for reward-conditioned regression in flow-based generators. What remains invariant is the attempt to preserve the analytic regression structure of flow matching while replacing reward-agnostic supervision by reward-aware credit assignment.
2. Flow-matching foundations and reward-tilted objectives
Flow matching models learn a time-dependent velocity field that transports a simple prior to a target data distribution by integrating an ODE,
In rectified-flow form, one commonly samples a clean datum and Gaussian noise, constructs a linear path such as or , and regresses the model to a closed-form target velocity, yielding the standard supervised objective
The notation differs across papers—some index clean data by , others by —but the structure is the same: pretraining is regression against an analytically known target induced by a fixed interpolation kernel (Deng et al., 28 Jan 2026, Bergmeister et al., 11 May 2026, Xue et al., 29 Sep 2025).
RWFM augments this foundation with a reward-tilted target distribution. A recurring formal objective is
0
or, equivalently in some papers, a KL-constrained optimum in which the aligned endpoint distribution is an exponential tilt of a reference model while the noising law or conditional bridge is preserved (Ouyang et al., 31 Jan 2026, Bergmeister et al., 11 May 2026, Holderrieth et al., 5 Feb 2026). Under this view, alignment is not an arbitrary modification of the sampler; it is sampling from a reward-weighted distribution, and the central technical question becomes how to encode that tilt into the velocity field.
A basic RWFM objective therefore takes the form
1
where 2 is derived from reward, advantage, or a stabilized transform of either (Deng et al., 28 Jan 2026). Later formulations replace the scalar multiplier by a corrected regression target rather than a plain weight, but the reward-tilting interpretation remains the organizing principle.
3. Principal formulations
Three main formulations recur in the literature: direct sample reweighting, analytically corrected regression targets, and advantage- or dense-reward-weighted updates. Direct sample reweighting is the simplest and most literal RWFM. The robotics RWFM scheme uses 3 to weight the ordinary flow-matching loss on action chunks, while ORW-CFM and ORW-CFM-W2 use 4 or 5 with online sampling from the current model (Pfrommer et al., 20 Jul 2025, Fan et al., 9 Feb 2025). HapticVLA’s safety-aware RWFM belongs to the same class, except that the weight is built from robustly normalized blends of chunk-level and episode-level tactile rewards, then clipped, renormalized, and combined with an anchor penalty toward the imitation-initialized parameter vector (Gubernatorov et al., 16 Mar 2026).
A second family derives the aligned regression target from KL-regularized optimal control. RAM gives the clearest closed-form expression:
6
where the target is reference-anchored and reward-corrected, with stop-gradient applied to the target branch (Bergmeister et al., 11 May 2026). FlowAWR derives a related but centered expression for the optimal field,
7
with
8
so the update is an advantage-weighted rectification of a reference velocity rather than a raw reward multiplier (Fu et al., 29 Jun 2026). AWM makes the same unification claim from the policy-gradient side: it keeps the pretraining flow-matching loss intact and weights it by group-relative advantage, thereby obtaining a lower-variance alternative to DDPO-style reverse-step likelihood training (Xue et al., 29 Sep 2025).
A third family emphasizes trajectory structure. DenseGRPO replaces sparse terminal rewards with per-step dense reward gains. It maps intermediate noisy states to clean-image proxies with an ODE rollout, defines
9
and then uses 0 in per-timestep GRPO updates; the paper explicitly presents this as a natural RWFM view of flow-model alignment (Deng et al., 28 Jan 2026). FlowTTS-GRPO adopts an SDE-path formulation instead: rewards are fused across speaker similarity, intelligibility, and perceptual quality, normalized by standard deviation, converted into group-relative advantages, and then used in a clipped GRPO objective on stochastic FM trajectories (Wang et al., 22 Jun 2026). ReFPO provides a control-theoretic variant in which FPO gradients are interpreted as an implicit advantage-weighted Reflow process; its central claim is that flow-matching policy gradients already constitute reward-/advantage-weighted flow matching, and that an additional unweighted Reflow regularizer is needed to stabilize the geometry (Wang et al., 19 Jun 2026).
| Method | Reward signal in the FM update | Distinguishing mechanism |
|---|---|---|
| RWFM for flow-matching policies | 1 | Direct sample-weighted FM on action chunks |
| ORW-CFM-W2 | 2 | Online sampling plus Wasserstein-2 regularization |
| DenseGRPO | 3 or 4 | Per-timestep dense rewards and reward-aware exploration |
| RAM | Reward-corrected target | Reference-anchored consistency regression |
| FlowAWR | Centered 5 | Advantage-weighted rectification of 6 |
| AWM | Group-relative advantage | Same FM loss as pretraining, lower-variance RL surrogate |
4. Representative systems and empirical record
Text-to-image alignment is the most developed RWFM application area. DenseGRPO reports improvements over Flow-GRPO on DrawBench prompts under the PickScore metric from 7 to 8, with Aesthetic improving from 9 to 0, ImageReward from 1 to 2, GenEval from 3 to 4, and OCR accuracy from 5 to 6; its ablations attribute the gains to valid dense rewards, timestep-specific stochasticity calibration, and longer ODE rollouts for intermediate clean estimation (Deng et al., 28 Jan 2026). RAM reports that, on Stable Diffusion 3.5 Medium, it achieves the highest reward on composability, text rendering, and human preference and reaches Flow-GRPO’s peak reward in up to 7 fewer training steps; the reported held-out scores are GenEval 8, OCR 9, and PickScore 0 (Bergmeister et al., 11 May 2026). FlowAWR reports improved alignment performance alongside a 1 to 2 convergence acceleration over DiffusionNFT, including reaching a 3 PickScore in 4k steps, versus 5 in 6k steps for DiffusionNFT and 7 in 8k steps for FlowGRPO (Fu et al., 29 Jun 2026). AWM, which treats advantage-weighted flow matching as the natural RL continuation of pretraining, reports up to a 9 speedup over Flow-GRPO on GenEval, OCR, and PickScore (Xue et al., 29 Sep 2025).
The training-free branch of RWFM pursues the same reward-tilted distribution without any parameter update. For flow matching, the alignment framework based on reward-weighted distributions reports that, on SD3.5 Large Turbo, its training-free velocity guidance improves PickScore from 0 to 1, HPSV2 from 2 to 3, ImageReward from 4 to 5, and Aesthetic from 6 to 7, with no extra training; for diffusion, the corresponding finetuning-free score-guidance network attains comparable performance to finetuning-based baselines with one-step generation and at least a 8 reduction in computational cost (Ouyang et al., 31 Jan 2026). Diamond Maps occupy a related but distinct position: rather than learning the reward-tilted field during post-training, they estimate the value function and its gradient efficiently at inference time via stochastic one-step maps, enabling guidance, search, and sequential Monte Carlo for arbitrary rewards (Holderrieth et al., 5 Feb 2026).
Outside image generation, RWFM has been adapted to sequence generation and embodied control. FlowTTS-GRPO reports objective and subjective preference gains on CosyVoice 3.0 and F5-TTS, with F5-TTS improving intelligibility and both models improving speaker similarity and perceptual quality under multi-objective reward fusion (Wang et al., 22 Jun 2026). In simulated unicycle tasks, the original flow-matching-policy paper reports that both RWFM and GRPO dramatically improve upon a suboptimal demonstrator and that the GRPO approach generally incurs between 9 and 0 less cost than a naive imitation-learning flow-matching approach (Pfrommer et al., 20 Jul 2025). For continuous-control flow policies, ReFPO reports improvements in average performance and discretization robustness across GridWorld, MuJoCo Playground, and Humanoid Control, including MuJoCo results of 1 at 10-step inference and 2 at 1-step inference for ReFPO versus 3 and 4 for FPO, alongside lower Straightness Error and Explosion Rate (Wang et al., 19 Jun 2026). In contact-rich manipulation, HapticVLA’s safety-aware RWFM plus tactile distillation attains a mean success rate of 5 on real-world tasks while outperforming baseline VLAs, including versions with direct tactile feedback at inference (Gubernatorov et al., 16 Mar 2026).
5. Limitations, controversies, and stabilization strategies
A persistent limitation is reward dependence. DenseGRPO states that quality hinges on the reward model 6, and that more accurate dense rewards can also increase reward-hacking risk; it recommends mixed metrics such as Aesthetic, DeQA, and UnifiedReward for monitoring (Deng et al., 28 Jan 2026). RAM makes a parallel point: the method aligns to whatever reward is provided, the path-cost correction is dropped for variance and computation reasons, and the resulting estimator is exact only at initialization and in linear-Gaussian cases, while remaining first-order accurate more generally (Bergmeister et al., 11 May 2026). FlowTTS-GRPO likewise relies on proxy models for speaker similarity, intelligibility, and perceptual quality, so reward fusion and standard-deviation normalization become part of the optimization problem rather than a neutral preprocessing choice (Wang et al., 22 Jun 2026).
A second limitation is collapse under repeated online reward weighting. The flow-matching-policy paper notes that iterative weighted FM evolves as
7
which concentrates on 8 as 9 grows and can reduce conditional specificity if 0 is too large (Pfrommer et al., 20 Jul 2025). ORW-CFM formalizes the same phenomenon in the online setting: without regularization, the induced distribution converges to a Dirac delta at the maximizer of the reward, so policy collapse is a theorem rather than merely an empirical hazard (Fan et al., 9 Feb 2025). The practical response is regularization. Different papers instantiate this with Wasserstein-2 penalties on vector fields, KL penalties to a reference policy, reference anchors in the regression target, parameter anchoring to the imitation model, clipped or normalized advantages, and adaptive temperatures or noise schedules (Fan et al., 9 Feb 2025, Bergmeister et al., 11 May 2026, Gubernatorov et al., 16 Mar 2026, Fu et al., 29 Jun 2026).
A more fundamental controversy concerns whether reward-weighted regression under training-time marginals can ever be a sufficient surrogate for inference-time quality. DRL argues that flow- and score-matching losses are measured under interpolation marginals 1 rather than model rollout marginals 2, and proves a “No reward certificate from FSM” result: arbitrarily small training error under 3 can coexist with poor reward at the rollout endpoint (Beltran-Velez et al., 17 Jun 2026). In the same paper, attempts to distill an RL teacher back into ordinary flow matching fail to recover the teacher’s reward and quality, which is used as evidence that pure reward-weighted regression can inherit a geometry mismatch between velocity-space 4 optimization and sample-space reward (Beltran-Velez et al., 17 Jun 2026). This critique does not invalidate RWFM, but it narrows its safe operating regime and motivates on-policy updates, reference anchoring, and explicit reward-gradient or control-theoretic corrections.
Compute is the final major constraint. DenseGRPO notes that ODE rollouts for intermediate clean-image estimates are expensive: for 5 training steps, 6, 7, and 8 require approximately 9, 0, and 1 GPU hours, respectively, with larger 2 improving reward fidelity and downstream results (Deng et al., 28 Jan 2026). This is one reason why RAM, FlowAWR, and AWM emphasize preserving the regression structure of pretraining while avoiding SDE rollouts, backward adjoint sweeps, or noisy reverse-step likelihoods (Bergmeister et al., 11 May 2026, Fu et al., 29 Jun 2026, Xue et al., 29 Sep 2025).
6. Related paradigms and future directions
RWFM sits between training-time fine-tuning and inference-time guidance. The training-free velocity-guidance framework for text-to-image generation treats alignment itself as sampling from a reward-weighted distribution and derives an exact velocity correction for flow matching, so that the probability-flow ODE can be modified without any finetuning (Ouyang et al., 31 Jan 2026). Diamond Maps pursue the same reward-tilted target from another direction: instead of modifying the model during training, they estimate the value function
3
and then assemble the guided drift 4 at inference time (Holderrieth et al., 5 Feb 2026). This suggests a continuum in which one endpoint is pure RWFM post-training and the other is exact or approximate reward guidance without retraining.
Several open problems are explicit in the recent literature. FlowAWR identifies tighter convergence guarantees beyond gradient-equivalence, improved estimators for 5 and the centered advantage in high-variance reward landscapes, extensions to other schedules and non-Gaussian kernels, joint multi-reward optimization without expert branching, and theoretical bounds on advantage dynamics and out-of-domain generalization (Fu et al., 29 Jun 2026). DenseGRPO highlights the need for reliable dense rewards and better exploration calibration over time (Deng et al., 28 Jan 2026). DRL raises a broader question: whether the most effective “reward-weighted flow matching” may, in some settings, require abandoning static sample reweighting in favor of on-policy KL-regularized control over model-generated states (Beltran-Velez et al., 17 Jun 2026).
Taken together, these developments show that RWFM is best understood not as a single algorithmic recipe but as a design principle: preserve the tractable supervised structure of flow matching, but make its regression signal responsive to reward. The literature differs on how that responsiveness should be encoded—through exponential sample weights, centered advantages, reward-corrected targets, dense per-step gains, path-likelihood weighting, or inference-time value guidance—but all variants seek the same outcome: an aligned velocity field that transports mass toward preferred, safer, or more semantically faithful generations while retaining the scalability of flow-matching pretraining (Pfrommer et al., 20 Jul 2025, Bergmeister et al., 11 May 2026, Fu et al., 29 Jun 2026).