HCLSM: Hierarchical Causal Latent State Machines for Object-Centric World Modeling

📄 arXiv: 2603.29090v1 📥 PDF

作者: Jaber Jaber, Osama Jaber

分类: cs.LG, cs.CV, cs.RO

发布日期: 2026-03-31

备注: 10 pages, 3 tables, 4 figures, 1 algorithm. Code: https://github.com/rightnow-ai/hclsm

🔗 代码/项目: GITHUB


💡 一句话要点

HCLSM:用于对象中心世界建模的分层因果隐状态机

🎯 匹配领域: 支柱一:机器人控制 (Robot Control) 支柱二:RL算法与架构 (RL & Architecture)

关键词: 世界模型 对象中心建模 分层时间动态 因果关系学习 机器人操作 状态空间模型 Transformer

📋 核心要点

  1. 现有世界模型使用扁平隐表示,难以解耦对象、建模因果关系和处理多尺度时间动态。
  2. HCLSM通过对象中心分解、分层时间动态建模和因果结构学习,构建更有效的世界模型。
  3. 在PushT机器人操作任务上,HCLSM实现了低至0.008的MSE下一状态预测损失,并涌现了空间分解能力。

📝 摘要(中文)

本文提出了一种名为HCLSM的世界模型架构,旨在克服现有世界模型在视频未来状态预测方面的局限性,这些局限性包括扁平的隐表示纠缠对象、忽略因果结构以及将时间动态压缩到单一尺度。HCLSM基于三个相互关联的原则:通过带有空间广播解码的Slot Attention实现对象中心分解;通过一个三层引擎实现分层时间动态,该引擎结合了用于连续物理的Selective State Space Models、用于离散事件的稀疏Transformer和用于抽象目标的压缩Transformer;以及通过图神经网络交互模式实现因果结构学习。HCLSM引入了一个两阶段训练协议,其中空间重建在动态预测开始之前强制进行Slot专业化。在Open X-Embodiment数据集的PushT机器人操作基准上训练了一个68M参数的模型,实现了0.008 MSE的下一状态预测损失,并涌现了空间分解(SBD损失:0.0075)和学习到的事件边界。用于SSM扫描的自定义Triton内核比顺序PyTorch加速了38倍。整个系统跨越51个模块的8,478行Python代码,包含171个单元测试。代码可在https://github.com/rightnow-ai/hclsm 获取。

🔬 方法详解

问题定义:现有世界模型通常采用扁平的隐表示,导致对象纠缠,难以有效建模对象之间的因果关系,并且无法捕捉不同时间尺度上的动态变化。这限制了模型在复杂环境中的预测能力和泛化性能。

核心思路:HCLSM的核心思路是将世界建模分解为三个关键部分:对象中心表示、分层时间动态和因果关系。通过对象中心表示,模型能够显式地表示和跟踪场景中的各个对象。分层时间动态建模允许模型捕捉不同时间尺度上的变化,从连续的物理运动到离散的事件发生。因果关系建模则帮助模型理解对象之间的相互作用,从而做出更准确的预测。

技术框架:HCLSM的整体架构包含三个主要模块:对象中心分解模块、分层时间动态模块和因果结构学习模块。对象中心分解模块使用Slot Attention和空间广播解码器将场景分解为多个对象槽。分层时间动态模块包含一个三层引擎,分别使用Selective State Space Models (SSM) 建模连续物理,稀疏Transformer建模离散事件,压缩Transformer建模抽象目标。因果结构学习模块使用图神经网络学习对象之间的交互模式。模型采用两阶段训练协议,首先进行空间重建以强制Slot专业化,然后进行动态预测。

关键创新:HCLSM的关键创新在于其分层因果隐状态机架构,该架构能够同时解决对象纠缠、时间尺度混合和因果关系缺失的问题。通过对象中心分解,模型能够显式地表示和跟踪场景中的各个对象。通过分层时间动态建模,模型能够捕捉不同时间尺度上的变化。通过因果结构学习,模型能够理解对象之间的相互作用。

关键设计:HCLSM使用了Slot Attention进行对象分割,并采用空间广播解码器进行重建。时间动态模块使用了三种不同的模型:SSM用于连续物理,稀疏Transformer用于离散事件,压缩Transformer用于抽象目标。模型采用两阶段训练策略,首先训练对象分割,然后训练时间动态模型。为了加速SSM的计算,作者还开发了一个自定义的Triton内核,实现了38倍的加速。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

HCLSM在PushT机器人操作基准上取得了显著的成果,实现了0.008 MSE的下一状态预测损失,并涌现了空间分解能力(SBD损失:0.0075)。此外,作者还开发了一个自定义的Triton内核,将SSM的计算速度提高了38倍。这些结果表明HCLSM是一种有效的世界建模方法,具有很高的实用价值。

🎯 应用场景

HCLSM具有广泛的应用前景,例如机器人操作、自动驾驶、游戏AI等领域。通过对环境进行建模和预测,HCLSM可以帮助机器人更好地理解和控制环境,从而完成更复杂的任务。在自动驾驶领域,HCLSM可以用于预测其他车辆和行人的行为,从而提高驾驶安全性。在游戏AI领域,HCLSM可以用于创建更智能、更逼真的游戏角色。

📄 摘要(原文)

World models that predict future states from video remain limited by flat latent representations that entangle objects, ignore causal structure, and collapse temporal dynamics into a single scale. We present HCLSM, a world model architecture that operates on three interconnected principles: object-centric decomposition via slot attention with spatial broadcast decoding, hierarchical temporal dynamics through a three-level engine combining selective state space models for continuous physics, sparse transformers for discrete events, and compressed transformers for abstract goals, and causal structure learning through graph neural network interaction patterns. HCLSM introduces a two-stage training protocol where spatial reconstruction forces slot specialization before dynamics prediction begins. We train a 68M-parameter model on the PushT robotic manipulation benchmark from the Open X-Embodiment dataset, achieving 0.008 MSE next-state prediction loss with emerging spatial decomposition (SBD loss: 0.0075) and learned event boundaries. A custom Triton kernel for the SSM scan delivers 38x speedup over sequential PyTorch. The full system spans 8,478 lines of Python across 51 modules with 171 unit tests. Code: https://github.com/rightnow-ai/hclsm