TLDR: Flow Matching (FM) policies in robotics suffer performance degradation with increased inference steps due to velocity field instability near the end of integration and overfitting to training actions at late times. A new method, Dense-Jump Flow Matching with Non-Uniform Time Scheduling (FM-DJβ), addresses this by using a U-shaped time sampling during training to emphasize early and late stages, and a Dense-Jump inference strategy that concentrates computation in stable regions and uses a single jump to avoid unstable late-time areas. This approach significantly improves performance and robustness across various robotic tasks, offering up to 23.7% gains over baselines.
In the rapidly evolving field of robotics, generative policies have shown immense promise for teaching robots complex tasks. Among these, Flow Matching (FM) has emerged as a particularly efficient framework, offering a faster alternative to traditional diffusion models by learning a direct velocity field to guide actions. However, a counterintuitive problem has plagued these policies: increasing the number of inference steps, which should ideally refine robot actions, actually degrades performance. This issue has limited the practical application of Flow Matching in real-time robotic scenarios.
The Problem with Multi-Step Inference
Researchers Zidong Chen, Zihao Guo, Peng Wang, ThankGod Itua Egbe, Yan Lyu, and Chenghao Qian, in their paper “Dense-Jump Flow Matching with Non-Uniform Time Scheduling for Robotic Policies: Mitigating Multi-Step Inference Degradation”, delve into the core reasons behind this performance drop. They identify two primary culprits. First, as the integration time approaches its end (t=1), the learned velocity field becomes mathematically unstable, losing a property called Lipschitz continuity. This means that small errors or perturbations get amplified significantly, leading to unpredictable and incorrect robot actions. Second, standard training methods tend to oversample the middle and late stages of the flow trajectory. This causes the policy to ‘memorize’ specific training actions rather than learning to generalize, effectively constraining the robot’s ability to adapt to new situations.
Introducing Dense-Jump Flow Matching
To tackle these challenges, the researchers propose a novel approach called Dense-Jump Flow Matching with Non-Uniform Time Scheduling (FM-DJβ). This method introduces two complementary solutions: a non-uniform time scheduling scheme during training and a Dense-Jump integration strategy during inference.
Non-Uniform Time Scheduling for Smarter Training
During training, instead of uniformly sampling time steps, FM-DJβ uses a U-shaped time schedule, specifically a Beta distribution. This schedule allocates more training emphasis to both the very early and very late stages of the integration process, while reducing focus on the intermediate times. By strengthening supervision at early times (t≈0), the policy learns to align its initial feature space more closely with observations, improving its conditional faithfulness. This helps prevent the policy from drifting towards specific training actions in the mid-interval. By emphasizing late times (t≈1), the network is exposed more frequently to this critical region, encouraging it to learn a smoother velocity field. This implicitly regularizes the model, making the terminal jump more stable.
Dense-Jump Integration for Stable Inference
The second key innovation is the Dense-Jump integration strategy for inference. This method intelligently manages the computational steps. It performs most of the integration steps (N-1 steps) over a stable early interval, typically from t=0 to a ‘jump point’ (tjump). After reaching this jump point, instead of continuing with many small, unstable steps, it performs a single, large ‘jump’ directly to t=1. This avoids repeatedly traversing the late-time region where the velocity field is highly unstable and errors would otherwise accumulate rapidly. The U-shaped training schedule, by smoothing the velocity field near t=1, makes this terminal jump more accurate and robust.
Significant Performance Gains
The effectiveness of FM-DJβ was rigorously tested across diverse robotic tasks, including Walker2D, Adroit Pen Sparse, and Humanoid Standup. The results were compelling. FM-DJβ consistently outperformed both vanilla Flow Matching and diffusion policies, especially in terms of robustness across varying numbers of inference steps. For instance, on the Adroit Pen task, the method achieved a substantial 23.7% performance gain over state-of-the-art baselines. Crucially, it also significantly improved one-step inference performance, making it a highly efficient and reliable option for real-time robotic control. While vanilla Flow Matching often saw performance degrade sharply with more steps, FM-DJβ maintained strong and stable performance, demonstrating its ability to mitigate the instability issues.
Also Read:
- Tenma: A New Approach to Versatile Robot Manipulation
- MEMBOT: Enhancing Robot Reliability in Unpredictable Environments
A Step Forward for Robotic Control
This research provides a principled explanation for the degradation observed in Flow Matching policies and offers a practical, effective solution. By combining intelligent training time scheduling with a strategic inference integration method, FM-DJβ delivers policies that are not only more accurate but also remarkably robust across different computational budgets. This advancement paves the way for more reliable and efficient generative policies in real-world robotic applications, addressing a critical limitation that has hindered their widespread adoption.
