机器学习笔记(25.1): RC-aux

可能更好的阅读体验 Jeefy's Blog: 机器学习笔记(25.1): RC-aux

在 LeWorldModel 的基础上，解决了这样一个问题：

Predictive \(\neq\) Plannable（可预测不等于可规划）

传统的免重建（Reconstruction-free）JEPA 类世界模型在训练和测试之间存在严重的时空错配（Spatiotemporal Mismatch）：

1. 时间错配（Temporal Mismatch）：训练时往往只监督单步或短期的局部预测，但测试规划时却要把预测出来的隐状态不断自回归喂给模型
2. 空间错配（Spatial Mismatch）：规划器在最后一步通常用隐空间的欧氏距离（\(L_2\) distance）来评估离 Goal 有多近。但这个相邻的状态可能根本不可达

---

• 时间错配：那就直接多轮迭代，修正误差。

\[\hat{z}_{t+1:t+K} = F_\theta^{(K)}(z_{t-L+1:t}, a_{t:t+K-1}) \]

然后用真实编码的未来隐状态 \(z_{t+1:t+K}\) 作为 Target（注意这些 Target 不参与 rollout 输入），计算加权的多步预测损失：

\[\ell_{\text{mh}} = \sum_{k=1}^{K} w_k \Vert{}\hat{z}_{t+k} - z_{t+k}\Vert{}_2^2 \]

个人批注：论文在 Appendix B 里给出了严格的 Lipschitz 上界证明（Lemma 1 & Lemma 2）。单步误差在开环下会按 \(\sum L_f^i\) 指数级累积，而直接约束 \(\ell_{\text{mh}}\) 能够直接控制规划器评估代价的失真上界 \(\vert{}C_{\text{base}}(\hat{z}) - C_{\text{base}}(z^*)\vert{} \le \frac{L_C}{\sqrt{w_{\min}}}\sqrt{\ell_{\text{mh}}}\)。

---

• 空间错配：那就多一个 “可达” 的旁路头，用于增加可达性的约束。

\[R_\phi(z, z', h) = \sigma(r_\phi(z, z', h)) \in [0, 1] \]

用来预估从源状态 \(z\) 出发，能否在 \(h\) 步之内到达目标状态 \(z'\)。注意这里 \((z, z')\) 是有序对，因为存在障碍物或不可逆动力学时，可达性天然是非对称的。

作者直接用离线轨迹里的步数差 \(\Delta = j - i\) (\(i < j\)) 作为经验代理，随机采样预算 \(h \in \{0, 1, \dots, H_{\max}\}\)：

\[y_{ijh} = \mathbb{1}[h \ge \Delta] \]

这里的样本构造是整篇论文最精妙的地方，一共包含四类样本：

1. Reachable Positives (\(y=1\))：同一条轨迹上，\(h \ge \Delta\)，说明在预算 \(h\) 内确实观测到了从 \(z_i\) 走到 \(z_j\)。
2. Temporal Hard Negatives (\(y=0\))：同一条轨迹上，但 \(h < \Delta\)。也就是说 \(z_j\) 最终是可达的，但在当前给定的预算 \(h\) 下还到不了！
3. Batch Negatives (\(y=0\))：不同轨迹之间随便拼的 \((z_i, z_j^-, h)\)，防止模型无脑输出全可达。
4. Predicted Rollout Pairs：因为测试时 \(R_\phi\) 接收的是动力学模型预测出来的中间状态 \(\hat{z}_{t+k}\)，为了消除分布偏移，把预测状态做 Stop-gradient 得到 \(\tilde{z}_{t+k} = \text{sg}(\hat{z}_{t+k})\)，和未来的真实状态 \(z_{t+\ell}\) (\(\ell > k\)) 配对，按 \(\Delta_{k,\ell} = \ell - k\) 同样打上 \(y = \mathbb{1}[h \ge \Delta_{k,\ell}]\) 的标签。

最终可达性损失就是编码特征对 \(\mathcal{B}_{\text{enc}}\) 和预测特征对 \(\mathcal{B}_{\text{pred}}\) 上的加权 BCE 损失：

\[\ell_{\text{reach}} = \frac{1}{\vert{}\mathcal{B}_{\text{enc}}\vert{}} \sum \omega_y \ell_{\text{bce}}(z, z', h, y) + \rho_{\text{pred}} \frac{1}{\vert{}\mathcal{B}_{\text{pred}}\vert{}} \sum \omega_y \ell_{\text{bce}}(\tilde{z}, z', h, y) \]

其中 \(\mathcal{B}_{\text{enc}}\) 的梯度会回传给编码器 \(e_\theta\) 来塑形隐空间几何，而 \(\mathcal{B}_{\text{pred}}\) 因为截断了梯度，专门用来在预测分布上校准判别头 \(R_\phi\)。

总训练目标保留了底层的正则项 \(\ell_{\text{reg}}\)（在 LeWM 中就是 SIGReg 高斯正则）：

\[\ell_{\text{RC-aux}} = \ell_{\text{mh}} + \alpha \ell_{\text{reg}} + \beta \ell_{\text{reach}} \]

---

• CEM solver++

训练好之后，在测试阶段用 CEM 规划长度为 \(H\) 的动作序列 \(\tau = a_{t:t+H-1}\) 时，基础规划代价通常是终点距离 \(C_{\text{base}}(\tau) = \Vert{}\hat{z}_{t+H} - z_g\Vert{}_2^2\)。

RC-aux 将 \(R_\phi\) 也用到了搜索打分里。对于一条预测轨迹 \(\hat{z}_{t+1:t+H}\)，先看中间状态能否在剩余步数 \(H-k\) 内到达目标 \(z_g\)：

\[R(\tau, z_g) = \max_{1 \le k < H} R_\phi(\hat{z}_{t+k}, z_g, H - k) \]

然后用这个可达性分数对基础代价做一个软折扣（Soft Feasibility Gate）：

\[C_{\text{RC}}(\tau) = C_{\text{base}}(\tau) \cdot \max\big(m, 1 - \lambda_{\text{plan}} R(\tau, z_g)\big) \]

这里 \(\lambda_{\text{plan}} \in [0, 1]\) 是耦合权重，\(m > 0\) 是防下溢的小常数。

• 当 \(\lambda_{\text{plan}} = 0\) 时，完全退化为标准的终点 \(L_2\) 规划器（可以用来单测训练表征带来的提升）。
• 当 \(\lambda_{\text{plan}} > 0\) 时，如果某条轨迹虽然基础 \(L_2\) 距离稍微大一点点，但中间状态的可达性 \(R \approx 1\)，它的总代价就会被打折（最高打 \(1 - \lambda_{\text{plan}}\) 折），从而击败那些“隐空间欧氏距离很近、但撞在墙上根本不可达（\(R \approx 0\)）”的虚假捷径。

---

• 目前的局限在于用离线轨迹里的步数差 \(\Delta\) 当作可达性标签只是一种偏保守的代理（如果数据集里的轨迹绕了远路，\(\Delta\) 就会偏大）。
• 另外，虽然 RC-aux 让 CEM 规划变准了，但它依然需要在测试时采样几百上千条序列做 Rollout 搜索。能不能连测试时的 CEM 搜索都直接干掉，把世界模型变成毫秒级直出的控制器？