JEPA 월드 모델로 슈퍼마리오 학습, 가까운 목표는 성공하지만 먼 목표는 실패

latent space에서 플래닝하는 게 문제인 듯. latent가 예측 위주로 학습되다 보니 실제 위치 정보를 정확히 못 담아서, 옵티마이저가 특정 feature만 보고 위치를 스킵할 수 있음.

(저자) 맞말임. latent가 제어보다 예측에 최적화돼 있어서 플래닝이 노이즈가 심함. 그래도 A100 한 대로 2시간도 안 학습시켰는데 수평 위치를 이렇게 잘 이해하는 거 보고 놀랐음.

JEPA는 예측 가능한 feature랑 노이즈는 잘 구분하는데, 예측 가능한 feature들 사이 중요도를 매기는 방법이 없음. HUD의 승리 조건 픽셀 몇 개랑 화면 대부분이 일치하는 장면 중 뭘 골라야 하는지 모델은 구분 못 함. linear probe로 X 위치만 따로 가중치 주는 식으로 보완한 거 같음.

이 예시가 장기 플래닝 문제를 잘 보여줌. 중간 목표로 쪼개는 것부터 이미 특정 해결 전략을 강제하는 셈이고, 목표를 이미지로 써야 하는 것도 이상함. 같은 목표를 만족하는 상태가 여러 개일 수 있는데 이미지 하나로는 그걸 다 못 담음.

(저자) 그 말이 딱 맞음. 중간 이미지 목표는 진단용으로 넣은 거지 일반적인 해법은 아니었음. 체크포인트 덕분에 마리오가 더 멀리 가긴 했지만 결국 해법을 강제로 주입한 거라 근본적인 representation 문제는 못 고쳤음. reward-free 유지하려면 목표를 reward나 value function 없이 표현해야 하는데 그게 어려움.

목표는 특정 위치가 아니라 액션이라고 봐야 함. 계속 오른쪽 방향키를 누르고 있는 상태, 혹은 그 상태로 돌아가기 위한 중간 액션이 진짜 목표인 거임.

나는 시스템 전체를 대상으로 JEPA 비슷한 모델 학습시키는 중임. BPF로 프로세스, 메모리, cpu, 네트워크, 스토리지, 화면 버퍼 다 모아서 Vulkan 기반 커스텀 브라우저로 3D 공간에 재현하는 프로젝트인데 기본 컨셉은 잘 작동함.

JEPA는 결국 일반화된 encoder일 뿐이니까 거기서 나온 latent를 유저 데이터나 RL policy로 학습한 transformer에 먹여야 할 것 같음. 근데 수직 위치도 잘 예측 못 했다니까 이것도 잘 안 될 수도 있겠음.