world2vec: Mechanistic world models through latent dynamics learning

AI summary

Overview

Research area: Machine learning — world models, latent dynamics learning, model-based control, and mechanistic interpretability (identifiability theory for self-supervised representation learning).

Technical level: Advanced. The paper assumes familiarity with latent world models (Dreamer, TD-MPC2, JEPA/LeWM), neural ODEs, contrastive learning, and non-linear ICA-based identifiability arguments.

Scope: The paper introduces world2vec, a framework that jointly learns an encoder and an action-conditioned latent dynamics model from observations and actions alone (no rewards), evaluates it on three MuJoCo DeepMind Control Suite tasks for planning-based control, and reports system-identification results on a benchmark of 71 attractors.

What This Paper Is About

World models can score well on control tasks while still getting the underlying physics wrong, which limits their use for interpretability and scientific discovery. The authors ask whether a latent world model can be trained from pixels and actions only, with no reward signal, such that its latent dynamics faithfully reproduce the true system over long horizons. Their goal is a model that simultaneously (i) controls well, (ii) predicts faithfully, and (iii) exposes a parameterization whose dynamics can be read out mechanistically.

Key Contributions

1. The world2vec framework and "latent dynamics matching." The authors formalize faithfulness as a condition where integrating a learned latent velocity field over a horizon must equal the encoding of the future observation, and turn this into a multi-step training objective over rollouts of length K (with explicit Euler integration in the main experiments).

2. Objective-agnostic design with two instantiations. The framework accepts any regularizer from the class of objectives known to yield identifiable representations; the authors instantiate it as world2vec-dcl (dynamics contrastive learning) and world2vec-sigreg (LeJEPA-style SIGReg regularizer).

3. A compact, continuous-time architecture with planning-based control. Action-conditioned continuous-time dynamics let the model be integrated at a different step size than the training data, enabling latent-space model-predictive control with the cross-entropy method using a goal latent derived from target frames — no reward model and no state decoder.

4. An interpretable dynamics parameterization. Replacing the MLP velocity field with a SINDy-style structured model (learned coefficient matrix acting on fixed library features, with an L1 penalty), in the spirit of DYSCO, allows the learned dynamics to be read out as equations, which the authors link to accurate system identification across a benchmark of 71 attractors.

Main Findings

• Parameter efficiency: world2vec achieves its results with at least one order of magnitude fewer parameters than DreamerV3, TD-MPC2 and LeWM. The main model uses a latent dimensionality of d=32 and totals 263k parameters (221k encoder, 42k dynamics).

• Control success: world2vec enables successful model-predictive control in 2 of 3 considered environments, none of which were solvable by the previous best JEPA-based approach, LeWM.

• Reward-free, offline training: world2vec and LeWM are trained offline on a random-policy dataset, whereas TD-MPC2 and DreamerV3 are trained online following their own protocols. No reward information is used by world2vec.

• Joint reward/forecast quality: across the three environments, world2vec jointly achieves high reward and high forecasting performance, as summarized by the reward-versus-Equivariance-Accuracy (EQA) comparison in Figure 1B.

• Long-horizon faithfulness metrics: the paper reports consistency (R² after optimal affine alignment between independently trained encoders), Latent Rollout Consistency (LRC, the R² between predicted and encoded futures at horizon j), and EQA (top-1 retrieval of the encoded future among 512 negatives). EQA is used for hyperparameter selection. The specific numerical values appear in Table 1, whose contents are not included in the provided text.

• Training cost: on a single NVIDIA A100 GPU, training takes 2–6 minutes for a one-step model, about 10 minutes for a 20-step model, and about 30 minutes for a 40-step model; closed-loop control takes about 20–50 ms per control step including simulation and rendering.

• System identification: the abstract reports accurate system identification across a benchmark of 71 attractors, presented as evidence that such world models are amenable to scientific discovery.

Methodology in Plain English

The setup assumes an underlying physical state that evolves under actions and is observed through an unknown, non-linear mixing function — so the model only ever sees images and actions, never the state itself.

The encoder takes a short history of H=3 frames (a single frame cannot reveal velocities), and maps it to a latent vector. A dynamics model predicts the rate of change of that latent vector given the current latent and action; a single Euler step then advances the latent by one time interval Δt. During training, the model starts from a sampled point and rolls its own predictions forward K steps, comparing each predicted latent with the encoding of the actual future frame at that offset. Because the target is the change of the encoded state rather than a reconstruction of pixels, the loss directly matches "latent velocities" — the authors call this latent dynamics matching — and small one-step errors become visible as they compound over the rollout. A regularizer (contrastive normalization in DCL, or SIGReg pulling embedding projections toward an isotropic Gaussian) prevents the trivial solution where the encoder outputs a constant.

Architecturally, the visual encoder is a small vision transformer (16×16 patches, 2 layers of width 64, 4 heads, with a keypoint readout) followed by an MLP head (2 hidden layers of 128 units); the dynamics model is an MLP with 3 hidden layers of 128 units. Models are trained for 20k iterations with Adam at learning rate 10⁻³ with cosine decay. Data consists of 200k RGB frames at 64×64 px per environment, collected with random policies — uniform random actions for Reacher, temporally correlated Ornstein–Uhlenbeck actions for Pendulum and Cartpole. A separate non-visual dataset passes positional and angular coordinates (not velocities) through a randomly initialized injective MLP to 50 dimensions for identifiability analysis.

For control, the desired target is supplied as a sequence of frames which is encoded into a goal latent; the cross-entropy method samples candidate action sequences, rolls out the latent dynamics, and ranks plans by weighted distance to the goal, subject to action bounds. A prefix of the best plan is executed and replanning occurs from new observations. For Reacher, the goal is computed per episode via inverse kinematics and rendering, and both elbow configurations are planned with the average reported.

Why This Matters

The paper argues that control performance alone is a weak test of a world model: a model can control well without capturing the environment's dynamics, which undermines its use for interpretability and for science. By tying the training objective to a formal faithfulness condition (Eq. 3), and by showing a compact model that both controls and supports readout of dynamics, the work connects identifiability theory for self-supervised learning to practical continuous control. It also reports performance with models an order of magnitude smaller than prevailing baselines, and training times of minutes on a single GPU.

Real-world applications suggested by the setting:

• Robotics: reward-free model-based control from camera images, where specifying a target observation is easier than designing a reward.
• Scientific discovery: recovering governing equations of dynamical systems from observations, motivated by the 71-attractor system-identification benchmark.
• Closed-loop experimental science: the authors explicitly note that the efficiency of training and 20–50 ms control steps enables closed-loop experimental settings.
• Interpretable control systems: models whose dynamics can be read out (e.g., via SINDy-style coefficients) rather than remaining opaque.

Industry relevance: the combination of offline training, no reward engineering, small parameter counts, and fast training/control loops is directly relevant to practitioners who need deployable model-based controllers and to teams interested in auditing or interpreting learned dynamics.

Future Directions

• Carrying identifiability guarantees into practice: the theory guarantees recovery up to at least affine transformations under idealized conditions; the gap between those conditions and numerical estimation is left as an open challenge, measured here only empirically via consistency.
• Scaling the observation model: all reported control experiments use 64×64 px frames and small transformers; the framework is described as allowing larger inputs, but this is not demonstrated.
• Broader mechanistic evaluation: the SINDy-style parameterization is presented as an interpretable alternative to the MLP velocity field, and the symbolic-regression/attractor benchmark is the natural place to extend equation recovery to harder systems.
• Understanding the multi-step inductive bias: the authors hypothesize, rather than prove, that modeling state change through a learned velocity field is a useful inductive bias for world modeling; isolating why multi-step dynamics matching helps relative to single-step prediction remains open.

Target Audience

Researchers working on world models, model-based reinforcement learning, and JEPA-style self-supervised learning; scientists interested in identifiable representation learning and equation discovery for dynamical systems; and engineers who need compact, reward-free, offline-trained controllers with interpretable dynamics. Readers without background in latent dynamics models, neural ODEs, or identifiability theory will find the paper demanding.

Authors’ 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.