Self-Improving World Modelling with Latent Actions — part 1/4

Cleaned-text code points [0, 15528). Bibliography omitted here; source.md includes it.

# Self-Improving World Modelling with Latent Actions

###### Abstract

Internal modelling of the world—predicting transitions between previous states X X and next states Y Y under actions Z 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 ) P_{\theta}(Y|X,Z) and an Inverse Dynamics Modelling (IDM) Q ϕ ​ ( Z | X , Y ) 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 Aurora-Bench , 28% on ByteMorph, 16% on WorldPredictionBench , and 14% on StableToolBench . 1 1 1 The code and models developed in this paper will be made available at https://github.com/yfqiu-nlp/swirl .

###### Keywords:

## 1 Introduction

Intrinsic world modelling is a model’s latent understanding of the environment and agent dynamics, i.e., trajectories of states and actions, which enables simulating possible futures without hallucinations. Large foundation models, such as Large Language Models (LLMs) and Vision-Language Models (VLMs), have arguably been shown to internalise world modelling to some extent during their training, thus enabling reasoning and planning without an external, specialised world model ( Vafa et al., 2024 ; Qiu et al., 2024 ; Liu et al., 2025b ; Chen et al., 2025b ; Wang et al., ; Xiong et al., 2026 ) . For example, explicitly training models to predict the results of invoking tools ( Guo et al., 2025b ) or executing code ( Copet et al., 2025 ) significantly enhances tool calling and coding tasks. This ability can also transfer to spatial reasoning tasks ( Qiu et al., 2025 ; Tehenan et al., 2025 ) and reduce hallucinations that contradict external regularities ( Liu et al., 2025a ; Chen et al., 2025b ) .

While promising, the development of robust internal world models faces a significant bottleneck in data scalability. Current approaches rely heavily on execution logs or trajectories where observations are densely annotated with specific actions, because such annotations are naturally available in restricted environments (e.g., logs for tool calling or coding). However, for open-world tasks, collecting manual annotations for every transition is prohibitively expensive and intractable. Moreover, an additional challenge is the inherent ambiguity of inverse dynamics : a transition between two states may be explained by multiple valid actions, making purely supervised learning brittle when data is sparse.

Inspired by recent advances in self-improving learning for other applications ( Wang et al., 2025a ; Wang et al., 2025c ; Jin et al., 2025 ; Mao et al., 2025 ; Jin et al., 2025 ) , we propose SWIRL (Self-improving World modelling with Iterative RL), a reciprocal optimisation framework to enhance the intrinsic world modelling of LLMs and VLMs from state-only sequences where the intermediate action is latent. We formalise world modelling as two components ( Wang et al., 2025b ; Qiu et al., 2025 ; Chen et al., 2025a ) : Forward World Modelling (FWM; predicting the next state y y given current state and latent action x , z x,z ), and Inverse Dynamics Model (IDM; inferring latent action z z given states x , y x,y ). In our reciprocal framework, we optimise the FWM P θ P_{\theta} to generate futures that are consistently identifiable by the IDM, and the IDM Q ϕ Q_{\phi} to infer actions that maximise the likelihood of the state dynamics predicted by the FWM. From a variational inference perspective, we theoretically prove that the optimisation of FWM is equivalent to maximising a lower bound on the Conditional Mutual Information I ⁡ ( Z ; Y ^ | X ) I(Z;\hat{Y}|X) ( Barber and Agakov, 2003 ) between the latent Z Z and future states Y ^ \hat{Y} predicted by the FWM; and that the optimisation of the IDM is equivalent to performing coordinate ascent on the Evidence Lower Bound (ELBO) of the log-likelihood log ⁡ P θ ​ ( Y | X ) \log P_{\theta}(Y|X) . Without relying on ground-truth action annotations, we optimise such a reciprocal framework with Group Relative Policy Optimisation (GRPO; Guo et al. 2025a ). FWM and IDM take alternating roles of policy and reward, iterating until both components converge.

We validate our framework across four distinct environments on six benchmarks: open-world visual dynamics for VLMs on Aurora-Bench , ByteMorph and WorldPredictionBench , synthetic textual worlds on ScienceWorld , web HTML on Mind2Web , and tool calling on StableToolBench for LLMs. Empirical results confirm that our reciprocal self-improving allows models to learn effective dynamics from unlabelled state sequences, outperforming supervised fine-tuning baselines and achieving parity with larger, state-of-the-art models. Our contributions are summarised as follows: • We propose a novel self-improving framework for world modelling in LLMs and VLMs, reciprocally reinforcing FWM and IDM without action annotations.

• We provide a rigorous theoretical proof that SWIRL corresponds to alternating between maximising Variational Mutual Information and Evidence Lower Bound.

• Empirical evidence on six benchmarks across visual, textual, web, tool calling environments demonstrates that SWIRL models effectively self-improve, leading to more enhanced forward world modelling capabilities.

## 2 Related Work

#### Intrinsic World Models.

Recent research has explored whether internalised world models emerge in LLMs and VLMs through careful evaluation. Vafa et al. (2024) assessed whether the models’ representations truly capture coherent world dynamics ( Xiong et al., 2026 ) . Similarly, Tehenan et al. (2025) and Qiu et al. (2024) found evidence that LLMs implicitly encode spatial and temporal relationships to some degree. World modelling also emerges naturally during the pre-training of large unified VLMs ( Deng et al., 2025 ; Cui et al., 2025 ) and from video-based training ( Chen et al., 2025b ; Qiu et al., 2025 ) .

In addition, previous work established that explicit world modelling can boost performance in downstream applications. For instance, Copet et al. (2025) proposed the Coding World Model for programming, and Lehrach et al. (2025) extended this approach to game-playing environments. Modelling the outcomes of function tool calls ( Guo et al., 2025b ) or planning ( Li et al., 2025 ) had similar effects. As a result, dedicated post-training pipelines have been proposed to endow LLMs ( Xie et al., 2025 ) and VLMs ( Xiang et al., 2024 ) with explicit world modelling capabilities. Extensive benchmarks for evaluation have been proposed, including forward and inverse dynamics prediction ( Chen et al., 2025a ; Wang et al., 2025b ; Gao et al., 2025 ) .

#### Self-Improving Learning.

Self-improving learning refers to models learning from their own (possibly curated) signals without external supervision. Huang et al. (2023) showed that LLMs can generate high-confidence answers and fine-tune themselves on these outputs to improve reasoning. Bensal et al. (2025) proposed a self-reflection and reinforcement learning loop where models analyse mistakes and retry, boosting performance on tasks like maths and function calling. Lee et al. () introduce a curriculum in which models iteratively generate and filter correct answers, progressively tackling harder problems, while Zhao et al. (2024) demonstrate that self-synthesised input–output pairs improve classification and generation quality. ( Wang et al., 2025a ) enhance self-improvement capabilities of agents with a skill library. LLM’s coding and unit test generation capabilities can also co-evolve iterating on each other’s outcomes ( Wang et al., 2025c ) . Similarly, VLMs can refine visual and language reasoning using self-generated corrections ( He et al., 2025 ) .

In unified VLMs ( Deng et al., 2025 ; Wu et al., 2024 ; Lin et al., 2025 ; Xiao et al., 2025 ) , understanding performance often exceeds generation ( Shi et al., 2025 ; Ma et al., 2025 ; Qu et al., 2025 ; Zheng et al., 2025 ; Zhang et al., 2025 ) . A common strategy is then to use the understanding head as a critic to guide generation with carefully designed rubrics or heuristics ( Mao et al., 2025 ; Jin et al., 2025 ; Qiu et al., 2026 ) . Unlike these approaches, SWIRL theoretically and empirically proves the effectiveness of utilising the predicted likelihoods from understanding (action prediction) and generation (next state prediction) to establish a reciprocal cycle for VLMs where improvements in the generation head also enhance the understanding head, and vice versa.

## 3 Methodology

### 3.1 Task Formulation

We consider transitions from a source state x ∈ 𝒮 x\in\mathcal{S} to a target state y ∈ 𝒮 y\in\mathcal{S} , mediated by a latent action z ∈ 𝒜 z\in\mathcal{A} . Following Qiu et al. (2025) and Wang et al. (2025b) , we parametrise world modelling using two components: i) Forward World Modelling (FWM) : P θ ​ ( y | x , z ) P_{\theta}(y|x,z) , which predicts the next state given previous state and action; and ii) Inverse-Dynamics Prediction (IDM): Q ϕ ​ ( z | x , y ) Q_{\phi}(z|x,y) , which infers the action given the state transition. The parameters θ \theta and ϕ \phi can be either disjoint or shared.

We consider four classes of environments with different observation–action formulations. In real-world visual environments, observations are pixel-level visual inputs and actions are specified in natural language. We study this setting using unified vision–language models (VLMs) capable of perceiving and generating interleaved image–text sequences. In synthetic textual environments, both observations and actions are expressed purely in language and are governed by an underlying simulator, which we model using large language models (LLMs). In web-based environments, states and actions correspond to raw HTML and interaction logs, respectively, while in tool-use settings, actions are tool calls and observations consist of the conversational context and tool execution outcomes.

### 3.2 SWIRL

#### Intuition.

SWIRL alternates between two tightly coupled optimisation phases. In the first phase, we optimise the Forward World Model (FWM), P θ P_{\theta} , to generate future states that are identifiable by the Inverse Dynamics Model (IDM). This enforces identifiability in forward prediction. In Eq. 2 below, we prove that this objective is equivalent to maximising a variational lower bound on the conditional mutual information I ⁡ ( Z ; Y ^ | X ) I(Z;\hat{Y}|X) ( Barber and Agakov, 2003 ) , encouraging the predicted futures Y ^ \hat{Y} to retain maximal information about the underlying latent actions Z Z .

In the second phase, we optimise the IDM, Q ϕ Q_{\phi} , to improve data fidelity by inferring actions that best explain the observed state transitions under the learned FWM. We formally justify this step by proving that the resulting objective corresponds to maximising the Evidence Lower Bound (ELBO) of the inverse model (see Eq. 5 below).

#### Objectives.

As depicted in Algorithm 1 , we rely only on state-only sequences ( x t , y t + 1 ) (x_{t},y_{t+1}) as data. In the first phase of each iteration of SWIRL , we optimise FWM as a policy based on the frozen IDM rewards. Let P θ P_{\theta} denote the FWM policy parametrised by θ \theta . For each training step in GRPO ( Guo et al., 2025a ) , we first sample a latent action z t ∼ Q ϕ ( ⋅ ∣ x t , y t ) z_{t}\sim Q_{\phi}(\cdot\mid x_{t},y_{t}) from the IDM. Then, we sample a group of G G rollouts { y ^ t + 1 ( k ) } k = 1 G \{\hat{y}_{t+1}^{(k)}\}_{k=1}^{G} from P θ ( ⋅ ∣ x t , z t ) P_{\theta}(\cdot\mid x_{t},z_{t}) . To evaluate their quality, we use a frozen IDM model Q ϕ Q_{\phi} prompted for the IDM task, which estimates the likelihood of the action z t z_{t} given the generated transition: P π ϕ ​ ( z t ∣ x t , y ^ t + 1 ( k ) ) . P_{\pi_{\phi}}(z_{t}\mid x_{t},\hat{y}_{t+1}^{(k)}). The reward for the k k -th rollout is defined as r k = log ⁡ Q ϕ ​ ( z t ∣ x t , y ^ t + 1 ( k ) ) . r_{k}=\log Q_{\phi}(z_{t}\mid x_{t},\hat{y}_{t+1}^{(k)}). This reward enforces that a predicted future is considered plausible only if IDM can retrospectively explain the action that led to the predicted state. We then optimise P θ P_{\theta} to maximise the advantage induced by these rewards.

After the FWM policy converges, we invert the optimisation direction to refine the IDM. In this phase, the optimised FWM model is frozen and used as the reward signal. The policy Q ϕ Q_{\phi} is prompted for the IDM task, generating action candidates { z ^ t ( i ) } i = 1 G \{\hat{z}_{t}^{(i)}\}_{i=1}^{G} given a state transition ( x t , y t + 1 ) (x_{t},y_{t+1}) . Each candidate action is evaluated by querying the FWM for the likelihood of the target observation: r k = log ⁡ P π θ ​ ( y t + 1 ∣ x t , z ^ t ( k ) ) . r_{k}=\log P_{\pi_{\theta}}(y_{t+1}\mid x_{t},\hat{z}_{t}^{(k)}). This reward favours data fidelity with respect to the state transition. We alternate between FWM and IDM optimisation phases, swapping the roles of policy and reward model until convergence. This reciprocal training paradigm ensures that improvements in forward prediction directly enhance action inference by the inverse dynamic model, and vice versa, driving the system toward a globally consistent both given all unlabelled observations.

### 3.3 Theoretical Analysis

Phase I encourages that generated futures are distinguishable (high mutual information), allowing the FWM to produce distinct outcomes for varied latent actions. Phase II encourages that inferred actions are plausible (high likelihood), allowing the IDM to map transitions to actions that the FWM can reproduce.

#### Phase I: FWM optimisation via Variational Information Maximisation.

In this phase, we freeze ϕ \phi and update θ \theta to generate next states y ^ \hat{y} that are identifiable by the inference model. We formalise this as maximising the Conditional Mutual Information (CMI) I ⁡ ( Z ; Y ^ | X ) I(Z;\hat{Y}|X) with respect to the empirical belief distribution of the model. Let P ~ ​ ( z | x ) ≜ 𝔼 y ∼ 𝒟 ⁡ ( y | x ) ​ [ Q ϕ ​ ( z | x , y ) ] \tilde{P}(z|x)\triangleq\mathbb{E}_{y\sim\mathcal{D}(y|x)}[Q_{\phi}(z|x,y)] denote the marginal distribution of latent actions inferred by the IDM over the dataset.

###### Theorem 3.1 (FWM Lower Bound) .

Optimising the FWM to maximise the log-probability assigned to generated samples by the frozen IDM maximises a variational lower bound on the Conditional Mutual Information I P ~ ​ ( Z ; Y ^ | X ) I_{\tilde{P}}(Z;\hat{Y}|X) defined over the empirical belief distribution P ~ ​ ( z | x ) \tilde{P}(z|x) .

###### Proof.