ID-WM은 시간적으로 확장된 하위 목표 생성과 국소적 도달 가능성을 분리하는 계층적 세계 모델링 및 계획 프레임워크입니다. 확산 노이즈를 암묵적 행동으로 취급해 계획 탐색 공간을 단위 가우시안으로 제한함으로써 신뢰성 높은 하위 목표 생성을 가능하게 합니다. ID-WM은 평가된 기준 방법보다 뛰어난 성능을 보이며 시각 기반 하위 목표 생성도 입증합니다.
실험
실험에서는 다음 네 가지 질문을 다룹니다.
- 하위 목표 생성: ID-WM은 다양하고 실현 가능하며 도달 가능한 상태를 신뢰성 있게 생성할 수 있는가?
- 일반화: ID-WM은 테스트 시점에 임의의 제로샷 목표에 도달하는 장기 계획 과제를 해결할 수 있는가?
- 확장성: 복잡도가 점차 증가하는 환경에서 ID-WM은 어떤 성능을 보이는가?
- 시각 기반 계획: 고차원 시각 관측만으로 장기 계획을 수행할 때 ID-WM은 효과적인가?
환경
ID-WM은 Wall, PushT, Maze의 세 가지 내비게이션 및 조작 환경 묶음에서 평가됩니다. 과제의 시간 범위는 수백에서 수천 환경 단계에 이릅니다.
Wall
Wall은 문이 있는 벽으로 두 방이 나뉜 단순한 2D 내비게이션 환경입니다. 에이전트는 벽 한쪽의 무작위 시작 위치에서 출발해 문을 통과하여 반대편의 무작위 위치로 이동해야 합니다. 데이터셋은 각각 50단계로 구성된 무작위 생성 궤적 1,920개로 이루어집니다.
기준 방법
ID-WM은 오프라인 학습 방법인 목표 조건부 행동 복제(GCBC), 오프라인 강화 학습(HIQL), 확산 계획기(Diffuser 및 Hierarchical Diffuser), 단일 계층 세계 모델과 모델 예측 제어(MPC)를 사용하는 계획 방법과 비교됩니다. 정성적 시각 기반 실험에서는 ID-WM을 AVDC와도 비교합니다.
단일 계층 세계 모델 및 계획
이 기준 방법은 계층적 추상화 없이 평면형 세계 모델을 사용해 교차 엔트로피 방법(CEM)으로 단일 계층 계획을 수행합니다.
목표 조건부 행동 복제
GCBC는 정책이 현재 상태와 목표 관측값을 입력으로 받아 해당 목표를 향해 이끄는 전문가 행동을 직접 출력하는 모방 학습 패러다임입니다. 구현에서는 궤적에서 무작위로 샘플링한 미래 프레임을 조건으로 정책에 제공하도록 Diffusion Policy를 확장합니다.
HIQL
HIQL은 장기 과제를 위해 상위 수준의 하위 목표 행동 정책과 하위 수준의 하위 목표 조건부 정책을 결합하여 계층적 표현을 학습하는 오프라인 강화 학습 방법입니다.
Diffuser
Diffuser는 전체 상태-행동 시퀀스를 샘플링하여 계획을 수립하는 궤적 수준 확산 모델입니다. 환경에서 각 행동을 실행한 뒤 계획을 다시 수립합니다.
Hierarchical Diffuser
Hierarchical Diffuser는 계획을 두 가지 시간 척도로 분해합니다. 상위 모듈은 중간 하위 목표를 생성하고, 하위 디노이징 네트워크는 이를 사용해 정제된 궤적을 생성합니다. 환경에서 각 행동을 실행한 뒤 두 계층 모두 계획을 다시 수립합니다.
AVDC
AVDC는 비디오 생성 모델로 미래 상태를 예측하고 역운동학(IK)을 사용해 필요한 행동을 도출합니다. 생성된 비디오 시퀀스를 ID-WM의 상위 수준 하위 목표와 비교하는 데 사용됩니다.
HWM
HWM은 상위 수준 모델의 프레임 건너뛰기를 늘리고 학습된 매크로 행동을 계획에 활용하는 동시대의 계층적 잠재 세계 모델입니다. 비교에는 HWM의 DINO-WM 변형을 사용합니다. 이 변형의 공식 구현은 오픈 소스가 아니므로, 보고된 설정과 명시되지 않은 세부 사항에 대한 DINO-WM을 사용해 재구현했습니다.
ID-WM으로 인접 상태 생성
실험에서는 ID-WM의 각 계층이 기저 데이터 분포를 얼마나 정확히 포착하고 서로 다른 시간 단위에 걸쳐 미래 상태를 예측하는지 살펴봅니다. 상위 수준 세계 모델의 경우, 잠재 사전 분포에서 샘플링하면 다양하고 도달 가능한 미래 상태를 얻을 수 있는지 평가합니다.
과제 복잡도에 따른 확장성
실험에서는 환경 크기와 목표 도달 시간이 증가할 때 ID-WM의 확장성을 조사합니다.
목표 도달 시간에 따른 확장성
목표가 25단계 앞에 있는지를 평가하는 PushT 계획 벤치마크를 바탕으로, 검증 세트에서 단계 수가 명시되지 않은 거리만큼 떨어진 목표를 샘플링해 복잡도를 높입니다.
데이터 증강
- 회전 (30%): 모든 위치와 T자 블록의 각도를 장면 중심을 기준으로 명시되지 않은 각도만큼 회전합니다. 속도와 행동도 이에 맞춰 회전합니다.
- 이동 + 회전 (30%): 이동 범위의 절반만큼 이동시키고 명시되지 않은 작은 회전을 함께 적용합니다.
컴퓨팅 자원
상태 기반 세계 모델은 NVIDIA H200 Tensor Core GPU 1~8개를 사용해 학습합니다. GPU 8개로 확장할 때는 분산 데이터 병렬 처리(DDP)를 사용해 학습을 가속합니다. 기준 모델은 단일 NVIDIA H200 Tensor Core GPU에서 수렴할 때까지 학습하고, 계획 평가에도 단일 NVIDIA H200 Tensor Core GPU를 사용합니다.
시각 기반 실험에서는 CDiT 백본을 NVIDIA H200 Tensor Core GPU 32개로 구성된 클러스터에서 DDP를 사용해 학습합니다.
상위 수준 세계 모델의 학습 수렴
학습 초기에는 예측 분포가 크게 흩어져 있으며, PushT 장면의 구조를 유지하지 못하는 경우가 많습니다. 초기 샘플은 물체의 국소 기하 구조나 에이전트와 물체의 상호작용으로 발생하는 물리적으로 타당한 변위를 일관되게 반영하지 못합니다.
학습이 진행될수록 생성된 샘플은 일관된 미래 구성 주변에 더 집중됩니다. 학습 단계가 약 200~600회에 이르면 모델은 각 시간 범위에서 대략적인 도달 가능 영역을 포착하기 시작합니다. 이후 체크포인트에서는 의미 있는 샘플 다양성을 유지하면서 예측 분포가 더욱 정교해집니다.
결론
ID-WM은 시간적으로 확장된 하위 목표 생성과 국소적 도달 가능성을 분리하는 계층적 세계 모델링 및 계획 프레임워크입니다. 확산 노이즈를 암묵적 행동으로 취급하면 계획 탐색 공간이 단위 가우시안으로 제한되어 신뢰성 높은 하위 목표 생성이 가능합니다. ID-WM은 평가된 최첨단 기준 방법보다 뛰어난 성능을 보이며 고품질 시각 기반 하위 목표 생성을 입증합니다. 계층적 추상화를 명시적 국소 동역학에 기반하게 하는 것은 견고한 장기 의사결정을 향한 한 걸음입니다.
한계
확산 잠재 공간에서의 검색 기반 최적화로 인해 ID-WM은 반응형 정책이나 결정론적 세계 모델보다 추론 시간이 더 오래 걸립니다. 향후 연구에서는 이 병목을 해결하기 위해 상각 또는 그래디언트 기반 계획을 살펴볼 예정입니다.
계층 구조를 통해 계획 복잡도가 궤적 길이에 따라 증가하지 않지만, 복잡도는 여전히 상태 공간의 차원에 따라 증가합니다. 압축적이고 구조화된 표현 학습을 통합하면 복잡도가 높은 환경에서 계획을 더욱 효율화할 수 있습니다.
감사의 말 및 연구비 공개
저자들은 소중한 논의와 피드백을 제공한 Amir Bar, Alfredo Canziani, Yilun Kuang, Hengkai Pan, Mike Rabbat, Ying Wang, Wancong Zhang에게 감사드립니다.
본 연구는 AFOSR의 연구비 FA95502310139와, Global AI Frontier Lab 국제 공동 연구 협약(No. RS-2024-00469482 & RS-2024-00509279)에 따라 대한민국 과학기술정보통신부(MSIT)가 지원하는 정보통신기획평가원(IITP)의 연구비로 지원되었습니다.
자주 묻는 질문
ID-WM은 무엇을 위해 설계되었나요? 계층적 세계 모델링과 계획을 결합해 확장된 하위 목표 생성과 국소적 도달 가능성을 분리합니다.
ID-WM은 계획 탐색 공간을 어떻게 제한하나요? 확산 노이즈를 암묵적 행동으로 취급해 탐색 공간을 단위 가우시안으로 제한합니다.
평가에는 어떤 환경이 사용되나요? 평가에는 Wall, PushT, Maze 내비게이션 및 조작 환경이 사용됩니다.
ID-WM의 명시된 한계는 무엇인가요? 검색 기반 최적화는 추론 시간을 늘리며, 계획 복잡도는 여전히 상태 공간의 차원에 따라 증가합니다.