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

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

The mutual information under the data-induced joint distribution P ⁡ ( x ) ​ P ~ ​ ( z | x ) ​ P θ ​ ( y ^ | x , z ) P(x)\tilde{P}(z|x)P_{\theta}(\hat{y}|x,z) is: I ⁡ ( Z ; Y ^ | X ) = 𝔼 x ∼ 𝒟 ​ [ H P ~ ​ ( Z | X ) − H ⁡ ( Z | Y ^ , X ) ] . I(Z;\hat{Y}|X)=\mathbb{E}_{x\sim\mathcal{D}}\left[H_{\tilde{P}}(Z|X)-H(Z|\hat{Y},X)\right]. (1) The entropy of the latent belief H P ~ ​ ( Z | X ) H_{\tilde{P}}(Z|X) depends only on the frozen IDM and data distribution, and is thus constant w.r.t. θ \theta . Maximising CMI is therefore equivalent to minimizing the conditional entropy H ⁡ ( Z | Y ^ , X ) H(Z|\hat{Y},X) . We use the variational posterior Q ϕ ​ ( z | x , y ^ ) Q_{\phi}(z|x,\hat{y}) as a proxy for the intractable true posterior P θ ​ ( z | y ^ , x ) P_{\theta}(z|\hat{y},x) ( Barber and Agakov, 2003 ) to bound the entropy term: − H \displaystyle-H ( Z | Y ^ , X ) = 𝔼 x , z ∼ P ~ , y ^ ∼ P θ ​ [ log ⁡ P θ ​ ( z | y ^ , x ) ] \displaystyle(Z|\hat{Y},X)=\mathbb{E}_{x,z\sim\tilde{P},\hat{y}\sim P_{\theta}}[\log P_{\theta}(z|\hat{y},x)] (2) ≥ 𝔼 x ∼ 𝒟 𝔼 z ∼ P ~ ​ ( z | x ) 𝔼 y ^ ∼ P θ ( ⋅ | x , z ) [ log Q ϕ ( z | x , y ^ ) ] . \displaystyle\geq\mathbb{E}_{x\sim\mathcal{D}}\mathbb{E}_{z\sim\tilde{P}(z|x)}\mathbb{E}_{\hat{y}\sim P_{\theta}(\cdot|x,z)}[\log Q_{\phi}(z|x,\hat{y})]. Substituting the definition of P ~ ​ ( z | x ) \tilde{P}(z|x) creates the following objective: 𝒥 ( θ ) = 𝔼 ( x , y ) ∼ 𝒟 𝔼 z ∼ Q ϕ ( ⋅ | x , y ) 𝔼 y ^ ∼ P θ ( ⋅ | x , z ) [ log Q ϕ ( z | x , y ^ ) ] . \mathcal{J}(\theta)=\mathbb{E}_{(x,y)\sim\mathcal{D}}\mathbb{E}_{z\sim Q_{\phi}(\cdot|x,y)}\mathbb{E}_{\hat{y}\sim P_{\theta}(\cdot|x,z)}[{\log Q_{\phi}(z|x,\hat{y})}]. (3) This matches Algorithm 1 : we sample trajectories ( x , y ) (x,y) , infer z z using the IDM, rollout y ^ \hat{y} using the FWM, and reward the FWM with the IDM’s log likelihood. The gradient of Eq. 3 is ∇ θ 𝒥 ( θ ) = 𝔼 ( x , y ) ∼ 𝒟 𝔼 z ∼ Q ϕ ( ⋅ | x , y ) 𝔼 y ^ ∼ P θ ( ⋅ | x , z ) [ log Q ϕ ( z ∣ x , y ^ ) ∇ θ log P θ ( y ^ ∣ x , z ) ] \nabla_{\theta}\mathcal{J}(\theta)=\mathbb{E}_{(x,y)\sim\mathcal{D}}\mathbb{E}_{z\sim Q_{\phi}(\cdot|x,y)}\mathbb{E}_{\hat{y}\sim P_{\theta}(\cdot|x,z)}[\log Q_{\phi}(z\mid x,\hat{y})\nabla_{\theta}\log P_{\theta}(\hat{y}\mid x,z)] , for which Group Relative Policy Optimisation (GRPO) offers an estimator: ∇ θ 𝒥 ^ GRPO ​ ( θ ) = 1 G ​ ∑ k = 1 G A k F ​ ∇ θ ​ log ⁡ P θ ​ ( y ^ k ∣ x , z ) , \widehat{\nabla_{\theta}\mathcal{J}}_{\text{GRPO}}(\theta)=\frac{1}{G}\sum_{k=1}^{G}A^{\text{F}}_{k}\nabla_{\theta}\log P_{\theta}(\hat{y}_{k}\mid x,z), (4) where A k A_{k} are the group-relative advantages derived from rewards R k = log ⁡ Q ϕ ​ ( z ∣ x , y ^ k ) R_{k}=\log Q_{\phi}(z\mid x,\hat{y}_{k}) . ∎

#### Phase II: IDM optimisation via ELBO Maximisation.

In the second phase, we freeze θ \theta and optimise ϕ \phi . The goal is to infer actions z z that explain the transition ( x , y ) (x,y) in the FWM dynamics, while staying close to the valid prior. We treat the initialised model at each iteration as an informative prior P ⁡ ( z | x ) ≜ π ref ​ ( z | x ) P(z|x)\triangleq\pi_{\text{ref}}(z|x) . This corresponds to maximising the Evidence Lower Bound (ELBO).

###### Theorem 3.2 (IDM Lower Bound) .

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

###### Proof.

We seek to maximise the marginal log-likelihood log ⁡ P θ ​ ( y | x ) \log P_{\theta}(y|x) . By introducing the variational distribution Q ϕ ​ ( z | x , y ) Q_{\phi}(z|x,y) and the informative prior P ⁡ ( z | x ) = π ref ​ ( z | x ) P(z|x)=\pi_{\text{ref}}(z|x) , we derive the ELBO: log ⁡ P θ ​ ( y | x ) = log ⁡ 𝔼 z ∼ Q ϕ ​ [ P θ ​ ( y | x , z ) ​ P ​ ( z | x ) Q ϕ ​ ( z | x , y ) ] ≥ 𝔼 z ∼ Q ϕ ( ⋅ | x , y ) [ log P θ ( y | x , z ) ] − D KL ( Q ϕ ( z | x , y ) | | P ( z | x ) ) ≜ ℒ ELBO . \begin{split}\log P_{\theta}(y|x)&=\log\mathbb{E}_{z\sim Q_{\phi}}\left[\frac{P_{\theta}(y|x,z)P(z|x)}{Q_{\phi}(z|x,y)}\right]\\ &\geq\mathbb{E}_{z\sim Q_{\phi}(\cdot|x,y)}\left[\log P_{\theta}(y|x,z)\right]\\ &-D_{\mathrm{KL}}(Q_{\phi}(z|x,y)\,||\,P(z|x))\triangleq\mathcal{L}_{\text{ELBO}}.\end{split} The KL-regularised policy gradient objective is defined as the expected reward subject to a KL constraint weighted with β = 1 \beta=1 against the reference policy: 2 2 2 In practice, we may allow β \beta to be tunable for adjusting the deviation against the reference policy, as the common practice in GRPO. 𝒥 ⁡ ( ϕ ) = 𝔼 z ∼ Q ϕ ​ [ R ⁡ ( x , z , y ) ] − D KL ( Q ϕ ( ⋅ | x , y ) | | π ref ( ⋅ | x ) ) . \displaystyle\begin{split}\mathcal{J}(\phi)&=\mathbb{E}_{z\sim Q_{\phi}}[R(x,z,y)]\\ &-D_{\mathrm{KL}}(Q_{\phi}(\cdot|x,y)||\pi_{\text{ref}}(\cdot|x)).\end{split} (5) By setting the reward R ⁡ ( x , z , y ) = log ⁡ P θ ​ ( y | x , z ) R(x,z,y)=\log P_{\theta}(y|x,z) and identifying the prior P ⁡ ( z | x ) P(z|x) with the reference policy π ref ​ ( z | x ) \pi_{\text{ref}}(z|x) , we observe that 𝒥 ⁡ ( ϕ ) ≡ ℒ ELBO \mathcal{J}(\phi)\equiv\mathcal{L}_{\text{ELBO}} . We again use the GRPO as the estimator to perform coordinate ascent on the ELBO. ∎

## 4 Experiments and Results

### 4.1 Experimental Setup

#### Models.

We choose Liquid ( Wu et al., 2024 ) as our base VLM as the publicly best 7B-size unified VLM with an autoregressive architecture. The medium size limits the scale requirements for the compute infrastructure, and the autoregressive nature allows for applying GRPO off-the-shelf (without ad-hoc adaptations for diffusion-based generation). For text-based environments, we use Qwen-2.5-3B-Instruct ( Qwen Team, 2024 ) , a competitive mid-size LLM.

#### SFT Warm-up.

Since our base VLMs/LLMs are general-purpose, they lack the specific interface capabilities required for environments (e.g., Liquid cannot natively predict image conditioning on both image and textual action). Prior to SWIRL , we conduct an initial SFT, which is strictly viewed as policy initialisation to ensure the model outputs valid actions and states; without this, a random policy would generate outputs rendering RL exploration impossible (as evidenced by Liquid’s poor zero-shot performance in GEdit-Bench in Appendix C ). For VLMs, we utilise image editing mixtures from Pico-banana-400K ( Qian et al., 2025 ) and Aurora ( Krojer et al., 2024 ) . 3 3 3 As a sanity check, we report the zero-shot and SFT performance of Liquid on GEdit-Bench compared to other VLMs in Appendix C . For LLMs, we fine-tune on the half of environment-specific episodes, and remain the rest discarding the annotated actions for SWIRL .

#### Iterative RL.

For SWIRL , we first conduct controlled experiments on the unlabelled video mixture from UCF-101 ( Soomro et al., 2012 ) , Movement-in-Times ( Monfort et al., 2019 ) , Kinetics700 ( Kay et al., 2017 ) , limiting training to the first phase IDM → \rightarrow FWM of SWIRL . This setup enables strict comparison with the bootstrapping baseline of Qiu et al. (2025) and allows us to observe stable convergence within a single epoch. After this controlled phase, we scale training by uniformly sampling 30K videos per iteration from a large-scale unlabelled video corpus, VidGen-1M ( Tan et al., 2024 ) . We extract the frame pairs from videos as in ( Chen et al., 2025d ) . Each iteration is trained for one epoch, and the model alternates between FWM and IDM optimisation phases as described in Section 3.2 . This protocol ensures a clean comparison to baselines in ( Qiu et al., 2025 ) while measuring the improvement across iterations.

#### Benchmarks.

We evaluate visual world modelling under two settings. Firstly, for single-turn next-observation prediction, we use Aurora-Bench and ByteMorph , both of which formulate dynamics prediction as action-conditioned image editing tasks with a strong emphasis on correctness in action-centric dynamics. Secondly, to evaluate long-horizon world modelling, we adopt WorldPredictionBench , which supports multi-step rollouts of up to four future observations across five subtasks, enabling a comprehensive assessment of FWM consistency.

To validate the generalisation of SWIRL , we conduct experiments across three grounded environments on LLM, each is a unique scheme of state transition. We utilise ScienceWorld ( Wang et al., 2022 ) to evaluate physical dynamics, where LLM must predict the textual consequences of scientific actions within a simulated world. We also employ Mind2Web ( Deng et al., 2023 ) , which challenges the model to forecast the updated HTML DOM tree following user interactions (e.g., clicks) on web elements. Finally, we assess functional tool dynamics via StableToolBench ( Guo et al., 2024 ) , where the objective is to simulate the execution output of API calls conditioned on the current conversational state and the invoked tool.

#### Baselines.

For visual world modelling, we compare SWIRL against a directly fine-tuned Liquid model ( SFT ). We reproduce two strong baselines from Qiu et al. (2025) : Test-time Verification , which selects the best sample from the fine-tuned Liquid based on IDM scores, and Bootstrap , which fine-tunes Liquid on silver data annotated by the IDM. We also include specialised diffusion-based image editing models, including GoT ( Fang et al., 2025 ) , SmartEdit ( Huang et al., 2024 ) , InstructPix2Pix ( Brooks et al., 2023 ) , and the Chameleon model family, following the protocol of Qiu et al. (2025) . To position our approach against recent (larger or diffusion-based) unified VLMs, we also evaluate BAGEL ( Deng et al., 2025 ) , OmniGen ( Xiao et al., 2025 ) , OmniGen2 ( Wu et al., 2025 ) , BLIP3o-NEXT ( Chen et al., 2025c ) , and UniWorld-V1 ( Lin et al., 2025 ) .

For text-based environments, prior work under this formulation is limited. We therefore primarily compare against SFT baselines and larger LLMs to ensure a fair comparison.

#### Evaluation Metrics.

For visual forward world modelling, we adopt GPT-4o-as-a-judge for holistic evaluation, following ( Qiu et al., 2025 ; Fang et al., 2025 ) with a 10-point scheme. We also report DiscEdit (DE) and CLIP scores for Aurora-Bench as in ( Krojer et al., 2024 ) . For text environments, we use BLEU, BERTScore, and ROUGE-L.

### 4.2 Single-turn Visual Dynamics Prediction

#### Aurora-Bench .

We present the quantitative evaluation of visual world modelling on Aurora-Bench in Table 1 . We compare SWIRL against state-of-the-art unified VLMs, specialised diffusion-based editing models, and ablations of our method. The primary comparison of interest is against Liquid-SFT, our direct supervised fine-tuning baseline. As highlighted in blue, SWIRL delivers consistent and significant improvements across all five benchmarks. Compared to the baselines from ( Qiu et al., 2025 ) , SWIRL surpasses the bootstrapping strategy using IDM to synthesise trajectories from unlabelled videos, and computation-heavy inference techniques like Test-time Verification (with up to N = 8 N=8 samples). Our best iterative strategy raises the average GPT-4o evaluation score from 4.36 (SFT) to 5.06, which also outperforms the non-iterative supervision from IDM only ( SWIRL (IDM → \rightarrow FWM) ) at 4.83, indicating the effectiveness of SWIRL by alternating policy and reward roles between FWM and IDM. Furthermore, despite starting from a weaker base model such as Liquid, and relying on a more lightweight post-training, 4 4 4 We use only around 400K samples from ( Qian et al., 2025 ; Krojer et al., 2024 ) to initialise the editing ability of Liquid. SWIRL remains highly competitive with other state-of-the-art unified VLMs such as OmniGen2 and BAGEL, while substantially outperforming diffusion-based editors like InstructPix2Pix. Qualitative examples for SWIRL are presented in Appendix J .

#### ByteMorph .

As the ByteMorph results in Table 2 show, all variants of SWIRL substantially raise the Average score over the Liquid-SFT baseline, e.g. SWIRL Iterative goes from 43.38 to 53.77 ( + 26.4 +26.4 %). Performance on global camera control ( Camera Zoom/Motion ) remains comparable. This is expected since the unlabelled in-the-wild videos used during the RL phase (e.g., VIDGEN-1M) are predominantly static and provide limited supervision for camera control. In contrast, SWIRL yields pronounced improvements on Object/Human Motion and Interaction , indicating effective learning of fine-grained dynamics. Notably, SWIRL matches the performance of larger or more heavily supervised VLMs (e.g., BAGEL and UniWorld-v1).

### 4.3 Multi-turn Visual Dynamics Prediction

#### WorldPredictionBench .

To assess long-horizon world modelling in addition to single-step prediction, we evaluate performance on WorldPredictionBench . The model must predict future observations autoregressively up to T = 6 T=6 time steps, using its own previous predictions as context. This setup tests the model’s ability to maintain physical consistency and resist the compounding covariate shift inherent in sequential generation. Table 3 summarises the GPT-4o-as-a-judge scores across all subsets in WorldPredictionBench .

The best variant of our models demonstrates a significant improvement in temporal consistency compared to Liquid-SFT. While SFT yields competitive performance at the immediate next step ( T = 1 T=1 ), it suffers from rapid degradation as the horizon increases, dropping from a score of 3.09 to 0.97 by T = 6 T=6 . In contrast, SWIRL ( Iterative ) maintains significantly higher fidelity throughout the rollout trajectory, achieving a +14.4% relative improvement over the baseline at T = 6 T=6 (1.11). We also remark that repeated self-improvement is not beneficial in this dataset, as performance peaks at the first iteration (IDM → \rightarrow FWM).

### 4.4 Textual Environments

Table 4 shows the quantitative evaluation on textual environments. When applied to Qwen-2.5-3B-Instruct in physical and digital simulation environments ( ScienceWorld and Mind2Web ), both SWIRL and SFT are near-saturation in semantic accuracy ( > > 92 BERTScore) and comparable in exact lexical matching (ROUGE-L). The advantage of our approach becomes most pronounced in StableToolBench , which requires the simulation of API execution outcomes. Here, we observe substantial improvements in generalisation capabilities, with our method surpassing SFT by +4.03 BLEU on ID-Low and +3.69 BLEU on ID-Medium splits. This establishes a new state-of-the-art for this model scale, outperforming also significantly larger open-weight instruction models (e.g., Qwen-2.5-32B, DeepSeek-7B). Overall, this suggests that the SWIRL can generalise to internalise complex API dynamics effectively for LLMs.

### 4.5 Analysis

#### Iterative RL.

We run an analysis to determine the potential for cumulative gain through iterative self-improvement beyond the first round. Figure 2 tracks the training dynamics across three iterations, where each iteration consists of updating the FWM using the current IDM as a reward (Phase I), followed by updating the IDM using the improved FWM as a reward (Phase II). We observe a clear virtuous cycle : enhancing the FWM’s forecasting capability enables it to provide a more robust verification for action prediction, which in turn yields a more precise reward signal for the subsequent FWM updates. The reward curves demonstrate that FWM and IDM in a separate-weight setup improve the training rewards effectively, and their optimisation converges with GRPO.

#### RL vs SFT.