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

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

To isolate the contribution of the optimisation objective, we compare SWIRL against direct SFT using the same set of unlabelled videos initially annotated by our IDM model. As shown in Figure 3 with details in F , SWIRL significantly outperforms the SFT baselines (both continued training and data merging), which stagnate or degrade as data scales. We attribute this disparity to the inherent ambiguity of visual dynamics and action verbalisation. SFT enforces a strict token-level imitation of the IDM’s specific pseudo-labels; however, a single visual transition often corresponds to multiple valid descriptions. Forcing the model to mimic one specific verbalisation can lead to overfitting noise and suppressing valid alternative predictions. In contrast, our GRPO framework effectively relaxes this constraint by encouraging consistency assessed by IDM, rather than exact replication. By exploring the FWM’s rollout space and action trajectories that the IDM model recognises as physically plausible, SWIRL learns a more robust and generalisable FWM that is not limited by the specific phrasing of the teacher annotations.

#### Sharing θ \theta and ϕ \phi .

We study the trade-off between parameter efficiency and training stability by comparing shared- and separate-weight designs (Figure 2 bottom and Table 1 ). Using separate weights yields stable optimisation, with the IDM score improving monotonically from 6.37 to 6.57 over three iterations (Figure 2 , Left). In contrast, shared weights ( θ = ϕ \theta=\phi ) reduce memory footprint but introduce instability, as reflected by a performance drop at Iteration 3 (to ∼ \sim 6.52).

This instability is also reflected in downstream performance: on Aurora-Bench , shared weights slightly underperform the separate variant (5.06 vs. 5; Table 1 ), and on long-horizon evaluation (Table 3 ), the shared model consistently lags behind (e.g., 1.74 vs. 1.68 overall). We attribute this gap to gradient interference between the heterogeneous objectives of FWM (visual generation) and IDM (inferring the linguistic action), suggesting that more robust unification mechanisms are needed to effectively share representations across modalities ( Shi et al., 2025 ; Ma et al., 2025 ; Qu et al., 2025 ; Zheng et al., 2025 ) .

## 5 Conclusion

We introduce SWIRL , a unified framework for enabling VLMs and LLMs to intrinsically model future states conditioned on the current state and latent actions, without relying on human-annotated trajectories. By interpreting actions as latent variables and alternately optimising forward world modelling and inverse dynamics modelling objectives with GRPO, our method induces self-improving world modelling purely from unlabelled data. We provide theoretical guarantees establishing the learnability of each optimisation phase, formally linking our objectives to variational mutual information bounds and evidence lower bounds. Empirically, we demonstrate the effectiveness of SWIRL across diverse settings, including real-world visual dynamics with VLMs and textual environments (physical and digital simulations and tool calling) with LLMs, validating its generality.

## Impact Statement

This paper presents SWIRL , a framework that enables LLMs and VLMs to self-improve their internal world models using unlabelled data. By reducing reliance on expensive human-annotated trajectories, our work advances the efficiency and accessibility of training capable reasoning agents. This has positive implications for the development of general-purpose assistants that can better understand physical dynamics and digital environments.

However, we acknowledge potential societal consequences associated with these capabilities. First, our method utilises self-improving loops on unlabelled “in-the-wild” data. Without human curation, there is a risk that the model may internalise, reinforce, or amplify biases and harmful patterns present in the raw distribution of web and video data. Future work employing this framework should incorporate safety filters or constitutional AI principles within the reward mechanism to mitigate this risk.

Second, the improvements in Visual Dynamics Prediction and Web HTML interaction imply a step forward in generative video capabilities and autonomous web agents. While these advancements aid in creative tools and automation, they also lower the barrier for generating deepfakes or creating agents capable of bypassing web-based security measures (e.g., CAPTCHAs) or conducting automated interactions at scale. We emphasise the necessity of developing robust detection tools for synthetic media and implementing strict access controls and guardrails for autonomous agents operating in real-world web environments.

## Acknowledgements

This work is supported by the ERC Starting Grant AToM-FM (101222956) awarded to Edoardo M. Ponti. The authors acknowledge the use of resources provided by the Isambard-AI National AI Research Resource (AIRR). Isambard-AI is operated by the University of Bristol and is funded by the UK Government’s Department for Science, Innovation and Technology (DSIT) via UK Research and Innovation; and the Science and Technology Facilities Council [ST/AIRR/I-A-I/1023] ( McIntosh-Smith et al., 2024 ) .

## Appendix A Theoretical Derivation

In this section, we provide the detailed theoretical justification for SWIRL . We formalise the interaction between the Forward World Modelling (FWM) and Inverse Dynamics Model (IDM) as an alternating maximisation of Identifiability (via Variational Mutual Information) and Data Fidelity (via the Evidence Lower Bound).

We define the following notation: • x ∈ 𝒮 x\in\mathcal{S} : The source observation (current state s t s_{t} ).

• y ∈ 𝒮 y\in\mathcal{S} : The ground-truth target observation (next state s t + 1 s_{t+1} ), distributed according to the data distribution 𝒟 ⁡ ( y | x ) \mathcal{D}(y|x) .

• y ^ ∈ 𝒮 \hat{y}\in\mathcal{S} : A generated target observation sampled from the forward model.

• z ∈ 𝒜 z\in\mathcal{A} : The latent action driving the transition.

We employ two parametrised models: 1. Forward World Modelling (FWM): P θ ​ ( y ^ | x , z ) P_{\theta}(\hat{y}|x,z) , parametrised by θ \theta . FWM parametrised the environment’s transition distribution.

2. Inverse Dynamics Model (IDM): Q ϕ ​ ( z | x , y ) Q_{\phi}(z|x,y) , parametrised by ϕ \phi . IDM parametrised an approximate posterior over latent actions given a state transition.

### A.1 Phase I: Learnability of the Forward Model (FWM)

In the first phase, we freeze the IDM parameters ϕ \phi and optimise the FWM parameters θ \theta . Our objective is to generate trajectories y ^ \hat{y} that are identifiable by the inference model. We formalise this as maximising the Conditional Mutual Information (CMI) between the latent action Z Z and the generated observation Y ^ \hat{Y} , conditioned on the source state X X .

Crucially, in our algorithm, the latent actions Z Z used for training are not sampled from a fixed uninformative prior, but are inferred from the ground-truth data using the current IDM. We denote this Empirical Belief Distribution as P ~ ​ ( z | x ) \tilde{P}(z|x) : P ~ ​ ( z | x ) ≜ 𝔼 y ∼ 𝒟 ⁡ ( y | x ) ​ [ Q ϕ ​ ( z | x , y ) ] \tilde{P}(z|x)\triangleq\mathbb{E}_{y\sim\mathcal{D}(y|x)}\left[Q_{\phi}(z|x,y)\right] (6)

###### Theorem A.1 .

Optimising the FWM to maximise the likelihood assigned by the frozen IDM to generated samples maximises a Variational Lower Bound of the Conditional Mutual Information I P ~ ​ ( Z ; Y ^ | X ) I_{\tilde{P}}(Z;\hat{Y}|X) defined over the empirical belief distribution.

###### Proof.

The Conditional Mutual Information under the joint distribution induced by the data and the frozen IDM is: I P ~ ​ ( Z ; Y ^ | X ) = H P ~ ​ ( Z | X ) − H ⁡ ( Z | Y ^ , X ) I_{\tilde{P}}(Z;\hat{Y}|X)=H_{\tilde{P}}(Z|X)-H(Z|\hat{Y},X) (7)

The term H P ~ ​ ( Z | X ) H_{\tilde{P}}(Z|X) is the entropy of the marginal distribution of actions inferred from the dataset. Since ϕ \phi is frozen in this phase and the dataset 𝒟 \mathcal{D} is fixed, P ~ ​ ( z | x ) \tilde{P}(z|x) is constant with respect to θ \theta . Therefore, maximising the mutual information is equivalent to minimizing the conditional entropy H ⁡ ( Z | Y ^ , X ) H(Z|\hat{Y},X) .

The conditional entropy is defined as: H ⁡ ( Z | Y ^ , X ) \displaystyle H(Z|\hat{Y},X) = − 𝔼 x ∼ 𝒟 𝔼 z ∼ P ~ ​ ( z | x ) 𝔼 y ^ ∼ P θ ( ⋅ | x , z ) [ log P ( z | y ^ , x ) ] \displaystyle=-\mathbb{E}_{x\sim\mathcal{D}}\mathbb{E}_{z\sim\tilde{P}(z|x)}\mathbb{E}_{\hat{y}\sim P_{\theta}(\cdot|x,z)}\left[\log P(z|\hat{y},x)\right] (8)

The true posterior P ⁡ ( z | y ^ , x ) P(z|\hat{y},x) is intractable as it requires marginalising over the action space. We utilise the frozen IDM Q ϕ ​ ( z | x , y ^ ) Q_{\phi}(z|x,\hat{y}) as a variational approximation. By the non-negativity of the KL divergence D KL ( P ( z | y ^ , x ) | | Q ϕ ( z | x , y ^ ) ) ≥ 0 D_{\mathrm{KL}}(P(z|\hat{y},x)||Q_{\phi}(z|x,\hat{y}))\geq 0 , we have the lower bound: 𝔼 y ^ ​ [ log ⁡ P ⁡ ( z | y ^ , x ) ] ≥ 𝔼 y ^ ​ [ log ⁡ Q ϕ ​ ( z | x , y ^ ) ] \mathbb{E}_{\hat{y}}[\log P(z|\hat{y},x)]\geq\mathbb{E}_{\hat{y}}[\log Q_{\phi}(z|x,\hat{y})] (9)

Substituting this into the entropy term yields the variational lower bound for the CMI: I P ~ ​ ( Z ; Y ^ | X ) \displaystyle I_{\tilde{P}}(Z;\hat{Y}|X) ≥ H P ~ ( Z | X ) + 𝔼 x ∼ 𝒟 𝔼 z ∼ P ~ ​ ( z | x ) 𝔼 y ^ ∼ P θ ( ⋅ | x , z ) [ log Q ϕ ( z | x , y ^ ) ] \displaystyle\geq H_{\tilde{P}}(Z|X)+\mathbb{E}_{x\sim\mathcal{D}}\mathbb{E}_{z\sim\tilde{P}(z|x)}\mathbb{E}_{\hat{y}\sim P_{\theta}(\cdot|x,z)}\left[\log Q_{\phi}(z|x,\hat{y})\right] (10)

Expanding the definition of P ~ ​ ( z | x ) \tilde{P}(z|x) from Eq. (1), the optimization objective becomes: 𝒥 FWM ( θ ) = 𝔼 x ∼ 𝒟 𝔼 y ∼ 𝒟 ⁡ ( y | x ) [ 𝔼 z ∼ Q ϕ ( ⋅ | x , y ) [ 𝔼 y ^ ∼ P θ ( ⋅ | x , z ) [ log Q ϕ ( z | x , y ^ ) ] ] ] \mathcal{J}_{\text{FWM}}(\theta)=\mathbb{E}_{x\sim\mathcal{D}}\mathbb{E}_{y\sim\mathcal{D}(y|x)}\left[\mathbb{E}_{z\sim Q_{\phi}(\cdot|x,y)}\left[\mathbb{E}_{\hat{y}\sim P_{\theta}(\cdot|x,z)}\left[\log Q_{\phi}(z|x,\hat{y})\right]\right]\right] (11)

∎

### A.2 Phase II: Learnability of the Inverse Model (IDM)

In the second phase, we freeze θ \theta and optimise the IDM parameters ϕ \phi . We seek to maximise the log-likelihood of the observed ground-truth dynamics log ⁡ P θ ​ ( y | x ) \log P_{\theta}(y|x) (Data Fidelity). We model this via the Evidence Lower Bound (ELBO), treating the initialised policy at the start of the iteration as the reference prior, denoted π ref ​ ( z | x ) \pi_{\mathrm{ref}}(z|x) .

###### Theorem A.2 .

Optimising the IDM via the KL-regularised policy gradient objective with reward R ⁡ ( x , z , y ) = log ⁡ P θ ​ ( y | x , z ) R(x,z,y)=\log P_{\theta}(y|x,z) , β = 1 \beta=1 , and reference policy π ref \pi_{\mathrm{ref}} maximises the Evidence Lower Bound (ELBO).

###### Proof.

We express the marginal log-likelihood of the data by introducing the variational distribution Q ϕ ​ ( z | x , y ) Q_{\phi}(z|x,y) : log ⁡ P θ ​ ( y | x ) \displaystyle\log P_{\theta}(y|x) = log ∑ z ∈ 𝒜 P θ ( y | x , z ) π ref ( z | x ) \displaystyle=\log\sum_{z\in\mathcal{A}}P_{\theta}(y|x,z)\pi_{\mathrm{ref}}(z|x) = log 𝔼 z ∼ Q ϕ ( ⋅ | x , y ) [ P θ ​ ( y | x , z ) ​ π ref ​ ( z | x ) Q ϕ ​ ( z | x , y ) ] \displaystyle=\log\mathbb{E}_{z\sim Q_{\phi}(\cdot|x,y)}\left[\frac{P_{\theta}(y|x,z)\pi_{\mathrm{ref}}(z|x)}{Q_{\phi}(z|x,y)}\right] (12)

Applying Jensen’s inequality (concavity of log): log ⁡ P θ ​ ( y | x ) \displaystyle\log P_{\theta}(y|x) ≥ 𝔼 z ∼ Q ϕ ​ [ log ⁡ P θ ​ ( y | x , z ) + log ⁡ π ref ​ ( z | x ) − log ⁡ Q ϕ ​ ( z | x , y ) ] \displaystyle\geq\mathbb{E}_{z\sim Q_{\phi}}\left[\log P_{\theta}(y|x,z)+\log\pi_{\mathrm{ref}}(z|x)-\log Q_{\phi}(z|x,y)\right] = 𝔼 z ∼ Q ϕ [ log P θ ( y | x , z ) ] − D KL ( Q ϕ ( z | x , y ) | | π ref ( z | x ) ) \displaystyle=\mathbb{E}_{z\sim Q_{\phi}}\left[\log P_{\theta}(y|x,z)\right]-D_{\mathrm{KL}}(Q_{\phi}(z|x,y)\,||\,\pi_{\mathrm{ref}}(z|x)) (13)

This is the standard Evidence Lower Bound. The KL-regularised policy gradient objective (Phase II) is defined as: 𝒥 ( ϕ ) = 𝔼 z ∼ Q ϕ ( ⋅ | x , y ) [ R ( x , z , y ) ] − β D KL ( Q ϕ ( ⋅ | x , y ) | | π ref ( ⋅ | x ) ) \mathcal{J}(\phi)=\mathbb{E}_{z\sim Q_{\phi}(\cdot|x,y)}\left[R(x,z,y)\right]-\beta\,D_{\mathrm{KL}}(Q_{\phi}(\cdot|x,y)\,||\,\pi_{\mathrm{ref}}(\cdot|x)) (14)

By setting the reward to the FWM log-likelihood, R = log ⁡ P θ ​ ( y | x , z ) R=\log P_{\theta}(y|x,z) , we observe that: 𝒥 ​ ( ϕ ) ≡ ELBO ​ ( ϕ ) \mathcal{J}(\phi)\equiv\text{ELBO}(\phi) (15)

Thus, the IDM update step performs coordinate ascent on the evidence lower bound of the data likelihood, ensuring the inferred actions explain the ground-truth transitions under the current forward dynamics. ∎

## Appendix B Implementation Details

### B.1 Hyperparameters

For the Liquid-SFT model, we fine-tune Liquid-7B ( Wu et al., 2024 ) on Pico-banana-400K ( Qian et al., 2025 ) and Aurora ’s training set. Training is performed for 5 epochs with a batch size of 128. We use a learning rate of 2 × 10 − 5 2\times 10^{-5} with a cosine learning rate schedule, allocating 5% of the total training steps for warm-up.

For the non-iterative setup ( SWIRL (IDM → \rightarrow FWM) ), the model is trained with a batch size of 64 and a learning rate of 1 × 10 − 6 1\times 10^{-6} . A cosine learning rate schedule with 100 warm-up steps is applied. GRPO is used with a rollout size of 8 and a KL regularization coefficient β = 0.1 \beta=0.1 .

For the iterative setup ( SWIRL (Iterative) ), we use a batch size of 128 and a learning rate of 2 × 10 − 7 2\times 10^{-7} , together with a cosine schedule and 50 warm-up steps. In the SWIRL (Iterative + Share) variant, the learning rate is increased to 5 × 10 − 7 5\times 10^{-7} while all other settings remain unchanged. In both iterative GRPO setups, we set the decoding temperature to 0.75 and Top- P P to 0.96 to control the quality of predicted futures. Additionally, we apply a logit processor to constrain Liquid’s generation to image tokens for FWM and textual tokens for IDM.

For LLM experiments on StableToolBench , we fine-tune models using GRPO with DeepSpeed for memory-efficient distributed training. Optimization is performed with a learning rate of 5 × 10 − 5 5\times 10^{-5} under a cosine learning rate schedule with 25 warm-up steps, using an effective batch size of 128. For each prompt, we sample 64 rollouts. The maximum prompt and completion lengths are set to 8,126 and 4,096 tokens, respectively. Unless otherwise specified, we use a KL regularization coefficient of β = 0.1 \beta=0.1 . During GRPO rollout, we set the decoding temperature to 0.7 and Top- P P to 0.96. For Mind2Web and ScienceWorld , we adopt the same configuration, except that the maximum prompt and completion lengths are both set to 4,096 tokens, and the number of GRPO rollouts is reduced to 16.

For VLM experiments, training is conducted using DeepSpeed on 32 NVIDIA H200 140GB GPUs for FWM and 8 NVIDIA H200 140GB GPUs for IDM. For LLM experiments, both FWM and IDM are trained on 8 NVIDIA Grace-Hopper (GH100) GPUs.

### B.2 Evaluation Prompt for GPT-4o as a Judge

For Aurora-Bench , we use the same prompt as in Deng et al. (2025) . For ByteMorph , we employ the official GPT-4o judge template provided by Chang et al. (2025) . Since WorldPredictionBench ( Chen et al., 2025a ) does not natively support multi-turn image editing evaluation, we adopt the prompt template shown in Figure 4 . This template is designed to balance the penalties for directly copying the source image and for excessive editing.

## Appendix C Sanity Check Results for General Image Editing.