Files
worldmodel/JEPA/math/06_planning_equivalence.md
T

8.8 KiB
Raw Blame History

Topic 6:正交不变性与最优规划(定理 4)

前置知识: Topic 3:谱分解与线性可识别性、基础控制理论(可选) 目标: 理解为什么线性可识别性足以保证在学到的潜空间中规划与在真实世界中规划完全等价


🎯 定理 4 的完整陈述

定理 4(最优潜空间规划):h(z) = QzQ ∈ O(n),由定理1保证)。对任意有限时域控制问题,若代价函数关于状态是 O(n)-不变的,则:

V̂*(h(z₀)) = V*(z₀)        (最优值函数相等)
â*_{1:T}(h(z₀)) = a*_{1:T}(z₀)   (最优动作序列相等)

白话翻译: 如果代价函数不区分旋转方向,那么在学到的潜空间 ĥ = Qz 中规划,与在真实潜空间 z 中规划,得到的最优策略完全相同。


🤔 为什么这个结论重要?

世界模型的终极目标

学习世界模型的目的是规划:给定当前状态,找到最优动作序列。

如果学到的表示 h(z) 不能支持正确的规划,那么世界模型就没有实用价值。

定理4说明:线性可识别性(正交等价)已经足够支持最优规划——不需要精确恢复 z,只需要恢复到旋转等价。


📐 关键概念:O(n)-不变代价函数

定义

代价函数 (z, a)O(n)-不变的,如果对所有正交矩阵 Q ∈ O(n)

(Qz, a) = (z, a)   对所有 z, a

直觉: 代价函数不依赖于坐标系的旋转方向,只依赖于状态的"本质"(如距离、范数等)。

常见的 O(n)-不变代价函数

代价函数 形式 不变性
欧氏距离到目标 ‖z - z_goal‖² (若 z_goal 也旋转)
线性二次调节(LQR z^T P z + a^T R a (若 P = cI
范数惩罚 ‖z‖²
目标到达 𝟙[‖z - z_goal‖ < r]
任意旋转不变量 f(‖z‖, ‖a‖, ...)

不满足 O(n)-不变性的代价函数

代价函数 形式 原因
坐标惩罚 z₁²(只惩罚第一维) 旋转后变成 (Qz)₁²
非对称目标 ‖z - [1,0,...,0]‖² 目标方向固定

📐 证明的核心思路

关键引理:代价等价

h(z) = QzQ ∈ O(n)。对任意 O(n)-不变代价函数

(h(z), a) = (Qz, a) = (z, a)

这一步是整个证明的核心! 正交变换不改变 O(n)-不变代价函数的值。

轨迹推前(Trajectory Pushforward

设真实动力学为 p(z'|z, a),学到的潜空间动力学为 p̂(ẑ'|ẑ, a)(其中 ẑ = Qz)。

由于 h(z) = Qz 是线性双射,学到的动力学是真实动力学的推前

p̂(ẑ'|ẑ, a) = p(Q⁻¹ẑ'|Q⁻¹ẑ, a) = p(z'|z, a)

(因为 Q⁻¹ = Q^T 对正交矩阵成立)

总代价等价

对任意动作序列 a_{1:T},从初始状态 z₀ 出发的总期望代价:

J(a_{1:T}; ẑ₀) = E[Σ_t (ẑ_t, a_t) + _T(ẑ_T) | ẑ₀ = Qz₀]
               = E[Σ_t (Qz_t, a_t) + _T(Qz_T) | z₀]
               = E[Σ_t (z_t, a_t) + _T(z_T) | z₀]   O(n)-不变性)
               = J(a_{1:T}; z₀)

结论: 对任意动作序列,两个空间中的总代价完全相同!

最优性等价

由于对所有 a_{1:T} 代价相等,最小化代价的动作序列也相同:

a*_{1:T}(ẑ₀) = argmin_a J(a; ẑ₀) = argmin_a J(a; z₀) = a*_{1:T}(z₀)

最优值函数也相等:

V̂*(ẑ₀) = min_a J(a; ẑ₀) = min_a J(a; z₀) = V*(z₀)

🎨 几何直觉

真实潜空间 z:              学到的潜空间 ẑ = Qz:

    z₂                          ẑ₂
    ↑                            ↑
    │  ●goal                     │      ●goal'
    │                            │
    │●start                      │  ●start'
    └──────→ z₁                  └──────→ ẑ₁

最优路径(蓝色):              最优路径(蓝色):
  start → goal                   start' → goal'
  (直线,欧氏距离最短)           (直线,欧氏距离最短)

两条路径在旋转意义下完全相同!

🔧 Lean 4 验证(Planning.lean

-- 核心:对任意动作序列,两个空间的总代价相等
theorem planning_equivalence
    (cp : ControlProblem n Action) (Q : Latent n  Latent n)
    (hinv : IsOrthogonalInvariant cp Q)
    (a : Plan Action T) (z : Latent n) :
    totalCost cp E_hat a (Q z) = totalCost cp E a z

-- 推论:最优动作序列相同
theorem minimizer_equivalence ... :
    ( a', cost_hat a (Q z)  cost_hat a' (Q z)) 
    ( a', cost a z  cost a' z)

-- 推论:最优值函数相等
theorem value_equivalence ... :
    totalCost cp E a z = V 
    totalCost cp E_hat a (Q z) = V

🔬 实验验证:DMC Reacher

实验设置

  • 环境DeepMind Control Suite 的 Reacher 任务
  • 输入:像素图像(64×64 RGB
  • 潜变量2D 关节角度 z = (θ₁, θ₂)
  • 编码器CNN(见 models.py
  • 规划方式:在潜空间中线性插值,用最近邻检索解码

两种训练数据

数据类型 生成方式 分布 可识别性
OU 采样 z' = ρz + √(1-ρ²)η 各向同性高斯 高(满足定理1
RL 轨迹 训练好的策略采样 非高斯、各向异性 低(违反假设)

实验结果

规划代价(路径长度,越低越好,理想值=1):

Oracle(关节空间直线):  ████░░░░░░  ~1.0(基准)
OU 编码器:              ████░░░░░░  ~1.0(与 oracle 无统计显著差异)
轨迹编码器:             ██████░░░░  ~1.5(显著偏高)

结论: OU 编码器(满足定理1条件)的规划质量与 oracle 相当;轨迹编码器(违反假设)的规划质量显著下降。

可视化

[顶行] Oracle:
  ●──────────────────●  (关节空间直线,平滑弧线)

[中行] OU 编码器(可识别):
  ●──────────────────●  (紧密跟随 oracle)

[底行] 轨迹编码器(不可识别):
  ●────╮╰──────────●  (偏离,因为潜空间扭曲)

🔗 与世界模型的联系

什么是"可证明地学到世界模型"?

论文的标题问题:"When Does LeJEPA Learn a World Model?"

答案(由定理4给出):

LeJEPA 学到世界模型,当且仅当它实现了线性可识别性。

因为:

  • 线性可识别性 → h(z) = Qz(正交等价)
  • 正交等价 → O(n)-不变代价函数下的规划等价(定理4)
  • 规划等价 → 可以在学到的潜空间中做最优规划
  • 最优规划 → 学到的表示是"可用的世界模型"

⚠️ 定理4的局限性

1. 只覆盖 O(n)-不变代价函数

如果代价函数依赖于特定坐标方向(如"向北走"),则定理4不适用。

实践中: 大多数物理任务的代价函数(距离、能量、时间)都是旋转不变的。

2. 只处理编码器侧

定理4假设动力学 p̂(ẑ'|ẑ, a) 是真实动力学的推前。但在实践中,还需要学习一个转移模型predictor)。

未来工作: 动作条件转移 p̂(ẑ'|ẑ, a) 的可识别性(与因果表示学习相关)。

3. 有限时域

定理4是有限时域(T 步)的结论。无限时域(折扣 MDP)的情况需要额外分析。


小结

  1. 定理4 证明线性可识别性足以保证最优规划等价
  2. 关键条件:代价函数是 O(n)-不变的(旋转不变)
  3. 证明核心:正交变换不改变 O(n)-不变代价函数的值
  4. 实验验证:OU 编码器的规划质量与 oracle 相当,轨迹编码器显著下降
  5. 世界模型含义:线性可识别性 = 可证明地学到世界模型

🏁 四个定理的完整图景

定理1(正向):高斯世界 + LeJEPA → 线性可识别性 h(z) = Qz
    ↕
定理2(逆向):高斯是唯一使线性可识别性成立的分布
    ↓
定理3(近似):条件近似满足时,误差 ≤ D + (ε+D)²
    ↓
定理4(应用):线性可识别性 → 最优潜空间规划

核心信息: LeJEPA 在高斯世界中可证明地学到世界模型,且这个保证对近似条件优雅降级,并直接支持最优规划。


➡️ 返回总览

README:数学 Topic 导航论文完整笔记