World model : 논문 리뷰 : Next Embedding Prediction Makes World Models Stronger

용어 설명

• MBRL (Model-Based Reinforcement Learning): 환경의 동작 방식(dynamics)을 모방하는 world model을 먼저 학습하고, 이 모델 안에서 가상으로 미래를 상상(imagination)하여 policy를 훈련하는 강화학습 기법.
• World model: agent가 경험을 바탕으로 환경의 상태 전이(state transition)와 보상을 예측하도록 학습된 내부 표현 모델.
• Decoder-free: 픽셀 단위의 이미지를 다시 그려내는(reconstruction) pixel decoder를 사용하지 않고, 잠재 공간(latent space)에서 직접 모델을 학습시키는 구조.
• Next-embedding prediction: 현재 상태가 아닌 '다음 시점'의 observation이 인코딩된 embedding 값을 예측하는 학습 목표.
• Temporal predictive alignment: 시간에 따른 흐름을 반영하여, 과거 기록을 바탕으로 예측한 미래의 상태와 실제 미래 상태의 representation이 일치하도록 맞추는 과정.
• Barlow Twins: 중복성 감소(redundancy-reduction)를 위한 self-supervised loss. Negative samples 없이도 모델의 출력값이 하나로 통일되어 버리는 현상(representation collapse)을 방지함.
• Partial observability: agent가 현재 시점의 관측치(observation) 하나만으로는 환경의 전체 상태를 파악할 수 없는 조건 (예: 1인칭 3D 맵 탐색).
• Stop-gradient target: 모델 학습 시 정답지(target) 역할을 하는 변수에 역전파(gradient)가 흐르지 않도록 차단하여 학습 붕괴를 막고 안정성을 높이는 기법.

---

Purpose of the Paper

• 기존의 한계: Dreamer와 같은 대표적인 world model들은 픽셀 수준의 reconstruction을 사용하여 계산 비용이 높고, 보상과 무관한 시각적 디테일에 모델 용량을 낭비하는 문제가 있음. 이를 해결하기 위해 등장한 decoder-free 모델들은 대부분 동일 시점(same-timestep)의 feature 일치에만 초점을 맞춰, partial observability 환경에서 필수적인 '시간적 맥락(long-horizon temporal structure)'을 제대로 포착하지 못하고 메모리 기반 task에서 실패함.
• 새로운 접근 방식: 본 논문은 픽셀 reconstruction을 완전히 제거하고, 대신 causal temporal transformer를 이용해 과거 history로부터 다음 시점의 상태를 예측(next-embedding prediction) 하도록 모델을 강제하는 새로운 방식의 decoder-free MBRL agent인 NE-Dreamer를 제안함. 이를 통해 동일 시점 매칭의 한계를 극복하고 시간에 따른 예측력(temporal predictiveness)을 모델에 직접적으로 부여하고자 함.

---

Key Contributions

• Next-embedding prediction objective 도입:
  • Contribution: Same-step 일치 방식이나 pixel decoder를 완전히 대체하고, 다음 시점의 encoder embedding을 예측하여 stop-gradient target과 맞추는 새로운 목적 함수를 설계함.
  • Novelty: Representation 학습을 단순한 이미지 특징 추출이 아닌, '인과적 미래 예측(causal next-step prediction)' 문제로 전환하여 temporal predictiveness를 명시적으로 최적화함.
• Causal temporal transformer 기반의 MBRL 파이프라인 구축:
  • Contribution: 기존 RSSM 기반 Dreamer 구조에 가벼운 causal temporal transformer를 결합하여 과거의 history를 효과적으로 압축함.
  • Novelty: 단순한 feedforward 네트워크가 아닌 시퀀스 모델링을 활용함으로써, 부분 관측 환경에서도 강건한 멀티스텝 예측(latent overshooting)을 가능하게 함.
• Data augmentation 없는 안정적인 representation 학습:
  • Contribution: Barlow Twins loss를 활용하여 예측된 미래 embedding과 실제 미래 embedding 간의 invariance를 높이고 redundancy를 줄임.
  • Novelty: 강력한 이미지 augmentation이나 negative samples 없이도 시간축(temporal) 기준의 alignment를 통해 representation collapse를 성공적으로 방지함.

---

Experimental Highlights

• 핵심 결과 (DMLab Rooms SOTA 달성): Long-horizon memory와 navigation 능력이 요구되는 부분 관측 환경인 DMLab Rooms (4개 task)에서 50M environment steps 기준, 강력한 baseline인 DreamerV3(decoder 기반) 및 R2-Dreamer, DreamerPro(decoder-free)를 압도적인 격차로 능가함.
• 표준 제어 성능 유지 (DMC benchmark): 20개의 DMC 연속 제어 task (1M steps)에서 기존 SOTA 모델들과 동등하거나 약간 우수한 성능을 달성함. 이는 reconstruction을 제거해도 standard continuous control 성능이 전혀 저하되지 않음을 증명함.
• 가설 검증 (Ablation Study): Causal transformer를 제거하거나, next-step 예측을 current-step 매칭(no next-step shift)으로 변경했을 때 DMLab 성능이 완전히 붕괴됨. 이 실험은 성능 향상의 핵심 원인이 부가적인 트릭(projector 등)이 아닌 predictive sequence modeling 자체에 있음을 명확하게 증명함.
• Representation 진단 (Post-hoc decoder): 학습 과정에서 배제된 pixel decoder를 얼어있는(frozen) latent state에 달아 이미지를 복원한 결과, 기존 모델들과 달리 NE-Dreamer는 시간이 지나도 환경의 핵심 객체와 공간 레이아웃(temporal consistency)을 일관되게 유지함을 시각적으로 확인함.

---

Limitations and Future Work

• Limitations: 논문의 실험이 주로 세밀한 시각적 묘사(fine visual detail)보다는 장기적인 구조 파악(long-term structure)이 중요한 도메인에 집중되어 있음. 따라서 텍스처나 복잡한 배경 등 시각적 복잡도가 극도로 높은(visually complex) 환경에서도 decoder-free 방식이 pixel reconstruction 기반 모델만큼 고해상도의 제어 정보를 유지할 수 있을지는 아직 입증되지 않음.
• Future Work:
  • 시각적으로 복잡한 고해상도 도메인으로 NE-Dreamer의 평가 환경을 확장하여 한계점을 테스트.
  • Barlow Twins 외에 모델의 표현력과 안정성을 극대화할 수 있는 새로운 형태의 alternative alignment loss들을 탐구하여, 다양한 환경에서의 범용성을 높이는 연구 진행.

---

Overall Summary
이 논문은 무거운 pixel reconstruction이나 의존적인 data augmentation 없이, causal temporal transformer를 활용한 next-embedding prediction만으로 강력한 world model을 구축하는 NE-Dreamer를 제안합니다. 이 방법론은 모델이 현재 상태에만 머물지 않고 과거의 history를 기반으로 미래의 핵심 상태를 예측(temporal predictive alignment)하도록 강제함으로써, 부분 관측 및 장기 기억이 필수적인 환경에서 압도적인 SOTA 성능을 달성했습니다. 이는 향후 decoder-free MBRL 연구에 있어 '시간적 예측성'이 효율적이고 강력한 representation 학습의 새로운 표준이 될 수 있음을 시사하는 중요한 의의를 가집니다.

---

쉬운 설명
기존의 인공지능 요원(agent)들이 방을 탐색할 때 "지금 눈앞에 보이는 장면의 픽셀들을 스케치북에 똑같이 따라 그리는 연습(reconstruction)"을 하거나, "방금 찍은 사진에 필터를 씌워 원본과 같은 사진인지 맞추는 연습(same-timestep alignment)"에 집중했습니다. 하지만 이 방식은 방을 한 바퀴 돌고 났을 때 내 뒤에 무엇이 있었는지 까먹기 쉽습니다.

반면 이 논문의 모델은 "지금까지 걸어온 길을 쭉 떠올려보고, 딱 1초 뒤에 눈앞에 나타날 중요한 장면이 무엇일지 미리 맞추는 연습(next-embedding prediction)" 을 합니다. '다음 순간'을 지속적으로 예측하도록 훈련받았기 때문에, 당장 눈에 보이지 않더라도 "내 뒤에는 사과가 있다"는 중요한 사실을 잊지 않고 머릿속에 일관되게 기억할 수 있게 되어, 복잡한 미로 찾기나 기억력 테스트에서 훨씬 똑똑하게 행동할 수 있게 된 것입니다.

 

 

Abstract

partially observable하고 high-dimensional한 domains에서 temporal dependencies를 포착하는 것은 model-based reinforcement learning (MBRL)에 있어 매우 중요합니다.

우리는 representation space에서 temporal predictive alignment를 직접 optimizing하면서, latent state sequences로부터 next-step encoder embeddings를 predict하기 위해 temporal transformer를 활용하는 decoder-free MBRL agent인 NE-Dreamer를 소개합니다.

이러한 접근 방식은 NE-Dreamer가 reconstruction losses나 auxiliary supervision 없이도 coherent하고 predictive한 state representations를 학습할 수 있게 해줍니다.

DeepMind Control Suite에서 NE-Dreamer는 DreamerV3 및 선도적인 decoder-free agents의 성능과 맞먹거나 이를 능가합니다.

또한 memory 및 spatial reasoning을 수반하는 까다로운 DMLab tasks의 하위 집합에서 NE-Dreamer는 상당한 성능 향상을 달성합니다.

이러한 결과는 temporal transformers를 활용한 next-embedding prediction이 복잡하고 partially observable한 environments에서 MBRL을 위한 효과적이고 scalable한 framework임을 입증합니다.

 

Figure 1. DMLab Benchmark Summary 설명

전문적인 요약 노트: 이 그래프는 memory 및 navigation 역량이 강하게 요구되는 partially observable 환경(DMLab Rooms)에서 NE-Dreamer와 기존 SOTA 모델들의 성능(Return)을 비교한 결과입니다.

• 핵심 포인트: 매칭된 compute와 model capacity(50M environment steps, 12M parameters)라는 동일 조건 하에서, 파란색 선인 NE-Dreamer만이 성공적으로 유의미한 Return 상승을 만들어냅니다.
• 기존 모델의 한계 증명: 반면 기존의 decoder-based 모델인 Dreamer(빨간선)나 decoder-free baseline인 R2-Dreamer(초록선), DreamerPro(보라선)는 완전히 학습에 실패하여 바닥에 머무르는 것을 볼 수 있습니다.
• 결론: 이는 단순히 reconstruction을 제거하는 것만으로는 부족하며, temporal dependencies를 명시적으로 학습하는 NE-Dreamer의 방식이 복잡한 환경의 long-horizon 문제를 해결하는 데 필수적임을 시각적으로 증명하는 강력한 결과입니다.

쉬운 설명: 이 그래프는 복잡한 미로 게임에서 인공지능들이 얼마나 점수를 잘 내는지 보여주는 성적표입니다. 조건은 모두 공평하게 '똑같은 크기의 뇌(파라미터 수)'와 '똑같은 연습 시간(환경 스텝)'을 주었습니다. 결과를 보면 파란색 선인 NE-Dreamer 혼자서만 점수가 쑥쑥 올라가고, 나머지 유명한 AI들은 아예 길을 찾지 못하고 0점대에 머물고 있습니다. NE-Dreamer가 다음 상황을 미리 예측하는 훈련 덕분에, 복잡한 공간의 지도를 머릿속에 성공적으로 그렸다는 뜻입니다.

 

 

Figure 2. Method overview 설명

전문적인 요약 노트:

Figure 2는 NE-Dreamer의 전반적인 아키텍처 흐름도를 보여줍니다. 기존 Dreamer 모델과 무엇이 다르고 무엇이 같은지를 명확히 보여주는 설계도입니다.

• 유지된 구조: $h_t$ (deterministic state)와 $z_t$ (stochastic state)로 이루어진 기존 Dreamer의 핵심인 RSSM dynamics 구조는 그대로 유지하고 있습니다.
• 핵심 변화 (Decoder의 부재와 Transformer의 도입): 아래쪽으로 향해야 할 pixel reconstruction 과정이 완전히 사라졌습니다. 대신 위쪽에 빨간색 블록인 causal temporal transformer가 추가되었습니다.
• Next-Embedding Prediction 메커니즘: 모델은 현재까지의 상태($h_{\le t}, z_{\le t}$)와 행동($a_t$)의 sequence를 Transformer에 통과시켜, 다음 시점($t+1$)의 encoder 출력값인 예측 임베딩 $ \hat{e}{t+1} $을 만들어냅니다. 그리고 이 예측값 $ \hat{e}{t+1} $과 실제로 다음 스텝의 관측치($x_{t+1}$)가 Encoder를 통과해 나온 실제 임베딩 간의 차이를 $L_{NE}$ (Next-Embedding loss)를 통해 최소화하도록 학습합니다.

쉬운 설명:

이 그림은 NE-Dreamer가 어떻게 세상을 이해하는지 보여주는 머릿속 작동 원리입니다. 왼쪽에서 오른쪽으로 시간이 흘러가면서 AI가 행동($a_t$)을 결정합니다. 기존 AI들은 자신이 본 장면($x_t$)을 그대로 다시 그려보는 연습을 했지만, 이 모델의 그림에는 그리는 과정이 없습니다.

대신 제일 위에 있는 빨간색 상자 'Transformer'가 핵심 역할을 합니다. 이 Transformer는 "지금까지 내가 겪은 상황들을 종합해 볼 때, 다음 순간($t+1$)에는 대충 이런 특징을 가진 풍경이 눈앞에 펼쳐지겠지?" 하고 미리 예측($\hat{e}_{t+1}$)해 보는 기능입니다. 그리고 실제로 다음 순간에 본 풍경과 자신이 예측했던 풍경의 특징이 얼마나 비슷한지 비교($L_{NE}$)하면서 스스로 똑똑해지는 과정을 나타냅니다.

 

 

1. Introduction

highdimensional observations로부터의 Model-based reinforcement learning (MBRL)은 long-horizon prediction과 control을 지원하는 compact한 latent state를 학습하는 것에 달려있습니다. 이러한 요구사항은 partial observability 하에서 더욱 중요해집니다: agent는 단일 frame에 반응하기보다는 시간에 걸쳐 정보를 통합해야 합니다.

지배적인 접근 방식은 Dreamer에서와 같이 pixel decoder를 사용하여 world model을 학습하며, 여기서 reconstruction은 풍부하고 control-effective한 features를 생성합니다. 그에 따른 비용은 modeling 부담입니다: reconstruction은 무거운 generative objective를 도입하고, optimization을 복잡하게 만들며, visually 디테일하지만 task-irrelevant한 측면에 capacity를 할당하게 만들 수 있습니다. Decoder-free methods는 pixel decoder를 제거하고 representations를 직접 training하여 파이프라인을 단순화하고 효율성을 향상시킵니다.

그러나 많은 decoder-free objectives는 주로 instantaneous (same-timestep) agreement를 강제합니다. partial observability 하에서 instantaneous agreement는 충분하지 않습니다: representation은 시간에 걸쳐 predictive해야 합니다. 명시적인 temporal constraint가 없다면, training은 drift되거나 collapse될 수 있으며, 이는 memory- 및 navigation-heavy tasks에서 나타나는 실패 모드인 약한 long-horizon structure로 이어집니다.

본 논문에서는 latent representations에서 temporal predictive alignment를 직접 optimizing하여 학습하는 decoder-free world model인 NE-Dreamer를 소개합니다. NE-Dreamer는 pixel-level reconstruction을 간단하면서도 강력한 objective로 대체합니다: 각 timestep에서 temporal transformer는 sequence의 next encoder embedding을 predict하며, 이 prediction은 redundancy-reduction metric(구체적으로 본 구현에서는 Barlow Twins)을 사용하여 실제 next-step embedding과 align됩니다. same-timestep matching에서 next-step prediction으로 초점을 이동함으로써, NE-Dreamer는 pixel reconstruction, data augmentation 또는 auxiliary regularization 없이도 temporally coherent한 latent states를 학습합니다. Figure 1에 설명된 바와 같이, 이 design은 NE-Dreamer가 동일한 model size의 이전 methods와 비교하여 partially observable한 DMLab environments에서 상당히 더 높은 성능을 달성할 수 있게 해줍니다.

우리의 주요 contributions는 다음과 같습니다:

1. 우리는 학습된 representation에서 temporal predictiveness를 명시적으로 강제하는 next-embedding prediction 기반의 decoder-free world-model objective를 제안합니다.
2. 우리는 표준 RSSM training 내에서 history로부터 next-step prediction을 구현하기 위해 Dreamer-style MBRL pipeline에 lightweight causal temporal transformer를 통합합니다.
3. 우리는 DeepMind Control Suite 및 DeepMind Lab에서 NE-Dreamer를 평가하여, 매칭된 compute 및 model size 하에서 DMC에서는 강력한 성능을 보여주고 memory/navigation-heavy DMLab Rooms에서는 상당한 성능 향상을 보여줍니다.
4. 타겟이 지정된 ablations 및 representation diagnostics를 통해, 우리는 성능 향상의 원인을 reconstruction이나 auxiliary tricks가 아닌 predictive sequence modeling (causal transformer + next-step target shift)으로 분리해냅니다.

 

 

 

 

 

NE-Dreamer: Introduction 정리 노트

1. 연구 배경 및 기존 방법론의 한계 (Problem Definition)

• 기존 MBRL의 한계 (Reconstruction 기반): Dreamer와 같은 지배적인 MBRL 모델들은 pixel decoder를 통한 reconstruction에 의존함. 이는 무거운 generative objective를 유발하고, task-irrelevant한 시각적 디테일까지 모델링하게 만들어 비효율적임.
• 기존 Decoder-free 방식의 한계: Pixel decoder를 없애 효율성을 높였으나, 대부분 동일 시간대(same-timestep)의 instantaneous agreement에만 집중함.
• Partial Observability에서의 치명적 문제: 명시적인 temporal constraint가 없으면, 모델이 시간에 따른 변화를 예측하지 못해 학습이 drift되거나 collapse됨. 결과적으로 memory 및 navigation-heavy task에서 요구되는 long-horizon structure 학습에 실패함.

2. 핵심 제안: NE-Dreamer (Proposed Architecture)

• Next-Step Prediction으로의 패러다임 전환: Pixel-level reconstruction을 폐기하고, latent representations 내에서 temporal predictive alignment를 직접 optimizing하는 새로운 decoder-free world model 제안.
• 작동 메커니즘: Temporal transformer를 도입하여 과거 history로부터 next encoder embedding을 predict함. 이 예측값과 실제 next-step embedding을 redundancy-reduction metric(구체적으로 Barlow Twins 사용)을 통해 align함.
• 효율성 및 강건성: Same-timestep matching을 next-step prediction으로 대체함으로써, 복잡한 data augmentation이나 auxiliary regularization 없이도 temporally coherent한 latent states를 성공적으로 학습함.

3. 주요 기여 요약 (Key Contributions)

• Learned representation에 temporal predictiveness를 명시적으로 강제하는 next-embedding prediction 기반 decoder-free objective 제안.
• 기존 Dreamer 스타일의 MBRL (표준 RSSM training) 파이프라인에 lightweight causal temporal transformer를 성공적으로 통합.
• 동일한 compute 및 model size 조건에서 DMC 환경은 물론, 특히 memory/navigation 중심의 부분 관찰 환경인 DMLab Rooms에서 기존 방식 대비 압도적인 성능 향상 증명.
• Ablation study를 통해 모델의 성능 향상이 단순한 트릭(auxiliary tricks)이 아닌, 명확한 predictive sequence modeling(causal transformer + next-step target shift)에서 비롯되었음을 검증.

---

쉬운 설명 :

기존의 똑똑한 AI 에이전트들(예: Dreamer)은 주변 환경을 이해하기 위해 자신이 본 모든 장면을 픽셀 단위로 일일이 다시 그려보는(reconstruction) 방식을 썼습니다. 하지만 이는 너무 많은 에너지를 소모하고, 벽지 색깔 같은 쓸데없는 정보까지 외우게 만드는 단점이 있었죠.

그래서 그림 그리는 과정을 과감히 생략한 모델들(Decoder-free)이 등장했지만, 이들은 단순히 "지금 내가 본 것"과 "지금 내가 생각하는 것"이 일치하는지만 확인했습니다. 이렇게 하면 당장 눈앞의 일은 잘 처리해도, 복잡한 미로 찾기처럼 과거의 기억을 바탕으로 미래를 예측해야 하는 상황에서는 길을 잃고 헤매기 일쑤였습니다.

NE-Dreamer는 이 문제를 해결하기 위해 **"다음 순간에 무엇이 보일지(next-embedding) 그 핵심 특징만 미리 예측해보자"**라는 아이디어를 도입했습니다. AI가 과거의 기억(history)을 바탕으로 바로 다음 스텝의 상태를 예측하고, 실제 다음 상태와 자신의 예측을 맞춰보며 학습하게 한 것입니다. 그림을 그리지 않으니 가볍고 빠르며, 계속해서 다음 상황을 예측하도록 훈련받았기 때문에 기억력과 공간 지각력이 필요한 복잡한 환경에서 훨씬 더 길을 잘 찾고 똑똑하게 행동할 수 있게 된 것입니다.

 

 

 

 

2. Related Work

World models for pixel control. Latent world models는 high-dimensional observations로부터 long-horizon prediction 및 decision-making을 지원하는 compact states를 학습하는 것을 목표로 합니다. 초기 연구는 latent space에서 dynamics를 학습하는 것이 pixels로부터 "in imagination"으로 행동함으로써 planning 및 control을 가능하게 할 수 있음을 보여주었습니다.

PlaNet은 images로부터의 planning을 위한 실용적인 latent dynamics backbone으로 recurrent state-space model (RSSM)을 도입했습니다. RSSMs를 기반으로, Dreamer family는 latent imagination을 통해 latent space의 imagined rollouts에서 actor-critic을 train합니다. NE-Dreamer는 이러한 RSSM-based control backbone을 유지하고 latent representation이 학습되는 방식을 변경합니다.

Reconstruction-based world models. world-model representations를 학습하는 일반적인 방법은 흔히 reward 및 termination/continuation prediction과 함께 observation likelihood (pixel reconstruction)를 maximize하는 것입니다. Reconstruction은 종종 optimization을 안정화하는 dense supervision을 제공하지만, reward와 약하게만 결합된 visually detailed factors (예: 텍스처 또는 배경)에 capacity를 할당할 수도 있습니다. 이는 decision-making을 위해 직접적으로 latent space를 형성하는 decoder-free objectives의 동기가 됩니다.

Decoder-free world models. pixel reconstruction을 제거하면 문제는 modeling observations에서 latent state를 고정하는 요소와 learning signal이 타겟으로 하는 시간 인덱스를 선택하는 것으로 이동합니다.

한 부류는 task-oriented입니다: MuZero 및 TD-MPC 변형에서와 같이 search 또는 TD learning에 의해 유도된 supervision을 사용하여, reward/value prediction 및 planning을 지원하도록 latents가 optimize됩니다; 관련된 Dreamer-style agents 또한 reconstruction을 control-centric prediction objectives (예: MuDreamer)로 대체합니다.

두 번째 부류는 representation-oriented입니다: models는 때로는 future steps에 걸쳐 (예: CPC, SPR), 때로는 per-timestep invariances 또는 clustering을 통해 self-supervised objectives와 함께 학습된 embeddings를 predict하거나 align합니다.

partially observable control의 경우, 강력한 same-step objectives조차도 시간 $t$의 state가 $t+1$에 일어날 일을 predictive하게 만들 필요는 없습니다. NE-Dreamer는 representation-oriented 부류에 속하지만 이러한 temporal requirement를 명시적으로 만듭니다: causal sequence model은 history로부터 next encoder embedding을 predict하고 이를 stop-gradient target에 align하여, representation learning을 per-timestep agreement가 아닌 causal next-step prediction으로 바꿉니다.

Representation prediction and collapse prevention. future embeddings를 predict하는 것은 self-supervised learning에서 reconstruction에 대한 점점 더 인기 있는 대안이 되고 있습니다. 예를 들어, NEPA는 stop-gradient targets와 함께 next-embedding prediction을 적용하는 반면, I-JEPA 및 data2vec은 masked prediction과 context modeling에 중점을 둡니다.

핵심적인 문제는 학습된 state가 퇴화되는 representational collapse를 방지하는 것입니다. reinforcement learning에서 augmentations를 통한 invariance는 일반적인 stabilizer이며, Distracting Control Suite와 같은 benchmarks는 이를 명시적으로 만듭니다. BYOL, SimSiam, Barlow Twins 또는 VICReg에서 사용되는 것과 같은 Bootstrapping 및 redundancy-reduction regularizers 또한 negatives 없이 collapse를 방지할 수 있지만, 일반적으로 same timestep의 paired views에 적용됩니다.

NE-Dreamer는 이러한 아이디어를 predictive context로 일반화합니다: 이 모델의 causal sequence model은 history로부터 forecasted embedding $e_{t+1}$을 생성하며, 이는 (예를 들어 Barlow Twins loss를 사용하여) stop-gradient target에 align됩니다. 이는 latent space에서 temporal coherence를 강제하여, redundancy reduction을 단순한 within-frame invariance가 아닌 future prediction으로 확장합니다.

 

NE-Dreamer: Related Work 정리 노트

1. RSSM Backbone 유지 & Pixel Reconstruction 폐기

• 기존 흐름: 기존 Dreamer 계열은 RSSM을 활용하여 latent space에서 dynamics를 학습하고 planning을 수행함. 하지만 observation likelihood를 극대화하는 pixel reconstruction 방식은 배경, 텍스처 등 task-irrelevant하고 visually detailed한 요소에 모델 capacity를 낭비하는 치명적인 단점이 있음.
• NE-Dreamer의 포지셔닝: Dreamer의 검증된 RSSM control backbone은 그대로 유지하되, 무거운 pixel reconstruction을 폐기하고 latent representation을 학습하는 방식을 근본적으로 재설계함.

2. Decoder-free 계열 내에서의 차별화: Temporal Requirement의 명시화

• 기존 Decoder-free의 한계: 크게 Task-oriented(보상/가치 예측 중심)와 Representation-oriented(자기지도학습 기반)로 나뉨. 그러나 기존의 representation-oriented 모델들은 주로 동일 시간대(same-step)의 invariance나 clustering에 의존함. 이는 partial observability 환경에서 상태 변화를 추적하기에 불충분함.
• NE-Dreamer의 돌파구: Representation-oriented 방식에 속하지만, 단순 per-timestep agreement를 거부함. 과거의 history를 바탕으로 causal sequence model이 next encoder embedding을 예측하도록 강제하여, 학습의 목표를 causal next-step prediction으로 전환함.

3. Representational Collapse 방지 기법의 시간적 확장 (Temporal Generalization)

• 기존의 한계: Self-supervised learning에서 모델이 모든 입력을 똑같은 상태로 출력해버리는 representational collapse를 막기 위해 BYOL, Barlow Twins 등의 redundancy-reduction 기법을 사용함. 하지만 이는 철저히 '동일한 timestep' 내의 paired views (within-frame invariance)에만 적용되어 왔음.
• NE-Dreamer의 핵심 아이디어: 이 아이디어를 predictive context로 확장함. 예측된 $e_{t+1}$을 stop-gradient target에 align 시키는 과정에 Barlow Twins loss와 같은 redundancy-reduction을 적용함. 즉, 공간적인 불변성이 아닌 temporal coherence(시간적 일관성)를 future prediction에 강제하여 collapse를 방지함.

---

쉬운 설명 :

이 섹션은 NE-Dreamer가 "기존 AI들의 어떤 단점을 고쳤는지" 그리고 "비슷한 목적을 가진 다른 AI들과 무엇이 다른지"를 설명하는 족보 정리 파트입니다.

1. 그림 그리기 포기: 기존의 똑똑한 AI들은 자기가 본 게임 화면을 사진처럼 똑같이 다시 그려보는 연습(Reconstruction)을 했습니다. 그런데 점수 얻는 데는 전혀 쓸모없는 벽지 무늬나 그림자까지 완벽하게 그리려고 에너지를 다 써버렸죠. 그래서 NE-Dreamer는 과감하게 그림 그리기를 포기했습니다(Decoder-free).
2. "지금"이 아니라 "다음"을 맞혀라: 그림 그리기를 포기한 다른 AI들도 있었지만, 그들은 보통 "화면이 살짝 노이즈가 껴도 원래 화면이랑 똑같은 거야"라는 식으로 지금 당장의 모습에만 집중했습니다. 하지만 NE-Dreamer는 다릅니다. "지금까지 진행된 상황을 보니, 다음 1초 뒤에는 이런 특징의 화면이 나올 거야!"라며 시간의 흐름을 예측하는 능력을 키우는 데 집중했습니다.
3. 꼼수 방지법의 진화: AI가 학습하기 싫어서 "다음에 뭐가 나오든 그냥 다 똑같은 화면이야"라고 우기는 꼼수(Collapse)를 부릴 수 있습니다. 기존에는 이 꼼수를 막기 위해 '같은 시간'에 찍은 사진 두 장을 비교하는 수학적 기법을 썼습니다. NE-Dreamer는 이 수학적 기법을 기발하게 변형해서, '내가 예측한 미래'와 '실제 미래'를 비교하는 데 사용했습니다. 덕분에 꼼수를 부리지 못하고 시간의 흐름을 완벽하게 이해하게 되었습니다.

 

3. Method

3.1. Problem setup

우리는 pixels로부터의 partially observable control을 연구합니다. 시간 $t$에서 environment는 image observation $x_t$를 방출합니다. agent는 action $a_t$를 선택하고 reward $r_t$를 받습니다.

우리는 또한 continuation indicator $c_t \in {0, 1}$를 사용하는데, 여기서 episode가 $t$에서 $t+1$로 계속되면 $c_t = 1$이고 terminal transitions에서는 $c_t = 0$입니다.

NE-Dreamer는 표준 Dreamer pipeline을 따르지만—(i) experience로부터 latent world model을 학습하고, (ii) latent space의 imagined rollouts에서 actor-critic을 train합니다—world model을 위한 representation objective를 변경합니다.

구체적으로, pixel reconstruction을 제거하고 대신 next-step encoder embedding을 predict합니다. 시간 $t$까지 사용 가능한 정보만 사용하여, model은 $\hat{e}_{t+1}$을 predict하고 이를 self-supervised loss(우리의 구현에서는 Barlow Twins)를 사용하여 stop-gradient target에 align합니다.

3.2. Latent world model (RSSM)

우리는 deterministic recurrent state $h_t$와 stochastic latent $z_t$를 가진 recurrent state-space model (RSSM)을 기반으로 합니다.

Encoder and latent inference. encoder는 observations를 embeddings로 매핑합니다:

$$e_t = f_{enc}(x_t)$$

이전의 latent state와 이전의 action이 주어지면, RSSM은 그것의 deterministic state를 업데이트합니다:

$$h_t = f_{rec}(h_{t-1}, z_{t-1}, a_{t-1})$$

그런 다음 stochastic latent에 대한 prior 및 posterior를 정의합니다:

$$p_\phi(z_t | h_t), \quad q_\phi(z_t | h_t, e_t)$$

world-model training 동안 우리는 $z_t \sim q_\phi(z_t | h_t, e_t)$를 샘플링하고; imagination 동안에는 $\hat{z}t \sim p\phi(z_t | h_t)$를 샘플링합니다.

Reward and continuation heads. Dreamer에서와 같이, world model은 reward와 continuation을 predict합니다:

$$p_\phi(r_t | h_t, z_t), \quad p_\phi(c_t | h_t, z_t)$$

표준 Dreamer는 또한 pixel decoder $p_\phi(x_t | h_t, z_t)$를 통해 observations를 predict합니다. NE-Dreamer는 이 decoder를 제거하고 Sec. 3.3의 next-embedding objective로 대체합니다.

World-model objective. world model은 reward 및 continuation likelihoods, prior-posterior regularizer, 그리고 제안된 next-embedding loss로 train됩니다:

$$L_{wm} = L_{rew} + L_{cont} + \beta_{kl} L_{kl} + \beta_{ne} L_{NE}$$

prediction losses는 negative log-likelihoods입니다:

$$L_{rew} = -\mathbb{E}[\log p_\phi(r_t | h_t, z_t)]$$

$$L_{cont} = -\mathbb{E}[\log p_\phi(c_t | h_t, z_t)]$$

KL 항은 posterior를 prior 방향으로 regularize합니다:

$$L_{kl} = \mathbb{E} \left[ \text{KL} \left( q_\phi(z_t | h_t, e_t) \parallel p_\phi(z_t | h_t) \right) \right]$$

우리는 $L_{kl}$에 대해 표준 Dreamer stabilizers (예: KL balancing / free-nats)를 채택하며; 세부 사항은 이전 Dreamer 관행을 따릅니다.

3.3. Next-embedding predictive alignment

NE-Dreamer는 latent dynamics가 representation space에서 predictive하도록 train합니다: 시간 $t$까지의 history로부터, next observation의 encoder embedding을 predict하고 그 prediction을 stop-gradient target에 align합니다.

Causal next-embedding predictor. causal mask가 있는 causal temporal transformer $T_\theta$는 next-step embedding prediction을 생성하기 위해 시간 $t$까지 사용 가능한 정보만 사용합니다:

$$\hat{e}_{t+1} = T_\theta(h_{\le t}, z_{\le t}, a_{\le t})$$

target은 next-step encoder embedding입니다:

$$e^\star_{t+1} = \text{sg}(e_{t+1}) = \text{sg}(f_{enc}(x_{t+1}))$$

우리는 stop-gradient를 $\text{sg}(\cdot)$로 표기합니다. Gradients는 $\hat{e}_{t+1}$을 통해 $T_\theta$와 RSSM으로 흐르지만, $e^\star_{t+1}$을 통해서는 흐르지 않습니다.

Alignment loss (Barlow Twins).

우리는 predicted 및 target embeddings 사이의 Barlow Twins redundancy-reduction objective로 $L_{NE}$를 구현합니다. $\tilde{\hat{e}}{t+1}$과 $\tilde{e}^\star{t+1}$을 각 minibatch 내의 valid transitions 집합에 대해 차원별로 정규화된(zero mean, unit variance) embeddings라고 합시다.

다음과 같이 정의합시다:

$$\mathcal{I} \doteq \{(b, t) : c^{(b)}_t = 1\}, \quad N \doteq |\mathcal{I}|$$

cross-correlation matrix는 다음과 같습니다:

$$C_{ij} = \frac{1}{N} \sum_{(b,t) \in \mathcal{I}} \tilde{\hat{e}}^{(b)}_{t+1,i} \tilde{e}^{\star(b)}_{t+1,j}$$

next-embedding loss는 다음과 같습니다:

$$L_{NE} = \sum_i (1 - C_{ii})^2 + \lambda_{BT} \sum_{i \neq j} C_{ij}^2$$

이 objective는 invariance (큰 대각선 상관관계)를 장려하는 동시에 redundancy (작은 비대각선 상관관계)를 억제하며, 여기서는 same-timestep matching이 아닌 next-step prediction에 적용됩니다.

3.4. Actor-Critic Learning

DreamerV3와 마찬가지로, NE-Dreamer는 world model을 사용하여 imagined trajectories를 생성함으로써 latent space에서 policy와 value function을 학습합니다. 이러한 imagined trajectories (horizon $H = 15$ steps)는 효율적인 batch actor-critic updates를 가능하게 합니다.

우리는 imagined full latent state를 $s_t = (h_t, \hat{z}_t)$로 나타내며, 여기서 $\hat{z}_t \sim p_\phi(z_t | h_t)$입니다. 각 imagination step에서 actions는 policy $\pi_\theta$로부터 샘플링되고 그 values는 critic $V_\psi$에 의해 추정됩니다:

$$a_t \sim \pi_\theta(a_t | s_t), \quad V_\psi(s_t) \approx \mathbb{E}_{p_\phi, \pi_\theta} [R^\lambda_t]$$

Critic: critic은 imagined rewards를 기반으로 $\lambda$-returns의 distribution을 predict합니다:

$$R^\lambda_t = r_t + \gamma c_t ((1 - \lambda)V_\psi(s_{t+1}) + \lambda R^\lambda_{t+1})$$

$$L_{critic}(\psi) = -\mathbb{E}_{p_\phi, \pi_\theta} \left[ \sum_{t=1}^H \log p_\psi(R^\lambda_t | s_t) \right]$$

Actor: actor는 EMA-based scale인 $S$를 사용하여 정규화된 advantages를 극대화합니다:

$$L_{actor}(\theta) = -\mathbb{E}_{p_\phi, \pi_\theta} \left[ \sum_{t=1}^H \text{sg} \left( \frac{R^\lambda_t - V_\psi(s_t)}{\max(1, S)} \right) \log \pi_\theta(a_t | s_t) + \eta \mathcal{H}[\pi_\theta(a_t | s_t)] \right]$$

여기서 $\text{sg}(\cdot)$는 stop-gradient 연산자를 나타내고 $\eta$는 entropy regularization coefficient입니다.

Policy gradients는 continuous actions를 위해 world model을 통해 backpropagate됩니다. learning procedure와 모든 hyperparameters는 DreamerV3와 일치하며, 이는 관찰된 성능 향상이 representation learning objective에서 비롯되었음을 보장합니다.

 

NE-Dreamer: Method 정리 노트

1. 기본 구조: Dreamer 파이프라인 유지 및 Decoder의 완전한 제거

• 기반 모델: 결정론적 상태 $h_t$와 확률론적 잠재 상태 $z_t$를 사용하는 RSSM (Recurrent State-Space Model) 구조를 채택함.
• 결정적 차이: 기존 Dreamer의 핵심이었던 observation을 복원하는 pixel decoder $p_\phi(x_t | h_t, z_t)$를 완전히 제거함.
• 새로운 World-model objective:(Reconstruction loss가 사라지고 제안하는 Next-embedding loss인 $L_{NE}$가 그 자리를 대체함.)
• $$L_{wm} = L_{rew} + L_{cont} + \beta_{kl} L_{kl} + \beta_{ne} L_{NE}$$

2. 핵심 제안: Next-Embedding Predictive Alignment ($L_{NE}$)

• 예측 (Predictor): 시간 $t$까지의 history $(h_{\le t}, z_{\le t}, a_{\le t})$만을 입력으로 받아, causal temporal transformer $T_\theta$가 다음 스텝의 임베딩 $\hat{e}_{t+1}$을 예측함.
• 정답 (Target): 실제 다음 스텝의 observation $x_{t+1}$을 Encoder에 통과시킨 값. 이때 모델 붕괴를 막기 위해 stop-gradient 연산 $\text{sg}(\cdot)$을 반드시 적용함.
• $$e^\star_{t+1} = \text{sg}(f_{enc}(x_{t+1}))$$
• 정렬 손실 함수 (Barlow Twins): 예측값과 Target 간의 cross-correlation matrix $C_{ij}$를 계산하여 $L_{NE}$를 산출함.대각 성분(invariance)은 1에 가깝게 유지하고, 비대각 성분(redundancy)은 0에 가깝게 억제함. 중요한 점은 이를 동일 시간대가 아닌 미래 예측(next-step prediction)에 적용하여 시간적 일관성을 강제했다는 것임.
• $$L_{NE} = \sum_i (1 - C_{ii})^2 + \lambda_{BT} \sum_{i \neq j} C_{ij}^2$$

3. Actor-Critic Learning: 통제된 변인

• 학습 방식: World model이 상상한 길이 $H=15$의 imagined trajectories 위에서 policy $\pi_\theta$ (Actor)와 value function $V_\psi$ (Critic)를 학습함.
• 의도적인 설계: 이 과정의 모든 learning procedure와 hyperparameters를 기존 SOTA인 DreamerV3와 완벽하게 동일하게 세팅함. 이는 모델의 성능 향상이 Actor-Critic 부분의 튜닝이 아니라, 오직 "Representation learning objective($L_{NE}$)의 변화"에서 비롯되었음을 엄밀하게 증명하기 위함임.

---

쉬운 설명 :

이 섹션은 NE-Dreamer가 실제로 어떻게 조립되어 작동하는지 구체적인 메커니즘을 설명합니다.

복잡한 게임 디자인이나 밸런스를 잡는 과정을 생각해 보면 이해가 쉽습니다. AI 에이전트가 게임 환경을 학습할 때, 기존 방식은 다음 프레임에 렌더링될 모든 그래픽 텍스처와 배경의 시각적 이펙트까지 일일이 상상해서 그려내려 했습니다(Pixel Reconstruction). 하지만 이렇게 시각적인 디테일에 집착하는 것은 무겁기만 할 뿐, 실제 게임을 플레이하는 데는 비효율적입니다.

NE-Dreamer는 화려한 그래픽을 복원하는 대신, **"다음 순간의 핵심 상태(Embedding)가 어떻게 변할 것인가"**에만 집중합니다.

Transformer라는 예측 엔진이 플레이어의 현재까지의 행동 이력을 바탕으로 다음 상태의 핵심 특징을 예측하고, 실제 다음 상태와 비교하여 정답을 맞춰갑니다. 이때 모델이 단순히 모든 상태를 똑같다고 우기는 꼼수(Collapse)를 쓰지 못하도록, Barlow Twins라는 수학적 장치를 통해 불필요한 정보는 버리고 '진짜 핵심적인 변화'만 남기도록 밸런스를 엄격하게 잡습니다.

마지막으로 행동을 결정하는 Actor-Critic 부분은 이미 성능이 입증된 기존 시스템(DreamerV3)의 코드를 그대로 가져와 썼습니다. "행동하는 방식"은 똑같이 두되 "세상을 이해하는 방식(Representation)"만 바꿨을 때 AI가 얼마나 더 똑똑해지는지 명확하게 비교하기 위한 연구자들의 치밀한 설계입니다.

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

DMLab이나 DMC 같은 공간이 환경
에이전트는 자기 관측치(x_t)만 볼 수 있음.

매 시간(t)마다 에이전트가 현재 이미지(x_t)를 보고 특정 행동(a_t)을 취하면, 환경은 그 결과로 에이전트에게 보상(r_t)을 주고 다음 단계로 계속 넘어갈지(c_t) 여부를 알려줌. 
이게 반복

보상과 다음 스테이지 여부는 환경이 하는 것, 우리의 목표는 에이전트를 학습하는 것.

큰 흐름은 처음에 G개 생성 후 배치사이즈 x 64개의 데이터로 한스텝 학습
그 뒤 d개 생성 학습
d개 생성 학습 반복

Replay Buffer에 데이터는 계속 쌓이는 구조 500만개까지 있다고 함.


이제 학습 과정

준비물
인코더 : 이미지를 임베딩으로
RSSM : 기억 변수인 z를 관리하는 모델
(이전스텝의 기억(h), 이전 스텝의 상태(z), 에이전트의 행동(a), 인코더가 만든 현재 임베딩)
을 넣어서
현재 기억 h와 현재 상태 z를 뱉음

Causal Temporal Transformer : 예언가 
(기억 h, 상태 z, 행동 a의 연속된 시퀀스) 

Predictors : 
RSSM이 만든 현재 기억(h)과 현재 상태(z)
현재의 정보만 가지고
이 상태에서 얻게 될 보상(r)과 게임이 안 끝나고 계속될 확률(c)을 예측하는군

이걸 가짜 점수와 보상으로 사용하기 위함 같음. 


동작을 살펴보자!

진짜 환경에서 x_t 이미지를 가져옴. 이는 현재 이미지.
이를 인코더에 통과시켜 e_t라는 현재 임베딩을 만듬

RSSM차례
과거 기억, 과거 상태, 과거 행동을 RSSM에 넣어 현재의 기억으로 갱신
현재의 기억과 현재 임베딩을 합쳐서 z를 뽑아냄 (학습 모델)
이때 현재 임베딩을 안쓰고도 z를 뽑음


트랜스 포머 차례
과거부터 지금까지의 기록, h, z ,a 시퀀스 를 예언가인 트랜스 포머에 넣는다 (전부 시퀀셜)
그래서 다음 스텝의 예측 임베딩을 뽑게됨.

채점타이밍
트랜스포머가 잘해야하는데. 채점을 해야함.
진짜 다음 스텝의 이미지를 인코더에 넣어서 정답 임베딩을 만듬.
(여기서 그걸로 채점해서 transformer부터 rssm, 인코더로 쭉 전파시킴, 정답용 인코더는얼음)


그리고 점수 예측 타이밍 
현재 기억, 상태를 가지고 보상과 게임 연속성 예측 (MLP차례)

모델이 추측한 보상과 연속성을 실제와 비교해서 loss 

KL loss : z끼리 비교 (안보고 푼애가 보고 푼 애를 따라가도록)

Next-Embedding Loss : 비슷해지게는 하지만 중복성제거

기존 모델은 복원을 하느라 시간을 썼지만 이건 안해도됨

즉 저 4개의 loss를 통해서 그냥 업데이트 해버림.


근데 이러면 무슨 행동할지를 그 누구도 예측하지 않음.
액터가 필요 : h와 z를 보고 무슨행동할지 결정
가치평가자도 필요 : 현재 상태를 보고 "이대로면 앞으로 받을 총 점수가 얼마일지"
(보상해킹 방지)

일단 순서는
세계 모델 업데이트 후 액터가 15스텝을 진행.

액터가 각 생성한 15개에 대해 프레딕터가 뽑아준 점수가 있을것이고.
크리틱이 뽑아준 점수가 있겠네. 

[끝점] 15스텝 총점: 크리틱이 15스텝 상태를 보고 "앞으로 게임 끝날 때까지 얻을 가치"를 예측

14스텝 총점: (14스텝 프레딕터 보상) + (방금 구한 15스텝 총점). 끝입니다.

13스텝 총점: (13스텝 프레딕터 보상) + (방금 구한 14스텝 총점). 끝입니다.

...

1스텝 총점: (1스텝 프레딕터 보상) + (방금 구한 2스텝 총점). 끝입니다.


이렇게 각 스텝에 대해 총점을 구함.

크리틱은 각 스텝별 총점을 기준으로, 자기가 예측한 추측점수와 비교 후 줄이는 방향으로 역전파 (이 이유는 나중엔 크리틱 혼자 점수를 내야하기때문)


액터는 

어드밴티지 = (방금 뼈빠지게 역순으로 계산한 람다 리턴 총점) - (크리틱이 원래 그 상태만 보고 "음, 이 정도 나오겠네" 하고 예측했던 추측 점수)

이점을 계산해서 높이는 그 행동을 하도록 가중치 업데이트

반복. 

그럼 월드모델 학습되고, 액터 학습되고, 크리틱 학습됨. 

인퍼런스때도


인코더 : 이미지 받아서 압축
RSSM : 상태 출력
액터 : 행동

반복.






인퍼런스 요약

Encoder: 실제 화면 $x_t$를 보고 특징 $e_t$로 압축.

RSSM (Posterior): 과거 기억 $h_{t-1}$과 방금 압축한 $e_t$를 합쳐서 현재의 완벽한 상태 $z_t$를 생성. (그리고 다음 스텝을 위해 $h_t$ 갱신).

Actor: $h_t$와 $z_t$를 보고 "이 행동 $a_t$를 해야겠다!"라고 결정.

결정된 행동을 환경에 던져주고, 다음 화면을 받아 1번부터 반복.