Mean Flow Policy Optimization
Abstract
Diffusion models have recently emerged as expressive policy representations for online reinforcement learning (RL). However, their iterative generative processes introduce substantial training and inference overhead. To overcome this limitation, we propose to represent policies using MeanFlow models, a class of few-step flow-based generative models, to improve training and inference efficiency over diffusion-based RL approaches. To promote exploration, we optimize MeanFlow policies under the maximum entropy RL framework via soft policy iteration, and address two key challenges specific to MeanFlow policies: action likelihood evaluation and soft policy improvement. Experiments on MuJoCo, DeepMind Control Suite and HumanoidBench benchmarks demonstrate that our method, Mean Flow Policy Optimization (MFPO), achieves performance comparable to or exceeding current diffusion-based baselines while considerably reducing training and inference time. Our code is available at https://github.com/dongxiaoyi-xyz/MFPO.
Lay Summary
Reinforcement learning is a way for AI systems to learn by trying actions, observing the results, and improving over time. Many existing methods make the AI choose actions in a relatively simple way, which can limit its ability to explore different possibilities. This can cause the AI to miss better solutions, especially in difficult control tasks where several different actions may lead to good outcomes. This paper proposes a more flexible and efficient way for AI agents to choose actions. Instead of relying on slow, multi-step generation methods, it uses a new type of model that can produce diverse actions in only a few steps. This allows the agent to explore more possible solutions while keeping the computational cost low. The proposed method, called Mean Flow Policy Optimization (MFPO), is designed to help reinforcement learning agents learn faster and perform better. Experiments on standard control tasks show that MFPO achieves strong performance comparable to or better than previous advanced methods, while using much less training and decision-making time.