행동을 구분하는 세계 모델 AD-WM은 반사실적 MPC를 위한 것이다.

무엇인가

이 논문이 겨냥하는 문제는 잠재 세계 모델의 학습 목표와 모델 예측 제어(MPC)가 실제로 요구하는 능력 사이의 불일치다. 잠재 세계 모델은 보통 데이터에 기록된 전이, 즉 실제로 취한 행동의 결과를 예측하도록 학습된다. 반면 MPC는 같은 상태에서 여러 후보 행동을 놓고 무엇이 더 나은지 비교해야 한다. 논문은 전이 z_{t+1} = z_t + δ_t에서 증가분 δ_t가 작으면, 행동을 아예 무시하는 예측기 F0(z,a)=z도 기대오차 E‖δ_t‖²밖에 내지 않는다는 점을 지적한다. 오차는 낮지만 모든 후보 행동에 같은 종료 비용을 매기므로 행동 선택에는 쓸모가 없다. 시각 표현이 장면의 지속적 내용에 지배되면 큐브에 다가가는 것과 멀어지는 것 같은 작은 행동 의존적 변화가 예측에서 묻힌다는 것이다.

어떻게 동작하나

제안 방법 AD-WM은 두 축을 결합한다. 첫째는 잔차 잠재 예측이다. 현재 표현 z_t가 공유 정보를 담고, 예측기는 행동 임베딩 e_t = ψ_ρ(a_t)를 받아 변화분만 추정한다. Δẑ_t = f_θ(z_t, e_t), ẑ_{t+1} = z_t + Δẑ_t로 두고, 인코딩된 다음 표현과의 제곱 오차 L_pred = ‖ẑ_{t+1} − z_{t+1}‖²로 지도한다. 절대 예측과 잔차 예측은 제약 없는 함수 클래스에서 같은 최적점을 갖지만, 잔차 매개화는 국소적 잠재 변화를 명시적으로 모델링해 학습 편향을 바꾼다. 다만 이것만으로는 그 변화가 행동 정보를 담도록 강제하지 않는다.

무엇과 다른가

둘째는 예측기 수준의 행동 복원이다. 관측된 상태 표현이 아니라 모델이 생성한 전이에 복원을 걸어, MPC가 실제로 굴리는 동역학을 직접 제약한다. 역동역학 헤드는 ê_t = g_ω(z_t, ẑ_{t+1})로 예측된 전이 양 끝에서 행동 임베딩을 복원하고 L_inv = ‖ê_t − sg(e_t)‖²를 최소화한다. 목표는 detach되어 있고 예측된 종점은 미분 가능하므로, 기울기가 예측기와 행동 인코더로 흘러 들어간다. 정규화 행동 복원은 임베딩을 배치 평균·표준편차로 성분별 표준화한 뒤, [z_t, ẑ_{t+1}]을 받는 별도 헤드가 q_η(·|z_t, ẑ_{t+1}) = N(μ_η, I)를 출력하게 한다. 목표는 L_MI = −log q_η(ē_t|z_t, ẑ_{t+1}) + β·D_KL(q_η ‖ N(0,I))이고, 단위 공분산 가정에서 상수항을 빼면 ½‖ē_t − μ_η‖² + (β/2)‖μ_η‖²로 줄어든다. 첫 항은 표준화된 행동 임베딩을 복원하고, 둘째 항은 예측 평균을 0으로 당긴다. 저자들은 이를 조건부 상호정보량 I(ē_t; ẑ_{t+1} | z_t)에 대한 Barber–Agakov 하한의 변분 항으로 동기화한다. 전체 손실은 L_pred + λ_sig·L_sig + λ_inv·L_inv + λ_MI·L_MI이며, SIGReg 표현 정규화는 LeWM과 동일하다. 배포 시 두 보조 헤드는 버려지고 MPC 절차는 전혀 바뀌지 않는다. CEM은 300개 후보, 30개 엘리트, 30회 반복, 지평 5, 리시딩 지평 5, 행동 블록 5로 돌아가고 종료 잠재 거리로 점수를 매긴다.

어떻게 쓰나

논문은 제어 성능을 예측 오차와 분리해 보기 위한 진단 지표도 제안한다. 공유 후보 뱅크에서 예측 비용과 실제 실행 후 인코딩한 종료 이미지로 얻은 실현 비용의 스피어만 순위 상관을 CAD(반사실적 행동 구분성)로 정의하고, CEM이 실제로 남기는 엘리트 집합의 품질을 정규화된 best-in-elite regret R_k와 elite-mean regret R̄_k로 측정한다. 실현 비용은 훈련과 MPC 바깥에서 진단용으로만 쓰인다.

전제와 한계

시뮬레이션 결과는 OGBench-Cube를 중심으로 보고된다. 동일한 시작 상태에서 매칭된 LeWM 재현 베이스라인 대비 하드 스타트 성공률(P00–P04 평균)이 3.7%에서 52.0%로 올랐다. 교차 환경에서는 다섯 개 중 네 개에서 평균 성공률이 개선됐다. Cube 73.3%→90.7%, Reacher 76.7%→83.3%, TwoRoom 90%→98%, Scene 35.5%→39.5%이고, PushT는 94%→92%로 떨어졌다. Scene의 이득은 Drawer 63.3%→69.3%, Window 46.7%→55.3%에 집중됐지만 세 시드 짝지은 검정에서는 p=0.13으로 유의하지 않다. 구성 요소 분해에서는 LeWM 3.7%에서 절대 예측에 Inv+MI를 더하면 14.4%, 잔차 예측만으로 34.7%, Res+Inv 37.1%, Res+MI 54.7%, AD-WM 52.0%가 나왔다. MI는 역동역학 입력을 바꿔도 평균 하드 스타트를 14.9/15.6/20.4 포인트 끌어올렸고, MI 가중치는 0.005–0.05 구간에서 50.9–65.2%로 37.1%(Inv만)보다 높았으며 0.03에서 정점을 찍었다. Inv의 효과는 더 작고 가중치에 의존적이어서, 기본값 0.10은 no-Inv 평균을 넘지 못했다.

진단 결과는 이 논문의 핵심 주장을 뒷받침한다. LeWM이 사실적 MSE는 가장 낮았지만 하드 스타트 성공률도 가장 낮았고, AD-WM은 MSE가 가장 높은데도 52.0%를 기록했다. 15개 모델–시드 관측에서 성공률은 CAD와 ρ=−0.399의 상관을 보인 반면, −R_30, −R̄_30과는 각각 0.863, 0.810으로 훨씬 강하게 연관됐다. 시드·케이스 계층적 짝지은 부트스트랩(20,000회)에서 R_30 감소는 LeWM→Res 0.028 [0.014, 0.043], Res→Res+Inv 0.002 [−0.003, 0.009], Res+Inv→AD-WM 0.014 [0.009, 0.023]였다. LeWM에서 AD-WM으로 가면 실현 최선 후보 보존율이 0.198에서 0.135로 떨어지는데, 거의 동점인 대안이 있으면 regret은 여전히 낮을 수 있다. 국소 동역학도 안정해져 잠재 증가분 MSE가 0.0067→0.0055로, 연속 예측 증가분의 코사인 유사도가 음수인 비율이 0.242→0.146으로 줄었고, CEM 중심과 최저 비용 격자점 거리는 0.827→0.130, 평균 엘리트-중심 거리는 0.448→0.082로 좁혀졌다.

실로봇 전이는 Franka에서 검증했다. 동결된 V-JEPA 2 ViT-G 인코더와 DROID(cadene/droid_1.0.1) 후속 학습을 V-JEPA 2-AC와 동일하게 맞추고 전이 매개화와 보조 목표만 바꿨으며, 실험실 이미지나 시연은 일절 쓰지 않았다. 기본 픽앤플레이스 성공률이 42.2%에서 71.1%로 올랐고, 세 물체 모두 개선됐다(33.3%→60.0%, 46.7%→73.3%, 46.7%→80.0%). 모델당 82회, 총 164회 시행이며 제외된 시행은 없고 안전 정지도 실패로 계산했다. 다만 이 평가는 수동으로 준 grasp/move/place 이미지 목표를 따르는 단거리 MPC 파이프라인 안에서 이뤄진다.

개발자 관점에서 이 논문의 실용적 메시지는 명확하다. JEPA류 잠재 세계 모델을 MPC나 CEM 플래너와 붙일 때 검증 지표로 사실적 예측 MSE만 보면 안 된다는 것이다. 학습 손실이 좋아지는 것과 후보 행동 순위가 좋아지는 것은 별개이며, 엘리트 집합의 실현 비용(regret)을 직접 재는 편이 폐루프 성공률과 더 잘 맞는다. 보조 헤드는 추론 시 제거되므로 배포 지연이나 플래너 코드를 건드리지 않고 학습 파이프라인에만 손을 대면 된다는 점도 도입 부담을 낮춘다. 다만 MI 가중치가 0.005–0.05에서 성능이 50.9–65.2%로 크게 흔들리고 Inv는 기본값에서 이득이 확인되지 않았으므로, 자체 데이터셋에서는 가중치 민감도를 먼저 확인해야 한다.

저자들이 밝힌 한계도 분명하다. 진단 실험은 Cube에 국한되고, Scene의 이득은 세 시드에서 결론이 나지 않았다. 외부 베이스라인 비교는 추론 방식과 체크포인트 수가 달라 LeWM–AD-WM 짝처럼 훈련·추론을 통제하지 못한다. 시뮬레이션과 로봇 후속 학습은 서로 다른 보조 가중치를 쓰고, 로봇 평가는 수동 서브골과 비무작위 시행, 단일 사이트·카메라·백본에 의존한다. 향후 과제로 더 넓은 설정과 수동 서브골 의존도 축소를 제시한다.