[논문 리뷰] Self-Improving World Modelling with Latent Actions

Yifu Qiu, Zheng Zhao외 5명|arXiv (Cornell University)|2026. 02. 05.

Multimodal Machine Learning Applications인용 수 0

한 줄 요약

SWIRL은 상태-단 시퀀스에서 행동을 잠재 변수로 다루고 순환 RL을 사용하여 Forward World Modelling과 Inverse Dynamics Models를 번갈아 최적화함으로써 세계 모델을 학습합니다. 이 접근법은 행동 주석 없이 시각적, 텍스트, 웹 및 도구 활용 환경에서 강력한 결과를 달성합니다.

ABSTRACT

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.

연구 동기 및 목표

• LLMs와 VLMs를 위한 행동 주석 없이도 강력한 내부 세계 모델링의 필요성 제기.
• GRPO로 학습되는 상호 최적화 프레임워크(FWM 및 IDM) 제안.
• Phase I를 변분 상호 정보의 하한과 Phase II를 ELBO 최적화에 연결하는 이론적 보장 제공.
• 개방형 세계의 시각적 다이나믹스, 텍스트 시뮬레이션, 웹 상호작용, 도구 사용에서 실험적 이득 시연.

제안 방법

• Forward World Modelling P_theta(y|x,z)와 Inverse Dynamics Q_phi(z|x,y)로 세계 모델링을 형식화합니다.
• 두 가지 단계를 번갈아 수행합니다: Phase I은 FWM을 최적화하여 조건부 상호 정보 I(Z; Ŷ|X)의 하한을 최대화하고, Phase II는 IDM을 최적화하여 log P_theta(Y|X)의 ELBO를 최대화합니다.
• 두 모델을 서로 보상 신호로 취급하여 업데이트하기 위해 Group Relative Policy Optimisation (GRPO)을 사용합니다.
• IDM에서 잠재 행동을 샘플링하고 FWM으로 롤아웃을 생성한 뒤 IDM의 로그 가능도로 보상합니다; 그런 다음 FWM을 고정시키고 IDM에 대한 행동을 샘플링하여 FWM의 로그 가능도로 보상합니다.
• CMI의 변분 하한 및 IDM의 ELBO 목표와의 등가를 보여주는 이론을 적용하고, 여러 환경에서의 실험으로 검증합니다.

실험 결과

연구 질문

• RQ1 잠재 행동 세계 모델이 ground-truth 행동 없이도 상태-전용 시퀀스에서 효과적으로 학습할 수 있는가?
• RQ2 전방 예측과 역 동역학 간의 상호 최적화가 식별 가능하고 충실한 세계 모델을 낳는가?
• RQ3 SWIRL은 행동 주석이 있는 또는 감독된 기준과 비교하여 시각 및 텍스트 다이나믹 환경에서 어떻게 성능을 보이는가?
• RQ4 SWIRL 업데이트에 수반되는 이론적 보장(CMI 하한 및 ELBO 최대화)은 무엇인가?

주요 결과

• SWIRL은 다수의 벤치마크(Aurora-Bench, ByteMorph, WorldPredictionBench, StableToolBench)에서 SFT 기초선 대비 상당한 개선을 보인다.
• 공유 또는 분리 가중치를 가진 반복적 SWIRL 변형은 시각적 다이나믹스 작업에서 비반복 기준선보다 더 높은 평가 점수를 달성한다.
• SWIRL은 SFT보다 장기적 충실도를 더 높게 유지하여 WorldPredictionBench의 horizon 턴까지 시간적 일관성을 향상시킨다.
• 텍스트 환경에서 SWIRL은 도구 다이나믹스의 일반화를 향상시키며 StableToolBench에서 BLEU 및 관련 지표에서 SFT를 능가한다.
• 벤치마크 전반에 걸쳐 SWIRL은 더 큰 모델과 비슷한 성능을 보이며 더 가벼운 사후 학습으로 데이터 효율적이고 자기 개선 학습을 보여준다.