JEPA-TTT, 배포 후에도 세계 모델의 동역학을 적응시키다

배경

세계 모델은 환경의 미래 상태를 예측하고, 그 결과를 이용해 행동을 계획한다. 그러나 학습 이후 마찰, 질량, 제어기 반응처럼 환경의 동역학이 달라지면 기존 예측은 빠르게 부정확해질 수 있다. JEPA-TTT는 이 문제를 해결하기 위해 사전 학습된 잠재 세계 모델을 테스트 단계에서도 계속 학습시키는 접근을 제시한다.

작동 방식

JEPA-TTT는 행동 조건부 JEPA(Joint-Embedding Predictive Architecture) 세계 모델을 기반으로 한다. 전체 모델을 다시 학습하는 대신, 환경 변화와 직접 관련된 구성 요소만 선택적으로 업데이트한다.

• 잠재 동역학 예측기만 업데이트: 시각 인코더는 고정해 테스트 데이터가 사전 학습된 시각 표현을 훼손하는 것을 막는다. 보상 헤드도 고정하므로 기존 과제 목표를 유지할 수 있다.
• 자기지도 신호로 적응: 에이전트가 실제로 관측한 이후 상태를 사용해 예측을 보정한다. 따라서 온라인 환경 보상이나 목표 이미지를 계획에 요구하지 않는다.
• 조밀한 리플레이 적용: 궤적에서 일부 샘플만 추출하지 않고, 모든 시간 오프셋에서 예측 윈도우를 만든다. 이를 계속 커지는 버퍼에 저장한 뒤 미니배치로 샘플링해 업데이트한다.
• 에피소드 사이에서도 학습 지속: 각 에피소드가 끝난 뒤 적응을 초기화하지 않고, 테스트 중 축적된 경험을 다음 업데이트에 활용한다.

이 구조는 안정적으로 유지해야 하는 지식과 새 환경에 맞춰 바꿔야 하는 예측 기능을 분리한다. 모델은 기존 시각 표현과 과제 목적을 보존하면서 잠재 공간의 미래 예측만 새로운 동역학에 맞게 조정할 수 있다. 전체 영상을 직접 재구성할 필요가 없다는 점도 특징이다.

결과와 영향

연구진은 4개 연속 제어 환경에서 8가지 동역학 변화를 평가했다. JEPA-TTT는 모든 변화 조건에서 업데이트하지 않은 JEPA 세계 모델보다 나은 계획 성능을 보였다. 테스트 에피소드 500회 후 잠재 자기회귀 예측 오차는 평균 83% 감소했고, 계획 성능은 고정 기준 모델보다 153% 향상됐다.

이 결과는 세계 모델을 한 번 학습한 뒤 고정해 사용하는 방식에서, 실제 운용 중 지속적으로 보정하는 방식으로 확장할 가능성을 보여준다. 온라인 보상이 없어도 관측 상태와 예측의 차이를 적응 신호로 활용할 수 있기 때문이다. 다만 오래된 경험과 최근 경험의 균형, 환경이 계속 변할 때의 망각 방지, 다양한 변화에 대한 안정성은 여전히 검증해야 한다. 이번 실험은 연속 제어 환경에 초점을 두었으므로 복잡한 시각 작업과 개방형 환경에서의 일반화는 추가 연구가 필요하다.

출처: Hugging Face Daily Papers