world2vec는 보상 없이 관측만으로 해석 가능한 세계 모델을 학습한다

무엇인가

세계 모델(world model)은 제어 성능만 좋으면 환경의 동역학을 충실히 담지 못한 채로도 높은 점수를 낼 수 있다. 논문은 이 점을 문제로 삼는다. 성능이 좋아도 내부가 실제 시스템과 대응하지 않으면, 모델을 읽어서 환경의 방정식을 알아내는 기계론적 해석(mechanistic interpretability)이나 과학적 발견에는 쓸 수 없다는 것이다. world2vec는 보상 신호 없이 관측과 행동만으로, 긴 예측 지평에서도 실제 시스템을 재현하는 잠재 동역학을 학습하는 것을 목표로 한다.

어떻게 동작하나

문제는 다음과 같이 정식화된다. 상태 s가 행동 a에 따라 ṡ = f(s, a)로 진화하고, 우리는 이 상태를 비선형 혼합 함수 x = g(s)를 통해 관측한다. 여기서 s도 f도 직접 볼 수 없다. 목표는 관측을 잠재 상태 z = h(x)로 보내는 인코더 h와, 잠재 공간에서의 변화율 ż = f̂(z, a)를 예측하는 동역학 모델을 관측·행동만으로 동시에 추정하는 것이다. 논문이 세우는 충실성 조건은 명확하다. 인코딩된 상태를 잠재 공간에서 적분해 나간 결과가 미래 시점에 관측을 인코딩한 것과 같아야 한다. 즉 z(t) + ∫f̂ dτ = h(x(t+T))가 모든 t, T에 대해 성립해야 한다. 중요한 점은 이 조건의 양변이 모두 관측과 행동만으로 계산된다는 것이다.

무엇과 다른가

실제 학습 절차는 이산 시점에서 이뤄진다. 부분 관측 가능성 때문에 프레임 한 장으로는 속도를 알 수 없으므로, 인코더는 길이 H의 짧은 관측 이력을 입력받아 잠재 상태를 만든다. 학습 시에는 시작점 n을 뽑아 K스텝 롤아웃을 만들고, 각 스텝의 잠재 예측과 실제 미래 관측의 인코딩 사이 제곱거리를 가중합해 최소화한다. 적분은 명시적 오일러 스텝 Φ(z, a) = z + Δt·f̂(z, a)로 구현된다. 여기에 상수 인코더처럼 자명한 해로 붕괴하는 것을 막는 정규화 항이 붙는다. 논문은 특정 목적함수에 종속되지 않다고 밝히며, 대조학습 방식의 DCL을 쓴 world2vec-dcl과 임베딩을 등방성 가우시안으로 정규화하는 SIGReg를 쓴 world2vec-sigreg 두 가지 인스턴스를 제시한다.

어떻게 쓰나

저자들이 강조하는 원리는 잠재 동역학 매칭(latent dynamics matching)이다. 한 스텝만 보면 이 목적함수는 유한차분으로 근사한 속도 $(finite-difference velocity)와 학습된 속도장을 맞추는 것과 같아진다(Neural ODE와 같은 발상). 여러 스텝으로 늘리면 모델이 예측 궤적을 따라 누적한 속도와 인코딩된 실제 변위를 맞추게 되므로, 작은 한 스텝 오차가 긴 롤아웃에서 어떻게 증폭되는지가 학습 신호에 그대로 드러난다. 동역학 함수 f̂의 형태는 자유롭게 고를 수 있는데, 제어 실험에서는 잠재 상태와 행동을 변화율로 보내는 MLP를 쓰고, 해석 가능한 파라미터화가 필요할 때는 SINDy·DYSCO 계열의 구조화 모델 f̂(z, a) = Ξθ(z, a)를 쓴다. θ는 고정된 라이브러리 특징(다항식)이고 Ξ가 학습되는 계수이며 L1 페널티를 건다.

전제와 한계

표현이 실제로 식별 가능한지도 정량화한다. 이론적으로 이 계열 목적함수는 이상적 조건에서 affine 변환까지 일관된 표현을 준다고 알려져 있는데, 논문은 이를 경험적으로 검증한다. 독립적으로 학습한 두 인코더를 최적 affine 정렬한 뒤의 R²를 일관성(consistency)으로 보고, 같은 인코더-동역학 쌍에서 인코딩된 미래와 롤아웃 예측의 R²를 LRC(Latent Rollout Consistency)로 측정한다. 다만 LRC는 자명한 임베딩으로도 높게 나올 수 있어, 부정 샘플 512개 중에서 인코딩된 미래를 top-1으로 찾아내는 정확도 EQA(Equivariance Accuracy)를 함께 쓰고 하이퍼파라미터도 검증 EQA로 고른다.

실험은 MuJoCo 기반 DeepMind Control Suite의 Pendulum, Reacher, Cartpole에서 이뤄졌다. 각 환경에서 랜덤 정책으로 64×64 픽셀 RGB 프레임 20만 장을 모았고(Reacher는 균등 랜덤, Pendulum과 Cartpole은 시간 상관이 있는 Ornstein-Uhlenbeck 행동), 식별 가능성 분석을 위해 위치·각도만 무작위 주입식 MLP로 50차원에 투영한 비시각적 비선형 관측 변형 데이터셋도 만들었다. 인코더는 16×16 패치, 폭 64의 2층·4헤드 비전 트랜스포머에 키포인트 판독과 MLP 헤드를 붙였고, 동역학은 128유닛 3층 MLP, 잠재 차원 32로 전체 26만 3천 파라미터(인코더 22만 1천, 동역학 4만 2천)다. 이는 DreamerV3, TD-MPC2, LeWM 대비 최소 한 자릿수 이상 적은 규모다.

제어는 보상 모델도 상태 디코더도 없이 잠재 공간 MPC로 한다. 목표 프레임 시퀀스를 인코딩해 목표 잠재 z*를 만들고, cross-entropy method(CEM)로 계획 지평 P만큼의 행동 후보를 샘플링해 목표까지의 가중 제곱거리를 최소화하는 계획을 고른 뒤 앞부분만 실행하고 다시 계획한다. 연속시간 동역학이므로 학습과 테스트의 샘플링 레이트를 다르게 쓸 수 있다는 점도 실무적으로 유용하다. TD-MPC2와 DreamerV3는 각자의 온라인 프로토콜로, LeWM과 world2vec는 랜덤 정책 데이터로 오프라인 학습했고 동일한 홀드아웃 에피소드에서 평가했다. 논문이 보고한 결과는 세 환경 중 두 곳에서 MPC 성공이며, 기존 최고 JEPA 계열인 LeWM은 어느 환경도 풀지 못했다는 것이다. 초록 기준으로 71개 어트랙터 벤치마크에서의 시스템 식별도 함께 달성했다고 주장한다. 학습 비용은 A100 한 장에서 1스텝 모델 2~6분, 20스텝 약 10분, 40스텝 약 30분이고, 폐루프 제어는 시뮬레이션·렌더링 포함 스텝당 20~50ms다. 다만 이 글에서 확인한 원문 범위에는 환경별 누적 보상이나 71개 어트랙터별 정확도가 표의 숫자로 제시되어 있지 않다.

개발자 관점에서 이 논문의 쓸모는 두 가지다. 첫째, 월드 모델을 제어 성능 하나로 평가하지 않고 롤아웃 충실도(LRC)와 표현 식별성(EQA, 일관성)으로 함께 검증하는 프레임을 제공한다. 시뮬레이터 대체나 모델 기반 로봇 제어를 붙일 때, 보상은 잘 나오는데 예측이 무너지는 모델을 걸러내는 데 바로 쓸 수 있는 지표다. 둘째, 보상 없이 랜덤 정책 데이터만으로 학습하고 목표 이미지만 주어 MPC를 돌리는 구성이라, 보상 함수를 설계하기 어려운 환경에서 출발점이 된다. 동역학을 MLP 대신 SINDy식 구조화 모델로 바꾸면 계수를 읽어 방정식 형태를 확인할 수 있다는 점도 시뮬레이터 검증 작업에 유용하다.

한계와 전제도 분명하다. 식별 가능성 보장은 이상화된 가정 아래 affine 변환까지의 일관성이며, 논문 스스로 이를 실제 수치 추정에서 확보하는 것이 도전이라고 말한다. 부분 관측 가능성을 전제로 하므로 상태를 복원할 수 있을 만큼의 관측 이력 H가 필요하고, 정규화 항이 없으면 상수 인코더 같은 자명해로 붕괴한다. 실험은 64×64 픽셀의 소규모 환경 세 개에 한정되며, 그중 한 환경에서는 MPC가 성공하지 못했다. 구조화 동역학과 71개 어트랙터 시스템 식별 결과는 원문 후반부에서 다뤄지므로, 실제 적용 전에는 그 부분의 수치와 가정을 직접 확인해야 한다.