다중과제 강화학습을 위한 모듈형 세계 모델 Mixture‑of‑World

초록

Mixture‑of‑World Models(MoW)는 시각적 멀티태스크 강화학습에서 샘플 효율성을 높이기 위해, 작업별 VAE와 작업 임베딩 기반 라우팅을 결합한 혼합 전문가 트랜스포머 구조를 제안한다. Atari‑100k와 Meta‑World 벤치마크에서 파라미터 절반 수준으로 기존 최첨단 모델과 동등하거나 우수한 성능을 달성한다.

상세 분석

본 논문은 시각적 멀티태스크 강화학습(MTRL)에서 기존 단일 세계 모델이 겪는 “시각·동역학 이질성” 문제를 해결하기 위해 세 가지 핵심 설계를 제시한다. 첫째, 작업마다 별도 VAE(카테고리형)를 배치해 고차원 픽셀 입력을 32×32개의 토큰으로 압축한다. 각 VAE는 학습 초기에 gradient‑based clustering을 통해 작업군을 자동으로 할당받으며, 학습 중에는 고정된 task‑specific encoder‑decoder 쌍을 사용해 재구성 손실을 최소화한다. 둘째, 동적 모델링은 “Mixture‑of‑Experts Transformer”로 구현한다. 작업 임베딩 eₖ를 입력으로 하는 라우터가 Top‑K softmax을 통해 nₖ개의 전문가 트랜스포머를 선택하고, 선택된 전문가들의 출력은 공유 트랜스포머에 연결돼 공통 지식을 학습한다. 이때 라우팅은 토큰‑레벨이 아닌 task‑level로 수행돼, 동일 작업 내에서 일관된 전문가가 활성화되어 시간적 연속성을 보존한다. 셋째, 전문가 균형을 위해 온도 스케줄링과 expert‑balance loss를 도입하고, task‑prediction loss를 추가해 은닉 상태 hₜᵏ가 작업을 구분하도록 유도한다. 손실 함수는 재구성, 보상·종료 예측, KL 기반 동역학 예측, 그리고 위의 보조 손실을 가중치 β₁=0.5, β₂=0.1로 결합한다.

학습 과정은 두 단계로 나뉜다. Warm‑up 단계에서는 각 작업의 gradient vector를 클러스터링해 VAE와 예측기(보상·연속·critic)를 공유할 그룹을 결정한다. 이후 전체 파라미터를 end‑to‑end로 최적화한다. 이렇게 하면 파라미터 수를 크게 늘리지 않으면서도 작업별 특화된 표현과 공통 역학을 동시에 학습할 수 있다.

실험 결과는 두 가지 주요 벤치마크에서 확인된다. Atari‑100k(26게임)에서는 단일 MoW 에이전트가 평균 인간 정규화 점수 110.4%를 기록했으며, 이는 26개의 task‑specific 모델을 앙상블한 STORM(114.2%)과 비슷하면서도 파라미터는 절반 수준이다. Meta‑World에서는 300k 스텝 내에 74.5%의 성공률을 달성해 기존 최고 기록을 넘어섰다. Ablation study에서는 VAE 모듈, task‑level 라우팅, 그리고 gradient‑based clustering 각각이 성능에 미치는 영향을 정량화해, 모든 구성 요소가 상호 보완적으로 작용함을 입증한다.

이러한 설계는 “전문가 과잉 활용”이나 “전문가 붕괴” 문제를 완화하고, 시각적 입력이 복잡한 실제 로봇 혹은 게임 환경에서도 확장 가능한 모델 기반 RL을 구현할 수 있음을 시사한다. 또한, 작업 임베딩 기반 라우팅은 향후 새로운 작업이 추가될 때 기존 전문가를 재활용하거나 최소한의 파라미터만으로 빠르게 적응할 수 있는 기반을 제공한다.