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

2026-05-20

AI Summary

This work addresses the limitations of Gaussian policies in online off-policy reinforcement learning, which struggle to represent multimodal action distributions, and existing generative policies that often lack tractable entropy computation or rely on iterative sampling, thereby compromising expressiveness, exploration, and stability. To overcome these challenges, the paper proposes Stochastic Mean-Flow Policy (SMFP), which introduces the MeanFlow architecture into policy modeling for the first time. SMFP generates actions via a single-step reparameterization, enabling exact entropy estimation and naturally integrating entropy regularization with mirror descent optimization. Evaluated on seven MuJoCo benchmarks, SMFP significantly outperforms both Gaussian and generative baselines while maintaining single-step inference efficiency, effectively breaking through the expressivity bottleneck of conventional policy representations.

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.

Problem

• reinforcement learning
• multimodal action distributions
• entropy estimation
• policy class
• mirror descent

Innovation

• Stochastic MeanFlow Policies
• one-step generative control
• entropic mirror descent
• off-policy reinforcement learning
• tractable entropy