기술 요약: BayesFP: Feynman-Kac 샘플링을 통한 Flow 기반 정책의 사후 확률 추정

문제 정의

실제 환경에서 작동하는 로봇은 멀티모달 전문가 행동을 포착할 수 있을 만큼 표현력이 풍부하면서도, 안전 제약 조건 및 작업 특정 목적을 충족할 수 있을 만큼 반응적인 궤적을 생성해야 합니다. 핵심적인 과제는 이러한 제약 조건(예: 장애물 회피, 관절 한계, 에너지 비용)이 학습 시에는 알려지지 않았거나 배포 사이에 변경될 때 발생합니다.

기존의 제약 조건이 있는 궤적 생성 방식은 일반적으로 다음 두 가지 범주로 나뉩니다:

1. 사후 필터링/최적화(Post-hoc filtering/optimization): 제약 조건이 없는 롤아웃을 경사 기반 플래너나 안전 필터를 사용하여 교정하는 방법입니다. 이러한 방식은 생성과 교정을 분리하기 때문에, 분포 외(out-of-distribution) 궤적, 지역 최적점(local minima) 문제, 높은 계산 오버헤드를 초래할 수 있습니다.
2. 생성 가이던스(Generative guidance): 샘플링 과정을 직접 수정하는 방법(예: 분류기 기반 가이던스, 투영, 반사 역학)입니다. 효과적이긴 하지만, 이는 대개 샘플링 역학에 대한 휴리스틱한 수정에 불과합니다. 이들은 학습된 데이터 사전 확률(prior)과 추론 시 제약 조건을 결합하여 존중하는 잘 정의된 사후 확률 분포와의 원칙적인 관계가 결여되어 있습니다.

핵심 질문은 다음과 같습니다: 사전 학습된 확산(diffusion) 또는 flow-matching 정책이 주어졌을 때, 학습된 데모 분포를 충실히 따르면서도 추론 시에 지정된 제약 최적화 문제를 준수하는 궤적을 어떻게 샘플링할 것인가?

방법론

저자들은 제약 조건이 있는 궤적 생성을 **베이지안 사후 확률 샘플링(Bayesian posterior sampling)**으로 정식화합니다. 학습된 데모 분포 $p_{data}(x)$는 사전 확률(prior) 역할을 하며, 추론 시 제약 조건을 인코딩하는 비용 함수 $J(x)$는 가능도(likelihood) 역할을 합니다. 목표는 다음의 사후 확률로부터 샘플링하는 것입니다:
$$p(x | O = 1) \propto p_{data}(x) \exp(\beta J(x))$$
여기서 $O=1$은 최적성을 나타내며, $\beta < 0$은 비용에 의한 기울기(tilt) 강도를 조절합니다.

기본 정책을 재학습하지 않고 이 사후 확률을 샘플링하기 위해, 본 논문은 Feynman-Kac Corrector (FKC) 프레임워크를 활용합니다. 주요 방법론적 기여는 다음과 같습니다:

1. Feynman-Kac 가중치 역학 (Weighted Dynamics)

정책의 스코어(score)나 속도장(velocity field)에 비용 그래디언트를 휴리스틱하게 더하는 대신, 저자들은 **가중치 부여된 확률 미분 방정식(weighted SDE)**을 도출합니다. 이 SDE는 각 입자가 로그 가중치 $w_t$를 지니는 입자 집단의 진화를 제어합니다. 역학은 타겟 Gibbs-tilted 분포의 Fokker-Planck PDE로부터 유도됩니다.

• 드리프트 교정(Drift Correction): 드리프트 항은 $\nabla J(x)$에 비례하는 비용 인식 교정 항을 포함합니다.
• 가중치 진화(Weight Evolution): 가중치는 시뮬레이션된 역학과 타겟 밀도 진화 사이의 불일치로부터 유도된 재가중치(reweighting) 항을 축적합니다. 이는 드리프트 교정이 불완전하더라도 입자 집단이 사후 확률의 점근적으로 편향되지 않은 추정치(asymptotically unbiased estimator)가 되도록 보장합니다.
• 재샘플링(Resampling): 긴 호라이즌 동안의 가중치 퇴화(weight degeneracy)를 완화하기 위해, 방법론은 순차적 몬테카를로(Sequential Monte Carlo, SMC) 기법, 특히 활성 시간 구간에 대한 체계적 재샘플링을 채택합니다.

2. 결정론적 Flow-Matching으로의 확장

FKC 프레임워크는 원래 확률적 확산 과정(stochastic diffusion processes)을 위해 정식화되었습니다. 저자들은 이를 **한계 보존 SDE(marginal-preserving SDE)**로 재구성함으로써 결정론적 flow-matching 정책으로 확장합니다.

• Singh와 Fischer [16]를 따라, 보조 확산 항 $\tilde{\sigma}(t)$와 그에 대응하는 스코어 기반 드리프트 교정을 도입합니다.
• 결정적으로, 가우시안 소스를 가진 rectified flows의 경우, 스코어 $\nabla \log q_t(x)$를 학습된 속도장 $v_t(x)$로부터 폐쇄형(closed form)으로 임퓨테이션(impute)할 수 있어 별도의 스코어 네트워크가 필요하지 않습니다.
• 이러한 구조를 통해 FKC 프레임워크를 확산 모델과 flow-matching 정책 모두에 일관되게 적용할 수 있습니다.

3. 이론적 보장

본 논문은 페널티 가중치와 역온도 $|\beta|$의 크기가 증가함에 따라, Gibbs-tilted 사후 확률이 실행 가능한(feasible) 근사 최적 궤적에 임의로 가깝게 집중된다는 이론적 보장(Corollary 4.2)을 제공합니다. 이는 이 방법이 외부 안전 필터나 백업 플래너 없이도 제약 조건을 만족할 수 있음을 의미합니다.

주요 기여

1. 원칙적인 정식화: 본 연구는 제약 조건이 있는 궤적 생성을 베이지안 추론으로 재정의하여, 휴리스틱한 가이던스 및 사후 투영(post-hoc projection)에 대한 원칙적인 대안을 제시합니다.
2. 통합 샘플러: 사후 확률을 Feynman-Kac 가중치 SDE로 구현하며, 한계 보존 SDE 재구성을 통해 확산 모델에서 결정론적 flow-matching 정책까지 프레임워크를 확장합니다.
3. 이론적 수렴성: 결과적인 사후 확률이 충분한 페널티 강도 하에서 실행 가능하고 최적에 가까운 솔루션에 집중됨을 증명합니다.
4. 실증적 검증: 사전 학습된 모델(Diffusion Policy, GR00T-N1.6, $\pi0.5$)을 사용하여 시뮬레이션 및 실제 환경의 벤치마크에서 검증하였으며, 추론 시 도입된 비볼록(non-convex) 장애물을 회피하는 데 성공함을 보여주었습니다.

실험 결과

저자들은 BayesFP를 선형 결합 가이던스(Linear-Combination Guidance), CASF (Constraint-Aware Streaming Flow), JM2D (Joint Model-based Model-free Diffusion) 등의 베이스라인과 비교하여 여러 벤치마크에서 평가했습니다.

• 시뮬레이션 작업:
  • Diffusion Policy: 원통형 장애물이 추가된 RoboMimic 작업(Can, Transport)에서 BayesFP는 **충돌률 0%**를 달 기록했습니다 (No Guidance는 30%, Linear Combination은 2%).
  • GR00T-N1.6: 원통형 및 비볼록 V자형 장애물이 있는 LIBERO-Object 작업에서, BayesFP는 84% 및 75%의 성공률을 달성하여 비볼록 기하 구조에서 어려움을 겪은 베이스라인들을 크게 앞질렀습니다.
  • $\pi0.5$: 제로샷 RoboLab 작업(예: FoodPacking, BBQSauce)에서, BayesFP는 선형 결합 베이스라인에 비해 성공률을 높이고(예: HammersInLeftBin에서 36% vs 28%) 충돌률을 줄였습니다.
• 실제 환경 실험:
  • 단 20개의 데모로 미세 조정된 SO101 암을 사용하여, BayesFP는 추론 시 도입된 머그컵 장애물이 있는 PickAndPlace 작업(충돌률을 98%에서 6%로 감소)과 원치 않는 모드를 억제하기 위해 언어 제약이 비용 제약으로 강화된 BimodalMugSelection 작업(베이스라인 72% 대비 100%의 정확한 머그 선택 달성)을 성공적으로 수행했습니다.
• Ablation Study:
  • 입자 수($K$)를 늘리면 성능이 향상되었으나 단순 장애물의 경우 $K=4$ 부근에서 포화되었습니다.
  • 가이던스 강도 $|\beta|$를 높이면 제약 위반이 지수적으로 감소했으며, 이는 이론적 예측과 일치했습니다.

의의 및 주장

본 논문은 BayesFP가 확산 및 flow 정책 모두를 위한 통합된, 추론 시 사용 가능한, 재학습이 필요 없는 샘플러라고 주장합니다. 주요 의의는 다음과 같습니다:

• 원칙적인 가이던스: 분포의 합으로부터 샘플링하는 선형 결합 방식과 달리, BayesFP는 사전 확률과 비용 가능도의 실제 곱(사후 확률)으로부터 샘플링하여 제약 조건에 대한 수학적으로 일관된 처리를 보장합니다.
• 비볼록성 처리: 로컬 그래디언트 교정을 적용하는 대신 사후 확률로부터 전체 궤적을 샘플링함으로써, 비볼록 장애물과 멀티모달 솔루션 공간을 자연스럽게 다룹니다.
• 범용성: 프레임워크는 베이스 정책의 아키텍처(diffusion 또는 flow)와 미분 가능한 비용 함수로 인코딩될 수 있는 제약의 성격에 관계없이 독립적입니다.

저자들은 SDF(Signed Distance Function) 구축의 순차적 특성으로 인한 계산 속도 저하(1.25~2배)와 입자 시뮬레이션 오버헤드 등의 한계를 인정하였으나, 이는 병렬 GPU 실행을 통해 완화될 수 있습니다. 또한 비용 함수가 미분 가능해야 한다는 점을 명시했습니다. 향후 연구로는 저지연 제어를 위한 스트리밍 flow 정책으로의 확장과 더 풍부한 제약 표현의 통합이 제안되었습니다.