🤖 AI

DWM: Separating World Effects from Actions in Latent World Models

This paper introduces DWM (Decomposed World Model), a supervision-level framework that explicitly separates action-driven transitions from action-invariant world effects in latent world models, thereby improving planning success in environments with persistent dynamics like gravity and inertia.

Original authors: Yi-Ge Zhang, Tianqi Du, Qi Zhang, Yisen Wang

Published 2026-07-22

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Yi-Ge Zhang, Tianqi Du, Qi Zhang, Yisen Wang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to play a game of pool. You want the robot to learn how to hit the cue ball so it knocks the eight-ball into a pocket. But there's a catch: the table isn't perfectly flat. It has a slight tilt, so even if the robot does nothing, gravity is constantly trying to slide the balls down the slope. To be a good player, the robot needs to understand two things at once: how its own cue stick changes the ball's path, and how the table's tilt is moving the ball all by itself. This is the heart of "embodied AI," a field where computers learn to act in the physical world by building mental maps of how things move.

For a computer to plan ahead, it uses something called a "latent world model." Think of this as a dream machine inside the robot's brain. Instead of watching the real world frame-by-frame, the robot creates a simplified, invisible sketch of the future. It asks, "If I push the ball this way, what will happen?" The problem is that current dream machines treat every change in the future as a single, messy blob. They don't know which part of the movement came from the robot's push and which part came from the table's tilt. They just see the ball moving and try to guess the whole thing at once, which gets very confusing when the environment has its own stubborn habits.

This paper, titled "DWM: Separating World Effects from Actions in Latent World Models," tackles that confusion head-on. The authors, researchers from Peking University, argue that to make robots smarter, we need to teach their dream machines to split the future into two distinct ingredients: the "Action Effect" (what the robot causes) and the "World Effect" (what the environment does on its own, like gravity or wind). They introduce a new training method called DWM (Decomposed World Model). Instead of forcing the robot to guess the whole future in one go, DWM gives it a special "World Head" that learns to predict only the stuff that happens even if the robot does nothing. Then, a separate "Action Head" learns the rest. By forcing these two heads to work together but stay distinct, the robot learns a much clearer picture of cause and effect.

The researchers tested this idea on several robotic tasks, including a version of a pushing game where a block slides down a slope due to gravity even when the robot isn't touching it. They found that their new method, DWM, was much better at planning than the old, "monolithic" way. On these tricky, gravity-filled tasks, DWM improved the robot's success rate by an average of 13.1%. For example, on one task called PushT-W, the old method only succeeded 32% of the time, while DWM got it up to 44%. The paper suggests that by untangling the robot's actions from the world's natural drift, the robot can imagine the future more accurately and make better decisions, even when the environment is trying to move things on its own. Importantly, this improvement didn't hurt the robot's performance on flat, easy tasks where gravity wasn't a factor; it just made the robot smarter when the world was doing its own thing.

Technical Summary: DWM: Separating World Effects from Actions in Latent World Models

Problem Statement

Latent world models are foundational to modern model-based control, enabling planning and policy learning directly from pixel observations. However, current action-conditioned formulations suffer from a supervision-level limitation: they predict the next latent state using a single, undifferentiated target. This forces the model to absorb all sources of state change—both those induced by the agent's action and those arising from the environment's intrinsic dynamics—into a monolithic learning signal.

In real-world physical systems, transitions often arise from two heterogeneous sources:

Action-driven effects: Changes induced by the agent's control.
World effects: Action-invariant changes that persist even under a null action (e.g., gravity-driven sliding, inertia, contact rebound, or persistent drift).

Fusing these into a single target entangles the two components within the latent transition. Consequently, the model cannot explicitly attribute observed changes to their underlying causes. This limitation is often hidden in static benchmarks (e.g., flat-table pushing) but becomes critical in environments with persistent autonomous dynamics, leading to inaccurate multi-step rollouts and suboptimal planning.

Methodology: Decomposed World Model (DWM)

The authors propose DWM (Decomposed World Model), a training framework that operationalizes the separation of world and action effects without altering the underlying architecture or the inference pipeline.

Core Architecture

DWM retains the standard components of an action-conditioned latent world model (encoder $\phi$ , action-conditioned predictor $g$ , and a projection head):

Pred Head ( $h_p$ ): Retained from the baseline. It predicts the full next latent state $\hat{z}$ and is the only head used during inference.
World Head ( $h_w$ ): An auxiliary, lightweight MLP attached in parallel to the predictor state $r_t$ . It is trained to capture the action-invariant component of the transition ( $\hat{z}_w \approx z_t + \Delta z^{world}_t$ ).
Action-Driven Component: Defined as the residual $\hat{z}_a := \hat{z} - \hat{z}_w$ .

The model enforces the identity $\hat{z} = \hat{z}_w + \hat{z}_a$ , explicitly decomposing the predicted transition into an action-invariant world component and a complementary action-driven component.

Training Objective

The total loss function combines the standard world model loss with two auxiliary terms:
$L = L_{WM} + \lambda_{wc}L_{wc} + \lambda_{orth}L_{orth}$

World-Contrastive Loss ( $L_{wc}$ ): An InfoNCE objective applied to the world head. It constructs positive pairs by feeding the same state history with different current actions (action perturbations) into the world head. This forces $\hat{z}_w$ to be invariant to the action while remaining discriminative across different states.
Orthogonality Constraint ( $L_{orth}$ ): Defined as $|\cos(\hat{z}_w, \hat{z}_a)|$ , this term encourages the action-invariant and action-driven components to span independent directions in the latent space, preventing the action effect from being absorbed into the world representation.

Crucially, the auxiliary world head is discarded after training, ensuring the inference-time pipeline remains identical to the baseline.

Key Contributions

Identification of a Supervision Limitation: The paper identifies that single-target supervision in latent world models fails to distinguish between agent-induced changes and autonomous environmental dynamics, leading to entangled representations.
New Benchmark Suite (W-Variants): The authors introduce controlled variants of three standard benchmarks—PushT-W, Reacher-W, and TwoRoom-W—which inject persistent action-invariant dynamics (gravity-driven sliding, vertical-plane gravity, and constant drift) into standard tasks. They also evaluate on Ball-in-Cup, a task with state-dependent oscillatory dynamics.
DWM Framework: A training framework that augments existing latent world models with an auxiliary world head and complementary separation objectives, requiring no changes to the backbone or inference.
Empirical Validation: Demonstration that DWM improves planning success and latent dynamics accuracy specifically in settings with world effects, while maintaining performance on standard tasks.

Experimental Results

The authors evaluated DWM against the LeWM baseline across the proposed benchmarks using Cross-Entropy Method (CEM) planning.

Planning Success: On the W-variant tasks, DWM achieved a mean absolute improvement of 13.1% in CEM planning success:
- PushT-W: 32.0% $\to$ 44.0% (+12.0 pp)
- Reacher-W: 60.0% $\to$ 70.7% (+10.7 pp)
- TwoRoom-W: 60.0% $\to$ 76.7% (+16.7 pp)
- Ball-in-Cup: 36.8% $\to$ 42.8% (+6.0 pp)
Baseline Performance: On the original "flat" counterparts (without world effects), DWM remained comparable to the single-head LeWM baseline, indicating the auxiliary losses do not degrade performance when world effects are negligible.
Representation Diagnostics: Analysis showed the world head became nearly action-invariant (variance reduced by orders of magnitude under action perturbation), while the pred head retained its sensitivity to actions.
Prediction Quality: DWM significantly reduced one-step and horizon-20 autoregressive prediction errors (e.g., -89.8% offline loss on PushT-W), demonstrating that the decomposition leads to more accurate latent dynamics.
Out-of-Distribution (OOD) Generalization: On PushT-W with shifted gravity vectors, DWM maintained high planning success (44.0%) compared to the baseline's degradation (30.0%), and the learned action-driven component remained stable across gravity shifts (cosine similarity $\approx$ 0.999).

Significance and Claims

The paper claims that disentangling action-invariant world effects from action-driven changes at the level of the training signal serves as a useful inductive bias for action-conditioned latent world models in physically realistic control settings.

The authors emphasize that their approach addresses a fundamental limitation in current supervision strategies without requiring architectural overhauls. By explicitly separating the "what the world does on its own" from "what the agent changes," DWM enables models to learn more accurate latent dynamics, which directly translates to improved multi-step planning and robustness to environmental dynamics that persist independently of agent actions. The results suggest that this decomposition is particularly vital for tasks where the environment evolves autonomously, a scenario often overlooked in standard benchmarks.