le-wm README

The project's README from github.com/lucas-maes/le-wm. As one markdown file: README.md.

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:

uv venv --python=3.10
source .venv/bin/activate
uv pip install stable-worldmodel[train,env]

Data

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

tar --zstd -xvf archive.tar.zst

Place the extracted .h5 files under $STABLEWM_HOME (defaults to ~/.stable-wm/). You can override this path:

export STABLEWM_HOME=/path/to/your/storage

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

Training

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

Before training, set your WandB entity and project in config/train/lewm.yaml:

wandb:
  config:
    entity: your_entity
    project: your_project

To launch training:

python train.py data=pusht

Checkpoints are saved to $STABLEWM_HOME upon completion.

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

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:

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:

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:

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 [email protected]

le-wm (lucas-maes/le-wm)

self-supervised-learning
Searching Database…

4.5K ·GitHub
Le-wm
What is this?This is the official codebase for LeWorldModel (LeWM), a Joint-Embedding Predictive Architecture (JEPA) that learns a world model directly from pixels using just two loss terms, avoiding the usual tricks like EMA, pretrained encoders, or auxiliary supervision.
Is this a standalone library or does it depend on other repos?It's a research codebase that relies on two companion packages: stable-worldmodel handles environments, planning and evaluation, while stable-pretraining handles training, leaving this repo focused on jepa.py (the model architecture and training objective).
How do I train a LeWM model?Training is driven by Hydra config files under config/train/ (set your WandB entity/project in config/train/lewm.yaml, then launch training); checkpoints are saved to $STABLEWMHOME when done.
Where do the training datasets come from?Datasets are distributed as HDF5 files on HuggingFace; you download and decompress them, then place the .h5 files under $STABLEWMHOME (defaults to ~/.stable-wm/).
Do I have to train everything from scratch?No - pretrained LeWM checkpoints for the pusht, cube, two-room and reacher environments are mirrored on the Hugging Face Hub, and a broader baseline suite (PLDM, LeJEPA, IVL, IQL, GCBC, DINO-WM) is available via Google Drive.
Do I need a big GPU cluster to use this?No - LeWM has only about 15M parameters and trains on a single GPU in a few hours, and it plans up to 48x faster than foundation-model-based world models.
What kinds of tasks has LeWM been tested on?The paper evaluates it on diverse 2D and 3D control tasks, including the pusht, cube, two-room, and reacher environments, where pretrained checkpoints and baselines are also provided.
What makes LeWM different from other JEPA-based world models?Compared to the only other end-to-end JEPA alternative, LeWM cuts tunable loss hyperparameters from six down to one, since it only needs a next-embedding prediction loss plus a Gaussian regularizer to avoid representation collapse.
Does the model learn more than just control policies?Yes - probing the latent space shows it encodes meaningful physical quantities, and a surprise-evaluation test shows the model reliably flags physically implausible events.
How do I load a pretrained checkpoint?Checkpoints come as an object.ckpt (a serialized Python object used by eval.py) plus a weight.ckpt (state dict); from the Hugging Face mirror you first convert the shipped weights.pt/config.json into an object.ckpt, then load it via swm.policy.AutoCostModel.
What are similar projects?transferlearning, lerobot, SolarWM, video-prediction-policy-2, Zero-WAM, wuji-ego-mint, Code-as-World, OraRL
Which topics is it in?self-supervised-learning, world-models, paper, simulation, dataset, embeddings, reinforcement-learning, academic-research, huggingface, computer-vision, pytorch, benchmark
Creator
lucas-maes
Size
5.3 MB
Stars
4.5K
Updated
CommitMay 22
Age
7mo
Commits · 90 Days