world2vec: Mechanistic world models through latent dynamics learning

Tobias Schmidt; Rodrigo Gonzalez Laiz; Krisztián Szuppinger; Steffen Schneider

cs.LGOct 29, 2026

world2vec: Mechanistic world models through latent dynamics learning

Authors: Tobias Schmidt, Rodrigo Gonzalez Laiz, Krisztián Szuppinger, Steffen Schneider

Organizations: Helmholtz Munich & Munich Center for Machine Learning (MCML)

Abstract

World models can achieve strong control performance without faithfully capturing the dynamics of their environment, limiting their potential for mechanistic interpretability and scientific discovery. Here, we introduce world2vec to learn mechanistic world models without reward information directly from observations and actions. World2vec learns latent dynamics that faithfully reproduce the true system over long prediction horizons, with at least one order of magnitude fewer parameters than DreamerV3, TD-MPC2 and LeWM. This enables state of the art performance in controlling continuous environments from pixels and accurate system identification across a benchmark of 71 attractors, demonstrating how world models are amenable to scientific discovery.

Figures & tables

Figure 1: A. World2vec jointly learns representations and dynamics in latent space; the dynamics model is an action-conditioned neural ODE which can be parameterized to provide an interpretable vector representation of the dynamics. The objective function incorporates assumptions about the underlying data-generating process. B. Reward vs. 1 s Equivariance Accuracy (EQA) of world2vec vs. baselines, averaged across environments.

Figure 2: Pendulum, Reacher and CartPole environments used in our study with their physical variables.

World Model Learning

World model learning builds predictive models of how environments change under actions to support planning, simulation, and decision-making.

See all 419 papers

CausalDreamer: Learning Predictive World Models with Latent Disentanglement
Oct 8, 2026Prince Jha, Nils Lukas, Kun Zhang +1World Model Learning Causal World Models
Mohamed bin Zayed University of Artificial Intelligence (MBZUAI) · Carnegie Mellon University
stable-worldmodel: A Platform for Reproducible World Modeling Research and Evaluation
May 20, 2026Lucas Maes, Quentin Le Lidec, Luiz Facury +9World Models ML Reproducibility
Mila & Université de Montréal · New York University · Universidade Federal de Minas Gerais +4
DWM: Separating World Effects from Actions in Latent World Models
Jul 21, 2026Yi-Ge Zhang, Tianqi Du, Qi Zhang +1World Model Learning Latent Dynamics Modeling
Peking University

Oct 8, 2026J/K move · Enter open · S save

CausalDreamer: Learning Predictive World Models with Latent Disentanglement

Prince Jha, Nils Lukas, Kun Zhang, Salem Lahlou

Mohamed bin Zayed University of Artificial Intelligence (MBZUAI) · Carnegie Mellon University

World Model Learning Causal World Models Disentangled Representation Learning Model-Based Planning Causal Representation Learning

Open paper AskPDF arXiv

World models for control must capture which aspects of the environment respond to the agent's actions and which are relevant to reward. Generative world models such as Dreamer 4 consist of a video tokenizer, which encodes each frame into a latent, and a dynamics model, which is pretrained to predict future latents from past latents and actions. Yet the tokenizer is trained with a reconstruction objective, without action or reward supervision, so its latent provides no explicit mechanism to separate controllable, uncontrollable, reward-relevant, and reward-irrelevant information. We propose CausalDreamer, which keeps the tokenizer frozen and re-encodes its latent into a factored representation of four groups along two axes: controllability, where only the two controllable groups receive the action, and reward relevance, learned by predicting the reward from the two reward-relevant groups. The pretrained dynamics model is then fine-tuned to predict the factored representation. We evaluate CausalDreamer and the pretrained world model it starts from with model-predictive planning on 20 MMBench2 tasks: 10 clean tasks seen during training and 10 unseen tasks, of which 6 are manipulated variants of clean tasks with a changed background, object, or maze layout, and 4 are new environments. We normalize returns so that a policy taking uniformly random actions scores 0 and an expert scores 1. CausalDreamer achieves a 14% higher normalized score than the pretrained world model on the clean tasks (0.199 vs.\ 0.175) and a 25% higher score on the manipulated variants (0.307 vs.\ 0.246), while neither model scores meaningfully above the random policy in the new environments. Additionally, our analysis shows that the factored representation separates reward-irrelevant changes, such as a changed background, from its reward-relevant groups.

Abstract

World Model Learning

Explore similar work

CausalDreamer: Learning Predictive World Models with Latent Disentanglement

stable-worldmodel: A Platform for Reproducible World Modeling Research and Evaluation

DWM: Separating World Effects from Actions in Latent World Models