다음 임베딩 예측을 통한 세계 모델 강화

기존의 세계 모델은 화면 전체를 다시 그리는 복잡한 픽셀 복원 방식에 의존했으나, 이 논문은 다음 장면의 핵심 특징인 임베딩만 예측해도 충분함을 입증했다. 이를 통해 불필요한 시각적 디테일에 낭비되는 연산량을 줄이고, 복잡한 미로 찾기처럼 장기 기억이 필수적인 환경에서 에이전트의 지능을 획기적으로 높였다.

강화학습 에이전트가 똑똑해지려면 미래를 상상할 수 있는 World Model이 필요하다. 기존 Dreamer 계열 모델들은 미래의 화면을 픽셀 단위로 다시 그리는 픽셀 복원(Pixel Reconstruction)을 통해 세상을 배웠다. 하지만 화면 속의 복잡한 배경이나 텍스처는 실제 보상을 얻는 것과는 무관한 경우가 많아, 모델의 학습 능력이 낭비되는 한계가 있었다.

NE-Dreamer는 화면을 그리는 대신, 인코더가 추출한 다음 장면의 핵심 요약본인 임베딩(Embedding)을 직접 맞추는 방식을 택했다. 이때 Causal Transformer를 사용하여 지금까지 에이전트가 겪은 모든 과거 이력을 참고하게 한다. 이는 에이전트가 현재 눈앞에 보이지 않는 정보(예: 아까 지나온 복도의 구조)를 머릿속에 유지하며 미래를 예측하도록 돕는다.

예측된 값과 실제 다음 장면의 임베딩을 비교할 때는 Barlow Twins라는 기법을 사용한다. 이는 정보의 중복을 제거하고 핵심적인 특징들만 남도록 유도하여, 모델이 장기적인 시계열 흐름 속에서 정말 중요한 정보가 무엇인지 스스로 깨닫게 만든다. 결과적으로 픽셀 하나하나에 집착하지 않고도 세상의 논리적 흐름을 더 정확히 파악하게 된다.

NE-Dreamer는 RSSM(Recurrent State-Space Model) 아키텍처를 기반으로 하며, 기존의 픽셀 디코더를 제거하고 Causal Temporal Transformer $T_\theta$를 도입했다. Transformer는 과거의 결정론적 상태 $h_{\le t}$, 확률적 잠재 상태 $z_{\le t}$, 그리고 취해진 행동 $a_{\le t}$ 시퀀스를 입력으로 받아 다음 시점의 예측 임베딩 $\hat{e}_{t+1}$을 출력한다.

학습 시에는 실제 다음 시점의 이미지 $x_{t+1}$을 인코더 $f_{enc}$에 통과시켜 얻은 $e^*_{t+1}$을 정답으로 사용한다. 이때 정답 임베딩에는 Stop-gradient 연산을 적용하여 예측 모델인 Transformer와 RSSM의 가중치만 업데이트되도록 제한한다. [예측 임베딩 $\hat{e}_{t+1}$과 실제 임베딩 $e^*_{t+1}$을 입력으로] → [Barlow Twins 손실 함수를 통해 교차 상관 행렬의 대각 성분은 1로, 비대각 성분은 0으로 만드는 연산을 수행해] → [최종적인 $L_{NE}$ 값을 얻고] → [이 값이 작아질수록 모델은 미래의 핵심 특징을 중복 없이 정확하게 예측하게 된다].

전체 손실 함수 $L_{wm}$은 보상 예측, 종료 여부 예측, KL Divergence 정규화 항에 제안된 $L_{NE}$를 더해 구성된다. 이를 통해 모델은 픽셀 복원 없이도 보상과 직결된 환경의 동역학을 잠재 공간에서 완벽하게 학습한다.

DMLab Rooms 벤치마크의 4가지 핵심 작업(Collect, Exploit, Select, Watermaze)에서 NE-Dreamer는 DreamerV3 및 R2-Dreamer 등 기존 모델들을 압도했다. 특히 장기적인 상태 유지가 필수적인 Rooms Exploit 작업에서는 다른 모델들이 10 미만의 점수를 기록할 때 NE-Dreamer는 30 이상의 높은 점수를 달성하며 독보적인 성능을 보였다.

Ablation Study 결과, 성능 향상의 핵심은 Transformer 구조와 다음 단계 예측(Next-step shift) 목표의 결합임이 확인됐다. Transformer를 제거하고 단순한 구조를 사용하거나, 다음 단계가 아닌 현재 단계의 임베딩을 맞추도록 학습할 경우 성능이 급격히 하락하여 기존 모델 수준으로 회귀했다.

DeepMind Control Suite(DMC)의 20가지 연속 제어 작업에서도 NE-Dreamer는 픽셀 복원 없이도 DreamerV3와 대등하거나 소폭 우수한 성능을 기록했다. 이는 제안된 방식이 장기 기억이 필요한 복잡한 환경뿐만 아니라 일반적인 제어 문제에서도 범용적으로 작동함을 증명한다.

NE-Dreamer 아키텍처는 DreamerV3의 RSSM 구조를 계승하되, 픽셀 디코더를 2개의 레이어와 4개의 헤드를 가진 경량 Causal Transformer로 교체했다. Transformer의 히든 차원은 256으로 설정되어 전체 파라미터 수 12M 내외의 효율적인 구조를 유지한다.

핵심 메커니즘인 Temporal Predictive Alignment는 잠재 공간에서의 시간적 일관성을 강제한다. 이는 단순히 현재 시점의 인코더와 잠재 상태를 일치시키는 기존의 Decoder-free 방식들과 달리, 과거의 궤적을 모두 고려하여 미래의 표현을 예측하게 함으로써 Partial Observability에 대한 강건성을 확보한다.

Barlow Twins 손실 함수는 차원별 정규화를 거친 후 $C_{ij} = \frac{1}{N} \sum \tilde{\hat{e}}_{i} \tilde{e}^*_{j}$ 연산을 통해 교차 상관 행렬을 계산한다. [행렬의 대각 성분 $C_{ii}$를 1에 가깝게 만들어] → [예측값과 실제값 사이의 불변성을 확보하고] → [비대각 성분 $C_{ij}$를 0에 가깝게 만들어] → [각 차원이 서로 독립적인 정보를 담게 하여 표현 붕괴를 방지한다].

구현 측면에서 $\lambda_{BT}$ 하이퍼파라미터는 $5 \times 10^{-4}$를 사용하여 중복 제거 항의 세기를 조절했다. 또한 정답 임베딩에 Stop-gradient를 적용함으로써 인코더가 Transformer의 예측에 맞춰 쉬운 정답을 내놓는 편법(Trivial Solution)을 방지하고, 오직 Transformer가 인코더의 출력을 따라가도록 유도했다.

본 연구는 주로 장기적인 구조 파악이 중요한 환경(DMLab)에 집중되어 있다. 따라서 아주 미세한 시각적 디테일(예: 작은 텍스처의 변화)이 보상과 직결되는 고해상도 제어 작업에서도 픽셀 복원 없이 동일한 수준의 정밀도를 확보할 수 있을지에 대해서는 추가적인 검증이 필요하다.