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

기술 요약: SUREFlow

문제 정의

생성적 시각-언어-행동(VLA) 정책은 로봇 조작 분야를 발전시켰으나, 노이즈, 부분 관측성, 확률적 초기 조건 하에서 빈번하게 불안정성을 보입니다. 핵심적인 실패 모드는 장기 실행(extended rollouts) 중에 발생하며, 작은 속도 오차가 누적되어 실행 신뢰도를 저하시킵니다. 기존의 확산(diffusion) 및 흐름 기반(flow-based) 정책은 일반적으로 동분산 잔차(homoscedastic residuals, 균일한 불확실성)를 가정하며, 행동 생성 과정 내에서 불확실성을 모델링하거나 교정하는 명시적인 메커니즘이 부족합니다. 또한, 많은 고성능 접근 방식은 이차 복잡도를 갖는 대규모 트랜스포머 백본이나 높은 비용을 초래하는 반복적인 샘플링 단계에 의존합니다. Mamba와 같은 상태 공간 모델(state-space models)은 효율적인 선형 시간 모델링을 제공하지만, 장기 실행의 안정성을 위해 불확실성 인지 정밀화(uncertainty-aware refinement)를 통합하는 측면에서 이들의 VLA 적용은 여전히 제한적입니다.

방법론: SUREFlow

저자들은 SUREFlow(State-space Uncertainty-aware REsidual Flow)를 제안하며, 이는 Mamba 기반 백본 내에서 조건부 흐름 매칭(conditional flow matching)과 불확실성 인지 잔차 정밀화(uncertainty-aware residual refinement)를 결합한 프레임워크입니다.

핵심 아키텍처

1. 백본 및 컨디셔닝(Backbone & Conditioning): 모델은 효율적인 장기 시계열 모델링을 위해 Mamba 상태 공간 백본을 활용합니다. 다중 모달 입력(다중 뷰 RGB 이미지, 로봇 고유 수용 감각 상태, 자연어 작업 임베딩)은 통합된 토큰 시퀀스로 인코딩됩니다.

   • 흐름 조건부 변조(Flow Conditioned Modulation, FCM): 연속적인 흐름 시간 임베딩을 어파인 변조(affine modulation)를 통해 중간 특징 표현에 직접 주입하는 메커니즘으로, 토큰 구조를 변경하지 않고도 노이즈에서 전문가 행동으로 가는 확률 경로를 따라 동작을 적응시킵니다.
   • 메모리 가이드 행동 디코더(Memory-Guided Action Decoder, MGAD): 행동 토큰을 학습 가능한 쿼리로 취급하는 경량 모듈로, 이를 통해 다중 모달 메모리(시각, 고유 수용 감각, 작업 토큰)에 어텐션을 수행하여 시계열 컨디셔닝과 구조화된 행동 정밀화를 강화합니다.

2. 불확실성 인지 잔차 흐름(Uncertainty-aware Residual Flow, URFlow):

   • 정책은 행동을 직접 회귀하는 대신, 가우시안 사전 분포를 전문가 행동 분포로 이동시키는 시간 의존적 속도장(velocity field)을 학습합니다.
   • 결정적으로, SUREFlow는 예측된 속도와 해석적 타겟 속도 사이의 차이를 **잔차 흐름(residual flow)**으로 모델링합니다.
   • 이 프레임워크는 예측된 속도($\hat{v}_t$)와 입력 의존적 불확실성 추정치($\hat{s}_t$, 로그 분산을 나타냄)를 공동으로 예측합니다. 이 공식은 잔차가 이분산성(heteroscedastic)임을 인정하며, 불확실성이 행동 차원 및 컨텍스트에 따라 달라짐을 명시합니다.

3. 추론 시 정밀화(Inference-Time Refinement):

   • 추론 과정에서 정책은 예측된 불확실성에 의해 유도되는 결정론적 정밀화 절차를 수행합니다.
   • $t=0$에서 모델은 각 행동 차원에 대한 예측 불확실성을 평가합니다. 사전 정의된 임계값($\tau$)을 초과하는 차원은 신뢰할 수 없는 것으로 식별됩니다.
   • 이러한 특정 차원들은 외부 환경 피드백 없이 예측된 속도를 사용하여 폐루프 업데이트($a_i \leftarrow a_i - \eta \hat{v}_{0,i}$)를 통해 선택적으로 정밀화됩니다. 이는 반복적인 흐름 통합 중 발생하는 오차 누적을 완화합니다.

주요 기여

1. 불확실성 인지 잔차 흐름(URFlow): 입력 의존적 속도 분산을 학습하는 조건부 흐름 매칭 내의 정밀화 프레임워크를 도입했습니다. 이를 통해 계산 오버헤드를 거의 발생시키지 않으면서 신뢰할 수 없는 행동 차원을 선택적으로 교정하여 반복적인 흐름 통합을 안정화합니다.
2. 메모리 가이드 행동 디코더(MGAD): 학습 가능한 행동 쿼리를 다중 모달 메모리에 재-어텐션(re-attend)하여, 생성적 행동 합성에서의 시계열 컨디셔닝과 강건성을 향상시키는 경량 디코더입니다.
3. 효율적인 상태 공간 생성 정책: 불확실성 인지 잔차 정밀화를 상태 공간 생성 정책에 통합한 첫 사례입니다. SUREFlow는 단 179M 파라미터만으로도 조건부 흐름 매칭과 URFlow, MGAD를 통합하여 확장 가능하고 효율적인 생성 제어를 달고합니다.

실험 결과

저자들은 LIBERO, Meta-World, LIBERO-PRO 벤치마크에서 SUREFlow를 트랜스포머 기반 VLA(예: OpenVLA, Octo), 확산 정책(diffusion policies) 및 기타 상태 공간 방법(예: MaIL)과 비교 평가했습니다.

• LIBERO 성능: SUREFlow는 평균 성공률(SR) **92.5%**를 달성하여, Mamba 기반의 MaIL보다 34.2%, 트랜스포머 기반의 Octo보다 17.4% 높은 성능을 보였습니다.
• Meta-World 성능: 모델은 평균 SR **88.32%**를 기록하며, Diffusion Policy(37.6%)와 TinyVLA-H(31.6%)를 크게 앞질렀습니다.
• LIBERO-PRO (강건성): 다양한 섭동(물체 위치, 의미론, 작업, 환경) 하에서 SUREFlow는 179M 파라미터만으로 0.49의 정규화된 SR을 달성했습니다. 이 성능은 OpenVLA(7B 파라미터, SR 0.52) 및 $\pi0.5$(3B 파라미터, SR 0.53)와 같은 대규모 VLA와 대등한 수준으로, 불확실성 가이드 정밀화가 모델 규모를 보완할 수 있음을 입증합니다.
• 실제 로봇 검증: Franka Emika 암을 사용한 실제 색상 매칭 및 컵 삽입 작업에서, SUREFlow는 색상당 20회의 롤아웃 동안 **88.5%**의 성공률을 달성하여 실제 환경의 지각 노이즈에 대한 강건성을 입증했습니다.
• 어블레이션 연구: 실험 결과는 URFlow가 가장 유의미한 개별 개선 효과(FCM 단독 0.72 대비 URFlow 0.81 SR)를 제공함을 확인했습니다. 모든 구성 요소를 결합했을 때 가장 높은 성능(0.92 SR)을 나타냈습니다. 민감도 분석을 통해 적절한 불확실성 가중치와 3단계의 정밀화가 최적의 안정성을 제공함을 보여주었습니다.

의의 및 주장

본 논문은 SUREFlow가 효율적인 상태 공간 모델링과 강건한 생성 제어 사이의 간극을 해결한다고 주장합니다. 불확실성 인지 잔차 정밀화를 통합함으로써, 이 방법은 대규모 트랜스포머 백본의 계산 비용이나 개방 루프(open-loop) 생성 정책의 불안정성 없이 안정적인 장기 실행을 가능하게 합니다. 저자들은 이 접근 방식이 불확실성 교정과 선택적 정밀화가 누적 오차를 완화하는 데 핵심적임을 입증하였으며, 이를 통해 경량 모델(179M 파라미터)이 복잡한 장기 조작 작업에서 수십억 개의 파라미터를 가진 VLA와 대등한 성능을 낼 수 있음을 보여줍니다. 이 연구는 성능 향상이 단순한 규모 확장(brute-force scaling)이 아닌, 입력 의존적 불확실성 모델링의 개선에서 비롯되었음을 강조합니다.