机器学习笔记(25): LeWorldModel

可能更好的阅读体验 Jeefy's Blog: 机器学习笔记(25): LeWorldModel

其实模型设计非常简单，是一个非常简单的范式。损失函数也非常简单：

\[{\mathcal L}_{LeWM} = {\mathcal L}_{pred} + \lambda {\rm SIGReg}({\bf Z}) \]

这个范式仅仅是学习 \(p(z_{t+1} | z_t, a_t)\)，不带任何策略的设计。
动作使用 CEM solver 求解。

---

众所周知，世界模型非常大的一个问题就是表征可能坍缩，也就是 \(enc(x_t)\) 可能坍缩为一个常数，或者变成一个分布极度单峰或者不均衡的分布。所以需要一种 “正则化” 的手段，防止表征坍缩。

维度	完全重建损失	SIGReg	VICReg
约束对象	嵌入必须能还原出输入本身（ground-truth 锚定）	嵌入分布须匹配 \(\mathcal{N}(0,I)\) 的完整分布形状	只约束二阶统计量：各维方差 ≥ 阈值 + 维度间去相关
防坍缩保证	结构性免疫——常数解的重建误差巨大，不存在捷径	充要条件：\(\mathrm{SIGReg}\to 0 \Leftrightarrow P_Z \to \mathcal{N}(0,I)\)	必要条件：方差够大且不相关即可，分布形状任意
信息量要求	无损：输入的全部信息必须保住	只锁统计形状，内容完全自由	最松，只锁尺度与去相关
需要额外模块	解码器 \(D\)（训练期激活）	随机投影 + Epps–Pulley 统计量	无
有效超参	权重 1 个 + decoder 容量（隐性超参）	1 个（\(\lambda\)；\(M\) 实验上不敏感）	每个正则项一个权重（PLDM 扩到 4~6 个）+ 铰链常数
计算开销	最高	\(O(M \cdot Nd)\)，与维度解耦	\(O(d^2)\) 协方差矩阵
与预测损失的梯度关系	重建梯度可能劫持表征	与 \(\mathcal{L}_{\text{pred}}\) 目标正交，loss 曲线互不拉扯	多项正则互相竞争，训练曲线非单调（PLDM 实测）
典型失效模式	容量浪费在无关细节上；decoder 瓶颈逼 E 有损编码	环境内在维度低时高斯先验错配（LeWM 在 Two-Room 翻车）；非高斯流形上持续产生梯度	不坍缩但允许病态分布：多峰、偏斜、重尾都合法
表征几何	继承输入空间几何	各向同性——隐空间 L2 距离可直接当代价做规划（LeWM 的 terminal cost 就靠这个）	弱白化，几何不规则，难直接规划

\(N\) 表示每个样本的时间步，\(M\) 表示随机投影的方向，\(d\) 表示嵌入维度。

---

• SIGReG

核心逻辑就是：我们是让某一个时间 t 的 B 个嵌入，在某个方向上的投影，符合正态分布

根据 Cramér-Wold 定理，要让高维点云真正成为一个各向同性（各个方向都均匀对称）的超球体，必须满足：从任意角度投影过去，投影都必须是正态分布。

所以可以采用 \(M\) 个随机方向的投影，几乎可以涵盖全部方向。

投影后得到 \(B\) 个标量，我们期望这个标量符合 \({\mathcal N}(0, 1)\) 的正态分布。

然后对每个一维投影计算 Epps–Pulley 统计量*。核心是特征函数空间里的加权 L2 距离——经验特征函数（ECF）与标准高斯特征函数之差：

\[T\big(h^{(m)}\big) = \int_{-\infty}^{\infty} w(t)\left\lvert \phi_N\big(t; h^{(m)}\big) - \phi_0(t) \right\rvert^2 dt \]

其中经验特征函数是嵌入的可微函数：

\[\phi_N(t; h) = \frac{1}{N}\sum_{n=1}^{N} e^{i t h_n} \]

\(\phi_0(t) = e^{-t^2/2}\) 是标准正态的特征函数，\(w(t)\) 是一个衰减权重函数（压制大 \(|t|\) 处的数值噪声，论文取指数衰减形式）。

补充：
这个 \(\sum e^{itX}\) 的形式，或者类似物，来自于概率分布的特征函数（CF）
CF 在频域去观察一个随机变量，通过傅里叶变换 \(\int e^{it f(x)} dx\) 的类似物变换，得到：

\[\phi_X(t) = {\mathbb E}[e^{itX}] \]

这里使用这个形式，就是在 CF 的维度上，衡量分布是否均匀。

经验分布 ECF $$\phi_N(t; h) = \frac 1 N \sum_{i = 1}^N e^{ith_n}$$ 就是经验分布的 CF，并且期望收敛到整个分布的 CF。所以可以这样去做。

但是这里是一个积分，是难以计算的，论文中采用了一种数值计算的技巧：梯形法。

\(\mathcal{T}\) 个节点均匀分布在 \([0.2, 4]\)。离散成

\[\hat T\big(h^{(m)}\big) \approx \sum_{j=1}^{\mathcal{T}} c_j\, w(t_j)\Big\lvert \underbrace{\tfrac{1}{B}\sum_{b=1}^{B} e^{i t_j h_b^{(m)}}}_{\text{ECF 在该节点的复数值}} -\, e^{-t_j^2/2}\Big\rvert^2 \]

其中 \(t_j\) 是均匀网格点，\(c_j\) 是梯形权重（端点半权）。每个节点的计算就是一次复指数：\(e^{i t_j h} = \cos(t_j h) + i\sin(t_j h)\)，对 \(B\) 个投影值取均值——全程元素级运算。

维度	方法	性质
频率 \(t\)	确定性求积（\(\mathcal{T}\) 个固定节点，梯形）	无随机性，loss 可复现
方向 \(m\)	蒙特卡洛（每步采样 \(M=1024\) 个单位方向）	有采样噪声，但消融显示 \(M\) 不敏感

总复杂度 \(O(M \cdot \mathcal{T} \cdot B)\)，三个维度都能向量化摊平，相对训练主路径的开销可以忽略。

还有一种方法，直接求闭式解。
如果权重取高斯型 \(w(t) = e^{-\lambda t^2/2}\)，积分有闭式：利用 \(\int e^{-\lambda t^2/2} e^{itx} dt \propto e^{-x^2/2\lambda}\)，整个统计量变成两两投影差的高斯核之和（\(\frac{1}{B^2}\sum_{j,k} e^{-(h_j-h_k)^2/2\lambda}\) 这类形式）——此时 EP 统计量显式地就是一个核均值嵌入距离的有限样本估计。代价是朴素实现 \(O(B^2)\)，工程上要做低秩或截断近似。

最后在对样本时间样本时间 N 做一次平均即可。

---

• CEM solver

训练出来还是要解决问题的，有了这个又该如何？

在本文的设计下，任务形式本身就是 goal-reaching，让 T 到达最终位置即可。
而巧合的是，有了这个 SIGReg 的约束，根据实验，发现可以从 \(z\) 中通过某种线性映射读出各种物理量，也就是说，隐空间中的距离和现实中物理状态的距离是相关的（不一定是完全一致或者说同构的）。由于约束到了 \({\mathcal N}(0, 1)\)

• 各方向方差归一（isotropic），没有哪个维度垄断距离；
• 没有坍缩，距离有动态范围
• 各向同性高斯下，欧氏几何就是该分布的自然几何

但是这个也只是说，近似正确，不一定是最优。不过确实可以拿到一个解。

这个方法和模拟退火非常像：当前解附近随机扰动 + 探索/开采权衡。

CEM solver 维护的是一个随机分布，每次在这个随机分布中采样 \(N\) 条路径，选择 \(K\) 条最优的路径，适应他们的分布情况，得到新的随机分布方式。

\[\mu_t \leftarrow \frac{1}{K}\sum_{i \in \mathcal{E}} a_{1:H}^{(i)}, \qquad \Sigma_t \leftarrow \mathrm{Var}_{i \in \mathcal{E}}\big(a_{1:H}^{(i)}\big) = \frac{1}{K}\sum_{i \in \mathcal{E}}\big(a^{(i)} - \mu_t\big)\big(a^{(i)} - \mu_t\big)^{\top} \]