Anisotropic Representations Improve Planning in JEPA World Models
Mingu Kang, Yoori Oh, Sookyung Kim, Joonseok Lee
cs.RO
2026-09-29
AnisoWM 用 ΛReg 把各向同性高斯目标换成可学对角协方差，预测器和欧氏规划器保持不变；四个视觉控制环境规划成功率全部超过 LeWM，TwoRoom 93% 对 87%、Cube 79% 对 74%。
JEPA 世界模型在表征空间里做动作条件预测。视觉目标规划时，规划器把候选动作序列滚出去，用预测终点到目标表征的欧氏距离打分。编码器既提供预测特征，也决定规划代价怎么给不同终点误差加权。
LeWorldModel（LeWM）联合训练编码器和动作条件预测器，再用 LeJEPA 的 SIGReg（sketched 各向同性高斯正则）防止表征坍缩。SIGReg 把特征推向各向同性高斯，动机是 probing 质量，不是欧氏距离适不适合当规划代价。预测准、表征不坍缩，规划排序仍可能和任务代价拧着。
线性高斯分析给出机制。各向同性特征协方差会在物理状态空间里按状态协方差的逆矩阵加权：训练方差小的方向，规划时更贵。目标在可达集外、必须在方向间权衡时，潜空间规划器和任务代价会挑不同的可达终点，留下严格为正的规划 regret。二维 MLP toy 上，过程噪声缩小 30 倍，held-out 条件均值预测误差明显下降，归一化物理规划 regret 却停在约 0.16（十个种子），贴着 0.1607 的 Σ⁻¹ 参照。预测变好，排错没有消失。
AnisoWM 只改正则目标。ΛReg 把固定各向同性高斯换成可学对角协方差，两条硬约束：
κ=1 退回各向同性，就是 LeWM。对角钉在表征坐标上，编码器仍可旋转这些轴相对物理几何的朝向。对特征做 Λ^{-1/2} 标准化后再跑原来的 SIGReg，因为 zN(0,Λ) 时标准化后回到 N(0,I)。标准化只进正则项，预测和规划继续用原始 z。
Λ 用截断 logit 参数化，迹和条件数始终可行。预测损失、预测器结构、CEM（交叉熵采样）欧氏规划器全部不动。Λ 训练完丢掉，推理期没有新度量。
可学目标把极限度量从 qΣ⁻¹ 扩成 qΣ^{-1/2} L Σ^{-1/2}，L 落在约束集合里。预测项把较大目标方差分给噪声较小的方向，对冲各向同性正则的逆协方差加权。二维闭式特例里，κ 等于状态协方差条件数时，极限状态度量回到各向同性，规划 regret 趋于 0；κ 再大，regret 会重新抬头。
四个视觉目标规划环境沿用 LeWM 的数据、架构和协议：TwoRoom、Reacher、PushT、Cube。表征 192 维，λ=0.09，主实验共用 κ=2，三个种子取均值。LeWM 基线同一套代码和设置，对应 κ=1。PLDM、DINO-WM、GCBC、GCIQL、GCIVL、Random 的成功率引自原 LeWM 论文。
| 环境 | AnisoWM | LeWM |
| TwoRoom | 93% | 87% |
| Reacher | 89% | 86% |
| PushT | 97% | 96% |
| Cube | 79% | 74% |
四个环境全部超过 LeWM。TwoRoom 和 Cube 上部分参考基线更高；AnisoWM 保住的是 LeWM 那套潜空间规划。
同一批已记录动作序列上，比较潜空间代价排序和任务结果排序。Jpred 是 CEM 的规划代价，Jenc 只看真实到达观测的编码距离：
| 环境 | Jpred LeWM | Jpred AnisoWM | Jenc LeWM | Jenc AnisoWM |
| TwoRoom | 0.575 | 0.646 | 0.564 | 0.661 |
| Reacher | 0.676 | 0.723 | 0.790 | 0.893 |
| PushT | 0.594 | 0.621 | 0.647 | 0.641 |
| Cube | 0.537 | 0.553 | 0.555 | 0.576 |
Jpred 四环境都升。Jenc 在 PushT 低了 0.005。Cube 上潜空间代价与任务代价的 Spearman ρ 从 0.15 升到 0.91，PushT 从 0.65 升到 0.92。
共用 κ=2，学到的谱仍不同：均值以上坐标数从 PushT 的 95 到 Reacher 的 124。κ 扫描非单调，看起来最好的各向异性因环境而异；κ>1 每个点只有一次训练，论文当敏感性诊断，不当调参对比。
改的是训练期正则目标。规划器和预测架构保持原样。已经在用 LeWM 或同类 JEPA 潜空间规划的人，这是一个推理成本为零的旋钮：多一组对角方差，规划时丢掉。
防坍缩用的几何会直接变成规划几何。各向同性高斯对 probing 可能合理，对欧氏距离规划不必合理。SCALE、TRM、DecisionMetric Alignment 改距离或加可达性度量；AnisoWM 把同一件事塞进高斯正则，欧氏规划器原样留下。
增益是渐进的。PushT 只多 1 个点，Reacher 多 3 个点，主实验没有按环境调 κ。该跟进的是「正则目标即规划度量」这条机制。这点提升是渐进改进。
正文没有独立 Limitations 小节。作者写进结果和分析的保留包括：各向异性过大会抬高规划 regret；κ 增大不单调提升成功率；κ 扫描每个设置只有单次训练；TwoRoom 和 Cube 上部分参考基线更高。
理论主体是线性高斯加仿射编码器。Lemma 1 依赖仿射假设，非线性编码器对局部度量约束更弱。视觉控制实验是非线性的，成功率和排序诊断对得上故事，但不等于定理直接搬过来。主对比只有三颗种子，PushT 的 Jenc 还略差。没有和 SCALE、TRM 做同设定对照，也没有真机。Λ 强制对角，各向异性方向交给编码器去转，这个偏置稳不稳，正文没有单独拆开验证。