LeWorldModel

Stable End-to-End Joint-Embedding Predictive Architecture from Pixels

Lucas Maes*, Quentin Le Lidec*, Damien Scieur, Yann LeCun and Randall Balestriero

Abstract: Joint Embedding Predictive Architectures (JEPAs) offer a compelling framework for learning world models in compact latent spaces, yet existing methods remain fragile, relying on complex multi-term losses, exponential moving averages, pretrained encoders, or auxiliary supervision to avoid representation collapse. In this work, we introduce LeWorldModel (LeWM), the first JEPA that trains stably end-to-end from raw pixels using only two loss terms: a next-embedding prediction loss and a regularizer enforcing Gaussian-distributed latent embeddings. This reduces tunable loss hyperparameters from six to one compared to the only existing end-to-end alternative. With ~15M parameters trainable on a single GPU in a few hours, LeWM plans up to 48× faster than foundation-model-based world models while remaining competitive across diverse 2D and 3D control tasks. Beyond control, we show that LeWM's latent space encodes meaningful physical structure through probing of physical quantities. Surprise evaluation confirms that the model reliably detects physically implausible events.

[ Paper | Checkpoints & Data | Website ]

If you find this code useful, please reference it in your paper:

@article{maes_lelidec2026lewm,
  title={LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels},
  author={Maes, Lucas and Le Lidec, Quentin and Scieur, Damien and LeCun, Yann and Balestriero, Randall},
  journal={arXiv preprint},
  year={2026}
}

Using the code

This codebase builds on stable-worldmodel for environment management, planning, and evaluation, and stable-pretraining for training. Together they reduce this repository to its core contribution: the model architecture and training objective.

Installation:

sudo apt install -y swig python3-dev zstd

cd lewm
uv sync

Data

Datasets use the HDF5 format for fast loading. Download the data from HuggingFace and decompress with:

zstd -d archive.tar.zst

Place the extracted .h5 files under $STABLEWM_HOME/datasets/ (defaults to ~/.stable_worldmodel/datasets/). You can override the root path:

export STABLEWM_HOME=/path/to/your/storage

Dataset names are specified without the extension. For example, config/train/data/pusht.yaml references pusht_expert_train, which resolves to $STABLEWM_HOME/datasets/pusht_expert_train.h5.

The training configs default to .h5. If you want to use the faster LanceDB format instead, convert with:

from stable_worldmodel.data import convert
convert('pusht_expert_train.h5', 'pusht_expert_train.lance', dest_format='lance')

Then update the name field in the relevant data config (e.g. config/train/data/pusht.yaml) to use the .lance suffix.

Training

jepa.py contains the PyTorch implementation of LeWM. Training is configured via Hydra config files under config/train/.

To launch training:

python train.py data=pusht

Additional Config

• Enable wandb tracking: wandb.enabled=true (default=false)
• Specify wandb entity: wandb.config.entity=myentity (default=lewm)
• Specify wandb project: wandb.config.project=myproject (default=lewm)
• Specify model checkpoint interval: checkpoint_epoch_interval=5 (default=1)
• Specify number of epochs: trainer.max_epochs=50 (default=100)
• Specify batch size: loader.batch_size (default=128)
• Enable torch.compile: compile.enabled=true (default=false)
• Specify compile mode: compile.mode=max-autotune (default=max-autotune)

Checkpoints are saved to $STABLEWM_HOME upon completion.

For baseline scripts, see the stable-worldmodel scripts folder.

Decoder (Visualization)

To visualize what the latent representation looks like as an image, train a lightweight decoder on top of a frozen backbone. Two variants are available:

• train_predictor_decoder.py — decodes the predictor's z_{t+1} embedding back to pixels, supervised against the actual next frame.
• train_encoder_decoder.py — decodes the encoder's own [CLS] embedding back to pixels, supervised against the frame it was encoded from. This is the diagnostic decoder described in the paper, used to visualize what visual information is retained in the [CLS] representation.

# Standard LeWM
python train_predictor_decoder.py data=pusht backbone_checkpoint=/path/to/lewm_weights.pt

# Flow-matching LeWM
python train_predictor_decoder.py --config-name=decoder_fm data=pusht backbone_checkpoint=/path/to/lewm_fm_weights.pt

# Encoder-only [CLS] decoder
python train_encoder_decoder.py data=pusht backbone_checkpoint=/path/to/lewm_weights.pt

Set backbone_checkpoint to the _weight.ckpt (state dict) file for your environment. The encoder and predictor are frozen throughout; only the decoder is trained.

Enable wandb to log target vs. reconstructed image grids at each validation step:

python train_predictor_decoder.py data=pusht backbone_checkpoint=... wandb.enabled=true

Planning

Evaluation configs live under config/eval/. Set the policy field to the checkpoint path relative to $STABLEWM_HOME, without the _object.ckpt suffix:

# ✓ correct
python eval.py --config-name=pusht.yaml policy=pusht/lewm

# ✗ incorrect
python eval.py --config-name=pusht.yaml policy=pusht/lewm_object.ckpt

Pretrained Checkpoints

Pretrained LeWM checkpoints for each environment are mirrored on the Hugging Face Hub (model repos), alongside the datasets (dataset repos) in the same collection:

• quentinll/lewm-pusht
• quentinll/lewm-cube
• quentinll/lewm-tworooms
• quentinll/lewm-reacher

The full baseline checkpoint suite (PLDM, LeJEPA, IVL, IQL, GCBC, DINO-WM, DINO-WM-noprop) is available on Google Drive:

Method	two-room	pusht	cube	reacher
pldm	✓	✓	✓	✓
lejepa	✓	✓	✓	✓
ivl	✓	✓	✓	—
iql	✓	✓	✓	—
gcbc	✓	✓	✓	—
dinowm	✓	✓	—	—
dinowm_noprop	✓	✓	✓	✓

Loading a checkpoint

From the Drive archive

Each tar archive contains two files per checkpoint:

• <name>_object.ckpt — a serialized Python object for convenient loading; this is what eval.py and the stable_worldmodel API use
• <name>_weight.ckpt — a weights-only checkpoint (state_dict) for cases where you want to load weights into your own model instance

Place the extracted files under $STABLEWM_HOME/ and load via:

import stable_worldmodel as swm

# Load the cost model (for MPC)
cost = swm.policy.AutoCostModel('pusht/lewm')

AutoCostModel accepts:

• run_name — checkpoint path relative to $STABLEWM_HOME, without the _object.ckpt suffix
• cache_dir — optional override for the checkpoint root (defaults to $STABLEWM_HOME)

The returned module is in eval mode with its PyTorch weights accessible via .state_dict().

From the Hugging Face mirror

The HF model repos ship the LeWM checkpoint as a weights.pt (state dict) plus a config.json describing the model. Convert once to produce the _object.ckpt that eval.py expects:

# download weights.pt + config.json
hf download quentinll/lewm-pusht --local-dir $STABLEWM_HOME/hf_pusht

# convert to object checkpoint under $STABLEWM_HOME/pusht/lewm_object.ckpt
python - <<'PY'
import json, torch, stable_pretraining as spt
from pathlib import Path
from jepa import JEPA
from module import ARPredictor, Embedder, MLP
import stable_worldmodel as swm

src = Path(swm.data.utils.get_cache_dir(), "hf_pusht")
out = Path(swm.data.utils.get_cache_dir(), "pusht", "lewm_object.ckpt")

cfg = json.loads((src / "config.json").read_text())
encoder = spt.backbone.utils.vit_hf(
    cfg["encoder"]["size"],
    patch_size=cfg["encoder"]["patch_size"],
    image_size=cfg["encoder"]["image_size"],
    pretrained=False, use_mask_token=False,
)
mlp = lambda k: MLP(input_dim=cfg[k]["input_dim"], output_dim=cfg[k]["output_dim"],
                    hidden_dim=cfg[k]["hidden_dim"], norm_fn=torch.nn.BatchNorm1d)
model = JEPA(
    encoder=encoder,
    predictor=ARPredictor(**cfg["predictor"]),
    action_encoder=Embedder(**cfg["action_encoder"]),
    projector=mlp("projector"),
    pred_proj=mlp("pred_proj"),
)
sd = torch.load(src / "weights.pt", map_location="cpu", weights_only=False)
model.load_state_dict(sd, strict=True)
out.parent.mkdir(parents=True, exist_ok=True)
torch.save(model, out)
PY

After conversion, load via swm.policy.AutoCostModel('pusht/lewm') as usual.

Contact & Contributions

Feel free to open issues! For questions or collaborations, please contact lucas.maes@mila.quebec