잠재적 행동을 활용한 자기 개선 세계 모델링

Self-Improving World Modelling with Latent Actions

LLM(대규모 언어 모델) 및 VLM(대규모 비전 모델)의 추론 및 계획 능력에 필수적인 요소는 이전 상태 $X$와 다음 상태 $Y$ 사이의 전이를 행동 $Z$ 하에서 예측하는 세계 모델 구축입니다. 일반적으로 이러한 모델을 학습하려면 비용이 많이 드는 행동 레이블이 부착된 데이터가 필요합니다. 본 논문에서는 SWIRL이라는 자기 개선 프레임워크를 제안합니다. SWIRL은 행동을 잠재 변수로 간주하고, 상태 정보만으로 구성된 시퀀스에서 학습하며, 순방향 세계 모델링(FWM) $P_θ(Y|X,Z)$과 역역학 모델링(IDM) $Q_φ(Z|X,Y)$을 번갈아 가며 사용합니다. SWIRL은 두 가지 단계를 반복합니다. (1) 변분 정보 최대화(Variational Information Maximisation): FWM을 업데이트하여, 사전 상태가 주어졌을 때 잠재적 행동과의 조건부 상호 정보량을 최대화하는 다음 상태를 생성하도록 합니다. 이를 통해 일관성을 높입니다. (2) ELBO(Evidence Lower Bound) 최대화: IDM을 업데이트하여 관찰된 전이를 설명합니다. 이는 효과적으로 좌표 상승(coordinate ascent)을 수행하는 것입니다. 두 모델 모두 강화 학습(특히 GRPO)을 사용하여 학습되며, 상대 모델의 로그 확률을 보상 신호로 사용합니다. 우리는 두 업데이트에 대한 이론적인 학습 가능성을 보장하며, 다양한 환경(단일 턴 및 다중 턴 개방형 시각적 동역 시스템 및 물리, 웹, 도구 호출을 위한 합성 텍스트 환경)에서 LLM 및 VLM에 대한 SWIRL의 성능을 평가했습니다. SWIRL은 AURORABench에서 16%, ByteMorph에서 28%, WorldPredictionBench에서 16%, StableToolBench에서 14%의 성능 향상을 달성했습니다.

Internal modelling of the world -- predicting transitions between previous states $X$ and next states $Y$ under actions $Z$ -- is essential to reasoning and planning for LLMs and VLMs. Learning such models typically requires costly action-labelled trajectories. We propose SWIRL, a self-improvement framework that learns from state-only sequences by treating actions as a latent variable and alternating between Forward World Modelling (FWM) $P_θ(Y|X,Z)$ and an Inverse Dynamics Modelling (IDM) $Q_φ(Z|X,Y)$. SWIRL iterates two phases: (1) Variational Information Maximisation, which updates the FWM to generate next states that maximise conditional mutual information with latent actions given prior states, encouraging identifiable consistency; and (2) ELBO Maximisation, which updates the IDM to explain observed transitions, effectively performing coordinate ascent. Both models are trained with reinforcement learning (specifically, GRPO) with the opposite frozen model's log-probability as a reward signal. We provide theoretical learnability guarantees for both updates, and evaluate SWIRL on LLMs and VLMs across multiple environments: single-turn and multi-turn open-world visual dynamics and synthetic textual environments for physics, web, and tool calling. SWIRL achieves gains of 16% on AURORABench, 28% on ByteMorph, 16% on WorldPredictionBench, and 14% on StableToolBench.