🌀
SWIRL: Self-Improving World Modelling with Latent Actions

Overview

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.

🛠️ Environment Setup

This repo is configured to use the Conda environment in environment.yaml.

conda env create -f environment.yaml
conda activate swirl

Optional (Hugging Face caches):

export HF_HOME=/path/to/hf_cache
export HF_DATASETS_CACHE=/path/to/hf_datasets_cache

📦 Assets and Paths

We keep paths configurable via environment variables and script arguments. By default, the scripts assume:

• STORAGE_PATH points to your local workspace (defaults to the repo root in the scripts).
• DATASET_ROOT=$STORAGE_PATH/datasets
• CKPT_ROOT=$STORAGE_PATH/rlwm-checkpoints
• VAE_PATH=$STORAGE_PATH/vision_tokenizers/chameleon

Required assets (not included in this repo):

• Vision tokenizer (VQGAN/Chameleon): vqgan.yaml and vqgan.ckpt
  • expected at: vision_tokenizers/chameleon/
• Datasets (place under datasets/ or pass --base_dir/--dataset_root):
  • datasets/pico-banana-400k/
  • datasets/AURORA/
  • datasets/ByteMorph/
  • datasets/unsupervised-frames-ucf-kinetics-mit/
  • datasets/VIDGEN-1M_images/ (for VidGen annotations)
• Model checkpoints (if not using HF model IDs):
  • rlwm-checkpoints/<experiment_name>/

📂 Data Preparation

Before training, datasets must be tokenized into discrete codes using the vision tokenizer.

Preprocessing (tokenization + annotations)

The main preprocessing entrypoint is:

sbatch scripts/preprocess_liquid.sh

Notes:

• Edit STORAGE_PATH, DATASET_ROOT, VAE_PATH, and dataset names inside the script as needed.
• preprocess/annotate_mit_kinetics_ucf.py and preprocess/annotate_vidgen_shards.py are called from the script.

Optional: Extract frames for VidGen

sbatch preprocess/extract_frames.sh

🚀 Training

1) Supervised Fine-Tuning (SFT)

Warm-start the Liquid model using FSDP:

sbatch scripts/sft_liquid.sh

Key knobs:

• EXPERIMENT_NAME
• --model_name_or_path
• --data_mixture
• --dataset_base_path

2) Reinforcement Learning (GRPO)

Train both directions (FDP and IDP) from the same script:

sbatch scripts/grpo_liquid.sh

Key knobs:

• RUN_NAME
• CKPT_PATH
• --dataset_paths

✅ Evaluation

AURORA

sbatch scripts/eval_liquid_aurora.sh

• Generates images and computes CLIP-based metrics.
• Optional GPT evaluation requires:

  export OPENAI_API_KEY=...
  

ByteMorph

sbatch scripts/eval_liquid_bytemorph.sh

• Optional GPT evaluation also uses OPENAI_API_KEY.

WorldPredictionBench

sbatch scripts/eval_worldprediction.sh

Notes:

• Set DATASET_ROOT and METADATA_PATH to point to the WorldPrediction dataset and its metadata JSON.
• If you have a frame map (segment_uid -> relative frame path), set FRAME_MAP_PATH.
• ffmpeg must be available in PATH to extract frames from videos.
• Optional GPT evaluation uses OPENAI_API_KEY.

🧪 Running without Slurm

All scripts/*.sh files are standard bash scripts. You can run them directly after:

1. activating the environment, and
2. setting the required path variables.