Leworldmodel分析解读与实验复现.md

LEWM世界模型分析解读

是不是看完LeWM论文还是一头雾水？是不是复现代码时被表征崩溃、SIGReg、高维正态分布绕晕？是不是搞不懂它为啥能比 JEPA、PLDM 更简单稳定？

这篇文章用最通俗的话、最清晰的结构，把LeWM从头到尾讲明白，一步到位带你吃透最新的LeWM世界模型算法，不管是入门世界模型还是深入科研都能用！

一、先搞懂最核心的痛点：表征崩溃到底是个啥？

LeWM 这篇论文从头到尾，就是为了解决世界模型训练里的表征崩溃问题。

举一个简单的例子来解释表征崩溃问题：

输入一个汽车运行时候的画面，让世界模型推断下一秒车在什么位置，然后输出对应的特征向量，学会环境里的物理规律。但是由于模型的编码器和预测器都是一起训练的，那模型就会想要“走捷径”，不管你的输入是什么画面，他编码器都编码输出一个固定的0向量，那接下来预测未来状态的时候，未来的状态也全是 0，这样预测值和真实值的误差直接完美等于0，损失函数直接拉到最低，模型直接 “躺赢”。当你把这个训练好的模型拿去做下游任务时，它完全无法区分不同的图像，因为它眼里所有的东西特征都一样，这就意味着模型白训练了，提取的特征完全是没有用的。

这个问题可以说是自监督学习里的一个比较常见的问题，只要你想让编码器和预测器一起端到端从头训，就大概率会踩这个坑。

二、现有解决方法都有啥？又藏着哪些坑？

第一类：靠外部方法 “治标不治本”

这类问题不从模型本身解决问题，全部采用外部手段来给模型“打补丁”，主要有这么几种：

1）启发式规则：靠经验定一些规则来限制模型，比如强行约束特征的均值、方差不能太离谱。但问题也很明显，换个任务、换个数据集，规则就失效了，基本没啥泛化性。

2）多目标损失函数：在主预测损失之外，加一堆辅助损失来约束特征。比如论文里提到的PLDM，直接用了7个损失函数，但是各个损失之间的权重得反复试，稍微调不好，训练曲线就会出现不稳定的情况。

3）直接用预训练模型：比如DINO-WM，直接拿一个在海量数据上预训练好的DINOv2视觉大模型当编码器，训练的时候直接把编码器参数冻住，完全不更新。这样确实绝对不会崩，因为特征提取能力是已经训练好的。但这完全违背了自监督学习从头学特征的初衷，属于 “逃避问题” 而不是解决问题。更关键的是，模型的上限直接被预训练模型锁死了，比如你做机械臂抓取的任务，预训练模型是在自然风景图上训的，根本不懂你任务里的关键特征，你想微调都没机会，只能重新换预训练模型，平白增加了一堆工作量。

第二类：靠内部方法 “管用但说不清”

这类方法就是JEPA系列模型（I-JEPA、V-JEPA）用的方案，比如采用Stop-Gradient（SG，停止梯度）+ Exponential Moving Average（EMA，指数移动平均）。

SG：模型分成两部分，一个是学生（在线编码器 + 预测器），一个是老师（目标编码器）。学生做错了题，只惩罚学生、更新学生的参数，绝对不碰老师的参数，这就是停止梯度。为啥要这么做？要是允许老师跟着学生的错误调整，老师为了让学生拿满分，直接会把所有答案都改成0，俩人一起摆烂，直接发生崩溃，切断梯度，就是断了模型走捷径的退路。那老师的参数怎么来的？就用到了EMA。

EMA：老师的参数是学生过去一段时间参数的缓慢平均值，学生在前面慢慢学，老师在后面稳稳地跟着，既保证老师能给学生提供有意义的学习目标，又保证老师的参数极其稳定，不会突然崩掉。

这套方法效果确实不错，现在也用得很多，但它有个致命的问题：数学上根本解释不通为什么管用。

传统机器学习的框架是什么？你有一个明确的损失函数，训练的时候严格按照微积分求导、算梯度，然后顺着梯度下降的方向更新参数，找损失的最低点。但这套方法里：

停止梯度，就是人为把反向传播的梯度流给剪断，你更新参数用的，根本不是整个系统真实的导数；EMA中老师的参数更新，根本不是算梯度、求导算出来的，是强行把学生过去的参数拿过来做加权平均（比如老师的新参数 = 0.99 × 老师旧参数 + 0.01 × 学生新参数）。

也就是说这就是个经验性的黑盒操作，管用是管用，但理论上站不住脚，调参还特别复杂，稍微改个超参数，模型就可能直接崩了。

三、LeWM 的核心解法：到底是怎么从根上解决崩溃问题的？

现在来看看LeWM的核心创新，它用一个特别巧妙的方法，直接从根上解决了编码器和预测器同时训练时的表征崩溃问题，而且数学上有严谨的理论支撑，实操还特别简单。

先搞懂 LeWM 的基础模型结构

Leworldmodel的基础架构特别简单，就两个核心模块：编码器 Encoder + 预测器 Predictor。

$$ \text{Encoder: } \quad z_t = \text{enc}_\theta(o_t) $$

$$ \text{Predictor: } \quad \hat{z}_{t+1} = \text{pred}_\phi(z_t, a_t) $$

1）编码器用的是Vit架构，输入当前时刻的环境画面 $o_t$，把它压缩编码成一个高维的潜在特征向量 $z_t$，给画面做特征提取。

2）预测器接收编码器输出的当前特征 $z_t$，还有智能体当前执行的动作 $a_t$，预测出下一时刻的特征表示 $z_{t+1}$，让模型在潜在空间里，推演环境接下来会怎么变，学会物理规律。

那有了模型，怎么训练呢？这就来到了lewm最最核心的创新点，lewm用一个很巧妙的方法，解决了编码器和预测器同时训练时，编码器会偷懒将所有的输入都映射到一个常数向量的问题。如以下代码：

# LeWM loss
output["pred_loss"] = (pred_emb - tgt_emb).pow(2).mean()
output["sigreg_loss"]= self.sigreg(emb.transpose(0, 1))
output["loss"] = output["pred_loss"] + lambd * output["sigreg_loss"]  

可以看出，相比于论文中同时提到的PLDM采用了7个损失函数，导致训练曲线不稳定lewm仅采用了两个loss函数。

pred_loss这个很好理解，就是世界模型训练的核心目标是预测器预测出来的下一时刻特征pred_emb，和编码器真实编码出来的下一时刻特征tgt_emb，算一个均方误差，这个损失越小，说明模型预测未来状态的能力越强，对环境物理规律的学得越准。

但只靠这个损失，模型一定会走捷径，发生表征崩溃，这时候就需要lewm这篇论文的核心创新点，SIGReg损失。

首先我们从前面了解到，表征崩溃的本质是什么？是模型把所有输入都映射成了同一个常数向量（比如全0），那这个特征的分布，就完全不是正态分布了，是一个固定的单点。那反过来想，如果我们能逼着模型输出的高维特征，必须符合一个完美的高维正态分布，那它就不可能把所有特征都变成同一个常数，从根本上就断了它走捷径的可能。

但这里有个难题：一维数据的正态分布检验，有很多成熟好用的方法，比如Epps-Pulley检验，但一个高维特征，比如论文里用的192维向量，想判断它是不是符合高维正态分布就很难。

LeWM 是怎么解决这个问题的？它用了一个特别巧妙的方法：随机投影 + 一维正态性检验，也就是SIGReg方法。

$$ \text{SIGReg}(Z) \triangleq \frac{1}{M} \sum_{m=1}^{M} T(h^{(m)}) $$

$$ \mathcal{L}_{\text{LeWM}} \triangleq \mathcal{L}_{\text{pred}} + \lambda \text{SIGReg}(Z) $$

只需要在d维的特征空间里，随机找M个不同的方向，把高维特征全部投影到这M条一维的线上，就能得到M组一维的特征数据。然后对这M组一维数据，分别用成熟的Epps-Pulley检验，逼着它们都符合一维正态分布就行。

可能读论文很难理解，这里给大家举一个很通俗易懂的例子：

你想检查一个悬浮在空中的复杂云团，是不是一个完美的球体（对应高维正态分布），直接去测量它的三维体积、曲面曲率，太难了。于是你拿手电筒，从M个随机的角度照向这个云团，去看墙上的影子，如果每个角度照出来的影子，都是完美的圆形（对应一维正态分布），那这个云团，大概率就是个完美的球体。

LeWM就是这个逻辑，训练的时候不断随机抽取M个投影方向，逼着特征在这些方向上都符合一维正态分布。随着训练的进行，这就等价于把整个高维特征空间，塑造成了完美的高维正态分布。模型想把所有特征都映射成0？不可能，因为全0的分布根本过不了正态性检验，SIGReg损失会直接拉满，模型想走捷径都不行。

最关键的是，整个方法超简单，你只需要调整两个超参数：随机投影的维度M，还有两个损失之间的权重lambda，大大简化了超参数的选择，新手也能快速上手。

同时，这里有个数学定理给咱们兜底，叫Cramér–Wold定理：一个高维数据，如果它在所有可能的一维投影方向上都符合正态分布，那这个高维数据本身，一定是高维正态分布。

四、训好的 LeWM 怎么用？如何用世界模型驱动智能体

训练世界模型，最终目的不是让它预测特征，而是要让它真正驱动智能体，在环境里完成任务，比如机器人导航、机械臂抓取。这部分就是为了给你讲明白，训练好的LeWM，怎么落地成真正能执行的动作。

世界模型的核心功能，就是通过当前的画面和动作，预测未来的状态，然后在脑海里模拟试错，找到能完成目标的最优动作。这里将论文中比较浓缩难懂的文字给你拆解成一步一步的具体流程，比论文里的描述好懂10倍。

第一步：先给目标和当前状态做编码

$$ \hat{z}_1 = \text{enc}_\theta(o_1), \quad z_g = \text{enc}_\theta(o_g) $$

使用训练好的世界模型编码器，将当前的真实图像以及任务目标的最终图像进行编码器编码，将两张图片压缩成192维的潜在特征向量，分别记作 $z_1$和 $z_g$。

最终目标就是找到一系列动作，让智能体从当前的 $z_1$状态，最终走到 $z_g$这个目标状态。

第二步：在脑海里疯狂试错，模拟未来轨迹

$$ \hat{z}_{t+1} = \text{pred}_\phi(\hat{z}_t, a_t) $$

随机生成一大批候选动作序列，每个序列里都包含了未来H步的动作（比如 H=10，就是规划未来10步要做什么）。

然后把每一个候选动作序列，都输入到LeWM的预测器里，从当前状态 $z_1$开始，结合第一个动作 $a_1$生成下一步的 $z_2$，再结合 $a_2$生成 $z_3$，就这样自回归地一直预测到第H步的最终状态 $z_h$。

这里的核心就是利用世界模型在潜在空间里，把每一套动作方案会带来什么结果，全都提前模拟一遍，不用真的在现实环境里试错。

第三步：筛选最优动作，迭代优化

$$ \mathcal{C}(\hat{z}_H) = |\hat{z}_H - z_g|_2^2 $$

$$ a_{1:H}^* = \arg\min_{\boldsymbol{a}_{1:H}} \mathcal{C}(\hat{z}_H), $$

模拟完所有方案之后，咱们计算每一套方案最终得到的 $z_h$，和目标状态 $z_g$之间的差距，差距越小，说明这套动作越好，越能帮我们完成任务。

然后我们挑出表现最好的前k个方案，根据这些优秀方案的均值和方差，更新动作采样的分布，就像进化一样，好的动作被保留下来，差的被淘汰，经过几次迭代，就能得到一个最优动作序列。

第四步：执行动作，走一步看一步

虽然规划了未来H步的最优动作，但不会把这H步一次性全执行完。因为世界模型的预测是自回归的，每一步的预测都会带一点微小的误差，随着预测步数增加，误差会不断累积，要是一次性执行完H步，最后智能体可能离目标越来越远。

五、LeWM 对比现有方案，优势到底在哪？

讲完了整个逻辑，总结一下，LeWM对比之前的方案，到底好在哪，为啥它值得你深入研究：

1）从根上解决表征崩溃，理论严谨：不像SG+EMA的黑盒操作，LeWM从特征分布的本质出发，用严谨的数学定理做支撑，减轻了模型走捷径的可能。

2）超参数极少，训练简单稳定：整个模型只需要调M和lambda两个超参数，对比多损失函数的方案、JEPA系列的复杂调参，友好很多，训练曲线也特别平滑。

3）端到端从头训练，不用依赖预训练模型：不需要采用训练好的大模型冻住编码器的方式，完全可以根据你的任务，从零开始训一个专属的世界模型，模型的表达能力和上限完全不受限，特别适合机器人、具身智能这种垂类任务。

4）架构极简，方便修改和拓展：没有复杂的老师-学生双架构，就是编码器 + 预测器的极简结构，你想加什么模块、做什么改进都特别方便。

六、当前的局限以及可以深挖的科研方向

当前存在的局限

1）潜在空间维度的适配问题：论文里用的是192维的潜在空间，面对 tworoom 这种维度很低的简单任务时，会出现 “杀鸡用牛刀” 的情况，反而可能破坏特征的自然表达，导致模型性能下降。

2）复杂真实环境的适配性待验证：目前的实验，基本都集中在比较干净、简单的仿真环境里。但现实生活中，会有光照变化、遮挡、动态物体等各种干扰，模型在这种复杂环境里的表现，还需要进一步的研究和验证，比如可以通过增大M的维度，来提升模型对复杂特征的拟合能力。

3）对动作标签的依赖：目前模型的训练，需要输入每一时刻对应的动作标签，但现实里的大量视频数据，都是没有动作标注的，这就限制了模型用海量互联网数据做预训练的可能。

可以深挖的科研方向

1）自适应潜在空间维度设计：可以设计一个能根据任务复杂度，自动调整潜在空间维度的方案，简单任务用低维度，复杂任务用高维度，解决 “杀鸡用牛刀” 的问题。

2）结合逆动力学建模，摆脱动作标签依赖：论文里也提到了这个方向，与其依赖外部给定的动作标签，不如给模型加一个逆动力学模块，让模型自己从连续的画面里，反推出对应的动作。这样就能直接用无标签的互联网视频，让模型学习通用的物理规律，大幅提升模型的泛化性。

七、实验复现避坑指南

1. 配置环境

有两种方案，一种是官方给出的uv管理：

uv venv --python=3.10
source .venv/bin/activate
uv pip install stable-worldmodel[train,env]

另一种大家也可以使用我们熟悉的conda管理：

conda create -n lewm python=3.10
conda activate lewm
pip install stable-worldmodel[train,env]

2. 下载数据

这边推荐下载tworoom和cube这两个数据，因为官方在google drive中提供的模型权重，由于进不去网站，所以下载不了，因此，我们从零开始训练了tworoom和cube这两个环境，得到了两个权重，由于卡和时间有限，tworoom训练了完整的轮次，但是cube只训练了12轮，大家想要尝试的可以自行训练。（官方提供的huggingface下的模型pt权重不是eval所需要的ckpt权重，不用下载，或者可以自行修改测试文件，将ckpt加载换成pt加载）

自己训练的模型权重下载地址：https://huggingface.co/Datawhale/lewm/tree/main

tworoom任务数据地址：https://huggingface.co/datasets/quentinll/lewm-tworooms/tree/main

cube任务数据地址：https://huggingface.co/datasets/quentinll/lewm-cube/tree/main

pusht任务数据地址：https://huggingface.co/datasets/quentinll/lewm-pusht/tree/main

reacher任务数据地址：https://huggingface.co/datasets/quentinll/lewm-reacher/tree/main

为每一个任务新建文件夹，例如tworoom任务的目录结构如下：

lewm
└── tworoom
    └── tworoom
        ├── lewm_object.ckpt
        └── tworoom.h5

其中tworoom.h5为解压的数据文件，lewm_object.ckpt为模型权重文件

cube任务有些不一样

lewm
└── cube
    ├── ogbench
    │   └── cube_single_expert.h5
    └── cube
        └── lewm_object.ckpt

其他的任务可以自行设置路径

3. 训练和测试

首先需要设置路径：

这里一定要设置成多一个文件夹的路径，否则模型权重会在你的lewm下保存，会显得很杂乱。后面每一步都需要加载这个路径。

export STABLEWM_HOME=/home/your-user/lewm/tworoom

训练部分：

在config/train/lewm.yaml中添加你自己的wandb：https://wandb.ai/home

wandb:
  config:
    entity: your_entity
    project: your_project

启动训练：

data后面填写自己想要训练的任务

python train.py data=tworoom

推理测试：

可以使用我们自己训练的权重，也可以使用你自己训练的权重。后续嗯更换任务仅需修改两处tworoom即可，同时，为了区分，上传的时候对tworoom和cube权重进行了重命名，测试的时候将tworoom和cube提供的权重名称将放入各自任务的文件夹之后，统一修改成lewm_object.ckpt这个名称。

python eval.py --config-name=tworoom.yaml policy=tworoom/lewm

4. 效果展示

tworoom效果展示

==== RESULTS ====
metrics: {'success_rate': 98.0, 'episode_successes': array([ True,  True,  True, False,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True]), 'seeds': None}
evaluation_time: 179.38973236083984 seconds

可以看到，tworoom简单导航任务的成功率较高

     

cube效果展示

==== RESULTS ====
metrics: {'success_rate': 64.0, 'episode_successes': array([False, False,  True,  True,  True,  True, False,  True,  True,
       False,  True, False,  True,  True, False,  True,  True, False,
        True,  True,  True,  True, False, False, False,  True, False,
       False,  True, False,  True, False,  True,  True,  True,  True,
       False,  True, False,  True,  True,  True, False,  True,  True,
        True, False,  True,  True,  True]), 'seeds': None}
evaluation_time: 191.06767058372498 seconds

可以看到，机械臂抓取任务的成功率较低，同时，我们发现，其实可能成功率还要更低，可能是只训练了12轮的缘故，后续可以试试训练官方提供的99轮，但是官方的源码中判断成功标准如以下代码：

def _compute_successes(self):
    cube_successes = []
    for i in range(self._num_cubes):
        obj_pos = self._data.joint(f'object_joint_{i}').qpos[:3]   # 方块当前位置
        tar_pos = self._data.mocap_pos[self._cube_target_mocap_ids[i]]  # 目标位置

        if np.linalg.norm(obj_pos - tar_pos) <= 0.04:  # 欧氏距离 ≤ 4cm
            cube_successes.append(True)
        else:
            cube_successes.append(False)
    return cube_successes

只要方块位置对就算成功，其实有时候机械臂的状态也是成功率判断的一个标准，因此可能会有些偏差。

     

八、致谢

本文内容基于LeWM原论文的核心思想展开，十分感谢原论文作者们的优秀工作。希望这篇通俗解读，能帮到更多想要入门世界模型、具身智能的同学，也欢迎大家一起交流，补充更多的复现经验和优化思路。