World2vec

A physical state $\boldsymbol s$ evolves under actions $\boldsymbol a$ as $\dot{\boldsymbol s}=\mathbf f(\boldsymbol s,\boldsymbol a)$ and is observed through a nonlinear mixing $\boldsymbol x=\mathbf g(\boldsymbol s)$. World2vec learns an encoder $\mathbf h$ that maps observations to latent states $\boldsymbol z=\mathbf h(\boldsymbol x)$, and a latent vector field $\hat{\mathbf f}$ with $\dot{\boldsymbol z}=\hat{\mathbf f}(\boldsymbol z,\boldsymbol a)$. The model is faithful if evolving the encoded state in latent space gives the same trajectory as encoding the evolved physical state:

The condition must hold for all times $t$ and horizons $T\geq0$. Both sides can be computed from observations and actions alone, without the physical state or the true dynamics. In training, the encoder maps a short history of frames $\boldsymbol X_n$ to the initial latent (a single frame does not show velocities), Euler steps advance it under the recorded actions, and the rollout is matched to the encodings of the observed future:

For a single step, this objective matches latent velocities up to a finite difference, as in neural ODEs. Over $K$ steps, it matches the encoded displacement with the velocities accumulated along the model's own predicted trajectory, so training exposes how small one-step errors compound. We call this principle latent dynamics matching.

A regularizer $\mathcal R$ prevents collapse, since a constant encoder satisfies the condition trivially. world2vec-dcl normalizes prediction scores against sampled negative embeddings (dynamics contrastive learning); world2vec-sigreg pulls embedding projections towards an isotropic Gaussian (SIGReg from LeJEPA).

We measure the quality of the representation with consistency and the quality of the rollouts with LRC and EQA. Consistency is the R² between two independently trained encoders after an optimal affine alignment. LRC (latent rollout consistency) is the R² between latent rollouts and the encodings of the observed future. EQA (equivariance accuracy) is the top-1 retrieval accuracy of the encoded future among sampled negatives; unlike LRC, a collapsed embedding cannot score high on it. Hyperparameters are selected on validation EQA.

Setup: RGB frames from random policies on each environment. A small vision transformer encodes a short history of frames into a low-dimensional latent, and an MLP parameterizes the vector field; the whole model is at least an order of magnitude smaller than the baselines and trains in minutes on a single GPU. Planning uses the cross-entropy method towards an encoded goal, without a reward model or state decoder.

Increasing the training horizon drastically improves both world2vec-dcl and world2vec-sigreg, most of all on CartPole, where one-step training fails entirely. The exact horizon matters little, as long as training goes well beyond single-step prediction.

We test the second component by replacing the Euler step with discrete transitions, $\hat{\boldsymbol z}_{n,j+1}=\mathrm{MLP}(\hat{\boldsymbol z}_{n,j},\boldsymbol a_{n+j})$, in which the MLP predicts the next latent state directly. The Euler models are better on consistency, EQA and state decoding for both variants. Discrete dynamics still reach a high LRC when the latent space is as small as the physical state.

The analysis of latent world models found that standard predictors do not capture velocity information. Modeling the latent change rather than the next latent state counteracts this failure.

For these objectives, theory predicts representations that are consistent up to an affine transformation and dynamics that are modeled faithfully over time. The state itself is not expected to be recovered linearly, since the noiseless systems lack variability along some dimensions.

As the mixing deepens, LRC and EQA decrease gradually, with world2vec-dcl degrading least, while the one-step baselines collapse as soon as the mixing has more than one layer. Reward declines only slightly, so both world2vec variants still solve CartPole from non-linearly mixed observations, whereas the baselines fail. The gap between EQA here and under pixel observations is modest, which suggests that the result can transfer to pixels with better estimation of the dynamics model.

Setup: the positions and angles of the state, but not its velocities, pass through a randomly initialized, injective MLP into a higher-dimensional observation. An MLP encodes the observation history in place of the vision transformer.

World2vec-SINDy improves substantially over an identity predictor and performs comparably to an MLP vector field, and all three give highly consistent ensembles. Process noise adds variation along every state dimension and enables linear state recovery: with more noise, affine recovery and consistency increase for all three models. Noise also caps LRC, so at high noise we judge the dynamics against the true rollouts (iDynR²), where MLP and SINDy stay far above the identity baseline.

Unlike the MLP, SINDy yields equations. For SprottA, a system whose coefficients agree well across seeds, the recovered equations approximate the true system after a least-squares alignment to the true states and removal of small coefficients. Expressing the equations in ground-truth coordinates uses the true states for this alignment; recovering them without the true states, for example with unsupervised symbolic regression, is left for future work.