Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

One Step Is Enough: Dispersive MeanFlow Policy Optimization

Guowei Zou, Haitao Wang, Hejun Wu, Yukun Qian, Yuhang Wang, Weibing Li

TL;DR

DMPO tackles real-time robotic control by enabling true one-step generative policies. It unites MeanFlow one-step inference, dispersive regularization to prevent representation collapse, and PPO-based RL fine-tuning to exceed expert demonstrations, achieving 5–20x inference speedups and >120 Hz performance on GPUs. The two-stage pipeline—MeanFlow pre-training with a lightweight Vision Transformer and stagewise PPO fine-tuning with BC regularization—produces robust, data-efficient policies that transfer to a Franka Panda robot with low latency. Across RoboMimic and OpenAI Gym benchmarks, DMPO matches or surpasses multi-step baselines, validating the practicality and impact of one-step generation for real-time robotic control.

Abstract

Real-time robotic control demands fast action generation. However, existing generative policies based on diffusion and flow matching require multi-step sampling, fundamentally limiting deployment in time-critical scenarios. We propose Dispersive MeanFlow Policy Optimization (DMPO), a unified framework that enables true one-step generation through three key components: MeanFlow for mathematically-derived single-step inference without knowledge distillation, dispersive regularization to prevent representation collapse, and reinforcement learning (RL) fine-tuning to surpass expert demonstrations. Experiments across RoboMimic manipulation and OpenAI Gym locomotion benchmarks demonstrate competitive or superior performance compared to multi-step baselines. With our lightweight model architecture and the three key algorithmic components working in synergy, DMPO exceeds real-time control requirements (>120Hz) with 5-20x inference speedup, reaching hundreds of Hertz on high-performance GPUs. Physical deployment on a Franka-Emika-Panda robot validates real-world applicability.

One Step Is Enough: Dispersive MeanFlow Policy Optimization

TL;DR

DMPO tackles real-time robotic control by enabling true one-step generative policies. It unites MeanFlow one-step inference, dispersive regularization to prevent representation collapse, and PPO-based RL fine-tuning to exceed expert demonstrations, achieving 5–20x inference speedups and >120 Hz performance on GPUs. The two-stage pipeline—MeanFlow pre-training with a lightweight Vision Transformer and stagewise PPO fine-tuning with BC regularization—produces robust, data-efficient policies that transfer to a Franka Panda robot with low latency. Across RoboMimic and OpenAI Gym benchmarks, DMPO matches or surpasses multi-step baselines, validating the practicality and impact of one-step generation for real-time robotic control.

Abstract

Real-time robotic control demands fast action generation. However, existing generative policies based on diffusion and flow matching require multi-step sampling, fundamentally limiting deployment in time-critical scenarios. We propose Dispersive MeanFlow Policy Optimization (DMPO), a unified framework that enables true one-step generation through three key components: MeanFlow for mathematically-derived single-step inference without knowledge distillation, dispersive regularization to prevent representation collapse, and reinforcement learning (RL) fine-tuning to surpass expert demonstrations. Experiments across RoboMimic manipulation and OpenAI Gym locomotion benchmarks demonstrate competitive or superior performance compared to multi-step baselines. With our lightweight model architecture and the three key algorithmic components working in synergy, DMPO exceeds real-time control requirements (>120Hz) with 5-20x inference speedup, reaching hundreds of Hertz on high-performance GPUs. Physical deployment on a Franka-Emika-Panda robot validates real-world applicability.
Paper Structure (53 sections, 1 theorem, 59 equations, 13 figures, 8 tables, 3 algorithms)

This paper contains 53 sections, 1 theorem, 59 equations, 13 figures, 8 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Multi-Step Generative Policies
  4. One-Step Generation Methods
  5. RL Fine-tuning for Generative Policies
  6. Methodology
  7. Framework Overview
  8. One-Step Generation Policy in DMPO
  9. Dispersive Regularization for Preventing Collapse in One-Step Generation Policy
  10. RL Fine-tuning for One-Step Generation Policies
  11. Experiments
  12. Stage1: Pre-Training Results of DMPO (RQ1 and RQ2)
  13. Robomimc pre-training tasks
  14. Gym Pre-training Tasks
  15. Stage2: Fine-Tuning Results of DMPO (RQ3)
  16. ...and 38 more sections

Key Result

Theorem 3.2

Let $\pi_\theta$ be a MeanFlow policy generating actions via $K$-step Markov chain as defined in Eq. eq:alg_multistep_chain. The policy gradient is: where $o$ is an observation sampled during policy rollouts, $a^{0:K} = (a^0, a^1, \ldots, a^K) \sim \pi_\theta(\cdot|o)$ is the full denoising trajectory sampled from the policy, $A^{\pi_\theta}(o, a^K)$ is the advantage of the final action$a^K$ whic

Figures (13)

  • Figure 1: From efficiency--performance trade-off to practical real-time control. Top: Existing methods lie on the trade-off curve: multi-step approaches (DPPO, ReinFlow, $\pi_{RL}$) achieve strong performance but slow inference, while one-step methods (CP, MP1, 1-DP) are fast but unstable. DMPO breaks this trade-off by occupying the upper-right region. Bottom: DMPO's two-stage approach: dispersive MeanFlow prevents representation collapse during pre-training, and PPO fine-tuning enables adaptation beyond demonstrations.
  • Figure 2: DMPO Framework Overview. Stage 1 (Top & Middle): Pre-training with dispersive MeanFlow. MeanFlow learns velocity fields that transform noise into actions via Vision Transformer encoding with dispersive losses to prevent representation collapse. Stage 2 (Bottom): PPO fine-tuning formulated as a two-layer policy factorization: the outer layer is the true environment MDP, while the inner latent chain only reparameterizes the action distribution.
  • Figure 3: Stage 1 (Pre-training): Inference efficiency vs. success rate trade-off across four RoboMimic tasks and three dispersive weights ($\alpha_{\text{disp}} \in \{0.1, 0.5, 0.9\}$). The upper-left region (fast + high success) is ideal. MF and MF+Disp lie on the Pareto frontier, achieving 6--10$\times$ speedup over ShortCut and 25--40$\times$ over ReFlow while maintaining superior success rates.
  • Figure 4: Stage 1 (Pre-training): Success rate vs. denoising steps on a 4$\times$3 grid. Rows correspond to Lift, Can, Square, and Transport tasks with increasing difficulty, and columns correspond to $\alpha_{\text{disp}} \in \{0.1, 0.5, 0.9\}$. MeanFlow variants, MF and MF+Disp, achieve near-saturated performance at 1--5 steps, while ReFlow and ShortCut require 32--128 steps. Dispersive regularization reduces variance on complex tasks.
  • Figure 5: Stage 2 (PPO Fine-tuning): Performance on RoboMimic manipulation tasks. MeanFlow in blue is compared with DPPO in yellow, Gaussian in green, ReinFlow-S in red, and ReinFlow-R in purple. Lift is omitted as MeanFlow achieves 100% success during Stage 1 pretraining.
  • ...and 8 more figures

Theorems & Definitions (2)

  • Definition 3.1: Multi-Step Markov Policy
  • Theorem 3.2: Multi-Step Policy Gradient