LeMario: JEPA 모델로 학습된 Super Mario Bros

저는 픽셀과 행동으로부터 월드 역학(world dynamics)을 학습하는 작은 Joint-Embedding Predictive Architecture (JEPA)인 LeWorldModel을 재현해보고 싶었습니다. 원본 논문에서는 이를 Push-T에서의 보상 없는 계획(reward-free planning)에 사용했습니다. 하지만 저는 비디오 게임을 좋아했고, 동시에 LeCun의 JEPA 아키텍처에 대해 더 깊이 배우고 싶었기에, 전체 아키텍처를 처음부터 직접 작성하여 Super Mario Bros로 학습시키기로 결정했습니다.

모델은 제가 처음에 중요하다고 생각했던 모든 테스트를 통과했습니다. 모델은 학습에 사용되지 않은 에피소드(held-out episodes)에도 일반화되었고, 행동(actions)을 사용했으며, 강력한 베이스라인(baselines)보다 더 나은 5단계 미래 예측 성능을 보여주었습니다. 가공되지 않은 보상 없는 계획(Raw reward-free planning)은 Mario를 근처의 이미지 목표물 쪽으로 이동시킬 수 있었고, 목표 지점의 2~5 픽셀 이내로 완료할 수 있었습니다. :D

잠시 동안은 모델이 게임을 플레이하는 법을 배운 것처럼 보였습니다. 그러다 제가 목표 지점을 레벨의 더 먼 곳으로 옮겼을 때... Mario는 첫 번째 주요 장애물을 안정적으로 뛰어넘지 못하거나 단일한 먼 목표 이미지 쪽으로 이동하지 못했습니다.

모델은 게임을 예측하는 법은 배웠지만, 그것이 게임을 통해 전진하는 법을 배웠다는 의미는 아니었습니다. D:

이 포스트는 기술적인 설명인 동시에 사후 분석(postmortem)이기도 합니다. 제가 무엇을 만들었는지, 어떻게 테스트했는지, 어떤 실수를 했는지, 그리고 실제 문제를 점진적으로 드러낸 실험들에 대해 다룹니다. (대부분의 교훈은 지나고 나면 당연해 보입니다 T^T )

전체 아키텍처

각 방정식을 개별적으로 소개하기 전에, 전체 메커니즘을 한눈에 보는 것이 도움이 됩니다:

녹색 경로부터 시작하겠습니다. 각 학습 샘플은 4개의 Mario 프레임을 포함합니다. **비전 인코더 (vision encoder)**는 모든 프레임을 **잠재 변수 (latent, $z$)**라고 불리는 192개의 숫자로 이루어진 표현으로 압축합니다:

잠재 변수는 스크린샷에 대한 모델만의 독자적인 설명이라고 생각하면 됩니다.

빨간색 경로는 컨트롤러 입력을 포함합니다. 각 관측(observations) 쌍은 5개의 에뮬레이터 프레임으로 구분되며, 각 프레임은 6가지 가능한 버튼 상태를 포함합니다:

frames: [batch, 4, 3, 224, 224]
actions: [batch, 4, 5, 6] # Left, Right, Up, Down, A, B

**액션 인코더 (action encoder)**는 각 5 × 6 버튼 시퀀스를 또 다른 192개 숫자로 이루어진 벡터로 압축합니다.

그 후 프레임(frame) 및 액션 잠재 변수(action latents)는 **인과적 예측기 (causal predictor)**로 전달됩니다. 이 예측기의 역할은 다음과 같은 질문에 답하는 것입니다:

이전 프레임들이 어떤 모습이었고 어떤 버튼이 눌렸을 때, 다음 프레임의 잠재 변수는 어떤 모습이어야 하는가?

예측기는 6개의 트랜스포머 블록 (transformer blocks)으로 구성됩니다. 여기서 각 프레임은 이전 프레임들을 참조(attend)하게 됩니다. 하지만 액션을 어떻게 주입할 수 있을까요?

액션은 **적응형 레이어 정규화 제로 (Adaptive LayerNorm Zero, AdaLN-Zero)**를 통해 이 트랜스포머 블록들에 입력됩니다.

단순히 액션 벡터를 프레임 벡터에 붙이는 대신, AdaLN-Zero는 각 액션을 세 가지 종류의 제어 신호로 변환합니다:

시프트 (Shift): 프레임 특징(features)에 액션에 의존적인 오프셋(offset)을 더합니다.
스케일 (Scale): 특정 특징들을 키우거나 줄입니다.
게이트 (Gate): 트랜스포머가 현재 상태를 얼마나 강하게 업데이트할지 제어합니다.

이제 이것들이 트랜스포머 블록에 어떻게 영향을 미칠까요? 일반적으로 일반적인 어텐션 (attention)은 각 프레임에 이전 컨텍스트를 부여한 다음, 해당 정보를 합성하기 위해 피드포워드 (feedforward, MLP)를 통과시키지만, AdaLN은 액션에 따라 두 단계 모두를 수정합니다.

예를 들어, 점프 액션은 수직 운동과 관련된 잠재 특징을 **스케일 업 (scale up)**하고, 점프를 예측하는 데 덜 중요한 특징은 스케일 다운 (scale down)할 수 있습니다. **시프트 (Shift)**는 정규화된 특징들을 액션에 의존적인 다른 베이스라인(baseline)으로 이동시킵니다. 마지막으로, **게이트 (gate)**는 어텐션 또는 MLP 업데이트가 예측된 상태에 얼마나 강하게 영향을 미칠지를 결정합니다.

이러한 제어 값들은 어텐션 브랜치와 MLP 브랜치에 대해 각각 별도로 생성되어, 블록당 총 6개의 값(각 브랜치당 시프트, 스케일, 게이트 하나씩)을 제공합니다.

"제로 (Zero)"라는 명칭은 이들의 가중치가 0에서 시작한다는 것을 의미하며, 따라서 예측기는 무작위적인 액션 효과 없이 시작하여 훈련 과정에서 어떤 게이트를 열어야 할지를 점진적으로 학습합니다.

이제 훈련 과정에서 6개의 트랜스포머 블록을 거친 후, 작은 프로젝션 헤드 (projection head)가 세 개의 예측된 미래 잠재 변수를 생성합니다:

이 값들은 세 개의 실제 다음 프레임(real next frames)에서 생성된 잠재 변수(latents)와 비교됩니다:

이제 우리는 이 손실(loss)을 0으로 낮추고 싶습니다... 하지만 모델이 속임수를 쓸 수 있는 아주 쉬운 방법이 하나 있습니다. 바로 모든 잠재 벡터(latent vectors)를 동일하게 만드는 것입니다! 그렇게 되면 마리오, 파이프, 새로운 월드가 모두 동일하게 보이기 때문에 예측이 완벽해질 것입니다 (표현 붕괴 (representation collapse)).

따라서 이러한 속임수를 방지하기 위해, 우리는 SIGReg1을 사용합니다! 이는 실제 프레임의 잠재 변수가 다양하고 유익한 상태를 유지하도록 장려함으로써 이러한 붕괴를 방지합니다. 따라서 우리의 새로운 손실 함수는 다음과 같습니다:

이것이 전체 아키텍처(architecture)입니다! 이제 실제 결과로 넘어가 보겠습니다.

하지만 실제로 학습했을까요?

LeMario는 32개의 마리오 레벨에 걸친 280개 에피소드, 총 737,134개의 프레임을 통해 학습되었습니다. 모델이 학습했는지 확인하기 위해 단순히 낮은 손실 값만을 확인할 수는 없었습니다. 낮은 손실만으로는 모델이 역학(dynamics)을 학습했음을 증명하기에 충분하지 않았습니다. 인접한 프레임들은 종종 매우 유사하기 때문에, "변화가 없다"라고 예측하는 것이 강력한 베이스라인(baseline)이 되기 때문입니다.

홀드아웃 에피소드(held-out episodes)에서, 저는 LeMario를 해당 지속성 베이스라인(persistence baseline) 및 셔플된 행동(shuffled actions)과 결합된 실제 프레임 이력과 비교했습니다:

행동을 셔플(shuffling)했을 때 1단계 오차는 20.2% 증가했습니다. 5단계의 재귀적 단계(recursive steps) 전체에 걸쳐, LeMario는 지속성 베이스라인보다 45.5% 우수한 성능을 보였으며, 셔플된 행동은 47.5% 더 나쁜 성능을 보였습니다. 예측 단계가 멀어질수록 버튼 조작의 중요성이 더 커졌습니다.

LeMario는 플레이어의 행동에 조건화된(conditioned on) 단기 마리오 역학(short-horizon Mario dynamics)을 학습했습니다!

컨트롤러를 만지게 하기

이제 재미있는 부분입니다. 모델이 미래를 상상할 수 있게 되면, 우리는 그 미래들을 탐색하여 마리오가 무엇을 해야 할지 선택하게 할 수 있습니다.

상상을 통한 탐색

행동이 주어졌을 때 단 한 프레임 앞만을 예측하는 모델을, 미래의 여러 행동과 단계를 예측할 수 있는 무언가로 바꾸기 위해, 저는 교차 엔트로피 방법 (Cross-Entropy Method, CEM)을 사용합니다!

현재 이미지와 목표 이미지가 주어지면, 인코더(encoder)는 $z_t$와 $z_{goal}$을 생성합니다. 그러면 CEM은 다음과 같이 수행합니다:

• 수백 개의 행동 시퀀스(action sequences)를 샘플링합니다.
• 각 시퀀스를 LeMario를 통해 앞으로 전개(roll forward)시킵니다.
• 예측된 최종 잠재 상태(latent)를 $z_{goal}$과 비교하여 점수를 매깁니다.
• 가장 좋은 후보들을 유지합니다.
• 그 주변에서 다시 샘플링(resample)하고 이 과정을 반복합니다.

CEM은 무작위 후보들보다 예측된 목표 거리(goal distance)가 훨씬 짧은 행동 시퀀스들을 찾아냈습니다. 저에게는 상상할 수 있는 모델, 그 상상 속을 탐색할 수 있는 최적화 도구(optimizer), 그리고 보상 설계(reward engineering)가 필요 없는 목표 이미지(goal image)가 있었습니다. 정말 장관이었습니다.

하지만 마리오는 거의 움직이지 못했습니다

저는 아주 작은 목표로 시작했습니다. 마리오는 x=40에서 시작했고, 목표 프레임은 그가 x=72에 있는 모습을 보여주었습니다. 순수한 JEPA+CEM 방식은 x=44에서 끝났습니다.

비전문가적인 표현을 빌리자면, 형편없었습니다.

이 시점에서 저는 어느 부분이 실패했는지 알 수 없었습니다. 예측기(predictor)가 틀렸을 수도 있고, CEM이 고장 났을 수도 있으며, 혹은 인코더(encoder)가 마리오를 완전히 무시했을 수도 있었습니다. 저는 가장 단순한 질문부터 시작해야 했습니다. 과연 그 192개의 숫자에 마리오의 위치 정보가 포함되어 있기는 한 걸까요?

192개의 숫자 안에는 무엇이 들어있을까?

잠재 상태(latent)는 단 192개의 숫자로만 이루어져 있습니다. 인코더를 변경하지 않고 그 안에 어떤 정보가 들어있는지 물을 수 있는 방법이 필요했습니다.

마리오를 잊어버린 걸까?

저는 JEPA를 동결(freeze)하고, 잠재 상태로부터 마리오의 에뮬레이터 좌표를 복구하기 위한 작은 프로브(probe)2를 학습시켰습니다. 인코더를 변경할 수 없었기 때문에, 프로브가 복구한 모든 위치 정보는 JEPA가 이미 학습한 정보여야만 했습니다.

가로 위치(horizontal position): MAE = 9.30 px, R² = 0.997
세로 위치(vertical position): MAE = 21.62 px, R² = 0.188

프로브가 작동했습니다! 마리오의 가로 위치는 거의 완벽하게 복구 가능했습니다. 세로 상태는 훨씬 약했지만, 인코더가 플레이어에 대한 유용한 정보를 분명히 학습했다는 사실은 확인되었습니다.

모든 것을 "해결한" 프로브

저는 프로브가 예측한 가로 위치를 사용하여 CEM이 상상한 미래들에 대해 일시적으로 점수를 매겼습니다. LeMario는 여전히 미래를 상상했고 CEM은 여전히 행동을 선택했지만, 프로브는 그 미래들이 순위가 매겨지는 방식만을 바꾸었습니다.

x=72를 목표로 했을 때, 프로브 점수가 매겨진 CEM은 마리오를 x=40에서 x=71로 이동시켰습니다. 국소적 재계획(local replanning)을 통해, 나중에는 x=177이 목표일 때 x=176에 도달했습니다.

이것이 성공한 첫 번째 롤아웃(rollout)이었습니다! JEPA 모델은 유용한 수평 이동을 상상할 수 있었고, 프로브(probe)는 이를 찾아낼 수 있었습니다.

이를 통해 잠재 계획(latent planning)이 이론적으로 작동할 수 있다는 확신을 얻었습니다. 만약 Mario의 위치가 이미 표현(representation) 내부에 있었고 학습된 역학(dynamics)이 이를 이동시킬 수 있다면, 아마도 첫 번째 목표가 시작 프레임과 너무 유사했던 것이 문제였을 것입니다.

레벨의 절반 정도 떨어진 목표 시도하기

플래너(planner)의 역할은 두 임베딩(embedding)을 연결하는 행동(action)을 찾는 것이었습니다. 만약 시작 프레임과 근처의 목표가 이미 유사한 임베딩을 가지고 있다면, 거의 아무것도 하지 않는 것이 성공처럼 보일 수 있습니다. 그래서 저는 지도 학습 기반의 프로브(supervised probe)를 제거하고, Worlds 1-1, 2-1, 3-1의 대략 중간 지점에 있는 도달 가능한 목표 이미지들을 사용하여 순수 잠재 계획(raw latent planning)을 다시 시도했습니다.

레벨 3-1의 사례를 들어보겠습니다. 이번에는 Mario가 훨씬 더 멀리 이동했습니다. x=40에서 x=44로 이동하는 대신, 세 번의 실행 결과 대략 x=290–307에 도달했습니다. 여전히 첫 번째 의미 있는 장애물 근처에서 사망했지만, 순수 잠재 계획이 더 이상 아무것도 하지 않는 상태는 아니었습니다.

CEM은 제가 요청한 대로 정확히 수행했습니다: 두 임베딩을 연결하는 행동을 탐색하는 것 말입니다. 더 먼 목표는 잠재 공간(latent space)에 충분한 압력을 가하여 Mario를 움직이게 만든 것으로 보입니다.

하지만 Mario는 여전히 목표로부터 1,442 월드 픽셀(world pixels) 떨어져 있었던 반면, 인코더(encoder)는 그의 최종 장면을 단 0.164의 잠재 거리(latent distance)로 할당했습니다. CEM은 0.153을 예측했으므로, 예측기(predictor)가 심하게 환각(hallucination)을 일으키고 있는 것은 아니었습니다. 인코더 자체가 잘못된 장면을 상당히 가깝다고 판단한 것입니다(위의 시작 및 목표 프레임에서 볼 수 있듯이).

Mario의 스크롤 카메라가 그 이유를 설명해 줍니다. 게임 내에서 두 개의 서로 다른 위치임에도 불구하고, 두 멀리 떨어진 위치가 매우 유사해 보일 수 있습니다! 따라서 모델은 사실 자신의 역할을 수행하고 있었습니다. 모델은

그래서 저는 만약 체크포인트(checkpoint)를 더 작은 단위로 나눈다면 프레임 간의 차이가 더 커질 것이라는 가설을 세웠고, 인간의 플레이를 중간 이미지 목표(intermediate image goals)로 나누었습니다. 이것은 도움이 되었으며, 가공되지 않은 잠재 계획(raw latent planning)은 이 프로젝트에서 탐사(probe) 없이 달성한 가장 큰 진전인 x=314에 도달했습니다.

하지만 여전히 더 짧은 목표들이 표현(representation)을 의미 있게 수정하지는 못했습니다.

Mario는 첫 번째 목표에 2픽셀 이내의 오차로 도달했습니다. 두 번째 목표의 경우, x=283까지 초과 달성한 후, 뒤로 수정하여 기준점으로부터 5픽셀 떨어진 x=239에서 멈췄습니다.

시각적으로 Mario는 올바른 위치에 도달했습니다. 하지만 벤치마크는 최종 임베딩(embedding)이 목표 임베딩과 충분히 가깝지 않았기 때문에(아마도 HUD나 다른 작은 세부 사항 때문일 것입니다) 해당 체크포인트를 실패로 표시했습니다.

이는 거리가 문제의 일부일 뿐이라는 것을 보여주었습니다. 더 작은 목표들은 Mario가 더 멀리 이동하는 데 도움을 주었지만, 잠재 거리(latent distance)는 그가 도착했는지 여부를 결정하는 취약한(brittle) 방식이었습니다.

다음 목표는 Mario가 점프할 것을 요구했고, 플래너(planner)는 다시 실패했습니다. 이는 이전의 탐사(probe) 결과와 일치했습니다. 수평 위치는 잘 표현되었지만, 수직 위치는 훨씬 더 약했습니다.

더 작은 목표들은 계획 지평(planning horizon)에는 도움이 되었습니다. 하지만 잠재 거리를 진전의 신뢰할 수 있는 척도로 만들지는 못했습니다.

이 시점에서 실패들은 더 이상 서로 무관해 보이지 않았습니다. 그것들은 세 가지 주요 문제를 가리키고 있었습니다.

예측 상태(Predictive state)는 제어 상태(control state)가 아니다

인코더(encoder)는 미래의 이미지를 예측하는 데 도움이 되는 무엇이든 표현하도록 보상받습니다. 카메라 위치, 적의 단계(enemy phase), 애니메이션, 타이머 상태 등이 모두 유용할 수 있습니다.

컨트롤러(controller)에는 다른 것, 즉 거리가 제어 가능한 진전과 일치하는 상태가 필요합니다.

CEM은 모델의 약점을 탐색한다

CEM은 제가 요청한 대로, 즉 LeMario가 목표에 도달할 것이라고 믿는 행동을 정확히 찾아냈습니다. 하지만 모델이 틀렸을 때는 이를 인식하지 못했습니다.

이로 인해 모델의 약점이 명확해졌습니다. 모델은 시각적으로 유사한 위치를 동일한 장소로 취급했고, 수직 이동에 대해 추론하는 데 어려움을 겪었습니다.

Mario는 Push-T의 가정을 변화시켰다