torch_world.models¶

JEPA ¶

Bases: Module

Joint-embedding predictive architecture: predict the embedding of a target from the embedding of a context.

Subclasses implement forward returning a dict with a scalar "loss" and its components.

make_target_encoder ¶

make_target_encoder(encoder: Module) -> Module

Frozen copy of encoder updated by exponential moving average.

WorldModel ¶

Bases: Module

Action-conditioned latent world model.

Latents are (B, T, ...). predict(z, actions) maps z[:, t] and actions[:, t] to a prediction of z[:, t + 1] for every t (teacher forcing). rollout and cost build on it for planning.

rollout ¶

rollout(z: Tensor, actions: Tensor, **cond) -> Tensor

Unroll from T0 context latents with T >= T0 actions; returns T + 1 latents.

cost ¶

cost(z: Tensor, goal: Tensor) -> Tensor

Squared distance between the final latent of each rollout (B, T, ...) and the goal latent (B, ...).

IJEPA ¶

IJEPA(
    encoder: VisionTransformer, predictor: MaskedPredictor
)

Bases: JEPA

I-JEPA: predict the target encoder's features of masked image blocks from a visible context block.

Targets are layer-normalized features of an EMA target encoder, and the loss is smooth L1. Call update_target(momentum) after each optimizer step (reference schedule: 0.996 to 1.0).

Assran et al., "Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture", CVPR 2023. Paper: https://arxiv.org/abs/2301.08243 Code: https://github.com/facebookresearch/ijepa

Example::

encoder = VisionTransformer.from_size("huge", img_size=224, patch_size=14)
model = IJEPA(encoder, MaskedPredictor(encoder.grid, encoder.embed_dim, 384, depth=12, num_heads=12))
masks_enc, masks_pred = MultiBlockMask(encoder.grid)(batch_size)
loss = model(images, masks_enc, masks_pred)["loss"]

VJEPA ¶

VJEPA(
    encoder: VisionTransformer,
    predictor: MaskedPredictor,
    loss_exp: float = 1.0,
    reg_coeff: float = 0.0,
)

Bases: JEPA

V-JEPA: predict EMA target features of masked spatio-temporal tubes of a video.

masks_enc and masks_pred hold one (B, K) context/target pair per masking strategy (short-range and long-range in the paper), and strategy i uses the predictor's mask token i. The loss is mean(|z - h| ** loss_exp) / loss_exp, averaged over strategies, plus reg_coeff times a variance hinge on the predictions. V-JEPA 2 keeps this objective with 3D RoPE: build the encoder and predictor with use_rope=True.

Bardes et al., "Revisiting Feature Prediction for Learning Visual Representations from Video", 2024. Paper: https://arxiv.org/abs/2404.08471 Code: https://github.com/facebookresearch/jepa

Assran et al., "V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning", 2025. Paper: https://arxiv.org/abs/2506.09985 Code: https://github.com/facebookresearch/vjepa2

LeJEPA ¶

LeJEPA(
    encoder: Module,
    projector: Module,
    lamb: float = 0.05,
    sigreg: SIGReg | None = None,
)

Bases: JEPA

LeJEPA: multi-view prediction with SIGReg, with no teacher, stop-gradient or schedules.

Each view's projection predicts the mean projection of the num_global global views, and SIGReg pushes the projections of every view toward an isotropic Gaussian. lamb trades the two terms: (1 - lamb) * prediction + lamb * sigreg.

Balestriero & LeCun, "LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics", 2025. Paper: https://arxiv.org/abs/2511.08544 Code: https://github.com/galilai-group/lejepa

forward ¶

forward(
    views: Tensor | Sequence[Tensor],
    num_global: int | None = None,
) -> dict[str, Tensor]

views is (B, V, ...) or a list of (B, ...) tensors whose first num_global are global views.

VLJEPA ¶

VLJEPA(
    x_encoder: Module,
    predictor: Module,
    query_embed: Embedding,
    y_encoder: Module,
    vision_dim: int,
    predictor_dim: int,
    text_dim: int,
    embed_dim: int = 1536,
    temperature: float = 0.07,
    y_lr_mult: float = 0.05,
)

Bases: JEPA

VL-JEPA: predict the embedding of a target text from visual tokens and an optional text query.

A frozen X-encoder turns images or videos into visual tokens. The predictor attends bidirectionally over the projected visual tokens followed by the query tokens, and the non-padding outputs are average-pooled into a predicted embedding. The Y-encoder embeds the target text. Linear heads map both into a shared space, and the predictor and Y-encoder are trained jointly with bidirectional InfoNCE. The Y-encoder trains at y_lr_mult times the base learning rate (see param_groups). Text is produced only on demand by a separate decoder; segment picks the decoding points of an embedding stream (selective decoding).

The paper does not report the InfoNCE temperature, so it is learnable here, initialized at 0.07 as in CLIP.

Chen et al., "VL-JEPA: Joint Embedding Predictive Architecture for Vision-language", ICLR 2026. Paper: https://arxiv.org/abs/2512.10942 Code: not released; this follows the paper.

Parameters:

Name	Type	Description	Default
`x_encoder`	`Module`	frozen module mapping visual input to tokens `(B, N, vision_dim)` (V-JEPA 2 ViT-L in the paper).	required
`predictor`	`Module`	module mapping `(tokens, attn_mask=...)` to tokens of the same shape (last 8 Llama-3.2-1B layers).	required
`query_embed`	`Embedding`	token embedding of the predictor's language model.	required
`y_encoder`	`Module`	module mapping `(ids, mask)` to pooled text embeddings `(B, text_dim)` (EmbeddingGemma-300M).	required

predict ¶

predict(
    x: Tensor,
    query_ids: Tensor | None = None,
    query_mask: Tensor | None = None,
) -> Tensor

Predicted target embedding (B, embed_dim); query_mask is True on non-padding query tokens.

retrieve ¶

retrieve(
    x: Tensor,
    bank: Tensor,
    query_ids: Tensor | None = None,
    query_mask: Tensor | None = None,
) -> Tensor

Index (B,) of the closest row of bank (candidate text embeddings from embed_text) by cosine.

param_groups ¶

param_groups(lr: float, **kwargs) -> list[dict]

Optimizer groups with the Y-encoder at y_lr_mult * lr; the frozen X-encoder is excluded.

segment `staticmethod` ¶

segment(stream: Tensor, num_segments: int) -> Tensor

Selective decoding points: midpoints of num_segments contiguous segments of a (T, D) stream.

Segments come from agglomerative clustering with temporal connectivity and Ward linkage.

LeWM ¶

LeWM(
    encoder: Module,
    action_dim: int,
    embed_dim: int = 192,
    history: int = 3,
    depth: int = 6,
    num_heads: int = 16,
    head_dim: int = 64,
    mlp_dim: int = 2048,
    proj_dim: int = 2048,
    dropout: float = 0.1,
    lamb: float = 0.09,
    sigreg: SIGReg | None = None,
)

Bases: WorldModel

LeWorldModel: an end-to-end JEPA world model trained from pixels with two terms, next-embedding MSE and SIGReg.

Each frame is encoded to one embedding (the encoder's CLS token) and projected. A causal transformer, conditioned on action embeddings through adaLN-zero, predicts the next embedding from the last history frames. There is no EMA target and no stop-gradient: loss = mse + lamb * sigreg.

Maes et al., "LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels", 2026. Paper: https://arxiv.org/abs/2603.19312 Code: https://github.com/lucas-maes/le-wm

Parameters:

Name	Type	Description	Default
`encoder`	`Module`	module mapping images `(N, C, H, W)` to embeddings `(N, embed_dim)` (ViT-tiny CLS in the paper).	required
`action_dim`	`int`	dimension of one action, times the frameskip when actions are stacked.	required

encode ¶

encode(pixels: Tensor) -> Tensor

(B, T, C, H, W) to (B, T, embed_dim).

forward ¶

forward(
    pixels: Tensor, actions: Tensor
) -> dict[str, Tensor]

pixels (B, history + 1, C, H, W) and actions (B, history + 1, action_dim).

DINOWM ¶

DINOWM(
    encoder: Module,
    num_patches: int,
    embed_dim: int,
    action_dim: int,
    proprio_dim: int = 0,
    action_emb_dim: int = 10,
    proprio_emb_dim: int = 10,
    history: int = 3,
    depth: int = 6,
    num_heads: int = 16,
    head_dim: int = 64,
    mlp_dim: int = 2048,
    dropout: float = 0.1,
)

Bases: WorldModel

DINO-WM: a world model over the patch features of a frozen pre-trained encoder (DINOv2).

Proprioception and action embeddings are tiled over every patch and concatenated on the feature dim. A ViT predictor with frame-level causal attention predicts the next frame's visual and proprio features, trained with MSE against detached targets. Latents from encode are (B, T, P, D + proprio_emb_dim).

Zhou et al., "DINO-WM: World Models on Pre-trained Visual Features enable Zero-shot Planning", ICML 2025. Paper: https://arxiv.org/abs/2411.04983 Code: https://github.com/gaoyuezhou/dino_wm

Parameters:

Name	Type	Description	Default
`encoder`	`Module`	frozen module mapping images `(N, C, H, W)` to patch tokens `(N, num_patches, embed_dim)`.	required

encode ¶

encode(
    pixels: Tensor, proprio: Tensor | None = None
) -> Tensor

pixels (B, T, C, H, W) and proprio (B, T, proprio_dim) to (B, T, P, D').

forward ¶

forward(
    pixels: Tensor,
    actions: Tensor,
    proprio: Tensor | None = None,
) -> dict[str, Tensor]

pixels (B, history + 1, C, H, W), actions (B, history + 1, action_dim).

RoboJEPA ¶

RoboJEPA(
    embed_dim: int,
    action_dim: int,
    state_dim: int = 0,
    num_views: int = 1,
    num_embodiments: int = 1,
    predictor_dim: int = 512,
    depth: int = 7,
    num_heads: int = 4,
    mlp_ratio: float = 4.0,
    window: int | None = 8,
    rollout_steps: int = 2,
    encoder: Module | None = None,
)

Bases: WorldModel

RoboJEPA: an action-conditioned predictor in the latent space of a frozen V-JEPA 2.1 encoder, for many robots.

Every timestep contributes the patch tokens of all camera views, one action token and one proprioceptive state token. Actions and states are projected per embodiment. Attention is causal over frames within a window of past frames, and positions use 3D RoPE (t, h, w). The loss is L1 on layer-normalized features, under teacher forcing and over rollout_steps autoregressive steps, as in V-JEPA 2-AC.

Zholus et al., "RoboJEPA: Scaling Robotic Latent World Models", 2026. Paper: https://arxiv.org/abs/2610.10515 Code: https://github.com/facebookresearch/robo_jepa

Assran et al., "V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning", 2025. Paper: https://arxiv.org/abs/2506.09985 Code: https://github.com/facebookresearch/vjepa2

Latents are (B, T, V, H, W, embed_dim): V views on an H x W patch grid.

encode ¶

encode(frames: Tensor, grid: tuple[int, int]) -> Tensor

frames (B, T, V, C, H, W) to layer-normalized features (B, T, V, h, w, D) with the frozen encoder.

rollout ¶

rollout(
    z: Tensor,
    actions: Tensor,
    states: Tensor | None = None,
    embodiment: Tensor | None = None,
) -> Tensor

States past the given ones repeat the last known state.

forward ¶

forward(
    z: Tensor,
    actions: Tensor,
    states: Tensor | None = None,
    embodiment: Tensor | None = None,
) -> dict[str, Tensor]

z (B, T, V, H, W, D) encoder features, actions (B, T, A) and states (B, T, S).

EmbodimentLinear ¶

EmbodimentLinear(
    in_dim: int, out_dim: int, num_embodiments: int = 1
)

Bases: Module

Linear map with separate weights per embodiment, selected by ids (B,).

torch_world.models¶

JEPA ¶

make_target_encoder ¶

WorldModel ¶

rollout ¶

cost ¶

IJEPA ¶

VJEPA ¶

LeJEPA ¶

forward ¶

VLJEPA ¶

predict ¶

retrieve ¶

param_groups ¶

segment staticmethod ¶

LeWM ¶

encode ¶

forward ¶

DINOWM ¶

encode ¶

forward ¶

RoboJEPA ¶

encode ¶

rollout ¶

forward ¶

EmbodimentLinear ¶

segment `staticmethod` ¶