계층형 월드모델이 확산 잠재공간 탐색으로 장기 계획의 누적 오차를 줄인다
Hierarchical World Models with Implicit Dynamics
무엇인가
월드모델 기반 계획은 행동 공간을 탐색해 학습 때 보지 못한 목표에도 테스트 시점에 적응할 수 있지만, 두 가지 벽에 부딪힌다. 계획 지평이 길어질수록 탐색 복잡도가 지수적으로 늘고, 예측 오차가 롤아웃 전반에 누적된다. 기존 계층형 방법들은 고정된 매크로 행동이나 잠재 스킬을 쓰거나, 사전 학습된 스킬·정책, 보상·가치 함수 같은 과제 특화 사전 지식에 의존해 하위 목표의 실현 가능성을 검증했다. 저자들은 그 이유를 계획의 적대적 성질로 설명한다. 이런 제약이 없으면 강력한 옵티마이저가 학습된 월드모델의 결함을 파고들어 비현실적인 전이, 이른바 '웜홀'을 찾아낸다는 것이다.
어떻게 동작하나
ID-WM의 핵심은 상태 공간의 도달 가능성(reachability)과 지역 제어(local control)를 분리하는 것이다. 고수준 모델은 조건부 확산 모델(DDPM)로, 현재 상태와 시간 반경 Δ를 조건으로 받아 Δ 스텝 안에 도달 가능한 미래 상태 분포를 생성한다. 추론 시에는 DDIM의 결정론적 역과정으로 초기 노이즈 벡터 z_t를 하위 목표로 매핑하는데, 이때 z_t 자체를 '암묵적 고수준 행동'으로 취급한다. 이산 스킬 인덱스나 학습된 잠재 코드가 모드 붕괴나 위상적 불연속을 겪는 것과 달리, 등방성 가우시안 잠재공간은 매끄럽고 다루기 쉽다는 것이 저자들의 주장이다. 즉 계획기는 '어디로 갈지'만 가우시안 공간에서 탐색하고, 실행 세부는 하위 모델에 위임한다.
무엇과 다른가
저수준 모델은 과거 k개의 상태·행동 이력을 조건으로 다음 상태를 예측하는 트랜스포머 기반 결정론적 월드모델로, L2 예측 오차로 학습된다. 추론 시에는 이 모델을 역동역학 솔버로 사용한다. 목표 상태 s*가 주어지면 저수준 계획 지평 H_lo에 걸쳐 예측 도달 상태와 s*의 거리를 최소화하는 행동 시퀀스를 탐색해, 상위 계층이 던진 하위 목표를 실제 실행 가능한 행동으로 접지시킨다.
어떻게 쓰나
계획 절차는 모델 예측 제어(MPC) 루프다. 먼저 고수준 월드모델에서 CEM(교차 엔트로피 방법)으로 잠재 노이즈 z_0,…,z_{H_hi-1}을 최적화해 하위 목표 시퀀스를 찾는다. 비용은 최종 예측 상태와 목표의 거리에 정규화 항을 더한 형태이며, 정규화 항은 잠재 벡터의 크기가 임계값 τ를 넘을 때 벌점을 준다. 잠재공간이 이론상 단위 가우시안이어도 제약 없는 최적화는 비현실적 전이를 만드는 적대적 잠재점을 선호할 수 있어, 사전 분포의 고밀도 영역 안에 머물게 하려는 장치다. 이어 저수준 월드모델에서 다시 CEM으로 각 하위 목표에 도달할 행동 시퀀스를 찾아 순차 실행하고, 계획이 끝나면 현재 상태에서 상위 궤적을 재계획한다. 실험에서는 과제별 튜닝 없이 Δ=5(프레임스킵 적용 시 25 환경 스텝), 상위 계획당 하위 목표 5개, 최대 15회 재계획을 사용했고, 목표까지의 거리를 사전에 알 필요가 없다고 밝힌다.
전제와 한계
실험은 Wall, PushT, Maze 세 가지 이동·조작 환경군에서 수행됐고, 과제 지평은 수백에서 수천 스텝에 이른다. 각 환경에서 초기-목표 상태 50쌍을 무작위 샘플링해 평균 성능을 보고했으며 최대 지평은 2000 스텝이다. 비교 대상은 모방학습(GCBC), 오프라인 강화학습(HIQL), 확산 플래너(Diffuser, Hierarchical Diffuser), 단일 계층 월드모델 MPC다. 저자들이 본문에서 제시한 수치는 전체 38.5% 개선, 단일 계층 월드모델 MPC 대비 135% 성능 향상이며, 무작위 목표 PushT에서 데이터 증강 없이 성공률 30%를 넘은 유일한 방법이라고 강조한다. 실패 양상도 정리했는데, GCBC는 분포 내 목표에서는 좋지만 분포 밖 목표에서 크게 나빠지고, HIQL은 가치 함수 학습을 위해 조밀한 데이터 커버리지를 요구하며, Diffuser 계열은 데이터셋 노이즈에 취약하다. 프레임스킵을 키우고 매크로 행동을 학습하는 계층형 잠재 월드모델 HWM(DINO-WM 변형)과의 PushT 비교에서도 ID-WM이 모든 변형에 큰 격차로 앞선다.
확장성 실험도 있다. 3×3 U-Maze부터 13×13 Giant Maze까지 네 개 미로에서 각 2000개의 100스텝 무작위 궤적 데이터를 썼고, 데이터 커버리지가 희박해질수록 베이스라인이 무너지는 반면 ID-WM은 Large Maze에서 최고 성능을 유지했다. 목표 거리를 25, 50, …, 150 스텝으로 늘린 PushT 실험에서는 매 스텝 전체 지평 CEM 탐색과 재계획을 하는 단일 계층 기준선 대비 거의 완벽한 성공률을 보였다. 시각 기반 PushT에서는 3×64×64 이미지를 사전 학습된 Stable Diffusion VAE로 4×8×8 잠재로 인코딩했고, 모든 하위 목표를 한 번에 생성하는 언어 조건부 비디오 플래너 AVDC는 분포 밖 초기-목표 쌍에서 의미 있는 계획을 만들지 못한 반면 ID-WM은 순차적 하위 목표를 안정적으로 생성했다. 정량 비교에서는 DINO-WM을 저수준 모델로 재사용해 계층 구조 자체의 이득을 분리했고, d=50, d=75, 무작위 목표의 난도 높은 설정에서 ID-WM이 가장 좋은 성능을 냈다.
개발자 관점에서 이 논문의 쓸모는 보상 라벨 없는 과제 무관 오프라인 데이터로 월드모델을 학습하고, 테스트 시점에 임의의 목표로 계획한다는 설정에 있다. 계층 구조를 도입해 계획 복잡도가 궤적 길이에 비례해 커지는 것을 막고, 확산 노이즈라는 매끄러운 탐색 공간을 쓴다는 설계는 장기 로봇 조작·내비게이션 파이프라인을 구성할 때 참고할 만하다. 다만 Δ, 상위 하위 목표 수, 잠재 노이즈 정규화 임계값 τ 같은 하이퍼파라미터와, 저수준 모델의 계획 지평이 실제 실행 가능성에 직결된다는 점은 자신의 환경에서 다시 검증해야 한다.
저자들이 밝힌 한계는 두 가지다. 첫째, 확산 잠재공간 안에서의 탐색 기반 최적화에 의존하기 때문에 반응형 정책이나 결정론적 월드모델보다 추론 시간이 길다. 향후 연구로 상각(amortized) 계획이나 그래디언트 기반 계획을 통해 이 병목을 줄이겠다고 한다. 둘째, 계층 구조가 계획 복잡도가 궤적 길이에 따라 늘어나는 것은 막아주지만 상태 공간 차원에 따라 커지는 것은 막지 못한다. 복잡한 환경에서는 압축된 구조적 표현 학습을 결합해 계획 과정을 더 간결하게 만드는 것이 과제로 남아 있다.