The first neural world model: sequential latent variables

Introduce latents because pixels are too big to plan in.

Photo by Shubham Dhage on Unsplash

In the last blog, we summarized the classical world models and introduced the need for a learnable latent variable in the presence of unknown dynamics.

Till now, our model was:

\(s_{t+1} = As_t+Ba_t\)

This is rigid.

Here, we accept that we don’t know the state equations. We must learn a latent dynamics model. The high-dimensional observations (pixels) force two changes:

• the observation model p(o_t|x_t) must be learned (decoder).

• inference p(x_t|o_<t,a_<t) must be amortized (encoder).

The core question becomes:

How do we learn a latent dynamical system whose latent state is (1) predictice, (2) trainable from pixels, and (3) usable for control?

We have sequences from an environment or simulator:

\(D={o_{1:T},a_{1:T-1},r_{1:T}}\)

Think of a replay buffer: chunks of length T, e.g., T = 50. Our job is to learn a model that can answer counterfactual questions: If I’m in “the situation implied by history”, and I take actions a_t, what distribution over next observations/rewards should I expect?

\(p(o_{t+1},r_{t+1}|h_t,a_t)\)

But we don’t want to store the whole history h_t, we want a compact “state”.

Why not just predict pixels directly?

A naive world model would do:

\(\hat{o}_{t+1} = F_{\theta}(o_{1:t},a_{1:t})\)

If we do this with a transformer, it can work for short horizons, but it has three deep problems.

Problem A: multi-modality becomes blurred

If from the same current frame, two futures are possible (e.g., an occluded car might appear left or right), the MSE-optional prediction is the average, which is not a valid future. So “pixel regression” learns means of the futures.

Problem B: partial observability requires “belief”, not a point

In partially observable Markov decision processes, the correct internal state is a distribution over latent states, not a deterministic vector. Direct pixel predictors usually hide ambiguity poorly.

Problem C: planning needs a state that evolves cleanly under actions

Planning requires rolling forward under imagined actions for many steps. Pixels’ space is huge - compounding errors explode. So we want a latent state that is small, predictive, action-conditioned, and (ideally) carries undertainty. This is exactly what the classical state-space model did, except now o_t is pixels and the observation model is unknown.

The minimal fix: introduce a latent state

We posit a latent variable z_t at every time step. You can interpret z_t as the compact internal state/belief that explains the observation and predicts the future under actions — just defined a generative model (the “neural world model”).

Building the generative story…

First of all, sample the initial latent:

\(z_1 \sim p_{\theta}(z_1)\)

For each time step t,

\(\text{transition (dynamics): }z_{t+1} \sim p_{\theta}(z_{t+1}|z_t,a_t)\)

\(\text{emission (rendering observation): }o_t \sim p_{\theta}(o_t|z_t)\)

This is just a latent state-space model. We still need to figure out: how do we learn it from data?

Before we do that, let’s have some concepts in order.

The obstacle: we never observe the latent state (z_t)

We only see o_t. The training process aims to maximize:

\(\log p_{\theta}(o_{1:T}|a_{1:T-1}) = \log \int p_{\theta}(z_{1:T},o_{1:T}|a)dz_{1:T}\)

This integral is the “intractable marginalization over latent trajectories”. So we do the standard pick: approximate the posterior with an inference model.

The single most important concept: two distributions

At each step, we will maintain two distributions over the latent state:

Prior (prediction-only): What the dynamics model predicts before seeing the next observation:

\(p_{\theta}(z_t|z_{t-1},a_{t-1})\)

This is what our model “imagines”.

Posterior (inference, corrected by observation): What we believe after incorporating the actual observation:

\(q_{\phi}(z_t|z_{t-1},a_{t-1},o_t)\)

This is our learned “filter update”.

World models are trained by forcing the prior (imagination) to match the posterior (reality-corrected belief), while also reconstructing observations.

That’s it, everything else is details!

Key equations: ELBO for sequential latents (deriving the training objective)

Please read the equations carefully, they are simple and understandable.

How do we train this? We cannot maximize p(x) over all time steps directly because it requires integrating over all possible latent trajectories (an intractable problem!). Instead, we maximize the Evidence Lower Bound (ELBO) OR Variational Lower Bound (ELBO): a lower bound on the log-likelihood of some observed data.

Why are we using ELBO? Read the first two paragraphs

For a single time step t, the loss function L has three parts:

\(L = \underbrace{E_q[\log p(x_t|h_t,z_t)]}_{\text{Reconstruction Loss}} + \underbrace{E_q[\log p(r_t|h_t,z_t)]}_{\text{Reward Prediction}} - \underbrace{\beta KL[q(z_t|h_t,x_t)||p(z_t,h_t)]}_{\text{Dynamics Consistency}}\)

Let’s decode the KL term. This is the most critical equation in World Models.

• q(.) (Posterior): the distribution given the current observation

• p(.) (Prior): the distribution given only the past

Essentially, we want the “prior” to be close to the “posterior”. Why? Because at runtime (deployment), we don’t have the observation (image) x_t yet. We only have the Prior. If the Prior accurately matches the Posterior, our “dreams” (p) match “reality” (q).

A simple derivation

We will start with what we want to maximize (check above):

\(\log p_{\theta}(o_{1:T}|a) = \log \int p_{\theta}(z_{1:T},o_{1:T}|a)dz\)

Multiply and divide by:

\(q_{\phi}(z_{1:T}|o,a) \implies = \log \int q_{\phi}(z|o,a)\frac{p_{\theta}(z,o|a)}{q_{\phi}(z|o,a)}dz\)

Apply Jensen’s inequality:

\(\geq E_{q_{\phi}}[\log p_{\theta}(z,o|a) - \log q_{\phi}(z|o,a)]\)

Now, let’s expand the first term using the model factorization:

\(\log p_{\theta}(z,o|a) = \sum_{t=1}^T \log p_{\theta}(o_t|z_t) + \log p_{\theta}(z_1) + \sum_{t=1}^{T-1} \log p_{\theta}(z_{t+1}|z_t, o_t)\)

So, the ELBO becomes:

\(L = \underbrace{\sum_{t=1}^T E_q[\log p_{\theta}(o_t|z_t)]}_{\text{reconstruction/likelihood}} - \underbrace{ KL[q_{\phi}(z_{1:T}|o,a)||p_{\theta}(z_{1:T}|a)]}_{\text{make posterior consistent with dynamics prior}}\)

Let’s understand these two terms intuitively

1. Reconstruction Loss: Forces the latent state (z_t) to carry enough information to render o_t.

2. KL to the prior: Forces the inferred latent trajectory to be explainable by the dynamics model. Essentially, this is the “physics constraint” in the latent space.

To deeply understand the KL term, let’s decompose it per time step.

Pick a filtering factorization:

\(q_{\phi}(z_{1:T}|o,a) = q(z_1|o_1)\prod_{t=2}^T q(z_t|z_{t-1},a_{t-1},o_t)\)

and the Markov prior:

\(p_{\theta}(z_{1:T}|a) = p(z_1)\prod_{t=2}^T p(z_t|z_{t-1},a_{t-1})\)

The, the trajectory KL decomposes as:

\(KL(q(z_{1:T})||p(z_{1:T})) = KL(q(z_1|o_1)||p(z_1)) + \sum_{t=2}^T E_q[KL(q(z_t|z_{t-1},a_{t-1},o_t)||p(z_t|z_{t-1},a_{t-1})]\)

This is exactly why our code can do per-step KL. The training objective in operational form:

\(L = \sum_t E_q[\log p(o_t|z_t)] - \sum KL(q(z_t|\cdot)||p(z_t|\cdot))\)

This looks abstract until we realize:

• q(.) is the encoder distribution (posterior)

• p(.) is the dynamics distribution (prior)

So we are literally penalizing the mismatch between what our model predicts and what reality forces us to infer.

To summarize, we discussed:

prior rollout (imagination)  ----->  p(z_{t+1}|z_t,a_t)
          |                                  |
          | (observe o_{t+1})                | (no observation)
          v                                  v
posterior correction  ----->  q(z_{t+1}|...,o_{t+1})

Following this, a paper came out on this: Variational Recurrent Neural Networks (VRNN).

Briefly discussing this, VRNN uses an RNN hidden state h_t as memory of the past, and samples z_t each step.

\(\text{Update a deterministic memory }h_{t-1}\)

\(\text{From } h_{t-1} \text{ and } o_t, \text{ infer } q(z_t)\)

\(\text{From } h_{t-1}, \text{ predict prior } p(z_t)\)

\(\text{Decode } o_t \text{ from } h_{t-1},z_t\)

\(\text{Update } h_t \text{ using } o_t,z_t\)

The RNN gives a flexible memory backbone; pure Markov can be too weak. The issue is that the RNN can learn to reconstruct o_t without needing z_t, causing posterior collapse (when the model can minimize loss while ignoring the latent state).

The field has converged on a specific architecture known as the Recurrent State Space Model (RSSM), which we are going to cover in the next blog!

~Ashutosh