\textit{Stochastic} MeanFlow Policies: One-Step Generative Control with Entropic Mirror Descent

Abstract: Online off-policy reinforcement learning (RL) is shaped by two coupled choices: the policy class and the update rule. Gaussian policies are fast and have tractable entropy, but struggle with multimodal action distributions. Generative policies are more expressive, but often require iterative sampling or lack tractable entropy estimates. On the optimisation side, SAC-style soft policy improvement and mirror descent (MD) can be viewed as minimising different KL divergences: the former moves the policy towards a value-induced Boltzmann distribution, while the latter regularises each update against the previous policy. Combining entropy regularisation with an MD constraint is therefore attractive, as it supports exploration while stabilising policy improvement; however, the resulting target can be multimodal and is poorly matched by unimodal Gaussian policies. We propose Stochastic MeanFlow Policies (SMFP), a one-step generative policy class that maps Gaussian noise to actions through a MeanFlow transformation. This stochastic reparameterisation yields a tractable entropy surrogate and allows MeanFlow policies to be trained within off-policy mirror descent under a unified objective for exploratory yet stable improvement. Across seven MuJoCo benchmarks, SMFP improves over Gaussian and generative baselines while retaining single-step inference efficiency.

• The paper introduces a novel generative policy architecture that uses one-step noise-to-action mappings combined with stochastic reparameterisation to enable tractable entropy maximisation.
• It integrates SAC-style entropy regularisation with mirror descent updates to stabilize training and improve policy expressivity in high-dimensional control tasks.
• Empirical results on MuJoCo tasks demonstrate that SMFP outperforms both unimodal Gaussian and advanced generative baselines in sample efficiency and inference speed.

Stochastic MeanFlow Policies: One-Step Generative Control with Entropic Mirror Descent

Motivation and Context

Policy expressivity and optimisation methodology are central axes in online off-policy RL, with unimodal Gaussian policies dominating due to their computational tractability and closed-form entropy evaluation. However, they fall short in representing multi-modal action distributions inherent to the intersection of entropic exploration (e.g., SAC) and stable update schemes (e.g., Mirror Descent/TRPO/PPO). Flow-based and diffusion policies offer expressivity but incur non-trivial computational and estimation costs—especially regarding density evaluation and entropy estimation, limiting their adoption in online RL. The crucial challenge remains: how to achieve high expressivity, stable policy improvement, and tractable entropy regularisation in a computationally efficient manner.

Methodological Innovations

Stochastic MeanFlow Policies (SMFP)

SMFP constitutes a distinct advance by combining one-step MeanFlow policy architectures (enabling direct, parallelizable noise-to-action mappings) with an explicit stochastic reparameterisation. The model parameterises the conditional policy as

$$g_\theta(a_t, b, t) = a_t - u_\theta(a_t, b, t) + \sigma_\theta(a_t, b, t) \odot \epsilon$$

with $a_t$ as the state-noise interpolation and $\epsilon \sim \mathcal{N}(0, I)$. This schema admits closed-form conditional entropy lower bounds, enabling efficient maximisation of a tractable surrogate even when the marginal density is intractable.

Unified Objective: Entropic Mirror Descent

SMFP integrates the SAC entropy-regularised objective and the MD trust-region regulariser, yielding:

$$\mathcal{L}(\theta) = \mathbb{E}_{s \sim D} \left[ -Q(s, g_\theta(s)) + \alpha \mathcal{L}_\text{Ent}(\theta) + \lambda \mathcal{L}_\text{MD}(\theta) \right]$$

• $\mathcal{L}_\text{Ent}(\theta)$: an entropy-floor regulariser active only below a threshold, controlling policy variance to avoid collapse, normed by hinge penalisation to guarantee lower-bounded entropy.
• $\mathcal{L}_\text{MD}(\theta)$: advantage-weighted MeanFlow regression to the energy-reweighted behaviour distribution, using robust advantage weighting instead of numerically unstable exponentiated weights.

Notably, this objective requires neither explicit marginal density computation nor policy-conditional likelihood ratios, a significant computational advantage over typical generative policies.

Value-Guided Action Selection

SMFP maintains single-step inference (NFE=1) by sampling $K$ candidates per state in parallel and selecting the action maximising the current Q-value. This Best-of-K strategy is efficiently implemented via JAX vmap, amortising the cost and stabilising Q-value estimation, which is particularly critical for high-dimensional environments.

Empirical Results

Extensive evaluations on seven standard MuJoCo control tasks reveal that SMFP consistently outperforms both classical (Gaussian) and state-of-the-art generative (diffusion/flow-based) baselines across most domains, especially on high-dimensional control problems (e.g., Humanoid, Ant). The empirical results exhibit:

• Superior sample efficiency and final episodic returns compared to both unimodal and iterative generative policies.
• Inference speed comparable to unimodal Gaussian policies ($<$0.5ms per action; NFE=1) while yielding higher expressivity and robustness.
• Notable robustness to choice of entropy temperature ($\alpha$) and mirror descent regulariser ($\lambda$), reducing practical hyperparameter tuning burden.
• Strong ablation evidence showing both the entropy surrogate and mirror descent regularisation are necessary: dropping either substantially degrades exploration and/or training stability.

Quantitatively, SMFP achieves the highest mean return in 5 out of 7 tasks (e.g., achieving $a_t$0 on Ant-v4 with $a_t$1ms inference), outperforming DIME, SAC Flow, and QVPO, and closing the sample-efficiency/performance gap with regards to high-cost iterative generative methods.

Theoretical and Practical Implications

The derivation of a tractable entropy surrogate for generative policies under stochastic reparameterisation is a salient contribution, as it enables entropy-regularised RL with policies far more expressive than Gaussian. The formal justification relies on information-theoretic decomposition: maximising the conditional entropy is a valid lower bound for the marginal, and asymptotically equivalent to KL-divergence minimisation against a diffuse prior.

The Best-of-K value-guided sampling inherently couples the representational flexibility of the generative policy with robust Q-filtered action selection, providing practical benefit in sample diversity and mode coverage, crucial for multi-modal tasks.

Critically, this work addresses the policy distribution mismatch problem when soft value targets are inherently multi-modal due to the SAC-MD composite objective, a setting where Gaussian policies are provably insufficient.

Limitations and Directions for Future Work

The tractability of SMFP's objective fundamentally leverages the MeanFlow reparameterisation; extension to more general non-Markovian or hierarchical generative policy architectures will likely require task-specific developments of tractable entropy bounds and regularisation controls. While training remains more costly than vanilla SAC due to repeated sampling from the behaviour policy for stability, it is still significantly more efficient than diffusion-based baselines.

Adaptive scheduling for the entropy temperature and mirror descent regulariser could further stabilise training in nonstationary or highly multimodal environments. Future research can also augment SMFP with distributional critics (e.g., CrossQ, DQNs) to enhance value estimation granularity.

Conclusion

Stochastic MeanFlow Policies represent a highly practical solution to the long-standing expressivity-tractability-stability trilemma in online RL. By merging tractable entropy regularisation with stable mirror descent updates within a one-step generative policy framework, SMFP unlocks new performance regimes, particularly in high-dimensional, multi-modal continuous control tasks. The unified SMFP objective substantiates the importance of rich action modelling for effective entropy-regularised RL and provides a reproducible framework for further work on structured, high-expressivity policy optimisation in both theory and practice.