Einstein World Models

Einstein World Models (EWMs) address the limitation that complex reasoning—particularly involving physical intuition, motion, and counterfactual scenarios—is often difficult to capture through language alone. Inspired by Albert Einstein’s use of visual thought experiments (such as imagining chasing a beam of light), the authors propose a framework where an LLM acts as a reasoner that can selectively invoke an external 'world-module' to generate visual-temporal rollouts. Unlike standard vision-language models that process pre-existing images, EWMs allow the model to generate and inspect these visual hypotheses dynamically during its reasoning process.

The architecture treats the world-module as a tool, similar to web search or code execution. The LLM generates a query, receives a short video sequence (the rollout), and incorporates this visual data into its reasoning trace. This process is designed to be selective: the model learns to decide when a visual thought experiment is necessary and how to formulate the query to best support its final answer. By externalizing these visualizations, the system makes the model's latent assumptions about the physical world visible and inspectable.

The authors propose a two-stage training protocol. First, the model undergoes supervised fine-tuning on EWM trace formats. Second, it is optimized using Reinforcement Learning from Verifiable Rewards (RLVR). The reward function is designed to balance final-answer correctness with the efficiency of tool use, penalizing unnecessary or unhelpful calls to the world-module. The authors suggest that this approach leverages the existing world knowledge within LLMs, focusing on post-training the model to effectively coordinate between linguistic reasoning and visual simulation.