Self-Improving World Modelling with Latent Actions

TL;DR

SWIRL is a self-improving framework that learns world models from state-only sequences by treating actions as latent variables, using alternating variational and ELBO maximization with reinforcement learning, and demonstrating significant performance gains.

Contribution

It introduces SWIRL, a novel approach for learning world models from state-only data by combining variational information maximization and ELBO maximization with reinforcement learning.

Findings

Achieves up to 28% improvement on benchmark tasks.

Effectively learns from state-only sequences without explicit action labels.

Provides theoretical guarantees for the learning process.

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_\theta(Y|X,Z)$ and an Inverse Dynamics Modelling (IDM) $Q_\phi(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.