💻 computer science

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

SUREFlow is a state-space uncertainty-aware residual flow matching framework built on a Mamba backbone that jointly predicts action velocities and input-dependent uncertainty to enable selective refinement of unreliable actions, achieving robust robot manipulation performance comparable to large vision-language-action models while using significantly fewer parameters.

Original authors: Md Tanvir Islam, Sai Navaneet Peddapalli, Sangmoon Lee, Sangtae Ahn

Published 2026-07-14

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Md Tanvir Islam, Sai Navaneet Peddapalli, Sangmoon Lee, Sangtae Ahn

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to juggle, but instead of throwing balls, it has to move its arm to pick up cups and stack them. Most robot brains today work like a student who memorizes a single path to the answer. If the student trips on a tiny pebble (a little bit of noise or a weird starting position), they panic, stumble, and the whole juggling act falls apart. This is because small mistakes in their speed predictions pile up, like a snowball rolling down a hill, until the robot is doing something completely wrong by the end of the task.

Enter SUREFlow, a new robot brain that acts less like a memorizer and more like a cautious, self-checking navigator.

The Problem: The "Snowball" Effect

The paper argues that current robot policies, especially those that use fancy "generative" methods (like diffusion or flow matching), often assume that their mistakes are the same size everywhere. It's like a driver who thinks they are equally likely to miss a turn by one inch or by a mile. In reality, some parts of a task are tricky and uncertain, while others are easy. When these robots run long sequences of moves without checking their work, those tiny speed errors accumulate, leading to a crash. The authors explicitly rule out the idea that simply making the robot's brain bigger (using massive "billion-parameter" models) is the only way to fix this; they show that a smaller, smarter brain can do just as well.

The Solution: A "Self-Correcting" GPS

The authors propose SUREFlow, which stands for State-space Uncertainty-aware REsidual Flow Matching. That's a mouthful, so let's break it down with an analogy.

Imagine the robot is driving a car on a foggy road.

The Mamba Backbone: Instead of using a giant, heavy engine (like a massive Transformer model that takes forever to compute), SUREFlow uses a lightweight, efficient engine called Mamba. It's like a nimble sports car that can zoom through long distances without getting tired.
The Flow Matching: The robot doesn't just guess the final destination. It learns a "velocity field," which is like a map of wind currents that pushes a boat from a starting point (random noise) to the finish line (the correct action).
The Secret Sauce (SURE): Here is the magic. SUREFlow doesn't just push the boat; it also carries a uncertainty detector. It constantly asks, "How sure am I about this specific part of the move?"
- If the robot is confident (low uncertainty), it keeps cruising.
- If the robot is shaky (high uncertainty), it hits the brakes and makes a tiny, precise correction before it even moves the arm.

This is called Uncertainty-aware Residual Flow (URFlow). It's like a co-pilot who whispers, "Hey, that turn looks slippery, let's adjust the steering wheel just a tiny bit," without needing to ask the human for help. This happens entirely inside the robot's brain, without needing to touch the real world to check if it's right.

The Results: Small Brain, Big Wins

The researchers tested this in a virtual world called LIBERO and Meta-World.

The Score: SUREFlow achieved a 92.5% success rate on the LIBERO tests. This is a huge jump, beating the previous best Mamba-based robot (MaIL) by 34.2%.
The Size: The best part? SUREFlow only has 179 million parameters. To put that in perspective, other top robots (like OpenVLA) need 3 to 7 billion parameters to get similar results. SUREFlow is like a compact car that gets the same mileage as a massive truck.
The Real World: They also tried it on a real robot arm (a Franka Emika) with a camera. The robot had to pick up cups of specific colors and stack them. SUREFlow succeeded 88.5% of the time, proving it isn't just a video game trick.

What They Didn't Do (and Why It Matters)

The paper is very careful about what it doesn't claim.

It does not claim that bigger models are useless; it just shows that for this specific type of long, tricky task, being "uncertainty-aware" is more important than just being "big."
It does not say the robot is perfect. The success rates (like 49% on the harder LIBERO-PRO test) show there is still room for improvement, but it's a massive step forward for such a small model.
It explicitly argues against the idea that you need to rely on huge, expensive backbones to get robust performance. They showed that a lightweight model with a "self-checking" mechanism can outperform the giants.

The Takeaway

In simple terms, SUREFlow teaches robots to be humble. Instead of blindly trusting every move, it learns to spot the moves it's unsure about and fix them instantly. This stops the "snowball" of errors from growing, allowing the robot to complete long, complex tasks reliably, even with a brain that is much smaller and faster than its competitors. The authors suggest that this approach could be the key to making robots that don't just work in perfect labs, but can handle the messy, unpredictable real world.

Technical Summary: SUREFlow

Problem Statement

Generative vision-language-action (VLA) policies have advanced robot manipulation but frequently exhibit instability under noise, partial observability, and stochastic initial conditions. A critical failure mode occurs during extended rollouts: small velocity errors accumulate, degrading execution reliability. Existing diffusion and flow-based policies typically assume homoscedastic residuals (uniform uncertainty) and lack explicit mechanisms to model or correct uncertainty within action generation. Furthermore, many high-performing approaches rely on large transformer backbones with quadratic computational complexity or iterative sampling steps that incur high costs. While state-space models like Mamba offer efficient linear-time modeling, their application to generative robot control remains limited, particularly regarding the integration of uncertainty-aware refinement to stabilize long-horizon execution.

Methodology: SUREFlow

The authors propose SUREFlow (State-space Uncertainty-aware REsidual Flow Matching), a framework that unifies conditional flow matching with uncertainty-aware residual refinement within a Mamba-based backbone.

Core Architecture

Backbone & Conditioning: The model utilizes a Mamba state-space backbone for efficient long-range temporal modeling. Multimodal inputs (multi-view RGB images, robot proprioceptive states, and natural language task embeddings) are encoded into a unified token sequence.
- Flow Conditioned Modulation (FCM): A mechanism that injects continuous flow-time embeddings directly into intermediate feature representations via affine modulation, allowing the network to adapt behavior along the probability path from noise to expert actions without altering token structure.
- Memory-Guided Action Decoder (MGAD): A lightweight module that treats the action token as a learnable query, enabling it to attend to multimodal memory representations (visual, proprioceptive, and task tokens) to enhance temporal conditioning and structured action refinement.
Uncertainty-aware Residual Flow (URFlow):
- Instead of directly regressing actions, the policy learns a time-dependent velocity field that transports a Gaussian prior to the expert action distribution.
- Crucially, SUREFlow models the discrepancy between the predicted velocity and the analytic target velocity as a residual flow.
- The framework jointly predicts the action velocity ( $\hat{v}_t$ ) and an input-dependent uncertainty estimate ( $\hat{s}_t$ , representing log-variance). This formulation treats the residual as heteroscedastic, acknowledging that uncertainty varies across action dimensions and contexts.
Inference-Time Refinement:
- During inference, the policy performs a deterministic refinement procedure guided by the predicted uncertainty.
- At $t=0$ , the model evaluates the predicted uncertainty for each action dimension. Dimensions exceeding a predefined threshold ( $\tau$ ) are identified as unreliable.
- These specific dimensions are selectively refined using the predicted velocity via a closed-loop update ( $a_i \leftarrow a_i - \eta \hat{v}_{0,i}$ ) without requiring external environment feedback. This mitigates error accumulation during iterative flow integration.

Key Contributions

Uncertainty-aware Residual Flow (URFlow): The introduction of a refinement framework within conditional flow matching that learns input-dependent velocity variance. This allows for the selective correction of unreliable action dimensions with negligible computational overhead, stabilizing iterative flow integration.
Memory-Guided Action Decoder (MGAD): A lightweight decoder that re-attends learnable action queries to multimodal memory, enhancing temporal conditioning and robustness in generative action synthesis.
Efficient State-Space Generative Policy: The first integration of uncertainty-aware residual refinement within a state-space generative flow policy. SUREFlow achieves scalable, efficient generative control with only 179M parameters, avoiding reliance on billion-parameter backbones while unifying conditional flow matching with URFlow and MGAD.

Experimental Results

The authors evaluated SUREFlow on LIBERO, Meta-World, and LIBERO-PRO benchmarks, comparing it against transformer-based VLAs (e.g., OpenVLA, Octo), diffusion policies, and other state-space methods (e.g., MaIL).

LIBERO Performance: SUREFlow achieved an average Success Rate (SR) of 92.5%, outperforming the Mamba-based MaIL by 34.2% and transformer-based Octo by 17.4%.
Meta-World Performance: The model attained an average SR of 88.32%, significantly surpassing Diffusion Policy (37.6%) and TinyVLA-H (31.6%).
LIBERO-PRO (Robustness): Under diverse perturbations (object position, semantics, task, environment), SUREFlow achieved a normalized SR of 0.49 with only 179M parameters. This performance is comparable to large-scale VLAs like OpenVLA (7B params, SR 0.52) and $\pi0.5$ (3B params, SR 0.53), demonstrating that uncertainty-guided refinement can compensate for model scale.
Real-Robot Validation: In a real-world color-matching and cup insertion task using a Franka Emika arm, SUREFlow achieved an 88.5% success rate over 20 rollouts per color, validating its robustness against real-world perception noise.
Ablation Studies: Results confirmed that URFlow provides the most significant individual improvement (0.81 SR vs. 0.72 for FCM alone). The combination of all components yielded the highest performance (0.92 SR). Sensitivity analysis showed that moderate uncertainty weighting and 3 refinement steps provided optimal stability.

Significance and Claims

The paper claims that SUREFlow addresses the gap between efficient state-space modeling and robust generative control. By integrating uncertainty-aware residual refinement, the method enables stable long-horizon rollout execution without the computational cost of large transformer backbones or the instability of open-loop generative policies. The authors assert that their approach demonstrates that uncertainty calibration and selective refinement are critical for mitigating compounding errors, allowing a lightweight model (179M params) to achieve performance comparable to billion-parameter VLAs in complex, long-horizon manipulation tasks. The work highlights that the performance gains stem from improved modeling of input-dependent uncertainty rather than brute-force scaling.