Reinforcement learning can align visual generators with reward signals beyond their pretraining objectives, but existing score-function methods typically introduce stochasticity by perturbing the reverse denoising trajectory, leading to costly rollout-based optimization. We instead ask whether stochastic exploration can be introduced inside the generator while preserving its original sampler. Routing Policy Optimization (RPO) answers this question by turning attention routing into a stochastic exploration action. At a few selected denoising levels, each query block keeps its own block and samples further key blocks in proportion to the generator's native attention mass, or keeps dense attention. The sampled route perturbs the velocity of that transition, and a Gaussian likelihood around the dense velocity scores the perturbation with a score-function objective. RPO further uses a controlled two-rollout estimator: two rollouts share the same prompt, initial latent, routed denoising levels, and dense anchor layers, differing only in their sampled attention routes. Their reward difference provides an advantage signal for the routed transitions. Because the likelihood differentiates the dense velocity itself, the learned updates remain effective after routing is removed. At inference, RPO restores dense attention and retains the original deterministic sampler. Experiments on text-to-image and text-to-video generation show that RPO improves generation quality while achieving nearly 4x higher training efficiency than Flow-GRPO-style methods. Please see our project page at https://velox-world.com/rpo for visual results.
Figures & tables
Figure 1: A controlled routing pair replaces a large sampler-noise group. Sampler-space GRPO injects continuous noise into the reverse process and estimates an advantage from many complete visual rollouts. RPO retains the original deterministic sampler and pairs two rollouts that share the prompt, initial visual latent, routed denoising levels, and anchor layers. They differ only in their sampled attention routes, so their reward difference measures the aggregate effect of the two recorded routed trajectories. Generated visual outputs are schematic.
Figure 2: A temporary route perturbs training rollouts, and a velocity likelihood updates the dense generator. At each routed level t , the native attention mass of every query block over the key blocks defines a categorical draw over the candidate blocks Cq and a dense option, while the safety block Sq stays active. Two matched rollouts record independent routes under the same prompt, content seed, routed levels, and anchor layers. Replay recomputes the dense velocity at every recorded routed latent, and the reward difference weights a Gaussian likelihood of the routed step around it. At inference, RPO removes all masks and runs dense attention with the original deterministic sampler.
Figure 3: RPO on Stable Diffusion 3.5 Medium. (a) Samples before and after RPO on DrawBench prompts. Each pair shares the prompt and initial noise. (b, c) DrawBench scores over 1,000 images from 200 prompts with 5 seeds, generated without classifier-free guidance. RPO improves desired property across all five reported metrics.
Figure 4: RPO learns faster than Flow-GRPO on Stable Diffusion 3.5 Medium. Both methods train the same LoRA adapters with HPSv3 reward and a 10-step training sampler for 50 steps from the same base model. (a) Mean raw HPSv3 of each step’s training rollouts (faint) and its exponential moving average with smoothing weight 0.6 (bold). (b) Raw HPSv3 of the dense model on 64 fixed validation prompts with a 28-step sampler, measured every 5 steps from step 0.
Method
Train GPU-h
Aesthetic quality
Imaging quality
Subject consistency
Motion smoothness
Object class
Multiple objects
Human action
Scene
Overall score
CogVideoX-5B
–
60.17
61.92
96.50
97.42
82.44
57.29
86.40
42.73
81.30
HunyuanVideo
–
63.05
63.78
98.76
99.42
65.79
56.46
83.20
27.01
78.69
Wan2.2-TI2V-5B
–
61.91
67.71
98.37
98.76
81.01
59.15
79.40
37.41
81.61
RPO
24.1
62.64
68.65
98.69
99.06
85.41
61.30
84.20
40.84
82.00
DanceGRPO
102.3
61.74
67.69
98.38
98.75
80.81
59.34
80.00
36.95
81.46
Flow-GRPO
75.4
60.09
67.43
97.59
97.88
79.10
59.98
81.80
36.73
81.19
Table 1: Comparison on VBench across multiple dimensions with DanceGRPO and Flow-GRPO, plus overall score and training cost. Scores are percentages. The Wan2.2 rows share one generation configuration. CogVideoX-5B and HunyuanVideo use the same prompts, frame count, denoising steps, and scoring, but keep their native resolution and guidance. Training GPU-hours are measured on H100 GPUs.
Figure 5: Text-to-video samples from Wan2.2-TI2V-5B before and after RPO. Each row shows one prompt. For each model we show the first and last of the 13 generated frames ( 1280×704 , 50 sampling steps). Both models use the same prompt and initial noise. RPO follows the prompt more closely and replaces distorted, painterly textures with coherent, realistic scenes.
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Gilbert blocks and the safety set on toy 16×16 frames with 8-token blocks. (a) Raster order turns each block into a horizontal strip and jumps back at every row end (red dotted segments). (b) The frame-sliced Gilbert order turns each block into a compact 2×4 or 4×2 patch. The curve continues from the end of frame f to the same position in frame f+1 (dashed arrow) and traverses that frame in a reflected direction. (c) The route of one query block q . The block is its own safety set Sq={q} , and a sampled route adds candidate blocks from any frame, here two (blue). All other blocks remain in the candidate set Cq but stay masked for this draw.
Wan2.2
Stable Diffusion 3.5
Output size
13×704×1280
512×512
Token grid F×H×W
4×22×40
1×32×32
Visual tokens M
3,520
1,024
Block size B
128
128
Blocks Q
28
8
Safety set ∣Sq∣
1
1
Appendix
Table 2: Token grids, blocks, and routing options of the reference configurations. Counts cover visual tokens only. Routed levels index the transitions of the 10-step training sampler, counted from 0.
Group-relative RL methods such as Flow-GRPO post-train image generators by exploring with isotropic Gaussian noise added at every denoising step. This noise decides which rollouts the model learns from, yet it perturbs every channel and spatial position of the latent equally. In this paper, we instead show that latent elements differ in how much they change the generated image, so exploration should adapt to these differences. We introduce EXPLORENET to learn an adaptive exploration distribution. EXPLORENET is a policy that predicts a noise scale for every latent element from the current latent, the denoising step, and the prompt, before any reward is observed; it is trained on the reward spread of each rollout group and discarded after training, leaving inference unchanged. On Stable Diffusion 3.5 Medium, EXPLORENET improves held-out GenEval2 by 14% over Flow-GRPO, transfers to two independent compositional benchmarks and five preference and image-quality models, and reaches a 67.2% human preference win-rate. Overall, across our group-relative diffusion RL experiments, we find that exploration is learnable, the shape of the exploration distribution outweighs its magnitude, and rollout quality is more effective than rollout quantity.
Aligning denoising generative models with human preferences or verifiable rewards remains a key challenge. While policy-gradient online reinforcement learning (RL) offers a principled post-training framework, its direct application is hindered by the intractable likelihoods of these models. Prior work therefore either optimizes an induced Markov decision process (MDP) over sampling trajectories, which is stable but inefficient, or uses likelihood surrogates based on the diffusion evidence lower bound (ELBO), which have so far underperformed on visual generation. Our key insight is that the ELBO-based approach can, in fact, be made both stable and efficient. By reducing surrogate variance and controlling gradient steps, we show that this approach can beat MDP-based methods. To this end, we introduce Variational GRPO (V-GRPO), a method that integrates ELBO-based surrogates with the Group Relative Policy Optimization (GRPO) algorithm, alongside a set of simple yet essential techniques. Our method is easy to implement, aligns with pretraining objectives, and avoids the limitations of MDP-based methods. V-GRPO achieves state-of-the-art performance in text-to-image synthesis, while delivering a 2× speedup over MixGRPO and a 3× speedup over DiffusionNFT.
Bingda Tang, Yuhui Zhang, Xiaohan Wang +3
Stanford University · Tsinghua University · Amazon FAR +1
Reinforcement learning is increasingly used to align image generators with reward signals, and Flow-GRPO recently extended this paradigm to flow-matching models by treating the denoising sampler as a stochastic policy that can be optimized from reward feedback. Training in this setting is unstable in a way specific to multi-step denoising: the policy update changes systematically across denoising steps, with importance ratios drifting below one, becoming increasingly dispersed, clipping at different rates, and leaving fewer usable samples late in training. Prior work treats these effects as separate failure modes and addresses each with a hand-tuned stabilizer. We show instead that they arise from a single per-step quantity, which we call path variance. This quantity is determined exactly by the sampler's Gaussian transition kernel and can be estimated cheaply during training. This reframes instability as a resource that can be measured and budgeted rather than a collection of symptoms to repair. Our method, λ-Controlled GRPO, calibrates importance-ratio behavior from this predicted law rather than from noisy empirical statistics, and allocates gradient effort across denoising steps according to their predicted cost. The two scales governing the update are fixed by standard policy choices rather than introduced as free tuning parameters. On a text-to-image model under two reward settings, rendering difficult target text scored by optical character recognition and matching human preferences scored by a preference model, λ-Controlled GRPO improves both text accuracy and preference reward over the strongest empirical stabilizer. It also keeps late-step path variance within its intended budget, precisely where the baseline systematically overshoots. The result is a Flow-GRPO update calibrated by its own transition law rather than stabilized after instability appears.
Yufeng Wang, Parivesh Priye, Meeshawn Marathe +1
Stony Brook University · Georgia Institute of Technology · Rivian and Volkswagen Group Technologies