A JEPA world model that trains end-to-end on one GPU without collapsing

What it does

LeWorldModel (LeWM) is a Joint-Embedding Predictive Architecture that learns to predict future latent embeddings directly from raw pixel observations. It serves as a compact world model for reinforcement learning and control, handling both 2D and 3D tasks. The model trains end-to-end with roughly 15 million parameters on a single GPU in a matter of hours, then plans actions significantly faster than foundation-model-scale alternatives.

The interesting bit

Instead of the usual JEPA bag of tricks—exponential moving averages, pretrained encoders, or auxiliary losses to prevent representation collapse—LeWM stays stable with just two loss terms: a next-embedding prediction loss and a Gaussian regularizer on latents. The authors also show the learned latent space captures actual physical quantities, and the model can flag physically implausible events through “surprise” evaluation.

Key highlights

• Trains stably end-to-end from pixels with only two loss terms, cutting tunable loss hyperparameters from six to one versus prior end-to-end JEPA work.
• Runs at ~15M parameters and trains on a single GPU in a few hours, yet plans up to 48× faster than foundation-model-based world models.
• Competitive performance across diverse 2D and 3D control benchmarks.
• Latent space probes reveal meaningful physical structure, and the model reliably detects physically implausible events.
• Pretrained checkpoints and datasets are available on Hugging Face for several environments (pusht, cube, two-rooms, reacher).

Caveats

• This repository is deliberately scoped to the core architecture and training objective; environment management, planning, and evaluation live in separate stable-worldmodel and stable-pretraining packages.
• Loading Hugging Face checkpoints requires a manual conversion script to produce the serialized object checkpoint that the evaluation code expects.

Verdict

Grab this if you are researching sample-efficient world models or want a lightweight, fast-planning JEPA baseline for control tasks. Skip it if you need a fully self-contained training and evaluation framework in one repo.