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

01 动机

训练 Joint Embedding Predictive Architecture (JEPA) 的核心难题是表征坍塌（representation collapse）：若不加约束，编码器会将所有输入映射到同一点或低秩流形，预测误差趋零但表征毫无意义。现有方案须依赖复杂的多项损失、指数移动平均（EMA）目标网络、预训练编码器或辅助监督，工程复杂度高且超参数众多。

"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 正则化通过统计检验强制潜在空间服从各向同性高斯分布，从而防止坍塌。

2训练所需损失项数量（vs. PLDM 的 7 项）

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

15M模型参数量

数小时单 GPU 上的训练时长

方法对比定位 — 图 2：LeWM 在端到端训练与基于基础模型方法之间的定位对比。相较于 PLDM 等端到端方法，LeWM 具备更简洁的训练目标；相较于 DINO-WM 等基础模型方法，LeWM 规划效率高出约 48 倍。

02 方法

LeWM 由一个轻量级 Vision Transformer 编码器和一个带 Adaptive Layer Normalization 动作注入的 Transformer 预测器组成，以两项损失端到端训练：预测损失（ℒ_pred） + SIGReg 正则化，无需 EMA 目标网络或预训练特征。规划阶段采用 Cross-Entropy Method (CEM) + Model Predictive Control (MPC)。

潜在规划流程 — 图 4：潜在规划工作流。给定初始帧与目标帧的嵌入，CEM 在潜在空间中优化动作序列，使终态嵌入与目标嵌入的距离最小化。MPC 每步只执行前 K 个动作后重新规划，缓解自回归误差累积。

编码器与预测器架构

编码器采用 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 在每层注入，初始化为零以保证训练初期稳定性。

两项训练损失

预测损失 ℒ_pred

对下一帧嵌入的均方误差：

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

其中 ẑ_t+1 为预测器输出，z_t+1 为编码器对真实下一帧的输出。

SIGReg 正则化

基于 Epps–Pulley 统计检验，对 M=1024 个随机投影方向逐一检验潜在表征是否服从高斯分布，并将偏差作为正则化损失：

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

默认 λ=0.1，是整个方法唯一需要调整的超参数（PLDM 需调 6 个）。

03 实验

实验在四个离线、无奖励的控制环境中评估：Push-T（2D 操作）、OGBench-Cube（3D 机器人）、Two-Room（2D 导航）、Reacher（运动规划）。基线包括端到端方法 PLDM 和基于 DINO 基础模型的 DINO-WM。

各环境规划性能对比 — 图 6：四个环境下的规划性能对比。LeWM 在 PushT 和 Reacher 任务上尤为突出，整体与 DINO-WM 相当，但规划速度快 48×。Two-Room 环境因内在维度低，SIGReg 较难匹配高维高斯先验，性能相对下降。

关键性能指标

评估维度	PLDM	DINO-WM	LeWM（本文）
PushT 成功率	基准	相当	高出 PLDM 18%
规划时间	—	较慢	比 DINO-WM 快 48×
损失项数量	7	N/A（冻结编码器）	2
Block Location MSE（Push-T）	0.011±0.066	0.009±0.052	0.001±0.006
训练损失曲线	嘈杂、非单调	—	平滑、单调收敛

解码潜在预测（OGBench-Cube） — 图 7：OGBench-Cube 上的解码潜在预测。解码结果保留了全局场景结构（物体位置、背景），但丢失了部分细节，表明潜在空间已编码场景的高层语义，而非精确像素重建。

物理理解评估

论文通过两项实验验证 LeWM 是否学习了物理先验：

线性探针（Probing）：对 Push-T 环境，LeWM 潜在空间的 Block Location MSE 为 0.001±0.006，显著优于 PLDM 的 0.011±0.066，说明位置信息被有效编码。
违反预期（Violation-of-Expectation）：LeWM 对物体瞬移（teleportation，物理违规）赋予显著更高的"惊讶度"（p < 0.01），而对颜色变化（视觉扰动）则反应较小，表明模型确实编码了物理约束而非纯视觉统计。

消融分析

实验表明：SIGReg 的投影数量 M 对性能影响可忽略；嵌入维度超过某阈值后性能饱和；用 ResNet-18 替换 ViT 编码器仍能获得有竞争力的结果。整个方法仅 λ 需要调整，可通过 O(log n) 的二分搜索确定（PLDM 需 O(n⁶) 网格搜索）。

04 局限性

Note: 以下局限性部分为作者在论文中明确陈述（标注 stated），部分由设计推断（标注 inferred）。

短视界规划限制（stated）

"Planning remains restricted to short horizons, motivating hierarchical world modeling for long-horizon reasoning." 论文指出当前方法仅能在短时间步内有效规划，长程任务需要分层世界模型。

低多样性数据集上的 SIGReg 困难（stated）

在 Two-Room 等内在维度较低的环境中，SIGReg 难以将低维数据的潜在分布对齐至高维各向同性高斯，导致性能相对下降。

对离线数据集覆盖率和动作标签的依赖（stated）

方法依赖具有充分覆盖率且包含动作标签的离线数据集。未来工作提出预训练于大规模视频数据并引入逆动力学建模，以减少这一约束。

解码器细节丢失（inferred）

由图 7 可见，重建的 OGBench-Cube 序列保留了场景全局结构但细节模糊，说明潜在空间是有损压缩，精细操作任务中的像素级精度可能不足（论文未明确讨论此局限）。