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

AI summary

Overview

Research area: Offline reinforcement learning (RL) and generative policy learning, sitting at the intersection of flow-based generative modelling (Flow Matching, MeanFlow) and value-based (Q-learning) policy optimisation.

Technical level: Advanced. The paper assumes familiarity with Markov Decision Processes, actor-critic methods, Bellman updates, flow matching, ordinary differential equation (ODE) solvers, and backpropagation through time (BPTT).

Scope: The paper proposes a residual reformulation of MeanFlow that turns a flow-based generative policy into a single-stage, one-step noise-to-action policy that can be trained directly with Q-functions, and evaluates it on 73 tasks across the OGBench and D4RL benchmarks.

What This Paper Is About

Offline RL agents must learn policies from a fixed dataset, and the policy needs to be both expressive (able to represent complex, multimodal action distributions) and efficient (fast to sample and stable to train). One-step Gaussian policies sample quickly but are unimodal, while flow- and diffusion-based policies are expressive but normally require multi-step ODE/SDE integration, which forces BPTT when combined with a critic, or a two-stage distillation pipeline that sacrifices expressivity.

The goal of this paper is to design a generative policy that generates an action from noise in a single forward pass while retaining the expressivity of flow-based models and remaining trainable end-to-end with Q-learning, without distillation and without iterative sampling.

Key Contributions

1. A residual reformulation of MeanFlow for direct noise-to-action mapping. The authors replace the original two-stage MeanFlow pipeline (estimate average velocity, then integrate it into an action) with a single mapping g(a_t, b, t) = a_t − u(a_t, b, t), where a_t is a linear interpolation between an offline-dataset action and noise. This reduces to the original MeanFlow output when b = 0, t = 1, and a_1 = e.

2. Identification and diagnosis of a failure mode in naive MeanFlow policies. The paper shows that decoupling velocity estimation from action generation produces actions outside valid bounds (e.g. [−1, 1]), especially early in training, requiring post-hoc clipping that breaks the alignment between policy outputs and the actions used in the Bellman target. It also reports that a naive residual reformulation a = ε − u(ε, b, t) underfits and fails to capture multimodal distributions in toy experiments.

3. A practical training scheme combining Q-learning with the reformulated policy. This includes a MeanFlow Identity (MFI) regression loss with a stop-gradient target, value-guided rejection sampling over K parallel candidate actions (implemented with vmap in JAX), and an adaptive behaviour-cloning coefficient α adjusted from a moving average of the Q loss.

4. A large empirical evaluation. Results are reported on 73 tasks — 50 state-based and 5 visual-based OGBench tasks across 10 environments, plus 18 D4RL antmaze and adroit tasks — against 10 baselines in both offline and offline-to-online settings, with code released at https://github.com/HiccupRL/MeanFlowQL.

Main Findings

• Offline performance is competitive or best on most task categories. The paper reports that its method achieves best or near-best performance on most of the 73 tasks, with averages over 8 seeds for state-based tasks and 4 seeds for pixel-based tasks. Selected results: OGBench antmaze-large-singletask (5 tasks) 81 ± 3 versus ReBRAC 81 ± 5 and FQL 79 ± 3; humanoidmaze-medium-singletask (5 tasks) 62 ± 1 versus IFQL 60 ± 14 and FQL 58 ± 5; humanoidmaze-large-singletask (5 tasks) 20 ± 3 versus IFQL 11 ± 2; antsoccer-arena-singletask (5 tasks) 62 ± 3 versus FQL 60 ± 2; scene-singletask (5 tasks) 60 ± 1 versus FQL 56 ± 2; puzzle-3x3-singletask (5 tasks) 66 ± 8 versus FQL 30 ± 1; puzzle-4x4-singletask (5 tasks) 40 ± 6 versus IDQL 29 ± 3.

• Not every category is won. On OGBench antmaze-giant-singletask (5 tasks) the method scores 0 ± 0 while ReBRAC reaches 26 ± 8 and FQL 9 ± 6. On cube-single-singletask (5 tasks) it scores 95 ± 2 while FQL scores 96 ± 1 and IDQL 95 ± 2. On cube-double-singletask (5 tasks) it scores 3 ± 2 versus FQL 29 ± 2. On visual manipulation (5 tasks) it scores 55 ± 2 versus FQL 65 ± 2 and ReBRAC 60 ± 2. On D4RL antmaze (6 tasks) it scores 83 ± 2 versus FQL 84 ± 3. On D4RL adroit (12 tasks) it scores 54 ± 3 versus ReBRAC 59.

• The paper attributes weak sparse-task performance to task difficulty and data limitations, not the model. It states that low results on giant maze navigation are consistent with trends across baselines, and that the antmaze-giant-navigate task improves substantially after online fine-tuning, suggesting the bottleneck may lie in the offline dataset.

• Strong offline-to-online transfer. Policies are trained offline for 1M steps and then fine-tuned online. Examples: antmaze-giant-navigate improves from 0 ± 0 to 82 ± 5; humanoidmaze-medium-navigate from 62 ± 1 to 100 ± 1; cube-double-play from 3 ± 2 to 95 ± 2; scene-play from 60 ± 1 to 100 ± 1; puzzle-4x4-play from 40 ± 6 to 100 ± 1; pen-cloned-v1 from 79 ± 3 to 151 ± 7; door-cloned-v1 from 3 ± 1 to 96 ± 4; hammer-cloned-v1 from 10 ± 5 to 132 ± 11; relocate-cloned-v1 from 1 ± 1 to 19 ± 8.

• The method is claimed to match or surpass flow-based baselines with a simpler pipeline. Compared with FQL, IFQL, and FBRAC, the paper reports matching or better performance while using a single-stage training procedure and one-step inference.

• Ablation and toy-experiment details are referenced but not contained in the provided content. The paper points to Appendix B.3 for derivation discussion, Table 4 in the appendix for full results, and an appendix for the toy multimodal-distribution experiments.

Methodology in Plain English

The starting point is MeanFlow, a generative framework that learns the average velocity of a noise-to-data transition over a time interval [b, t], which allows sampling in a single step by evaluating a = e − u(e, b = 0, t = 1). Adapting this directly to RL means predicting a velocity and then integrating it into an action, which the authors find unstable in Q-learning.

Their fix is to fold velocity estimation and action generation into one network output. Rather than predicting velocity and then subtracting it, the policy network directly outputs g(a_t, b, t) = a_t − u(a_t, b, t). Here a_t = (1 − t)·a + t·e is an interpolation between a dataset action a and Gaussian noise e, matching the convention used in flow matching. The authors argue this formulation keeps the representational capacity of the original MeanFlow velocity field (justified via the Universal Approximation Theorem) and, with zero initialisation or small-variance Kaiming initialisation, produces well-bounded one-step outputs (typically within [-1, 1]) even early in training, removing the need for post-hoc clipping.

To train the network, they use the MeanFlow Identity, which relates the average velocity to the instantaneous velocity and its time derivative. The training loss is a squared error between the network output g_θ(a_t, b, t) and a regression target g_tgt, with a stop-gradient applied to the target. The target is g_tgt = a_t + (t − b − 1)·v(a_t, t) − (t − b)·[v(a_t, t)·∂_{a_t}g + ∂_t g], where the instantaneous velocity defaults to v = e − a and the gradient terms come from the current policy network. These derivative terms are computed with a Jacobian-vector product (jvp), so no numerical integration is needed.

On top of the MeanFlow objective, the policy is trained with a Q-learning loss. Because a one-step policy is still stochastic through its input noise, the authors add value-guided rejection sampling: they sample K candidate actions in parallel using vmap in JAX and keep the one with the highest Q-value for the Bellman target. They also adapt the behaviour-cloning coefficient α during training: multiply it by 1.2 if the Q loss exceeds 5 times its moving average, multiply by 0.8 if it falls below 0.2 times the moving average, and leave it unchanged otherwise — a scheme the authors say is inspired by the adaptive KL penalty in PPO. The final objective balances maximising Q_φ(s, a^π) against the MeanFlow Identity loss weighted by α.

Baselines span three policy classes: Gaussian policies (BC, IQL, ReBRAC), diffusion-based policies (IDQL, SRPO, CAC), and flow-based policies (FAWAC, FBRAC, IFQL, FQL), with Cal-QL and PLPD included for the online setting.

Why This Matters

Impact on research. The paper offers a way

Authors’ 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.