Reproduce “HCLSM: Hierarchical Causal Latent State Machines for Object-Centric World Modeling” on OpenResearch

About this paper

Jaber Jaber, Osama Jaber

HCLSM unifies object-centric decomposition, hierarchical temporal dynamics, and causal structure learning into a differentiable world model, employing a two-stage training protocol to foster specialized object representations. It demonstrates structured understanding of robot manipulation tasks on the PushT dataset and features a custom Triton kernel for improved computational efficiency.

Key insights

• It is possible to unify object-centric decomposition, hierarchical temporal dynamics, and explicit causal structure learning within a single, differentiable world model architecture.
• A two-stage training protocol is critical for successfully learning structured object-aware representations, as it prevents predictive losses from dominating and collapsing specialized slot codes.
• Custom engineering optimizations, such as a Triton kernel for selective State Space Models, are essential for making complex hierarchical dynamics computationally tractable, achieving significant speedups.
• The learned event detector can identify salient state transitions, and the Spatial Broadcast Decoder effectively encourages slots to claim distinct spatial regions, leading to object-aware representation.

Summary by alphaXiv.