Optimizing Stochastic Policy Geometry for Few-Step Flow Reinforcement Learning

Zhuguanyu Wu; Congjia Chen; Ruihao Gong; Yang Yong; Shiqiao Gu; Yi Xu; Xianglong Liu

cs.AIOct 30, 2026

Optimizing Stochastic Policy Geometry for Few-Step Flow Reinforcement Learning

Authors: Zhuguanyu Wu, Congjia Chen, Ruihao Gong, Yang Yong, Shiqiao Gu, Yi Xu, Xianglong Liu

Organizations: Beihang University · Sensetime Research

Abstract

Few-step flow models generate samples through a handful of large denoising transitions. Applying GRPO to such models requires introducing stochasticity into the denoising transitions, and when a single transition spans a wide range of noise levels, the choice of the resulting stochastic policy becomes a major factor in optimization. Recent constructions constrain the policy to preserve the target flow coefficients, but this constraint admits a family of policies and offers no criterion for choosing among them. We introduce Geo-FSRL (Geometry-Optimized Few-Step Reinforcement Learning), built on a selection principle: at a fixed conditioning state and timestep pair, choose the admissible policy that minimizes the conditional KL divergence induced by any given nonzero velocity-output update. Within a coefficient-preserving family of isotropic Gaussian transitions, this criterion reduces to minimizing a single scalar geometry factor that determines the conditional KL for every velocity-output update. The same factor governs velocity-space Fisher information and conditional transition log-likelihood-ratio variance. We derive the unique closed-form optimizer within this family, shared across all nonzero update magnitudes and directions, with transition coefficients determined by both the source and target noise levels. Because the minimized geometry factor still varies across timesteps, we further adopt gradient-scale matching from prior work and derive closed-form policy-gradient and KL weights that normalize explicit score scales and reference-KL curvature. Geo-FSRL trains on the original few-step schedule and retains the deterministic solver at inference. Our method achieves higher final reward scores than Flow-GRPO and Flow-CPS on both image generation and video generation models under the 4-step setting. Code is available at https://github.com/ModelTC/LightX2V/tree/main/lightx2v_rl

Figures & tables

Figure 1: Four-step image validation rewards versus optimizer updates: (a) HPSv3-only; (b,c) joint HPSv3 and CLIPScore. Flow-CPS curves use matched weights for their respective $\phi$ values. DiffusionNFT denotes the tested four-step adaptation.

	Single-Reward	Multi-Reward
Method	HPSv3	CLIPScore	HPSv3
Qwen-Image-Lightning	12.17	0.971	12.17
Flow-GRPO ( Liu et al., 2025 )	13.62 (+1.45)	0.996 (+0.025)	13.45 (+1.28)
Flow-CPS ( Wang and Yu, 2025 )	13.90 (+1.73)	0.988 (+0.017)	13.58 (+1.41)
Geo-FSRL (Ours)	14.05 (+1.88)	1.008 (+0.037)	13.80 (+1.63)

Table 1: Final validation reward scores for single-reward setting and multi-reward setting text-to-image training.

	Visual Quality (%)	Overall Preference (%)
Flow-CPS	46	42
Geo-FSRL (Ours)	54	58

Table 2: Human evaluation of 20 image pairs by five evaluators blinded to method identity. We report the percentage of preference votes for each method.

Method	Without policy reweighting	With policy reweighting	Gain
Flow-CPS	13.69	13.90	+0.21
Geo-FSRL (Ours)	13.61	14.05	+0.44

Table 3: Single-reward policy-weighting ablation on the validation set. Each method uses its own matched weights; KL weighting remains enabled throughout. Gain is the change from enabling policy weighting.

Figure 2: Four-step video validation rewards versus optimizer updates under joint optimization. Flow-CPS uses matched weights; DiffusionNFT denotes the tested four-step adaptation.

Method	HPSv3-General	HPSv3-Percentile	VideoAlign-MQ
Wan2.1-T2V-1.3B	4.13	8.85	0.25
Flow-GRPO ( Liu et al., 2025 )	7.29 (+3.16)	12.10 (+3.25)	1.05 (+0.80)
Flow-CPS ( Wang and Yu, 2025 )	8.15 (+4.02)	11.94 (+3.09)	1.26 (+1.01)
Geo-FSRL (Ours)	9.24 (+5.11)	12.90 (+4.05)	2.07 (+1.82)

Table 4: Final validation reward scores for joint text-to-video training.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Transition	$\alpha_{t,s}$	$\beta_{t,s}$	$\alpha_{t,s}^{2}+\beta_{t,s}^{2}-s^{2}$
ODE Euler	$0$	$s$	$0$
Flow-GRPO (Euler–Maruyama)	$\sigma_{t}\sqrt{\Delta_{t,s}}$	$s-\dfrac{\Delta_{t,s}\sigma_{t}^{2}}{2t}$	$\dfrac{\Delta_{t,s}^{2}\sigma_{t}^{2}}{t}\left(1+\dfrac{\sigma_{t}^{2}}{4t}\right)$
Flow-CPS	$s\sin(\pi\phi/2)$	$s\cos(\pi\phi/2)$	$0$

Appendix

Table 5: Finite transitions from

t

s

(

0<s<t<1

) in the form of Eq. 4 . The last column is the deviation from the coefficient-preserving condition

\alpha_{t,s}^{2}+\beta_{t,s}^{2}=s^{2}

. The Flow-GRPO row is the standard uncapped Euler–Maruyama construction at interior timesteps; endpoint and noise-cap variants are given in Appendix A.1 .

Setting	Text-to-image	Text-to-video
Optimizer updates	4,000	2,000
Optimizer	AdamW	AdamW
Learning rate	$2\times 10^{-6}$	$2\times 10^{-6}$
Adam $(\beta_{1},\beta_{2})$	$(0.9,0.999)$	$(0.9,0.999)$
Adam $\epsilon$	$10^{-8}$	$10^{-8}$
Weight decay	$10^{-4}$	$10^{-4}$

Appendix

Table 6: Training hyperparameters for the image and video GRPO-based experiments. The geometry multiplier

c_{\mathrm{geo}}

applies to the matched weights of Geo-FSRL and Flow-CPS.

	Selected split ( $\phi^{\star}$ )			Flow-CPS ( $\phi=0.3$ )			Flow-CPS ( $\phi=0.7$ )
$(t,s)$	$\alpha^{\star}$	$\beta^{\star}$	$G^{\star}$	$\alpha$	$\beta$	$G$	$\alpha$	$\beta$	$G$
$(1,15/16)$	0.93750	0.00000	0.06667	0.42562	0.83532	0.14685	0.83532	0.42562	0.07482
$(15/16,5/6)$	0.78567	0.27778	0.17678	0.37833	0.74251	0.29034	0.74251	0.37833	0.17859
$(5/6,5/8)$	0.58926	0.20833	0.47140	0.28374	0.55688	0.77424	0.55688	0.28374	0.47624

Appendix

Table 7: Analytic Wan coefficients (five decimal places). Flow-CPS with

\phi=0.3

is the reported video baseline;

\phi=0.7

is shown for analytic comparison.

Figure 3: Validation HPSv3 with and without policy reweighting for (a) Geo-FSRL and (b) Flow-CPS. All curves use the same 256 validation prompts; horizontal axes count optimizer updates. KL reweighting remains enabled throughout. Scores are reported in Table 3 .

Figure 4: Qwen-Image-Lightning, Flow-CPS, and Geo-FSRL on the displayed prompts. Multi-reward training optimizes HPSv3 and CLIPScore; HPSv3-only training optimizes HPSv3 alone.

Figure 5: Additional text-to-image examples from the Qwen-Image-Lightning base model and Geo-FSRL under joint HPSv3 and CLIPScore training or HPSv3-only training. Each row uses the displayed prompt.

Figure 6: Selected frames from videos generated by the four-step distilled Wan2.1 base model and Geo-FSRL under joint reward optimization. For each prompt, the upper row shows the base model and the lower row (“After FSRL”) shows Geo-FSRL, with four frames per video. Some prompts are shortened for display.

Figure 7: Additional video comparisons under joint reward optimization. Each prompt has two rows of four selected frames: the four-step distilled Wan2.1 base model above and Geo-FSRL (“After FSRL”) below. Some prompts are shortened for display.

RL for Flow Matching Models

See all 22 papers

Flow-Map GRPO: Reinforcement Learning for Few-Step Flow-Map Generators via Anchored Stochastic Composition
Jul 1, 2026Zhiqi Li, Wen Zhang, Bo ZhuImage Generation Flow Map Learning
Georgia Institute of Technology, Atlanta, GA, USA.
$λ$ -Controlled GRPO: Turning Flow-Matching Ratio Instability into a Budgeted Resource
Sep 18, 2026Yufeng Wang, Parivesh Priye, Meeshawn Marathe +1Flow Matching T2I Generation
Stony Brook University · Georgia Institute of Technology · Rivian and Volkswagen Group Technologies
AdvantageFlow: Regularized Advantage-Weighted RL in Flow Models
May 25, 2026Branislav Kveton, Anup Rao, Subhojyoti Mukherjee +2Rectified Flow RL for Image Generation
Adobe Research

Jul 1, 2026J/K move · Enter open · S save

Flow-Map GRPO: Reinforcement Learning for Few-Step Flow-Map Generators via Anchored Stochastic Composition

Zhiqi Li, Wen Zhang, Bo Zhu

Georgia Institute of Technology, Atlanta, GA, USA.

Image Generation Flow Map Learning Flow-Based Generative Modeling RL for Generative Models

Open paper AskPDF arXiv

Few-step flow-map generators, such as consistency models and MeanFlow, accelerate sampling by learning long-range transport maps between noise and data. However, their deterministic transitions do not directly provide the stochastic trajectories and tractable likelihood ratios required by reinforcement learning (RL) post-training. Existing SDE-based stochasticization techniques target velocity-based samplers and do not directly extend to long-range flow-map transitions. We propose Flow-Map GRPO, an online RL post-training framework for deterministic few-step flow-map generators. Its key component, Anchored Stochastic Flow Map Composition (ASFMC), combines deterministic transport with anchor-based conditional resampling. We establish the conditions under which this construction preserves the marginal probability path and develop tractable local- and endpoint-anchor policies for two-time and single-time flow maps. These policies enable a unified GRPO training procedure. Experiments on FLUX-based MeanFlow and sCM generators demonstrate substantial improvements in OCR, PickScore, and GenEval at different numbers of inference steps, including joint OCR--PickScore gains with mixed rewards. Controlled ablations show that the stochastic transition design is essential for translating training rewards into generation quality. Flow-Map GRPO enables effective RL alignment of pretrained deterministic flow-map generators while retaining their original parameterization, without retraining them as native stochastic models.

Abstract

RL for Flow Matching Models

Explore similar work

Flow-Map GRPO: Reinforcement Learning for Few-Step Flow-Map Generators via Anchored Stochastic Composition

λλλ-Controlled GRPO: Turning Flow-Matching Ratio Instability into a Budgeted Resource

AdvantageFlow: Regularized Advantage-Weighted RL in Flow Models

$λ$ -Controlled GRPO: Turning Flow-Matching Ratio Instability into a Budgeted Resource