Einstein World Models Munachiso Samuel Nwadike1,2 , Zangir Iklassov1 , Ali Mekky1 Zayd M. Kawakibi Zuhri1 , Kentaro Inui1,2,3 1 MBZUAI 2 RIKEN AIP, Japan 3 Tohoku University [email protected]
arXiv:2606.26969v1 [cs.AI] 25 Jun 2026
Abstract
Does intelligence require the ability to reason about phenomena beyond direct experience? It is natural to suspect that some complex thought cannot be captured through language alone. However, of particular concern to this work, is whether visualising counterfactual events can complement language as a mechanism for complex thought. We ask whether LLMs can be trained to utilise such visualisation mechanisms, in a way that benefits their reasoning abilities. Motivated by this question, we propose Einstein World Models. EWMs are a blueprint for LLM-based reasoning systems that place visual-temporal rollouts inside the reasoning trace, allowing them to reason in ways that text alone may not support well. In an EWM, the LLM calls a world-module (not to be confused with a world model), to produce short rollouts of scenes under consideration. The returned rollout is treated not as the answer, but as an inspectable hypothesis that can support later reasoning. Einstein World Models extend the capability of LLMs for tool calling (such as web search or code execution), into the domain of visual thought experiments.
Prologue A call for datasets becomes meaningful after the desirable capability has been specified. This culminates in Section 5. The predominance of this work therefore aims to motivate a learnable format for visual thought experimentation, proposing the architecture and training objectives necessitating their training data. This work may therefore be understood as an operationalisation of a promising capability whose data requirements are, per this moment, a conceivably fertile frontier.
1
Introduction
Scientific invention often makes abstract ideas tractable by turning them into thought experiments. In a thought experiment, we visualise a scene, let it unfold, and notice what changes. Einstein’s recollection of special relativity begins with precisely such a thought experiment, later popularised through the image of Bern’s famous clock tower. In his Autobiographical Notes (Einstein, 1949), he recalls imagining what it would be like to chase a beam of light. If he could accelerate until he matched the speed of the beam, would it eventually appear to hang motionless be-
Figure 1: Einstein World Models (proposed) build upon traditional LLM reasoning traces. However, in addition to generating tokens across N autoregressive steps, the model may, at a sparse set of M intermediate steps, invoke a callable world-module . The returned visual-temporal rollout becomes part of the trace as an inspectable hypothesis.
side him? In those notes, Einstein writes that such a stationary light wave seemed impossible both empirically and according to Maxwell’s equations. His simple thought experiment had transformed an abstract tension between electrodynamics and intuitions about motion, into a concrete scene upon which reasoning could be built. It is precisely in this spirit that “Einstein World Models” invoke “Einstein”. Vitally, the “E” in EWM also carries a dual reading, overloaded to mean “Externalised”. Externalisation brings the thought experiment into view as an inspectable reasoning trace component, and therefore, as a measurable one (Nwadike et al., 2026), as per chain-of-thought (Wei et al., 2022). LLMs can already search the web when they lack reliable facts (Nakano et al., 2021), run code when numerical calculation is needed (Gao et al., 2023), and call external tools when a task is easier to act out than to solve from tokens alone (Schick et al., 2023; Yao et al., 2023; Patil et al., 2024). What we propose is to render LLMs capable of imagining a scene, when quantitatively beneficial. Hadamard (Hadamard, 1945), in his inquiry into the inner mental processes behind mathematical invention, recorded Einstein’s self-description of the thought experiment process in strikingly imagistic terms: — Professor Albert Einstein (Hadamard, 1945) “The psychical entities which seem to serve as elements in [my] thought are certain signs and more or less clear images which can be ‘voluntarily’ reproduced and combined. . . . The above mentioned elements are, in my case, of visual and some of muscular type. Conventional words or other signs have to be sought for laboriously only in a secondary stage, when the mentioned associative play is sufficiently established and can be reproduced at will.”
In a thought experiment, the combination of visuals, albeit subject to meaning, takes precedence, while symbols and words play an augmentative role. Although LLMs already excel with words and symbols in chain-of-thought, their visual-temporal reasoning faculties are yet to capture this dynamic (kindly see Section 4). Thus, in an effort towards precisely this dynamic, prior work has popularised world models (LeCun et al., 2022). These works often characterise world modelling as either passive visual prediction from experience
(Bardes et al., 2024; Garrido et al., 2025), or actionconditioned prediction for embodied agents (MurLabadia et al., 2026; Nam et al., 2026; Maes et al., 2026). However, if world models are treated mainly as high-fidelity simulators of futures tethered either to observed states or to chosen actions, then thought experiments become correspondingly limited either to experience, or to intervention (see Supplementary Notes, Part C). Einstein World Models reconsider this correspondence. Much like Einstein could visualise riding beside a beam of light without first having to ride one, a language model should be able to benefit from imagining how a described scene could unfold, without necessarily acting within it. A decisive element in Einstein World Models is the choice of a world-module capable of producing useful video thought experiment rollouts, as discussed in Section 3. Equally central is the ability to integrate these rollouts back into the LLM’s reasoning trace, as discussed in Section 2. Our analysis of this element forms one of three contributions: First, we propose Einstein World Models as a mechanism for selective visualtemporal thought experiments instantiated by tooluse behavior in LLM reasoning. Second, we distinguish Einstein World Models as reasoning systems, from the world-modules they call, treating generated rollouts as inspectable intermediate artifacts. Third, we offer dataset and training recommendations for training LLMs into Einstein World Model reasoners, capable of utilising visualtemporal thought-experiment traces.
2
Einstein World Models
2.1
Overview
The objectives of Einstein World Models are twofold. First, we wish to imbue language models with the ability to construct visualisations when answering questions that require physical intuition or scene-level visualisation. Second, we wish to obtain a window into those visualisations. An EWM rollout, as illustrated in Figure 1, is in effect, a visible hypothesis about how a described scene might unfold. Let T denote a reasoning trace, and let N denote the number of autoregressive text-generation steps in that trace. A standard CoT trace uses these N steps to generate text alone. An Einstein World Model also generates text autoregressively, but at a sparse set of M intermediate generation steps,
Answer
Answer Thoughts
LLM
×N
LLM
“If I chased light, would it stand still?”
Vision Encoder
Question
“If I chased light, would it stand still?”
Question Visuals
CoT
VLM
Predictor
—
Visuals
Minimise Difference
Answer “If I chased light, would it stand still?”
Vision Encoder
Text Encoder
Thoughts World- module
Question
×M
×N
Answer
LLM Visuals “If I chased light, would it stand still?”
Question
VL-JEPA
EWM
Figure 2: If Einstein were travelling away from the Bern clock tower at the speed of light, would the clock’s hands appear frozen in time? This figure compares how different reasoning paradigms would support such a thought experiment. Top left: In CoT, the LLM reasons only through N autoregressive text-generation steps. Top right: In VLMs, visual information can condition generation, but the relevant visuals are externally provided, rather than visualised during the reasoning trace. Bottom left: In VL-JEPA-style world models, predictive visual representations are learned from observed visual inputs, but the visual prediction itself is not selectively invoked as an intermediate reasoning artifact. Bottom right: In Einstein World Models, by contrast, the LLM (the reasoner) preserves autoregressive text generation while taking responsibility for deciding when to invoke a world-module M times, how to query it, and how to incorporate each returned rollout into subsequent reasoning. Each rollout therefore enters the trace as a visible thought experiment rather than as a pre-given input or final answer.
queries a world-module. Each call to the worldmodule produces a short video sequence, which is then returned to the reasoner and used to condition subsequent generation. Figure 2 contrasts this structure with CoT, VLMs, and VL-JEPA-style world models. For problems in which Einstein World Models
are useful, we expect that 0 < M ≪ N − 1, as the purpose of a rollout is not to replace the chain-of-thought, but to externalise a visualised scene, at moments where doing so facilitates subsequent reasoning.
Once returned, the rollout becomes part of the reasoning trace. The choice of prompt used to query the world-module is therefore part of the reasoning problem as well. As with web search, the query helps determine what information is returned. In Einstein World Models, the reasoner must learn not only when to visualise, but how to query for the visualised scene that will support later reasoning. 2.2
Rollouts as Inspectable Hypotheses
Externalised thought makes latent assumptions visible (Nwadike et al., 2026). Einstein World Models extend this principle from language to visualtemporal reasoning. Instead of leaving the model’s visualised scene implicit, the system renders that scene as an examinable video artifact, for instance as a sequence of frames. Because this rollout is externalised, it can be shared and studied without access to model weights or hidden activations. In this sense, Einstein World Models turn an otherwise private visualised episode into a public object of analysis, even when the underlying model is not itself open to inspection. Key here is the distinction between the inspectability of the world-module’s rollout, and the physical plausibility of the rollout. A rollout is not considered more plausible solely because it is inspectable. Furthermore, the visual-temporal rollout need not begin from physically ordinary premises in order to be informative. Einstein’s own light-chasing scene, for example, was not valuable because it was itself a realisable experiment. To the contrary, it was conspicuously counterfactual. However, it was valuable because it made a counterintuitive possibility precise enough to reason about. Similarly, we neither treat the rendered visibility of an EWM rollout as evidence of its plausibility, nor its imperfections as evidence that it is uninformative. What matters is whether the hypothesis it exposes can be inspected, tested, and improved. 2.3
Inference
At inference time, an Einstein World Model is specified by two core components πθ
and
W.
Here, πθ denotes the Einstein reasoner, a trainable LLM policy parameterised by θ, and W denotes the world-module. πθ generates a reasoning trace and may query W for visual-temporal rollouts. If queried, W generates these video rollouts, and returns them to πθ for further autoregression.
For a text-only input problem x, inference constructs a thought trace with the initialisation T0 = x, such that Tt then denotes the partial trace after step t. At each step, the Einstein reasoner defines a conditional distribution over the next model-generated segment, πθ (· | Tt ). A generated segment is either a non-tool segment st , such as language reasoning or a final answer, or a world-module query segment qt . If a query qt is generated, the world-module returns a visualtemporal rollout vt ∼ W(qt ). The trace update is therefore ( st , if W is not queried, Tt+1 = Tt ⊕ (1) [qt , vt ], if W is queried. Thus, ordinary text segments are appended directly to the trace, while world-module calls append the reasoner-generated query, concatenated with the returned rollout observation. If the generated segment is a final answer, inference terminates and returns that answer. In implementation, these trace segments can be serialised with special tags, formatted similarly to recent RL-based tool-use systems for search and agentic tool interaction (Jin et al., 2025; Singh et al., 2025). An EWM trace may use <think>...</think> for language reasoning, <tool_call>{“name”: “world_module”, “query”: q}</tool_call> for a world-module query, <visual_rollout>... </visual_rollout> for the returned visualtemporal rollout, and <answer>...</answer> for the final answer. The <visual_rollout> segment is returned by W. In practice, W may render a short sequence of frames, which can be encoded into visual tokens for incorporation into the reasoning trace. The rendered frames remain available for inspection, while the visual tokens provide the representation consumed by the reasoner. 2.4
Training
At the initial stage, the reasoner undergoes supervised fine-tuning on EWM trace formats with standard next-token cross-entropy, masking returned
rollout observations from the loss as detailed in the Supplementary Notes, part A. Consistent with standard practice in recent reasoning-model training (Guo et al., 2025), we then recommend following SFT with RLVR-style training over complete EWM trajectories, since target tasks provide verifiable final answers even when intermediate visual thought experiments remain unlabeled. RL-based tool use already provides a standard methodology for training LLMs to call fixed external systems, such as search engines, and incorporate their returned outputs (Jin et al., 2025; Qian et al., 2025; Singh et al., 2025). Formally, let D = {(xi , yi⋆ )}ni=1 be a dataset of text-only problems xi with verifiable final answers yi⋆ . A sampled EWM trajectory is a completed thought trace T containing all language-reasoning segments, world-module queries, returned rollouts, and the final answer ŷ. Let r(ŷ, y ⋆ ) denote the verifier reward for the final answer. For exact-answer tasks, this may simply be r(ŷ, y ⋆ ) = 1[ŷ = y ⋆ ]. Since world-module calls may be computationally expensive and should be used selectively, we define an EWM reward that combines final-answer correctness with an optional additional reward term for world-module usage. rM (T , y ⋆ ) = r(ŷ, y ⋆ ) + rW (T ).
(2)
Here, rW (T ) denotes the optional implementationdependent reward for world-module calling behavior in the trace. We address this reward in Section 2.4.1. A GRPO-style implementation (Shao et al., 2024) can optimise this reward with a clipped surrogate. For each training pair (x, y ⋆ ) ∼ D, sample G complete EWM rollouts τ1:G = {τi }Gi=1 using the frozen old reasoner πold with access to W. Let ri = rM (τi , y ⋆ ). The group-relative advantage of rollout i is Ai =
ri − r̄ , sr + ϵadv
where r̄ and sr are the mean and standard deviation of the ri within τ1:G . Let τi = (zi1 , . . . , ziLi ) denote the serialized trajectory, where Li is its length and zit is the token at position t. Also let 1it = 1 for policy-generated tokens and 1it = P0Lifor returned rollout observations. g Define Li = t=1 1it , and, for generated tokens, ρit =
πθ (zit | τi,<t ) . πold (zit | τi,<t )
Writing clipϵ (r) = clip(r, 1 − ϵ, 1 + ϵ), the EWM training objective is " JE (θ) = Ex, τ1:G
G
L
i=1
t=1
i 1X 1 X 1it G Lgi
min(ρit Ai , clipϵ (ρit )Ai ) # − βDKL (πθ ∥πref ) . (3) Here, πref is the original pre-RL LLM used as a reference policy, and β ≥ 0 controls the KL penalty that discourages drift from it. Feasibility. Prior work suggests that language models can acquire structured representations of space, time, colour, and other real-world variables from language alone (Gurnee and Tegmark, 2024; Huh et al., 2024). The LLM already contains much of the world knowledge needed to query a worldmodule effectively. Thus, the remaining task is not necessarily fresh pretraining from scratch, but targeted post-training. Furthermore, autoregressive reasoning and video rollouts share a forward temporal structure. A chain of thought advances token by token, while a visualised scene advances frame by frame. The unidirectional nature of autoregression is surprisingly robust in practice (Nwadike et al., 2025). It is also reflected in models that combine autoregressive prediction with latent diffusion sampling (Parker-Holder et al., 2024). 2.4.1
Selective Thought Experiments
The rW term is intended to encourage selective world-module use. If rW (T ) = 0, then selectivity can be learned solely through the final-answer reward, as in Jin et al. (2025) and Qian et al. (2025). Emerging evidence, however, suggests that shaping rewards beyond final-answer correctness can encourage more specific behaviours, from improved reasoning to more efficient tool use (Guo et al., 2025; Chen et al., 2026; Wang et al., 2025). For example, one may set rW (T ) = −λM (T )/B, where M (T ) counts world-module calls, B is a call budget, and λ ≥ 0 controls the penalty for excessive calling. Upon choice of rW , Algorithm 1 summarises how the resulting reward rM may be utilised during training. Optimising JE in Eq. 3 converts this tracelevel reward into group-relative advantages over complete EWM trajectories. Traces where a visual thought experiment proves useful enough to justify its world-module usage receive higher advantages,
while unnecessary or unhelpful calls receive lower advantages. The reasoner therefore learns when to query W and how to use the returned rollout. Algorithm 1 Proposed EWM RLVR protocol Require: Dataset D, reasoner πθ , reference policy πref , world-module W , group size G 1: for each GRPO update round do 2: Set frozen old policy πold ← πθ 3: Sample training batch from D 4: for each (x, y ⋆ ) in the batch do 5: Sample G trajectories τ1:G using πold with access to W 6: for each trajectory τi do 7: Generate <think>, <tool_call>, or <answer> segments 8: if the world-module is invoked then 9: Call W with the query in <tool_call> 10: Append rollout as a <visual_rollout> segment 11: end if 12: Compute reward ri = rM (τi , y ⋆ ) 13: end for 14: Compute group-relative advantages Ai 15: end for 16: Compute importance weights ρit over reasoner-generated tokens 17: Mask returned visual-rollout observations from the policy loss 18: Update πθ by maximising JE in Eq. 3 19: end for
3
World-Module Selection
Since RLVR optimises the reasoner’s use of the world-module rather than the world-module itself, the question becomes how to select a module that produces useful visual-temporal thought experiments. 3.1
Architecture
World-modules may come in several forms. We modify recent informal taxonomies (Li and Labs, 2026) of world models to distinguish candidate world-modules: 1. Renderers: return observations, such as images or video frames. In EWMs, renderers are the default world-modules. Text-to-video generators, or image-to-video generators used within a text-to-image pipeline, are of particular interest. This includes (but is not limited to) diffusion and flow-matching video models, provided their rollouts expose information that can support later
reasoning (Kong et al., 2024; Wan Team et al., 2025; HaCohen et al., 2024). 2. Simulators: allow the reasoner to intervene in a visualised world and observe what follows. Interactive world models such as Genie-style systems provide one example of this interface (Bruce et al., 2024; Parker-Holder et al., 2024), and benchmarks such as WBench suggest that the coherence of such interactive rollouts can be measured (Ying et al., 2026). However, in practice, repeated renderer calls may play the role of a simulator, making a renderer (aforementioned) the primary world-module of interest. In particular, the LLM can inspect one visualised consequence, revise its hypothesis, and then request another rollout from a modified condition or counterfactual premise. Simulators are therefore useful when a thought experiment requires explicit intervention, but they are not required for the central Einstein World Model mechanism. 3. Planners: In Einstein World Models, planning remains the role of the LLM reasoner, since the aim is visual-temporal reasoning over thought experiment sequences, rather than embodied robotic action. 3.2
World-Module Quality
Einstein’s visual thought experiments were disciplined by strong physical intuition. A generated video rollout may therefore only be as useful as the physical intuitions sustained by its underlying video model. Diffusion models remain strong candidates because their intuitive-physics quality can be measured through human-verifiable likelihood estimates derived from the denoising objective of the diffusion model. Yuan et al. (2025) use precisely this technique to find substantial differences between video diffusion models, with stronger performance from recent systems. Some video generators remain unreliable physical simulators (Bansal et al., 2025; Zhang et al., 2025; Motamed et al., 2026). However, not all diffusion models are equal, and the frontier is fast-improving. Recent physicsaware generators further suggest that explicit dynamical priors can improve physical consistency and controllability (Yuan et al., 2026). Einstein World Models also require rollout faithfulness. A faithful rollout exposes information that the reasoner actually uses and that the final answer
depends on. An unfaithful rollout, much like in traditional chain-of-thought, may not reveal the computation that produced the answer (Turpin et al., 2023; Lanham et al., 2023). The response is not to abandon the need for video rollout traces, but to evaluate and improve their faithfulness (Nwadike et al., 2026). 3.3
Ensembling
Einstein World Models already allow repeated world-module calls within a single reasoning trace. Here, ensembling extends beyond repeated calls to a single module to describe the comparison of visual hypotheses generated by different worldmodules, much as different humans may visualise the same problem differently. One module may favour visual realism, another physical consistency, and another temporal continuity. Because rollouts are externalised, these assumptions can be compared rather than left hidden. The problem is therefore not only to select a strong world-module, but to select modules whose inductive biases are usefully different. In ensembling, different LLM reasoners, attached to different world-modules, exchange rollouts and critiques before answering. Each rollout proposes a different interpretation of the scene, and disagreement reveals what must be inspected next.
4
Related Work
Chain-of-thought (Wei et al., 2022) made intermediate reasoning visible, but made it visible only as language. Many commonsense questions depend on variables that text traces represent poorly, including object identity, containment, contact, heat, motion, and material state. Humans often reason about such variables through visualisation. Einstein World Models seek to provide an analogous capability for language models by treating visualised visual episodes as intermediate reasoning artifacts. Whiteboard-of-Thought (Menon et al., 2024) is a notable procedural predecessor for Einstein World Models. It gives a multimodal model a visual scratchpad, asks it to draw intermediate reasoning steps as an image, often through code, and then feeds that image back into the model for final reasoning. Einstein World Models preserve this feedback loop, but change the artifact from a static drawing to a short video rollout that serves as a visual play-through of a hypothesis.
Hu et al. (2024) similarly give multimodal LMs an external visual sketchpad for intermediate reasoning. However, the resulting artifacts are mainly static visual annotations (auxiliary lines, bounding boxes, segmentation masks, etc.) produced through programmatic tools, rather than visual-temporal rollouts. Wu et al. (2024) propose Visualization-ofThought (VoT), a related method for spatial reasoning in LLMs. However, VoT focuses on controlled 2D grid-world tasks rather than real-world visualtemporal thought experiments. Its visualisations remain text-form grids or maps, not separate visualtemporal rollouts. Thus, VoT is best understood as a prompting strategy for symbolic state tracking, whereas the aim of Einstein World Models is to let a reasoner call a world-module and incorporate the resulting rollout into its trace. In other work, Chern et al. (2025) showed that when generating an image from a prompt, intermediate visual subgoals can guide the model toward a better final image. However, their work focused on the objective of image generation as the end goal, rather than enhanced reasoning as the end goal (the latter being merely facilitated by visualisation as a means to an end). Tong et al. (2025) study whether video generation models can reason by producing answerbearing videos. Einstein World Models pursue a different research objective. Rather than asking whether the video generator can serve as the reasoner, EWMs ask whether an LLM can use video generation as a thought experiment tool. In other words, rather than replacing the LLM with a video generator, we ask how frontier LLMs can decide when visualisation is useful, selectively invoke the world-module, and integrate the resulting rollout back into reasoning. Yang et al. (2025) and Yu et al. (2026) explore the importance of visualisation for 3D spatial reasoning in VQA-style settings, where the model begins from an observed image and generates additional views to answer questions about that scene. This differs from the setting of interest here, where thought experiments are used to support text-based reasoning problems and the model must decide, as part of its own reasoning process, when visualtemporal rollouts are useful. We discuss the relationship between Einstein World Models and contemporary agentic LLM systems in the Supplementary Notes, Part B.
5
Future Work: A Call for Datasets
A central bottleneck for Einstein World Models is data. Few existing datasets explicitly target the behaviour Einstein World Models require. The missing setting is neither ordinary text reasoning, nor visual question answering over an already provided image. It is a setting in which an LLM is fully capable of taking in text alone as input (even without the aid of visual inputs, when requested), and can perform a visual thought experiment in its reasoning trace before outputting final answer tokens. This makes dataset construction the immediate experimental bottleneck for Einstein World Models, since the core learning problem still lacks a suitable benchmark. SimpleBench is one of the rare public datasets pointing in this direction. However, the full dataset contains only a little over 200 questions, while its public release exposes only 10, making it useful as an illustration, rather than as a training corpus. The questions in SimpleBench are short and text-only, yet difficult, because answering them correctly often depends on performing a thought experiment about how a described scene unfolds. Consider the following example:
for ordinary solid balls to land. An EWM rollout would externalise such a missing visual-temporal computation, allowing it to become part of the reasoning trace. Existing physical reasoning datasets are valuable, but many begin with the relevant scene already available. Some focus on two-dimensional puzzle images (Bakhtin et al., 2019). Others ask models to explain, predict, or judge physical events in supplied video clips (Yi et al., 2020; Bear et al., 2021; Riochet et al., 2018; Bordes et al., 2025). Still, others focus on the physical fidelity of video generators. Even then, the visual scene is supplied in advance, either as initial frames to complete, or as rendered samples to score (Upadhyay et al., 2026; Yuan et al., 2025). By contrast, Einstein World Models need datasets where the problem begins as language, and the model must decide, for itself, whether to generate a visualisation. We therefore make an open call for datasets in this setting. Ideally, such datasets should contain both problems that benefit from visual thought experiments, and problems that do not, allowing Einstein World Models to learn not only how to visualise, but also when not to.
SimpleBench (Philip and Hemang, 2024) Test Sample #2
Epilogue
Question: A juggler throws a solid blue ball a meter in the air and then a solid purple ball (of the same size) two meters in the air. She then climbs to the top of a tall ladder carefully, balancing a yellow balloon on her head. Where is the purple ball most likely now, in relation to the blue ball? A. at the same height as the blue ball B. at the same height as the yellow balloon C. inside the blue ball D. above the yellow balloon E. below the blue ball F. above the blue ball
This paper proposed Einstein World Models, for treating visual thought experiments as a tool-use behaviour in LLM reasoning. The central idea is that some questions require a model to visualise how a described scene unfolds, and that this may be poorly supported by language alone. EWMs keep the LLM as the reasoner, but allow it to call a world-module, render a short visual-temporal rollout, and return that rollout to the reasoning trace before answering. A video rollout is not assumed to be a perfect simulation, but rather a visual hypothesis about how a described situation may unfold. Its distinctive usefulness comes from being externalised, making the model’s reasoning process available for inspection and debugging. The path forward thus requires both better world-module curation, as well as better datasets to allow for training LLMs on when and how thought experiments should be invoked. In this sense, Einstein World Models point toward language models that reason not only through words and symbols, but through externalised visual walk-throughs.
Correct Answer: A.
What makes this question difficult for LLMs is that it requires visualising the scene, and realising that, although the purple ball is thrown higher than the blue ball, enough time passes for both solid balls to fall back down before any juggler could possibly finish climbing a ladder while carefully balancing a balloon on their head. A text-only LLM may over-formalise the prompt and conclude that the answer depends on unspecified variables such as launch velocity or climbing time. However, a one- or two-metre throw lasts only moments, while carefully climbing a tall ladder takes long enough
Supplementary Notes A
Supervised Finetuning
Before reinforcement learning, the Einstein reasoner may be warm-started with supervised finetuning on valid EWM trace formats using a standard cross-entropy loss. This stage teaches the syntax and role structure of EWM reasoning traces, including ordinary language reasoning, worldmodule query segments, returned visual-rollout observations, and final-answer segments. For example, traces may use <think> for language reasoning, <tool_call> for world-module queries, <visual_rollout> for returned observations, and <answer> for final answers. Let DSFT = {(xi , Ti⋆ )}i be a supervised dataset of input problems and target EWM traces. Each trace Ti⋆ contains both reasonergenerated segments and observation segments returned by the world-module. The model is trained with standard next-token cross-entropy, but only on tokens generated by the reasoner. Returned visual rollouts are observations, not policy actions, and are therefore masked from the supervised loss. Let zt⋆ denote the target token at position t in the serialized trace T ⋆ . Let 1t = 1 if zt⋆ is a target token that the reasoner is expected to produce, and let 1t = 0 if zt⋆ is part of a returned rollout observation. The masked supervised fine-tuning loss is: LSFT (θ) = −
E
(x,T ⋆ )∼DSFT
1 P
t 1t
X
1t log πθ zt⋆ | T<t
⋆
! .
t
Thus, supervised fine-tuning prepares the Einstein reasoner πθ for RLVR by teaching it the format of valid EWM traces, while RLVR teaches when and why such traces should contain visual thought experiments. Crucially, DSFT should include both call and no-call traces, so that the model learns the world-module interface without learning to invoke W by default.
B
Relation to Agentic Systems
Modern LLM agents are commonly understood as language models augmented with additional machinery for tool use, planning, environment interaction, or action execution on behalf of a user (Yao et al., 2023; Shen et al., 2023). However, although
an EWM can be incorporated into an agentic system, it is not necessarily an agent in itself. Such systems may qualify as Einstein World Models when they use a world-module to generate visualtemporal rollouts as intermediate reasoning artifacts, but not merely by virtue of being agents. Existing agentic systems can call tools, but they do not by default perform externalised visual thought experiments in the sense proposed here. Conversely, an Einstein World Model need not have the broader infrastructure of a general-purpose agent for executing user-facing actions, such as managing software repositories or querying databases. It may be implemented as a language reasoner that calls a world-module only to imagine a described scene, inspect the resulting rollout, and incorporate it into its reasoning trace. In this sense, EWMs describe a reasoning capability that can either stand alone, or be added to agentic systems.
C
Relation to Other World Models
The term world model has a long history in AI, especially in model-based reinforcement learning, where learned models of environment dynamics support planning, control, and interaction with an environment (Sutton, 1991; Ha and Schmidhuber, 2018). In this lineage, a world model helps an agent anticipate possible futures before acting in the external world. More recent proposals similarly treat world models as central components of autonomous systems that learn to predict, reason, and plan over future states (LeCun et al., 2022). Einstein World Models shift the emphasis from the notion of a world model to that of a worldmodule. An EWM is not itself a learned dynamics model, simulator, or video generator. It is a reasoning system in which the LLM remains the reasoner and calls the world-module when a visual-temporal thought experiment may support its reasoning trace. The world-module is the component most closely aligned with what is traditionally the predictive component of a world model. It need not be a full 3D simulator, and is treated only as a component of a reasoning system. Its role in an EWM is to supply an inspectable rollout rather than to replace the LLM’s reasoning process. This distinction also shifts what must be learned. Whereas many world-model architectures focus on modelling environmental dynamics, EWMs focus on how a reasoner should selectively invoke visualisation to facilitate reasoning.
References Anton Bakhtin, Laurens van der Maaten, Justin Johnson, Laura Gustafson, and Ross Girshick. 2019. PHYRE: A new benchmark for physical reasoning. In Advances in Neural Information Processing Systems, volume 32.
Luyu Gao, Aman Madaan, Shuyan Zhou, Uri Alon, Pengfei Liu, Yiming Yang, Jamie Callan, and Graham Neubig. 2023. PAL: Program-aided language models. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 10764–10799. PMLR.
Hritik Bansal, Zongyu Lin, Tianyi Xie, Zeshun Zong, Michal Yarom, Yonatan Bitton, Chenfanfu Jiang, Yizhou Sun, Kai-Wei Chang, and Aditya Grover. 2025. VideoPhy: Evaluating physical commonsense for video generation. In International Conference on Learning Representations.
Quentin Garrido, Nicolas Ballas, Mahmoud Assran, Adrien Bardes, Laurent Najman, Michael Rabbat, Emmanuel Dupoux, and Yann LeCun. 2025. Intuitive physics understanding emerges from selfsupervised pretraining on natural videos. Preprint, arXiv:2502.11831.
Adrien Bardes, Quentin Garrido, Jean Ponce, Xinlei Chen, Michael Rabbat, Yann LeCun, Mido Assran, and Nicolas Ballas. 2024. Revisiting feature prediction for learning visual representations from video. Transactions on Machine Learning Research. Featured Certification.
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Peiyi Wang, Qihao Zhu, Runxin Xu, Ruoyu Zhang, Shirong Ma, Xiao Bi, and 1 others. 2025. DeepSeekR1 incentivizes reasoning in LLMs through reinforcement learning. Nature, 645:633–638.
Daniel M. Bear, Elias Wang, Damian Mrowca, Felix J. Binder, Hsiao-Yu Fish Tung, R. T. Pramod, Cameron Holdaway, Sirui Tao, Kevin A. Smith, Fan-Yun Sun, Li Fei-Fei, Nancy Kanwisher, Joshua B. Tenenbaum, Daniel L. K. Yamins, and Judith E. Fan. 2021. Physion: Evaluating physical prediction from vision in humans and machines. In Advances in Neural Information Processing Systems, Datasets and Benchmarks Track. Florian Bordes, Quentin Garrido, Justine T. Kao, Adina Williams, Michael Rabbat, and Emmanuel Dupoux. 2025. IntPhys 2: Benchmarking intuitive physics understanding in complex synthetic environments. Preprint, arXiv:2506.09849. Jake Bruce, Michael D. Dennis, Ashley Edwards, Jack Parker-Holder, Yuge Shi, Edward Hughes, Matthew Lai, Aditi Mavalankar, Richie Steigerwald, Chris Apps, Yusuf Aytar, Sarah Maria Elisabeth Bechtle, Feryal Behbahani, Stephanie C. Y. Chan, Nicolas Heess, Lucy Gonzalez, Simon Osindero, Sherjil Ozair, Scott Reed, and 6 others. 2024. Genie: Generative interactive environments. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 4603–4623. PMLR. Liuji Chen, Dianxing Tang, Xing Shi, Dingshuo Chen, Qiang Liu, Shu Wu, and Liang Wang. 2026. Learning when not to act: Mitigating tool abuse in agentic reinforcement learning. arXiv preprint arXiv:2606.02132. Ethan Chern, Zhulin Hu, Steffi Chern, Siqi Kou, Jiadi Su, Yan Ma, Zhijie Deng, and Pengfei Liu. 2025. Thinking with generated images. arXiv preprint arXiv:2505.22525. Albert Einstein. 1949. Autobiographical notes. In Paul Arthur Schilpp, editor, Albert Einstein: Philosopher-Scientist, pages 1–95. Open Court.
Wes Gurnee and Max Tegmark. 2024. Language Models Represent Space and Time. In International Conference on Learning Representations. David Ha and Jürgen Schmidhuber. 2018. Recurrent world models facilitate policy evolution. In Advances in Neural Information Processing Systems, volume 31. Yoav HaCohen, Nisan Chiprut, Benny Brazowski, Daniel Shalem, Dudu Moshe, Eitan Richardson, Guy Shiran, Nir Zabari, Ori Gordon, Poriya Panet, and 1 others. 2024. Ltx-video: Realtime video latent diffusion. arXiv preprint arXiv:2501.00103. Jacques Hadamard. 1945. An Essay on the Psychology of Invention in the Mathematical Field. Princeton University Press. Yushi Hu, Weijia Shi, Xingyu Fu, Dan Roth, Mari Ostendorf, Luke Zettlemoyer, Noah A. Smith, and Ranjay Krishna. 2024. Visual sketchpad: Sketching as a visual chain of thought for multimodal language models. In Advances in Neural Information Processing Systems 37. Minyoung Huh, Brian Cheung, Tongzhou Wang, and Phillip Isola. 2024. Position: The Platonic Representation Hypothesis. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 20617–20642. PMLR. Bowen Jin, Hansi Zeng, Zhenrui Yue, Jinsung Yoon, Sercan O. Arik, Dong Wang, Hamed Zamani, and Jiawei Han. 2025. Search-r1: Training llms to reason and leverage search engines with reinforcement learning. In Proceedings of the 2nd Conference on Language Modeling (COLM 2025). Weijie Kong, Qi Tian, Zijian Zhang, Rox Min, Zuozhuo Dai, Jin Zhou, Jiangfeng Xiong, Xin Li, Bo Wu, Jianwei Zhang, and 1 others. 2024. Hunyuanvideo: A systematic framework for large video generative models. arXiv preprint arXiv:2412.03603.
Tamera Lanham, Anna Chen, Ansh Radhakrishnan, Benoit Steiner, Carson Denison, Danny Hernandez, Dustin Li, Esin Durmus, Evan Hubinger, Jackson Kernion, Kamilė Lukošiūtė, Karina Nguyen, Newton Cheng, Nicholas Joseph, Nicholas Schiefer, Oliver Rausch, Robin Larson, Sam McCandlish, Sandipan Kundu, and 11 others. 2023. Measuring faithfulness in chain-of-thought reasoning. Preprint, arXiv:2307.13702. Yann LeCun and 1 others. 2022. A path towards autonomous machine intelligence version 0.9. 2, 202206-27. Open Review, 62(1):1–62. Fei-Fei Li and World Labs. 2026. A functional taxonomy of world models. https://www.worldlabs. ai/blog/taxonomy-of-world-models. Accessed: 2026-06-08. Lucas Maes, Quentin Le Lidec, Damien Scieur, Yann LeCun, and Randall Balestriero. 2026. Leworldmodel: Stable end-to-end joint-embedding predictive architecture from pixels. arXiv preprint arXiv:2603.19312. Sachit Menon, Richard Zemel, and Carl Vondrick. 2024. Whiteboard-of-thought: Thinking step-by-step across modalities. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 20016–20031, Miami, Florida, USA. Association for Computational Linguistics. Saman Motamed, Laura Culp, Kevin Swersky, Priyank Jaini, and Robert Geirhos. 2026. Do generative video models understand physical principles? In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision. Lorenzo Mur-Labadia, Matthew Muckley, Amir Bar, Mido Assran, Koustuv Sinha, Mike Rabbat, Yann LeCun, Nicolas Ballas, and Adrien Bardes. 2026. V-JEPA 2.1: Unlocking dense features in video selfsupervised learning. Preprint, arXiv:2603.14482. Reiichiro Nakano, Jacob Hilton, Suchir Balaji, Jeff Wu, Long Ouyang, Christina Kim, Christopher Hesse, Shantanu Jain, Vineet Kosaraju, William Saunders, Xu Jiang, Karl Cobbe, Tyna Eloundou, Gretchen Krueger, Kevin Button, Matthew Knight, Benjamin Chess, and John Schulman. 2021. WebGPT: Browser-assisted question-answering with human feedback. arXiv preprint arXiv:2112.09332. Heejeong Nam, Quentin Le Lidec, Lucas Maes, Yann LeCun, and Randall Balestriero. 2026. Causal-JEPA: Learning world models through object-level latent interventions. Preprint, arXiv:2602.11389. Munachiso S. Nwadike, Zangir Iklassov, Toluwani Aremu, Tatsuya Hiraoka, Benjamin Heinzerling, Velibor Bojkovic, Hilal AlQuabeh, Martin Takáč, and Kentaro Inui. 2025. Library-Like Behavior In Language Models is Enhanced by Self-Referencing Causal Cycles. In Proceedings of the 63rd Annual
Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 25365– 25377, Vienna, Austria. Association for Computational Linguistics. Munachiso Samuel Nwadike, Zangir Iklassov, Kareem Ali, Rifo Genadi, and Kentaro Inui. 2026. Measuring ai reasoning: A guide for researchers. Preprint, arXiv:2605.02442. Jack Parker-Holder, Philip Ball, Jake Bruce, Vibhavari Dasagi, Kristian Holsheimer, Christos Kaplanis, Alexandre Moufarek, Guy Scully, Jeremy Shar, Jimmy Shi, Stephen Spencer, Jessica Yung, Michael Dennis, Sultan Kenjeyev, Shangbang Long, Vlad Mnih, Harris Chan, Maxime Gazeau, Bonnie Li, and 13 others. 2024. Genie 2: A large-scale foundation world model. Google DeepMind Blog. Shishir G. Patil, Tianjun Zhang, Xin Wang, and Joseph E. Gonzalez. 2024. Gorilla: Large language model connected with massive apis. In Advances in Neural Information Processing Systems, volume 37. Philip and Hemang. 2024. Simplebench: The text benchmark in which unspecialized human performance exceeds that of current frontier models. Technical report. Accessed: 2026-05-18. Cheng Qian, Emre Can Acikgoz, Qi He, Hongru Wang, Xiusi Chen, Dilek Hakkani-Tür, Gokhan Tur, and Heng Ji. 2025. Toolrl: Reward is all tool learning needs. In Advances in Neural Information Processing Systems. NeurIPS 2025. Ronan Riochet, Mario Ynocente Castro, Mathieu Bernard, Adam Lerer, Rob Fergus, Véronique Izard, and Emmanuel Dupoux. 2018. IntPhys: A framework and benchmark for visual intuitive physics reasoning. Preprint, arXiv:1803.07616. Timo Schick, Jane Dwivedi-Yu, Roberto Dessì, Roberta Raileanu, Maria Lomeli, Eric Hambro, Luke Zettlemoyer, Nicola Cancedda, and Thomas Scialom. 2023. Toolformer: Language models can teach themselves to use tools. In Advances in Neural Information Processing Systems, volume 36, pages 68539–68551. Curran Associates, Inc. Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, Y. K. Li, Y. Wu, and Daya Guo. 2024. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. CoRR, abs/2402.03300. Yongliang Shen, Kaitao Song, Xu Tan, Dongsheng Li, Weiming Lu, and Yueting Zhuang. 2023. HuggingGPT: Solving AI tasks with ChatGPT and its friends in Hugging Face. In Advances in Neural Information Processing Systems, volume 36. Joykirat Singh, Raghav Magazine, Yash Pandya, and Akshay Nambi. 2025. Agentic reasoning and tool integration for LLMs via reinforcement learning. arXiv preprint arXiv:2505.01441.
Richard S. Sutton. 1991. Dyna, an integrated architecture for learning, planning, and reacting. SIGART Bulletin, 2(4):160–163. Jingqi Tong, Yurong Mou, Hangcheng Li, Mingzhe Li, Yongzhuo Yang, Ming Zhang, Qiguang Chen, Tianyi Liang, Xiaomeng Hu, Yining Zheng, Xinchi Chen, Jun Zhao, Xuanjing Huang, and Xipeng Qiu. 2025. Thinking with video: Video generation as a promising multimodal reasoning paradigm. Preprint, arXiv:2511.04570. Accepted to CVPR 2026. Project page: https://thinking-with-video.github.io/. Benchmark: https://huggingface.co/datasets/OpenMOSSTeam/VideoThinkBench. Miles Turpin, Julian Michael, Ethan Perez, and Samuel R. Bowman. 2023. Language models don’t always say what they think: Unfaithful explanations in chain-of-thought prompting. In Advances in Neural Information Processing Systems, volume 36, pages 74952–74965. Curran Associates, Inc. Rishi Upadhyay, Howard Zhang, Jim Solomon, Ayush Agrawal, Pranay Boreddy, Shruti Satya Narayana, Yunhao Ba, Alex Wong, Celso M. de Melo, and Achuta Kadambi. 2026. WorldBench: Disambiguating physics for diagnostic evaluation of world models. Preprint, arXiv:2601.21282. Wan Team, Ang Wang, Baole Ai, Bin Wen, Chaojie Mao, Chen-Wei Xie, Di Chen, Feiwu Yu, Haiming Zhao, Jianxiao Yang, and 1 others. 2025. Wan: Open and advanced large-scale video generative models. arXiv preprint arXiv:2503.20314. Hongru Wang, Cheng Qian, Wanjun Zhong, Xiusi Chen, Jiahao Qiu, Shijue Huang, Bowen Jin, Mengdi Wang, Kam-Fai Wong, and Heng Ji. 2025. Acting less is reasoning more! teaching model to act efficiently. arXiv preprint arXiv:2504.14870. Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed H. Chi, Quoc V. Le, and Denny Zhou. 2022. Chain-of-thought prompting elicits reasoning in large language models. In Advances in Neural Information Processing Systems, volume 35. Wenshan Wu, Shaoguang Mao, Yadong Zhang, Yan Xia, Li Dong, Lei Cui, and Furu Wei. 2024. Mind’s eye of LLMs: Visualization-of-thought elicits spatial reasoning in large language models. In Advances in Neural Information Processing Systems. Also available as arXiv:2404.03622. Yuncong Yang, Jiageng Liu, Zheyuan Zhang, Siyuan Zhou, Reuben Tan, Jianwei Yang, Yilun Du, and Chuang Gan. 2025. Mindjourney: Test-time scaling with world models for spatial reasoning. In Advances in Neural Information Processing Systems. Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik R. Narasimhan, and Yuan Cao. 2023. ReAct: Synergizing reasoning and acting in language models. In International Conference on Learning Representations.
Kexin Yi, Chuang Gan, Yunzhu Li, Pushmeet Kohli, Jiajun Wu, Antonio Torralba, and Joshua B. Tenenbaum. 2020. CLEVRER: Collision events for video representation and reasoning. In International Conference on Learning Representations. Kaining Ying, Hengrui Hu, Siyu Ren, Jiamu Li, Fengjiao Chen, Ziwen Wang, Xuezhi Cao, Xunliang Cai, and Henghui Ding. 2026. Wbench: A comprehensive multi-turn benchmark for interactive video world model evaluation. arXiv preprint arXiv:2605.25874. Shoubin Yu, Yue Zhang, Zun Wang, Jaehong Yoon, Huaxiu Yao, Mingyu Ding, and Mohit Bansal. 2026. When and how much to imagine: Adaptive test-time scaling with world models for visual spatial reasoning. arXiv preprint arXiv:2602.08236. Jianhao Yuan, Fabio Pizzati, Francesco Pinto, Lars Kunze, Ivan Laptev, Paul Newman, Philip Torr, and Daniele De Martini. 2025. Likephys: Evaluating intuitive physics understanding in video diffusion models via likelihood preference. arXiv preprint arXiv:2510.11512. Yu Yuan, Xijun Wang, Tharindu Wickremasinghe, Zeeshan Nadir, Bole Ma, and Stanley H. Chan. 2026. Newtongen: Physics-consistent and controllable textto-video generation via neural newtonian dynamics. In International Conference on Learning Representations. Chenyu Zhang, Daniil Cherniavskii, Andrii Zadaianchuk, Antonios Tragoudaras, Antonios Vozikis, Thijmen Nijdam, Derck W. E. Prinzhorn, Mark Bodracska, Nicu Sebe, and Efstratios Gavves. 2025. Morpheus: Benchmarking physical reasoning of video generative models with real physical experiments. Preprint, arXiv:2504.02918.