Ben Bolte

In this post, I'll provide a framing for various different formalizations of world models (the cake génoise of machine learning).

What is a world model?

A world model is a model of the world.

This definition is so nebulous that the term "world model" is often co-opted to mean whatever a particular author wants it to mean. I personally like the "energy minimization" framing:

w⋆∈arg min⁡w∈WEθ(w∣c)w^\star \in \argmin_{w \in \mathcal{W}} E_\theta(w \mid c)

where:

  • EθE_\theta is the learned energy function (our "world model")
  • cc is the conditioning context: observations, actions, text, goals, or history
  • ww is a candidate world
  • W\mathcal{W} is the space of possible candidate worlds
  • w⋆w^\star is the world judged most compatible with the context

This framing is general enough that it can be extended to represent the various different types of world models, as I'll do below.

More reading:

Observation-predictive models

We can reframe the energy-minimization framing above as a simple "predict the future" objective:

w⋆∈arg min⁡xt+1:t+h∈XhEθ(xt+1:t+h∣x≤t,at:t+h−1)w^\star \in \argmin_{x_{t+1:t+h} \in \mathcal{X}^h} E_\theta(x_{t+1:t+h} \mid x_{\leq t}, a_{t:t+h-1})

where:

  • xt+1:t+hx_{t+1:t+h} is a candidate trajectory (video frames, actions, propioceptive tokens, or some other readings from the environment)
  • Xh\mathcal{X}^h is the space of possible futures (of length hh, meaning that we assume some discretized time), with X\mathcal{X} being the observation space
  • (x≤t,at:t+h−1)(x_{\leq t}, a_{t:t+h-1}) is the combined.observational history, with a proposed sequence of future actions that we will take
  • tt is the current timestep
  • w⋆w^\star is the lowest-energy future for the past and intervention

Thus, we can learn an observation-predictive model by minimizing some learned energy function over the future, parametrized by a set of actions we intend to take. Note that this energy function is over concrete future states rather than abstract beliefs about the future. We can then convert this to an action-taking model via model-predictive control.

More reading:

Latent dynamics models

One of the main issues with the above approach is that the observation space x1:hx_{1:h} can be extremely large. For instance, 30 seconds of 1080p video at 30 Hz can take up 5.6 GB when unpacked to RGB pixel space. The straightforward fix for this is to instead model some lower-dimensional manifold.

zt=fϕ(x≤t),w⋆∈arg min⁡zt+1:t+h∈ZhEθ(w∣zt,at:t+h−1)z_t = f_\phi(x_{\leq t}), \qquad w^\star \in \argmin_{z_{t+1:t+h} \in \mathcal{Z}^h} E_\theta(w \mid z_t, a_{t:t+h-1})

where:

  • zt+1:t+hz_{t+1:t+h} is a candidate latent rollout
    • Keen observers will note that this looks suspiciously similar to xt+1:t+hx_{t+1:t+h}
    • Remember, zz comes after xx, as in, "iz not ze xact representation, iz ze latent variable"1
  • Zh\mathcal{Z}^h is the space of latent rollouts, analogous to Xh\mathcal{X}^h
  • (zt,at:t+h−1)(z_t, a_{t:t+h-1}) is- okay, at this point you should basically get the gist
  • fϕf_\phi is some encoder which maps a sequence of observation histories into a latent state
  • w⋆w^\star is the lowest-energy rollout from our energy function
    • There often exists some quasi-inverse of the encoder which provides x~=gϕ(w⋆)\tilde{x} = g_\phi(w^\star), meaning the approximate reconstruction of the original world state from our predicted latents
    • An autoencoder can provide a compact, regularized latent space in which a diffusion model learns to generate new representations

There are basically two reasons for taking this latent dynamics approach:

  1. It makes the observation space smaller
  2. It simplifies the manifold, making learning easier

The downside is that you need to learn an encoder fϕf_\phi which can correctly model the world state.

  • This introduces a Pareto curve for possible encoders, trading off representation capacity for compression
  • Often, the stuff you care about modeling well is in the tail of the data distribution, meaning that highly-compressed points on this Pareto curve effectively wash out important world details

We can sometimes jointly learn the encoder as part of the modeling problem, for example:

  • Self-distillation aligns representations across different views
  • Predictive learning, including temporal difference learning, aligns representation across different timesteps

Both approaches need to be trained carefully to avoid mode collapse. There is a significant amount of additional literature on representation learning which is worth exploring, and in my own view this is obviously the most interesting direction for scaling world models because it leverages a small number of robust priors.

More reading:

Object-centric and relational models

We can similarly represent ww as st+1:t+hs_{t+1:t+h}, where sts_t is some set of entities and relations:

st=(ot1,…,otn,rt)s_t = (o_t^1, \ldots, o_t^n, r_t)

In this framing, the model is selecting among compositional hypotheses whose energy decomposes over entities and interactions. In other words, the energy function is representing the energy of various possible futures as a function of different objects and their relations.

World models of this flavor are typically highly structured, and have a symbolic AI flavor to them. This is an interesting framing for large world models today insofar as the latent dynamics can be shown to be implicitly learning these types of objects and their relations; as far as I am aware this type of analysis has yet to be done, since in the LLM regime it's typically easier to instead just ask the model relational questions and see if it returns the right answers.

More reading:

Successor and value-equivalent models

For some variations of world models, particularly for learned control, the relevant world is the future occupancy measure induced by a policy rather than the full space of potential future observations. In other words, if the energy function in our original equation is over the futures that some agent is likely to see rather than over the set of all possible world states, then updates to the agent's policy will shift the energy function meaningfully. This bridges the gap nicely with the Bellman equation and provides a framing by which it is possible to adapt RL algorithms:

ψπ(s,a)=Eπ[∑k=0∞γkϕ(st+k,at+k)∣st=s,at=a]\psi^\pi(s, a) = \mathbb{E}_\pi\left[\sum_{k=0}^{\infty} \gamma^k \phi(s_{t+k}, a_{t+k}) \mid s_t=s, a_t=a\right]

where:

  • ϕ(s,a)\phi(s, a) is a feature vector
  • γ\gamma is the discount factor
  • ψπ\psi^\pi is the expected discounted accumulation of future features

for some policy π\pi.

We can then adapt this into our original energy minimization framing via:

w⋆=ψ^θπ(s,a)∈arg min⁡w∈WEθ(w∣s,a,π)w^\star = \widehat{\psi}_\theta^\pi(s, a) \in \argmin_{w \in \mathcal{W}} E_\theta(w \mid s, a, \pi)

where:

  • ww is a candidate successor-feature vector that estimates ψπ(s,a)\psi^\pi(s, a).
  • W\mathcal{W} is the space of admissible successor-feature vectors.
  • π\pi is the policy used to continue the trajectory
  • ψ^θπ(s,a)\widehat{\psi}_\theta^\pi(s, a) estimates ψπ(s,a)\psi^\pi(s, a) by minimizing the learned energy

The basic idea here is that our various perceptual futures may differ substantially in the observation space, but they can be world-model equivalent if they imply the same downstream successor features and therefore the same value under many reward functions.

More reading:

Generative imagination models

A distinct variant of the observation-space and latent-space world models is to introduce latent stochasticity to the energy function to model future uncertainty:

ϵ∼p(ϵ),w⋆(ϵ)∈arg min⁡xt+1:t+h∈XhEθ(xt+1:t+h∣x≤t,at:t+h−1,ϵ)\epsilon \sim p(\epsilon), \qquad w^\star(\epsilon) \in \argmin_{x_{t+1:t+h} \in \mathcal{X}^h} E_\theta(x_{t+1:t+h} \mid x_{\leq t}, a_{t:t+h-1}, \epsilon)

where:

  • xt+1:t+hx_{t+1:t+h} is a candidate future observation trajectory.
  • Xh\mathcal{X}^h is the space of length-hh observation trajectories.
  • x≤tx_{\leq t} is the observation history.
  • at:t+h−1a_{t:t+h-1} is the proposed sequence of future actions.
  • ϵ\epsilon is noise sampled once per rollout and held fixed during the minimization.
  • p(ϵ)p(\epsilon) is a simple prior over noise, such as a standard Gaussian.
  • w⋆(ϵ)w^\star(\epsilon) is one plausible future selected for the sampled noise, observation history, and proposed actions.

We can then resample ϵ\epsilon, inducing a distribution over the future. Minimizing the energy function then becomes something like a diffusion process.

More reading:

World-action models

If we have a lot of paired world-action data, rather than building a policy (as in the Successor Feature approach), you can simply consider actions to be part of the world:

w⋆∈arg min⁡w∈WEθ(w∣x≤t,a<t,g)w^\star \in \argmin_{w \in \mathcal{W}} E_\theta(w \mid x_{\leq t}, a_{<t}, g)

where:

  • w=(st+1:t+h,at:t+h−1)w = (s_{t+1:t+h}, a_{t:t+h-1}) pairs a candidate future state trajectory with a candidate action sequence.
  • W=Sh×Ah\mathcal{W} = \mathcal{S}^h \times \mathcal{A}^h is the space of paired length-hh state and action trajectories.
  • x≤tx_{\leq t} is the observation history, including the current observation.
  • a<ta_{<t} is the history of actions already taken.
  • S\mathcal{S} is the state space.
  • A\mathcal{A} is the action space.
  • gg is a goal, instruction, reward description, or task embedding.
  • w⋆w^\star is the paired future and action sequence judged most compatible with the observation and action histories and the goal.

This has the distinctive advantage that the model is inherently both counterfactual and teleological. You can sample actions directly from the world model to pass them to your agent, which fits very nicely into the LLM pre-training / post-training paradigm.

More reading:

Footnotes

  1. This is best read in Yann Lecun's voice. ↩