ar27111994/lewm-implementation

LeWorldModel (LeWM): Stable End-to-End JEPA from Pixels

This repository contains a clean, self-contained PyTorch implementation of LeWorldModel (LeWM) from the paper:

LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels Lucas Maes, Quentin Le Lidec, Damien Scieur, Yann LeCun, Randall Balestriero arXiv: 2603.19312 — https://arxiv.org/abs/2603.19312 Official repo: https://github.com/lucas-maes/le-wm

🚀 Quick Start: Free GPU Training on Google Colab

The easiest way to train LeWM on a free GPU is via our Colab-ready notebook:

📓 [Open in Colab](https://colab.research.google.com/github/ar27111994/lewm-implementation/blob/main/lewm_colab.ipynb) (upload the notebook from this repo)

Or read the step-by-step guide: 📖 [COLAB_GUIDE.md](https://huggingface.co/ar27111994/lewm-implementation/blob/main/COLAB_GUIDE.md)

What you need:

• —A Google account (free)
• —~30–60 minutes for 10 epochs on synthetic data
• —Optional: Hugging Face token (free) to push trained models

Hardware: Free Colab T4 GPU (15 GB VRAM) — LeWM's ~18M parameters fit comfortably.

What is LeWorldModel?

LeWorldModel (LeWM) is a Joint-Embedding Predictive Architecture (JEPA) world model that learns directly from raw pixels with a single tunable hyperparameter. It is the first end-to-end JEPA that trains stably without:

• —Stop-gradient / EMA mechanisms
• —Pre-trained encoders (e.g., DINOv2)
• —Complex multi-term losses (e.g., VICReg variants)

Key Innovations

Architecture (from paper section 3.1 & Appendix D)

Raw Pixels (224x224) ---> ViT-Tiny Encoder ---> [CLS] + MLP+BN ---> Latent z_t
                                              |    (192-dim)
                                              |
                                              v
                                    +-------------------+
                                    |  AR Predictor     |  <--- Actions (AdaLN-zero)
                                    |  6 layers, 16h    |
                                    |  Causal masking   |
                                    +-------------------+
                                              |
                                              v
                                    Predicted z_{t+1}
                                              |
                                              v
                                    MSE(z_{t+1}, pred) + lambda * SIGReg(z)

Components:

• —Encoder: ViT-Tiny (patch 14, 12 layers, 3 heads, hidden 192) -> [CLS] token -> MLP + BatchNorm1d projector
• —Predictor: 6-layer transformer with AdaLN-zero action conditioning, causal temporal masking
• —SIGReg: Sketch Isotropic Gaussian Regularizer - anti-collapse via Epps-Pulley test on random 1-D projections
• —Planner: Cross-Entropy Method (CEM) in latent space for goal-conditioned control

SIGReg: The Anti-Collapse Engine

SIGReg is the critical component that makes stable end-to-end training possible.

Problem: Prediction-only loss causes representation collapse (encoder maps everything to a constant).

Solution: SIGReg forces latent embeddings to match an isotropic Gaussian N(0, I).

How it works:

1. 1.Collect latent tensor Z in R^(TxBxd) (time x batch x dim)
2. 2.Sample M=1024 random unit-norm directions u^(m) on the hypersphere S^(d-1)
3. 3.Project: h^(m) = Z dot u^(m) -> (T, B) 1-D marginals
4. 4.Apply the Epps-Pulley test statistic T(h^(m)) using the characteristic function
5. 5.Trapezoid quadrature on nodes uniformly in [0, 3] with weighting w(t) = exp(-t^2/2)
6. 6.By the Cramer-Wold theorem: matching all 1-D marginals <=> matching the full joint distribution

Key insight: The projector uses BatchNorm1d (not LayerNorm) because the ViT final layer already applies LayerNorm - this is essential for SIGReg optimization.

Training

Free GPU Training (Google Colab T4)

# In a Colab notebook with GPU runtime enabled:
!pip install -q transformers einops huggingface_hub matplotlib numpy tqdm

# Download implementation
from huggingface_hub import hf_hub_download
hf_hub_download("ar27111994/lewm-implementation", "lewm_model.py", local_dir="/content")
hf_hub_download("ar27111994/lewm-implementation", "lewm_train.py", local_dir="/content")

# Train with synthetic data (no 12GB download needed)
!python /content/lewm_train.py --use_synthetic \
    --n_episodes 2000 --epochs 10 --batch_size 128 \
    --lambd 0.1 --history_size 3 --seq_len 4 \
    --frameskip 5 --action_dim 2 --output_dir /content/drive/MyDrive/lewm

See [COLAB_GUIDE.md](https://huggingface.co/ar27111994/lewm-implementation/blob/main/COLAB_GUIDE.md) for the full notebook, troubleshooting, and real dataset download instructions.

Synthetic Smoke Test (CPU-friendly)

python lewm_train.py --use_synthetic \
    --n_episodes 2000 --epochs 10 \
    --batch_size 128 --lr 1e-3 \
    --lambd 0.1 --history_size 3 --seq_len 4 \
    --frameskip 5 --action_dim 2

Real PushT Dataset

1. 1.Download official dataset:

python -c "from huggingface_hub import hf_hub_download; \
    hf_hub_download('quentinll/lewm-pusht', 'pusht_expert_train.h5.zst', repo_type='dataset')"

1. 1.Decompress and train:

python lewm_train.py \
    --h5_path /path/to/pusht_expert_train.h5 \
    --epochs 10 --batch_size 128 \
    --lambd 0.1 --history_size 3

Hyperparameters (from paper)

Only lambda needs tuning - performance is insensitive to number of projections (M=1024) and integration knots (17).

Planning with CEM

from lewm_model import build_lewm, cem_plan

model = build_lewm(action_dim=10, history_size=3)
# ... load trained weights ...

best_actions = cem_plan(
    model,
    initial_pixels=context_frames,  # (1, H, C, 224, 224)
    goal_pixels=goal_frame,         # (1, 1, C, 224, 224)
    action_dim=10,
    horizon=5,        # 5 latent steps = 25 env steps (frame_skip=5)
    n_samples=300,
    n_iters=30,       # 30 for PushT, 10 for others
    n_elites=30,
    history_size=3,
)

Results (from paper)

• —48x faster planning than DINO-WM due to ~200x fewer tokens in latent space
• —Single GPU (L40S) training in "a few hours"
• —No stop-gradient, no EMA, no pre-trained encoders

Project Structure

lewm_model.py         - Core model (Encoder, Predictor, SIGReg, CEM)
lewm_train.py         - Training script (HDF5 + synthetic datasets)
lewm_mini_test.py     - Minimal sanity check
lewm_colab.ipynb      - Full Colab-ready training notebook
COLAB_GUIDE.md        - Step-by-step free GPU training guide
config.json           - Verified architecture config from official model
EXPLANATION.md        - 16KB deep-dive technical explanation

Interactive Demo

Try the explainable interactive Space (no training required):

🔗 https://huggingface.co/spaces/ar27111994/lewm-explainable

Features:

• —Architecture tab: Full pipeline schematic
• —SIGReg Explorer: Adjust collapse level and see real-time distributional analysis
• —CEM Planning: Visualize Cross-Entropy Method convergence
• —Key Results: Paper results and hyperparameters

Citation

@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}
}

License

MIT (same as the official repository).

This implementation is self-contained in standard PyTorch + transformers + einops, with no dependency on the private `stable-pretraining` or `stable-worldmodel` packages for the core model logic.

<!-- ml-intern-provenance -->

Generated by ML Intern

This model repository was generated by ML Intern, an agent for machine learning research and development on the Hugging Face Hub.

• —Try ML Intern: https://smolagents-ml-intern.hf.space
• —Source code: https://github.com/huggingface/ml-intern

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "ar27111994/lewm-implementation"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)

For non-causal architectures, replace AutoModelForCausalLM with the appropriate AutoModel class.