오프라인 강화학습을 위한 차별가능 세계 모델 기반 모델 예측 제어

초록

본 논문은 오프라인 데이터만으로 학습된 정책과 보상·전이 모델을 활용해, 추론 시에 차별가능한 세계 모델(DWM)을 통해 가상 롤아웃을 수행하고, 그 결과를 기반으로 정책 파라미터를 즉시 미세조정하는 MPC‑형 프레임워크를 제안한다. MuJoCo와 AntMaze 벤치마크에서 기존 오프라인 RL 방법들을 일관적으로 능가한다.

상세 분석

이 연구는 오프라인 강화학습(Offline RL)에서 흔히 발생하는 “분포 이동” 문제를 세계 모델을 이용해 근본적으로 회피한다는 점에서 혁신적이다. 기존 오프라인 RL은 고정된 정책 혹은 가치 함수를 학습한 뒤, 추론 단계에서는 추가 연산 없이 바로 실행한다. 그러나 정책이 학습된 데이터 분포와 실제 실행 시 마주하는 상태‑행동 쌍이 다르면, Q‑값 추정 오류와 행동 선택 오류가 급격히 증폭된다. 저자들은 이러한 한계를 극복하기 위해 두 가지 핵심 아이디어를 도입한다.

첫째, 차별가능 세계 모델(Differentiable World Model, DWM) 파이프라인을 구축한다. DWM은 (i) 조건부 확산 샘플러 fθ — sₜ, aₜ를 입력으로 다음 상태 sₜ₊₁을 샘플링하고, 샘플링 과정 전체를 미분 가능하게 만든다; (ii) 보상 모델 rξ — sₜ, aₜ에 대한 즉시 보상을 예측한다; (iii) 터미널 가치 함수 Qϕ — 프리트레인된 정책의 가치 네트워크를 활용한다. 특히 확산 모델을 사용해 전이 확률 pθ(sₜ₊₁|sₜ,aₜ)를 학습함으로써 복잡한 다중모달 동역학을 정확히 포착하면서도, 역전파가 가능한 연산 그래프를 유지한다는 점이 주목할 만하다.

둘째, 모델 예측 제어(MPC)와 차별가능 롤아웃을 결합한 추론‑시간 최적화 메커니즘을 설계한다. 현재 상태 sₜ에서 정책 πψ를 이용해 행동을 생성하고, fθ와 rξ를 통해 H‑step 가상 트랙을 여러 개(예: N=10) 생성한다. 각 트랙은 즉시 보상과 터미널 가치 Qϕ를 합산해 “대리 반환”(surrogate return)을 계산한다. 이후 이 대리 반환을 손실로 삼아, 정책 파라미터 ψ에 대해 경사 하강법을 수행한다. 핵심은 전체 과정이 미분 가능하다는 점이며, 저자들은 이를 수학적으로 정리한 정리 4.1에서 정책 Jacobian과 확산 Jacobian이 어떻게 연쇄법칙을 통해 결합되는지를 명시한다. 이렇게 얻어진 업데이트된 정책은 실제 환경에 첫 번째 행동만 적용하고, 다음 타임스텝에서 다시 동일 과정을 반복한다.

기술적 강점은 다음과 같다. (1) 전이와 보상을 로컬하게 학습함으로써 장기 가치 추정의 불확실성을 회피한다. (2) 확산 기반 전이 모델은 복잡한 연속 제어 환경에서 멀티모달 행동을 자연스럽게 표현한다. (3) 추론 시 추가 연산(가상 롤아웃·역전파)을 허용함으로써, 고정된 정책보다 상황에 맞는 미세조정을 가능하게 한다. 실험에서는 D4RL의 24개 데이터셋(MuJoCo 18, AntMaze 6)에서 기존 CQL, IQL, Decision Diffuser 등 강력한 베이스라인을 모두 앞선 성능을 기록했으며, 특히 데이터가 희소하거나 보상이 드문 AntMaze에서 큰 이득을 보였다.

한계점도 존재한다. 차별가능 세계 모델을 학습하려면 충분한 전이 데이터와 고품질의 확산 모델 학습이 전제된다. 데이터가 매우 제한적이거나 노이즈가 심한 경우, fθ의 샘플링 정확도가 떨어져 정책 업데이트가 오히려 해를 끼칠 수 있다. 또한 추론 시 여러 롤아웃과 역전파를 수행하므로 실시간 제어에 필요한 계산량이 증가한다는 점에서, 경량화된 구현이나 하드웨어 가속이 필요하다.

전반적으로 이 논문은 “오프라인 RL → 추론‑시간 모델 기반 적응”이라는 새로운 패러다임을 제시하며, 차별가능 세계 모델과 MPC를 결합한 프레임워크가 실제 로봇·자율주행 등 안전이 중요한 연속 제어 분야에 적용될 가능성을 크게 열어준다.