Accurate and Efficient World Modeling with Masked Latent Transformers

기술 요약: 마스크드 잠재 트랜스포머(Masked Latent Transformers)를 통한 정확하고 효율적인 월드 모델링

문제 정의
모델 기반 강화 학습(RL) 알고리즘, 특히 Dreamer 시리즈는 잠재 월드 모델 내에서 시뮬레이션된 궤적을 통해 에이전트를 학습시킴으로써 상당한 성공을 거두었습니다. 그러나 표준 방식(예: DreamerV3)의 압축된 잠재 공간은 결정적인 한계점을 가집니다. 즉, 잠재 공간의 압축된 특성으로 인해 중요한 환경적 세부 사항이 손실되어 Crafter와 같은 복잡한 환경에서 에이전트의 성능이 저하됩니다. 반대로, $\Delta$-IRIS 및 DIAMOND와 같이 높은 정확도를 목표로 하는 최근 방법들은 에로 재구성된 픽셀 이미지로부터 에이전트를 직접 학습시켜야 합니다. 이러한 방식은 훈련 효율성을 떨어뜨리고, 에이전트가 월드 모델로부터 학습된 풍부한 내부 표현(예: 장기 기억)을 활용하는 것을 방해합니다. 과제는 정확하게 미세한 세부 사항을 포착하면서도 훈련과 추론 측면에서 효율적인 월드 모델링 접근 방식을 개발하는 것입니다.

방법론: EMERALD
저자들은 정확도와 효율성 사이의 절충안을 해결하기 위해 설계된 트랜스포머 기반 RL 알고리즘인 EMERALD(Efficient MaskEd latent tRAnsformer worLD model)를 제안합니다. 이 아키텍처는 세 가지 주요 구성 요소로 이루어져 있습니다:

1. MaskGIT를 이용한 공간 잠재 상태(Spatial Latent State): 이전 방식들이 관측치를 단일 벡터로 압축하는 것과 달리, EMERALD는 공간 잠재 상태($z_t$)를 사용합니다. 이 상태는 평면 벡터가 아닌 토큰의 그리드 형태로 표현됩니다. 이 토큰들을 효율적으로 디코딩하기 위해, 모델은 MaskGIT(Masked Generative Image Transformer)를 채택합니다. 순차적인 자기회귀(autoregressive) 디코딩은 속도가 느리기 때문에, MaskGIT는 스케줄링된 병렬 디코딩을 사용합니다. 훈련 중에는 코사인 스케줄에 따라 토큰이 마스킹되며, 모델은 이를 병렬적으로 예측합니다. 상상(추론) 단계에서 모델은 일관성을 보장하고 환각 현상을 방지하기 위해 여러 단계(예: $S=3$)에 걸쳐 예측을 정교화합니다.
2. 시계열 트랜스포머 네트워크(Temporal Transformer Network): 모델은 장기 의존성을 처리하기 위해 시계열 마스크드 트랜스포머(TSSM)를 통합합니다. 이는 마스크드 셀프 어텐션을 사용하여 공간 잠재 상태를 처리하고 시계열 은닉 상태($h_t$)를 생성합니다. 이는 DreamerV3에서 사용된 GRU 기반의 순환 상태 공간 모델(RSSM)을 대체하며, 순환 처리의 병목 현상 없이 과거의 맥락을 더 잘 모델링할 수 있게 해줍니다.
3. 잠재 공간 액터-크리틱 학습(Latent Space Actor-Critic Learning): 에이전트(액터 및 크리틱 네트워크)는 월드 모델의 내부 표현을 사용하여 전적으로 잠재 공간 내에서 훈련됩니다. 이를 통해 에이전트는 픽셀 기반 접근 방식처럼 재구성된 이미지로부터 표현을 다시 학습할 필요 없이, 월드 모델이 학습한 메모리와 공간적 특징의 이점을 누릴 수 있습니다.

훈련 목적 함수는 다음을 포함하는 복합 손실 함수를 최소화합니다:

• 마스킹 손실 ($L_{mask}$): 타겟 잠재 상태와 예측된 마스크드 토큰 사이의 KL 발산.
• 역학 손실 ($L_{dyn}$): 인코더 표현이 은닉 상태로부터 예측 가능하도록 보장하는 정규화.
• 재구성 손실 ($L_{rec}$): 입력 관측치의 픽셀 수준 재구성.
• 보상 및 지속 손실 ($L_{rew}, L_{con}$): 환경 보상 및 에피소드 지속 플래그 예측.

주요 기여

• 새로운 아키텍처: 공간 잠재 상태와 MaskGIT 예측을 결합하여, 자기회귀 디코딩의 계산 비용 없이 고충실도 궤적 생성을 가능하게 하는 월드 모델을 도입했습니다.
• 효율성 및 정확도: 이 방법은 낮은 픽셀 오차와 다이아몬드나 해골 화살 같은 세부 사항의 보존을 통해 거의 완벽한 이미지 재구성을 달면서도, 에이전트를 잠재 공간에 유지함으로써 높은 훈련 효율성을 유지합니다.
• 최첨단 성능(SOTA): 이 접근 방식은 Crafter 환경에서 새로운 벤치마크를 설정하며, 이전의 모델 기반 방법들과 인간 전문가의 성능을 능가했습니다.

결과
1,000만 환경 스텝의 예산으로 Crafter 벤치마크에서 평가한 결과:

• 달성 점수(Achievement Score): EMERALD는 **58.1%**를 기록하여 이전 최고 기록(DreamerV3 XL, 39.6%)을 넘어섰으며, 인간 전문가의 성능(50.5%)을 초과했습니다.
• 달성 범위(Achievement Coverage): 에이전트는 평가 중 최소 한 번씩 22개의 모든 업적을 달성하는 데 성공했으며, 이는 강력한 장기 계획 및 인지 능력을 입증합니다.
• 훈련 효율성: EMERALD는 RTX 3090에서 27 FPS를 기록했는데, 이는 $\Delta$-IRIS(12 FPS) 및 DreamerV3 XL(23 FPS)보다 훨씬 빠르면서도 이들보다 높은 점수를 기록했습니다.
• 어블레이션 연구(Ablation Studies): 실험을 통해 공간 잠재 상태와 트랜스포머 아키텍처의 결합이 매우 중요하다는 것을 확인했습니다. 공간 잠재 상태를 제거하거나 단일 디코딩 단계를 사용할 경우 성능과 예측 정확도가 크게 감소했습니다.

의의 및 주장
본 논문은 EMERALD가 월드 모델의 정확도와 훈련 효율성 사이의 긴장을 해결함으로써 모델 기반 RL의 중요한 진전을 나타낸다고 주장합니다. 공간 잠재 공간에서 병렬 디코딩을 위해 MaskGIT를 활용함으로써, 에이전트는 기존의 압축된 벡터 표현에서 손실되었던 결정적인 환경 세부 사항을 인지할 수 있게 되었습니다. 저자들은 이것이 Crafter 벤치마크에서 1,000만 스텝 이내에 인간 전문가의 성능을 초과한 첫 번째 방법임을 강조합니다. 나아가, 본 연구는 마스크드 어텐션을 갖춘 트랜스포머 기반 아키텍처가 복잡하고 메모리 집약적인 환경에서 장기 의존성과 공간적 관계를 효과적으로 모델링할 수 있음을 보여주며, 순환 모델에 대한 확장 가능한 대안을 제시합니다. 저자들은 이 작업을 마인크래프트와 같은 복잡한 도메인에서 더 유능한 자율 에이전트로 나아가기 위한 단계로 규정합니다.