Flow-Based Single-Step Completion for Efficient and Expressive Policy Learning
Prajwal Koirala, Cody Fleming
Abstract
Generative models such as diffusion and flow-matching offer expressive policies for offline reinforcement learning (RL) by capturing rich, multimodal action distributions, but their iterative sampling introduces high inference costs and training instability due to gradient propagation across sampling steps. We propose the Single-Step Completion Policy (SSCP), a generative policy trained with an augmented flow-matching objective to predict direct completion vectors from intermediate flow samples, enabling accurate, one-shot action generation. In an off-policy actor-critic framework, SSCP combines the expressiveness of generative models with the training and inference efficiency of unimodal policies, without requiring long backpropagation chains. Our method scales effectively to offline, offline-to-online, and online RL settings, offering substantial gains in speed and adaptability over diffusion-based baselines. We further extend SSCP to goal-conditioned RL (GCRL), enabling flat policies to exploit subgoal structures without explicit hierarchical inference. SSCP achieves strong results across standard offline RL and GCRL benchmarks, positioning it as a versatile, expressive, and efficient framework for deep RL and sequential decision-making.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 1c54ebcf-2a0b-4c17-a071-674f3c4720ecCited by top-tier papers3
- One-Step Generative Policies with Q-Learning: A Reformulation of MeanFlowZeyuan Wang, Da Li, Yulin Chen, Ye Shi et al.AAAI 2026 · 6 citations
- Offline Reinforcement Learning with Generative Trajectory PoliciesXinsong Feng, Leshu Tang, Chenan Wang, Haipeng ChenICML 2026 · 1 citation
- RAMAC: Multimodal Risk-Aware Offline Reinforcement Learning and the Role of Behavior RegularizationKai Fukazawa, Kunal Mundada, Iman SoltaniICML 2026
Builds on22
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- Conservative Q-Learning for Offline Reinforcement LearningAviral Kumar, Aurick Zhou, George Tucker, Sergey LevineNeurIPS 2020 · 2,881 citations
- Consistency ModelsYang Song, Prafulla Dhariwal, Mark Chen, Ilya SutskeverICML 2023 · 1,720 citations
- Offline Reinforcement Learning with Implicit Q-LearningIlya Kostrikov, Ashvin Nair, Sergey LevineICLR 2022 · 1,402 citations
- A Minimalist Approach to Offline Reinforcement LearningScott Fujimoto, Shixiang Shane GuNeurIPS 2021 · 1,292 citations
Related papers
- Consistency Models as a Rich and Efficient Policy Class for Reinforcement LearningZihan Ding, Chi JinICLR 2024 · 73 citations
- Scaling Offline RL via Efficient and Expressive Shortcut ModelsNicolas A. Espinosa Dice, Yiyi Zhang, Yiding Chen, Bradley Guo et al.NeurIPS 2025 · 28 citations
- One-Step Flow Q-Learning: Addressing the Diffusion Policy Bottleneck in Offline Reinforcement LearningXuan Thanh Nguyen, Chang Dong YooICLR 2026 · 11 citations
- Flow Q-LearningSeohong Park, Qiyang Li, Sergey LevineICML 2025
- Generative Online Reinforcement LearningChubin Zhang, Zhenglin Wan, Feng Chen, Fuchao Yang et al.ICML 2026 · 1 citation
