Dense-Jump Flow Matching with Non-Uniform Time Scheduling for Robotic Policies: Mitigating Multi-Step Inference Degradation
Zidong Chen, Zhihao Guo, Peng Wang, ThankGod Itua Egbe, Yan Lyu, Chenghao Qian
AI summary
Problem
Flow matching policies for robotics suffer from counterintuitive performance degradation as inference steps increase, caused by unstable late-time dynamics and overfitting to training trajectories.
Approach
The authors introduce a U-shaped non-uniform time sampling scheme during training and a dense-jump inference strategy that concentrates computation in stable early regions while taking a single controlled jump to the end to avoid late-time instability.
Key results
- First systematic demonstration that increasing FM inference steps degrades robotic policy performance
- Identification of non-Lipschitz velocity fields and training-action drift as root causes
- Development of a non-uniform time scheduling method to stabilize training phases
- Up to 23.7% performance improvement over state-of-the-art baselines across diverse tasks
Why it matters
It enables reliable, real-time robotic control with flow matching by resolving a critical inference stability issue that limits generalization and efficiency.
Abstract
Flow matching has emerged as a competitive frame- work for learning high-quality generative policies in robotics; however, we find that generalisation arises and saturates early along the flow trajectory, in accordance with recent findings in the literature. We further observe that increasing the number of Euler integration steps during inference counter-intuitively and universally degrades policy performance. We attribute this to (i) additional, uniformly spaced integration steps oversample the late-time region, thereby constraining actions towards the training trajectories and reducing generalisation; and (ii) the learned velocity field becoming non-Lipschitz as integration time approaches 1, causing instability. To address these issues, we propose a novel policy that utilises non-uniform time scheduling (e.g., U-shaped) during training, which emphasises both early and late temporal stages to regularise policy training, and a dense-jump integration schedule at inference, which uses a single- step integration to replace the multi-step integration beyond a jump point, to avoid unstable areas around 1. Essentially, our policy is an efficient one-step learner that still pushes forward performance through multi-step integration, yielding up to 23.7% performance gains over state-of-the-art baselines across diverse robotic tasks. Code is open-sourced at https:// github.com/DenseJumpFM/DenseJump_FlowMatching