LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels

01 动机Motivation

训练 Joint Embedding Predictive Architecture (JEPA) 的核心难题是表征坍塌（representation collapse）：若不加约束，编码器会将所有输入映射到同一点或低秩流形，预测误差趋零但表征毫无意义。现有方案须依赖复杂的多项损失、指数移动平均（EMA）目标网络、预训练编码器或辅助监督，工程复杂度高且超参数众多。The core difficulty in training a Joint Embedding Predictive Architecture (JEPA) is representation collapse: without constraints, the encoder maps all inputs to a single point or a low-rank manifold, so the prediction error goes to zero while the representation carries no meaning. Existing solutions must rely on complex multi-term losses, exponential moving average (EMA) target networks, pre-trained encoders, or auxiliary supervision, which brings high engineering complexity and a large number of hyperparameters.

"Existing methods for training JEPAs from pixels rely on complex multi-term losses, exponential moving averages, pre-trained encoders, or auxiliary supervision."

训练流程概览 — 图 1：LeWM 的训练流程。编码器将视频帧映射到潜在表征，预测器自回归地建模动态，SIGReg 正则化通过统计检验强制潜在空间服从各向同性高斯分布，从而防止坍塌。Figure 1: The LeWM training pipeline. The encoder maps video frames to latent representations, the predictor models the dynamics autoregressively, and SIGReg regularization uses a statistical test to force the latent space to follow an isotropic Gaussian distribution, thereby preventing collapse.

2训练所需损失项数量（vs. PLDM 的 7 项）Loss terms needed for training (vs. 7 for PLDM)

48×相比 DINO-WM 的规划加速比Planning speedup over DINO-WM

15M模型参数量Model parameter count

数小时Hours单 GPU 上的训练时长Training time on a single GPU

方法对比定位 — 图 2：LeWM 在端到端训练与基于基础模型方法之间的定位对比。相较于 PLDM 等端到端方法，LeWM 具备更简洁的训练目标；相较于 DINO-WM 等基础模型方法，LeWM 规划效率高出约 48 倍。Figure 2: Where LeWM is positioned between end-to-end training and foundation-model-based methods. Compared with end-to-end methods such as PLDM, LeWM has a simpler training objective; compared with foundation-model methods such as DINO-WM, LeWM is about 48 times more efficient at planning.

02 方法Method

LeWM 由一个轻量级 Vision Transformer 编码器和一个带 Adaptive Layer Normalization 动作注入的 Transformer 预测器组成，以两项损失端到端训练：预测损失（ℒ_pred） + SIGReg 正则化，无需 EMA 目标网络或预训练特征。规划阶段采用 Cross-Entropy Method (CEM) + Model Predictive Control (MPC)。LeWM consists of a lightweight Vision Transformer encoder and a Transformer predictor with action injection through Adaptive Layer Normalization, trained end-to-end with two losses: a prediction loss (ℒ_pred) + SIGReg regularization, requiring neither an EMA target network nor pre-trained features. Planning uses the Cross-Entropy Method (CEM) + Model Predictive Control (MPC).

潜在规划流程 — 图 4：潜在规划工作流。给定初始帧与目标帧的嵌入，CEM 在潜在空间中优化动作序列，使终态嵌入与目标嵌入的距离最小化。MPC 每步只执行前 K 个动作后重新规划，缓解自回归误差累积。Figure 4: The latent planning workflow. Given the embeddings of the initial and goal frames, CEM optimizes an action sequence in latent space so that the distance between the terminal-state embedding and the goal embedding is minimized. MPC executes only the first K actions at each step and then replans, mitigating the accumulation of autoregressive error.

编码器与预测器架构Encoder and Predictor Architecture

编码器采用 Vision Transformer Tiny，约 5M 参数（12 层、3 个 attention head、192 维 hidden，14 像素 patch），[CLS] token 经过单层 MLP + Batch Normalization 投影至潜在向量。预测器为 6 层 Transformer（16 个 attention head、10% dropout，约 10M 参数），动作通过 Adaptive Layer Normalization 在每层注入，初始化为零以保证训练初期稳定性。The encoder is a Vision Transformer Tiny with about 5M parameters (12 layers, 3 attention heads, 192-dimensional hidden size, 14-pixel patches); its [CLS] token is projected to the latent vector through a single-layer MLP + Batch Normalization. The predictor is a 6-layer Transformer (16 attention heads, 10% dropout, about 10M parameters); actions are injected at every layer through Adaptive Layer Normalization and initialized to zero to ensure stability in the early stage of training.

两项训练损失The Two Training Losses

预测损失 ℒ_predPrediction loss ℒ_pred

对下一帧嵌入的均方误差：The mean squared error on the next-frame embedding:

ℒ_pred ≜ ‖ẑ_t+1 − z_t+1‖²₂

其中 ẑ_t+1 为预测器输出，z_t+1 为编码器对真实下一帧的输出。where ẑ_t+1 is the predictor output and z_t+1 is the encoder output on the true next frame.

SIGReg 正则化SIGReg regularization

基于 Epps–Pulley 统计检验，对 M=1024 个随机投影方向逐一检验潜在表征是否服从高斯分布，并将偏差作为正则化损失：Based on the Epps–Pulley statistical test, the latent representation is tested for Gaussianity along each of M=1024 random projection directions, and the deviation is used as a regularization loss:

ℒ_LeWM ≜ ℒ_pred + λ · SIGReg(Z)

默认 λ=0.1，是整个方法唯一需要调整的超参数（PLDM 需调 6 个）。The default is λ=0.1, the only hyperparameter the whole method needs to tune (PLDM needs to tune 6).

03 实验Experiments

实验在四个离线、无奖励的控制环境中评估：Push-T（2D 操作）、OGBench-Cube（3D 机器人）、Two-Room（2D 导航）、Reacher（运动规划）。基线包括端到端方法 PLDM 和基于 DINO 基础模型的 DINO-WM。Experiments are evaluated in four offline, reward-free control environments: Push-T (2D manipulation), OGBench-Cube (3D robotics), Two-Room (2D navigation), and Reacher (motion planning). Baselines include the end-to-end method PLDM and DINO-WM, which builds on the DINO foundation model.

各环境规划性能对比 — 图 6：四个环境下的规划性能对比。LeWM 在 PushT 和 Reacher 任务上尤为突出，整体与 DINO-WM 相当，但规划速度快 48×。Two-Room 环境因内在维度低，SIGReg 较难匹配高维高斯先验，性能相对下降。Figure 6: Planning performance across the four environments. LeWM is particularly strong on the PushT and Reacher tasks and is overall comparable to DINO-WM, while planning 48× faster. In the Two-Room environment the intrinsic dimensionality is low, making it harder for SIGReg to match a high-dimensional Gaussian prior, so performance degrades relatively.

关键性能指标Key Performance Metrics

评估维度Evaluation aspect	PLDM	DINO-WM	LeWM（本文）LeWM (ours)
PushT 成功率PushT success rate	基准Baseline	相当Comparable	高出 PLDM 18%18% higher than PLDM
规划时间Planning time	—	较慢Slower	比 DINO-WM 快 48×48× faster than DINO-WM
损失项数量Number of loss terms	7	N/A（冻结编码器）N/A (frozen encoder)	2
Block Location MSE（Push-T）Block Location MSE (Push-T)	0.011±0.066	0.009±0.052	0.001±0.006
训练损失曲线Training loss curve	嘈杂、非单调Noisy, non-monotonic	—	平滑、单调收敛Smooth, monotonic convergence

解码潜在预测（OGBench-Cube） — 图 7：OGBench-Cube 上的解码潜在预测。解码结果保留了全局场景结构（物体位置、背景），但丢失了部分细节，表明潜在空间已编码场景的高层语义，而非精确像素重建。Figure 7: Decoded latent predictions on OGBench-Cube. The decoded results preserve the global scene structure (object positions, background) but lose some detail, indicating that the latent space encodes the high-level semantics of the scene rather than an exact pixel reconstruction.

物理理解评估Evaluating Physical Understanding

论文通过两项实验验证 LeWM 是否学习了物理先验：The paper verifies whether LeWM has learned physical priors through two experiments:

线性探针（Probing）：对 Push-T 环境，LeWM 潜在空间的 Block Location MSE 为 0.001±0.006，显著优于 PLDM 的 0.011±0.066，说明位置信息被有效编码。Linear probing: in the Push-T environment, the Block Location MSE of the LeWM latent space is 0.001±0.006, markedly better than PLDM's 0.011±0.066, showing that positional information is effectively encoded.
违反预期（Violation-of-Expectation）：LeWM 对物体瞬移（teleportation，物理违规）赋予显著更高的"惊讶度"（p < 0.01），而对颜色变化（视觉扰动）则反应较小，表明模型确实编码了物理约束而非纯视觉统计。Violation-of-Expectation: LeWM assigns a significantly higher "surprise" score to object teleportation (a physical violation) (p < 0.01), while reacting far less to color changes (a visual perturbation), indicating that the model does encode physical constraints rather than purely visual statistics.

消融分析Ablation Analysis

实验表明：SIGReg 的投影数量 M 对性能影响可忽略；嵌入维度超过某阈值后性能饱和；用 ResNet-18 替换 ViT 编码器仍能获得有竞争力的结果。整个方法仅 λ 需要调整，可通过 O(log n) 的二分搜索确定（PLDM 需 O(n⁶) 网格搜索）。The experiments show that the number of SIGReg projections M has negligible influence on performance; performance saturates once the embedding dimension exceeds a certain threshold; and replacing the ViT encoder with ResNet-18 still gives competitive results. The whole method needs only λ to be tuned, which can be determined by an O(log n) binary search (PLDM requires an O(n⁶) grid search).

04 局限性Limitations

Note: 以下局限性部分为作者在论文中明确陈述（标注 stated），部分由设计推断（标注 inferred）。Some of the limitations below are explicitly stated by the authors in the paper (marked stated), while others are inferred from the design (marked inferred).

短视界规划限制（stated）Planning limited to short horizons (stated)

"Planning remains restricted to short horizons, motivating hierarchical world modeling for long-horizon reasoning." 论文指出当前方法仅能在短时间步内有效规划，长程任务需要分层世界模型。The paper notes that the current method can only plan effectively over short time steps, and that long-horizon tasks require hierarchical world models.

低多样性数据集上的 SIGReg 困难（stated）SIGReg difficulties on low-diversity datasets (stated)

在 Two-Room 等内在维度较低的环境中，SIGReg 难以将低维数据的潜在分布对齐至高维各向同性高斯，导致性能相对下降。In environments with low intrinsic dimensionality such as Two-Room, SIGReg has difficulty aligning the latent distribution of low-dimensional data with a high-dimensional isotropic Gaussian, which leads to a relative drop in performance.

对离线数据集覆盖率和动作标签的依赖（stated）Dependence on offline dataset coverage and action labels (stated)

方法依赖具有充分覆盖率且包含动作标签的离线数据集。未来工作提出预训练于大规模视频数据并引入逆动力学建模，以减少这一约束。The method depends on an offline dataset with sufficient coverage that also contains action labels. Future work proposes pre-training on large-scale video data and introducing inverse-dynamics modeling to relax this constraint.

解码器细节丢失（inferred）Loss of decoder detail (inferred)

由图 7 可见，重建的 OGBench-Cube 序列保留了场景全局结构但细节模糊，说明潜在空间是有损压缩，精细操作任务中的像素级精度可能不足（论文未明确讨论此局限）。As Figure 7 shows, the reconstructed OGBench-Cube sequences preserve the global structure of the scene but are blurry in detail, indicating that the latent space is a lossy compression and that pixel-level precision may be insufficient for fine-grained manipulation tasks (the paper does not explicitly discuss this limitation).