No history yet

Latent Dynamics Architectures

From Pixels to Predictions

World models sidestep the immense computational cost of raw observation simulation by operating within a compact, learned latent space. Instead of predicting the next frame pixel by pixel, they predict how the latent state evolves. This shift from generative modeling of observations to predictive modeling of dynamics is the core architectural principle enabling long-horizon imagination and planning. Two dominant architectures exemplify this approach: the Recurrent State-Space Model (RSSM) and the Joint-Embedding Predictive Architecture (JEPA).

Lesson image

The (RSSM), famously used in the Dreamer series, builds a structured probabilistic model of the environment's dynamics. It factorizes its latent state into two components: a stochastic categorical state (sts_t) and a deterministic recurrent state (hth_t). The deterministic state, managed by an RNN like a GRU, aggregates information over time, providing a summary of history. The stochastic state captures the inherent randomness and ambiguity of the environment, representing a belief over possible world states consistent with the history.

The transition model learns p(st+1st,at,ht+1)p(s_{t+1} | s_t, a_t, h_{t+1}), predicting the next stochastic state from the current one and the action, conditioned on the recurrent state. During imagination, the agent rolls out this transition model without any real observations, generating sequences of latent states (st,ht)(s_t, h_t) and predicted rewards. A key innovation in DreamerV3 was the use of discrete categorical variables for the stochastic state, combined with ideas from VQ-VAE, which prevents the posterior collapse common in VAE-based world models and stabilizes training.

Embedding Space Prediction

In contrast, the (JEPA), pioneered by Meta AI, eliminates the pixel-reconstruction loss entirely. Instead of decoding the latent state back into an image, JEPA focuses solely on prediction in the embedding space. The architecture learns a representation by predicting the embedding of a future, masked patch of the input, given the embedding of a visible context patch.

Joint-Embedding Predictive Architectures (JEPAs) are a family of models that learn by predicting high-level features rather than pixels.

The model consists of a context encoder and a target encoder. The context encoder processes an unmasked region of the input (e.g., a video clip), producing a representation. A predictor module then takes this representation and an action to forecast the representation of a future, masked-out region. The target encoder, which is an exponentially moving average of the context encoder's weights, provides the target representation for this prediction. The loss is computed directly between the predicted embedding and the target embedding. This avoids modeling high-frequency, unpredictable details, focusing the model's capacity on capturing the underlying dynamics and semantics.

Both RSSM and JEPA aim for a sufficient representation, one that captures all the necessary information to predict future states and rewards without needing the full history of observations. By operating in a compressed latent space and focusing on dynamics, these architectures create efficient, scalable world models that form the cognitive backbone for advanced robotic agents.

Quiz Questions 1/5

What is the primary architectural difference between a Recurrent State-Space Model (RSSM) and a Joint-Embedding Predictive Architecture (JEPA)?

Quiz Questions 2/5

Within the RSSM framework, what is the specific role of the deterministic recurrent state (hth_t)?