DWM: Separating World Effects from Actions in Latent World Models
Abstract
Latent world models provide an effective way to model environment dynamics for control, but existing action-conditioned formulations typically supervise the next-latent transition using a single prediction target, requiring the model to explain all sources of state change through one unified learning signal. In real-world environments, however, state transitions often arise from two distinct sources: changes caused by the agent’s actions and changes induced by the environment’s intrinsic dynamics. The latter can occur even under a null action and include effects such as gravity-driven sliding, inertia, contact rebound, and persistent drift. Modeling these two sources jointly can entangle their effects in the learned latent transition, making it difficult to attribute observed changes to their underlying causes and limiting the transferability of the learned dynamics. In this paper, we introduce DWM (Decomposed World Model), which is a supervision-level framework that explicitly separates these two sources of transition. DWM augments a latent world model with an auxiliary world head that captures action-invariant changes and is regularized by a normalized world-contrastive objective. The original prediction head is coupled with the world head through an orthogonality constraint, encouraging the predicted transition to decompose into an action-invariant component and a complementary action-driven component. This formulation requires no modification to the underlying world-model architecture or inference pipeline. To evaluate DWM in environments with persistent world effects, we construct W-variants of three standard control benchmarks, PushT-W, Reacher-W, and TwoRoom-W, each introducing a distinct form of action-invariant dynamics. Across multiple representative latent world models, DWM consistently improves both prediction quality and downstream control performance on these benchmarks.
est. 32% chance this paper gets accepted at ICLR 2027.
What do you think this paper will get?
All positions stay anonymous.