One-Step Generative Policies with Q-Learning: A Reformulation of MeanFlow

2025-11-17

AI Summary

In offline reinforcement learning, one-step policy generation struggles to model multimodal action distributions, while existing flow-based approaches rely on knowledge distillation and two-stage training. This paper proposes MeanFlow-Recon: a single-stage, end-to-end flow-based policy that unifies velocity field modeling and noise-to-action mapping within a single residual network—eliminating separate estimation steps and enabling stable, efficient one-shot sampling. Crucially, MeanFlow-Recon natively integrates with Q-learning without requiring distillation. Evaluated across 73 tasks from OGBench and D4RL, it significantly improves both offline performance and offline-to-online transfer. It achieves high sampling efficiency while enhancing training stability. MeanFlow-Recon establishes a concise, robust, and scalable paradigm for multimodal action modeling in offline RL.

Technology Category

• Machine Learning: Reinforcement Learning
• Multiagent Systems: Multiagent Learning
• Search and Optimization: Learning to Search

Application Category

• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

We introduce a one-step generative policy for offline reinforcement learning that maps noise directly to actions via a residual reformulation of MeanFlow, making it compatible with Q-learning. While one-step Gaussian policies enable fast inference, they struggle to capture complex, multimodal action distributions. Existing flow-based methods improve expressivity but typically rely on distillation and two-stage training when trained with Q-learning. To overcome these limitations, we propose to reformulate MeanFlow to enable direct noise-to-action generation by integrating the velocity field and noise-to-action transformation into a single policy network-eliminating the need for separate velocity estimation. We explore several reformulation variants and identify an effective residual formulation that supports expressive and stable policy learning. Our method offers three key advantages: 1) efficient one-step noise-to-action generation, 2) expressive modelling of multimodal action distributions, and 3) efficient and stable policy learning via Q-learning in a single-stage training setup. Extensive experiments on 73 tasks across the OGBench and D4RL benchmarks demonstrate that our method achieves strong performance in both offline and offline-to-online reinforcement learning settings. Code is available at https://github.com/HiccupRL/MeanFlowQL.

Problem

• Enabling direct noise-to-action mapping for efficient policy generation
• Capturing complex multimodal action distributions in offline RL
• Achieving stable Q-learning integration without multi-stage training

Innovation

• One-step generative policy mapping noise to actions
• Residual MeanFlow reformulation for single-stage training
• Expressive multimodal action modeling with Q-learning