Valdi: 가치 확산 월드 모델(Value Diffusion World Models)

잠재 기반 월드 모델은 실시간 Model Predictive Control에 필요한 저지연 예측을 제공하면서도 불확실한 미래를 표현할 수 있어야 한다. Valdi는 잠재 공간에서 확산 모델을 동적으로 학습해 가치와 보상을 함께 예측하는 구조를 온라인 플래닝 루프에 결합함으로써 이 두 요구를 동시에 충족하는 방안을 제시했다. 단일 확산 단계만으로도 결정론적 MLP 기반 대비 동등한 제어 성능을 얻는 점이 실시간 플래닝에서 확산 모델의 실용성을 시사한다.

강화학습에서 Model Predictive Control은 현재 상태로부터 여러 단계의 행동 시퀀스를 가정하고 이들에 따른 미래 보상을 평가해 최적 행동을 선택한다. 실시간 제어에서는 매 플래닝 사이클의 계산 비용이 제한되므로 픽셀 공간 대신 저차원 잠재공간에서 빠르게 롤아웃할 수 있는 dynamics 모델이 선호된다. 기존 TD-MPC 계열은 단일 스텝을 예측하는 결정론적 MLP를 사용해 빠른 추론을 확보했지만, 이러한 결정론적 전이는 다중 모드의 미래를 포착하지 못하는 한계가 있다.

Valdi는 네 가지 구성 요소로 이루어져 있다: CNN 기반의 인코더 Eθ, bidirectional encoder-only Transformer로 구현된 diffusion dynamics Dθ, MLP 보상기 Rθ, MLP 가치기 Vθ. 전체 접근 방식은 관측 시퀀스와 행동 시퀀스를 받아 잠재 시퀀스 z_{t+1:t+H}의 노이즈화된 버전을 조건으로 velocity parameterization을 사용해 전체 H-스텝 잠재 경로를 한 번에 샘플링하는 것이다. 이 샘플은 복원(denoising)되어 보상 및 가치 예측에 투입되며, 예측된 가치 합을 CEM으로 최대화하는 행동 시퀀스가 선택된다.

CarRacing 환경의 예비 실험에서 모델들은 동일한 하이퍼파라미터와 동일한 학습 장치(RTX 4080)로 두 번씩 학습되어 고정된 100개 트랙에서 평가되었다. Valdi는 훈련과 평가에서 단일 확산 단계 사용 시 MLP 베이스라인과 평균 제어 점수에서 동등한 성능을 보였으며 이는 확산 기반 dynamics가 저지연 플래닝 조건에서도 경쟁력이 있음을 시사한다. 추가 실험에서 추론 단계 수를 늘리면 LPIPS가 유의하게 증가했으나 평가 점수는 오히려 떨어져 예측의 시각적 다양성과 제어 성능 사이에 명확한 트레이드오프가 관찰되었다.

전체 아키텍처는 인코더 Eθ, EMA 타깃 인코더 E_{¯θ}, diffusion dynamics Dθ, 보상/가치 헤드로 구성된다. Dθ는 velocity parameterization을 사용하며 노이즈화된 잠재 시퀀스 z^{τ}{t+1:t+H}와 액션 시퀀스 a{t:t+H-1}을 입력으로 받아 denoised velocity hat{u}^{ au}{t+1:t+H}를 예측한다. 이로부터 denoised 잠재 hat{z}{t+1:t+H}를 계산하는 식은 식(1)로 주어지며, 여기서 hat{z}는 sqrt{alpha^{ au}}ar{z}^{ au}-sqrt{1-alpha^{ au}}hat{u}^{ au} 꼴로 얻어진다.

논문에서 명시한 첫 번째 한계는 현재의 단일-step denoising 추정이 더 복잡한 동역학을 가진 환경에서는 정확도를 확보하기 어렵다는 점이다. 그런 경우 다단계 denoising이 필요할 수 있으며 이를 실시간 추론 비용 안에서 사용 가능하게 하려면 multi-step teacher를 single-step student로 distill하는 추가 연구가 요구된다. 두 번째 한계는 추론 시 diffusion steps의 증가가 예측 다양성을 높이지만 오히려 플래닝 성능을 저하시킨다는 점으로, 다양한 추론 스케줄을 허용하려면 훈련 절차의 대응 변화가 필요하다고 명시되어 있다.