Routing Policy Optimization for Online Diffusion Reinforcement Learning

Qiuhong Shen; Yiding Yang; Yuxin Zhang; Xinwei Huang; Minxuan Lin; Bo Liu; Hao Kang; Haotian Yang

cs.CVOct 15, 2026

Routing Policy Optimization for Online Diffusion Reinforcement Learning

Authors: Qiuhong Shen, Yiding Yang, Yuxin Zhang, Xinwei Huang, Minxuan Lin, Bo Liu, Hao Kang, Haotian Yang,

Organizations: National University of Singapore · ByteDance

Abstract

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 $C_{q}$ and a dense option, while the safety block $S_{q}$ 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.

Panel of Figure 4 — 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\times 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\times 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\times 4

4\times 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

S_{q}=\{q\}

, and a sampled route adds candidate blocks from any frame, here two (blue). All other blocks remain in the candidate set

C_{q}

but stay masked for this draw.

	Wan2.2	Stable Diffusion 3.5
Output size	$13\times 704\times 1280$	$512\times 512$
Token grid $F\times H\times W$	$4\times 22\times 40$	$1\times 32\times 32$
Visual tokens $M$	3,520	1,024
Block size $B$	128	128
Blocks $Q$	28	8
Safety set $\|S_{q}\|$	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.

RL for Image Generation

Reinforcement learning for image generation trains visual models with reward feedback to improve prompt alignment, quality, and control.

See all 93 papers

ExploreNet: Learning Where to Explore in Diffusion GRPO
Sep 29, 2026Shuyue Stella Li, Xiaochuang Han, Yulia Tsvetkov +1Diffusion Model Fine-Tuning RL Exploration
University of Washington
V-GRPO: Online Reinforcement Learning for Denoising Generative Models Is Easier than You Think
Apr 25, 2026Bingda Tang, Yuhui Zhang, Xiaohan Wang +3Diffusion Model Alignment Group Relative Policy Optimization
Stanford University · Tsinghua University · Amazon FAR +1
$λ$ -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

Sep 29, 2026J/K move · Enter open · S save

ExploreNet: Learning Where to Explore in Diffusion GRPO

Shuyue Stella Li, Xiaochuang Han, Yulia Tsvetkov, Luke Zettlemoyer

University of Washington

Diffusion Model Fine-Tuning RL Exploration RL for Image Generation RL for Diffusion Models T2I Diffusion Models

Open paper AskPDF arXiv

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.

Abstract

RL for Image Generation

Explore similar work

ExploreNet: Learning Where to Explore in Diffusion GRPO

V-GRPO: Online Reinforcement Learning for Denoising Generative Models Is Easier than You Think

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

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