Where Should Action Generation Begin? A Learnable Source Prior for Generative Robot Policies

기술 요약: LeaP - 생성적 로봇 정책을 위한 학습 가능한 소스 사전 분포 (A Learnable Source Prior for Generative Robot Policies)

문제 정의
디퓨전(diffusion) 기반 및 플로우 매칭(flow-matching) 접근 방식을 포함한 생성적 로봇 정책은 일반적으로 행동 생성을 조건부 샘플링으로 정식화합니다. 아키텍처의 다양성에도 불구하고, 이러한 방법들은 공통된 한계를 공유합니다: 바로 관측값에 독립적인 표준 가우시안 분포 $N(0, I)$로부터 행동 생성을 초기화한다는 점입니다. 이 소스 분포는 현재 로봇의 상태나 태스크 컨텍스트에 관한 정보를 전혀 포함하고 있지 않습니다. 결과적으로, 생성 모델은 이 정보가 없는 노이즈 소스로부터 태스크 관련 행동으로 샘플을 이동시키기 위해 상당한 양의 "통합 예산(integration budget)"을 소모해야 합니다. 최근 연구들은 이 표준 가우시안을 결정론적인 관측 유도형 소스(예: 고유 수용 감각 또는 시각적 잠재 변수로부터 얻은 점 추정치)로 대체하려는 시도를 해왔으나, 소스 분포에 내재된 불확실성은 대체로 간과해 왔습니다. 본 연구가 다루는 핵심 질문은 다음과 같습니다: 행동 생성은 어디서 시작되어야 하는가? 구체적으로, 소스가 결정론적이어야 하는지 아니면 분포적이어야 하는지, 그리고 만약 분포적이라면 그 분산이 고정되어야 하는지 아니면 상태에 따라 적응적이어야 하는지를 조사합니다.

방법론: LeaP
저자들은 표준 가우시안 소스를 **고유 수용 감각에 조건화된, 학습 가능한 대각 가우시안(proprioception-conditioned, learnable diagonal Gaussian)**으로 대체하는 프레임워크인 LeaP(Learnable source Prior)를 제안합니다. 주요 구성 요소는 다음과 같습니다:

1. 사전 네트워크 ($g_\phi$): 고유 수용 상태 특징($z_s$)을 대각 가우시안 분포의 평균($\mu_\phi$) 및 상태 적응형 분산($\sigma^2_\phi$)으로 매핑하는 경량 2층 MLP입니다. 결정적으로, 이 사전 분포는 오직 고유 수용 감각에만 조건화되며, 전체 컨디셔닝(시각 + 고유 수용)은 다운스트림 생성기(generator)에 맡깁니다.
2. 확률적 초기화: 추론 시, 소스 샘플 $z_0$는 $N(0, I)$ 대신 학습된 가우시안 $q_\phi(z_0 | z_s)$로부터 추출됩니다. 이는 (상태 의존적 근방에 닻을 내린) 관측 정보가 반영된 동시에 확률적인 초기화를 제공합니다.
3. 플러그 앤 플레이(Plug-and-Play) 설계: LeaP는 오직 소스 분포만을 수정합니다. 다운스트림 생성기 아키텍처(예: 플로우 매칭 또는 디퓨전 브릿지), 생성 역학(generative dynamics), 그리고 추론 솔버는 변경되지 않습니다.
4. 결합 학습 목적 함수: 사전 네트워크와 생성기는 다음의 복합 손실 함수를 사용하여 엔드 투 엔드(end-to-end)로 학습됩니다:
   • 플로우 매칭 손실 ($L_{flow}$): $z_0$를 타겟 행동으로 이동시키기 위해 속도장(velocity field) $v_\theta$를 감독합니다.
   • 음의 로그 가능도 ($L_{NLL}$): 사전 분포의 조건부 밀도를 감독하여, 예측된 평균과 분산이 전문가의 행동 청크(expert action chunks)와 일치하도록 합니다.
   • 대조 정렬 손실 ($L_{align}$): 학습된 임베딩 공간에서 소스 샘플과 행동 간의 인스턴스 수준 대응 관계를 강제하는 CLIP 스타일의 대칭 손실입니다.

주요 기여

• 소스 분포의 재정의: 본 논문은 소스 분포를 생성적 역학의 선택과는 별개로, 독립적이고 재사용 가능한 설계 축(design axis)으로 식별합니다.
• 상태 적응형 분산: 결정론적 소스나 고정 분산 휴리스틱과 달리, LeaP는 평균과 상태 적응형 분산을 공동으로 학습하여 소스의 불확실성을 명시적으로 모델링합니다.
• 모달리티 선택: 저자들은 고유 수용 감각만으로도 사전 분포를 매개변수화하기에 충분하며, 사전 분포에 시각적 특징을 통합하는 것이 오히려 성능을 저하시킨다는 것을 입증했습니다.
• 플러그 앤 플레이 호환성: LeaP는 내부 아키텍처를 변경하지 않고도 다양한 디퓨전 브릿지 생성기의 성능을 개선합니다.

실험 결과
저자들은 RoboTwin 2.0 시뮬레이션 환경에서의 15가지 양손 조작 태스크와 Franka Research 3 로봇을 이용한 3가지 실제 환경 태스크에서 LeaP를 평가했습니다.

• 시뮬레이션 성능: LeaP는 평균 성공률 **81.6%**를 달 기록하며, 4개의 대표적인 베이스라인(결정론적 소스 방식인 A2A 및 VITA, 사전 분포가 없는 NoPrior 대조군, 디퓨전 브릿지 정책)보다 6.5 ~ 25.5 퍼센트 포인트(pp) 높은 성능을 보였습니다. 가장 큰 격차(25.5 pp)는 동일한 생성기 아키텍처를 사용하면서 $N(0, I)$를 사용한 "NoPrior" 베이스라인에서 관찰되었으며, 이는 이 이득이 학습 가능한 사전 분포 자체에서 기인함을 확인시켜 줍니다.
• 효율성: LeaP는 ACT나 BridgePolicy와 같은 베이스라인에 비해 훨씬 적은 파라미터(총 13.22M, 사전 헤드는 0.21M에 불과)를 사용합니다. 또한 더 빠르고 안정적인 학습 수렴을 보여줍니다.
• 실제 환경 배포: 실제 실험에서도 LeaP는 **80.0%**의 평균 성공률로 최고 성능을 유지하며, 차순위 베이스라인인 A2A보다 11.7 퍼센트 포인트를 앞섰습니다.
• 절제 연구(Ablation Studies):
  • 입력 모달리티: 사전을 고유 수용 감각에만 조건화했을 때 최상의 결과를 얻었으며, 시각적 특징을 사전과 융합하면 성능이 감소했습니다.
  • 매개변수화: 평균과 상태 적응형 분산을 공동으로 학습하는 것이 매우 중요했습니다. 결정론적 또는 고정 분산 대안들은 더 낮은 성능을 보였습니다.
  • 생성기: 사전은 다양한 생성기 백본 전반에 걸쳐 일관되게 성능을 향상시켰습니다.

의의 및 주장
본 논문은 소스 분포가 생성적 로봇 정책에서 중요하지만 아직 충분히 탐구되지 않은 설계 차원이라고 주장합니다. 정보가 없는 표준 가우시안에서 학습 가능한 고유 수용 감각 조건부 확률 분포로 전환함으로써, 생성기는 먼 노이즈 소스로부터 샘플을 이동시켜야 하는 부담에서 벗어납 수 있습니다. 대신, 사전 분포는 초기 샘플을 "고유 수용적으로 타당한 영역(proprioceptively plausible region)"에 배치하여, 생성기가 정밀하고 관측치에 조건화된 정교화(refinement)에 집중할 수 있도록 합니다. 저자들은 이 접근 방식이 기존의 생성 역학 선택을 보완하면서, 정책 성능, 수렴 속도 및 강건성을 향상시킬 수 있는 낮은 오버헤드의 레버(lever)를 제공한다고 단언합니다.