기술 요약: 듀얼링 월드 모델 (Dueling World Models)

문제 정의

잠재 월드 모델(Latent world models), 특히 JEPA(Joint-Embedding Predictive Architecture)를 따르는 모델들은 현재 상태와 행동으로부터 미래의 상태 임베딩을 예측함으로써 계획을 세운다. 그러나 장면 내에 에이전트가 제어할 수 없는 움직임(방해 요소, distractors)이 포함될 경우, 이 모델들은 "침묵하는" 실패 모드(silent failure mode)를 겪는다. 즉, 전체적인 훈련 손실(재구성 또는 예측 오차)은 계속 개선됨에도 불구하고, 서로 다른 행동에 대한 예측값이 구별 불가능해지는 현상이 발생한다.

기존의 해결책들은 재구성 손실, 작업 보상, 부가적인 역동학(inverse-dynamics) 목적 함수, 또는 별도의 네트워크 브랜치를 추가하여 제어 가능한 역동성과 외생적 방해 요소를 분리하려고 시도한다. 이러한 접근 방식들은 대개 방해 요소의 과정에 대한 특정 가정을 필요로 하거나 추가적인 감독(supervision)을 요구한다. 저자들은 표준적인 단일 구조 예측기 $P(z, a)$에서, 실제 행동 라벨에 의해 명시적으로 접지(grounded)되거나 세그멘테이션 마스크를 통해 제거되지 않는 한, 행동 채널이 방해 요소의 역동성과 얽히게 된다는 점을 지적한다. 결정적으로, 표준적인 검증 손실 기반의 모델 선택은 행동 신호가 붕괴된 체크포인트를 적극적으로 선택하게 되므로, 이러한 "행동 맹목성(action-blindness)"을 감지하는 데 실패한다.

방법론: AD-JEPA

본 논문은 보상, 재구성, 또는 방해 요소 특화 부가 손실을 요구하지 않고 행동 채널을 격리하는, 행동 조건부 잠재 예측기의 최소한의 구조적 재구성인 **AD-JEPA (Action-Decomposed JEPA)**를 제안한다.

핵심 메커니즘: 행동 평균 중심화 (Action-Mean Centering)

이 방법은 예측기를 두 개의 스트림으로 분해한다:

1. 수동 스트림 ($B(z)$): 행동이 마진화(marginalized out)되었을 때 장면이 어떻게 진화하는지를 예측한다.
2. 중심화된 행동 스트림 ($u(z, a)$): 원시 행동 오프셋에서 행동에 대한 평균 예측을 뺀 값으로 계산되는 행동 특유의 오프셋을 나타낸다.

업데이트 규칙은 다음과 같이 정의된다:
$$\hat{z}_{t+1} = B(z_t) + \underbrace{[\Delta(z_t, a_t) - \bar{\Delta}(z_t)]}_{u(z_t, a_t)}$$
여기서 $\bar{\Delta}(z) = \mathbb{E}_{a' \sim \nu}[\Delta(z, a')]$는 행동 제안 분포 $\nu$에 대한 평균 오프셋이다.

이론적 토대

이 접근 방식은 스칼라 값에 대한 Dueling DQN 분해($Q = V + (A - \bar{A})$)를 벡터 값 형태의 잠재 역동성으로 확장 적용한다. 저자들은 이 뺄셈이 학습된 경향성이 아니라 **대수적 항등식(algebraic identity)**임을 증명한다:

• 공통 모드 상쇄 (Common-Mode Cancellation): 모든 행동에 대해 동일하게 유입되는 변동(공통 모드 변동), 즉 행동과 독립적인 방해 요소를 포함한 모든 요소는 정확하게 상쇄된다.
• 정확성 (Exactness): 이 상쇄는 이산 행동 집합(정확한 열거를 통해) 및 연속 행동 집합(몬테카를로 샘플링을 통해) 모두에서 성립하며, 이는 평균이 유한한 후보 집합에 대해 추정되는 한 유효하다.
• 외생성 가정 불필요: 이 방법은 프로세스 수준에서 방해 요소가 에이전트의 정책과 독립적이어야 한다는 프로세스 수준의 외생성 가정을 요구하지 않는다. 단, 전이(transition)에 대한 방해 요소의 기여가 행동들에 대해 공통 모드여야 한다는 조건만 필요하다.

훈련 목적 함수

모델은 스톱-그래디언트(stop-gradient) EMA 타겟을 대상으로 하는 표준 JEPA 예측 손실과 함께 다음을 사용하여 훈련된다:

• InfoNCE ($L_{act}$): 실행된 행동의 예측이 반사실적(counterfactual) 행동들 사이에서 타겟과 가장 가깝도록 유도한다.
• 압축 패널티 ($L_{off}$): 행동 채널이 상태 지속성(state persistence)을 흡수하는 것을 방지한다.
• VICReg 정규화: 표현 붕괴(representation collapse)를 방지한다.

특히, 저자들은 차원을 선택하기 위한 학습된 게이팅 메커니즘을 추가하는 것이 무력(inert)했으며, 균일한 다운 스케일링으로 퇴화했다는 점을 보고하며, 중심화(centering) 자체가 핵심적인 메커니즘임을 확인했다.

주요 기여

1. 행동 평균 중심화: 행동 채널을 공통 모드 불변(common-mode invariant)으로 만드는, 행동 조건부 예측기의 한 줄짜리 재구조화이다. 저자들은 이 상쇄가 이산 및 샘플링된 행동 집합 모두에 대해 정확한 상쇄를 제공한다는 증명(Proposition 1, 2)을 제공한다.
2. 경험적 검증: 이 방법은 네 가지 서로 다른 환경에서 검증되었다:
   • Gridworld (FourRooms): 회전하는 이진 방해 세포(binary distractor cells)가 있는 환경.
   • 합성 생성기 (Synthetic Generators): 알려진 제어 가능 요소와 노이즈 요소가 있는 환경.
   • 연속 제어 (DMC): 픽셀 수준의 움직이는 폐쇄물(occluders)이 있는 환경.
   • Atari (Freeway): 자연스러운 교통 방해 요소가 있는 환경.
     모든 경우에서, 중심화된 채널은 에이전트의 효과를 보존하면서 방해 요소의 누출을 거의 제로에 가깝게 줄였으나, 얽힌 예측기(entangled predictors)는 실패하였다.
3. 읽기 전용 적용 가능성 (Readout-Only Applicability): 이 메커니즘은 아키텍처의 속성이 아니라 읽기(readout)의 속성이므로, 동결된(frozen) 사전 훈련된 월드 모델(특히 RePo 및 TIA)에 사후적으로 적용할 수 있다. 저자들은 이러한 동결된 모델들의 원시 행동 델타 읽기 값이 거의 0에 가까웠음에도 불구하고, 이 뺄셈을 적용함으로써 이전에 잠재되어 있고 경로가 지정되지 않았던 행동 채널을 성공적으로 복구했음을 입증했다.
4. 경계 정의: 논문은 명시적으로 한계를 식별한다: 이 방법은 방해 역동성이 행동과 상관관계가 있는 경우(예: 재입력(reafferent) 방해 요소 또는 에이전트가 사용자에게 반응하는 경우) 실패한다. 이러한 경우 방해 요소는 공통 모드가 아니며 상쇄되지 않는다.

결과

• 복구 vs 거부 (Recovery vs. Rejection): 합성 테스트에서, 중심화된 채널은 높은 제어 가능 요소 복구율($R^2 \approx 0.87\text{--}0.93$)을 달성하는 동시에 노이즈 요소를 효과적으로 거부($R^2 \approx -0.00$)했다. 비중심화 버전은 신호를 복구하긴 했으나 노이즈를 심하게 유출했다.
• 제어 성능: Gridworld에서 중심화된 채널을 사용하는 모델 예측 제어기(MPC)는 30개의 방해 세포가 있음에도 높은 목표 도달 성공률($\approx 0.92$)을 유지한 반면, 표준 예측기는 $\approx 0.57$로 붕괴되었다.
• 사후 적용 성공: RePo 및 TIA 모델에 뺄셈을 적용했을 때, 원시 행동 델타 읽기 값은 0에 가까웠으나 중심화된 채널은 0.19에서 0.52 사이의 $R^2$ 값을 보이며 성공적으로 복구되었다.
• 실패 모드: 이 방법은 행동 상관관계가 있는 스트레스 테스트에서 정확히 실패하여 우연 수준으로 붕괴되었으며, 이는 이론적 경계와 일치한다.

의의 및 주장

본 논문은 중심화가 행동 조건부 월드 모델을 위한 보편적인 플러그인이라고 주장한다. 그 의의는 다음과 같다:

• 단순성: 재구성, 보상 기반 팩터화, 또는 부가적인 역동학 손실의 복잡성 없이 방해 요소 제거를 달성한다.
• 구조적 보장: 분리는 학습된 경향성에 의존하는 것이 아니라 파라미터화의 대수적 속성이다.
• 개조 능력 (Retrofitting Capability): 기존의 동결된 월드 모델을 재학습 없이 추론 시점에 "탈취득(de-distracted)"할 수 있게 하여, 얽힌 역동성에 의해 숨겨져 있던 행동 채널을 드러내 준다.

저자들은 제어 확장성에 대해 겸허한 태도를 보인다. 이 채널이 그리드월드에서는 제어를 가능하게 하지만, 이를 DMC 연속 제어 작업으로 확장하려는 시도는 랜덤 베이스라인을 이기지 못했다. 저자들은 이를 채널 자체의 결함이라기보다 플래너(planner)에서의 모델 정확도 병목 현상 때문이라고 설명한다. 따라서 본 논문의 주요 기여는 **표현 수준(representation-level)**에 있다: 즉, 얽힌 예측기가 "행동 맹목" 상태에 빠지는 곳에서 깨끗하고 공통 모드 불변인 행동 채널을 제공하는 것이다.