Diffusion Policy Improvement with Proposal-Conditioned Refinement Flows

Junhyun Ha; Juho Lee; Byoungwoo Park

cs.LGSep 29, 2026

Diffusion Policy Improvement with Proposal-Conditioned Refinement Flows

Authors: Junhyun Ha, Juho Lee, Byoungwoo Park

Organizations: KAIST

Abstract

Diffusion and flow policies can model complex behaviors in offline reinforcement learning (RL). However, penalizing their KL divergence from the behavior policy can discourage actions having high critic values with low behavior density. Directly refining behavior proposals may be an alternative, yet Gaussian or deterministic editors limit expressiveness to represent multiple separated modes for the same proposal. In this work, we introduce Proposal-Conditioned Refinement Flows (PReFlow), a policy extraction method combining critic-based proposal selection with a conditional refinement flow. To optimize proposal selection and refinement together, we formulate a KL-regularized objective whose optimum induces a Gibbs policy over final actions under a Gaussian-smoothed behavior prior. The refinement flow can represent multiple high value modes, while a proposal-centered Gaussian reference regulates large action changes. This Gaussian reference further enables us to make use of simulation-free, closed form adjoint matching targets from sampled endpoints and critic gradients, yielding a single velocity regression loss without a backward adjoint solve. On 50 OGBench tasks, PReFlow achieves competitive offline performance and the highest aggregate score among the compared methods after online fine-tuning, reaching 91% after 500K environment steps.

Figures & tables

Figure 1: PReFlow: proposal selection and stochastic refinement. A behavior flow generates proposals, the critic selects one, and a conditional flow refines it to produce final action

\tilde{a}=a+\Delta a

. The performance panel reports mean success rates across the 50 OGBench tasks (Table 1 ).

Figure 2: Toy comparison of policy extraction methods. Fine-tuning the full flow (QAM) undercovers the two modes farther from the proposal, and Gaussian refinement concentrates on a single mode. PReFlow captures all three modes. Appendix D.4 provides details of the toy experiment.

	al	ag	hm	hl	scene	p33	p44	c2	c3	c4	all
Method	5 tasks	5 tasks	5 tasks	5 tasks	5 tasks	5 tasks	5 tasks	5 tasks	5 tasks	5 tasks	50 tasks
Gaussian
ReBRAC	[94,95] [97,98] [-1pt] 0 94 $\scriptstyle\to$ 98	[53,60] [69,79] [-1pt] 0 57 $\scriptstyle\to$ 75	[65,74] [52,63] [-1pt] 0 69 $\scriptstyle\to$ 58	[15,19] [18,28] [-1pt] 0 17 $\scriptstyle\to$ 23	[61,69] [98,99] [-1pt] 0 65 $\scriptstyle\to$ 99	[73,84] [100,100] [-1pt] 0 79 $\scriptstyle\to$ 100	[0,0] [0,0] [-1pt] 0 0 $\scriptstyle\to$ 0	[8,10] [96,98] [-1pt] 0 9 $\scriptstyle\to$ 97	[0,1] [9,20] [-1pt] 0 1 $\scriptstyle\to$ 14	[6,11] [79,80] [-1pt] 0 9 $\scriptstyle\to$ 80	[39,41] [63,65] [-1pt] 0 40 $\scriptstyle\to$ 64
Guidance
QSM	[88,92] [91,96] [-1pt] 0 90 $\scriptstyle\to$ 94	[19,29] [83,91] [-1pt] 0 24 $\scriptstyle\to$ 87	[79,84] [78,81] [-1pt] 0 82 $\scriptstyle\to$ 80	[5,7] [1,3] [-1pt] 0 6 $\scriptstyle\to$ 2	[77,80] [81,83] [-1pt] 0 78 $\scriptstyle\to$ 82	[52,63] [82,90] [-1pt] 0 57 $\scriptstyle\to$ 86	[0,0] [94,100] [-1pt] 0 0 $\scriptstyle\to$ 98	[31,34] [98,99] [-1pt] 0 33 $\scriptstyle\to$ 98	[5,6] [25,28] [-1pt] 0 6 $\scriptstyle\to$ 27	[18,19] [62,71] [-1pt] 0 19 $\scriptstyle\to$ 67	[38,40] [71,73] [-1pt] 0 39 $\scriptstyle\to$ 72
QPILOTS-U	[76,88] [94,98] [-1pt] 0 82 $\scriptstyle\to$ 96	[0,4] [85,94] [-1pt] 0 1 $\scriptstyle\to$ 89	[54,75] [98,100] [-1pt] 0 63 $\scriptstyle\to$ 99	[2,14] [40,72] [-1pt] 0 8 $\scriptstyle\to$ 59	[80,95] [96,100] [-1pt] 0 89 $\scriptstyle\to$ 98	[100,100] [100,100] [-1pt] 0 100 $\scriptstyle\to$ 100	[6,25] [100,100] [-1pt] 0 16 $\scriptstyle\to$ 100	[71,80] [99,100] [-1pt] 0 75 $\scriptstyle\to$ 100	[4,10] [57,74] [-1pt] 0 7 $\scriptstyle\to$ 69	[0,6] [76,80] [-1pt] 0 3 $\scriptstyle\to$ 78	[42,47] [87,90] [-1pt] 0 44 $\scriptstyle\to$ 89

Table 1: OGBench: offline

\rightarrow

online success rates (%). Each cell shows end-of-offline (left) and end-of-online (right) success rates with 95% CIs over seeds in small gray brackets. Best results per domain are highlighted . Per-task results are in Appendix E.6 .

Figure 3: Learning curves during online fine-tuning. PReFlow, RQL, and Q-Flow are fine-tuned for 0.5M online environment steps after 1M offline gradient steps. Curves show means over three seeds; shaded regions show pointwise 95% confidence intervals.

		Refinement Type
Domain	Base	Gaussian	Flow
al	2.3 $\pm$ 0.8	76.3 $\pm$ 7.2	79.0 $\pm$ 7.5
hm	2.7 $\pm$ 1.0	57.1 $\pm$ 8.5	69.2 $\pm$ 1.2
scene	3.9 $\pm$ 1.5	68.1 $\pm$ 3.2	73.1 $\pm$ 0.8
c2	2.0 $\pm$ 0.0	40.8 $\pm$ 5.2	53.9 $\pm$ 3.3

Table 2: Offline success rates (%) after 1M steps with

K=1

Figure 4: Mean wall-clock time per full training update.

Figure 5: Effect of proposal selection. PReFlow with

K\in\{1,2,4,8\}

behavior proposals on hm , scene , and c2 . All settings use the learned flow to refine the proposal with the highest critic value. Curves show means over three seeds; shaded regions show pointwise 95% confidence intervals.

Figure 6: Sensitivity to the refinement scale and inverse temperature. Offline performance on OGBench domains across different combinations of the refinement scale $\delta$ and inverse temperature $\tau$ . Each cell reports the mean performance over three seeds.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Domain	$\bm{\tau}$	$\bm{\delta}$	$\bm{K}$	Notes
al	auto ( $c{=}2.5$ )	0.1	1	–
ag	auto ( $c{=}1.5$ )	0.1	1	$\delta_{\mathrm{online}}{=}0.2$
hm	30	0.1	4	–
hl	100	0.05	4	–
scene	3	0.1	4	sparse
p33	3	0.1	4	sparse

Appendix

Table 3: Per-domain PReFlow settings. “auto (

c

)” denotes the temperature adaptation of Appendix C.1 with target

c

Figure 7: KL divergence to the target on the toy problem. Lines show means over 3 seeds, and shaded regions show one standard deviation. Both axes are logarithmic.

	al	hm	scene	c2
Method	5 tasks	5 tasks	5 tasks	5 tasks
SPAR anchor	[43,52] [-1pt] 0 48	[5,8] [-1pt] 0 7	[12,15] [-1pt] 0 14	[2,4] [-1pt] 0 3
SPAR	[55,64] [-1pt] 0 60	[6,9] [-1pt] 0 8	[26,29] [-1pt] 0 27	[2,4] [-1pt] 0 3
PReFlow	[79,85] [-1pt] 0 82	[60,69] [-1pt] 0 65	[96,97] [-1pt] 0 97	[56,62] [-1pt] 0 59

Appendix

Table 4: Comparison with SPAR. Offline success rates (%) averaged over eight seeds and five tasks per domain, with 95% confidence intervals in small gray brackets above the means. Best results per domain are highlighted. SPAR is evaluated after 1M gradient steps of critic and anchor training followed by 1M steps of residual training.

Method	Actor grad.	Critic grad.	Full step
QAM	3.40	2.13	8.22
QAM-E	5.08	2.71	9.97
PReFlow (ours)	2.65	2.43	7.18
QPILOTS-U	0.84	6.99	10.93
QPILOTS-M	0.91	49.69	51.20

Appendix

Table 5: Training-update cost on c2 (ms). Median times over 30 calls after warm-up, measured on the same GPU. Actor and critic columns measure loss-gradient computation without optimizer updates; the full step includes all training updates. The lowest time in each column is highlighted .

		$\bm{K=4}$			$\bm{K=8}$
Domain	$\bm{K=1}$	Uniform	Softmax	Argmax	Uniform	Softmax	Argmax
hm ( $\tau=30$ )	69.2 $\pm$ 1.2	68.8 $\pm$ 2.6	70.1 $\pm$ 3.1	70.0 $\pm$ 2.9	68.7 $\pm$ 0.8	66.0 $\pm$ 3.3	69.2 $\pm$ 3.4
scene ( $\tau=3$ )	73.1 $\pm$ 0.8	74.0 $\pm$ 1.7	89.5 $\pm$ 0.8	94.9 $\pm$ 1.0	74.1 $\pm$ 0.8	90.7 $\pm$ 1.7	92.8 $\pm$ 0.6
c2 ( $\tau=3$ )	53.9 $\pm$ 3.3	50.3 $\pm$ 6.9	61.3 $\pm$ 1.0	63.3 $\pm$ 4.2	49.6 $\pm$ 8.0	60.1 $\pm$ 3.9	58.7 $\pm$ 4.1

Appendix

Table 6: Effect of the selection rule. Offline success rates (%) after 1M gradient steps. Entries report mean

\pm

standard deviation over three seeds, with each seed’s score averaged over five tasks. The listed

\tau

values apply to softmax selection. Best results are highlighted among the three rules for each domain and

K

Figure 8: Proposal selection with and without refinement. Offline success rates average on c2, hm, scene, al for

K\in\{1,2,4,8\}

, with refinement (

\delta=0.1

, left) and without refinement (

\delta=0

, right). Without refinement, the selected proposal is executed directly. Curves show means over three seeds; shaded regions show pointwise 95% confidence intervals.

Figure 9: Effect of proposal selection across all ten OGBench domains. PReFlow with

K\in\{1,2,4,8\}

behavior proposals. All settings use the learned refinement flow to refine the proposal with the highest critic value. Curves show means over three seeds; shaded regions show pointwise 95% confidence intervals.

Figure 10: Full training curves on the 50 OGBench tasks. Online fine-tuning begins after 1M offline gradient steps and continues for 500K environment steps.

Explore similar work

Refinement-based Flow Policy Optimization
Sep 14, 2026Bumgeun Park, Hyukjun Yang, Donghwan LeeFlow Policies Distributional Reinforcement Learning
Korea Advanced Institute of Science and Technology (KAIST)
ReFPO: Reflow Regularization for Flow Matching Policy Gradients
Jun 19, 2026Ge Wang, Yibo Peng, Fan Feng +10Flow Policies Offline Reinforcement Learning
Ising AI · CUHK-Shenzhen · FNii-Shenzhen
Discrete Flow Matching for Offline-to-Online Reinforcement Learning
May 12, 2026Fairoz Nower Khan, Nabuat Zaman Nahim, Peizhong JuFlow Policies Offline Reinforcement Learning
Department of Computer Science, University of Kentucky

Sep 14, 2026, cs.LGJ/K move · Enter open · S save

Refinement-based Flow Policy Optimization

Bumgeun Park, Hyukjun Yang, Donghwan Lee

Korea Advanced Institute of Science and Technology (KAIST)

Flow Policies Distributional Reinforcement Learning Multimodal Action Distributions Offline Reinforcement Learning

Open paper AskPDF arXiv

Flow-based policies offer an expressive representation for online reinforcement learning, but conventional flow matching requires samples drawn from the distribution to be modeled. This poses a challenge when the desired action distribution is defined only implicitly by a Q-function, since directly sampling actions from the resulting distribution is generally intractable. We propose Refinement-Based Flow Policy Optimization (RFPO), a novel framework for training a flow policy in online reinforcement learning by alternating between Q-guided sample refinement and self-target flow matching. RFPO first generates actions from Gaussian noise using the current flow policy and then uses a finite-step stochastic refinement procedure to move them toward an energy-based distribution induced by the Q-function. Each refined action is then paired with its corresponding initial noise sample and used as a fixed target for flow-matching training. By repeatedly refining its own outputs and learning from the resulting targets, RFPO incorporates Q-guidance into the policy without requiring direct samples from the target distribution, while retaining the capacity to represent multiple action modes. We further provide a theoretical analysis of the distributional dynamics induced by RFPO. Across six continuous-control tasks, RFPO matches or outperforms a standard Gaussian-policy baseline on almost every task. Experiments on six synthetic two-dimensional target distributions with diverse geometries demonstrate that RFPO captures complex multimodal structure without mode collapse.

Abstract

Explore similar work

Refinement-based Flow Policy Optimization

ReFPO: Reflow Regularization for Flow Matching Policy Gradients

Discrete Flow Matching for Offline-to-Online Reinforcement Learning