LeWorldModel (LeWM)

Summary: LeWM is an end-to-end JEPA world model that learns directly from pixels using only a prediction loss + SIGReg regularizer — no pretrained visual features (unlike DINO-WM), no reconstruction decoder, no temporal smoothness loss. Achieves 96% success on Push-T with ViT-S predictor, significantly outperforming PLDM (78%) and DINO-WM (80.67%).

Details

Architecture

  • Encoder: ViT (ViT-S default) — processes 3 context frames → latent embeddings z₁:₃
  • Predictor: ViT-S (optimal), autoregressive, action-conditioned — predicts future latents zₜ given zₜ₋₁ and action aₜ₋₁
  • Decoder: Only for visualization/probing; NOT used during training (no reconstruction loss)
  • Regularizer: SIGReg (Sketched-Isotropic-Gaussian Regularizer) — single scalar λ

SIGReg (Only Regularizer)

  • Enforces zₜ ~ 𝒩(0, I) at each timestep via sketched Epps-Pulley distance
  • Random projections + characteristic function matching
  • Robust hyperparameters: K (projections) and M (knots) insensitive; only λ matters (works in [0.01, 0.2])

Training

  • Objective: ℒ = ℒ_pred + λ ℒ_SIGReg
  • Data: 10K episodes per env, 10 epochs
  • Stability: Low variance across seeds (Push-T: 96.0±2.83% vs PLDM 78.0±5.0%)

Planning

  • Solver: CEM (Cross-Entropy Method) — best for LeWM (96%) and PLDM (78%)
  • Budget: 50 steps, goal 25 steps ahead
  • Other solvers tested: SGD, RMSProp, Adam — all significantly worse

Benchmarks

Environment Task LeWM (ViT-S) PLDM DINO-WM
TwoRoom 2D Navigation 94% 94% 90%
Push-T 2D Manipulation 96% 78% 80.67%
OGBench-Cube 3D Manipulation 82% 84% 80%
Reacher 2-link Control 94% 86% 88%

Probing Results (Linear / MLP)

  • Position quantities: LeWM ≥ PLDM (often better)
  • Rotational/dynamic (velocity, yaw): DINO-WM > LeWM ≈ PLDM (pretrained visual priors help)
  • Block orientation: All struggle

Key Ablations

Ablation Push-T SR Insight
Full LeWM (ViT-S pred, λ=0.09) 96%
No decoder / no recon loss 96% → 86% Recon loss hurts planning
ViT-T predictor 80.67% Too small
ViT-B predictor 86.7% Too big, overfits
Predictor dropout p=0.1 96% Optimal regularization
ViT encoder vs ResNet-18 96% vs 94% ViT slightly better
λ ∈ [0.01, 0.2] >80% Robust
K (projections), M (knots) Stable Insensitive

Temporal Straightening (Emergent)

  • Latent velocity vectors v_t = z_{t+1} - z_t become increasingly collinear during training
  • Measured by mean cosine similarity: S_straight = 𝔼[cos(v_t, v_{t+1})]
  • Emerges without temporal regularization (PLDM has explicit temporal smoothness but lower straightness)
  • Hypothesis: SIGReg per-timestep leaves temporal dimension unconstrained → temporal collapse → straighter paths → better planning

Significance

  1. Simplest JEPA world model yet: 1 loss term (prediction) + 1 regularizer (SIGReg)
  2. End-to-end from pixels: No DINOv2 dependency (DINO-WM needs pretrained features)
  3. No reconstruction: Proves JEPA prediction alone captures planning-relevant info
  4. No temporal loss: Temporal structure emerges as beneficial geometry
  5. Robust hyperparameters: Only λ needs tuning; K, M, predictor dropout, embedding dim all have wide sweet spots

Related Concepts

Sources

  • leworldmodel-paper: Abstract, Section 1, Section 3 — End-to-end JEPA from pixels with SIGReg; no pretrained features, no recon, no temporal loss
  • leworldmodel-paper: Section 4.1, Table 1 — Push-T 96% vs PLDM 78% vs DINO-WM 80.67%
  • leworldmodel-paper: Appendix D, Table 7 — No decoder loss improves planning (96% vs 86%); SIGReg only regularizer