Skip to main content
Back to timeline
arXivSource publication:

world2vec learns world models without rewards: an order of magnitude fewer parameters, and systems read out across 71 attractors

Synopsis

world2vec learns, from observations and actions alone and without reward information, an encoder that maps physical state into latent coordinates plus latent dynamics describing how that state evolves. It requires the encoded future to match what rolling out the latent dynamics predicts, a condition computable from observations and actions alone. On Pendulum, Reacher and Cartpole it performs pixel-based continuous control with at least an order of magnitude fewer parameters than DreamerV3, TD-MPC2 and LeWM, and identifies systems across a benchmark of 71 attractors.

AI-generated editorial illustration: world2vec: Mechanistic world models through latent dynamics learning

Interpretation

World models can control well without faithfully reproducing the dynamics of their environment, which limits their use for mechanistic interpretability and scientific discovery; world2vec makes faithful dynamics reproduction the learning target rather than treating control score alone as success. Earlier routes include reconstruction-based forecasting, reconstruction-free planning in latent space trained with rewards, and reward-free joint-embedding predictive architectures; none of these guarantees that the learned latent dynamics match the true system. Across the three MuJoCo environments Pendulum, Reacher and Cartpole, world2vec jointly achieves high reward and good forecasting, and succeeds at model-predictive control in 2 of the 3 environments, none of which the previous best JEPA-based approach LeWM could solve.

The encoder maps a short history of observations to a latent state, and a dynamics model predicts the latent rate of change, which numerical integration steps forward. Training compares the encoded future observation with the latent trajectory obtained by rolling out from the starting point; both sides of this latent dynamics matching condition can be computed from observations and actions without true states or ground-truth dynamics. The main world2vec model has 263k parameters in total, 221k for the encoder and 42k for the dynamics; a small vision transformer encodes 16×16 patches and the dynamics is a 3-hidden-layer MLP of width 128, trained for 20k iterations.

The dynamics function class can be swapped for an interpretable structured form that yields explicit equations. The main experiments map latent state and action to the rate of change with an MLP; another instantiation follows SINDy, fitting dynamics as a fixed library of features with learned coefficients under a penalty. On a benchmark of 71 attractors, world2vec performs accurate system identification, presented as evidence that world models are amenable to scientific discovery.

Perspective

The demonstrated scope is offline training and latent-space model-predictive control on the three MuJoCo environments Pendulum, Reacher and Cartpole, plus system identification on a benchmark of 71 attractors; control experiments use 64×64-pixel frames from 200k random-policy samples, and goals are supplied as target observation sequences.

Identifiability guarantees are theoretical and hold under idealized conditions; whether they are reached in practice depends on numerical estimation. Consistency is measured as the difference between independently trained encoders after optimal affine alignment, and latent rollout consistency can be gamed by collapsed embeddings, so a retrieval-based equivariance accuracy is also used. The Reacher goal is itself underspecified, so both elbow configurations are planned and the average is reported.