LeWorldModel is a world model architecture that learns stable end-to-end joint-embedding predictive representations directly from raw pixels.
The project addresses the fragility of existing Joint Embedding Predictive Architectures, which typically require complex multi-term losses, exponential moving averages, pretrained encoders, or auxiliary supervision to prevent representation collapse. LeWorldModel simplifies this by training stably with only two loss terms: a next-embedding prediction loss and a regularizer that enforces Gaussian-distributed latent embeddings. This reduces tunable hyperparameters significantly compared to existing end-to-end alternatives. The model contains approximately fifteen million trainable parameters and can be trained on a single GPU in a few hours, making it computationally efficient for research and practical applications.
The tool suits researchers and practitioners building world models for control tasks who want a lightweight, interpretable alternative to foundation-model-based approaches. It plans substantially faster than larger foundation models while remaining competitive across diverse two-dimensional and three-dimensional control benchmarks. Beyond control, the latent space encodes meaningful physical structure that can be probed for physical quantities, and the model reliably detects physically implausible events through surprise evaluation. This makes it valuable for applications requiring both planning capability and physical understanding.
The codebase is built on top of two existing repositories for environment management, planning, and evaluation, and for training infrastructure, which reduces the repository to its core contribution in model architecture and training objective. The project provides pretrained checkpoints and datasets on Hugging Face, alongside configuration files for training and evaluation via Hydra. Installation and usage are documented with clear paths for data setup, training launch, and planning evaluation, with baseline scripts available in the upstream stable-worldmodel repository.
