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

The gist: DeFlow presents a decoupled offline RL framework that separates manifold modeling from value maximization to achieve superior policy extraction by leveraging multi-step flow matching for behavior cloning and a lightweight refinement module for Q-value maximization.

How it works

DeFlow is structured around two distinct components: a Flow Matching Policy optimized exclusively via supervised learning to capture the behavior manifold, and a lightweight Action Refinement Module that maximizes Q-values via deterministic policy gradients (DDPG). The refinement module learns a residual shift on the generated action rather than distorting the flow model’s weights to find high rewards.

The framework utilizes a composite policy where the final action is derived from a Base Flow adjusted by a Refinement Module. Specifically, the action is formulated as a function of the Base Flow and the refinement module output, with stop-gradient applied to the flow model during refinement updates. This design choice bypasses backpropagating through ODE solvers by treating the base flow as a fixed stochastic generator when updating fϕ.

Decoupling and Constraints

The core philosophy of DeFlow is a structural disentanglement of policy learning, which prevents the competing objectives of density estimation and value maximization from creating a conflicting objective landscape. The Base Flow model is updated solely via the Flow Matching objective LFM(ψ) to ensure it remains an unbiased estimator of the behavior distribution.

To prevent the refinement module from drifting off the valid data manifold (Out-of-Distribution actions), DeFlow employs an automated Lagrangian constraint mechanism. This mechanism dynamically adjusts a penalty based on the divergence between the refined action and the flow support, unlike standard regularization which requires fixed, brittle hyperparameters.

Advantages of Decoupling

This decoupled design yields three significant advantages over prior joint-optimization methods:

** Preservation of Manifold Geometry:**

By isolating the flow training from Q-learning gradients, DeFlow prevents mode collapse, ensuring the policy remains grounded in the dataset support even during aggressive value maximization.

** Optimization Stability & Efficiency:**

DeFlow bypasses the need for computationally expensive BPTT, utilizing efficient, standard gradients for the refinement module to make policy extraction significantly faster and more stable.

** Seamless Offline-to-Online Adaptation:**

The Lagrangian mechanism acts as an auto-tuning gatekeeper; during online fine-tuning, the multiplier automatically adjusts, allowing the refinement module to explore more aggressively without manual hyperparameter retuning.

Hyperparameter Management and Adaptivity

DeFlow introduces a learnable Lagrangian multiplier α that is automatically adjusted to satisfy the constraint by minimizing L(α). This mechanism acts as a feedback controller, where if the refinement becomes too aggressive (∥∆a∥ squared > δ), α increases to enforce manifold adherence, and if the policy is conservative, α decreases to prioritize value maximization.

The target divergence δ is set based on the estimated Intrinsic Action Variance (IAV) of the dataset, calculated by averaging the variance of actions among k-nearest neighbors. This heuristic contrasts sharply with the exhaustive grid search required for the behavioral cloning weight α in prior works.

Experimental Validation

Empirical evaluations on 73 offline RL tasks and 15 O2O tasks demonstrate that DeFlow achieves SOTA or near-SOTA performance on complex environments such as cube-double-play and puzzle-3x3-play. The results show that DeFlow matches or exceeds state-of-the-art methods relying on complex joint optimization.

Furthermore, the analysis of inference scaling reveals a Trust Region Saturation phenomenon, indicating that the refinement module efficiently consumes the optimization budget in a single forward pass, suggesting that simply “sampling more” is futile when the valid optimization landscape is already saturated. The method also shows success in combining with value modeling methods like floq’s value function, observing synergistic improvements without any tuning.

The paper concludes that DeFlow positions itself as a lens to re-examine the interplay between generative diversity and critic-guided optimization, suggesting that navigating the signal-to-noise ratio of the learned Q-function through a deterministic, constrained gradient step could be safer than blind maximization. The next breakthrough is posited to lie in Manifold-Aware Sampling to explore the value landscape while explicitly adhering to the complex topology of the behavior manifold.

References

Agrawalla, B., Nauman, M., Agrawal, K., and Kumar, A. floq: Training critics via flow-matching for scaling compute in value-based rl. arXiv preprint arXiv:2509.06863, 2025

Ankile, L., Simeonov, A., Shenfeld, I., Torne, M., and Agrawal, P. From imitation to refinement – residual rl for precise assembly. In 2025 IEEE International Conference on Robotics and Automation (ICRA), pp. 01–08. IEEE, 2025

Ball, P. J., Smith, L., Kostrikov, I., and Levine, S. Efficient online reinforcement learning with offline data. In International Conference on Machine Learning (ICML), 2023

Chen, H., Lu, C., Wang, Z., Su, H., and Zhu, J. Score regularized policy optimization through diffusion behavior. In International Conference on Learning Representations (ICLR), 2024

Cho, K., van Merrienboer, B., Çaglar Gülçehre, Bahdanau, D., Bougares, F., Schwenk, H., and Bengio, Y. Learning phrase representations using rnn encoder–decoder for statistical machine translation. In Conference on Empirical Methods in Natural Language Processing (EMNLP), 2014

Ding, Z. and Jin, C. Consistency models as a rich and efficient policy class for reinforcement learning. In International Conference on Learning Representations (ICLR), 2024

Dong, P., Zheng, C., Finn, C., Sadigh, D., and Eysenbach, B. Value flows, 2025. URL https://arxiv.org/abs/2510.07650

Espinosa-Dice, N., Zhang, Y., Chen, Y., Guo, B., Oertell, O., Swamy, G., Brantley, K., and Sun, W. Scaling offline rl via efficient and expressive shortcut models. arXiv preprint arXiv:2505.22866, 2025

Fu, J., Kumar, A., Nachum, O., Tucker, G., and Levine, S. D4rl: Datasets for deep data-driven reinforcement learning. ArXiv, abs/2004.07219, 2020

Fujimoto, S. and Gu, S. S. A minimalist approach to offline reinforcement learning. In Neural Information Processing Systems (NeurIPS), 2021

Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. Soft actorcritic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In arXiv, 2018. URL https://arxiv.org/pdf/1801.01290.pdf

Hansen-Estruch, P., Kostrikov, I., Janner, M., Kuba, J. G., and Levine, S. Idql: Implicit q-learning as an actor-critic method with diffusion policies. ArXiv, abs/2304.10573, 2023

He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770–778, 2016

Janner, M., Du, Y., Tenenbaum, J. B., and Levine, S.

1. Bold header: Decoupled Modeling for Complex Manifolds

This system can faithfully capture complex, non-Gaussian behavior manifolds using a multi-step flow matching policy while simultaneously allowing for value maximization via a separate refinement module, thereby resolving the expressivity-optimization conflict.

2. Bold header: Stable Policy Improvement via Residual Refinement

The AI can perform precise value maximization without destroying the manifold geometry by utilizing a lightweight Action Refinement Module that learns a residual shift on the generated action, bypassing computationally prohibitive backpropagation through ODE solvers.

3. Bold header: Dynamic Constraint Management for Robustness

The system employs an automated Lagrangian constraint mechanism where the Lagrangian multiplier dynamically adjusts based on divergence, ensuring preservation of manifold geometry and providing stability during both offline training and online adaptation without requiring fixed, brittle hyperparameters.

4. Bold header: Seamless Offline-to-Online Adaptation

The AI can transition from offline to online training by freezing the Base Flow (behavioral prior) while allowing the refinement module to evolve via the auto-tuned multiplier α, enabling online fine-tuning without manual hyperparameter retuning or structural changes.

5. Bold header: Efficient and Scalable Policy Extraction

The framework achieves this by bypassing BPTT through the Stop-Gradient operator during refinement, making policy extraction significantly faster and more stable compared to methods that require differentiating through the solver.

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.