Domain Analysis: Deep Learning & Robotics Research
Persona: Senior Principal Research Scientist (AI/Robotics)
Abstract
This research introduces LeWorldModel (LeWM), a streamlined Joint-Embedding Predictive Architecture (JEPA) designed for stable, end-to-end world modeling from raw pixels. Traditional JEPA implementations frequently suffer from representation collapse, necessitating complex heuristics like stop-gradients, exponential moving averages (EMA), or multi-term loss functions (e.g., VICReg’s seven-term objective). LeWM simplifies this into a two-term objective: a mean-squared error (MSE) next-embedding prediction loss and a Sketched-Isotropic-Gaussian Regularizer (SIGReg). By enforcing Gaussian-distributed latent embeddings, LeWM achieves training stability with only one tunable loss hyperparameter ($\lambda$).
Experimental results across 2D and 3D control benchmarks (PushT, OGBench-Cube, Reacher) demonstrate that LeWM, despite its compact 15M-parameter footprint, achieves planning speeds up to 48× faster than foundation-model-based alternatives like DINO-WM while maintaining competitive success rates. Probing analyses confirm the latent space recovers high-fidelity physical quantities (agent/object coordinates), and Violation-of-Expectation (VoE) tests indicate the model identifies physical anomalies with higher sensitivity than visual perturbations. Furthermore, the model exhibits an emergent "path straightening" in its latent trajectories, suggesting an optimized internal representation of temporal dynamics.
LeWorldModel (LeWM): Synthesis of Architectural Innovations and Performance Metrics
[Sec 1] The End-to-End JEPA Challenge: Current World Models (WMs) typically bifurcate into generative models (pixel-space prediction) or JEPAs (latent-space prediction). JEPAs are theoretically superior for planning but historically unstable. LeWM provides a robust solution that is task-agnostic, reconstruction-free, and trainable on a single GPU in hours.
[Sec 3.1] Architectural Configuration:
Encoder: Vision Transformer (ViT-Tiny, ~5M parameters) mapping observations ($o_t$) to a 192-dimensional latent space ($z_t$).
Predictor: 6-layer Transformer (~10M parameters) using Adaptive Layer Normalization (AdaLN) to integrate actions ($a_t$) and autoregressively predict future states ($z_{t+1}$).
[Sec 3.1] The Anti-Collapse Objective (SIGReg): To prevent the model from mapping all frames to a single point (collapse), LeWM employs SIGReg. This projects embeddings onto $M=1024$ random directions and applies the Epps–Pulley normality test. This encourages the latent distribution to match an isotropic Gaussian, ensuring feature diversity without the need for complex contrastive pairs or EMA encoders.
[Sec 3.2] Latent Planning via MPC: The model utilizes the Cross-Entropy Method (CEM) for trajectory optimization. By rolling out predictions in a compact latent space rather than pixel space, the system performs Model Predictive Control (MPC) with significantly reduced computational overhead.
[Sec 4.2] Planning Efficiency and Latency:
Speed: LeWM demonstrates a 48× speedup in planning time compared to DINO-WM (foundation-model-based), largely due to the use of fewer tokens and a more compact representation.
Performance: On the PushT benchmark, LeWM’s pixels-only training outperformed DINO-WM configurations that included additional proprioceptive data.
[Sec 5.1] Representation Fidelity (Probing): Probing tests reveal that LeWM’s latent space accurately encodes "Agent Location" and "Block Position" (Pearson $r > 0.97$). While a decoder was not used during training, a post-hoc decoder successfully reconstructed visual scenes from the 192-dim [CLS] token, proving the latent space retains essential environmental structure.
[Sec 5.2] Violation-of-Expectation (VoE): The model was tested for "surprise" (prediction error) when encountering unphysical events. LeWM assigns significantly higher surprise values to physical teleportation (continuity violations) than to simple visual color changes, indicating a learned understanding of environmental "laws."
[Appendix H] Emergent Path Straightening: A notable discovery is that LeWM’s latent trajectories become naturally "straighter" (higher cosine similarity between velocity vectors) over time. This mimics biological visual processing and occurs as an emergent property without explicit temporal smoothing losses.
[Key Takeaway] Resource Efficiency: LeWM lowers the barrier to entry for World Model research by enabling SOTA-level planning and physical reasoning using minimal parameters and single-GPU training, bypassing the need for massive "foundation" vision encoders.
Target Review Audience
The ideal reviewers for this work include:
Robot Learning Researchers: Focus on the MPC planning speed and success rates in OGBench and PushT.
Self-Supervised Learning (SSL) Specialists: Evaluation of SIGReg versus VICReg or contrastive methods for collapse prevention.
Computational Neuroscientists: Analysis of the "path straightening" phenomenon and its alignment with biological representation learning.
Hardware/Optimization Engineers: Assessment of the 48× planning speedup and inference-time FLOP efficiency.