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
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 ± 1.2
101.3 ± 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 ± 13.5
113.5 ± 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 ± 2.3
113.7 ± 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 ± 1.2
61.7 ± 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 ± 15.8
76.3 ± 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 ± 1.2
94.9 ± 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 ± 5.3
38.3 ± 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 ± 2.9
70.3 ± 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 ± 2.9
16.8 ± 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 ± 4.7
47.9 ± 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 ( τ )
VRD-R ( σ )
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 τ of VRD-Q and σ of VRD-R ( λ=0.05 , full results see Appendix C.3 ).
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
Task
BC
IQL
ReBRAC
IDQL
SRPO
CAC
FAWAC
FBRAC
IFQL
FQL
GFP actor πω
VRD-Q (Ours)
VRD-R (Ours)
antmaze-large-navigate-singletask
task1-v0
0±0
48±9
91±10
0±0
0±0
42±7
1±1
70±20
24±17
80±8
95.4±0.8
96.7±0.4
97.2±0.6
task2-v0
6±3
42±6
88±4
14±8
4±4
1±1
0±1
35±12
8±3
57±10
92.2±3.0
96.2±0.1
93.1±3.3
task3-v0
29±5
72±7
51±18
26±8
3±2
49±10
12±4
83±15
52±17
93±3
95.6±2.7
98.4±0.7
96.4±1.7
task4-v0
8±3
51±9
84±7
62±25
45±19
17±6
10±3
37±18
18±8
80±4
90.6±2.6
90.7±3.8
90.3±2.9
task5-v0
10±3
54±22
90±2
2±2
1±1
55±6
9±5
76±8
38±18
83±4
95.0±1.3
95.4±0.9
97.5±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
FQL
GFP actor πθ
GFP VaBC πω
VRD-Q (Ours)
VRD-R (Ours)
Minari halfcheetah-simple-v0
59.2±0.3
72.5±0.5
64.4±0.4
71.4±0.1
70.3±0.2
Minari halfcheetah-medium-v0
100.2±6.7
121.3±5.8
108.6±5.3
117.4±3.1
123.9±2.3
Minari halfcheetah-expert-v0
134.1±2.3
133.4±1.3
136.4±0.8
96.7±10.7
139.0±0.1
Minari hopper-simple-v0
57.2±5.9
91.6±4.3
87.4±6.6
99.1±0.1
99.2±0.6
Minari hopper-medium-v0
81.9±23.9
79.6±14.5
78.2±24.2
107.3±4.4
104.3±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×10−4
Batch Size
256
Discount Factor ( γ )
0.99
Appendix
Table 6: General network architectures and optimization hyperparameters shared across all benchmarks.
Table 7: Task-specific hyperparameter selections for VRD-Q ( τ ) and VRD-R ( σ ). Highly multi-modal tasks with sparse expert data (e.g., OGBench and AntMaze) utilize a smaller τ and a larger σ (e.g., 10−2 ), whereas near-expert, structured datasets utilize a larger τ and a strictly smaller σ (e.g., 10−5 ).
Specification
One-Step Policy ( Actor )
Flow Policy ( ActorVectorField )
Critic ( Value )
Primary Inputs
s⊕z (Obs, Latent z)
s⊕zt⊕t (Obs, Latent z , Time t )
s⊕a (Obs, Actions)
Latent Variable ( z ) Role
Input Feature: Nonlinearly mapped via MLP to generate one-step action a=π(s,z)
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.
Jaehun Shon, Jinha Choi, Jongwook Jeon +1
Department of Artificial Intelligence, Yonsei University
We present DeFlow, a decoupled offline RL framework that leverages flow matching to faithfully capture complex behavior manifolds. Optimizing generative policies is computationally prohibitive, typically necessitating backpropagation through ODE solvers. We address this by learning a lightweight refinement module within an explicit, data-derived trust region of the flow manifold, rather than sacrificing the iterative generation capability via single-step distillation. This way, we bypass solver differentiation and eliminate the need for balancing loss terms, ensuring stable improvement while fully preserving the flow's iterative expressivity. Empirically, DeFlow achieves superior performance on the challenging OGBench benchmark and demonstrates efficient offline-to-online adaptation.
Zhancun Mu, Chi Zhang
Peking University, Institute of Artificial Intelligence, China, Beijing.
The dataset distributions in offline reinforcement learning (RL) often exhibit complex and multi-modal distributions, necessitating expressive policies to capture such distributions beyond widely-used Gaussian policies. To handle such complex and multi-modal datasets, in this paper, we propose Flow Actor-Critic, a new actor-critic method for offline RL, based on recent flow policies. The proposed method not only uses the flow model for actor as in previous flow policies but also exploits the expressive flow model for conservative critic acquisition to prevent Q-value explosion in out-of-data regions. To this end, we propose a new form of critic regularizer based on the flow behavior proxy model obtained as a byproduct of flow-based actor design. Leveraging the flow model in this joint way, we achieve new state-of-the-art performance for test datasets of offline RL including the D4RL and recent OGBench benchmarks.