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

TL;DR

This work addresses offline reinforcement learning by enabling expressive, multimodal action generation in a single, end-to-end policy without distillation. It reformulates MeanFlow into a residual, direct noise-to-action mapping, $g(a_t,b,t)=a_t-u(a_t,b,t)$, and trains it via a MeanFlow Identity objective to preserve expressive capacity while keeping outputs well-bounded. The method integrates value-guided candidate sampling and an adaptive behaviour cloning coefficient to maintain training stability within Q-learning, achieving strong results across 73 tasks and robust offline-to-online transfer. Practically, this approach provides a scalable, single-stage framework for expressive, one-step generative policies that pair well with value-based learning, improving robustness and adaptation in diverse tasks.

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.

Figures (11)

• Figure 1: Existing generative policy approaches v.s. our proposed solution. Existing methods rely on velocity-based modelling, requiring either multi-step ODE integration or two-stage distillation training. In contrast, our method enables one-step action generation directly from noise, trained via a single-stage, end-to-end process—improving both inference efficiency and training stability.
• Figure 2: Comparison between traditional and proposed flow-based policy network schemes. (a) Given state $s$ and noisy input $a_t$, traditional flow-based methods estimate the instantaneous velocity$v(s, a_t, t)$ at time $t$, as a one-step approximation of multi-step ODE integration. (b) Our method reformulate MeanFlow to estimate the average velocity over the interval $[b, t]$, enabling direct one-step action generation.
• Figure 3: Comparison of learning curves. Naive MeanFlow exhibits significantly larger bound losses in the early training stages. This instability directly undermines convergence of its flow loss, leading to a substantial performance gap compared to ours on the antsoccer-arena-singletask.
• Figure 4: Sensitivity to the behaviour cloning coefficient $\alpha$. Effect of different values of $\alpha$ on offline RL performance.
• Figure 5: Training dynamics under our framework in humanoidmaze-large-navigate-singletask. Despite potential sensitivity, the adaptive coefficient $\alpha$ adjusts dynamically and evolves smoothly, contributing to stable and robust policy learning throughout the training process.

Theorems & Definitions (1)

• proof