Semantic World Models

TL;DR

This paper tackles the limitation of pixel-based future-frame predictions for planning by proposing Semantic World Models (SWMs), which predict task-relevant semantic outcomes about the future via a visual question answering framework. SWMs fine-tune a vision-language model on a SAQA dataset to answer questions about future states conditioned on current observations and action sequences, enabling planning through QA likelihood with both sampling-based and gradient-based methods. The approach demonstrates substantial generalization and planning improvements on LangTable and OGBench, including robustness to out-of-distribution configurations and effectiveness with suboptimal data, outperforming pixel-based baselines and offline RL. By leveraging large pretrained VLMs and QA abstractions, SWMs offer a scalable path toward grounded, flexible robotic control with strong generalization and robustness properties.

Abstract

Planning with world models offers a powerful paradigm for robotic control. Conventional approaches train a model to predict future frames conditioned on current frames and actions, which can then be used for planning. However, the objective of predicting future pixels is often at odds with the actual planning objective; strong pixel reconstruction does not always correlate with good planning decisions. This paper posits that instead of reconstructing future frames as pixels, world models only need to predict task-relevant semantic information about the future. For such prediction the paper poses world modeling as a visual question answering problem about semantic information in future frames. This perspective allows world modeling to be approached with the same tools underlying vision language models. Thus vision language models can be trained as "semantic" world models through a supervised finetuning process on image-action-text data, enabling planning for decision-making while inheriting many of the generalization and robustness properties from the pretrained vision-language models. The paper demonstrates how such a semantic world model can be used for policy improvement on open-ended robotics tasks, leading to significant generalization improvements over typical paradigms of reconstruction-based action-conditional world modeling. Website available at https://weirdlabuw.github.io/swm.

Figures (11)

• Figure 1: Comparison between Vision-Language Models, Video World Models, and Semantic World Models. While Vision-Language Models answer questions about static observations and Video World Models predict future observations given actions, Semantic World Models take observations and actions as input to directly answer questions about the future outcomes of those actions.
• Figure 2: Overview of Semantic World Models. SWM is a VLM adapted to answer questions about the future realized by the actions used to condition the model. Using a set of questions and desired answers, its predictions can be converted into a planning signal and iteratively refine the action sequence.
• Figure 3: Example state entry in the SAQA dataset with two action horizons and six QA pairs.
• Figure 4: Examples of each evaluation task. The top frame represents the initialization, and the bottom frame represents task completion. The first three tasks are for LangTable and the last two are for OGBench.
• Figure 5: Policy Improvement across LangTable and OGBench across multiple tasks. The average success rates of the base policies ($14.4\%$ on LangTable and $45.33\%$ on OGBench) increase to $81.6\%$ and $76.0\%$, respectively. SWM further outperforms the IDQL and AVD baselines across all evaluated tasks and environments. Reported success rates over $n=50$ seeds with 95% confidence intervals (normal approximation).