Value-Rectified Distillation for Flow-based Offline Reinforcement Learning

Ke Jiang; Wen Jiang; Yoshinobu Kawahara; Xiaoyang Tan

cs.LGOct 21, 2026

Value-Rectified Distillation for Flow-based Offline Reinforcement Learning

Authors: Ke Jiang, Wen Jiang, Yoshinobu Kawahara, Xiaoyang Tan

Organizations: College of Computer Science and Technology, Nanjing University of Aeronautics and Astronautics · MIIT Key Laboratory of Pattern Analysis and Machine Intelligence · Graduate School of Information Science and Technology, The University of Osaka · Center for Advanced Intelligence Project, RIKEN

Abstract

We aim to address the problem of Out-of-Distribution (OOD) Mode Averaging, a critical problem that emerges when distilling Flow Matching (FM) based behavior policies in offline Reinforcement Learning (RL): When learning over multi-modal datasets, standard distillation inadvertently forces the policy to interpolate across different behavioral modes, causing the generated actions to collapse into invalid, OOD regions. To resolve this, we propose Value-Rectified Distillation (VRD), a novel framework that reformulates offline policy distillation as a value-guided generative trajectory alignment problem. By employing a dynamic coupling mechanism, VRD actively rectifies flow dynamics to explicitly decouple distinct behavioral modes, cleanly isolating high-value actions. Specifically, we instantiate VRD through two complementary algorithms: VRD-Q, which utilizes value-quantile behavior distillation to statistically filter out low-value behaviors, and VRD-R, which reorganizes the latent space to geometrically isolate high-value generative trajectories from the FM-based behavior policy. Extensive empirical evaluations demonstrate that VRD improves existing generative baselines across diverse continuous control and planning tasks on the D4RL, OGBench, and Minari benchmarks, especially over those multi-modal, non-expert datasets.

Figures & tables

Figure 1: Problem of Latent Fusion (left) and the OOD Mode Averaging problem caused by FQL-style behavior distillation (right).

Figure 2: Implementations for Value-Rectified Distillation (VRD) : Statistical rectification via quantile behavior distillation (left), Geometric rectification via reshaping behavior distillation (right).

Figure 3: Visualizing results on toy examples with 2D multi-modal data.

	Gaussian Policies			Diffusion Policies			Flow Policies
Task	BC	IQL	ReBRAC	IDQL	SRPO	CAC	FAWAC	FBRAC	FQL	GFP	DF	VRD-Q	VRD-R
halfcheetah-med-expert	55.2	86.7	101.1	95.9	92.2	84.3	44.5	54.4	86.1	92.3	87.6	93.7 $\pm$ 1.2	101.3 $\pm$ 1.3
hopper-med-expert	52.5	91.5	107.0	108.6	100.1	100.4	88.4	76.5	105.4	110.1	112.6	96.4 $\pm$ 13.5	113.5 $\pm$ 0.4
walker2d-med-expert	107.5	109.6	111.6	112.7	114.0	110.4	57.3	53.5	102.7	106.4	111.3	112.1 $\pm$ 2.3	113.7 $\pm$ 0.2
halfcheetah-medium	42.6	47.4	65.6	51.0	60.4	69.1	48.3	53.1	60.0	62.3	49.3	65.7 $\pm$ 1.2	61.7 $\pm$ 4.5
hopper-medium	52.9	66.3	102.0	65.4	95.5	80.7	81.2	52.1	67.3	84.7	97.1	65.8 $\pm$ 15.8	76.3 $\pm$ 12.1

Table 1: Results on D4RL MuJoCo, AntMaze and Adroit Benchmarks. The highest results of each row are bolded. Each result of our method is calculated over 8 seeds. Part of the results were reported in [ 21 ] .

	Gaussian Policies			Diffusion Policies			Flow Policies
Task	BC	IQL	ReBRAC	IDQL	SRPO	CAC	FAWAC	FBRAC	IFQL	FQL	GFP	VRD-Q	VRD-R
antmaze-large	11.0	53.0	81.0	21.0	11.0	33.0	6.0	60.0	28.0	79.0	93.8	95.5 $\pm$ 1.2	94.9 $\pm$ 1.9
antmaze-giant	0.0	4.0	26.0	0.0	0.0	0.0	0.0	4.0	3.0	9.0	27.9	49.8 $\pm$ 5.3	38.3 $\pm$ 7.2
humanoidmaze-medium	2.0	33.0	22.0	1.0	1.0	53.0	19.0	38.0	60.0	58.0	72.0	75.3 $\pm$ 2.9	70.3 $\pm$ 7.2
humanoidmaze-large	1.0	2.0	2.0	1.0	0.0	0.0	0.0	2.0	11.0	4.0	17.8	19.7 $\pm$ 2.9	16.8 $\pm$ 5.9
antsoccer-arena	1.0	8.0	0.0	12.0	1.0	2.0	12.0	16.0	33.0	60.0	57.9	63.8 $\pm$ 4.7	47.9 $\pm$ 9.7

Table 2: Results on OGBench Tasks. The highest results of each row are bolded. Each result of our method is calculated over 8 seeds. The full results are shown in Table 4 in Appendix, where part of the results were reported in [ 21 ] .

Figure 4: Results over different scales of multi-modal datasets.

	VRD-Q ( $\tau$ )					VRD-R ( $\sigma$ )
Values	0.5	0.6	0.7	0.8	0.9	1	1e-2	1e-3	1e-4	1e-5
halfcheetah-m	60.4	63.4	65.7	65.3	58.1	45.7	60.6	60.4	61.7	61.2
hopper-m	64.3	65.1	65.8	65.1	56.4	50.2	73.9	75.5	76.3	74.2
walker2d-m	76.5	81.7	82.3	81.1	77.4	53.9	89.1	88.6	89.7	88.3

Table 3: Sensitivity analysis for main hyperparameters

\tau

of VRD-Q and

\sigma

of VRD-R (

\lambda=0.05

, full results see Appendix C.3 ).

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Task

IQL

ReBRAC

IDQL

SRPO

CAC

FAWAC

FBRAC

IFQL

FQL

GFP actor

\pi_{\omega}

VRD-Q (Ours)

VRD-R (Ours)

antmaze-large-navigate-singletask

task1-v0

0\pm 0

48\pm 9

91\pm 10

0\pm 0

0\pm 0

42\pm 7

1\pm 1

70\pm 20

24\pm 17

80\pm 8

95.4\pm 0.8

\mathbf{96.7}\pm 0.4

\mathbf{97.2}\pm 0.6

task2-v0

6\pm 3

42\pm 6

88\pm 4

14\pm 8

4\pm 4

1\pm 1

0\pm 1

35\pm 12

8\pm 3

57\pm 10

92.2\pm 3.0

\mathbf{96.2}\pm 0.1

93.1\pm 3.3

task3-v0

29\pm 5

72\pm 7

51\pm 18

26\pm 8

3\pm 2

49\pm 10

12\pm 4

83\pm 15

52\pm 17

93\pm 3

95.6\pm 2.7

\mathbf{98.4}\pm 0.7

96.4\pm 1.7

task4-v0

8\pm 3

51\pm 9

84\pm 7

62\pm 25

45\pm 19

17\pm 6

10\pm 3

37\pm 18

18\pm 8

80\pm 4

\mathbf{90.6}\pm 2.6

\mathbf{90.7}\pm 3.8

\mathbf{90.3}\pm 2.9

task5-v0

10\pm 3

54\pm 22

90\pm 2

2\pm 2

1\pm 1

55\pm 6

9\pm 5

76\pm 8

38\pm 18

83\pm 4

{95.0}\pm 1.3

{95.4}\pm 0.9

\mathbf{97.5}\pm 0.8

Appendix

Table 4: Detailed Offline RL Results on OGBench Tasks: Performance comparison across different offline RLalgorithms. All prior work results are taken from [ 21 ] to compare a wide range of previous methods against VRD.

Task	Offline RL algorithms
Task	FQL	GFP actor $\pi_{\theta}$	GFP VaBC $\pi_{\omega}$	VRD-Q (Ours)	VRD-R (Ours)
Minari halfcheetah-simple-v0	$59.2\pm 0.3$	$\textbf{72.5}\pm 0.5$	$64.4\pm 0.4$	$71.4\pm 0.1$	$70.3\pm 0.2$
Minari halfcheetah-medium-v0	$100.2\pm 6.7$	$121.3\pm 5.8$	$108.6\pm 5.3$	$117.4\pm 3.1$	$\textbf{123.9}\pm 2.3$
Minari halfcheetah-expert-v0	$134.1\pm 2.3$	$133.4\pm 1.3$	$136.4\pm 0.8$	$96.7\pm 10.7$	$\textbf{139.0}\pm 0.1$
Minari hopper-simple-v0	$57.2\pm 5.9$	$91.6\pm 4.3$	$87.4\pm 6.6$	$\textbf{99.1}\pm 0.1$	$\textbf{99.2}\pm 0.6$
Minari hopper-medium-v0	$81.9\pm 23.9$	$79.6\pm 14.5$	$78.2\pm 24.2$	$\textbf{107.3}\pm 4.4$	$104.3\pm 8.6$

Appendix

Table 5: Comparisons on Minari MuJoCo benchmarks. The best scores are highlighted in bold.

Figure 5: Sensitivity analysis for VRD-Q.

Figure 6: Sensitivity analysis for VRD-R.

Figure 7: Validation study for Quantile Distillation in VRD-Q.

Figure 8: Validation for VRD-R: Q-value sensitivity to latent radius

Category	Hyperparameter	Value
Architecture	Actor / Critic Hidden Layers	$[256,256,256]$
	Activation Function	ReLU / SiLU
	Time Embedding Dimension	$64$
Optimization	Actor / Critic Learning Rate	$3\times 10^{-4}$
	Batch Size	$256$
	Discount Factor ( $\gamma$ )	$0.99$

Appendix

Table 6: General network architectures and optimization hyperparameters shared across all benchmarks.

Benchmark	Specific Datasets / Environments	VRD-Q Threshold ( $\tau$ )	VRD-R Margin ( $\sigma$ )
OGBench	humanoid-medium , ant-large , cube-single , puzzle	$0.7$	$10^{-3}$
OGBench	humanoid-large , ant-giant , cube-double , scene	$0.6$	$10^{-2}$
D4RL	antmaze-umaze/medium/large	$0.6\sim 0.7$	$10^{-2}\sim 10^{-3}$
	halfcheetah/hopper/walker2d-medium/medium-replay	$0.7\sim 0.8$	$10^{-4}$
	halfcheetah/hopper/walker2d-expert/medium-expert	$0.8\sim 0.9$	$10^{-5}$
Minari	Mixed / Sub-optimal datasets (e.g., pointmaze-medium-v1 )	$0.7\sim 0.8$	$10^{-4}$

Appendix

Table 7: Task-specific hyperparameter selections for VRD-Q (

\tau

) and VRD-R (

\sigma

). Highly multi-modal tasks with sparse expert data (e.g., OGBench and AntMaze) utilize a smaller

\tau

and a larger

\sigma

(e.g.,

10^{-2}

), whereas near-expert, structured datasets utilize a larger

\tau

and a strictly smaller

\sigma

(e.g.,

10^{-5}

Specification	One-Step Policy ( Actor )	Flow Policy ( ActorVectorField )	Critic ( Value )
Primary Inputs	$s\oplus z$ (Obs, Latent z)	$s\oplus z_{t}\oplus t$ (Obs, Latent $z$ , Time $t$ )	$s\oplus a$ (Obs, Actions)
Latent Variable ( $z$ ) Role	Input Feature: Nonlinearly mapped via MLP to generate one-step action $a=\pi(s,z)$	Input Feature : Actively processed to compute flow velocity ( $\dot{z}_{t}$ )	N/A : Evaluates the final outcome action $a$
Backbone	MLP (Optional Layer Norm)	Time-conditioned MLP (Sinusoidal Embedding)	Ensemblized MLP ( $N\geq 1$ via vmap )
Output Type	Deterministic Action $a$ (conditioned on $s$ and $z$ )	Vector Field $v_{\theta}(s,z_{t},t)$	Scalar Value $Q(s,a)$ or $V(s)$

Appendix

Table 8: Refined Architectural Specifications Including Latent Variable Roles

Offline RL

Offline reinforcement learning learns decision-making policies from fixed datasets without interacting with the environment during training.

See all 298 papers

Learning Multimodal One-step Flow Policy via Value-weighted Optimal Transport
Sep 14, 2026Jaehun Shon, Jinha Choi, Jongwook Jeon +1Offline RL Flow-Based Policy Learning
Department of Artificial Intelligence, Yonsei University
DeFlow: Decoupling Behavior-Prior Modeling and Value Maximization for Offline Policy Extraction
Jan 15, 2026Zhancun Mu, Chi ZhangOffline RL Policy Learning
Peking University, Institute of Artificial Intelligence, China, Beijing.
Flow Actor-Critic for Offline Reinforcement Learning
Feb 20, 2026Jongseong Chae, Jongeui Park, Yongjae Shin +3Flow-Based Policy Learning Conservative RL
KAIST · UNIST

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

Learning Multimodal One-step Flow Policy via Value-weighted Optimal Transport

Jaehun Shon, Jinha Choi, Jongwook Jeon, Jongmin Lee

Department of Artificial Intelligence, Yonsei University

Offline RL Flow-Based Policy Learning

Open paper AskPDF arXiv

Offline reinforcement learning aims to learn a policy solely from fixed datasets, which often contain multimodal action distributions. Flow policies can naturally represent such multimodal behaviors, but learning an efficient one-step flow policy remains challenging: standard value guidance often leads to mode collapse or exploits overestimation bias in out-of-distribution regions. To address this, we introduce One-step Flow policy via Optimal Transport (OptiFlow), a framework for one-step flow policy learning as a structured sample-allocation problem. OptiFlow jointly trains a value-aware reference flow policy and an efficient one-step policy, coupling their action samples through state-wise entropic optimal transport. For each state, critic-estimated values define the priority of distillation target actions, while the action-distance cost ensures geometrically compatible pairings. By avoiding direct critic maximization, our transport-guided approach enables in-distribution exploitation by anchoring the one-step policy to high-value, dataset-supported modes without the risk of out-of-distribution divergence. Experimental results demonstrate that OptiFlow effectively captures optimal multimodal behaviors and achieves strong performance across diverse offline RL benchmarks. Our code is available at https://github.com/Yonsei-DILLab/OptiFlow.

Abstract

Offline RL

Explore similar work

Learning Multimodal One-step Flow Policy via Value-weighted Optimal Transport

DeFlow: Decoupling Behavior-Prior Modeling and Value Maximization for Offline Policy Extraction

Flow Actor-Critic for Offline Reinforcement Learning