Awesome World Models / OpenReview

Improving World Models using Supervision with Co-Evolving Linear Probes

作者：作者信息待提取 · 单位：Department of Electrical Engineering and Computer Science, MIT, Department of Brain and Cognitive Sciences & McGovern Institute for Brain Research, MIT · 来源：General Approaches to World Models · 发布日期：2025 · 分类：1. Foundation World Models / Classics / Other Awesome Models

基础世界模型

精读笔记

Problem Setting

论文标题：Improving World Models using Supervision with Co-Evolving Linear Probes（Awesome World Models / 2025）。

这篇论文真正解决的不是“如何在 Flappy Bird 上做预测”，而是一个更窄但重要的问题：end-to-end predictive RNN 在只优化下一步观测预测时，是否会自然形成可用的 internal world state；如果不会，能否用少量真实状态监督把 hidden state 推向更接近环境状态的结构。

困难点在于 next-observation loss 对 latent state 的约束是不充分的。一个模型可以通过短期观测连续性、动作-观测局部关联、甚至训练分布上的 trajectory regularity 获得低预测误差，但 hidden state 未必对应环境的稳定因果变量。这在 open-loop rollout 时尤其明显：closed-loop teacher forcing 下看起来预测好，真正让模型用自己的预测递推时容易 drift。

以前方法卡在“预测能力”和“状态结构”之间的松耦合：Dreamer/World Models 系列证明 latent dynamics 可以支持控制，但并不保证 hidden state 以线性、可组合、可解释的方式编码真实世界变量。本文试图把这件事从事后 probe 变成训练约束。关键矛盾是：如果加入真实状态监督，会不会破坏 predictive objective 的自由度？实验显示没有，反而改善预测，这正是论文最值得关注的点。

Motivation

已有路线默认 observation prediction 足以逼出 world model，但这个假设在部分可观测环境里并不稳。预测目标只要求输出正确，不要求内部变量以可维护、可组合的方式存在；因此模型可能学到的是短期 filter，而不是长期 state estimator。

作者的核心观察是：linear probing 通常被用来判断模型是否已经学到了某种 world representation，但如果 probe 能揭示 representation 的好坏，也许它可以反过来作为训练信号，直接塑造 representation。也就是说，缺口不是缺一个更强预测头，而是缺一个对 recurrent memory 的结构性约束。

这个方向的动机很 practical：如果用少量 privileged variables 或额外传感器信号，就能让小模型获得接近大模型的预测效果，那么在 compute-limited 或嵌入式 setting 下很有价值。但从研究角度看，更重要的是它把“world model 是否有世界状态”这个问题转化为“是否能通过辅助监督让 hidden state 对齐到真实状态坐标系”。

Core Idea

核心思想很简单：在训练 world model 预测下一 latent observation 的同时，从 RNN hidden state 上接一个线性 probe 去回归部分真实世界变量，并把 probe 的 MSE 加到训练损失里。probe 不是冻结后分析，而是和主模型共同演化；因此它给 hidden state 施加了一个持续压力：真实状态变量必须以线性可读的形式存在于 memory 中。

这改变的是建模方式，而不是输出形式。普通 predictive RNN 可以把历史信息编码成任意高维非线性缠绕表示，只要预测头能解开即可；加入 linear probe 后，至少一部分状态信息必须被组织成低复杂度、稳定、线性可访问的子空间。这相当于引入了一个 representation alignment bias：把 latent dynamics 的内部坐标系向环境真实状态坐标系拉近。

和 prior 的本质区别在于：它不是用 probe 解释模型，也不是单纯多任务预测额外标签，而是用 probe 约束 recurrent state 的几何结构。真正新增的信息来自 privileged world features；真正新增的 bias 是“好的 world model 应该维护可线性读出的状态变量”。

Method

方法层面可以压缩成三个机制。

第一，主干仍然是 latent observation dynamics prediction。autoencoder 把 LIDAR 压到低维 latent，LSTM 根据当前 latent 和 action 预测下一 latent distribution。这个主任务保证模型优化目标仍然是环境动态预测，而不是退化成监督状态估计器。

第二，linear probe 作用在 LSTM hidden state 上。这一点比接在输出上重要得多：它约束的是 memory，而不是 readout。换句话说，模型必须在递归状态里保存真实变量，而不是只在最后一步为辅助任务临时计算一个答案。这直接针对 world model 的核心瓶颈：长期状态维护。

第三，probe 只监督部分 world features，例如位置、速度、角度。论文真正想证明的是：即便只监督一小部分变量，也可能重塑整个 hidden state，使未监督但相关的变量也更容易解码。这是本文比普通 auxiliary loss 更有意思的地方；如果成立，说明 probe loss 在诱导更全局的状态空间整理，而非仅仅学习三个额外回归头。

需要注意，λ=64 这类选择和 MDN-LSTM/autoencoder 细节不是核心贡献，更多是实现载体。真正机制是“用线性可解码性作为训练时 inductive bias”。

Key Insight / Why It Works

最可能有效的原因是 representation alignment，而不是简单 scaling。next-step prediction 的解空间太大，很多 hidden representation 都能达到类似 teacher-forced loss；linear probe 等价于在这些解中选择一个更接近真实状态坐标、更容易递推的解。这个选择降低了优化难度，也降低了 rollout 时误差在 hidden dynamics 中被非线性放大的概率。

这可以理解为一种弱形式的 state estimation regularization。RNN 在部分可观测环境中必须从历史 LIDAR 和 action 中推断 bird 的速度、位置、姿态以及管道相对位置。probe loss 告诉模型哪些变量值得稳定维护，从而避免 hidden state 把信息压成仅对一步预测有用的局部 code。未监督变量 decodability 的提升，可能来自环境变量之间的动力学耦合：如果位置、速度、角度被线性化，相关的管道相对位置、时序阶段等变量也更容易落入线性子空间。

训练稳定性的提升也合理。辅助 MSE 给 recurrent state 提供更直接、更 dense 的梯度路径，尤其在早期训练时帮助模型建立基础坐标系。这更像 curriculum / representation shaping，而不是模型容量增加。作者讨论 annealing 是有道理的：probe 的主要价值可能在早期把 hidden state 拉到好 basin，后期继续加监督未必必要，甚至可能限制主任务最优表示。

但增益归因仍不完全清楚。它可能混合了三件事：一是 latent structure 更好；二是额外标签提供了更强监督和更密集梯度；三是 probe loss 改善了 RNN 优化稳定性，减少 divergence。论文没有充分消融这些因素，例如没有比较非线性 probe、随机辅助标签、错误状态变量、只加辅助 loss 但 stop-gradient、或 annealed probe。因此不能断言所有提升都来自“更真实的 world model”。

另外，distribution drift 的结果并非全局成立，在穿越管道等高变化阶段不明显。这说明该方法更像改善状态估计的局部稳定性，而不是解决长期多模态未来预测。它没有引入 test-time planning、retrieval 或额外 rollout compute；收益主要来自训练时的 inductive bias 和 hidden supervision。

Relation To Prior Work

它最接近三条线：World Models / Dreamer 式 latent recurrent dynamics，deep supervision，以及 mechanistic interpretability 中的 linear probing。本文的新意不在任何单个组件，而在把 interpretability probe 变成训练目标，用来塑造 world model 的 hidden state。

和 Ha & Schmidhuber / Dreamer 的差异是：后者主要依赖 reconstruction/prediction/reward 等任务信号，让 latent state 自组织；本文显式引入真实状态变量作为 privileged supervision。因此它不是纯粹自监督 world model，而是 semi-supervised / privileged-supervised world model training。

和传统 deep supervision 的差异是监督位置和目的不同。视觉 deep supervision 通常为了解决梯度传播、加速收敛或提升中间层语义；本文关心的是 recurrent hidden state 是否形成动态系统状态表示。也就是说，它不是给中间层加分类头，而是给 memory 加状态坐标约束。

和 linear probing 文献的差异更直接：probe 通常是冻结模型后的诊断工具，用来回答“模型里有没有某变量”；这里 probe 是 co-evolving 的，回答的是“如果要求模型里有某变量，world model 会不会更好”。这是一种从 interpretability 到 training objective 的迁移。

看似新的部分其实是已有思想的重组：auxiliary state prediction、privileged information、deep supervision 都已有。实质创新在于把这些放到 world model hidden-state geometry 上，并用 decodability、drift、scaling shift 一起论证它确实改变了内部表示。

Dataset / Evaluation

评估环境是带 LIDAR 的 Flappy Bird，训练 rollouts 来自 expert DQN 加少量随机扰动，测试包含 DQN rollouts 和 random policy rollouts。这个设置足以检验一个受控问题：在低维规则动力学下，真实状态 probe 是否改善 recurrent latent dynamics。但它远不能支撑广义 foundation world model 的 claim。

任务覆盖范围很窄：单环境、单游戏、低维真实状态、有限动作、短期动力学，且环境变量可直接访问。没有跨场景、多任务、真实世界或机器人真机验证。random policy 测试提供了一点 distribution shift，但仍在同一环境机制和观测生成分布内，不是强泛化。

实验比较了 predictive loss、world-feature decodability、open-loop drift、scaling with model size/time、training stability。整体上支持“linear probe supervision 有助于这个 setting 的 world model 表征和优化”。其中最有说服力的是 predictive loss 没被 auxiliary loss 损害，反而下降；以及未监督变量也更可解码。但 decodability 本身不是 world-model quality 的充分证据，尤其在变量相关性强的 toy environment 中。

open-loop drift 结果有价值，但并不彻底：只在某些时段明显，穿管道阶段不成立。说明 evaluation 支持的是局部稳定性改善，而不是稳健长期 rollout 或复杂规划能力。若要证明真正 world model 改善，还需要 downstream control、counterfactual prediction、OOD dynamics、partial state corruption、多环境迁移等实验。

Limitation

最根本的前提是存在可用的真实 world features。现实中这通常是 privileged information：仿真器给得出，真实世界未必给得出；机器人可通过额外传感器近似获得，但那已经改变了数据采集设定。因此方法不是免费提升，而是用额外监督/传感器换取模型容量和稳定性。

第二，增益来源不清。论文把效果解释为更好的 world model，但同样可能来自更密集梯度、更强 regularization、更好的 early optimization basin，甚至是 auxiliary target 与 next observation 高相关带来的 supervised shortcut。没有充分消融就很难区分“学到了因果状态”与“被标签拉近了训练分布上的相关变量”。

第三，scalability 上限未知。Flappy Bird 的状态空间低维且动力学简单，线性 probe 是合理约束；在高维、多对象、接触丰富、隐藏变量复杂的环境中，真实状态未必可线性组织，或者需要监督的变量数量本身爆炸。linear decodability 可能成为过强约束，限制模型学习对预测更有用但非人类命名的 latent factors。

第四，所谓泛化证据较弱。测试仍在同一环境，同一观测形式，同一动力学规则下；random policy 不是强 OOD。未监督变量 decodability 提升也可能只是变量相关性传播，不一定代表模型获得了可组合的世界结构。

第五，该方法可能把问题从“如何自监督学习 world model”转移为“如何选择和获取好的辅助状态变量”。如果监督变量选错，模型可能被拉向非因果或任务无关表示。文中未充分说明 probe feature selection 的原则，也没有分析监督变量质量下降时的鲁棒性。

Takeaway

  • 1. 这篇论文最值得记住的不是 linear probe 本身，而是一个训练范式：把可解释性诊断工具反向用作 representation-shaping objective，直接约束 world model memory。
  • 2. 对小模型和不稳定 recurrent dynamics 来说，少量 privileged state supervision 可能比单纯加模型规模更划算；论文中的“相当于模型变大一倍”应理解为 scaling curve 下移，而不是能力范式跃迁。
  • 3. 未来更有价值的方向是 probe loss 的选择和调度：哪些变量值得监督、是否只在早期 anneal、能否用自发现 latent factors 或额外传感器替代 simulator state。
  • 这里比继续在 Flappy Bird 上调 λ 更重要。

一句话总结

这篇论文把 linear probing 从事后解释工具改造成训练时的 hidden-state 结构约束，证明在一个低维受控 world-model setting 中，少量真实状态监督可以改善预测、稳定性和表示可读性；它属于 privileged supervision / deep supervision 向 world model representation alignment 的一次小而清晰的演化。