SUREFlow: State-space Uncertainty-aware REsidual Flow Matching for Robust Robot Manipulation

2026-07-11

AI Summary

This work addresses the long-horizon execution instability of generative vision-language action policies under noisy, partially observable, and stochastically initialized conditions—instability primarily caused by accumulated velocity errors. To mitigate this, the authors propose a Mamba-based state-space framework that integrates uncertainty-aware residual flow matching. This approach introduces, for the first time in flow matching, input-dependent heteroscedastic uncertainty modeling to jointly predict action velocities and their residual uncertainties, enabling selective optimization of unreliable action dimensions without requiring environmental feedback. Evaluated on the LIBERO benchmark, the method achieves a 92.5% average success rate, surpassing MaIL by 34.2%. On the more challenging LIBERO-PRO benchmark, it attains approximately 49% success with only 179 million parameters, matching the performance of much larger models ranging from 3 to 7 billion parameters.

Abstract

Generative vision-language-action policies have advanced robot manipulation, but they often exhibit instability under noise, partial observability, and stochastic initial conditions. During extended rollouts, small velocity errors accumulate, degrading execution reliability. Existing diffusion and flow-based policies typically assume homoscedastic residuals and lack explicit uncertainty modeling within action generation, limiting robustness during iterative rollout. We propose SUREFlow, a state-space uncertainty-aware residual flow matching framework built on a Mamba backbone. The method jointly predicts action velocities and input-dependent residual uncertainty, enabling selective refinement of unreliable action dimensions without environment feedback while preserving computational efficiency. On LIBERO, SUREFlow achieves 92.5% average success rate (SR), outperforming the Mamba-based MaIL by 34.2%. On LIBERO-PRO, it attains around 49% SR using only 179M parameters, achieving performance comparable to large VLAs with 3-7B parameters. SUREFlow source code is available on: https://github.com/tanvirnwu/SUREFlow

Problem

• robot manipulation
• uncertainty modeling
• flow matching
• partial observability
• action generation

Innovation

• uncertainty-aware
• residual flow matching
• state-space modeling
• Mamba backbone
• robot manipulation