论文
HCLSM:用于以对象为中心的世界建模的分层因果潜在状态机
HCLSM: Hierarchical Causal Latent State Machines for Object-Centric World Modeling
摘要
从视频中预测未来状态的世界模型仍然受到平面潜在表示的限制,这些表示会纠缠对象、忽略因果结构并将时间动态压缩为单一尺度。我们提出了 HCLSM,一种按照三个相互关联的原则运行的世界模型架构:通过时隙注意和空间广播解码进行以对象为中心的分解,通过结合连续物理的选择性状态空间模型的三级引擎的分层时间动态,用于离散事件的稀疏 Transformer 和用于抽象目标的压缩 Transformer,以及通过图神经网络交互模式进行因果结构学习。 HCLSM 引入了一个两阶段训练协议,其中空间重建在动态预测开始之前强制槽专门化。我们在来自 Open X-Embodiment 数据集的 PushT 机器人操作基准上训练了 68M 参数模型,通过新兴空间分解(SBD 损失:0.0075)和学习事件边界实现了 0.008 MSE 下一状态预测损失。用于 SSM 扫描的定制 Triton 内核比顺序 PyTorch 提供 38 倍的加速。完整的系统跨越 51 个模块、8,478 行 Python 代码和 171 个单元测试。代码:https://github.com/rightnow-ai/hclsm