시각 세계 모델을 위한 병렬 확률적 그래디언트 플래너 GRASP

초록

본 논문은 비디오 기반 세계 모델의 미분 가능성을 활용하여, 상태를 최적화 변수로 두고 부드러운 동역학 제약을 적용한 병렬 플래닝 방법을 제안한다. 상태에 라플라스 잡음을 주어 탐색성을 높이고, 상태 입력에 대한 그래디언트를 차단함으로써 고차원 시각 모델의 불안정한 기울기를 회피한다. 제안된 GRASP(Gradient RelAxed Stochastic Planner)는 기존 CEM·GD 대비 긴 계획 horizon에서 성공률과 수렴 속도가 크게 향상됨을 실험으로 입증한다.

상세 분석

GRASP는 기존의 시리얼 롤아웃 기반 플래닝이 갖는 두 가지 근본적인 문제—시간에 따라 누적되는 미분 연산으로 인한 수치 불안정성과 고차원 시각 입력에 대한 기울기 민감도—를 해결한다. 첫 번째로, 원래의 최적화 목표인 ( \min_a |F_\theta^T(a,s_0)-g|^2 ) 를 “가상 상태”(virtual states) (z_1,\dots,z_T) 를 도입한 라그랑주식 형태로 변형한다. 이때 동역학 일관성은 (L_{\text{dyn}}=\sum_{t=0}^{T-1}|F_\theta(\bar s_t,a_t)-s_{t+1}|^2) 라는 페널티로 부드럽게 제약되며, (\bar s_t) 는 그래디언트 전파를 차단한 stop‑gradient 복사본이다. 이렇게 하면 각 (F_\theta) 평가가 독립적으로 병렬 실행될 수 있어 GPU 활용도가 크게 상승한다.

두 번째 핵심은 라플라스 스타일의 확률적 상태 업데이트이다. 상태 변수 (s_t) 에 (\sigma_{\text{state}}\xi_t) (표준 정규분포 잡음)를 추가함으로써, 비선형 손실 지형의 지역 최소점에 빠지는 현상을 완화한다. 잡음은 상태 공간을 넓게 탐색하게 하면서도, 동역학 제약을 완전히 포기하지는 않는다. 실제로 (s_{t+1}=s_t-\eta_s\nabla_{s_t}L_{\text{dyn}}+\sigma_{\text{state}}\xi_t) 와 같은 업데이트는 “노이즈가 섞인 라플라스 흐름”을 구현한다.

또한, 상태 입력에 대한 그래디언트를 차단함으로써 발생하는 “목표로 바로 점프하는” 비현실적인 최적화 경로를 방지하기 위해, (L_{\text{goal}}=\sum_{t=0}^{T-1}|F_\theta(\bar s_t,a_t)-g|^2) 라는 밀집 목표 손실을 추가한다. 이 손실은 각 타임스텝마다 행동 (a_t) 에 대한 유용한 신호를 제공하면서도, 상태에 대한 역전파는 차단한다. 결과적으로 (L = L_{\text{dyn}} + \gamma L_{\text{goal}}) 는 행동에 대한 명확한 경사와, 상태에 대한 잡음 기반 탐색을 동시에 제공한다.

알고리즘 흐름은 다음과 같다. (1) 초기 관측 (o_0) 와 목표 (o_g) 를 인코더로 잠재 상태 (s_0, s_T) 로 변환한다. (2) 임의 초기 가상 상태와 행동을 설정하고, (3) 매 반복마다 (a)와 (s)에 대해 공동 그래디언트 스텝을 수행한다. (4) 상태에 라플라스 잡음을 주입하고, 일정 주기마다 전체 롤아웃을 동기화하여 (s_{t+1}=F_\theta(s_t,a_t)) 를 재계산한다. (5) 마지막 동기화 단계에서 목표 손실에 대한 추가 GD 스텝을 적용한다.

이 설계는 두 가지 이점을 제공한다. 첫째, 병렬화 덕분에 (T) 개의 (F_\theta) 평가를 동시에 수행해 연산 시간을 (O(1)) 에 가깝게 감소시킨다. 둘째, 라플라스 잡음과 목표 손실의 조합이 비선형 손실 지형을 부드럽게 만들어, 기존 CEM이 수백 번의 샘플링을 필요로 하는 상황에서도 몇 십 번의 반복만에 수렴한다. 실험에서는 D4RL 및 DeepMind Control Suite의 비디오 기반 환경에서 50~200 스텝 horizon을 대상으로, 성공률이 기존 CEM 대비 평균 +8 ~ +12% 향상되고, 평균 수렴 시간은 절반 이하로 감소하였다.

이론적으로는 (L_{\text{dyn}}) 와 (L_{\text{goal}}) 의 결합이 “동역학 일관성 + 목표 지향성”이라는 두 개의 라그랑주 승수를 동시에 만족시키는 최적화 문제를 정의함을 보이며, stop‑gradient 처리와 라플라스 잡음이 각각 (\nabla_s F_\theta) 와 (\nabla_a L) 에 미치는 영향을 정량화한다. 논문 부록에서는 (\nabla_s F_\theta) 가 고차원 시각 입력에서 고라스코프(Lipschitz) 상수가 크게 되면서 발생하는 “adversarial state gradient” 현상을 정리하고, 제안된 방법이 이를 효과적으로 억제함을 실험적으로 증명한다.

요약하면, GRASP는 (1) 상태를 독립 변수로 두어 병렬화, (2) 상태 입력에 대한 그래디언트 차단, (3) 라플라스 잡음 기반 탐색, (4) 목표 손실을 통한 행동‑중심 경사 제공이라는 네 가지 핵심 아이디어를 결합해, 고차원 비디오 세계 모델에 최적화된 효율적이고 안정적인 플래너를 구현한다.