PaperlayerView exact revision on arXiv

Source-linked AI summary

Self-Improving World Modelling with Latent Actions

Yifu Qiu, Zheng Zhao, Waylon Li, Yftah Ziser, Anna Korhonen, Shay B. Cohen, Edoardo M. Ponti

arXiv:2602.06130v2cs.LGcs.AIcs.CL

TL;DR

Action-labelled trajectories make robust world modelling expensive, especially in open-world settings where transitions can have multiple valid inverse-dynamics explanations. SWIRL learns from state-only sequences by alternating latent-action forward and inverse models, and empirical evaluations show effectiveness across diverse visual and textual environments. The paper also flags risks from uncurated data and from enabling stronger generative-video and autonomous-web capabilities.

  • Problem

    Robust internal world models require costly action-annotated trajectories, while open-world transitions may have multiple valid actions that make sparse supervised learning brittle.

  • Method

    SWIRL treats actions as latent variables and alternates GRPO-based optimisation of a Forward World Model and an Inverse Dynamics Model using state-only sequences.

  • Results

    SWIRL effectively self-improves from unlabelled state sequences across visual dynamics, physical and digital simulations, web interaction, and tool calling, outperforming supervised fine-tuning baselines and matching larger state-of-the-art models.

  • Takeaways & Limitations

    The framework reduces reliance on human-annotated trajectories while providing a unified approach to intrinsic world modelling for LLMs and VLMs.

  • Takeaways & Limitations

    Uncurated in-the-wild data may reinforce harmful patterns, while stronger capabilities may lower barriers to deepfakes and web-security bypasses.

Abstract

from arXiv · show

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.

1. Introduction

SWIRL addresses the data bottleneck and ambiguity of action-labelled world-model training by learning from state-only sequences with reciprocal forward and inverse models. It theoretically justifies this optimisation and empirically validates it across visual, textual, web, and tool-calling environments.

  • Motivation: Robust world-model development is constrained by the expense of densely annotating actions in open-world transitions, while inverse dynamics can admit multiple valid explanations.These factors make purely supervised learning brittle when data is sparse.
  • Approach: SWIRL reciprocally optimises a Forward World Model and an Inverse Dynamics Model from state-only sequences with latent intermediate actions.The framework targets LLMs and VLMs without requiring action annotations.
  • Approach: The two components are iteratively trained with GRPO: the forward model generates identifiable futures using the inverse model as reward, then the inverse model improves data fidelity using the forward model as reward.The phases alternate policy and reward roles between the models.
  • Theory: SWIRL’s forward-model update corresponds to variational mutual-information maximisation, while its inverse-model update corresponds to coordinate ascent on the ELBO.These results provide a theoretical basis for the two optimisation phases.
  • Evaluation: The framework is evaluated across six benchmarks spanning open-world visual dynamics, synthetic textual worlds, web HTML, and tool calling.The reported environments cover both VLM and LLM settings.
  • Results: Empirical results show effective learning from unlabelled state sequences, outperforming supervised fine-tuning baselines and reaching parity with larger state-of-the-art models.The evidence spans the visual, textual, web, and tool-calling environments described in the evaluation.

2. Related Work

Prior work studies intrinsic world models, explicit world-model applications, self-improving learning, and critic-guided generation in LLMs and VLMs. SWIRL builds on these directions while positioning its reciprocal latent-action framework as distinct from understanding-head critics.

  • Intrinsic World Models: Research has evaluated whether LLMs and VLMs internalise coherent spatial, temporal, and broader world dynamics through pre-training and video-based training.These studies report evidence that such relationships are encoded to some degree.
  • Intrinsic World Models: Explicit world modelling has been applied to programming, game playing, function-tool outcomes, and planning, with reported benefits for downstream applications.These applications motivate dedicated post-training pipelines for LLMs.
  • Self-Improving Learning: Self-improving learning uses model-generated or self-reflective signals to improve reasoning, function calling, classification, generation, coding, and visual-language reasoning.Prior methods include iterative answer filtering, retry loops, skill libraries, and self-generated corrections.
  • Critic-Guided Generation: Unified VLMs commonly use an understanding head as a critic to guide generation with rubrics or heuristics because understanding performance often exceeds generation.The passage contrasts this strategy with SWIRL’s approach.

3. Methodology

SWIRL learns forward and inverse world models from state-only transitions by alternating reciprocal optimisation phases. The FWM makes predicted futures identifiable by latent actions, while the IDM infers actions that explain observed transitions under the learned dynamics.

  • Model components: SWIRL parameterises world modelling with an FWM that predicts next states from states and latent actions, and an IDM that infers actions from transitions.The FWM and IDM parameters may be disjoint or shared.
  • Reciprocal optimisation: The two phases alternate policy and reward roles using GRPO until convergence or a maximum iteration limit, returning jointly optimised FWM and IDM parameters.The training algorithm operates on an unlabelled dataset of state pairs and freezes the opposite model in each phase.
  • Phase I: FWM optimisation: Phase I freezes the IDM, samples latent actions and FWM rollouts, then rewards futures according to the frozen IDM’s log-likelihood of the sampled action.This makes predicted futures identifiable by the inverse model and corresponds to maximising a variational lower bound on conditional mutual information.
  • Phase II: IDM optimisation: Phase II freezes the FWM, samples candidate actions from the IDM, and rewards actions according to the FWM’s likelihood of the observed target state.The phase improves data fidelity by selecting actions that best explain observed transitions under the learned forward dynamics.
  • Theoretical basis: Theoretical analysis identifies Phase I with variational information maximisation and Phase II with ELBO maximisation, establishing lower-bound guarantees for both updates.The FWM update targets conditional mutual information, while the IDM update performs coordinate ascent on the ELBO under a KL-regularised objective.

4. Experiments and Results

SWIRL is evaluated across visual, long-horizon, textual, web, and tool-calling environments against supervised and larger-model baselines. Results show stronger data efficiency, improved long-horizon consistency, and gains in several dynamics tasks, while separate weights provide greater stability than sharing parameters.

  • Experimental Setup: SWIRL is evaluated on visual, long-horizon, physical, web, and tool-calling world-modelling environments using Liquid and Qwen-2.5-3B-Instruct.The visual evaluations include AURORA-BENCH, BYTEMORPH, and WORLDPREDICTIONBENCH; textual evaluations include SCIENCEWORLD, MIND2WEB, and STABLETOOLBENCH.
  • AURORA-BENCH: SWIRL consistently improves over Liquid-SFT across all five AURORA-BENCH datasets and remains competitive with larger unified VLMs.It also substantially outperforms diffusion-based editors such as InstructPix2Pix.
  • BYTEMORPH: 43.38 to 53.77 (+26.4%) Average score marks the improvement of SWIRL ITERATIVE over Liquid-SFT on BYTEMORPH.Gains are pronounced on Object/Human Motion and Interaction, while global camera control remains comparable because the unlabelled videos provide limited camera-control supervision.
  • WORLDPREDICTIONBENCH: +14.4% relative improvement over Liquid-SFT at T = 6 shows that SWIRL ITERATIVE preserves higher long-horizon fidelity.Liquid-SFT declines from 3.09 at T = 1 to 0.97 at T = 6, whereas SWIRL ITERATIVE reaches 1.11 at T = 6.
  • Textual Environments: +4.03 BLEU on ID-Low and +3.69 BLEU on ID-Medium show SWIRL’s strongest textual-environment advantage on STABLETOOLBENCH.SCIENCEWORLD and MIND2WEB are near saturation in semantic accuracy and comparable in exact lexical matching.
  • Analysis: SWIRL significantly outperforms continued-training and data-merging SFT baselines, which stagnate or degrade as the number of training samples increases.The comparison uses the same unlabelled videos initially annotated by the IDM model.

5. Conclusion

SWIRL enables LLMs and VLMs to model future states from unlabelled data by treating actions as latent variables and alternately optimising forward and inverse dynamics models. The framework is theoretically grounded and performs effectively across visual, physical, digital, and tool-calling environments.

  • Conclusion: SWIRL enables intrinsic future-state modelling in LLMs and VLMs without human-annotated trajectories.It treats actions as latent variables and learns from unlabelled data.
  • Conclusion: Alternating GRPO optimisation of FWM and IDM creates a self-improving loop that links predicted states with inferred actions.The framework provides theoretical guarantees connecting its objectives to variational mutual-information and evidence-lower-bound formulations.
  • Conclusion: SWIRL generalises across real-world visual dynamics, physical and digital simulations, and tool-calling environments.The reported evaluation covers both VLMs and LLMs across diverse state-transition settings.

Impact Statement

SWIRL enables LLMs and VLMs to improve internal world models from unlabelled data, reducing reliance on expensive annotated trajectories while raising safety and governance concerns. Its alternating FWM–IDM optimisation is theoretically linked to conditional mutual-information maximisation and ELBO ascent.

  • Impact Statement: SWIRL uses unlabelled data to help LLMs and VLMs self-improve their internal world models, reducing reliance on expensive human-annotated trajectories.The framework alternates between a Forward World Model and an Inverse Dynamics Model.
  • Societal Consequences: Unlabelled in-the-wild data may transmit or amplify biases and harmful patterns, motivating safety filters or constitutional-AI principles in future reward mechanisms.The paper also identifies risks from stronger generative-video and autonomous-web capabilities, including deepfakes and scaled security bypasses.
  • Phase I: Forward World Modelling: The FWM generates future states identifiable by the IDM, which provides inferred latent actions from observed transitions.During this phase, the IDM remains frozen while the FWM is optimised using the IDM’s likelihood of generated samples.
  • Phase I: Forward World Modelling: FWM optimisation maximises a variational lower bound on conditional mutual information between latent actions and generated future states given the source state.The bound uses the frozen IDM as a variational approximation to the posterior over latent actions.
  • Phase II: Inverse Dynamics Modelling: The IDM update maximises an ELBO for observed dynamics by using the FWM log-likelihood as its reward while regularising against a reference policy.This update performs coordinate ascent and encourages inferred actions to explain ground-truth transitions under the current FWM.

C. Sanity Check Results for General Image Editing.

The SFT warm-up pipeline gives a base model functional image-editing ability, enabling subsequent world-modelling training. PICO-BANANA-400K is important for maintaining this capability across metrics.

  • General Image Editing: Liquid-SFT acquires functional image-editing behaviour after training on AURORA and PICO-BANANA-400K, despite Liquid lacking native editing support.The evaluation reports semantic alignment, perceptual quality, and average performance against image-editing models.
  • General Image Editing: Liquid without SFT completely fails on GEDIT-BENCH, indicating that the warm-up stage is necessary before SWIRL.The paper treats this evaluation as capability verification rather than a competitive image-editing comparison.
  • Role in World Modelling: The warm-up capability supports later training for next-observation prediction and inverse dynamics modelling over visual transitions.The paper describes instruction-followed image transformations as a minimal but necessary foundation for these objectives.
  • Data Ablation: PICO-BANANA-400K removal lowers the average score from 3.06 to 2.43 and consistently degrades all reported metrics.The ablation attributes the difference to the component’s contribution to stabilising and enriching visual instruction-following behaviour.

D. IDM Analysis

The IDM provides a strong action-prediction discriminator and supplies a stable reward signal for SWIRL. Analyses also test whether latent actions remain interpretable rather than collapsing into short ciphers.

  • IDM Evaluation: The IDM reaches 6.38 out of 10 on AURORA-BENCH, supporting the premise that action inference is more tractable than high-dimensional future-state generation.The score is reported as a GPT evaluation of action prediction quality.
  • IDM Evaluation: Iterative SWIRL improves the AURORA-BENCH GPT4o score from 6.38 to 6.58 by leveraging the FWM reward signal.The paper links this stability to informative rewards that penalise physical inconsistencies.
  • Latent Action Interpretability: Generated latent actions retain greater than 94% uniqueness across iterations for both shared and separate weight configurations.The analysis examines action uniqueness, naturalness, and length as indicators of reward collapse or shortcutting.
  • Latent Action Interpretability: The actions remain natural under GPT-2 perplexity measurements and do not collapse into brief-token shortcuts during SWIRL training.These observations argue against the specific reward-hacking concern tested in the analysis.

F. Detailed Results for Comparing SFT and SWIRL.

Across controlled data budgets, reinforcement learning is more data-efficient than both supervised fine-tuning baselines, with advantages appearing early and widening as training data increases.

  • Overall Comparison: RL consistently outperforms SFT-Continue and SFT-Merge across all five benchmarks and every evaluated training-sample count.The comparison uses GPT-4o evaluation scores and an aggregated average.
  • Data Efficiency: RL begins outperforming SFT at 6.4K samples, and its performance gap widens as training proceeds.SFT baselines often plateau or degrade as more samples are added, whereas RL improves with additional data.
  • Structured Physical Dynamics: On Action Genome, Something, and Whatsup, Kubric, RL achieves the strongest or near-strongest performance at every data scale.These benchmarks emphasise structured physical dynamics.
  • Aggregated Results: Aggregated RL scores are highest at every data scale, with gains increasing from 4.27 to 4.73 at larger budgets.The authors interpret this pattern as evidence that RL uses limited data more effectively than supervised objectives.

G. Detailed Results for WORLDPREDICTIONBENCH

WorldPredictionBench evaluates temporal consistency across six autoregressive prediction turns, where early errors can compound over the rollout horizon. SWIRL degrades less sharply than the direct supervised baseline.

  • Six consecutive autoregressive turns stress-test temporal consistency because early prediction errors compound over the horizon.
  • 3.09 to 1.17: Liquid-SFT’s overall score falls from Turn 1 to Turn 4.
  • 1.59 at Turn 4: SWIRL’s Ours (Best) configuration retains a higher score than Liquid-SFT at the same turn.
  • The evaluation compares SWIRL variants with Liquid-SFT, Bagel, and other state-of-the-art VLMs across six prediction turns.

H. Detailed Results for Iterative Results

The iterative ablations compare separate and shared FWM–IDM weights, showing that separate weights produce the strongest peak forward-modelling performance. Rollout-size experiments further show a trade-off between absolute performance and computational cost.

  • Iterative training dynamics: Separate weights achieve the highest peak FWM performance, reaching a best average FWM score of 5.06 at Iteration 1.The configuration also performs strongly on MagicBrush, Action Genome, and Something.
  • Iterative training dynamics: Shared weights remain competitive and stable but do not surpass the separate-weight model’s peak generative performance.IDM accuracy improves slightly in later iterations for both settings, while separate weights reach their optimal FWM–IDM balance earlier.
  • Iterative training dynamics: Separate FWM and IDM parameters are used in the main experiments because peak forward world-modelling capability is the primary objective.
  • GRPO rollout-size ablation: 4.90 average at G = 64: larger GRPO rollout groups improve peak performance, while G = 16 reaches a competitive 4.80 average with less memory and compute.

J. Qualitative Examples

Qualitative examples illustrate SWIRL’s predictions across action-centric image editing, camera and object dynamics, and long-horizon procedural tasks. The examples span AURORA-BENCH, BYTEMORPH, and WORLDPREDICTIONBENCH.

  • AURORA-BENCH: AURORA-BENCH examples cover image editing, background replacement, geometric transformations, and spatial reasoning across several subsets.
  • BYTEMORPH: BYTEMORPH examples demonstrate camera zoom, camera motion, and object motion through coherent view and localised movement changes.
  • WORLDPREDICTIONBENCH: WORLDPREDICTIONBENCH examples show multi-step predictions conditioned on prior predicted images and textual actions in procedural tasks.Examples include replacing a car key battery and arranging bedding over multiple steps.
  • Figure 7 presents qualitative SWIRL predictions for AURORA-BENCH, BYTEMORPH, and WORLDPREDICTIONBENCH.
Loading 2602.06130v2…