JEPA 월드모델이 테스트 시점에 바뀐 동역학을 스스로 적응한다

무엇인가

월드모델은 관측과 행동으로부터 미래 상태를 예측해 에이전트가 계획을 세우도록 돕지만, 학습 시점과 테스트 시점의 전이 동역학이 다르면 예측이 체계적으로 틀어진다. 논문이 다루는 문제는 정확히 이 지점이다. 액추에이터 응답, 접촉 동역학, 물체 운동이 바뀌면 학습 동역학으로 훈련된 모델은 잘못된 후보 행동 순위를 매기고 계획 성능이 무너진다. 저자들은 상태 공간, 행동 공간, 관측 매핑, 점수 함수, 초기 상태 분포, 에피소드 길이는 그대로이고 전이 동역학 P_test만 P_train과 달라지는 상황을 설정한다. 목표 이미지도, 테스트 시점의 환경 보상도 주어지지 않는다.

어떻게 동작하나

JEPA-TTT는 사전학습된 action-conditioned JEPA 월드모델에서 잠재 동역학 예측기 F_θ만 테스트 시점 내내 업데이트한다. 여기서 θ는 행동 인코더, 자기회귀 예측기, 예측 프로젝션을 포함한다. 시각 인코더 E_φ와 오프라인에서 학습한 보상 헤드 R_ω는 고정된다. 보상 헤드는 배포 전에 학습 동역학에서 수집한 (관측, 점수) 쌍으로 평균제곱오차를 최소화해 적합하며, 인코더를 고정하는 이유는 테스트 시점 학습 목표에 안정적인 표현 공간을 주기 위해서고 보상 헤드를 고정하는 이유는 동역학 적응과 과제 목적 적응을 분리하기 위해서다. 학습 손실은 실행한 행동 블록으로 예측기를 굴린 뒤 고정 인코더가 만든 다음 잠재 표현과 일치시키는 단순한 L2 손실이며, 타깃에는 stop-gradient가 걸린다. 완전한 미니배치가 준비된 재계획 단계마다 AdamW 업데이트를 최대 한 번 수행한다. 계획은 CEM 기반 모델예측제어로, 적응된 예측기와 고정 보상 헤드로 후보 행동 블록 시퀀스를 잠재 공간에서 자기회귀 롤아웃해 할인된 예측 리턴을 최대화하고 첫 행동 블록만 실행한 뒤 다시 계획한다.

무엇과 다른가

적응은 에피소드 경계를 넘어 지속된다. 예측기 파라미터, 옵티마이저 상태, 온라인 버퍼가 다음 에피소드로 그대로 이어진다. 온라인 경험을 학습에 쓰기 위해 dense replay를 도입하는데, 모든 시간 오프셋에서 예측 윈도우를 만들어 계속 커지는 버퍼에 보관하고 미니배치를 균일 샘플링한다. 각 윈도우는 C+1개의 관측과 그 사이 C개의 행동 블록으로 구성되며, 프레임 스킵 K 때문에 dense 구성은 액션 블록 경계에서만 시작하는 sparse 구성보다 약 K배 많은 윈도우를 만든다. 저자들은 sparse, 연산량을 맞춘 sparse, dense 스트림, dense replay를 순차 비교해 추가 최적화, 시간적 커버리지, 리플레이의 효과를 분리한다.

어떻게 쓰나

실험은 PushT, Two-Room, DeepMind Control Suite의 Reacher, OGBench-Cube 네 개 연속제어 환경에서 진행됐고, 각 환경마다 관측 매핑과 점수, 행동 공간, 에피소드 길이를 유지한 채 두 개의 동역학 변화를 설계해 총 8개 변화를 만든다. 변화는 액추에이터 매핑, 접촉 응답, 상태 의존 제어 변환, 축 간 결합, 지연된 구동을 포함한다. 과제마다 학습 환경에서 3,000개 완전 에피소드를 모아 LeWorldModel 기반 action-conditioned JEPA 월드모델을 사전학습하고 인코더를 고정한 뒤 보상 헤드를 적합한다. 배포 한 번은 고정된 테스트 동역학 아래 500 에피소드로 구성되고, 0, 50, …, 500 에피소드 시점에 동일한 100개 홀드아웃 에피소드로 평가하며 서로 다른 시드로 3회 실행해 평균과 표준편차를 보고한다. 결과적으로 JEPA-TTT는 8개 변화 모두에서 계획 성능을 개선했고, 고정 JEPA 월드모델 대비 best score는 153%, mean AUC는 113% 향상됐다. 평균 정규화 홀드아웃 계획점수 AUC는 0.571로 고정 모델의 0.267을 크게 앞섰고, 평균 best 홀드아웃 점수는 0.267에서 0.678로 올랐다. 500 에피소드 후 자기회귀 잠재 예측 오차는 평균 83% 줄었다. 비교 대상은 고정 JEPA, 오라클 환경 보상으로 적응하는 PPO-TTT, 에피소드마다 모델과 버퍼를 초기화하는 동시대 연구 AdaJEPA(목표 이미지 없는 설정으로 이식)였고, JEPA-TTT는 모든 변화에서 두 계획 지표 모두 세 베이스라인을 앞섰다. 참고로 테스트 동역학으로 직접 학습한 데이터 풍부 오프라인 기준과 비교하면 8개 변화 중 6개에서 500 에피소드만으로 그에 근접하거나 약간 능가했다.

전제와 한계

어블레이션은 dense replay의 세 요인을 분리한다. 연산량을 맞춘 sparse 스트림은 sparse보다 best score와 mean AUC 모두 개선돼 추가 최적화의 이득을 보였고, 같은 예산에서 dense 스트림은 두 지표를 더 올려 시간적 커버리지의 이득을 시사했다. dense replay가 네 규칙 중 가장 높은 집계 성능을 냈으며 dense 스트림 대비 best score는 0.678 대 0.671로 비슷하면서 mean AUC는 0.571 대 0.548로 더 좋았다. 에피소드 내부 윈도우만 쓰는 통제 실험에서 예측기와 옵티마이저 상태를 에피소드 간 유지하면 같은 업데이트 횟수로 best score가 0.289에서 0.729, mean AUC가 0.273에서 0.637로 뛰었다.

실무 관점에서 이 논문이 주는 신호는 명확하다. 시뮬레이터에서 학습한 월드모델을 실제 로봇이나 제어 환경에 올릴 때, 전체 모델을 다시 학습하거나 보상 신호를 새로 설계하지 않고도 예측기만 온라인으로 적응시켜 계획 성능을 회복할 수 있다는 구성이 구체적인 수식과 절차로 제시된다. 특히 인코더와 보상 헤드를 고정해 표현과 과제 목적이 흔들리지 않게 하고, 적응 상태를 에피소드 간에 누적하는 설계는 배포 파이프라인에 그대로 옮길 수 있는 형태다. 다만 dense replay는 버퍼가 계속 커지는 구조이므로 메모리와 업데이트 예산을 실제 시스템에서 어떻게 관리할지, 그리고 보상 헤드가 학습 동역학에서만 적합됐다는 전제가 자신의 환경에서도 성립하는지를 먼저 확인해야 한다.

저자들이 밝힌 한계도 분명하다. 이 연구는 테스트 시점 내내 고정된 단일 동역학 변화만 다루며, 시각 관측 매핑이나 과제 점수, 행동 인터페이스, 에피소드 길이는 변하지 않는다. 인코더와 보상 헤드를 고정하기 때문에 시각적 변화나 보상 변화가 동시에 일어나는 상황은 설계 범위 밖이다. 또한 연속적인 동역학 변화나 적응 후 학습 동역학으로 되돌아가는 경우를 연구하지 않았고, 지속적인 예측기 업데이트가 이전에 학습한 동역학(학습 동역학 포함)의 예측을 악화시킬 수 있다는 점을 인정한다. 따라서 망각, 순방향·역방향 전이, 연속학습에서의 성능 유지 여부는 측정되지 않았다.