PaperlayerView exact revision on arXiv

Source-linked AI summary

Value-Rectified Distillation for Flow-based Offline Reinforcement Learning

Ke Jiang, Wen Jiang, Yoshinobu Kawahara, Xiaoyang Tan

arXiv:2610.29075v1cs.LGcs.AI

TL;DR

FQL-style distillation of flow-based behavior policies can average across modes in multimodal offline data, generating OOD actions. VRD aligns generative trajectories with value guidance through statistical filtering or geometric reshaping. The paper reports that VRD outperforms baselines across benchmarks and is robust to suboptimal-data noise, while relying on accurate value estimates.

  • Problem

    FQL-style distillation can interpolate across distinct behavioral modes in multimodal offline data, producing invalid out-of-distribution actions.

    Source: §1 · pp. 1–2

  • Method

    VRD reframes policy distillation as value-guided trajectory alignment, with VRD-Q statistically filtering low-value behaviors and VRD-R geometrically reshaping latent trajectories.

    Source: §1 · p. 2

  • Results

    VRD decouples high-value behavior modes and outperforms baselines across varied benchmarks and datasets, with reported robustness to suboptimal-data noise.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    The reported results support value-guided alignment as a way to extract high-value policies from multimodal data without mode averaging.

    Source: §1 · p. 2, §6 · p. 10

    VRD depends on value-network accuracy, and severe Q-value overestimation may mislead alignment.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

Flow Matching policies model multimodal behavior, but FQL-style distillation can interpolate across entangled modes and generate out-of-distribution actions. VRD reframes distillation as value-guided trajectory alignment, with statistical and geometric variants, and the paper reports improved or competitive performance against generative baselines.

1 Introduction

  • FQL-style distillation can interpolate across entangled behavioral modes, producing low-density out-of-distribution actions and degrading performance.The paper identifies this as OOD Mode Averaging and notes that Guided Flow Policy retains the issue when using FQL-style distillation.Source: §1 · pp. 1–2
  • VRD rectifies flow dynamics through value-guided trajectory alignment to decouple behavioral modes and exclude suboptimal modes during inference.Its dynamic coupling mechanism is intended to keep generated actions on the optimal action manifold.Source: §1 · p. 2
  • VRD-Q uses value-quantile weighting to statistically filter suboptimal regions without changing the flow topology, while VRD-R reshapes topology to isolate high-value trajectories.VRD-R transports trajectories toward a compact optimal support set.Source: §1 · p. 2
  • The paper reports that VRD-Q and VRD-R mitigate OOD mode averaging and outperform or perform competitively with generative baselines across continuous-control and planning tasks.The reported evaluations span diverse tasks; no benchmark-specific values are supplied here.Source: §1 · p. 2

2 Preliminaries and motivation

Offline RL learns from a fixed dataset, and FQL combines a flow-based behavior policy with a one-step actor and Q-value guidance. On isolated multimodal data, continuity can entangle optimal and suboptimal modes, while Euclidean distillation can drive the actor through low-density OOD regions.

2.1 Offline RL and Flow Q-Learning

  • FQL uses a critic, a behavior-cloned flow policy, and a one-step actor trained to maximize Q-values while distilling from the frozen flow policy.Its latent prior is an unstructured isotropic Gaussian, so latent position does not encode action quality.Source: §2.1 · pp. 2–3

2.2 Latent fusion and OOD mode-averaging problems of FQL in multi-modal data

  • For isolated multimodal data with a gap between modes, a continuous flow mapping from connected latent space cannot isolate optimal modes without intersecting OOD regions.The analysis says the optimal latent set is entangled with suboptimal or OOD regions at its boundary, without quantifying the probability mass there.Source: §2.2 · p. 3
  • Under the stated strong-concavity and smoothness assumptions, FQL-style distillation with any nonzero value penalty has expected OOD error bounded below by D^2/4 when modes are separated by gap D.The result follows from Euclidean regularization forcing interpolation across the action-space gap.Source: §2.2 · p. 4
  • The paper’s motivation is to replace action-space regularization with dynamic alignment of flow dynamics, enabling high-value actions to be extracted along the data manifold.The stated goal is to avoid interpolation through empty regions between modes.Source: §2.2 · p. 4

3 Methodology

VRD casts policy extraction as value-guided alignment between latent distributions, with a kernel controlling sample weights and latent coordinates. It yields two implementations: quantile-based statistical filtering and geometric latent reshaping.

3.1 Unified framework: kernel-coupled value rectification

  • The unified VRD framework replaces action-space Euclidean penalties with a latent alignment kernel, which guarantees a positive lower bound on optimal-action probability under its leakage and coupling conditions.The stated bound is C/(C+ϵ), and the kernel is intended to keep trajectories on the data manifold while bypassing OOD regions.Source: §3.1 · p. 5

3.2 Statistical rectification: quantile behavior distillation

  • VRD-Q keeps the target latent geometry fixed and uses a value-quantile-modulated diagonal kernel to reweight source latents.Its quantile threshold trades off suboptimal filtering against robustness to critic error; under the stated mass conditions, the kernel-induced distribution has the framework’s optimal-action guarantee, while standard-prior inference has a separate bound.Source: §3.2 · pp. 5–6

3.3 Geometric rectification: reshaped behavior distillation

  • VRD-R gives samples equal fitting importance but dynamically transports each source latent to a Q-guided target coordinate.The reshaping function acts as a gravitational field toward high-value actions, and the transport kernel is a deterministic shifted Dirac map.Source: §3.3 · p. 6
  • Under the stated conservative-critic and concave dominant-mode assumptions, VRD-R’s truncation-radius condition yields zero suboptimal leakage under the kernel-induced marginal.The bound on the radius depends on the optimality relaxation factor and state-wise value range.Source: §3.3 · pp. 6–7
  • The variants differ in their trade-offs: VRD-Q preserves expert-behavior diversity for multimodal tasks, whereas VRD-R isolates a stable peak mode for precision-focused control.The paper gives OGBench navigation and standard D4RL locomotion as examples, respectively.Source: §3.3 · p. 7

4 Related works

The related-work section contrasts Gaussian, diffusion, and flow-based policies, emphasizing their respective expressivity and optimization trade-offs.

4 Related works

  • Gaussian policies are efficient but can average modes and perform poorly on complex, skewed, non-Gaussian offline data.The section lists BC, IQL, and ReBRAC as examples.Source: §4 · p. 7
  • Diffusion policies model complex multimodal distributions, but Q-optimization can require unstable backpropagation through sampling or complex reward-guided sampling.IDQL, SRPO, and CAC are named examples.Source: §4 · p. 7
  • Flow matching offers simpler training and faster inference, while FQL distills a one-step policy and later methods target high-value action extraction.FAWAC, FBRAC, GFP, and Decision Flow are discussed as related approaches.Source: §4 · p. 7

5 Experiments

The experiments examine latent rectification on a toy dataset, benchmark performance, robustness to increasingly dominant sub-optimal modes, and hyperparameter sensitivity.

5.1 Visualizing latent rectification on toy examples (RQ1)

  • On the 2D bimodal toy dataset, baseline flow exhibits latent fusion and OOD action interpolation, while VRD-Q and VRD-R extract the optimal mode through statistical filtering and geometric reshaping.VRD-R attracts optimal latents toward the origin and expels sub-optimal ones to the periphery; low-temperature Gaussian sampling isolates optimal trajectories.Source: §5.1 · p. 8, Figure 3 · p. 8

5.2 Continuous control (D4RL) and complex planning (OGBench) tasks (RQ2)

  • VRD-R achieves a 92.8 average score on D4RL, while VRD-Q scores 88.8; VRD-R also scores 113.5 on hopper-med-expert and 111.3 on pen-human.On pen-human, VRD-R's 111.3 compares with GFP's 64.6.Source: §5.2 · pp. 8–9, Table 1 · p. 9
  • VRD-Q achieves the highest average score, 56.4, across OGBench's 50 tasks; on antmaze-giant it scores 49.8 versus GFP's 27.9.VRD-R excels on spatial reasoning tasks, with a 50.8 average and scores of 34.7 on puzzle-3x3 and 55.8 on scene.Source: Table 2 · p. 9, §5.2 · p. 9

5.3 Deep dive into multi-modal datasets (RQ3)

  • As the random-sample proportion rises from 50% to 90%, both VRD variants degrade more gradually and outperform CQL and FQL on halfcheetah-mix and walker2d-mix.The experiment keeps the dataset buffer at 1,000,000 samples while reducing the expert-mode proportion.Source: §5.3 · pp. 9–10, Figure 4 · p. 10

5.4 Hyperparameter analysis (RQ4)

  • VRD-Q peaks at τ = 0.7, while VRD-R performs best at σ = 10^-4; the recommended settings balance filtering and distribution coverage, and mode decoupling without distorting latent topology.The authors recommend selecting τ based on sub-optimal data proportions and tailoring σ to the environment.Source: §5.4 · p. 10, Table 3 · p. 10

6 Conclusion

VRD reframes offline policy distillation as value-guided trajectory alignment to address latent fusion. The paper reports that VRD separates high-value modes and outperforms baselines across benchmarks and datasets, while Proposition 1 characterizes the topology of latent regions.

6 Conclusion

  • VRD separates high-value behavior modes and outperforms baselines across benchmarks and datasets, with reported robustness to sub-optimal data noise.The conclusion describes VRD-Q as statistical filtering and VRD-R as geometric reshaping.Source: §6 · p. 10
  • VRD relies on value-network accuracy, and severe Q-value overestimation may mislead its alignment process.The authors identify more robust value estimation and extension to discrete action spaces as future work.Source: §6 · p. 10

A.1 Proof of Proposition 1

  • When both optimal and sub-optimal modes are represented, Proposition 1 states that the optimal latent boundary and OOD latent set are non-empty.It also states that every neighborhood of a boundary point intersects non-optimal regions.Source: §A.1 · p. 13

A.2 Proof of Theorem 1

Theorem 1 analyzes FQL-style distillation under separated multimodal data and shows how its regularized update can traverse OOD regions. The proof connects interpolation between modes to a lower bound on prediction error.

A.2 Proof of Theorem 1

  • Under the stated strong-concavity, smoothness, and multimodal-gap assumptions, FQL generates OOD actions with minimum expected error bounded below by D^2/4.The result concerns FQL-style distillation with a non-zero value regularization penalty.Source: §A.2 · pp. 14–15
  • As regularization varies, the FQL action moves continuously from the optimal mode at α → 0 to the behavior-cloning mode at α → +∞.The proof represents the update as a matrix-weighted interpolation of the two modes.Source: §A.2 · p. 14
  • Because the valid modes are separated by a gap D, the continuous interpolation must pass through the low-density OOD region.The proof places an interpolated action geometrically within the gap between modes.Source: §A.2 · p. 15

A.3 Proof of Theorem 2

Theorem 2 analyzes value-guided alignment through a non-merging kernel. Its argument distinguishes in-place alignment from contractive reshaping and relates optimal coupling mass to optimal-action output.

A.3 Proof of Theorem 2

  • Theorem 2 considers a conservative critic and a non-merging alignment kernel, with conditions on suboptimal leakage and optimal coupling mass.The supplied theorem statement identifies these as requirements for its guarantee.Source: §A.3 · p. 15
  • In-place alignment preserves each optimal latent point, while contractive reshaping can combine optimal latents at a target coordinate whose conditional barycenter remains optimal.The proof describes VRD-Q as the in-place case and VRD-R as the reshaping case.Source: §A.3 · pp. 16–17
  • Under either alignment regime, the value lower bound and critic conservatism exclude OOD actions for the covered test latents.The proof states this consequence for latents coupled to the optimal region.Source: §A.3 · p. 17

A.4 Proof of Proposition 2

The Proposition 2 proof analyzes VRD-Q’s quantile-based weights and derives guarantees under both its kernel-induced distribution and the standard Gaussian inference prior.

A.4 Proof of Proposition 2

  • For VRD-Q, the kernel-induced probability of an optimal action is bounded below by C/(C+ϵ), while the standard-prior guarantee is stated separately.The proposition also gives a standard-prior probability expression, but its full formula is not present in the supplied passage.Source: §A.4 · pp. 17–18
  • VRD-Q uses a stop-gradient fixed point with T(z)=z, making its regularization coefficient depend on the quantile weight ωτ(z).The proof defines αreg(z) as ωτ(z)/λ.Source: §A.4 · p. 17
  • The proof states that ∆(α) decreases strictly as α increases when v is nonzero, informing its analysis of the quantile-weight cases.It then checks candidate weights for latents in the specified value band.Source: §A.4 · pp. 17–18

A.5 Proof of Proposition 3

Under the stated critic, optimal-mode, and truncation-radius conditions, Proposition 3 guarantees that VRD-R samples optimal in-distribution actions with probability 1. The appendix algorithms instantiate value rectification through quantile weighting in VRD-Q or latent-space reshaping in VRD-R.

A.5 Proof of Proposition 3

  • With a conservative critic, a convex dominant optimal mode, and σ < ϵopt/(Qmax(s) − Qmin(s) + ϵm), VRD-R outputs optimal in-distribution actions with probability 1.The guarantee is under the kernel-induced marginal distribution over the test ball.Source: §A.5 · pp. 18, 20
  • VRD-R maps optimal behaviors inside radius δ and suboptimal or OOD behaviors outside it, separating their target images.Peak-value actions map to the origin; the proof also establishes zero suboptimal leakage when the test radius σ is below δ.Source: §A.5 · p. 19

B Algorithms

  • VRD-Q uses value-quantile weights and an in-place alignment kernel, while VRD-R uses a value-guided latent transport kernel to reshape trajectories.VRD-Q samples from a standard Gaussian at inference; VRD-R samples from a truncated Gaussian prior.Source: B Algorithms · p. 20

C.1 Full results over OGBenchmarks (50 tasks)

On OGBench, VRD-Q has the highest reported average score across 50 tasks; on Minari MuJoCo, VRD-R has the highest reported average. The Minari results also report strong performance on simple datasets and expert tasks.

C.1 Full results over OGBenchmarks (50 tasks)

  • 56.4 is VRD-Q’s OGBench average, above GFP actor (51.8) and FQL (43.5); VRD-R scores 50.8.OGBench comprises 50 tasks across 10 domains, including navigation, manipulation, and puzzles.Source: §C.1 · pp. 21–22, Table 4 · p. 21

C.2 Minari: External comparison study

  • 112.8 is VRD-R’s Minari average, above FQL (99.7) and the strongest GFP variants (108.4).In hopper-simple, both VRD variants score approximately 99.1, compared with FQL’s 57.2; VRD-R reaches 139.0 in halfcheetah-expert.Source: §C.2 · pp. 22–23, Table 5 · p. 23
  • Minari results report significant gains on Simple and Medium datasets, which contain more sub-optimal trajectories, and characterize VRD as robust on non-expert multi-modal data.The source also reports competitive performance on expert-level tasks.Source: §C.2 · pp. 22–23

C.3.1 Analysis of VRD-Q Sensitivity (Figure 5)

VRD-Q sensitivity tests on D4RL medium datasets find that performance depends on the quantile threshold τ and balance coefficient λ.

C.3.1 Analysis of VRD-Q Sensitivity (Figure 5)

  • VRD-Q performance follows an inverted U-shaped pattern as τ varies: low τ filters suboptimal modes insufficiently, while high τ restricts weights and may cause vanishing gradients.The reported peak-performance region is τ ∈ [0.7, 0.8].Source: §C.3.1 · p. 24
  • VRD-Q is relatively stable for λ from 0.02 to 0.1, but λ = 1.0 tends to degrade performance as rectification dominates Q-value maximization.Source: §C.3.1 · p. 24

C.3.2 Analysis of VRD-R Sensitivity (Figure 6)

VRD-R performance is sensitive to the inference truncation radius, and the supplied analyses also report results for VRD-Q latent-code comparisons, VRD-R latent-radius trends, task-specific hyperparameters, and network roles.

C.3.2 Analysis of VRD-R Sensitivity (Figure 6)

  • As σ decreases toward 10^-4, VRD-R performance improves and stabilizes; σ = 1 performs suboptimally because sampling reaches peripheral suboptimal regions.The analysis reports λ = 0.05 and σ = 10^-4 as a robust configuration across the tested locomotion tasks.Source: §C.3.2 · p. 24

C.3.3 Practical Hyperparameter Tuning Suggestions

  • For noisy or exploratory datasets, the tuning guidance recommends τ around 0.6–0.7 and σ = 10^-4 by default, with adjustments for mode collapse or lingering suboptimal actions.The guidance also recommends λ in [0.05, 0.1], with task-dependent adjustments.Source: §C.3.3 · pp. 24–25

C.4.1 Quantile Distillation in VRD-Q

  • In t-SNE visualizations of three D4RL medium tasks, VRD-Q produced more higher-Q-value actions than the flow model: 4,764 vs 236 in Halfcheetah, 4,068 vs 932 in Hopper, and 4,046 vs 954 in Walker2d.The points represent sampled latent codes, with colors indicating which policy produced the higher estimated Q-value.Source: §C.4.1 · p. 25

C.5 Reshaping Distillation in VRD-R

  • Across the tested environments, VRD-R’s Normalized ∆Q declines as latent radius increases, with high-value modes near the origin and low-value modes toward the periphery.The figure analysis covers Halfcheetah, Hopper, and Walker2d; its horizontal axis is latent radius and its vertical axis is Normalized ∆Q.Source: §C.5 · p. 26

C.6 Hyperparameters

  • Hyperparameter choices depend on dataset quality and task complexity: noisy, multi-modal tasks use moderate τ and relatively larger σ, while cleaner expert settings favor higher τ and smaller σ.The reported searched σ grid is {10^-2, 10^-3, 10^-4, 10^-5}.Source: §C.6 · p. 27

C.7 Network structures

  • The one-step actor takes latent noise with the state as input, while the flow actor uses state, time, and intermediate latent state to compute its vector field; the critic evaluates state-action pairs.Source: §C.7 · pp. 27–28

C.8 Computational Resources

  • Experiments and baseline reproductions used RTX 3090 GPUs, Python with PyTorch/JAX, and an Intel Xeon W-2295 CPU with 128GB RAM.Source: §C.8 · p. 28

D Discussion

The discussion outlines VRD’s limitations and notes the paper’s claim-support and reproducibility details. It also states that LLMs are not important, original, or non-standard components of the core method development.

D Discussion

  • VRD’s Q-value-derived rectification may steer flows toward falsely optimistic, low-density regions when offline Q-networks overestimate OOD actions.The authors suggest incorporating conservative value estimation, such as CQL-style penalties or uncertainty-aware ensembles, to mitigate this risk.Source: D Discussion · p. 29
  • VRD’s continuous vector-field and mapping foundations restrict the current framework to continuous control, making extension to discrete actions non-trivial.The authors suggest exploring Categorical Flow Matching or discrete diffusion, with alignment adapted to categorical probabilities.Source: D Discussion · p. 29
  • VRD-R’s centripetal mapping of high-value behaviors toward the latent origin may collapse distinct equally optimal modes, whereas VRD-Q’s in-place filtering avoids this risk.The authors propose multi-centric gravitational fields, such as Gaussian Mixture priors, to preserve expert multi-modality; they say VRD-Q generally exhibits broader stability across environments.Source: D Discussion · p. 29

NeurIPS Paper Checklist

  • The checklist says the abstract and introduction’s claims are supported by theoretical and experimental evidence and that the paper discusses its limitations.It points to a dedicated limitations paragraph in the Conclusion and Appendix.Source: NeurIPS Paper Checklist · p. 30
  • The checklist reports that experimental settings and reproducibility details are provided, with code supplied to reviewers and public links planned for the camera-ready version.It says the main text references detailed settings, expanded in Appendix C with hyperparameters, optimizers, and architectural configurations.Source: NeurIPS Paper Checklist · pp. 31–32
  • The authors mark the LLM disclosure as N/A because LLMs are not important, original, or non-standard components of the method.They refer readers to the NeurIPS handbook for guidance on what should be described.Source: NeurIPS Paper Checklist · p. 36

Related papers

Loading 2610.29075v1…