Law, State, and Body: Object-Centric World Models for Counterfactual Reasoning

Most learned world models compress a scene into a single flat vector and roll it forward with a recurrent network or a transformer. This works well for prediction, but it hides the objects that a physical scene is actually made of, so the model has no natural way to answer a question such as "what would happen if this object moved differently." We present LSB, a world model that keeps physical law, the visible state of a scene, and the body that acts on it as three separate, swappable pieces. State is a set of per-object slots produced by slot attention. Law is a self-attention transformer that updates every slot using only the slots themselves, so the same law applies regardless of which actuator produced the action. Body is a small embodiment-specific encoder that turns raw actions and proprioception into a control signal that Law can consume. Because objects are explicit and individually addressable, we can edit one object's velocity and roll the edited state forward through the same law used for ordinary prediction, giving a direct and cheap way to test counterfactual reasoning. Across a three-embodiment, six-physics-regime benchmark, LSB beats a strong non-slotted baseline by large, statistically robust margins on next-state prediction, on future-contact queries, and on physical intervention, and the advantage survives zero-shot transfer to six unseen physics regimes, to embodiments never trained with intervention, to object counts far outside the training range, and to rollouts sixteen times longer than the training horizon. We further show that the same advantage appears when Law is instantiated as a completely different architecture, a relational graph network in the style of interaction networks, and when Law is trained with a contrastive energy objective instead of regression, which indicates that the gain comes from the explicit object-centric design rather than from any one choice of network or loss.