기술 요약: WarmPrior: 시간적 사전 분포를 활용한 흐름 정합 정책의 직선화

문제 제기

로봇 조작을 위한 생성형 정책, 특히 확산 (diffusion) 과 흐름 정합 (flow matching) 기반 정책들은 시각-운동 제어 분야에서 지배적인 패러다임을 확립해 왔습니다. 이러한 프레임워크에서 신경 필드는 고정된 소스 분포 (일반적으로 등방성 가우시안 $N(0, I)$) 에서 행동 청크의 데이터 매니폴드로 샘플을 운반합니다. 네트워크 아키텍처, 보간자 (interpolants), 적분기 (integrators) 측면에서 진전이 이루어졌음에도 불구하고, 소스 분포는 여전히 거의 정적인 상태로 남아 있습니다.

저자들은 이러한 표준 접근법의 결정적인 한계를 지적합니다. 등방성 가우시안 소스는 상태가 없으며 정보가 부족하여, 로봇 운동의 연속적이고 시간적으로 상관된 특성을 무시합니다. 결과적으로 정책은 매 추론 단계에서 매번 행동 청크를 처음부터 다시 구축해야 합니다. 디노이징 스케줄이 단축됨에 따라 (추론 단계 감소) 시작점에 대한 부담은 커지지만, "맹목적인" 소스는 필요한 시간적 맥락을 제공하지 못해 곡선 확률 경로와 비최적 성능을 초래합니다. 이는 특히 낮은 추론 예산이나 복잡한 작업에서 두드러집니다.

방법론: WarmPrior

본 논문은 readily available 한 최근 행동 이력으로부터 구성된 시간적으로 기반을 둔 사전 분포로, 상태가 없는 표준 소스 분포를 대체하는 WarmPrior를 소개합니다. 이 개입은 최소한으로, 신경망, 보간자, 학습 목적 함수는 변경하지 않고 소스 분포 $p_0$만 수정합니다.

WarmPrior 는 두 가지 변형으로 구현됩니다:

1. WarmPrior-Past (WP-Past): 사전 평균을 직전 실행된 행동 청크에 고정합니다. 소스는 $a_0 = a_{\text{prev}} + \sigma \epsilon$로 샘플링되며, 여기서 $\epsilon \sim N(0, I)$입니다.
2. WarmPrior-Preview (WP-Preview): 정책이 각 단계에서 청크 길이의 두 배 ($2H$) 를 예측하도록 훈련시키지만, 실제로는 첫 $H$만 실행합니다. 나머지 $H$ 단계는 다음 청크에 대한 "미리보기" 역할을 합니다. 추론 시, 다음 단계의 사전 평균은 현재 청크에 대한 모델의 이전 예측치로 설정됩니다 ($a_0 = \hat{a}_{\text{prev}}[H:2H] + \sigma \epsilon$).

두 변형 모두 소스가 결정론적 점 질량이 아닌 적절한 확률 분포로 유지되도록 스케일 $\sigma$를 가진 잔여 가우시안 섭동 (residual Gaussian perturbation) 을 추가합니다. 고정되지 않은 위치 (cold region) 는 표준 가우시안 사전 분포를 유지합니다.

주요 기여 및 메커니즘

1. 흐름 경로의 기하학적 직선화

개선의 주요 메커니즘은 확률 경로의 직선화입니다. 목표 매니폴드 (최근 행동 이력 또는 예측치) 에 가깝게 운송을 초기화함으로써, 학습된 흐름 경로는 더 짧고 직선적으로 변합니다.

• 분기 비용 감소: 저자들은 이를 "분기 비용 (branching cost)" 분석을 통해 공식화합니다. 독립적인 결합 (independent couplings) 을 사용하는 표준 흐름 정합은 속도 네트워크가 모호한 종점들을 평균화하도록 강요하여 곡선 궤적을 생성합니다. WarmPrior 는 암시적 최적 수송 (OT) 결합으로 작용하여 중간 점에 대한 종점의 조건부 분산을 줄입니다. 이는 경로 곡률을 유발하는 회복 불가능한 잔류 오차를 억제합니다.
• 실증적 증거: 실험 결과, WarmPrior 변형들은 베이스라인에 비해 경로 곡률 ($\kappa$) 이 현저히 감소했으며, 이는 성공률 향상과 직접적으로 상관관계가 있었습니다.

2. 조절 가능한 시간적 일관성

WarmPrior 는 시간적 일관성과 다중 모드 표현력 사이의 균형을 맞추는 연속적인 조절 변수 $\sigma$를 도입합니다.

• 암시적 약속: 더 작은 $\sigma$는 새로운 행동 청크가 이전 모드의 분지 (basin) 내에 머물게 하여, 전체 이력에 명시적으로 조건을 부여하지 않고도 "모드 전환 (chunk 경계에서 유효한 행동 모드 간 진동)"을 방지합니다.
• 강건성: 이는 명시적 행동 청크링이 비활성화된 경우 ($H=1$) 에도 정책이 일관성을 유지할 수 있게 하며, 표준 베이스라인이 일반적으로 치명적으로 실패하는 영역에서도 작동합니다.

3. 향상된 사전 공간 강화 학습

이 방법은 사전 공간 (예: DSRL) 의 강화 학습 (RL) 으로 확장됩니다. WarmPrior 로 사전 학습된 정책을 고정된 베이스로 사용하면, RL 에이전트는 원점이 아닌 시간적으로 기반을 둔 평균을 중심으로 한 탐색 공간에서 작동합니다. 이는 탐색 공간을 축소하여 에이전트가 유능한 앵커 주변에서 제한된 잔류 보정을 학습하게 하여, 더 빠른 수렴과 더 높은 점근적 성능을 이끌어냅니다.

실험 결과

시뮬레이션 벤치마크

Diffusion Policy 백본과 GR00T N1.5 VLA 모델을 사용하여 Robomimic(상태 및 이미지 관측) 과 MimicGen(이미지 관측) 에서 평가되었습니다:

• 성공률: WarmPrior 는 모든 작업에서 $N(0, I)$ 베이스라인을 일관되게 능가했습니다.
• 추론 예산: 흐름의 곡률이 가장 중요한 낮은 추론 예산 (NFE = 1) 에서 향상 효과가 가장 두드러졌습니다. 예를 들어, NFE=1 인 어려운 Transport-MH 작업에서 WP-Preview 는 성공률을 23.3% 에서 34.5% 로 향상시켰습니다.
• 실제 로봇: Franka Research 3 에 GR00T N1.5 모델을 탑재하여 배포한 결과, WarmPrior 는 네 가지 탁상 조작 작업 (음식 폐기물 처리, 컵 쌓기, 블록 쌓기, 케이블 삽입) 에서 성공률을 향상시켰으며, 케이블 삽입과 같이 정밀도가 요구되는 작업에서 가장 큰 향상을 보였습니다.

RL 효율성

Robomimic 작업에 대한 사전 공간 RL 실험에서, WarmPrior 변형 (WP-Past 및 WP-Preview) 은 vanilla DSRL 베이스라인 (약 90% 에서 정체됨) 에 비해 더 빠르게 학습하여 더 높은 점근적 성공률에 도달했습니다 (Square 작업에서는 99% 초과, Transport 작업에서는 약 97%).

실시간 청크링 (RTC) 과의 비교

본 논문은 추론 시 겹치는 행동을 고정하는 추론 시간 절차인 실시간 청크링 (RTC) 과 WarmPrior 를 비교합니다. 결과는 두 방법 모두 성능을 향상시키지만 서로 다른 실패 모드를 해결함을 나타냅니다. RTC 는 추론 시 청크 간 불연속성을 수정하는 반면, WarmPrior 는 학습 시 학습된 흐름을 직선화합니다. 두 방법을 결합하면 상승적 개선이 이루어져 상호 보완적임을 확인시켜 줍니다.

중요성 및 주장

본 논문은 소스 분포가 생성형 로봇 제어에서 중요하고 미탐구된 설계 축이라고 주장합니다. "상태가 없는 기본값"에서 "시간적으로 기반을 둔 사전 분포"로 패러다임을 전환함으로써, WarmPrior 는 핵심 네트워크 아키텍처나 학습 손실을 변경하지 않고도 일관된 성능 향상을 달성합니다.

저자들은 WarmPrior 를 복잡한 아키텍처 변경을 대체하는 것이 아니라, 생성 과정 자체에 대한 근본적인 개선으로 위치시킵니다. 이는 다음과 같은 간단하고 플러그 앤 플레이 방식의 메커니즘을 제공합니다:

1. 속도 네트워크에 대한 부담을 줄이기 위해 흐름 궤적을 직선화합니다.
2. 명시적 행동 청크링 없이도 작동하는 시간적 일관성을 위한 조절 가능한 메커니즘을 제공합니다.
3. 탐색 공간을 재구성하여 하류 강화 학습의 샘플 효율성을 향상시킵니다.

이 연구는 생성형 로봇 제어의 향후 진보는 네트워크와 적분기를 넘어 사전 공간 $p_0$ 자체의 구축에 초점을 맞춰야 함을 시사합니다.