Stochastic MeanFlow Policies: One-Step Generative Control with Entropic Mirror Descent
cs.LG, cs.AI
Submitted: 2026-05-20
Updated: 2026-10-09
Code: https://github.com/ikostrikov/pytorch-a2c-ppo-acktr-gail
License: http://creativecommons.org/licenses/by/4.0/
The gist: Combining entropy regularisation with mirror descent in online off-policy reinforcement learning requires expressive policies with efficient sampling and explicit entropy control.
Terminology
Abstract
Combining entropy regularisation with mirror descent in online off-policy reinforcement learning requires expressive policies with efficient sampling and explicit entropy control. Gaussian policies struggle to represent multimodal targets, whereas generative policies often require iterative sampling or lack tractable entropy estimates. We introduce Stochastic MeanFlow Policies (SMFP), a one-step generative policy class that combines MeanFlow-based noise-to-action mappings with Gaussian reparameterisation. The latent mapping captures diverse action modes, while a learnable conditional Gaussian scale provides an analytic entropy lower bound. We use this bound to regulate stochasticity alongside direct value guidance and advantage-weighted MeanFlow regression, enabling entropy-regularised mirror-descent learning with one-step policies. Across seven MuJoCo benchmarks, SMFP achieves strong performance against Gaussian and generative baselines with single-step inference. Additional experiments demonstrate multimodal action modelling and effective exploration in sparse-reward navigation.
Sources
- OpenAI Gym
- One-Step Flow Policy Mirror Descent
- How Many Random Seeds? Statistical Power Analysis in Deep Reinforcement Learning Experiments
- IDQL: Implicit Q-Learning as an Actor-Critic Method with Diffusion Policies
- Adam: A Method for Stochastic Optimization
- Reinforcement Learning and Control as Probabilistic Inference: Tutorial and Review
- Back to Basics: Let Denoising Generative Models Denoise
- FLAC: Maximum Entropy RL via Kinetic Energy Regularized Bridge Matching
- Proximal Policy Optimization Algorithms
- Policy Representation via Diffusion Probability Model for Reinforcement Learning
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks