le-wm

LeWorldModel (LeWM) is the official code base for the paper "LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels" by Lucas Maes, Quentin Le Lidec, Damien Scieur, Yann LeCun and Randall Balestriero. It implements a Joint-Embedding Predictive Architecture (JEPA) that, according to the abstract, 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. The authors state this reduces tunable loss hyperparameters from six to one compared with the only existing end-to-end alternative, and that the model has roughly 15M parameters trainable on a single GPU in a few hours. The repository builds on two external projects: stable-worldmodel for environment management, planning and evaluation, and stable-pretraining for training. The stated goal is to reduce this repository to its core contribution, the model architecture and training objective, implemented in jepa.py as a PyTorch module. Installation uses uv with Python 3.10 and installs stable-worldmodel with the train and env extras. Datasets are distributed in HDF5 format via Hugging Face and are decompressed with tar --zstd; extracted .h5 files are placed under $STABLEWM_HOME (default ~/.stable-wm/), which can be overridden with an environment variable. Dataset names are referenced without the .h5 extension in Hydra config files under config/train/. Training is launched with python train.py data=pusht after setting WandB entity and project in config/train/lewm.yaml. Checkpoints are saved to $STABLEWM_HOME on completion. Evaluation configs live under config/eval/, and the policy field must point to a checkpoint path relative to $STABLEWM_HOME without the _object.ckpt suffix. Pretrained LeWM checkpoints are mirrored on the Hugging Face Hub for the pusht, cube, tworooms and reacher environments, alongside datasets in the same collection. A broader baseline checkpoint suite (PLDM, LeJEPA, IVL, IQL, GCBC, DINO-WM, DINO-WM-noprop) is available via Google Drive, with a table indicating which methods cover which environments. Checkpoints come in two forms: an _object.ckpt serialized Python object used by eval.py and the stable_worldmodel API, and a _weight.ckpt weights-only state dict. Loading is done through swm.policy.AutoCostModel with a run name relative to $STABLEWM_HOME; the returned module is in eval mode and exposes .state_dict(). For Hugging Face mirrors, the README provides a conversion script that rebuilds the JEPA model from config.json and weights.pt and saves an object checkpoint. The README also links the paper, a project website and the Hugging Face collection, and invites issues and collaboration via email. It reports that LeWM plans up to 48x faster than foundation-model-based world models while remaining competitive across diverse 2D and 3D control tasks, that its latent space encodes physical structure as shown by probing physical quantities, and that surprise evaluation detects physically implausible events. These are claims from the authors' abstract, not independently verified here.