芯耀
与非AI
674
转载自公众号：敢敢AUTOHUB
0. 简介
如果你已经在做世界模型、模仿学习、机器人离线强化学习，那么过去两年大概率被这样一组矛盾的事实困扰过：一边是 DreamerV4、IRIS 把"在像素空间生成未来视频"做到极致，能在 Atari、Crafter、Counter-Strike 上端到端训出会玩游戏的 agent；另一边是 Yann LeCun 在 2022 年的蓝图里一再强调，真正可扩展的世界模型不应该再去预测每一个像素，而应该把环境的"内在状态"压缩到一个紧致的 latent 空间里，再在这个空间里预测未来。前者算力贵到几十亿参数起步，规划一次要几十秒；后者听起来美，但训练极其脆弱，模型经常退化到"把所有图像编码成同一个常数向量"这种平凡解。LeWorldModel 想做的，就是给后者一个干净、稳定、单 GPU 友好的训练配方。
如果你完全没接触过这条线，也没关系。这篇文章会从"世界模型到底要做什么"开始解释，一步步拆到 LeWM 的两项损失、SIGReg 正则的数学原理、AdaLN-zero 注入动作的工程技巧、以及推理时的隐空间规划流程。读完之后你应该能在脑子里画出整张训练—规划流水线，并且看着仓库里的train.py、module.py、jepa.py 知道每段代码对应论文哪条公式。
1. 开始阅读--十五兆参数、一张 GPU
LeWorldModel 由 Mila（蒙特利尔大学合作实验室）、纽约大学、Samsung SAIL、Brown 大学的研究人员合作完成，第一作者是 Lucas Maes 与 Quentin Le Lidec，资深作者包括图灵奖得主 Yann LeCun 和 Brown 大学的 Randall Balestriero。它发布于 2026 年初，是 LeCun 团队"自主机器智能（autonomous machine intelligence）"路线下一系列工作的延续：I-JEPA 把 JEPA 用在图像表征学习上，V-JEPA 与 V-JEPA-2 把它扩展到视频，LeJEPA 在 2025 年底给出了 SIGReg 这套理论保证，而 LeWorldModel 则是把 SIGReg 第一次"嫁接"到带动作的世界模型场景。
模型规模刻意压到了大约 15M 参数：一个 ViT-Tiny 编码器约 5M、一个六层因果 Transformer 预测器约 10M，加上轻量的动作编码器和投影头。论文明确说所有实验都是在单张 NVIDIA L40S GPU 上跑出来的，PushT 数据集训练 10 个 epoch 大约几小时即可完成。这个体量对工业落地非常友好——很多 JEPA 系列工作（V-JEPA-2 是十亿级、DINOv2 也是亿级）光是装载预训练权重就要分布式系统支持，而 LeWM 在一台单卡工作站上就能复现，社区里已经出现abdelstark/lewm-rs-pusht 这种纯 Rust 推理实现(https://huggingface.co/abdelstark/lewm-rs-pusht) 和多份独立解读，例如 Medium 上 Adnan Masood 的逐节读后感(https://medium.com/@adnanmasood/leworldmodel-and-the-case-for-stable-latent-world-models-0e4c33ca0f3c) 与 arxivIQ 的深度分析(https://arxiviq.substack.com/p/leworldmodel-stable-end-to-end-joint)。 这本身已经说明了一件事：好的论文不需要靠规模碾压，它需要的是别人愿意复现。
2. 什么是"世界模型"，为什么值得做
把"世界模型（World Model）"翻译成一句工程语言：它是一段神经网络，给定当前观测和未来若干步的候选动作，能够预测出未来若干步的状态或观测，从而让 agent 可以在不真正与环境交互的前提下"想象"接下来会发生什么。最早把这个概念形式化的是 Ha 与 Schmidhuber 2018 年的World Models 论文，他们用 VAE + RNN 在赛车游戏里做 imagination rollout；后来 Dreamer 系列把它扩展到通用强化学习，并在 Atari、DMC、Crafter 上接连刷新基准；到 DreamerV4 已经能在 Minecraft 这种 3D 沙盒里训练出会挖钻石的 agent。世界模型对工程师真正有吸引力的地方有两点：第一是离线强化学习——拿固定数据训出一个动力学模型，就可以用模拟数据做 policy improvement 而不必再触发真实环境；第二是测试时规划（test-time planning）——不用学一个固定 policy，而是每次决策时用模型展开多条候选动作序列，挑代价最低的那条执行，本质上是把 Model Predictive Control 这套经典控制论方法搬到了深度学习里。
世界模型有两条主要技术路线，差别在于"在哪个空间预测未来"。第一条是生成式路线：模型直接预测未来的像素帧（或离散 token），代表方法包括 IRIS、DIAMOND、Genie、DreamerV4 等。这条路的优点是表达力足、可视化直观、容易和大规模视频数据耦合；缺点是算力贵到夸张——Genie 训练用了几千张 TPU，DreamerV4 在 Minecraft 上用了上百块 GPU，而且推理时每一步都要解码图像，规划开销巨大。第二条是联合嵌入预测路线（Joint-Embedding Predictive Architecture，简称 JEPA）：先把观测编码成一个紧致的 latent 向量，再在 latent 空间里预测下一步向量，完全不重建像素。LeCun 在A Path Towards Autonomous Machine Intelligence 里把这条路定义为通往"会规划的 AI"的关键，因为它把"建模世界"和"渲染世界"解耦：模型只需要捕捉对决策有用的抽象状态，不必学会画出每一片树叶。
3. JEPA 美在哪里、又卡在哪里
JEPA 的吸引力可以用一句话概括：只学预测有用的部分。如果你的下游任务是"机器臂把方块推到目标位置"，那么图像里墙壁的纹理、光照的微小波动、远处的反光都是噪音；JEPA 通过编码器把图像压成低维向量后，模型只需要在这个向量空间里区分"方块在 A 位置"和"方块在 B 位置"，剩下的视觉细节自然被压缩掉。这种"任务无关、表征驱动"的设计哲学和当代自监督学习一脉相承——I-JEPA 用它在 ImageNet 上学到媲美 DINO 的特征，V-JEPA-2 用它在大规模无标注视频上学到了能用于规划的视频表征。理论上看，JEPA 还有一个非常优雅的性质：训练目标里没有像素重建项，因此模型不会被"如何精确画出未来图像"这件事拖累，可以把所有容量花在"如何预测未来语义"。
但 JEPA 的命门也在这里：预测损失自己会鼓励模型坍塌。把这件事讲清楚需要一点直觉。JEPA 的训练目标是
$$ \left\| \hat{z}_{t+1} - z_{t+1} \right\|^2 $$
其中$ \left\| \hat{z}_{t+1} \right\| $ 是预测器给出的未来 latent，z*t+1 是编码器给出的真实未来 latent。如果编码器学会把所有图像都映射到同一个常数向量z*，那么 $ \hat{z}_{t+1} = z_{t+1} = z^* $ 永远成立，损失瞬间归零。这就是 JEPA 圈反复出现的representation collapse（表征坍塌）：模型确实把损失压到了零，但学到的表征完全无用——任何两张图编码成同一个向量，下游探针自然恢复不出任何信息。多篇综述（如A Lightweight Library for Energy-Based JEPA 与JEPA-based World Models for LiDAR）都把"如何不坍塌"列为 JEPA 训练的头号工程难题。
围绕这个难题，社区诞生了一整套"防坍塌补丁"，每一种都是工程上不那么优雅的折中。
第一种是stop-gradient + EMA 教师：把目标 encoder 设为学生 encoder 的指数滑动平均（EMA），并在反向传播时切断目标分支梯度，让学生不能通过修改目标来让损失变小。这是 BYOL、I-JEPA、V-JEPA 系列的主流做法，但 Ponce 等人 2026 年的分析指出 stop-gradient + EMA 不对应任何良定义的目标函数，更像是经验上"碰巧能稳定"的工程技巧。
第二种是冻结预训练编码器：DINO-WM 直接用 DINOv2 的 patch 特征当 encoder，只学预测器，避免了坍塌但也意味着表征上限被预训练模型卡死，且模型无法适应特殊领域（医学影像、卫星图、游戏画面等）。第三种是多项方差—协方差正则：VICReg 强制 latent 的每一维方差不小于 1、各维相关性接近 0，PLDM 在此基础上又加了三项时间正则（time-sim、time-var、time-cov）和一项反向动力学损失，形成"七项 loss、六个权重"的复杂目标。后者的代价是显著的：作者必须为每个新环境扫描数百组超参，跨环境根本不可迁移。
LeWorldModel 的核心论点正是冲着这第三类做法去的：端到端 JEPA 不需要七项 loss，只需要一项预测损失加一项分布匹配正则就够了。下面三节会逐步把这件事拆开讲清楚。
4. LeWorldModel 的训练目标：把七项 loss 砍到两项
LeWorldModel 的全部训练目标可以写成一行公式：
$$ \mathcal{L}_{\mathrm{LeWM}} = \left\| \hat{z}_{t+1} - z_{t+1} \right\|_2^2 + \lambda \cdot \mathrm{SIGReg}(Z) $$
第一项是下一步嵌入预测损失（next-embedding prediction loss），第二项是SIGReg（Sketched Isotropic Gaussian Regularizer）高斯正则项， 是唯一需要调的权重，论文里默认取 0.1。这看似简单的两项目标，对应的工程实现就是train.py里lejepa_forward函数最核心的几行代码——下面这段是节选并加了中文注释的版本，可以直接对着论文公式逐行看：
# train.py:25-69（节选）
def lejepa_forward(self, batch, stage, cfg):
ctx_len = cfg.history_size # 历史窗口长度，PushT 默认 3
n_preds = cfg.num_preds # 预测偏移，默认 1（即预测下一步）
lambd = cfg.loss.sigreg.weight # SIGReg 权重，默认 0.1
output = self.model.encode(batch) # 像素 -> latent；动作 -> action embedding
emb = output["emb"] # (B, T, D) 状态 latent
act_emb = output["act_emb"] # (B, T, A_emb) 动作 latent
ctx_emb = emb[:, :ctx_len] # 历史状态
ctx_act = act_emb[:, :ctx_len] # 历史动作
tgt_emb = emb[:, n_preds:] # 真实未来 latent（监督信号）
pred_emb = self.model.predict(ctx_emb, ctx_act) # 预测未来 latent
# 第一项：next-embedding prediction loss
output["pred_loss"] = (pred_emb - tgt_emb).pow(2).mean()
# 第二项：SIGReg 高斯正则（注意把 batch 与时间维度交换）
output["sigreg_loss"] = self.sigreg(emb.transpose(0, 1))
# 总损失 = 预测损失 + λ · SIGReg
output["loss"] = output["pred_loss"] + lambd * output["sigreg_loss"]
return output把这段实现和 PLDM 的训练目标摆在一起对比，差距会非常刺眼。PLDM 来自 Sobal 等人 2025 年的Stress-testing offline reward-free RL，它继承自 VICReg 的"方差—不变性—协方差"三项正则，又在时间维度上叠加了 time-sim、time-var、time-cov 三项时序项，再加上一项反向动力学损失（IDM），总共七项 loss、六个超参权重。论文 §C.2 明确写着：作者必须在 PushT 上做 256 组配置的网格搜索，才能找到一组能把训练撑住的系数 $ (\alpha,\beta,\gamma,\zeta,\nu,\mu) = (18,12,0.2,0.7,0,0) $，而且这组系数无法迁移到其他环境，每个新数据集都要重新调。这种成本在工业落地上几乎是致命的——任何一个机器人新数据集都要扫一遍超参，而 LeWM 只需要二分搜索一个λ。
LeWM 之所以能把这些项压成一项，关键在于一个想法的转换：与其在统计矩层面（方差、协方差）做硬约束，不如直接在分布层面要求 latent 服从各向同性高斯分布。这件事乍听抽象，但物理意义非常直接——如果一个分布已经匹配了各向同性高斯，那么它的均值为 0、每一维方差为 1、各维不相关、彼此独立、所有边缘分布N（0,I）都对称——所有 PLDM 想用 var 项和 cov 项强制达到的统计特性，全都自动满足。一项分布约束等价于无穷多个矩约束的叠加，时间维度上的 time-var、time-cov 也自动随之消失，因为只要每个时间步的 latent 都已经是高斯，时间维度上的二阶统计量就不再需要单独维护。这就是论文反复强调的"七项变两项"的真正来源：不是工程上挑出最重要的两项，而是用一个数学上更紧的分布约束替换掉所有矩约束。
但这里立刻冒出一个问题：怎么在 192 维空间里检验一个分布是不是高斯？ 经典的多元正态性检验（Mardia 检验、Henze–Zirkler 检验）在维度超过 50 后就开始失效，要么需要不可承受的样本数，要么数值上崩溃。如果不能高效、可微地计算"分布偏离高斯多远"，整个想法就只是空中楼阁。LeWM 真正的技术贡献其实在这里——它复用了 Balestriero 与 LeCun 在 LeJEPA 里提出的SIGReg 正则化方案，把"高维分布是否高斯"这个看似不可解的问题，转化成一个可微、可缩放、还能用随机投影来无偏估计的目标函数。这正是下一节要拆开讲的核心。
5. SIGReg 是怎么把高维分布检验做成可微正则的
SIGReg 的全名是Sketched Isotropic Gaussian Regularizer——"sketched"在这里不是"草图"，而是统计文献里"用随机投影做高维降采样"的术语，类似于 sketching algorithm 这一支数据结构方法。它的整体思路可以拆成三步：
第一步，从单位球面 SD-1上随机采 M个方向u(m)；
第二步，把整个 latent 张量Z 投到每个方向上，得到 M个一维变量h(m)=Zu(m)；
第三步，对每个一维变量做 Epps–Pulley 正态性检验，并把统计量平均。
这样原本的"高维分布是否高斯"就被转化成" M个一维分布是否标准高斯"，而后者有非常成熟的数学工具可以处理。
这个转换合法的数学依据是Cramér–Wold 定理——这是 1936 年 Cramér 与 Wold 在Journal of the London Mathematical Society 上发表的一个经典结果。它说：如果两个 维分布在所有一维投影上的分布都相同，那这两个分布本身就相同。换句话说，匹配所有 1D 边缘分布等价于匹配整个联合分布。所以当随机方向数M→∞时，"M 个 1D 投影都是N(0,1)"在弱收敛意义下等价于"原 D维分布是N(0,I)"。LeWM 的实现里默认M=1024，论文消融实验显示这个数从 256 加到 4096 对下游成绩几乎没有影响——这恰恰是 Cramér–Wold 定理的弱收敛性质在工程上的体现。
第二步里那个Epps–Pulley 检验也值得说一句。这是 1983 年 Epps 与 Pulley 在Biometrika 上提出的正态性检验方法，它不直接比较概率密度函数，而是比较特征函数（characteristic function）：标准高斯的特征函数有闭式表达 $ \phi_0(t) = e^{-t^2/2} $，而样本的经验特征函数可以直接用 $ \phi_N(t) = \frac{1}{N}\sum_n e^{it h_n} $ 计算。检验统计量是这两个特征函数在加权L2意义下的差距：$ T = \int_{-\infty}^{\infty} w(t) \left| \phi_N(t) - \phi_0(t) \right|^2 dt $。
这个公式有两个 LeWM 看中的好处：第一，避免了概率密度估计需要核带宽这种额外超参；第二，整个统计量对样本是可微的，可以直接进梯度下降。SIGReg 在工程上把这个积分用梯形法离散到 17 个节点，于是计算变得非常便宜。
下面这段是module.py 里SIGReg 类的核心forward，整个类不到 50 行——你可以一边对着论文公式一边看代码：
# module.py:48-72（节选 + 注释）
def forward(self, proj):
# proj: (T, B, D)，由 emb.transpose(0, 1) 传入
# 第一步：随机采 num_proj 个单位向量
A = torch.randn(proj.size(-1), self.num_proj, device=proj.device)
A = A.div_(A.norm(p=2, dim=0)) # 单位化每个投影方向
# 第二步：投影到每个方向，并准备积分坐标 t
# x_t 形状 (T, B, num_proj, knots)
x_t = (proj @ A).unsqueeze(-1) * self.t
# 第三步：用经验特征函数与标准高斯特征函数的差做 Epps-Pulley 统计量
# cos / sin 分别对应特征函数的实部 / 虚部
err = (x_t.cos().mean(-3) - self.phi).square()
+ x_t.sin().mean(-3).square()
# 用梯形法把积分离散化，乘以样本数 B 完成标准缩放
statistic = (err @ self.weights) * proj.size(-2)
return statistic.mean()这段代码里有三个工程细节是值得专门点出来的，因为它们解释了为什么这套方案"能用、且能稳定用"。第一，投影方向 是 forward 时新生成的随机张量，不是注册的 buffer——这意味着每个 mini-batch 看到的随机方向都不一样，等价于一个无偏的蒙特卡洛估计器，跨 batch 累积起来就能逼近"对所有方向取期望"。第二，Epps–Pulley 用cos().mean() 和sin().mean() 分别算实部、虚部，避免了任何复数运算，PyTorch 原生算子全程 GPU 友好。第三，积分用梯形法在 上离散到 17 个节点，外层再乘上一个高斯权重窗口；论文消融显示节点数从 9 加到 33 对结果都没影响。这三个细节合起来就是 SIGReg 比 VICReg 那套硬约束更"稳"的原因——它的数学本质是分布匹配，工程实现却几乎没有数值陷阱。
最后还有一个容易被忽略但非常重要的细节，藏在config/train/model/lewm.yaml 里：encoder 输出后的投影头用的是BatchNorm1d，而不是 ViT 默认的LayerNorm。论文 §3.1 解释得很直白——ViT 最后一层会施加 LayerNorm，把每个样本的特征 z-score 到单位方差，这就把 SIGReg 想约束的"特征分布形状"提前压平了，导致正则项学不动。换成 BatchNorm 之后，归一化是跨样本而非跨特征维做的，不会破坏 SIGReg 想约束的整体分布。这是把 LeJEPA 的 SIGReg 思想搬到带动作的世界模型时，最容易踩坑、又最容易被忽视的一个工程细节。
6. 模型架构：ViT-Tiny 看图、AdaLN-zero 注入动作
LeWorldModel 的整体网络由四部分组成：图像编码器、动作编码器、自回归预测器、两个轻量投影头。论文用最小尺寸的搭配把总参数量压到了大约 15M——一个 Vision Transformer Tiny（12 层、3 头、隐藏维度 192）约 5M、一个六层 16 头的因果 Transformer 预测器约 10M、加上动作编码器和投影头总共不到 1M。所有模块都从零开始训练，没有任何预训练权重——这一点是与 DINO-WM 最根本的区别，也是论文标题里"end-to-end from pixels"的字面承诺。下表把四个模块和它们对应的论文角色、代码位置整理出来，方便对着仓库读代码：
| 模块 | 论文中对应角色 | 实现位置 | 关键参数 |
|---|---|---|---|
| Encoder | 把像素观测压成 latent | jepa.py:62-95encode() |
ViT-Tiny, patch=14, dim=192, ~5M |
| Action Encoder | 把动作向量映射到 latent 同尺度 | module.py:295-338Embedder |
Conv1d(k=1) + 2-layer MLP |
| Predictor | 在 latent 空间预测下一步 | module.py:382-440ARPredictor |
6 层 × 16 头 Transformer + AdaLN-zero |
| Projector / pred_proj | 整理 latent 维度供 SIGReg 与预测损失使用 | module.py:341-379MLP |
1 层 MLP + BatchNorm1d |
整个架构里最值得说的一件事，是预测器如何把动作信息注入 Transformer。最直接的做法是把动作 token 和状态 token 拼起来一起送进 attention，但这会让序列长度翻倍、attention 复杂度变成，规划时尤其吃亏。LeWM 选了一条更优雅的路——借用 Peebles 与 Xie 在 DiT（Diffusion Transformer）里发明的 AdaLN-zero：动作不占序列位置，而是通过一个 SiLU + Linear 把每一步的动作 embedding 映射成 6 组调制参数（attention 的 shift / scale / gate，加上 MLP 的 shift / scale / gate），这些参数再去调制 LayerNorm 之后的状态特征。下面这段是ConditionalBlock 的核心 forward，注释里把六个参数的作用都标了出来：
# module.py:155-193 ConditionalBlock 节选
class ConditionalBlock(nn.Module):
def __init__(self, dim, heads, dim_head, mlp_dim, dropout=0.0):
super().__init__()
self.attn = Attention(dim, heads, dim_head, dropout)
self.mlp = FeedForward(dim, mlp_dim, dropout)
# LayerNorm 不带可学习参数，scale/shift 完全由动作条件 c 生成
self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
# 一次线性投影同时生成 6 组调制参数
self.adaLN_modulation = nn.Sequential(
nn.SiLU(), nn.Linear(dim, 6 * dim, bias=True)
)
# AdaLN-zero 关键：最后一层初始化为 0
nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
def forward(self, x, c):
# 从动作条件 c 里同时切出 6 组参数
shift_a, scale_a, gate_a, shift_m, scale_m, gate_m =
self.adaLN_modulation(c).chunk(6, dim=-1)
# attention 分支：归一化 → 用动作调制 → 因果注意力 → 门控 → 残差
x = x + gate_a * self.attn(modulate(self.norm1(x), shift_a, scale_a))
# MLP 分支：同样的调制流程
x = x + gate_m * self.mlp (modulate(self.norm2(x), shift_m, scale_m))
return x这段代码里最关键的两行是初始化部分——最后一层 Linear 的权重和偏置都被强制初始化为 0。这意味着训练刚开始时shift / scale / gate 全部为 0，调制操作x * (1 + 0) + 0 = x 完全不改变特征，整个 ConditionalBlock 退化成一个恒等映射。换句话说，动作条件不是一上来就大力影响预测器，而是在训练过程中由模型自己慢慢"打开"。这个 zero 初始化技巧最早出现在 DiT 论文里，对扩散模型来说能让训练前期更稳定；对 JEPA 来说同样有效——预测器在初期不被强动作信号扰动，可以先把 encoder 的 latent 表征学稳，再逐步引入动作的影响。这是一个看似不起眼、但对端到端训练稳定性很关键的设计。
预测器的另一个关键设计是因果掩码（causal mask）。Attention.forward 里直接传is_causal=True 给 PyTorch 的scaled_dot_product_attention，第 步的 query 只能看 的 key，看不到未来。这件事让训练和推理共用同一份参数变得自然——训练时是 teacher-forcing，每一步都监督下一步预测；推理时是自回归 rollout，把预测出的 latent 接回输入继续往前滚。如果没有因果掩码，训练时模型会"作弊"看到未来 latent，推理时性能就会断崖式下降，这是序列建模里非常经典的训练-推理不一致问题。
这里还要补一个有意思的细节：模型架构里其实还有一个没有被使用的部件——像素解码器（decoder）。LeWM 的训练目标里完全没有重建项，但jepa.py 里仍然定义了一个轻量的 cross-attention 解码器。它只在论文里做"可视化诊断"——给定一个 latent，看模型能不能解码出大致正确的图像。论文 Fig 8 展示了一个有趣的现象：训练前期解码出的图像主要是"慢特征"（背景、墙体），后期才出现快速变化的物体（agent、方块）。论文 §G 还做了一个对照实验：如果把这个 decoder 加进训练目标里做联合优化，PushT 上的成功率反而从 96% 降到 86%。这印证了 JEPA 的核心哲学——重建像素是一种过强的监督，会让 encoder 浪费容量记忆视觉细节，而不是抽象任务相关的状态。
7. 训练好之后怎么用：在 latent 空间里规划动作
世界模型训练好之后，并不直接用来"做决策"——它只是一段动力学函数。真正在测试时把模型用起来的，是一套围绕它的测试时规划（test-time planning）流程。LeWM 选用的是经典的Model Predictive Control（MPC）+ Cross-Entropy Method（CEM） 组合，这套思路可以追溯到 1978 年 Richalet 等人提出 MPC 的工业控制论传统，2004 年 Rubinstein 与 Kroese 把 CEM 系统化为通用零阶优化算法，最近几年又被 TD-MPC2、DINO-WM、Navigation World Models 等工作搬到 latent 世界模型场景。整个流程的数学描述非常简洁：给定初始观测o1 和目标观测og，把它们都用 encoder 映射到 latent 空间得到z1和zg ；然后求解，
$$ a^*_{1:H} = \arg\min_{a_{1:H}} \left\| \hat{z}_H - z_g \right\|_2^2 $$
其中$ \hat{z}_H $ 是从z1出发、按候选动作序列在 latent 空间 rollout H步后得到的预测终点。整套规划过程完全不需要重建图像，所有比较都在 latent 距离上完成。
CEM 是这里用来求解$ a_{1:H}^* $ 的核心算法。它本质上是一种"采样—精英化—重采样"的迭代过程：每轮从一个高斯分布$ \mathcal{N}(\mu, \Sigma) $里采N条候选动作序列，全部用世界模型 rollout 评估代价，挑代价最低的K条作为"精英"，用这些精英的均值与方差更新分布参数，进入下一轮。LeWM 的默认超参是N=300，K=30，iter=30，规划 horizonH=5（再乘以 frame-skip 5，相当于规划 25 个真实环境步）。整体规划流程用伪代码表达就是这样，可以直接对照论文 Algorithm 2 阅读：
# CEM 在 latent 空间规划动作序列（伪代码，对应论文 Algorithm 2）
mu, sigma = zeros(H, A_dim), ones(H, A_dim) # 初始动作分布参数
for it in range(30): # CEM 迭代 30 次
# 第一步：采 300 条候选动作序列
actions = normal(mu, sigma, size=(300, H, A_dim))
# 第二步：在 latent 空间 rollout，得到每条候选的最终 latent 状态
z_init = encoder(o_1) # 仅一次编码初始观测
z_final = world_model.rollout(z_init, actions)
# 第三步：计算每条候选的代价（与目标 latent 的 L2 距离）
z_goal = encoder(o_g)
costs = ((z_final - z_goal) ** 2).sum(dim=-1)
# 第四步：挑代价最小的 30 条精英，更新分布参数
elites = actions[costs.argsort()[:30]]
mu = elites.mean(0)
sigma = elites.std(0) + 1e-3 # 防退化
# MPC：执行 mu 的前 K 步，再用新观测重新规划
execute_actions(mu[:5])这段代码里隐藏的工程亮点是world_model.rollout 这一步——它需要在一个批次里同时滚动 300 条候选轨迹。如果用 Python 循环逐条算，整套规划会慢到不可用；LeWM 的jepa.py:121-201 里的解法是把候选维度S 直接合并进 batch 维：原本(B, S, T, D) 的张量 reshape 成(B*S, T, D)，让 PyTorch 一次推完所有候选。下面这段是 rollout 函数最核心的循环结构：
# jepa.py 节选，自回归 rollout 的核心循环
emb = info["emb"].unsqueeze(1).expand(B, S, -1, -1)
emb = rearrange(emb, "b s ... -> (b s) ...") # 候选维并入 batch
act = rearrange(act_0, "b s ... -> (b s) ...")
act_future = rearrange(act_future, "b s ... -> (b s) ...")
HS = history_size # 截断历史窗口长度
for t in range(n_steps):
act_emb = self.action_encoder(act)
emb_trunc = emb[:, -HS:] # 只看最近 HS 步
act_trunc = act_emb[:, -HS:]
pred_emb = self.predict(emb_trunc, act_trunc)[:, -1:]
emb = torch.cat([emb, pred_emb], dim=1)
next_act = act_future[:, t:t+1, :]
act = torch.cat([act, next_act], dim=1)
pred_rollout = rearrange(emb, "(b s) ... -> b s ...", b=B, s=S)这里最关键的一行是emb_trunc = emb[:, -HS:]——它把每一步预测时模型看到的历史限制在最近HS 步内（PushT 与 OGBench-Cube 取 3，TwoRoom 取 1）。如果不做这个截断，序列会随着 rollout 步数越滚越长，attention 复杂度从 退化成，规划成本会随 horizon 平方增长。截断之后 rollout 的成本变成，整套规划在 L40S 上不到 1 秒就能完成。这件事看似简单，但它是 LeWM 能比 DINO-WM 快 48 倍的两个直接原因之一（另一个是单 CLS token vs 多 patch token，下一节讨论）。
最后还要提一下论文用的MPC 重规划策略——这是工业控制论里成熟的标准做法，被称为 receding-horizon control。简单说就是：每次 CEM 解出 步的最优动作序列后，只执行前 步（论文里），然后用环境的新观测重新启动一轮 CEM。这样做的目的是抵消 latent 空间 rollout 的累积误差——预测器虽然能展开 5 步，但越往后预测误差越大；只执行少数步并定期"重启"规划，相当于把 latent 误差控制在可接受范围。这是 latent 世界模型几乎所有工作（TD-MPC2、DINO-WM、NWM）都采用的策略，区别只在 horizon 与重规划间隔的具体数字。
8. 复现路径、适用边界与一份工程清单
如果你看到这里想动手跑一遍 LeWM，整套环境配置非常简单。仓库依赖stable-worldmodel 这个独立包来管理环境、规划器和评估，依赖stable-pretraining 来管理训练循环和数据 transform。换句话说，LeWM 仓库本身只剩"模型架构 + 训练目标"这两块核心代码——这是论文工程组织上的另一个亮点：把通用基础设施抽出去，让贡献本身可见。下面是从空环境到跑出 PushT 第一个 checkpoint 的最短路径：
# 1. 创建 Python 3.10 虚拟环境（推荐 uv，比 pip 快很多）
uv venv --python=3.10
source .venv/bin/activate
uv pip install stable-worldmodel[train,env]
# 2. 从 HuggingFace 下载 PushT 数据集，解压到 ~/.stable-wm/
# 数据格式是 HDF5；可通过 STABLEWM_HOME 环境变量改默认目录
tar --zstd -xvf pusht_expert_train.tar.zst -C ~/.stable-wm/
# 3. 训练（单卡 L40S 大约 4-6 小时，10 个 epoch）
python train.py data=pusht
# - 如果不想上传 WandB，确保 config/train/lewm.yaml 里 wandb.enabled=False
# - 训练完 checkpoint 自动保存到 $STABLEWM_HOME/pusht/lewm_object.ckpt
# 4. 评估：用 CEM 规划器跑 50 个目标条件 episode
python eval.py --config-name=pusht.yaml policy=pusht/lewm
# 5. 跳过训练直接用作者预训练 checkpoint
hf download quentinll/lewm-pusht --local-dir $STABLEWM_HOME/hf_pusht
# 再按 README 里的转换脚本生成 _object.ckpt 即可读代码时，仓库里最有"信噪比"的四处可以重点看——按推荐顺序是：先看train.py:25-76 的lejepa_forward 函数（看两项 loss 怎么拼起来），再看module.py:21-72 的SIGReg.forward（看高斯正则的实现），然后是module.py:155-193 的ConditionalBlock（看动作如何 AdaLN-zero 注入），最后是jepa.py:121-201 的JEPA.rollout（看规划时如何把候选维并入 batch）。这四段代码加起来不到 250 行，已经覆盖了论文 90% 的核心贡献。如果你不熟悉 PyTorch Lightning 或 Hydra，可以跳过train.py 后半段的 trainer 设置——那部分是 stable-pretraining 的标准模板，不影响理解 LeWM 本身。
什么时候应该用 LeWM、什么时候不应该？基于论文实验和方法本身的性质，可以给出一份相对明确的工程清单：适合的场景包括"只有一张 GPU、想跑端到端 JEPA"、"任务以 2D 位置或低维状态为主（导航、推物、平面机械臂）"、"需要快速规划接近实时控制"、"reward 不可得只有观测和动作"。不适合或需谨慎的场景包括"任务依赖细粒度旋转或 6-DoF 姿态"（CLS token 容量不够，建议用 patch token 或 DINO-WM）、"数据量极小且环境内禀维度低"（SIGReg 在低复杂度任务上反而过度正则化）、"训练目标必须包含 reward"（用 DreamerV4 或 TD-MPC2 更合适）、"需要 long-horizon 规划"（latent 自回归 rollout 误差累积，论文也承认这是当前局限）。这份清单不是绝对边界，但能帮你避开方法的明显陷阱。
仓库 Issue #13 还提到了一个值得关注的复现注意事项：用 HuggingFace 上的预训练 checkpoint 直接跑评估时，PushT 和 OGBench-Cube 上的成功率可能与论文报告值有差距——作者解释这是因为不同 random seed 与评估数据子集采样导致的合理波动。如果你做对比实验，建议要么用论文同一份评估配置（config/eval/pusht.yaml），要么从头训练自己的 checkpoint 而非直接用 HF 上的版本。社区里abdelstark/lewm-rs-pusht 这个纯 Rust 推理实现做到了与 PyTorch 参考实现激活级别 L∞ < 1e-4 的数值对齐，是另一个交叉验证基础设施。
9. 总结
读完所有实验数字之后，回过头看 LeWorldModel，会发现它最有价值的部分不在性能数字本身——四个环境上赢两个、平一个、输一个，这种成绩单在顶会论文里并不算亮眼。它真正改变了什么，是JEPA 训练的复杂度账本。社区这几年围着 JEPA 转的所有"防坍塌补丁"——EMA、stop-gradient、VICReg 七项正则、PLDM 的时间项叠加——本质上都是在对一个根本问题做局部修补：怎么让一个会自己鼓励坍塌的损失函数还能学到有用表征。LeWM 用 SIGReg 给出的回答是：问题不在损失函数，而在我们一直在用错的语言来描述"非坍塌"——用矩约束去逼近一个其实属于分布约束的目标。Cramér–Wold 定理 + Epps–Pulley 检验把这个观察落地成了一个具体、可微、能跑的训练目标。
这个观察的工程后果是连锁的。第一层后果是超参搜索复杂度从 降到——PLDM 的六个权重必须联调，LeWM 只剩一个，二分搜索一次就扫完。论文 §G 的消融实验显示 这个两个数量级的范围内成功率都保持在 80% 以上，这种宽容度在 JEPA 文献里几乎没有先例。第二层后果是训练曲线的可解释性——两项 loss 单调下降，没有多项 loss 互相拉扯的震荡，工程师可以像看普通监督学习一样判断训练是否健康。第三层后果是把 JEPA 的入门门槛压到了单 GPU——15M 参数、几小时训完，意味着研究生、独立研究者、甚至机器人初创公司的工程师都能在自己的工作站上复现论文，这本身就改变了这条研究路线的参与人群。
参考链接
• 主论文: Maes, Le Lidec, Scieur, LeCun, Balestriero.
LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels, arXiv:2603.19312, 2026.
• 项目主页: le-wm.github.io — 含 TL;DR、视频演示、checkpoint 集合
• 代码仓库:lucas-maes/le-wm — 本仓库中文逐行注释见docs/中文注释导读.md
• 预训练 checkpoint: HuggingFace Collectionquentinll/lewm
• 播客访谈: Embodied AI Podcast — LeWorldModel
• Balestriero & LeCun.LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics, arXiv:2511.08544, 2025. — SIGReg 的原始论文
• LeCun.A Path Towards Autonomous Machine Intelligence, 2022. — JEPA 蓝图
• Zhou et al.DINO-WM: world models on pre-trained visual features enable zero-shot planning, ICML 2025.
• Sobal et al.Stress-testing offline reward-free RL: a case for planning with latent dynamics models, Robot Learning Workshop 2025. — 即 PLDM
• Hansen et al.TD-MPC2: scalable, robust world models for continuous control, ICLR 2024.
• Hafner et al.Training agents inside of scalable world models, 2025. — DreamerV4
•When Does LeJEPA Learn a World Model? — 给 LeJEPA + 世界模型这条路线加形式化保证
•What Drives Success in Physical Planning with JEPA World Models? — 系统研究 JEPA latent 与物理规划成功率的关系
•Variance-Invariance-Sketching Regularization for JEPA training — 给 SIGReg 补 vanishing-gradient 修复
• arxivIQ Substack: 深度分析(https://arxiviq.substack.com/p/leworldmodel-stable-end-to-end-joint)
• Adnan Masood 在 Medium 的逐节读后感(https://medium.com/@adnanmasood/leworldmodel-and-the-case-for-stable-latent-world-models-0e4c33ca0f3c)
