Skip to content

Concepts

Predict representations, not pixels

A joint-embedding predictive architecture (JEPA) encodes a context and a target and trains a predictor to map the context's embedding to the target's. It never reconstructs pixels, so capacity goes to what is predictable rather than to texture and noise.

The hard part is collapse: a constant embedding predicts itself perfectly. Models in torch_world avoid it in one of two ways:

strategy how models
EMA teacher targets come from a slowly updated copy of the encoder, with stop-gradient I-JEPA, V-JEPA
distribution regularizer SIGReg pushes embeddings toward an isotropic Gaussian LeJEPA, LeWM
frozen encoder a pre-trained encoder is not trained at all DINO-WM, RoboJEPA
contrastive InfoNCE pushes non-matching pairs apart VL-JEPA

From JEPA to world model

When the context is the present and the target is the future, and the predictor also reads the action taken, the JEPA becomes a latent world model: it imagines the consequence of an action in embedding space. Planning then searches for the actions whose imagined outcome lands closest to the embedding of a goal.

A latent world model, planning toward a goal The observation is encoded to a latent z_t. The predictor maps z_t and an action a_t to a predicted next latent, and rolls out over several steps. A cost compares the final predicted latent with the encoded goal image, and the planner adjusts the actions to lower it. imagine goal observation encode zt predict ẑt+1 rollout ẑT cost at plan: lower the cost goal image encode zgoal
Encode the present, imagine the future one action at a time, and plan by searching for the actions whose imagined outcome lands closest to the goal.

Two base classes

Every model subclasses one of two classes in torch_world.models.

base class you write you get
JEPA forward(...) -> {"loss": ...} make_target_encoder, update_target(momentum)
WorldModel encode, predict(z, actions), forward rollout, cost, planning with CEM

Conventions

Loss dicts. forward returns a dict with a scalar "loss" to backpropagate, plus its named components ("pred_loss", "sigreg_loss", ...) for logging.

Time is dim 1. Episodes and latents are (B, T, ...). For WorldModel, predict(z, actions) maps z[:, t] and actions[:, t] to a prediction of z[:, t + 1], for every t at once (teacher forcing).

Actions align with states. In collected episodes, actions[:, t] is applied at states[:, t] and produces states[:, t + 1]. The last action is padding.

History. WorldModel.history is the number of past latents the predictor sees. rollout slides a window of that length; None means the full past.

Masks are index tensors. A mask is a (B, K) tensor of token indices to keep, and a masking strategy returns lists of them, one per block. apply_masks gathers them and stacks multiple masks on the batch dim.

Videos are channels-first. VisionTransformer takes images (B, C, H, W) or videos (B, C, T, H, W); episodes from environments are (B, T, C, H, W), so transpose dims 1 and 2 before a video encoder.