Adaptive Linear Path Model-Based Diffusion

기술 요약: 적응형 선형 경로 모델 기반 확산 (Adaptive Linear Path Model-Based Diffusion)

문제 정의
확산 모델(Diffusion Models)을 모델 기반 제어(Model-based Control)와 통합하는 방식은 비선형 역학 및 비볼록 제약 조건이 포함된 복잡한 궤적 최적화 문제를 해결하는 데 유망한 가능성을 보여주었다. 그러나 기존의 모델 기반 확산(Model-Based Diffusion, MBD) 접근 방식은 노이즈 스케줄링 파라미터 설계에 매우 민다민하다. 표준 분산 보존(Variance-Preserving, VP) 스케줄에서는 노이즈 엔드포인트(endpoints)와 총 확산 단계(diffusion steps)와 같은 파라미터들이 서로 밀접하게 결합되어 있다. 이러한 결합성은 개별 파라미터가 탐색-정제(exploration–refinement) 트레이드오프에 미치는 영향을 분리하여 파악하기 어렵게 만들며, 이는 종종 지나치게 보수적인 선택을 강요하거나 광범위한 시행착오 기반의 튜닝을 요구하게 된다. 또한, 최적의 파라미터는 시스템 상태와 작업 복잡도에 따라 빈번하게 변하지만, 전통적인 MBD 프레임워크는 이러한 변화하는 조건에 적응할 수 없는 고정된 스케줄을 사용한다.

방법론
저자들은 이러한 한계를 해결하기 위해 두 단계 프레임워크인 선형 경로 모델 기반 확산(Linear Path Model-Based Diffusion, LP-MBD)과 적응형 LP-MBD(Adaptive LP-MBD, ALP-MBD)를 제안한다.

1. 선형 경로 모델 기반 확산 (LP-MBD):
   LP-MBD는 Flow Matching에서 영감을 받은 선형 확률 경로(linear probability path)를 사용하여 표준 분산 보존 노이즈 스케줄을 대체한다. 복잡한 보간법을 사용하는 VP-MBD와 달리, LP-MBD는 깨끗한 궤적 $Y^{(0)}$와 표준 가우시안 노이즈 $\epsilon$ 사이의 선형 보간을 정의한다:
   $$Y_t = (1 - t)Y^{(0)} + t\epsilon, \quad t \in [0, 1]$$
   이 공식은 기하학적으로 해석 가능하고 디커플링된(decoupled) 파라미터화를 제공한다. 세 개의 결합된 파라미터($\beta_0, \beta_1, T$)를 튜닝해야 하는 VP-MBD와 달리, LP-MBD는 최대 노이즈 스케일($\sigma_{max}$)과 확산 단계 수($T$)라는 두 개의 독립적인 파라미터에만 의존한다. 최대 노이즈 스케일은 허용 가능한 제어 입력 범위에 의해 결정되므로, 탐색 한계를 설정하는 데 있어 직접적이고 해석 가능한 방법을 제공한다. 효율성을 보장하기 위해, 가우시안 제안(Gaussian proposal)의 분산이 발산하는 것을 방지하고자 스케줄은 $t_{max} < 1$인 최대값에서 절단(truncate)된다.

2. 적응형 LP-MBD (ALP-MBD):
   LP-MBD의 디커플링된 특성을 바탕으로, 저자들은 환경 상태에 따라 확산 단계($T$)와 노이즈 레벨($\sigma_{max}$)을 동적으로 조정하는 강화 학습(RL) 기반의 적응형 스케줄러를 도입한다.

• 정식화: 적응형 스케줄러는 마르코프 결정 과정(MDP) 내의 정책 $\pi_\phi$로 취급된다. 상태 공간은 환경 변수(예: 장애물 거리, 궤적 오차)를 포함하며, 액션 공간은 스케줄링 파라미터 $(T, \sigma_{max})$로 구성된다.
• 학습: 시스템은 PPO(Proximal Policy Optimization)를 사용하여 학습된다. 보상 함수는 작업 성능과 더불어, 계산 효율성을 장려하기 위해 큰 확산 단계에 대한 페널티를 결합한다.
• 메커니즘: 복잡한 시나리오(예: 장애물 회피)에서 에이전트는 탐색을 넓히고 솔루션을 정제하기 위해 $T$와 $\sigma_{max}$를 증가시키는 법을 배운다. 단순하거나 구조화된 조건에서는 솔루션 품질을 저해하지 않으면서 샘플링을 가속화하기 위해 이 파라미터들을 줄인다.

주요 기여

1. LP-MBD: 기하학적으로 해석 가능하고 디커플링된 파라미터화를 제공하는 Flow Matching 영감의 선형 경로 스케줄러를 도입하였다. 이는 VP 스케줄과 비교하여 하이퍼파라미터 튜닝 부담을 크게 줄여준다.
2. ALP-MBD: 스케줄링 파라미터를 온라인으로 최적화하기 위해 강화 학습을 사용하는 적응형 확장 모델을 개발하였다. 이를 통해 시스템은 작업 복잡도와 환경 조건에 맞춰 조정함으로써 견고성과 효율성의 균형을 맞출 수 있다.
3. 포괄적 평가: 수치 연구, Brax 벤치마크(Hopper 및 HalfCheetah와 같은 로코모션 태스크 포함), 그리고 모바일 로봇 궤적 추적 태스크를 통한 두 방법론의 엄격한 검증을 수행하였다.

결과

• 수치 연구: 1D 및 2D 최적화 태스크에서 LP-MBD는 특히 확산 단계가 변할 때 VP-MBD에 비해 우수한 수렴 안정성을 보여주었다. VP-MBD는 $T$가 변할 때 종종 수렴에 실패하거나 $\beta$ 파라미터의 광범위한 재튜닝을 요구한 반면, LP-MBD는 디커플링된 파라미터 덕분에 일관된 성능을 유지하였다. 적응형 실험에서 ALP-MBD는 결합된 파라미터 공간으로 인해 어려움을 겪은 적응형 VP-MBD(AVP-MBD) 베이스라인보다 더 빠르게 수렴하고 더 높은 품질의 솔루션을 찾아냈다.
• Brax 벤치마크: LP-MBD는 대부분의 연속 제어 태스크(Ant, Hopper, HalfCheetah, Walker2D, Reacher)에서 CEM, MPPI 및 VP-MBD와 비교하여 경쟁력 있거나 더 우수한 단계당 보상(per-step rewards)을 달성하였다. Pusher 태스크에서는 VP-MBD가 LP-MBD보다 약간 더 높은 성능을 보였는데, 저자들은 이를 고차원성이 분산 보존 스케줄의 표현 능력(expressive capacity)을 더 유리하게 만들었기 때문이라고 설명한다.
• 모바일 로봇 궤적 추적: 장애물이 있는 S자형 및 타원형 경로를 가진 커스텀 환경에서, ALP-MBD가 모든 변형 모델 중 가장 높은 평균 보상을 달성하였다. 시스템은 장애물 근처를 통과하거나 급격한 회전을 할 때 확산 단계와 노이즈 분산을 높이고, 장애물이 없는 직선 구간에서는 이를 줄이는 방식으로 실시간 적응을 성공적으로 수행하였다. RL 정책의 포워드 패스(forward pass)로 인한 오버헤드에도 불구하고, ALP-MBD는 제어 애플리케이션에 적합한 실시간 효율성을 유지하였다.

의의
본 논문은 LP-MBD가 노이즈 스케줄 튜닝의 악명 높은 어려움을 완화하는, 기하학적으로 해석 가능한 안정적인 기초를 제공한다고 주장한다. 노이즈 스케일을 이산화 지평(discretization horizon)으로부터 분리함으로써, LP-MBD는 적응을 위한 강화 학습의 효과적인 적용을 가능하게 한다. 제안된 ALP-MBD 프레임워크는 확산 파라미터의 동적 조절이 복잡한 환경에서의 견고성과 단순한 환경에서의 계산 효율성을 모두 향상시킬 수 있음을 입증하며, 작업 조건이 크게 변하는 실제 로봇 제어 분야에 실질적인 솔루션을 제공한다.