world2vec: Stable inference of mechanistic world models
Abstract
World models learned from raw pixels have become a practical basis for planning and control, and recent joint-embedding predictive architectures can be trained stably end-to-end. Yet these models are typically evaluated in environments whose dynamics are largely driven by the agent’s actions, and stable training alone does not ensure that learned latent representations recover the underlying state or dynamics of the environment. Here, we introduce world2vec, a framework for learning mechanistic world models directly from observations. We show that world2vec recovers the ground-truth state under nonlinear observation mixing and learns latent dynamics that faithfully reproduce the true system over long prediction horizons while using one order of magnitude fewer parameters than existing world models. We show state of the art performance in controlling continuous control environments from pixels. world2vec treats the representation objective and dynamics model as hyperparameters, allowing their inductive biases to be automatically matched to the data-generating process. This enables recovery of ground-truth dynamical systems and their governing equations across simple physical environments and a dataset of 71 strange attractors and paves the way towards using world2vec for scientific discovery.
est. 32% chance this paper gets accepted at ICLR 2027.
What do you think this paper will get?
All positions stay anonymous.