PlaNet 논문 리뷰

Learning Latent Dynamics for Planning from Pixels

---

World model은 Model Base RL의 가능성에 대해서 지평을 연 논문 중 하나이지만, Controller(Policy)을 gradient descent가 아닌 Evolutional strategy를 이용한 학습을 한다는 점에서 한계점이 많았다. 기껏 미분이 되는 신경망으로 세계를 구현해놓고는 정작 그 안에서 애이전트를 학습시킬 때에는 미분을 사용하지 않은 것이기 때문이다.

 

이 지점에서 PlaNet가 등장한다. 논문의 저자들은 World Model의 그러한 비효율성을 지적하면서, 더 강력하고 미분 친화적인 모델인 RSSM을 제시한다. 그러면 Policy까지 미분이 이어지게 했냐? 아니다 ㅋㅋ. 당시의 기술적 한계로 수십 스텝 미래로 이어지는 궤적에 대해 직접 역전파를 수행하는 것이 불안정했기 때문이다. 실제로 미분으로 정책을 학습하는 로직은 Dreamer에서 나오므로, 이 논문은 RSSM에 집중하면서 보면 되겠다.

 

Recurrent State Space Model

1. RNN (Recurrent Neural Network) 기반의 접근과 MDN-RNN

우선 RNN부터 살펴보자. RNN은 이전 단계의 히든 스테이트($h_{t-1}$)와 현재의 행동($a_{t-1}$)을 입력으로 받아, 현재의 상태($h_t$)를 갱신하는 구조다. 기본적으로 과거의 정보를 압축하여 다음 상태 추론의 문맥(Context)으로 사용하기 때문에 순차적 의사결정(Sequential Decision Making)에 적합한 설계다.

 

그러나 RNN이 모델 기반 강화학습(MBRL)이나 생성 모델에서 치명적인 약점을 드러내는 지점은 바로 멀티모달리티를 다룰 때다. 여기서 멀티모달리티란 데이터의 타입이 다양하다는 뜻이 아니라, 미래의 결과가 여러 개의 봉우리(Mode)를 가진 확률 분포로 나타난다는 통계적 의미다. 예를 들어, 운전 중인 자동차가 갈림길을 만났다고 가정해보자. 차는 왼쪽으로 갈 수도 있고, 오른쪽으로 갈 수도 있다. 두 가지 선택지가 모두 정답일 수 있는 상황이다.

하지만 바닐라(Vanilla) RNN은 본질적으로 **결정론적(Deterministic)**이다. 입력($x$)과 이전 상태($h$)가 주어지면, 출력은 고정된 하나의 값으로 떨어진다. 통상적으로 RNN을 학습시킬 때 사용하는 MSE(Mean Squared Error) 손실 함수는 정답과 예측값의 '거리'를 줄이려 노력한다.

 

이 과정에서 평균의 함정이 발생한다.RNN은 왼쪽 길(Mode A)과 오른쪽 길(Mode B) 중 하나를 선택하는 것이 아니라, 에러를 최소화하기 위해 **그 둘의 평균값인 '중앙'**을 예측해버린다. 결과적으로 에이전트는 갈림길 한가운데에 있는 벽을 들이받거나, 생성된 이미지는 두 가능성이 겹쳐진 뿌연 형체가 된다.

이것이 바로 RNN이 멀티모달리티를 처리하지 못한다는 증거다. 고차원 데이터의 복잡한 확률 분포를 단순한 결정론적 벡터 하나($h_t$)에 구겨 넣으려다 보니, 정보의 손실이 일어나고 미래의 불확실성을 표현할 능력을 상실하는 것이다.

 

따라서 초기 World Model(Ha et al., 2018)은 이 RNN에 **MDN(Mixture Density Network)**을 결합하는 방식을 택했다. 미래는 불확실하기 때문에, 다음 상태를 하나의 값으로 찍는 것이 아니라 가우시안 혼합 분포(GMM)를 통해 확률적으로 예측하기 위함이다. 즉, RNN은 "과거의 기억($h_t$)"을 담당하고, MDN은 그 기억을 바탕으로 "미래의 불확실성($z_{t+1}$)"을 확률 분포($\mu, \sigma, \pi$)로 뱉어낸다. 하지만 이 구조에는 치명적인 구조적 한계가 존재한다.

첫째, 학습의 분리(Decoupling) 문제다. World Model은 시각 정보를 압축하는 VAE(Visual Encoder)와 미래를 예측하는 RNN을 따로 학습시킨다. VAE는 단순히 "이미지를 잘 복원하는 것"이 목표지, "미래 예측에 중요한 정보를 남기는 것"이 목표가 아니다. 만약 VAE가 날아오는 야구공처럼 작지만 중요한 객체를 노이즈로 취급해 지워버린다면, 뒷단의 RNN은 아무리 성능이 좋아도 공의 궤적을 예측할 수 없다.

둘째, 결정론적(Deterministic) 상태의 한계다. RNN의 히든 스테이트 $h_t$는 본질적으로 결정론적이다. 확률적 요소($z$)는 매 스텝 외부에서 샘플링되어 들어올 뿐, 상태 전이 자체는 고정된 연산을 따른다. 이는 모델이 복잡한 확률적 환경을 내재화(Internalize)하는 데 방해가 되며, 장기 예측(Dreaming) 시 오차가 누적되어 환각(Hallucination)을 일으키기 쉽다.

 

2. SSM (State Space Model)

SSM은 샘플링 기반으로 다음 스테이트를 제공하기 때문에 멀티모달리티에 대한 걱정을 할 필요가 없다. 그저 확률분포 내에서 하나의 아웃풋이 나오기 때문이다. 그러나 문제는 이러한 sequential sampling은 과거의 정보를 보존하지 못하기 때문에(그렇다고 RNN이 정보를 전부 다 보전하는건 아니다 ㅋㅋ) 순차 의사 결정에 문제가 생긴다.

 

Recurrent State Space Model

결국 RSSM(Recurrent State-Space Model)은 앞서 언급한 RNN의 결정론적(Deterministic) 성질과 SSM의 확률적(Stochastic) 성질을 하나로 합친 모델이다.

이것이 왜 강력한가? 바로 '역할의 완벽한 분업화' 때문이다.

기존 RNN에서는 하나의 hidden state가 과거의 문맥도 기억해야 하고, 동시에 미래의 불확실한 분포도 표현해야 하는 과부하 상태였다. 하지만 RSSM은 이 짐을 둘로 나눈다.

• 결정론적 상태($h_t$): 오직 정보의 보존(Information Preservation)에 집중한다. 과거의 긴 데이터를 잃어버리지 않고 시간 축을 따라 튼튼하게 전달하여 장기 의존성(Long-term dependency)을 확보한다.
• 확률적 상태($s_t$): 멀티모달리티(Multimodality)의 해결을 전담한다. 미래가 여러 갈래로 나뉠 때, 어설픈 평균값이 아니라 샘플링(Sampling)을 통해 구체적인 하나의 가능성을 명확하게 선택해준다.

즉, "기억은 확실하게($h_t$), 상상은 유연하게($s_t$)" 가져가는 전략이다. 덕분에 RSSM은 hidden state가 단독으로 미래를 예측하며 겪던 표현력의 한계를 벗어나, 안정적인 기억 위에서 다양한 미래를 선명하게 시뮬레이션할 수 있게 된다.

 

RSSM의 objective는 어떻게 될까? 본 논문이 "Planning"임을 잊지말자. 즉 내가 어떤 행동을 했을 때 나올 실제 결과(Observation, Reward)를 예측해야 하므로, 행동에 대한 결과의 우도(Likelihood, $\log p(o_{1:T}|a_{1:T})$)를 최대화해야 한다.

그러나 단순히 우도만 최대화하려고 하면 계산이 불가능하거나, 계획에 쓸 수 없는 모델이 되어버린다. 왜 그런지 생각해보자:

우선 $p(o)$를 구할 수 있어야 한다. RSSM 모델에서 $p(o)$를 구하기 위해서는 

$$p(o) = \int p(o|s) p(s) \, ds$$

와 같은 식을 풀어야 하지만, $p(s)$를 직접적으로 도출할 수도 없으며, ,$p(o|s)$에서 결과 $o$에 대해서 원인 $s$의 가능성을 전부 찾는 것 자체도 불가능에 가깝기 때문이다. 간단하게 예시를 들면, "집이 어질러져 있다"는 결과에 대해서 원인은 도둑이 들었을 수도 있고, 내가 그냥 청소를 안하는 더러운 인간일 수도 있고, 청소는 잘 하는데 시간이 없어서 급하게 서류를 찾느라 난장판을 쳐놨을 수도 있기 때문이다.

어떻게든 잘 모델링을 해서 도출을 했다고 해보자. 그래도 문제가 없을까? 그냥 state만 잘 보고 이미지를 예측하는 모델이 되어버린다. Planning에서는 과거의 상태를 바탕으로 state를 예측해야 하는데, 그 상태 자체를 추론할 수 없게 되는 것이다.

따라서 본 논문에서는 ELBO를 통해 likelihood의 lower bound를 높임으로써 이 문제를 해결한다. 유도과정을 간단하게 확인해보자.

RSSM에 맞게 다시 보면 더 좋다.

정말 우아하다고 말하지 않을 수 없는 식이다. $q$는 현재까지의 memory($h$)와 현재의 observation을 통해서 $s$의 확률분포를 주는 함수(Posterior)이다. 실제 batch processing을 할 때, 이 RSSM에서 나오는 샘플을 통해 현재 상태와 $h$를 받아 이미지를 다시 예측하도록 학습도 되며, KL 발산(Divergence) 항을 통해 '이미지를 보고 추론한 분포($q$)'와 '과거 기억만으로 예측한 분포($p$)'의 차이를 최소화한다.

이 지점에서 RSSM이 채택한 Variational Inference(변분 추론)와 ELBO(Evidence Lower Bound)의 도입이 왜 강력한지 주목해야 한다. 이것은 단순한 수학적 테크닉을 넘어, 모델 기반 강화학습의 고질적인 난제를 해결하는 열쇠가 되기 때문이다.

첫째, '정답지'를 통한 자가 학습(Self-Correction) 메커니즘을 완성한다. ELBO의 KL 항은 마치 **'선생님($q$)과 학생($p$)'**의 관계와 같다.

• 학생(Prior $p$): 과거의 기억($h$)만 가지고 미래($s$)를 상상해본다. (Planning 상황)
• 선생님(Posterior $q$): 실제 정답($o$, 이미지)을 보고 정확한 상태를 파악한다. (Training 상황)
• 학습 과정에서 학생은 선생님의 답을 보며 자신의 예측을 끊임없이 교정한다. 이 덕분에 나중에 정답지(이미지)가 없는 실전(Planning)에서도, 모델은 선생님이 옆에 있는 것처럼 정확한 상상을 할 수 있게 된다.

둘째, 계산 불가능한 문제(Intractability)를 최적화 가능한 문제로 바꾼다.

원래대로라면 관측 확률 $p(o)$를 최대화하려면 모든 가능한 잠재 상태 $s$에 대해 적분을 해야 한다. 이는 신이 아닌 이상 불가능한 계산이다. 하지만 ELBO를 도입함으로써 우리는 이 적분 문제를 미분 가능한 최적화 문제로 우아하게 치환했다. 덕분에 딥러닝의 역전파 (Backpropagation)를 이용해 거대한 모델을 End-to-End로 학습시킬 수 있는 길이 열린 것이다.

 

셋째, 정보의 누수 없는 '완전체' 학습이다.

World Model(2018)은 시각(VAE)과 기억(RNN)을 따로 학습시켰다. 이로 인해 시각 모델이 예측에 필요한 정보를 멋대로 지워버려도 손쓸 방법이 없었다. 반면, RSSM은 ELBO라는 하나의 목적 함수 안에서 이미지 압축, 미래 예측, 상태 추론이 동시에 맞물려 돌아간다. 예측에 필요한 정보라면 압축 과정에서 반드시 살아남게 되고, 압축된 정보는 다시 정교한 예측의 재료가 된다.

결국 ELBO는 "상상(Dynamics)이 현실(Observation)에 단단히 발을 붙이게(Grounding) 만드는" 가장 수학적으로 아름다운 도구인 셈이다.

 

본 논문에서는 추가로 Latent Overshooting에 대해서도 다룬다. 아직까지는 1-Step prediction이기 때문에, multi step planning에는 에러가 누적되어 미래에 대한 hallucination이 일어나기 때문이다. 따라서 본 논문은 단순히 즉각적인 다음 스텝만 예측하는 것이 아니라, 앞으로의 $d$ 스텝 후도 예측하도록 한다.

Overshooting loss를 적용하다보면, 아무래도 더 먼 미래를 예측하는 테스크다보니 Prior와 Posterior간의 오차가 심해져 representation이 무너질 가능성이 존재한다. 따라서 여기서는 posterior의 그래디언트가 전파되지 않도록 한다. 한가지 그럼 의문은... 1-step prediction은 그럴 가능성이 없느냐? 없지는 않다. 다만 1-step라 덜 고려해도 된다는 정도만 알아두면 되겠다. 이후 Dreamer에서는 KL balancing을 통해 1-step prediction에서 나타날 수 있는 문제를 해결하기도 한다.

 

사실 그리고 논문의 저자들이 고백하는 것은... 이거 안해도 좋더라 이거다 ㅋㅋ RSSM 자체가 성능이 충분히 좋았나보다.

 

이제 어떻게 실제로 모델이 돌아가는지 확인해보자:

나노바나나 참 좋다 ㅋㅋ. 현재 관측 이미지 $o_t$를 보고, 이전의 정보 $h_t$를 참고하여 posterior가 현재 state를 인지한다. 이를 바탕으로 prior는 앞으로의 state와 reward에 대해서 샘플링(CEM) 기법으로 생성된 수많은 가상 행동 궤적(Action Trajectory)들을 상상 속에서 미리 시뮬레이션한다. 즉, 실제 환경에서 위험하게 부딪혀보지 않고도 학습된 Dynamics 모델(Prior)을 통해 '내가 이렇게 행동하면 미래가 어떻게 변하고 얼마만큼의 보상을 받을지'를 수천 번 빠르게 계산해본 뒤, 그중 누적 보상이 가장 높은 최적의 행동을 선택해 실행하는 것이다. 이것이 바로 PlaNet이 보여준 잠재 공간에서의 MPC(Model Predictive Control)이다.

 

보듯이 control에 대해서 미분 전략은 없으며 그저 샘플링을 통해 가장 좋은 것을 찾아내는 것이기 때문에 다음 행동을 결정하는 쪽에서 bottleneck가 생긴다. 나중에 이것을 actor-critic 구조로 미분을 완전히 이어서 bottleneck도 없앤 것이 dreamer이니, 나중에 리뷰해보도록 하겠다.

 

오늘의 리뷰는 여기서 마무리!