GRASP: 장기 지평선에서 세계 모델을 위한 Gradient 기반 계획
BAIR는 GRASP (Gradient RelAxed Stochastic Planner) 를 소개했습니다. 이는 학습된 세계 모델에서 장기 지평선 계획을 실용적으로 만들기 위해 설계된 새로운 Gradient 기반 플래너입니다. GRASP는 궤적을 가상 상태로 올려 병렬 최적화를 수행하고, 상태 반복에 확률성을 주입해 탐색을 촉진하며, 불안정한 상태‑입력 Gradient 대신 안정적인 행동 Jacobian에 의존하도록 Gradient를 재구성합니다.
세계 모델에서 장기 지평선 계획의 도전 과제
대형 세계 모델은 범용 시뮬레이터 역할을 할 수 있지만, 이를 제어와 계획에 활용하는 것은 종종 취약합니다. 장기 지평선 계획은 주로 다음 세 가지 요인 때문에 실패합니다:
- 조건이 나쁜 계산 그래프: 모델을 자체에 반복 적용하면서 미분(시간을 통한 역전파)하면 Gradient가 폭발하거나 소멸합니다. Jacobian의 조건 수는 시간 지평선 $T$에 대해 지수적으로 증가합니다.
- 비탐욕적 지형: 장기 작업은 종종 비탐욕적 행동을 요구합니다(예: 장애물을 우회하기 위해 목표에서 멀어지는 경우). 지평선이 길어질수록 최적화 공간이 확대되어 지역 최소점에 갇힐 가능성이 높아집니다.
- 상태‑입력 Gradient 민감도: 딥러닝 기반 세계 모델에서 상태 입력에 대한 Gradient ($D_s F_\theta$)는 신뢰할 수 없는 경우가 많습니다. 세계 모델은 저차원 데이터 매니폴드 위에서 학습되기 때문에, 상태에 작은 변동을 주어도 원하는 출력을 예측하도록 속일 수 있는 “적대적” 예제에 취약합니다. 이로 인해 최적화 지형이 “끈적”하고 신뢰성이 떨어집니다.
GRASP의 기술적 접근 방식
GRASP는 콜로케이션 기반 계획과 특수한 Gradient 재구성 및 탐색 전략을 결합하여 이러한 실패를 극복합니다.
1. 동역학 제약 올리기 (Collocation)
연속적인 rollout $s_{t+1} = F_\theta(s_t, a_t)$ 대신, GRASP는 동역학을 부드러운 제약으로 취급합니다. 행동과 상태를 동시에 최적화하는 패널티 함수를 정의합니다:
$$\min_{\mathbf{s},\mathbf{a}} \mathcal{L}(\mathbf{s}, \mathbf{a}) = \sum_{t=0}^{T-1} \big|F_\theta(s_t,a_t) - s_{t+1}\big|_2^2, \quad \text{with } s_0 \text{ fixed and } s_T=g$$
이 “올린” 형태는 두 가지 주요 개선을 가능하게 합니다: 모든 $T$ 항을 시간 축을 따라 병렬로 계산할 수 있고, Gradient 신호가 깊은 $T$‑step 합성을 역전파할 필요가 없어집니다.
2. Stop‑Gradient를 통한 Gradient 재구성
상태 Jacobian의 적대적 민감성을 피하기 위해, GRASP는 최적화가 행동 Jacobian ($D_a F_\theta$)에만 의존하도록 보장합니다. 행동 Jacobian은 일반적으로 차원이 낮고 더 밀집하게 학습됩니다.
GRASP는 stop‑gradient 동역학 손실을 구현하여 세계 모델의 상태 입력으로 Gradient가 흐르는 것을 차단합니다. 시스템이 사소한 최소점에 수렴하는 것을 방지하기 위해, GRASP는 밀집 목표 shaping 항을 추가합니다:
$$\mathcal{L}(\mathbf{s},\mathbf{a}) = \mathcal{L}{\text{dyn}}^{\text{sg}}(\mathbf{s},\mathbf{a}) + \gamma \mathcal{L}{\text{goal}}^{\text{sg}}(\mathbf{s},\mathbf{a})$$
이 목표는 취약한 상태 Gradient에 의존하지 않으면서도 목표를 향한 일관된 신호와 실현 가능한 동역학을 균형 있게 유지합니다.
3. 탐색을 위한 상태‑반복 확률성
장기 계획의 비볼록 지형을 탐색하기 위해, GRASP는 최적화 중 가상 상태 업데이트에 직접 Gaussian 노이즈를 주입합니다:
$$s_t \leftarrow s_t - \eta_s \nabla_{s_t}\mathcal{L} + \sigma_{\text{state}} \xi, \qquad \xi\sim\mathcal{N}(0,I)$$
행동이 아니라 상태에 노이즈를 추가함으로써, 플래너는 올린 공간에서 서로 다른 basin 사이를 “점프”하면서도 행동 업데이트는 가이드된 형태를 유지합니다.
4. 주기적 동기화
올린 stop‑gradient 목표는 근사치이므로, GRASP는 매 $K_{\text{sync}}$ 반복마다 “동기화” 단계를 수행합니다. 현재 행동을 사용해 $s_0$부터 rollout을 수행하고, 원래 연속 손실에 대해 몇 번의 작은 Gradient 스텝을 취해 계획이 실제 궤적에 기반을 두도록 보장합니다.
성능 및 결과
GRASP는 Push‑T 환경에서 Cross‑Entropy Method (CEM)와 표준 Gradient Descent (GD)와 같은 기존 플래너에 비해 성공률과 계산 속도에서 크게 향상된 결과를 보여줍니다. 지평선 $H$가 증가함에 따라 GRASP는 더 높은 성공률과 더 낮은 중간 성공 시간을 유지합니다.
| 시간 지평선 | CEM | GD | LatCo | GRASP |
|---|---|---|---|---|
| H=40 | 61.4% / 35.3s | 51.0% / 18.0s | 15.0% / 598.0s | 59.0% / 8.5s |
| H=50 | 30.2% / 96.2s | 37.6% / 76.3s | 4.2% / 1114.7s | 43.4% / 15.2s |
| H=60 | 7.2% / 83.1s | 16.4% / 146.5s | 26.2% / 49.1s | |
| H=70 | 7.8% / 156.1s | 12.0% / 103.1s | 0.0% / — | 16.0% / 79.9s |
| H=80 | 2.8% / 132.2s | 6.4% / 161.3s | 0.0% / — | 10.4% / 58.9s |
(성공률 % / 중간 성공 시간)
향후 방향
연구진은 GRASP를 확장하기 위한 여러 경로를 제시합니다. 여기에는 diffusion 기반 세계 모델에 적용, 보다 정교한 노이징 전략 개발, 플래너를 폐쇄‑루프 시스템이나 적응형 장기 계획을 위한 강화학습 정책 학습에 통합하는 것이 포함됩니다.