LeMario: 슈퍼 마리오 브라더스에서 JEPA 세계 모델 훈련 — 기술적 워크스루 및 사후 분석

Benjamin Bai의 LeMario 프로젝트는 LeWorldModel — 소규모 JEPA(Joint-Embedding Predictive Architecture) — 을 처음부터 재현하여 Super Mario Bros에서 훈련합니다. 상세한 기술 워크스루는 전체 아키텍처, 훈련 설정, 그리고 작동한 것과 작동하지 않은 것에 대한 솔직한 포스트모템을 다룹니다.

아키텍처 개요

모델은 연속된 네 개의 마리오 프레임을 처리합니다. 비전 인코더는 각 프레임을 192차원 잠재 벡터로 압축합니다: z_t = E_θ(x_t), z_t ∈ R^192. 액션(다섯 에뮬레이터 프레임 동안 Left, Right, Up, Down, A, B 버튼 상태)은 액션 인코더를 통해 별도로 192차원 벡터로 인코딩됩니다.

프레임과 액션 잠재는 6개의 트랜스포머 블록을 가진 인과적 예측기에 공급됩니다. 액션은 Adaptive LayerNorm Zero(AdaLN-Zero)를 통해 주입되며, 이는 어텐션과 MLP 각 분기별로 shift, scale, gate 제어를 생성합니다. Gate는 업데이트가 예측된 상태에 얼마나 강하게 영향을 미치는지 제어합니다. 가중치는 0에서 시작하여 예측기가 점진적으로 학습합니다.

예측기는 세 개의 예측된 미래 잠재를 출력합니다: ẑ₁, ẑ₂, ẑ₃. 이들은 세 개의 실제 다음 프레임의 잠재와 MSE 손실을 사용하여 비교됩니다: L_pred = MSE([ẑ₁,ẑ₂,ẑ₃], [z₁,z₂,z₃]). 표현 붕괴를 방지하기 위해 SIGReg 정규화가 추가됩니다: L = L_pred + 0.1 * L_SIGReg.

작동한 것

• 모델이 보이지 않는 에피소드에 일반화되었습니다.
• 액션 정보를 효과적으로 사용했습니다.
• 강력한 기준선보다 5단계 미래를 더 잘 예측했습니다.
• 원시 보상 없는 계획으로 마리오를 가까운 이미지 목표까지 2-5픽셀 이내로 이동시킬 수 있었습니다.

실패한 것

• 목표가 레벨의 더 먼 곳에 배치되면 마리오가 첫 번째 주요 장애물을 안정적으로 점프하지 못했습니다.
• 단일 먼 목표 이미지를 향해 탐색할 수 없었습니다.
• 결론: 게임 예측 ≠ 진행 학습. 모델은 픽셀 수준의 역학을 학습했지만 장기 계획이나 목표 지향 행동은 학습하지 못했습니다.

주요 교훈

Bai는 대부분의 교훈이 나중에 보면 당연해 보인다고 언급합니다. 단기 예측과 장기 계획 사이의 격차는 중요하며, 표준 예측 손실로는 포착되지 않습니다. 게시물에서는 이 한계를 점진적으로 드러낸 실험들을 자세히 설명합니다.