Awesome World Models / OpenReview

Recurrent world model with tokenized latent states

作者：作者信息待提取 · 单位：Technical University of Munich, Munich Center for Machine Learning · 来源：General Approaches to World Models · 发布日期：2025 · 分类：1. Foundation World Models / Classics / Other Awesome Models

基础世界模型潜空间世界模型

精读笔记

Problem Setting

《Recurrent world model with tokenized latent states》（Awesome World Models / 2025）关注的是复杂视觉操控中的 latent world model 表达瓶颈，而不是重新定义 POMDP 或提出新的 IL benchmark。

真正的问题是：RSSM 这类 recurrent state-space model 在低维控制、Atari、简单视觉任务中有效，但当观测包含高分辨率视觉、多物体关系、语言条件和长阶段操作时，单个 latent vector 必须同时承担 scene representation、belief update、dynamics prediction 和 policy-relevant abstraction，容量和结构都不够。它不是单纯“模型不够大”，而是信息组织方式不合适：观测被过早压缩成一个全局向量，状态更新缺少显式的对象/区域/语义槽位，历史也只能通过 recurrent hidden state 隐式保留。

这类任务的关键矛盾是：world model 需要足够强的表示能力来保留复杂视觉细节和任务相关关系，但又不能退化成一个只做 observation prediction 的大 Transformer；它仍然需要一个可递推、可用于控制的 compact latent belief。TokenWM 的目标就是在“RSSM 的 sequential belief update”与“token model 的高容量结构化表示”之间找一个中间点。

Motivation

作者对已有路线的判断比较明确：

1. RSSM-style SSM 的优点是有 prior/posterior、latent dynamics、recurrent update，因此适合做 belief learning 和 imagination；但 latent state 是单向量，面对复杂环境时表达容量和信息分解能力不足。

2. Transformer / S4 / Mamba 替换 recurrent backbone 可以增强 temporal modeling，但如果 latent 仍然是单个 vector，本质瓶颈仍在。动态模型更强，不等于 state representation 更可扩展。

3. 另一类 token-based observation prediction world model 直接在 observation tokens 或 feature tokens 上预测未来，容量强，但容易变成 video/feature prediction；它们不一定学习一个面向决策的 recurrent latent state。

因此关键缺口是：缺少一种既保留 SSM-style latent belief update，又允许 latent state 以 token set 形式承载复杂信息的 world model。作者的核心观察是，复杂机器人任务中“状态”天然不是一个均质向量，而更像一个可被局部写入、局部读取、跨时间复用的 token workspace。

Core Idea

TokenWM 的核心不是“用了 attention”，而是改变了 latent state 的形态：从一个全局向量变成一组 latent tokens。这个变化带来的不是简单容量增加，而是新的信息组织 bias：不同 latent token 可以分别吸收不同视觉区域、物体关系、任务阶段或动作相关因素；posterior update 可以通过 cross-attention 从视觉 tokens 中选择性读取，而不是把整张图压成一个固定 embedding。

这使得 recurrent world model 的信息流发生变化。RSSM 中，观测编码器先把图像压成向量，然后 posterior 把这个向量写入 state；TokenWM 中，观测仍以 patch/token 形式进入，latent tokens 作为 query 主动选择需要的信息。这更接近一个 recurrent latent workspace：每一步用 prior state 预测当前 belief，再用 observation tokens 校正 belief，之后通过 memory reader 查询短期历史。

本质差异在于：TokenWM 不只是增强 temporal backbone，而是把 state representation 的容量和结构显式 token 化。它引入的 inductive bias 是“状态由多个可交互槽位组成”，这比单向量 RSSM 更适合高维视觉和多物体操控；同时又比纯 Transformer rollout 更保留 SSM 的递推结构和 prior/posterior 正则。

Method

方法层面可以压缩成几个必要机制：

1. Tokenized latent state：解决 RSSM 单向量 state 的硬瓶颈。latent state 变成 Nh × D 的 token set 后，模型不需要把所有场景、任务阶段和动作条件挤进一个向量。核心变化是 state capacity 从“宽向量”转为“多槽位 workspace”。

2. Observation-to-state cross-attention：解决高维视觉信息如何写入 latent。posterior state 用 latent tokens 作为 query，视觉 encoder 输出的 patch tokens 作为 key/value。这让 belief update 成为选择性读取，而不是全局压缩。它也是 TokenWM 相比普通 ViT+RNN 更关键的地方。

3. Action-conditioned prior：解决 recurrent dynamics 如何在 token space 中推进。动作不是简单拼接到 state，而是通过 FiLM 类条件调制影响 latent token transition。作者认为这也减少了 policy cheating，即 state 直接复制上一动作供 policy 解码的问题。这个解释合理但文中未充分证明。

4. Memory bank / memory reader：解决 recurrent attention block 的短记忆问题以及机器人任务中的执行延迟。它把最近若干 posterior states 存成队列，当前 state 查询历史。机制上更接近短期 retrieval，而不是学到无限长记忆。

5. Deterministic information bottleneck loss：作者放弃 probabilistic latent sampling，用 posterior-prior L2 替代 KL。这个选择很关键，因为它意味着 TokenWM 当前不是严格意义上的 probabilistic SSM，而是 deterministic recurrent representation model。这样更稳定，但也削弱了不确定性建模和 imagination 的理论基础。

Key Insight / Why It Works

最可能真正有效的是 tokenized latent + pretrained visual token alignment，而不是完整 world-modeling 框架本身。

第一，LIBERO 这类语言条件操控任务对视觉语义、物体定位和阶段识别高度敏感。RSSM baseline 使用低分辨率输入和 CNN-style vector bottleneck，很容易丢失小物体、相对位置和语言指代相关信息。TokenWM 使用 DINOv2 token features，并让 latent tokens cross-attend 到这些 features，本质上是把强视觉 foundation representation 接到了 recurrent policy state 上。这个增益很可能大。

第二，token latent 给 policy 提供了更丰富的可读状态。对 offline imitation learning 来说，这未必需要模型学会长程 rollout dynamics；只要 latent state 能整合当前图像、短期历史和语言条件，就足以提升行为克隆表现。因此这里的“world model”能力和“policy representation”能力有混叠。论文没有展示 imagination-based planning 或多步 rollout prediction 的质量，不能确认性能来自 dynamics reasoning。

第三，memory bank 的作用更像 retrieval / phase disambiguation。LIBERO 任务常有阶段切换：靠近、抓取、移动、释放。相似观测下动作取决于历史阶段，短期 state retrieval 可以显著缓解 partial observability。这个机制有效，但并不等同于学到了长期因果 dynamics。

第四，posterior-prior L2 bottleneck 实际上把模型训练成“可预测的 latent feature flow”。它鼓励 prior 追 posterior，但没有 probabilistic KL，也没有显式不确定性。因此方法成立更多依赖 representation smoothness 和 demonstration coverage，而不是严格生成式建模。

我会把贡献归因为：更好的 latent structure + 更强视觉 representation + 短期 memory reuse。FiLM、Perceiver decoder、post-norm 等是稳定训练的重要工程件，但不是概念核心。所谓 scalable world model 的 claim 目前证据不足；更保守的说法是：TokenWM 是一个更适合复杂视觉 IL 的 recurrent token-state representation model。

Relation To Prior Work

和 RSSM / Dreamer 系列相比，TokenWM 保留了 prior/posterior/decoder 的 SSM 训练模板，但替换了 state 的基本类型。RSSM 的 state 是 deterministic hidden + stochastic latent 的向量组合；TokenWM 的 state 是 deterministic token set。真正不同点不是 transition 用 attention，而是 belief state 不再是单一压缩向量。

和 TransDreamer、STORM、KalMamba 等 backbone 替换路线相比，TokenWM 更关注 latent interface。那些工作主要问“用什么序列模型推进 state”；TokenWM 问“state 本身是否应该 token 化”。这是更接近表征结构层面的改变。

和 token-based video/world prediction 模型相比，TokenWM 没有完全放弃 recurrent latent belief。它不是直接预测未来 observation tokens，而是维护一个 recurrent latent workspace，再从中解码 feature/observation。这让它更贴近 control-oriented world model，而不是纯生成式预测器。

和 Perceiver IO、RT-1、VLA 类模型相比，很多模块思想并不新：learnable query、cross-attention readout、action grouping、pretrained vision-language encoders 都是已有设计。实质创新在于把这些机制嵌入 RSSM-style prior/posterior world model 框架，并把 latent belief token 化。可以说这是一种已有 attention/perceiver/VFM 思想对 SSM world model 的结构重组，而不是全新算法范式。

Dataset / Evaluation

实验只在 LIBERO 上做 preliminary evaluation，本质是 offline language-conditioned imitation learning，而不是完整 world model evaluation。LIBERO 覆盖多套操控任务，有多物体、空间关系、目标条件和较长任务序列，确实比简单控制任务更能暴露 RSSM 表示瓶颈。

但这个 evaluation 只能支持较窄的 claim：TokenWM 在 LIBERO-style 视觉语言模仿任务上比同规模 RSSM baseline 更强。它不能充分支持“foundation world model”“more scalable decision-making”“complex environments 下通用 world modeling”这些更大的说法。

主要问题有三点：

1. baseline 不完全对等。TokenWM 使用 224×224 输入和 DINOv2-L pretrained encoder；RSSM 因不易扩展被缩到 64×64。这反映 RSSM 架构局限，但也让增益混入了视觉分辨率和 pretrained feature 的优势。

2. 评价目标是 policy success，不是 world model quality。没有系统报告 multi-step latent rollout、open-loop prediction、planning performance 或 uncertainty calibration。

3. 没有真实机器人、没有 online RL、没有跨数据集泛化。LIBERO demonstration distribution 相对封闭，policy 可能主要学到 dataset-covered behavior patterns。所谓 generalization 是否真实存在，文中未充分说明。

Limitation

最核心的限制是归因不清。TokenWM 一次性改变了太多变量：latent tokenization、DINOv2 encoder、高分辨率输入、attention prior/posterior、memory bank、deterministic loss、policy 联训、两阶段去 decoder 训练。现有 ablation 只展示 memory 有帮助，远远不足以说明主要增益来自 tokenized latent state。

第二，方法对 pretrained representation 依赖很重。DINOv2 已经提供了强语义视觉 token，TokenWM 很可能是在其上做 temporal aggregation 和 policy readout。如果移除 DINOv2 或换到非自然图像、低数据机器人域，性能是否保持，文中未充分说明。

第三，world model 的“推理”能力没有被真正验证。LIBERO success 可以来自 behavior cloning + history retrieval，而不需要模型能准确 imagination。论文没有展示用 learned prior 做 planning，也没有证明 prior rollout 长期稳定。所谓 reasoning 更可能是 representation alignment + retrieval，而非显式 dynamics reasoning。

第四，memory module 把 Markov state 学习问题部分外包给历史窗口。它有效，但也意味着 latent state 本身可能没有形成充分的 Markov belief。窗口大小固定后，长期依赖上限明确；更长任务或延迟更复杂时会遇到 scaling 问题。

第五，放弃 probabilistic latent 是双刃剑。训练稳定了，但不确定性、multi-modal dynamics 和 model-based RL 中常见的 epistemic/aleatoric handling 都缺失。TokenWM 当前更像 deterministic recurrent encoder-policy backbone，而不是成熟的 probabilistic world model。

第六，policy cheating 的讨论很有价值，但也暴露了 world model + policy 联训的脆弱性。作者说 TokenWM 较少受影响，可能与 FiLM action injection 有关，但这是 conjecture。没有更细 ablation 时，这个结论不能当作机制证明。

Takeaway

  • 1. 对复杂视觉控制，latent state 的结构比 temporal backbone 本身更关键。
  • 把 state 从 vector 改成 token workspace 是值得沿着做的方向。
  • 2. Tokenized recurrent world model 的有效路径可能不是替代 Transformer world model，而是作为 VFM/VLA 与 SSM 之间的接口：视觉 foundation model 提供 token observation，SSM 提供递推 belief，policy 从 belief tokens 读取动作。
  • 3. 未来真正该做的是 disentangle：token latent vs pretrained encoder vs memory vs policy architecture。

一句话总结

TokenWM 是把 RSSM-style recurrent belief learning token 化的一次有价值结构改造，真正贡献在于提出 token latent workspace 作为复杂视觉控制中的状态表示，但当前证据更支持其作为强表示型 IL backbone，而非已验证的通用 foundation world model。