You can deploy this model and start earning money today!

kevin510/
swm-next-paligemma-wm

Name: swm-next-paligemma-wm
Availability: InStock

google/paligemma-3b-pt-224

finetune

apache-2.0

Views

No views yet

SWM Reproduction — PaliGemma-3B Semantic World Model (OGBench cube stacking)

Reproduction of the Semantic World Model from Berg et al., Semantic World Models (arXiv 2510.19818) on OGBench visual-cube-quadruple: a PaliGemma-3B VLM fine-tuned to answer questions about future states conditioned on action sequences, used as a differentiable reward for gradient-based planning.

Repo root holds checkpoint 60,000 of 64,000 (the parity checkpoint) from run swm-repro-v1 (paper recipe: full fine-tune, lr 1e-5 linear decay, effective batch 96, AdamW). Subfolders checkpoint-4000/ … checkpoint-64000/ hold every 4,000-step checkpoint of the run (weights only). Suite results at n=50: 64k = 70.67, 56k = 70.67, 60k = 75.33.

Results (3 cube-stacking tasks × 50 seeds, authors' eval code)

	blue→green	blue→yellow	yellow→red	suite
Base diffusion policy	52	64	50	55.3
This checkpoint	86	72	68	75.33
Authors' released ckpt (same harness)	84	76	68	76.0

Usage

Weights only — the processor/tokenizer ships with the authors' released model:

python

from swm.paligemma_wm.modeling_paligemma_wm import PaliGemmaWMForConditionalGeneration
from swm.paligemma_wm.processing_paligemma_wm import PaliGemmaWMProcessor

model = PaliGemmaWMForConditionalGeneration.from_pretrained("kevin510/swm-next-paligemma-wm")
processor = PaliGemmaWMProcessor.from_pretrained("jacob3333/paligemma_wm_ogbench")

Model classes: krohling/swms @ swm-next-repro.

Training data

Self-collected per paper A.3.1 with the authors' generation scripts (krohling/ogbench @ swm-next-cli): 300 noisy + 300 play episodes, 200 steps, native 224×224, seeds 0–300; SAQA tuples with all 6 question types, 4 horizons/state over 0–20 (h=0 included), per-(type × answer) balanced sampling. Datasets: noisy · play

Training curves: wandb