FlowRL: A Taxonomy and Modular Framework for Reinforcement Learning with Diffusion Policies

Chenxiao Gao; Edward Chen; Tianyi Chen; Bo Dai

cs.LGMar 29, 2026

FlowRL: A Taxonomy and Modular Framework for Reinforcement Learning with Diffusion Policies

Authors: Chenxiao Gao, Edward Chen, Tianyi Chen, Bo Dai

Organizations: Georgia Institute of Technology

Abstract

Thanks to their remarkable flexibility, diffusion models and flow models have emerged as promising candidates for policy representation. However, efficient reinforcement learning (RL) upon these policies remains a challenge due to the lack of explicit log-probabilities for vanilla policy gradient estimators. While numerous attempts have been proposed to address this, the field lacks a unified perspective to reconcile these seemingly disparate methods, thus hampering ongoing development. In this paper, we bridge this gap by introducing a comprehensive taxonomy for RL algorithms with diffusion/flow policies. To support reproducibility and agile prototyping, we introduce a modular, JAX-based open-source codebase that leverages JIT-compilation for high-throughput training. Finally, we provide systematic and standardized benchmarks across Gym-Locomotion, DeepMind Control Suite, and IsaacLab, offering a rigorous side-by-side comparison of diffusion-based methods and guidance for practitioners to choose proper algorithms based on the application. Our work establishes a clear foundation for understanding and algorithm design, a high-efficiency toolkit for future research in the field, and an algorithmic guideline for practitioners in generative models and robotics. Our code is available at https://github.com/typoverflow/flow-rl.

Figures & tables

Guidance Method	Reference Policy	Algorithm Name
BoN Sampling	$\pi_{\mathcal{D}}$	IDQL ( Hansen-Estruch et al., 2023 ) , SfBC ( Chen et al., 2022 )
Q-value Guidance	$\text{Unif}(\mathcal{A})$	QSM ( Psenka et al., 2024 ) , iDEM ( Akhound-Sadegh et al., 2024 ) , DPS ( Jain et al., 2024 )
Q-value Guidance	$\pi_{\mathcal{D}}$	DAC ( Fang et al., 2024 ) , QGPO ( Lu et al., 2023 ) , Diffusion-DICE ( Mao et al., 2024 )
Weighted Matching	$\text{Unif}(\mathcal{A})$	SDAC ( Ma et al., 2025 ) , MaxEntDP ( Dong et al., 2025 ) , QVPO ( Ding et al., 2024 )
	$\pi^{k-1}$	DPMD ( Ma et al., 2025 ) , FPMD ( Chen et al., 2025 ) , GeMPO ( Ma et al., 2026 )
	$\pi_{\mathcal{D}}$	QIPO ( Zhang et al., 2024 )

Table 1: Summarization of existing DPRL algorithms, based on their guidance mechanism and reference policy.

Figure 2: Training curves of several off-policy DPRL algorithms across Gym-Locomotion tasks.

Figure 3: Training curves of on-policy DPRL algorithms on IsaacLab tasks.

Figure 4: Performance profile on Gym-Locomotion tasks.

Dataset	IQL	IDQL †	QGPO †	QIPO †	Diffusion-QL	DAC	BDPO
halfcheetah-m	47.4	51.0	54.1	54.2 ±1.3	50.7 ±0.7	59.1 ±0.6	71.2 ±0.9
hopper-m	66.3	65.4	98.0	94.0 ±13.3	80.4 ±15.7	103.5 ±0.3	100.6 ±0.7
walker2d-m	78.3	82.5	86.0	87.6 ±1.5	86.7 ±1.6	97.9 ±0.9	93.4 ±0.5
halfcheetah-m-r	44.2	45.9	47.6	48.0 ±0.8	47.4 ±0.6	55.4 ±0.5	58.9 ±0.9
hopper-m-r	94.7	92.1	96.9	101.3 ±2.2	101.0 ±0.3	103.1 ±0.1	101.4 ±0.5
walker2d-m-r	73.9	85.1	84.4	75.6 ±25.1	95.7 ±1.4	98.4 ±0.5	95.5 ±1.6

Table 2: Normalized returns of the final policy checkpoint for offline DPRL methods. Methods marked with † denote results reported in the original paper. Numbers within 95% of the highest scores are bolded.

Figure 5: Effect of diffusion steps on performance. Performance of QSM, SDAC, and DACER when varying the number of diffusion steps (2, 5, 10, and 50). Returns are normalized by task-specific reference scores in Table 5 and averaged across Ant, HalfCheetah, Humanoid, and Walker2d.

Figure 6: Effect of network architecture. Comparison between MLP and SimBa backbones for QSM (blue) and DACER (green) on DMControl hard tasks. Solid lines denote MLP policies and dashed lines denote Simba policies.

Figure 7: Effect of Action Dimensionality. Performance is normalized by SAC across tasks of increasing action dimension.

Figure 8: Effect of noise schedule. Performance of SDAC, QSM, and DACER with a cosine (black) or linear noise schedule.

Algorithm	Training (gradient steps/s)	Inference (actions/s)
SAC	1242.3	2623.0
QSM	675.9	1285.6
DACER	603.3	728.3
SDAC	256.3	1219.6
DPMD	126.2	1313.3
QVPO	150.6	1397.4

Table 3: Efficiency Comparison. Training and inference speed of off-policy algorithms. Training throughput is measured in gradient steps per second; inference throughput in single-action inferences per second.

Panel of Figure 9 — Figure 9: Representative tasks across three continuous-control suites. Left: MuJoCo Gymnasium. Middle: DeepMind Control Suite. Right: IsaacLab.

Setting	Gym-Locomotion	DMC	IsaacLab
Training frames	1M	1M	200M
Batch size	256	512	24576
Replay buffer size	$10^{6}$	$10^{6}$	—
Discount $\gamma$	0.99	0.99	0.99
Warmup frames	5K	10K	—
Random frames	5K	10K	—

Table 4: Shared training settings per benchmark.

Environment	Max Return
Ant-v5	6000
HalfCheetah-v5	11000
Hopper-v5	4000
Humanoid-v5	6000
Swimmer-v5	150
Walker2d-v5	6000

Table 5: Reference returns used for normalization in performance profiles (Figures 4 and 5 ). For each environment, we select a value approximately matching the highest returns observed across the evaluated algorithms.

Flow-Based Policy Learning

See all 37 papers

Designing Reinforcement Learning for Diffusion Models: A Unified Path-Space View
Aug 14, 2026Yixian Xu, Yuanrui Zhang, Shengjie Luo +2Diffusion Model Alignment Policy Gradient Methods
State Key Laboratory of General Artificial Intelligence, School of Intelligence Science and Technology, Peking University · ByteDance Seed
Mean Flow Policy Optimization
Apr 16, 2026Xiaoyi Dong, Xi Sheryl Zhang, Jian ChengMeanFlow Mean Flow
C2DL, Institute of Automation, Chinese Academy of Sciences · School of Artificial Intelligence, University of Chinese Academy of Sciences · AiRiA +1
Rethinking the Design Space of Reinforcement Learning for Diffusion Models: On the Importance of Likelihood Estimation Beyond Loss Design
Feb 4, 2026Jaemoo Choi, Yuchen Zhu, Wei Guo +6T2I Generation Diffusion Model Fine-Tuning
Georgia Institute of Technology · National University of Singapore · Nanjing university.

Aug 14, 2026J/K move · Enter open · S save

Designing Reinforcement Learning for Diffusion Models: A Unified Path-Space View

Yixian Xu, Yuanrui Zhang, Shengjie Luo, Liwei Wang+1

State Key Laboratory of General Artificial Intelligence, School of Intelligence Science and Technology, Peking University · ByteDance Seed

Diffusion Model Alignment Policy Gradient Methods RL for Diffusion Models Variance Reduction

Open paper AskPDF arXiv

Reinforcement learning (RL) post-training provides a direct way to align diffusion models with human preferences and task-specific rewards. However, current RL algorithms for diffusion models remain fragmented: reverse-trajectory methods rely on discretized likelihood ratios, whereas forward-matching methods train on reward-labeled noising versions of the rollout samples. This paper shows that these seemingly different losses arise from a single path-space principle. Starting from the regularized diffusion-RL objective, we use importance sampling between sampling SDEs to obtain an explicit policy-gradient estimator on trajectory space. The estimator contains the stochastic Itô integral underlying Flow-GRPO-type updates; we derive an equivalent variance-reduced value-gradient form that recovers the forward-matching structure of AWM and DiffusionNFT. This identifies the empirical gap between these method families as a variance-reduction effect rather than a difference in RL principle. The derivation yields a unified design space organized by value-gradient estimation, weight functions, and sampling choices. Within this space, we propose a multi-sample KDE value-gradient estimator that reuses rollout groups, together with scale-bounded weight families that retain stable existing recipes while excluding singular ones. Experiments on SD3.5-M and Qwen-Image models validate the variance-reduction explanation and show that the resulting recipe improves over prior diffusion-RL baselines.

Abstract

Flow-Based Policy Learning

Explore similar work

Designing Reinforcement Learning for Diffusion Models: A Unified Path-Space View

Mean Flow Policy Optimization

Rethinking the Design Space of Reinforcement Learning for Diffusion Models: On the Importance of Likelihood Estimation Beyond Loss Design