OMP: One-step Meanflow Policy with Directional Alignment

TL;DR

The paper addresses the latency-generalization trade-off in robotic policy learning by presenting OMP, a one-step MeanFlow policy augmented with Directional Alignment via a Cosine Loss and memory-efficient JVP optimization through a Differential Derivation Equation. The approach narrows the gap between high-fidelity diffusion methods and fast-flow methods, delivering real-time, single-step trajectory generation while improving few-shot generalization on Adroit and Meta-World tasks. Key contributions include the Cosine-based directional alignment, the DDE-based JVP approximation, and extensive empirical validation showing superior average success and robustness across challenging tasks. This work offers a practical, scalable solution for high-precision robotic manipulation with improved generalization capabilities in diverse environments.

Abstract

Robot manipulation, a key capability of embodied AI, has turned to data-driven generative policy frameworks, but mainstream approaches like Diffusion Models suffer from high inference latency and Flow-based Methods from increased architectural complexity. While simply applying meanFlow on robotic tasks achieves single-step inference and outperforms FlowPolicy, it lacks few-shot generalization due to fixed temperature hyperparameters in its Dispersive Loss and misaligned predicted-true mean velocities. To solve these issues, this study proposes an improved MeanFlow-based Policies: we introduce a lightweight Cosine Loss to align velocity directions and use the Differential Derivation Equation (DDE) to optimize the Jacobian-Vector Product (JVP) operator. Experiments on Adroit and Meta-World tasks show the proposed method outperforms MP1 and FlowPolicy in average success rate, especially in challenging Meta-World tasks, effectively enhancing few-shot generalization and trajectory accuracy of robot manipulation policies while maintaining real-time performance, offering a more robust solution for high-precision robotic manipulation.

Figures (4)

• Figure 1: Inference Speed vs. Success Rate Trade-off. We compare our proposed OMP against state-of-the-art Diffusion-based methods (DP3, Simple DP3) and Flow-based methods (FlowPolicy, MP1). The x-axis represents control frequency in Frames Per Second (FPS), while the y-axis shows the average success rate (%) across Adroit and Meta-World tasks. The radius of each circle denotes the standard deviation across three random seeds.
• Figure 2: Schematic Comparison of Generative Policy Trajectories. This diagram contrasts the denoising processes of mainstream paradigms. DP3 requires multi-step denoising (10-NFE). FlowPolicy uses segmented straight-line flows but requires consistency constraints. Mean Policy predicts interval-averaged velocity but may suffer from misalignment between the predicted and target velocities due to training discrepancies. Our proposed OMP introduces Directional Alignment via a Cosine Loss, forcing the predicted velocity vector to align explicitly with the true mean velocity direction, ensuring trajectory accuracy in a single step.
• Figure 3: Simulation Benchmark Environments. Visualizations of the diverse manipulation tasks used for evaluation. We test on the high-precision Adroit benchmark and the multi-task Meta-World benchmark, covering a total of 37 distinct tasks ranging from simple household interactions to complex tool use.
• Figure 4: Training Stability and Convergence Analysis. Success rate curves during training for (a) Adroit Hammer and (b) Meta-World Reach Wall.