DWM: Separating World Effects from Actions in Latent World Models

# DWM: Separating World Effects from Actions in Latent World Models
Source: [https://arxiv.org/html/2607.18715](https://arxiv.org/html/2607.18715)
Yi\-Ge Zhang Tianqi Du Qi Zhang Yisen Wang Peking University

###### Abstract

Latent world models underpin much of modern model\-based control, yet current action\-conditioned formulations supervise the next\-latent transition with a*single, undifferentiated target*, forcing a monolithic learning signal to absorb every source of state change\. In real world, however, transitions arise from two heterogeneous sources: an*action\-driven*component induced by the agent, and an*action\-invariant world effect*—the change that would still occur under a null action, dictated by the environment’s intrinsic dynamics \(e\.g\., gravity\-driven sliding, inertia, contact rebound, and persistent drift\)\. Fusing them into a single target entangles the two inside the latent transition, prevents the model from attributing observed changes to their underlying causes, and undermines the transferability of the learned dynamics\. We introduceDWM\(*Decomposed World Model*\), a supervision\-level framework that operationalizes this decomposition\. DWM augments the predictor of a latent world model with an auxiliary*world head*, regularized by a normalized world\-contrastive objective to be action\-invariant, while the original*pred head*is coupled to it via an orthogonality constraint; together, the two signals induce an explicit additive decomposition of the predicted transition into an action\-invariant and a complementary action\-driven component, without altering the underlying architecture or inference pipeline\. To evaluate DWM under persistent world effects, we construct*W\-variants*of three standard control benchmarks—PushT\-W,Reacher\-W, andTwoRoom\-W—each instantiating a distinct action\-invariant dynamic\. DWM matches strong baselines on the flat counterparts and delivers a mean absolute improvement of13\.1%13\.1\\%in CEM planning success across the W\-variants\.

## 1Introduction

Latent world models have gained increasing attention in embodied AI: by predicting future states in a learned latent space, they enable model\-based planning, imagination\-based policy learning, and long\-horizon reasoning directly from pixel observations\(Hafneret al\.,[2019](https://arxiv.org/html/2607.18715#bib.bib4);[2020](https://arxiv.org/html/2607.18715#bib.bib5);[2021](https://arxiv.org/html/2607.18715#bib.bib6);[2023](https://arxiv.org/html/2607.18715#bib.bib7); Hansenet al\.,[2022](https://arxiv.org/html/2607.18715#bib.bib8);[2024](https://arxiv.org/html/2607.18715#bib.bib9)\)\. Recent action\-conditioned latent\-dynamics formulations, represented by LeWM\(Assranet al\.,[2023](https://arxiv.org/html/2607.18715#bib.bib10); Maeset al\.,[2026](https://arxiv.org/html/2607.18715#bib.bib11)\), further strengthen this line by learning end\-to\-end action\-conditioned latent dynamics that generalize across manipulation, navigation, and control tasks, while showing that a compact latent transition model is often sufficient for planning competitive policies from raw images\.

A physical transition, however, is rarely caused by the agent’s action alone\. Consider a block sliding on a tilted surface\. Part of its motion is produced by the agent’s push, while another part would occur even if the agent did nothing, because gravity continues to move the block\. We refer to these two sources of change as theaction effectand theworld effect, respectively\. More precisely, the world effect is the part of the transition that would persist under the same state and history if the current action were replaced by a null action; the action effect is the remaining change attributable to the chosen action\. Similar decompositions arise when objects move because of inertia, articulated bodies swing under gravity, or agents drift under persistent external forces\.

Existing action\-conditioned latent world models generally do not distinguish between these two sources\. They predict the next latent state using a single target that contains both the action effect and the world effect\. Consequently, the model only observes their combined outcome and receives no explicit training signal indicating which part of the transition is caused by the action and which part reflects the environment’s autonomous dynamics\. This limitation can remain hidden in nearly static benchmarks, such as flat\-table pushing, where objects move primarily in response to the agent\. It becomes more consequential when the environment continues to evolve independently of the current action\. In such settings, a model that fails to represent the world effect reliably may produce inaccurate multi\-step rollouts, causing a planner to select actions based on an incorrect prediction of how the environment will evolve\.

Figure[1](https://arxiv.org/html/2607.18715#S1.F1)illustrates this issue on PushT\-W, a gravity\-perturbed variant of the standard PushT task\. Even when the pusher remains idle, the block continues to slide along the direction of gravity\. A conventional single\-head latent world model fails to reproduce this passive motion accurately, and its predicted trajectory consequently deviates from the goal\. In contrast, our model explicitly encourages the latent transition to separate the gravity\-induced world effect from the change induced by the pusher, resulting in a more accurate rollout and a successful plan\. These observations motivate the central question of this work:

![Refer to caption](https://arxiv.org/html/2607.18715v1/Figs/fig1.png)Figure 1:World/action entanglement on PushT\-W\. In every panel, thegray Tis the current block, thegreen Tis the goal pose, and theblue circleis the pusher \(i\.e\., the end\-effector controlled by the agent\)\. The arrow labelled𝐠\\mathbf\{g\}in the top\-left of each panel indicates the direction of the persistent gravity that defines the PushT\-W world effect: the block keeps sliding along𝐠\\mathbf\{g\}even when the pusher is idle\.Top:zero\-action rollout in which the pusher is held still, isolating the action\-invariant world effect—the gray T slides along𝐠\\mathbf\{g\}toward the lower\-right without any push\.Middle:on the same episode, the single\-head LeWM baseline fuses the action\-driven and the action\-invariant components into a single next\-latent target and fails to reproduce the sliding motion: its predicted block trajectory drifts away from the green goal\.Bottom:our proposed DWM on the same episode disentangles the two components and recovers both the gravity\-induced sliding of the block and the pusher\-driven change, bringing the gray T back onto the green T\.> *Can latent world models be improved by explicitly separating what the world does on its own from what the agent’s action changes?*

We answer this question with Decomposed World Model \(DWM\), a world/action\-decomposed training framework for action\-conditioned latent world models\. DWM retains the original prediction head, which continues to predict the complete next latent state and remains the only head used at inference\. During training, we introduce an auxiliary world head that is encouraged to capture the action\-invariant component of the transition\. For the same observation history, we perturb the current action and require the world\-head prediction to remain stable, while preserving its ability to distinguish different states\. We then define the complementary action\-driven component as the residual between the complete prediction and the world\-head prediction, and introduce an orthogonality regularizer to encourage the two components to encode complementary transition information\. Because the additional branch is used only during training, DWM leaves the inference\-time architecture and planning pipeline of the underlying world model unchanged\.

To evaluate DWM, we construct controlled variants of three standard benchmarks, i\.e\., PushT\-W, Reacher\-W, and TwoRoom\-W, in which the environment exhibits a persistent source of action\-invariant motion\. The three tasks respectively introduce gravity\-driven sliding, vertical\-plane gravity, and constant environmental drift\. Their original counterparts serve as controls in which autonomous world dynamics are weak or absent\. We additionally evaluate on Ball\-in\-Cup, a more challenging simulated control task with state\-dependent pendulum\-like dynamics\. Across the three W\-variant benchmarks, DWM improves CEM planning success by 12\.0%, 10\.7%, and 16\.7%, respectively, corresponding to an average absolute improvement of 13\.1%\. At the same time, it remains comparable to the single\-head baseline on the original tasks\. Representation diagnostics, multi\-step rollout evaluations, ablation studies, and out\-of\-distribution tests further show that the gains are associated with increased action invariance in the world head and more accurate latent dynamics\.

Our main contributions are as follows:

- •We identify a supervision\-level limitation of action\-conditioned latent world models: a single next\-latent target merges changes caused by the agent with changes that would occur without the current action, providing no explicit signal for separating the two sources of transition\.
- •We introduce a controlled benchmark suite consisting of PushT\-W, Reacher\-W, and TwoRoom\-W, which adds persistent action\-invariant dynamics to standard control tasks while retaining their original objectives and flat counterparts as controls\.
- •We propose DWM, a training framework that augments an existing latent world model with an auxiliary action\-invariant world head and a complementary world/action separation objective, without changing the base backbone or inference\-time pipeline\.
- •DWM consistently improves CEM planning on all three W\-variant tasks by an average of 13\.1%, improves performance on Ball\-in\-Cup by 6\.0%, and preserves performance on tasks without substantial world effects\.

## 2Related Work

#### Latent world models for planning\.

Latent world models learn compact predictive representations for imagination, planning, and policy learning\(Hafneret al\.,[2021](https://arxiv.org/html/2607.18715#bib.bib6);[2023](https://arxiv.org/html/2607.18715#bib.bib7); Hansenet al\.,[2022](https://arxiv.org/html/2607.18715#bib.bib8);[2024](https://arxiv.org/html/2607.18715#bib.bib9); Maeset al\.,[2026](https://arxiv.org/html/2607.18715#bib.bib11); Micheliet al\.,[2023](https://arxiv.org/html/2607.18715#bib.bib47); Robineet al\.,[2023](https://arxiv.org/html/2607.18715#bib.bib48)\)\. Despite differences in model architecture and downstream use, most existing methods train latent dynamics with a single next\-state prediction objective, where one predictor output is matched to the future latent representation\. Under this formulation, state changes induced by the agent’s actions and those governed by the environment’s own dynamics are absorbed into the same supervision signal\. Our method is complementary to prior advances in planners\(Rubinstein,[1997](https://arxiv.org/html/2607.18715#bib.bib24); Chuaet al\.,[2018](https://arxiv.org/html/2607.18715#bib.bib25); Hansenet al\.,[2024](https://arxiv.org/html/2607.18715#bib.bib9)\), backbones\(Bardeset al\.,[2024](https://arxiv.org/html/2607.18715#bib.bib52); Caronet al\.,[2021](https://arxiv.org/html/2607.18715#bib.bib72)\), and policy objectives\(Hafneret al\.,[2020](https://arxiv.org/html/2607.18715#bib.bib5);[2023](https://arxiv.org/html/2607.18715#bib.bib7)\)\. Instead, it modifies the supervision of a shared latent predictor by introducing two projections with complementary objectives, explicitly separating action\-driven and action\-invariant dynamics during training\.

#### Decomposing latent transitions in world models\.

Most prior work has focused on disentangling what a latent state represents, rather than what causes its transition\. Object\-centric models factorize states into entities or spatial components\(Kipfet al\.,[2020](https://arxiv.org/html/2607.18715#bib.bib19); Ferraroet al\.,[2025](https://arxiv.org/html/2607.18715#bib.bib20); Locatelloet al\.,[2020](https://arxiv.org/html/2607.18715#bib.bib64)\), while task\-oriented methods separate task\-relevant from irrelevant factors\(Wanget al\.,[2022](https://arxiv.org/html/2607.18715#bib.bib15); Fuet al\.,[2021](https://arxiv.org/html/2607.18715#bib.bib81)\)\. Exogenous\-state methods often discard uncontrollable factors\(Efroniet al\.,[2022](https://arxiv.org/html/2607.18715#bib.bib82)\), whereas such dynamics must be retained and predicted for planning in our setting\. Recent work on softly state\-invariant world models\(Saanumet al\.,[2024](https://arxiv.org/html/2607.18715#bib.bib99)\)simplifies latent dynamics by regularizing how much action effects depend on the current state\. This makes action effects more parsimonious, but still leaves the next\-latent target to mix action\-invariant world motion with action\-driven change\. Iso\-Dream\(Panet al\.,[2022](https://arxiv.org/html/2607.18715#bib.bib83)\)uses separate latent branches and forward models for controllable and noncontrollable dynamics\. Since both world effects and action effects jointly determine the next latent state, the goal is not to discard the action\-invariant component but to allocate supervision between the world\- and action\-driven parts\. In contrast, we retain a shared encoder and predictor, and separate these components only through two shallow projection heads with complementary objectives and an orthogonality constraint\. Our method therefore treats disentanglement as a supervision problem rather than an architectural decomposition, making it applicable to existing action\-conditioned latent world models without modifying their backbone\.

## 3World And Action Effects In Latent Dynamics

### 3\.1Preliminaries

#### Action\-conditioned latent world model\.

We adopt the general recipe shared by essentially all action\-conditioned latent world models\(Assranet al\.,[2023](https://arxiv.org/html/2607.18715#bib.bib10); Maeset al\.,[2026](https://arxiv.org/html/2607.18715#bib.bib11)\): instead of predicting future pixels, the model predicts future*latents*, and all supervision happens in that latent space\. Concretely, we assume access to trajectories of raw observationsoto\_\{t\}\(e\.g\., pixels\) and executed actionsata\_\{t\}, and introduce three components on top of them:

- •an*encoder*ϕ\\phithat maps each raw observationoto\_\{t\}to a compact latent statezt=ϕ​\(ot\)z\_\{t\}=\\phi\(o\_\{t\}\), i\.e\., the representation on which all downstream prediction and planning operates;
- •an*action\-conditioned predictor*ggthat summarizes a short context window of past latents and actions\(zt−h\+1:t,at−h\+1:t\)\(z\_\{t\-h\+1:t\},\\,a\_\{t\-h\+1:t\}\)into an intermediate rollout statert=g​\(zt−h\+1:t,at−h\+1:t\)r\_\{t\}=g\(z\_\{t\-h\+1:t\},\\,a\_\{t\-h\+1:t\}\), wherehhis the context length;rtr\_\{t\}can be read as “what the model believes about the world after observing history up tottand takingata\_\{t\}”;
- •a*projection head*that reads outrtr\_\{t\}and produces the predicted next latentz^t\+1\\hat\{z\}\_\{t\+1\}111We denote the generic one\-step predictionz^t\+1\\hat\{z\}\_\{t\+1\}simply byz^\\hat\{z\}when no ambiguity arises, while retaining full time indices for multi\-step rollouts\.\.

The predictor is trained by regressingz^t\+1\\hat\{z\}\_\{t\+1\}towards the true next latentzt\+1z\_\{t\+1\}with a mean\-squared lossℒpred\\mathcal\{L\}\_\{\\mathrm\{pred\}\}, together with a latent\-space regularizerℒsigreg\\mathcal\{L\}\_\{\\mathrm\{sigreg\}\}that keeps the representation from collapsing to a trivial constant; we denote the standard training objective byℒWM=ℒpred\+λsig​ℒsigreg\\mathcal\{L\}\_\{\\mathrm\{WM\}\}=\\mathcal\{L\}\_\{\\mathrm\{pred\}\}\+\\lambda\_\{\\mathrm\{sig\}\}\\mathcal\{L\}\_\{\\mathrm\{sigreg\}\}\.

#### CEM planning in latent space\.

For control, we use model\-predictive planning with the Cross\-Entropy Method \(CEM\) directly in the latent space induced by the encoderϕ\\phiand the predictorgg\(Rubinstein,[1997](https://arxiv.org/html/2607.18715#bib.bib24); Chuaet al\.,[2018](https://arxiv.org/html/2607.18715#bib.bib25)\)\. Given a raw image goalogo\_\{g\}, the goal is embedded once into a latent targetzg=ϕ​\(og\)z\_\{g\}=\\phi\(o\_\{g\}\), and planning is carried out entirely by rolling the world model forward from the current context\. Concretely, at each planning step the current history\(zt−h\+1:t,at−h\+1:t−1\)\(z\_\{t\-h\+1:t\},\\,a\_\{t\-h\+1:t\-1\}\)is fixed, and CEM optimizes an action sequenceat:t\+H−1a\_\{t:t\+H\-1\}of horizonHHunder a Gaussian proposal𝒩​\(μi,Σi\)\\mathcal\{N\}\(\\mu\_\{i\},\\Sigma\_\{i\}\)that is refined over iterationsi=1,…,Ii=1,\\dots,I:

\{at:t\+H−1\(k\)\}k=1N∼𝒩​\(μi,Σi\),z^t\+j\+1\(k\)=π​\(g​\(z^t\+j−h\+1:t\+j\(k\),at\+j−h\+1:t\+j\(k\)\)\),\\\{a^\{\(k\)\}\_\{t:t\+H\-1\}\\\}\_\{k=1\}^\{N\}\\sim\\mathcal\{N\}\(\\mu\_\{i\},\\Sigma\_\{i\}\),\\qquad\\hat\{z\}^\{\(k\)\}\_\{t\+j\+1\}=\\pi\\\!\\left\(g\\\!\\left\(\\hat\{z\}^\{\(k\)\}\_\{t\+j\-h\+1:t\+j\},\\,a^\{\(k\)\}\_\{t\+j\-h\+1:t\+j\}\\right\)\\right\),\(1\)whereπ\\pidenotes the projection head reading outrtr\_\{t\}\(in Section[4](https://arxiv.org/html/2607.18715#S4),π=hp\\pi=h\_\{p\}\), andz^t\(k\)=zt\\hat\{z\}^\{\(k\)\}\_\{t\}=z\_\{t\}initializes the rollout with the encoded context\. Each candidate sequence is scored by a cost functionc​\(z^t\+1:t\+H\(k\);zg\)c\(\\hat\{z\}^\{\(k\)\}\_\{t\+1:t\+H\};z\_\{g\}\)that we instantiate either as a terminal latent distance

c​\(z^t\+1:t\+H\(k\);zg\)=∥z^t\+H\(k\)−zg∥22c\\\!\\left\(\\hat\{z\}^\{\(k\)\}\_\{t\+1:t\+H\};z\_\{g\}\\right\)\\;=\\;\\lVert\\hat\{z\}^\{\(k\)\}\_\{t\+H\}\-z\_\{g\}\\rVert\_\{2\}^\{2\}\(2\)for tasks where the raw latent distance is informative, or asc​\(z^t\+1:t\+H\(k\);zg\)=−s​\(z^t\+1:t\+H\(k\);zg\)c\(\\hat\{z\}^\{\(k\)\}\_\{t\+1:t\+H\};z\_\{g\}\)=\-s\(\\hat\{z\}^\{\(k\)\}\_\{t\+1:t\+H\};z\_\{g\}\)for Ball\-in\-Cup \(Section[5\.6](https://arxiv.org/html/2607.18715#S5.SS6)\), withssa learned success head that scores the entire predicted rollout\. The elite setℰi\\mathcal\{E\}\_\{i\}of the top\-KKlowest\-cost samples is then used to update the proposal,

μi\+1=1K​∑k∈ℰiat:t\+H−1\(k\),Σi\+1=1K​∑k∈ℰi\(at:t\+H−1\(k\)−μi\+1\)2,\\mu\_\{i\+1\}=\\frac\{1\}\{K\}\\sum\_\{k\\in\\mathcal\{E\}\_\{i\}\}a^\{\(k\)\}\_\{t:t\+H\-1\},\\qquad\\Sigma\_\{i\+1\}=\\frac\{1\}\{K\}\\sum\_\{k\\in\\mathcal\{E\}\_\{i\}\}\\bigl\(a^\{\(k\)\}\_\{t:t\+H\-1\}\-\\mu\_\{i\+1\}\\bigr\)^\{\\\!2\},\(3\)and the first action of the resulting mean,μI​\[0\]\\mu\_\{I\}\[0\], is executed in the environment; the context is then advanced and planning repeats in receding\-horizon fashion\.

### 3\.2Defining World and Action Effects

The introduction described the world effect as the part of a transition that would persist if the current action were replaced by a null action\. We now formalize this counterfactual notion in latent space\. Fixing the same historyzt−h\+1:tz\_\{t\-h\+1:t\}and past actionsat−h\+1:t−1a\_\{t\-h\+1:t\-1\}, but replacing the current actionata\_\{t\}with the zero action𝟎\\mathbf\{0\}, we define the action\-invariant component of the latent transition as

Δ​ztworld:=𝔼​\[zt\+1\|zt−h\+1:t,at−h\+1:t−1,at=𝟎\]−zt\.\\Delta z^\{\\mathrm\{world\}\}\_\{t\}\\;:=\\;\\mathbb\{E\}\\\!\\left\[\\,z\_\{t\+1\}\\;\\bigl\|\\;z\_\{t\-h\+1:t\},\\,a\_\{t\-h\+1:t\-1\},\\,a\_\{t\}\{=\}\\mathbf\{0\}\\,\\right\]\\;\-\\;z\_\{t\}\.\(4\)Intuitively,Δ​ztworld\\Delta z^\{\\mathrm\{world\}\}\_\{t\}is what the world would do on its own if the agent chose not to intervene at steptt\. It captures the persistent phenomena that our benchmarks are designed to expose, such as gravity\-driven sliding, inertia, and drift\. Crucially, Eq\.[4](https://arxiv.org/html/2607.18715#S3.E4)depends only on the environment dynamics and the current context, not on the actualata\_\{t\}, soΔ​ztworld\\Delta z^\{\\mathrm\{world\}\}\_\{t\}is well\-defined for every transition regardless of whether the two effects combine additively in the true dynamics\.

GivenΔ​ztworld\\Delta z^\{\\mathrm\{world\}\}\_\{t\}, we define the action\-driven component as the residual,Δ​ztaction:=Δ​zt−Δ​ztworld\\Delta z^\{\\mathrm\{action\}\}\_\{t\}\\;:=\\;\\Delta z\_\{t\}\\,\-\\,\\Delta z^\{\\mathrm\{world\}\}\_\{t\}, so that

Δ​zt=Δ​ztworld\+Δ​ztaction\\Delta z\_\{t\}\\;=\\;\\Delta z^\{\\mathrm\{world\}\}\_\{t\}\\;\+\\;\\Delta z^\{\\mathrm\{action\}\}\_\{t\}\(5\)holds by construction\. This is a modelling identity rather than a claim that the two effects combine additively at the physical level\. In real environments they may couple nonlinearly, for instance when the same push produces different net motion depending on the current gravity\-induced velocity\. Our formulation does not assume such coupling away; instead, the residual absorbs whatever interaction the two effects have, and the interaction itself is later carried by the shared latent representation that both prediction paths read from \(Section[4\.1](https://arxiv.org/html/2607.18715#S4.SS1)\)\.

### 3\.3Diagnosing Monolithic Transition Supervision

Existing high\-performing latent world models such as LeWM\(Maeset al\.,[2026](https://arxiv.org/html/2607.18715#bib.bib11)\)and Dino\-WM\(Zhouet al\.,[2024](https://arxiv.org/html/2607.18715#bib.bib102)\)already report strong planning success on standard control benchmarks like PushT and TwoRoom, often reaching an8080–90%90\\%success rate\. These benchmarks, however, share a common limitation: their environments do not involve any non\-trivial physical dynamics, and the core object in the scene moves essentially only in response to the agent’s action\. As a result, the reported performance does not faithfully reflect how these methods behave in more realistic settings where the environment itself keeps evolving\. To better evaluate the effectiveness of these methods, we therefore consider datasets that carry a persistent world effect, such as gravity, which continues to influence the scene dynamics regardless of the action, and observe how existing methods behave when the next state is shaped not only by the agent’s action but also by this action\-invariant world effect\. Concretely, we introduce*PushT\-W*, a gravity\-perturbed variant of the standard 2D pushing task in which the block keeps sliding under gravity irrespective of the action \(its full construction is deferred to Section[5\.1](https://arxiv.org/html/2607.18715#S5.SS1)\), and use it to diagnose LeWM, a representative latent world model trained to convergence under the standard single\-target supervision\. In the remainder of this section, we first show that this paradigm degrades markedly once a world effect is present, and then show that augmenting the data with pure action\-invariant transitions does not remove the degradation\. Together, these two findings locate the difficulty in the training signal itself: because a single next\-latent target provides no mechanism to tell apart change caused by the agent from change that occurs regardless of it, the two sources of state change must instead be*disentangled at the level of the training objective*\. This is the core design principle of our method in Section[4](https://arxiv.org/html/2607.18715#S4)\.

#### Monolithic supervision degrades under a world effect\.

We first ask how a single\-target latent world model behaves once the transition contains a non\-trivial action\-invariant world effect\. On standard PushT, where the block barely moves without the agent’s input and the transition is essentially action\-driven, LeWM trained under the standard single\-target supervision reaches a CEM planning success rate of94\.0%94\.0\\%\(Fig\.[3](https://arxiv.org/html/2607.18715#S5.F3), right\)\. On PushT\-W, in contrast, gravity continually drags the block regardless of the action, so a substantial share of each transition is action\-invariant; the same model, trained to convergence on this variant, only achieves a CEM success rate of32%32\\%\. This sharp drop indicates that LeWM does not learn a latent transition that adequately accounts for the world effect: motion produced by gravity rather than by the agent is folded into the same next\-latent target as action\-driven motion, and the resulting predictor is not accurate enough to support CEM planning\. In the remainder of this section we examine whether this shortfall can be attributed to the training data or to the training objective\.

#### Injecting zero\-action data does not separate world from action\.

A natural hypothesis for the failure above is that it stems from insufficient exposure to purely action\-invariant transitions: if the training set rarely contains frames in which the block moves under gravity alone, the model may never see the world effect in isolation and thus never learn to represent it\. We test this hypothesis by augmenting the PushT\-W training set with a30%30\\%mixture of zero\-action frames, in which the block slides under gravity while the agent stays idle, and retrain the same single\-target model on this enriched dataset\. Despite being explicitly exposed to pure gravity\-driven transitions at scale, the retrained model attains a CEM planning success rate of30%30\\%, statistically indistinguishable from the32%32\\%baseline trained without the additional zero\-action data\. Exposure alone therefore does not force the model to disentangle the world component from the action\-driven one: the two remain fused inside the single next\-latent target regardless of how the data is composed\. This locates the bottleneck not in data coverage but in the training objective itself, which we address next\.

## 4DWM: World/Action Disentanglement for Latent World Models

Section[3\.2](https://arxiv.org/html/2607.18715#S3.SS2)defined the action\-invariant componentΔ​ztworld\\Delta z^\{\\mathrm\{world\}\}\_\{t\}and the complementary action\-driven residualΔ​ztaction\\Delta z^\{\\mathrm\{action\}\}\_\{t\}, and Section[3\.3](https://arxiv.org/html/2607.18715#S3.SS3)showed that a single next\-latent target does not induce this decomposition on its own\. We use the identityΔ​zt=Δ​ztworld\+Δ​ztaction\\Delta z\_\{t\}=\\Delta z^\{\\mathrm\{world\}\}\_\{t\}\+\\Delta z^\{\\mathrm\{action\}\}\_\{t\}as an inductive bias, and reallocate the supervision of the base latent world model so that the two sources of transition are learned by separate prediction paths\. Section[4\.1](https://arxiv.org/html/2607.18715#S4.SS1)instantiates this decomposition at the level of the predicted next latent, and Section[4\.2](https://arxiv.org/html/2607.18715#S4.SS2)introduces the two auxiliary losses that implement it\.

### 4\.1Decomposing the Predicted Transition

We operationalize Eq\.[5](https://arxiv.org/html/2607.18715#S3.E5)at the level of the model’s predicted next latent\. The original prediction head is retained as the*pred head*: its outputz^\\hat\{z\}predicts the full next latentzt\+1z\_\{t\+1\}and remains the only head used at inference\. Alongside it, we attach an auxiliary*world head*to the same predictor statertr\_\{t\}; its outputz^w\\hat\{z\}^\{\\,w\}is regularized to be invariant to the current actionata\_\{t\}and serves as the model’s readout of the action\-invariant component \(z^w≈zt\+Δ​ztworld\\hat\{z\}^\{\\,w\}\\\!\\approx\\\!z\_\{t\}\+\\Delta z^\{\\mathrm\{world\}\}\_\{t\}in the sense of Eq\.[4](https://arxiv.org/html/2607.18715#S3.E4)\)\. The complementary action\-driven component is defined as the residualz^a:=z^−z^w\\hat\{z\}^\{\\,a\}:=\\hat\{z\}\-\\hat\{z\}^\{\\,w\}, so thatz^=z^w\+z^a\\hat\{z\}=\\hat\{z\}^\{\\,w\}\+\\hat\{z\}^\{\\,a\}mirrors Eq\.[5](https://arxiv.org/html/2607.18715#S3.E5)at the level of the model’s predictions\.

![Refer to caption](https://arxiv.org/html/2607.18715v1/Figs/framework.png)Figure 2:DWM architecture\.The upper lane, enclosed by the dashed grey box, is the unchangedLeWM baseline: the current observationoto\_\{t\}is encoded byφ\\varphiintoztz\_\{t\}, the actionata\_\{t\}is injected through the predictorggto produce the rollout statertr\_\{t\}, and the pred headhph\_\{p\}readsrtr\_\{t\}to produce the predicted next latentz^:=hp​\(rt\)\\hat\{z\}\\\!:=\\\!h\_\{p\}\(r\_\{t\}\)\(this is exactly thez^t\+1\\hat\{z\}\_\{t\+1\}of Section[3\.1](https://arxiv.org/html/2607.18715#S3.SS1.SSS0.Px1); the subscript is dropped for brevity\)\. The lower lane encodes the next observationot\+1o\_\{t\+1\}with the*same*encoderφ\\varphito produce the targetzt\+1z\_\{t\+1\}, which supervisesz^\\hat\{z\}through the world\-model lossℒwm\\mathcal\{L\}\_\{\\mathrm\{wm\}\}\.Our addition is a single training\-time branch \(orange\): a*world head*hwh\_\{w\}is attached in parallel to the samertr\_\{t\}\(the*first*representation in which the actionata\_\{t\}, injected bygg, has already been fused with world context\), producingz^w:=hw​\(rt\)\\hat\{z\}^\{\\,w\}\\\!:=\\\!h\_\{w\}\(r\_\{t\}\); the action\-driven component is then obtained by subtraction,z^a:=z^−z^w\\hat\{z\}^\{\\,a\}\\\!:=\\\!\\hat\{z\}\-\\hat\{z\}^\{\\,w\}\. Two auxiliary losses shape this decomposition:ℒwc\\mathcal\{L\}\_\{\\mathrm\{wc\}\}forcesz^w\\hat\{z\}^\{\\,w\}to be invariant toata\_\{t\}\(Section[4\.2](https://arxiv.org/html/2607.18715#S4.SS2)\), andℒorth=\|cos⁡\(z^w,z^a\)\|\\mathcal\{L\}\_\{\\mathrm\{orth\}\}=\|\\cos\(\\hat\{z\}^\{\\,w\},\\hat\{z\}^\{\\,a\}\)\|forces the two components to span independent directions\. At inference the orange branch is discarded and the pipeline reduces exactly to LeWM\.Figure[2](https://arxiv.org/html/2607.18715#S4.F2)summarizes the architecture\. The LeWM backbone—encoderφ\\varphi, action\-conditioned predictorgg, and the pred headhph\_\{p\}—is kept unchanged; our only structural change is the orange branch:

- •World headhwh\_\{w\}: a lightweight MLP attached tortr\_\{t\}in parallel tohph\_\{p\}, trained byℒwc\\mathcal\{L\}\_\{\\mathrm\{wc\}\}to be action\-invariant\. It shares the encoder and predictor withhph\_\{p\}, so the auxiliary signal reshapes the representation thathph\_\{p\}reads from\.
- •Pred headhph\_\{p\}: unchanged; its outputz^:=hp​\(rt\)\\hat\{z\}\\\!:=\\\!h\_\{p\}\(r\_\{t\}\)is the samez^t\+1\\hat\{z\}\_\{t\+1\}of Section[3\.1](https://arxiv.org/html/2607.18715#S3.SS1.SSS0.Px1)\(subscript dropped for brevity\) and is regressed tozt\+1z\_\{t\+1\}underℒWM\\mathcal\{L\}\_\{\\mathrm\{WM\}\}\. This is the only head used at inference, so the pipeline reduces exactly to LeWM once training is done\.

The action\-driven component is then defined asz^a≜z^−z^w\\hat\{z\}^\{\\,a\}\\\!\\triangleq\\\!\\hat\{z\}\-\\hat\{z\}^\{\\,w\}, giving the explicit decompositionz^=z^w\+z^a\\hat\{z\}=\\hat\{z\}^\{\\,w\}\+\\hat\{z\}^\{\\,a\}that operationalizes Eq\.[5](https://arxiv.org/html/2607.18715#S3.E5):z^w\\hat\{z\}^\{\\,w\}carries what is stable across actions, andz^a\\hat\{z\}^\{\\,a\}carries the variation needed to makez^\\hat\{z\}action\-conditioned\.

### 4\.2Learning Complementary World and Action Components

The objective of the original world model,ℒWM\\mathcal\{L\}\_\{\\mathrm\{WM\}\}\(Section[3\.1](https://arxiv.org/html/2607.18715#S3.SS1.SSS0.Px1)\), is retained without modification and continues to supervise the pred head\. Our method adds two terms, highlighted in blue, that impose the world/action separation:

ℒ=ℒWM\+λwc​ℒwc\+λorth​ℒorth\.\\mathcal\{L\}\\;=\\;\\mathcal\{L\}\_\{\\mathrm\{WM\}\}\\;\+\{\\color\[rgb\]\{0,0,1\}\\lambda\_\{\\mathrm\{wc\}\}\\mathcal\{L\}\_\{\\mathrm\{wc\}\}\+\\lambda\_\{\\mathrm\{orth\}\}\\mathcal\{L\}\_\{\\mathrm\{orth\}\}\}\.\(6\)The two new terms are:

- •ℒwc\\mathcal\{L\}\_\{\\mathrm\{wc\}\}: the*world\-contrastive*loss onhwh\_\{w\}, implemented as an InfoNCE objective\(van den Oordet al\.,[2018](https://arxiv.org/html/2607.18715#bib.bib28); Chenet al\.,[2020](https://arxiv.org/html/2607.18715#bib.bib29); Heet al\.,[2020](https://arxiv.org/html/2607.18715#bib.bib30)\)whose positive and negative pairs are constructed from action perturbations of the same state\. For each contextiiin a mini\-batch of sizeBB, we*additionally*sample a second current actiona~t\(i\)\\tilde\{a\}^\{\(i\)\}\_\{t\}, drawn independently from the empirical action distribution \(in practice a random permutation ofata\_\{t\}within the batch\)\. Passing the two action choices through the shared predictor—while keeping the rest of the context fixed—and then throughhwh\_\{w\}yields two views of the same state:z^w,\(i\)=hw​\(g​\(zt−h\+1:t\(i\),at−h\+1:t\(i\)\)\)\\hat\{z\}^\{\\,w,\(i\)\}=h\_\{w\}\(g\(z^\{\(i\)\}\_\{t\-h\+1:t\},a^\{\(i\)\}\_\{t\-h\+1:t\}\)\)andz^~w,\(i\)=hw​\(g​\(zt−h\+1:t\(i\),at−h\+1:t−1\(i\),a~t\(i\)\)\)\\tilde\{\\hat\{z\}\}^\{\\,w,\(i\)\}=h\_\{w\}\(g\(z^\{\(i\)\}\_\{t\-h\+1:t\},a^\{\(i\)\}\_\{t\-h\+1:t\-1\},\\tilde\{a\}^\{\(i\)\}\_\{t\}\)\)\. Letu\(i\)u^\{\(i\)\}andu~\(i\)\\tilde\{u\}^\{\(i\)\}be theirℓ2\\ell\_\{2\}\-normalized versions\. We then take positive pair:\(u\(i\),u~\(i\)\)\(u^\{\(i\)\},\\tilde\{u\}^\{\(i\)\}\), i\.e\., same stateii, different current actions, and negatives:\{u~\(j\)\}j≠i\\\{\\tilde\{u\}^\{\(j\)\}\\\}\_\{j\\neq i\}, i\.e\., different states, different actions, and minimize the symmetric InfoNCE ℒwc=−12​B​∑i=1B\[log⁡eu\(i\)⊤​u~\(i\)/τ∑jeu\(i\)⊤​u~\(j\)/τ\+log⁡eu~\(i\)⊤​u\(i\)/τ∑jeu~\(i\)⊤​u\(j\)/τ\],\\mathcal\{L\}\_\{\\mathrm\{wc\}\}=\-\\frac\{1\}\{2B\}\\sum\_\{i=1\}^\{B\}\\\!\\left\[\\log\\frac\{e^\{u^\{\(i\)\\top\}\\tilde\{u\}^\{\(i\)\}/\\tau\}\}\{\\sum\_\{j\}e^\{u^\{\(i\)\\top\}\\tilde\{u\}^\{\(j\)\}/\\tau\}\}\+\\log\\frac\{e^\{\\tilde\{u\}^\{\(i\)\\top\}u^\{\(i\)\}/\\tau\}\}\{\\sum\_\{j\}e^\{\\tilde\{u\}^\{\(i\)\\top\}u^\{\(j\)\}/\\tau\}\}\\right\],\(7\)with temperatureτ=0\.07\\tau=0\.07\. Intuitively, the positive term forcesz^w\\hat\{z\}^\{\\,w\}to be*invariant*toata\_\{t\}, while the negative term keeps it*discriminative*across states so that it does not collapse to a constant\. The objective is self\-supervised: no external label of the world component is ever used\.
- •ℒorth=\|cos⁡\(z^w,z^a\)\|\\mathcal\{L\}\_\{\\mathrm\{orth\}\}=\|\\cos\(\\hat\{z\}^\{\\,w\},\\,\\hat\{z\}^\{\\,a\}\)\|: an orthogonality constraint that forcesz^w⟂z^a\\hat\{z\}^\{\\,w\}\\\!\\perp\\\!\\hat\{z\}^\{\\,a\}, so the action\-invariant content ofz^\\hat\{z\}is aligned withz^w\\hat\{z\}^\{\\,w\}instead of being absorbed into an undifferentiated target\. Sinceℒwc\\mathcal\{L\}\_\{\\mathrm\{wc\}\}acts on the shared encoder and predictor, combining the two losses yields a clean decompositionz^=z^w\+z^a\\hat\{z\}=\\hat\{z\}^\{\\,w\}\+\\hat\{z\}^\{\\,a\}, which we verify empirically in Section[5](https://arxiv.org/html/2607.18715#S5)\.

## 5Experiments

### 5\.1Benchmarks and Experimental Setup

To evaluate whether world/action disentanglement matters in settings closer to physical control than idealized flat benchmarks, we pair each of three standard tasks—PushT, Reacher, and TwoRoom—with a*W\-variant*that introduces a persistent action\-invariant source of motion \(gravity\-driven sliding, vertical\-plane gravity, or constant environmental drift, respectively\)\. Each W\-variant keeps the same goal as its flat counterpart but adds a source of motion that persists even when the agent is idle; the flat counterpart serves as a comparability control\. We further evaluate on Ball\-in\-Cup as an additional real\-world setting whose action\-invariant component is a state\-dependent pendulum\-like oscillation rather than a constant or gravitational bias\. Full physical parameterizations, dataset sizes, and OOD splits are deferred to Appendix[A](https://arxiv.org/html/2607.18715#A1)\(Table[5](https://arxiv.org/html/2607.18715#A1.T5)\)\.

#### Model and optimization\.

The encoderφ\\varphiis instantiated as a ViT\-tiny \(hidden dimension192192,1212layers, patch size1414, input resolution2242224^\{2\}\), and the action\-conditioned predictorggas a66\-layer Transformer \(1616attention heads, head dimension6464, MLP dimension20482048\) with AdaLN action conditioning; the context length is set toh=3h\\\!=\\\!3\(Section[3\.1](https://arxiv.org/html/2607.18715#S3.SS1.SSS0.Px1)\), and observations are subsampled at a stride of55frames\. The pred headhph\_\{p\}and the world headhwh\_\{w\}\(Section[4\.1](https://arxiv.org/html/2607.18715#S4.SS1)\) are each instantiated as a two\-layer MLP192→2048→192192\\\!\\rightarrow\\\!2048\\\!\\rightarrow\\\!192with BatchNorm\. The encoder, predictor, and both heads are initialized randomly and optimized jointly with AdamW \(lr=5×10−5\\mathrm\{lr\}\\\!=\\\!5\\\!\\times\\\!10^\{\-5\}across all modules, weight decay10−310^\{\-3\}\) under a linear\-warmup cosine\-annealing schedule applied at the epoch level\. We train inbfloat16on NVIDIA H20 GPUs with batch size128128\. The number of epochs is fixed per task family prior to evaluation:55for PushT / PushT\-W,2020for Reacher / Reacher\-W,100100for TwoRoom / TwoRoom\-W, and2020for Ball\-in\-Cup\. This protocol ensures that each task family reaches convergence before evaluation while keeping LeWM and DWM under an identical training budget\.

#### Evaluation metrics\.

- •Prediction quality:pred\_loss, i\.e\., the one\-step prediction error betweenz^t\+1\\hat\{z\}\_\{t\+1\}andzt\+1z\_\{t\+1\}, reported both offline on the held\-out validation set and under autoregressive rollout\. For PushT\-W and TwoRoom\-W it is the MSE in the latent space; for Reacher\-W it is the qpos error \(rad\) in the physical space in which the task goal is defined\.
- •Disentanglement diagnostics: on each W\-variant task we compare how much the world head and the pred head change under controlled action perturbations\. A useful decomposition should make the world head nearly action\-invariant while preserving action sensitivity in the pred head used for rollout\.
- •OOD gravity: on PushT\-W, the cosine similarity of the learned action\-driven componentz^a\\hat\{z\}^\{\\,a\}across held\-out gravities, together with CEM planning success at gravity settings distinct from the training gravity\.
- •Planning: CEM success rate on a fixed set of start–goal pairs, averaged over 3 CEM planner seeds\. Each W\-variant task uses a task\-specific planning protocol \(goal offset, evaluation budget, CEM sample and refinement counts, and success criterion\), specified in Appendix[A](https://arxiv.org/html/2607.18715#A1)\.

### 5\.2Main Results: Planning with and without World Effects

![Refer to caption](https://arxiv.org/html/2607.18715v1/x1.png)Figure 3:CEM planning success rate on our benchmark suite\.Left:Main results on the three W\-variant tasks \(with a persistent action\-invariant world effect\), averaged over 3 CEM planner seeds \(n=50n\{=\}50start–goal pairs per seed\); error bars denote across\-seed standard deviation\.Right:Performance on the flat counterparts \(tasks without world effect\), reported under the same training recipe and CEM evaluation\. In both panels we compare the single\-projection LeWM baseline against our world/action\-decomposed model\. Numbers above each bar are mean success rates in percent\.The three W\-variant tasks are the primary testbed for our claim: they take standard control benchmarks whose action\-invariant world effect is negligible and augment each with a persistent source of motion that the agent does not command\. This is precisely the regime in which our supervision\-level hypothesis is falsifiable—if fusing world and action into a single next\-latent target harms planning, the cost should be largest here\. We evaluate every model under an identical CEM protocol, holding out the flat counterparts as a comparability control\. Full CEM configurations are given per task in Appendix[A](https://arxiv.org/html/2607.18715#A1); all bars in Fig\.[3](https://arxiv.org/html/2607.18715#S5.F3)share the training recipe of Section[5\.1](https://arxiv.org/html/2607.18715#S5.SS1.SSS0.Px1)and average over33CEM planner seeds withn=50n\\\!=\\\!50start–goal pairs per seed\.

DWM improves CEM planning success over the LeWM baseline on all three tasks by a large margin \(Fig\.[3](https://arxiv.org/html/2607.18715#S5.F3), left\):32\.0%→44\.0%32\.0\\%\\\!\\to\\\!\\bm\{44\.0\\%\}on PushT\-W \(\+12\.0\+12\.0pp\),60\.0%→70\.7%60\.0\\%\\\!\\to\\\!\\bm\{70\.7\\%\}on Reacher\-W \(\+10\.7\+10\.7pp\), and60\.0%→76\.7%60\.0\\%\\\!\\to\\\!\\bm\{76\.7\\%\}on TwoRoom\-W \(\+16\.7\+16\.7pp\), for an average absolute improvement of\+13\.1\+13\.1pp\. The consistency of these improvements across three qualitatively distinct action\-invariant world effects—gravity\-driven sliding, vertical\-plane gravity, and constant environmental drift—supports our central claim: whenever the transition contains a genuine action\-invariant component, explicitly disentangling it from the action\-driven component yields more accurate latent dynamics and, through them, better plans\. Section[5\.3](https://arxiv.org/html/2607.18715#S5.SS3)tightens this claim by showing that the same model also exhibits a clean representational separation, more accurate multi\-step rollouts, and robustness of the learned decomposition to out\-of\-distribution gravity\.

The right panel of Fig\.[3](https://arxiv.org/html/2607.18715#S5.F3)shows that these gains do not come at the cost of performance in the absence of a world effect\. On the flat counterparts of the three W\-variant tasks—standard PushT, DMC Reacher, and TwoRoom—DWM remains comparable to the LeWM baseline under the same training recipe and CEM protocol\. When the transition is essentially action\-driven and no world effect needs to be isolated, the additional supervision imposed byℒwc\\mathcal\{L\}\_\{\\mathrm\{wc\}\}andℒorth\\mathcal\{L\}\_\{\\mathrm\{orth\}\}is inert, and the underlying world model behaves as in the single\-head baseline; the improvement appears specifically in the regime where a persistent action\-invariant component must be predicted for planning\.

### 5\.3Does DWM Learn the Intended Decomposition?

The main results of Section[5\.2](https://arxiv.org/html/2607.18715#S5.SS2)show that DWM improves CEM planning across all three W\-variant tasks; we now examine whether the intended world/action decomposition emerges inside the model at the level of the learned representation\. On each of the three W\-variant tasks \(PushT\-W, TwoRoom\-W, and Reacher\-W\), a useful decomposition should leave the pred head as responsive toata\_\{t\}as the original single\-head LeWM output, while making the world head essentially unresponsive toata\_\{t\}\.

For each validation window we hold the observation historyzt−h\+1:tz\_\{t\-h\+1:t\}fixed and replace the current actionata\_\{t\}withKKrandomly sampled alternatives, obtainingKKpaired outputsz^w\\hat\{z\}^\{\\,w\}of the world head andz^\\hat\{z\}of the pred head\. From these paired outputs we compute, on each head, the per\-sample variance and the maximum–minimum spread across theKKaction perturbations, averaged over validation windows; we quantify both effects in Table[1](https://arxiv.org/html/2607.18715#S5.T1)\.

Table 1:Disentanglement diagnostics on the three W\-variant tasks\. For each validation window the current actionata\_\{t\}is replaced withKKrandom alternatives while the history is held fixed; we report the per\-sample variance and max–min spread of each head across these perturbations, averaged over windows\. Since LeWM has no world head, only the pred\-head variance column is populated for it\.Table[1](https://arxiv.org/html/2607.18715#S5.T1)supports two consistent observations across the three W\-variants\. First, the pred head is not degraded by our auxiliary training: its response toata\_\{t\}remains within the same order of magnitude as the single\-head LeWM output on every task \(e\.g\.,0\.8640\.864vs\.0\.8470\.847on PushT\-W, and comparable on TwoRoom\-W and Reacher\-W\), which is what allows CEM planning throughz^\\hat\{z\}to remain well\-posed\. Second, the world head is close to a constant under action perturbation on all three tasks: its variance is two to three orders of magnitude smaller than the pred head’s, and both the world/pred variance ratio and the world/pred spread ratio stay well below one across tasks\. The auxiliary head therefore does not learn to mimic the action\-conditioned predictor but instead concentrates on the component of the transition that survives when the current action is varied, i\.e\., the action\-invariant world effect that our objective was designed to isolate\. At the level of the learned representation itself, the world/action decomposition posited in Section[4](https://arxiv.org/html/2607.18715#S4)is realized inside the model that produced the planning gains of Section[5\.2](https://arxiv.org/html/2607.18715#S5.SS2)\.

### 5\.4Multi\-Step Prediction and OOD Dynamics Generalization

Section[5\.3](https://arxiv.org/html/2607.18715#S5.SS3)verified that the decomposition emerges at the representation level\. We next ask whether it also manifests in the learned latent dynamics, in two complementary ways: by improving multi\-step prediction along the horizon that CEM actually queries, and by remaining stable when the source of the world effect is shifted out of distribution\.

#### Latent prediction and rollout quality\.

The planning gains above should be reflected in the latent dynamics that CEM optimizes, rather than only in downstream success rates\. We therefore evaluate the predictor on held\-out validation windows of each W\-variant at two levels: one\-step prediction, which tests local transition accuracy, and horizon\-2020autoregressive rollout, which tests whether errors remain controlled over the planning horizon\. For PushT\-W and TwoRoom\-W we report the latent\-space MSE; for Reacher\-W we report the qpos error \(rad\) in the physical space in which the task goal is defined \(Section[5\.1](https://arxiv.org/html/2607.18715#S5.SS1.SSS0.Px2)\)\.

Table 2:Latent prediction quality on the three W\-variant tasks\. “Offline” reports one\-step prediction on the held\-out validation set; “Rollout@20” reports autoregressive prediction at horizon2020on the same set\. For PushT\-W and TwoRoom\-W the error is the latent\-space MSE; for Reacher\-W it is the qpos error \(rad\) in the physical space in which the task goal is defined\.Table[2](https://arxiv.org/html/2607.18715#S5.T2)reports two complementary quantities on each W\-variant task: an*offline*one\-step prediction error, which measures the local accuracy of the learned transition, and a horizon\-2020*autoregressive rollout*error, which measures whether that accuracy is preserved once errors are composed over the horizon that CEM actually queries\. On all three tasks and at both horizons, DWM improves over the LeWM baseline\. At one step, the prediction error drops by89\.8%89\.8\\%on PushT\-W,31\.6%31\.6\\%on TwoRoom\-W, and29\.8%29\.8\\%on Reacher\-W, indicating that the auxiliary training signal shapes the shared representation so that the pred head resolves the local transition more accurately than a single\-target predictor, regardless of whether the world effect is a gravity\-driven slide, a constant environmental drift, or a vertical\-plane gravity\. The improvement further survives autoregressive composition: at horizon2020, DWM reduces the error by66\.7%66\.7\\%,31\.7%31\.7\\%, and26\.1%26\.1\\%on the three tasks respectively, so the one\-step advantage is not amplified into a larger downstream error when the model is used recurrently\. On PushT\-W the effect is particularly pronounced: DWM’s horizon\-2020error \(0\.3440\.344\) is already lower than the LeWM baseline’s one\-step error \(0\.3720\.372\), which means that a2020\-step rollout of the world/action\-decomposed model is more accurate than a single\-step prediction of the single\-projection baseline\. Taken together, these results indicate that the planning gains of Section[5\.2](https://arxiv.org/html/2607.18715#S5.SS2)are supported by demonstrably more reliable latent dynamics along the horizons queried by CEM, rather than by a favourable interaction between the CEM proposal distribution and a merely locally accurate predictor\.

#### OOD gravity generalization\.

We probe robustness of the learned decomposition on PushT\-W, whose world effect is a two\-dimensional gravity vectorworld\.gravityand can therefore be varied continuously to produce a well\-defined shift in the underlying dynamics while every other planning input is held fixed\. We train the model at a single gravity setting,\(45,25\)\(45,25\), and evaluate it at gravity values distinct from training\. For each setting we report two quantities: the cosine similarity between the learned action\-driven componentz^a\\hat\{z\}^\{\\,a\}and its in\-distribution counterpart, which probes whether the learned decomposition is stable across gravity, and CEM planning success on the same start–goal pairs and CEM budget as in Section[5\.2](https://arxiv.org/html/2607.18715#S5.SS2), which probes whether that stability translates into control\.

![Refer to caption](https://arxiv.org/html/2607.18715v1/x2.png)Figure 4:OOD gravity CEM planning success on PushT\-W \(n=50n\{=\}50pairs per gravity\); OOD\-pooled combines the two OOD gravities\.Two consistent observations emerge from this evaluation\. First, the action\-driven component is largely invariant to the tested gravity shifts: pooled over\(90,0\)\(90,0\)and\(0,45\)\(0,45\), the cosine similarity between the out\-of\-distribution and in\-distributionz^a\\hat\{z\}^\{\\,a\}attains a mean of0\.99910\.9991, with a mean effect MSE of0\.00150\.0015\. This indicates that DWM does not fit to a specific value ofworld\.gravity; rather, it isolates a direction of action effects that remains stable when the world component varies\. Second, this representational stability propagates to downstream control \(Fig\.[4](https://arxiv.org/html/2607.18715#S5.F4)\): the LeWM baseline degrades to30\.0%30\.0\\%pooled across the two OOD gravities, while DWM retains44\.0%44\.0\\%, yielding a\+14\+14pp improvement that is statistically significant \(p≈0\.040p\\\!\\approx\\\!0\.040, two\-proportionzz\-test\) and no smaller than the gap observed at the training gravity\. Taken together, these results indicate that the planning gains of Section[5\.2](https://arxiv.org/html/2607.18715#S5.SS2)are attributable to the learned world/action decomposition, rather than to overfitting to the training\-time world dynamics\.

### 5\.5Ablation Studies

Sections[5\.3](https://arxiv.org/html/2607.18715#S5.SS3)and[5\.4](https://arxiv.org/html/2607.18715#S5.SS4)attribute the planning gains to the intended world/action decomposition; it remains to verify that each auxiliary loss in Eq\.[6](https://arxiv.org/html/2607.18715#S4.E6),ℒwc\\mathcal\{L\}\_\{\\mathrm\{wc\}\}andℒorth\\mathcal\{L\}\_\{\\mathrm\{orth\}\}, is necessary rather than redundant\. We probe each weight by setting the other to its default value and reducing it to0\.10\.1and then to0, comparing the resulting CEM planning success on PushT\-W with our full objective and with the LeWM baseline\. All rows in Table[3](https://arxiv.org/html/2607.18715#S5.T3)are trained from scratch under the protocol of Section[5\.2](https://arxiv.org/html/2607.18715#S5.SS2)and evaluated with the same CEM configuration\.

Table 3:Loss\-weight ablation on PushT\-W\. Each row reports CEM planning success for a single choice of the two auxiliary loss weights\(λwc,λorth\)\(\\lambda\_\{\\mathrm\{wc\}\},\\lambda\_\{\\mathrm\{orth\}\}\)\. The row\(0,0\)\(0,0\)recovers the LeWM baseline and the row\(0\.3,0\.5\)\(0\.3,0\.5\)is our default; the two other rows in each block reduce one weight while holding the other at its default value\. Same scratch training and CEM evaluation as Fig\.[3](https://arxiv.org/html/2607.18715#S5.F3)\(left\) \(n=50n\{=\}50\)\.Table[3](https://arxiv.org/html/2607.18715#S5.T3)yields two observations\. First, CEM success is monotonically increasing inλwc\\lambda\_\{\\mathrm\{wc\}\}along the world\-contrastive axis: at fixedλorth=0\.5\\lambda\_\{\\mathrm\{orth\}\}=0\.5, decreasingλwc\\lambda\_\{\\mathrm\{wc\}\}from0\.30\.3to0\.20\.2to0\.10\.1produces CEM success of44\.0%44\.0\\%,38\.0%38\.0\\%and36\.0%36\.0\\%respectively, so a weaker contrastive signal lets the world head drift away from action\-invariance and recovers only about half of the gain over the LeWM baseline\. Second, CEM success also degrades along the orthogonality axis whenλorth\\lambda\_\{\\mathrm\{orth\}\}is reduced: at fixedλwc=0\.3\\lambda\_\{\\mathrm\{wc\}\}=0\.3, decreasingλorth\\lambda\_\{\\mathrm\{orth\}\}from0\.50\.5to0\.10\.1to0\.00\.0produces44\.0%44\.0\\%,42\.0%42\.0\\%and40\.0%40\.0\\%, indicating that a world\-invariant head alone does not couple its component intoz^\\hat\{z\}tightly enough to close the remaining gap\. Neither loss reproduces the full\+12\+12pp gain in isolation, soℒwc\\mathcal\{L\}\_\{\\mathrm\{wc\}\}andℒorth\\mathcal\{L\}\_\{\\mathrm\{orth\}\}act as complementary rather than redundant terms in Eq\.[6](https://arxiv.org/html/2607.18715#S4.E6)\.

### 5\.6Beyond Controlled World Effects: Ball\-in\-Cup

The three W\-variants in Section[5\.2](https://arxiv.org/html/2607.18715#S5.SS2)isolate an action\-invariant world effect that is either a constant per\-step drift or a static gravitational pull\. Real physical environments, however, often exhibit action\-invariant dynamics that are neither constant nor state\-independent\. Ball\-in\-Cup makes this regime concrete: the ball continues to swing under gravity and string tension even when the cup is held stationary, so the action\-invariant component of the transition is a state\-dependent oscillation whose direction and magnitude vary with the ball’s current position and velocity\. This constitutes a strictly harder test than the W\-variants and probes whether the world/action decomposition learned by DWM continues to be useful when the world component is no longer a simple parametric quantity\.

We train both LeWM and DWM on Ball\-in\-Cup under the recipe of Section[5\.1](https://arxiv.org/html/2607.18715#S5.SS1.SSS0.Px1)and evaluate them with CEM under the protocol of Section[3\.1](https://arxiv.org/html/2607.18715#S3.SS1.SSS0.Px2)\. Because raw latent goal distance is uninformative on Ball\-in\-Cup, we score CEM rollouts with a learned*success head*on top of the predicted rollout latents\. Since the two models induce different rollout\-latent distributions, a single shared success head would systematically favour one of them; to eliminate this confound, we train a dedicated success head for each model on its own rollouts and score each model with its own head, so that any planning gap cannot be attributed to one method being paired with a stronger evaluator\.

Table 4:Ball\-in\-Cup CEM planning success and autoregressive rollout prediction loss\. Same training recipe as Section[5\.1](https://arxiv.org/html/2607.18715#S5.SS1.SSS0.Px1)and same CEM budget across the two rows\.Table[4](https://arxiv.org/html/2607.18715#S5.T4)reports two matched\-evaluation quantities on Ball\-in\-Cup: CEM planning success under the per\-model success\-head protocol above, and the autoregressive rolloutpred\_lossof the underlying world model over the same horizon used for planning\. DWM improves CEM planning success from36\.8%36\.8\\%to42\.8%42\.8\\%\(\+6\+6pp\), and simultaneously reduces the rolloutpred\_lossby roughly an order of magnitude, from0\.6710\.671to0\.0890\.089\. The two quantities move together, mirroring the mechanistic account of Section[5\.3](https://arxiv.org/html/2607.18715#S5.SS3): better latent\-dynamics quality translates into better CEM plans\. The planning gap is smaller in absolute terms than on the three W\-variants, which is expected— Ball\-in\-Cup exposes a qualitatively harder, state\-dependent oscillatory world effect than the ones DWM was designed around, and noise from action sampling and evaluator scoring accumulates over the longer horizons this task requires\. Even so, the same objective, unchanged, continues to yield a consistent gain in both prediction and control, indicating that the utility of the world/action decomposition is not confined to the constant or static world effects of the W\-variant benchmark\.

## 6Conclusion

We identified a supervision\-level limitation of action\-conditioned latent world models: their next\-latent transition is treated as a single, undifferentiated target, fusing action\-driven and action\-invariant components into one learning signal\. To address this, we proposed a world/action\-decomposed training framework that keeps a single shared encoder and predictor and separates the two components through a world head trained by a world\-contrastive objective, a pred head retained as the next\-latent predictor, and an orthogonality constraint that disentanglesz^w\\hat\{z\}^\{\\,w\}fromz^a\\hat\{z\}^\{\\,a\}\. On a benchmark of three flat/W\-variant task pairs plus Ball\-in\-Cup, the resulting objective improves CEM planning success by an average of\+13\.1\+13\.1pp on the W\-variant tasks and transfers to Ball\-in\-Cup, while remaining comparable to the single\-head baseline on the flat counterparts\. These results suggest that disentangling action\-invariant world effects from action\-driven change at the level of the training signal is a useful inductive bias for action\-conditioned latent world models in physically realistic control settings\.

## References

- M\. Assran, Q\. Duval, I\. Misra, P\. Bojanowski, P\. Vincent, M\. Rabbat, Y\. LeCun, and N\. Ballas \(2023\)Self\-supervised learning from images with a joint\-embedding predictive architecture\.arXiv preprint arXiv:2301\.08243\.Cited by:[§1](https://arxiv.org/html/2607.18715#S1.p1.1),[§3\.1](https://arxiv.org/html/2607.18715#S3.SS1.SSS0.Px1.p1.2)\.
- A\. Bardes, Q\. Garrido, J\. Ponce, X\. Chen, M\. Rabbat, Y\. LeCun, M\. Assran, and N\. Ballas \(2024\)Revisiting feature prediction for learning visual representations from video\.Transactions on Machine Learning Research\.Cited by:[§2](https://arxiv.org/html/2607.18715#S2.SS0.SSS0.Px1.p1.1)\.
- M\. Caron, H\. Touvron, I\. Misra, H\. Jégou, J\. Mairal, P\. Bojanowski, and A\. Joulin \(2021\)Emerging properties in self\-supervised vision transformers\.InICCV,Cited by:[§2](https://arxiv.org/html/2607.18715#S2.SS0.SSS0.Px1.p1.1)\.
- T\. Chen, S\. Kornblith, M\. Norouzi, and G\. Hinton \(2020\)A simple framework for contrastive learning of visual representations\.InICML,Cited by:[1st item](https://arxiv.org/html/2607.18715#S4.I2.i1.p1.2)\.
- C\. Chi, S\. Feng, Y\. Du, Z\. Xu, E\. Cousineau, B\. Burchfiel, and S\. Song \(2023\)Diffusion policy: visuomotor policy learning via action diffusion\.InRSS,Cited by:[Appendix A](https://arxiv.org/html/2607.18715#A1.SS0.SSS0.Px1.p1.12)\.
- K\. Chua, R\. Calandra, R\. McAllister, and S\. Levine \(2018\)Deep reinforcement learning in a handful of trials using probabilistic dynamics models\.InNeurIPS,Cited by:[§2](https://arxiv.org/html/2607.18715#S2.SS0.SSS0.Px1.p1.1),[§3\.1](https://arxiv.org/html/2607.18715#S3.SS1.SSS0.Px2.p1.9)\.
- Y\. Efroni, D\. Misra, A\. Krishnamurthy, A\. Agarwal, and J\. Langford \(2022\)Provable rl with exogenous distractors via multistep inverse dynamics\.InICLR,Cited by:[§2](https://arxiv.org/html/2607.18715#S2.SS0.SSS0.Px2.p1.1)\.
- S\. Ferraro, P\. Mazzaglia, T\. Verbelen, and B\. Dhoedt \(2025\)FOCUS: object\-centric world models for robotic manipulation\.Frontiers in Neurorobotics19,pp\. 1585386\.Cited by:[§2](https://arxiv.org/html/2607.18715#S2.SS0.SSS0.Px2.p1.1)\.
- X\. Fu, G\. Yang, P\. Agrawal, and T\. Jaakkola \(2021\)Learning task informed abstractions\.InICML,Cited by:[§2](https://arxiv.org/html/2607.18715#S2.SS0.SSS0.Px2.p1.1)\.
- D\. Hafner, T\. Lillicrap, J\. Ba, and M\. Norouzi \(2020\)Dream to control: learning behaviors by latent imagination\.InICLR,Cited by:[§1](https://arxiv.org/html/2607.18715#S1.p1.1),[§2](https://arxiv.org/html/2607.18715#S2.SS0.SSS0.Px1.p1.1)\.
- D\. Hafner, T\. Lillicrap, I\. Fischer, R\. Villegas, D\. Ha, H\. Lee, and J\. Davidson \(2019\)Learning latent dynamics for planning from pixels\.InICML,Cited by:[§1](https://arxiv.org/html/2607.18715#S1.p1.1)\.
- D\. Hafner, T\. Lillicrap, M\. Norouzi, and J\. Ba \(2021\)Mastering atari with discrete world models\.InICLR,Cited by:[§1](https://arxiv.org/html/2607.18715#S1.p1.1),[§2](https://arxiv.org/html/2607.18715#S2.SS0.SSS0.Px1.p1.1)\.
- D\. Hafner, J\. Pasukonis, J\. Ba, and T\. Lillicrap \(2023\)Mastering diverse domains through world models\.arXiv preprint arXiv:2301\.04104\.Cited by:[§1](https://arxiv.org/html/2607.18715#S1.p1.1),[§2](https://arxiv.org/html/2607.18715#S2.SS0.SSS0.Px1.p1.1)\.
- N\. Hansen, H\. Su, and X\. Wang \(2024\)TD\-mpc2: scalable, robust world models for continuous control\.InICLR,Cited by:[§1](https://arxiv.org/html/2607.18715#S1.p1.1),[§2](https://arxiv.org/html/2607.18715#S2.SS0.SSS0.Px1.p1.1)\.
- N\. Hansen, X\. Wang, and H\. Su \(2022\)Temporal difference learning for model predictive control\.InICML,Cited by:[§1](https://arxiv.org/html/2607.18715#S1.p1.1),[§2](https://arxiv.org/html/2607.18715#S2.SS0.SSS0.Px1.p1.1)\.
- K\. He, H\. Fan, Y\. Wu, S\. Xie, and R\. Girshick \(2020\)Momentum contrast for unsupervised visual representation learning\.InCVPR,Cited by:[1st item](https://arxiv.org/html/2607.18715#S4.I2.i1.p1.2)\.
- T\. Kipf, E\. van der Pol, and M\. Welling \(2020\)Contrastive learning of structured world models\.InICLR,Cited by:[§2](https://arxiv.org/html/2607.18715#S2.SS0.SSS0.Px2.p1.1)\.
- F\. Locatello, D\. Weissenborn, T\. Unterthiner, A\. Mahendran, G\. Heigold, J\. Uszkoreit, A\. Dosovitskiy, and T\. Kipf \(2020\)Object\-centric learning with slot attention\.InNeurIPS,Cited by:[§2](https://arxiv.org/html/2607.18715#S2.SS0.SSS0.Px2.p1.1)\.
- L\. Maes, Q\. Le Lidec, D\. Scieur, Y\. LeCun, and R\. Balestriero \(2026\)LeWorldModel: stable end\-to\-end joint\-embedding predictive architecture from pixels\.arXiv preprint arXiv:2603\.19312\.Cited by:[§1](https://arxiv.org/html/2607.18715#S1.p1.1),[§2](https://arxiv.org/html/2607.18715#S2.SS0.SSS0.Px1.p1.1),[§3\.1](https://arxiv.org/html/2607.18715#S3.SS1.SSS0.Px1.p1.2),[§3\.3](https://arxiv.org/html/2607.18715#S3.SS3.p1.2)\.
- V\. Micheli, E\. Alonso, and F\. Fleuret \(2023\)Transformers are sample\-efficient world models\.InICLR,Cited by:[§2](https://arxiv.org/html/2607.18715#S2.SS0.SSS0.Px1.p1.1)\.
- M\. Pan, X\. Zhu, Y\. Wang, and X\. Yang \(2022\)Isolating and leveraging controllable and noncontrollable visual dynamics in world models\.InNeurIPS,Cited by:[§2](https://arxiv.org/html/2607.18715#S2.SS0.SSS0.Px2.p1.1)\.
- J\. Robine, M\. Höftmann, T\. Uelwer, and S\. Harmeling \(2023\)Transformer\-based world models are happy with 100k interactions\.InICLR,Cited by:[§2](https://arxiv.org/html/2607.18715#S2.SS0.SSS0.Px1.p1.1)\.
- R\. Y\. Rubinstein \(1997\)Optimization of computer simulation models with rare events\.European Journal of Operational Research99\(1\),pp\. 89–112\.Cited by:[§2](https://arxiv.org/html/2607.18715#S2.SS0.SSS0.Px1.p1.1),[§3\.1](https://arxiv.org/html/2607.18715#S3.SS1.SSS0.Px2.p1.9)\.
- T\. Saanum, P\. Dayan, and E\. Schulz \(2024\)Simplifying latent dynamics with softly state\-invariant world models\.InNeurIPS,Cited by:[§2](https://arxiv.org/html/2607.18715#S2.SS0.SSS0.Px2.p1.1)\.
- Y\. Tassa, Y\. Doron, A\. Muldal, T\. Erez, Y\. Li, D\. d\. L\. Casas, D\. Budden, A\. Abdolmaleki, J\. Merel, A\. Lefrancq,et al\.\(2018\)DeepMind control suite\.arXiv preprint arXiv:1801\.00690\.Cited by:[Appendix A](https://arxiv.org/html/2607.18715#A1.SS0.SSS0.Px2.p1.12),[Appendix A](https://arxiv.org/html/2607.18715#A1.SS0.SSS0.Px4.p1.2)\.
- S\. Tunyasuvunakool, A\. Muldal, Y\. Doron, S\. Liu, S\. Bohez, J\. Merel, T\. Erez, T\. Lillicrap, N\. Heess, and Y\. Tassa \(2020\)Dm\_control: software and tasks for continuous control\.Software Impacts6,pp\. 100022\.Cited by:[Appendix A](https://arxiv.org/html/2607.18715#A1.SS0.SSS0.Px2.p1.12),[Appendix A](https://arxiv.org/html/2607.18715#A1.SS0.SSS0.Px4.p1.2)\.
- A\. van den Oord, Y\. Li, and O\. Vinyals \(2018\)Representation learning with contrastive predictive coding\.arXiv preprint arXiv:1807\.03748\.Cited by:[1st item](https://arxiv.org/html/2607.18715#S4.I2.i1.p1.2)\.
- T\. Wang, S\. S\. Du, A\. Torralba, P\. Isola, A\. Zhang, and Y\. Tian \(2022\)Denoised mdps: learning world models better than the world itself\.InICML,Cited by:[§2](https://arxiv.org/html/2607.18715#S2.SS0.SSS0.Px2.p1.1)\.
- G\. Zhou, H\. Pan, Y\. LeCun, and L\. Pinto \(2024\)Dino\-wm: world models on pre\-trained visual features enable zero\-shot planning\.arXiv preprint arXiv:2411\.04983\.Cited by:[§3\.3](https://arxiv.org/html/2607.18715#S3.SS3.p1.2)\.

## Appendix ATask and Dataset Details

This appendix provides the full physical parameterization, dataset scale, and out\-of\-distribution \(OOD\) splits for the tasks introduced in Section[5\.1](https://arxiv.org/html/2607.18715#S5.SS1)\. Table[5](https://arxiv.org/html/2607.18715#A1.T5)summarizes the resulting benchmark suite; the paragraphs below describe each pair in turn, together with Ball\-in\-Cup as an additional real\-world setting\.

Table 5:Control tasks used in this paper, grouped into three flat/W\-variant pairs \(“\-W” for*world effect*, and*flat*for the counterpart*without*a world effect\) plus Ball\-in\-Cup as an additional real\-world setting\. Each W\-variant introduces a persistent action\-invariant world effect; the flat counterpart shares the same goal but removes that effect\. Tasks in the same block share a goal \(block pushing, qpos matching, 2D target reaching, catching the ball in the cup\)\.TaskAction\-invariant world effectPhysical settingPushTNoneHorizontal tablePushT\-WGravity\-driven slidingGravity\(45,25\)\(45,25\), damping0\.920\.92ReacherNoneStandard DMC ReacherReacher\-WVertical\-plane gravityg=−9\.81​sin⁡60∘g\{=\}\{\-\}9\.81\\sin 60^\{\\circ\}alongzzTwoRoomNoneΔ​p=a⋅s\\Delta p\{=\}a\\cdot sTwoRoom\-WConstant drift biasΔ​p=a⋅s\+b\\Delta p\{=\}a\\cdot s\+b,b=\(4,2\)b\{=\}\(4,2\)Ball\-in\-CupPendulum\-like oscillationStandard DMC Ball\-in\-Cup#### PushT / PushT\-W\.

Standard PushT\(Chiet al\.,[2023](https://arxiv.org/html/2607.18715#bib.bib38)\)is a horizontal\-table pushing task: the block does not drift under zero action, so action\-invariant world effects are weak\. Its W\-variant counterpart keeps the same pushing objective but includes gravity\(45,25\)\(45,25\)in thepymunkspace \(with damping0\.920\.92\), so the block continues to slide even when the agent is idle\. We use standard PushT as the flat comparability control, and train PushT\-W ontilted\_pusht\_train\_4xwith 3200 episodes×\\times200 steps under a random policy; for OOD stress tests, we additionally hold out validation episodes at gravity\(60,30\)\(60,30\),\(25,60\)\(25,60\)and\(0,0\)\(0,0\)\(Section[5\.3](https://arxiv.org/html/2607.18715#S5.SS3)\)\. CEM planning uses a goal offset of1010frames, an evaluation budget of5050environment steps, and a CEM configuration of300×30×top30300\\\!\\times\\\!30\\\!\\times\\\!\\mathrm\{top\}30\(300300samples,3030refinement iterations, top\-3030elite set\); an attempt is counted as a success if the block matches the goal configuration at the final step of the budget\.

#### Reacher / Reacher\-W\.

Standard DMC Reacher\(Tassaet al\.,[2018](https://arxiv.org/html/2607.18715#bib.bib39); Tunyasuvunakoolet al\.,[2020](https://arxiv.org/html/2607.18715#bib.bib86)\)operates on a horizontal plane, where action\-invariant world effects are negligible\. Reacher\-W keeps the same qpos\-matching objective but applies a60∘60^\{\\circ\}scaled\-zzgravityg=\(0,0,−9\.81​sin⁡60∘\)≈\(0,0,−8\.50\)g\\\!=\\\!\(0,0,\-9\.81\\sin 60^\{\\circ\}\)\\\!\\approx\\\!\(0,0,\-8\.50\)m/s2, so the arm is continually pulled downward regardless of the action\. Both variants use10,00010\{,\}000episodes×200\\times\\,200steps, with a held\-out validation set of500500episodes×200\\times\\,200steps generated under the same physics and policy mix\. CEM planning uses a goal offset of2525frames, an evaluation budget of200200environment steps, and a CEM configuration of300×10×top30300\\\!\\times\\\!10\\\!\\times\\\!\\mathrm\{top\}30; an attempt succeeds if the joint configuration reaches within0\.150\.15rad of the goal qpos at any step within the budget\.

#### TwoRoom / TwoRoom\-W\.

Standard TwoRoom is a 2D navigation environment with transitionpt\+1=pt\+at⋅sp\_\{t\+1\}=p\_\{t\}\+a\_\{t\}\\cdot s, so the agent does not drift under zero action and action\-invariant world effects are absent\. TwoRoom\-W keeps the same reach\-target objective but adds a persistent environmental drift,pt\+1=pt\+at⋅s\+bp\_\{t\+1\}=p\_\{t\}\+a\_\{t\}\\cdot s\+bwithb=\(4,2\)b=\(4,2\)pixels/step\. We train the W\-variant on 2000 episodes×\\times200 steps with a random policy under the same drift, yielding324,498324\{,\}498training windows and a matched validation set of60×10060\\\!\\times\\\!100transitions\. CEM planning on TwoRoom\-W adopts a longer\-horizon CEM\-MPC protocol to accommodate the persistent drift: goal offset9090frames, evaluation budget of2525MPC decisions, each planning over a rollout horizon of55steps and executing an action block of55environment steps before re\-planning, with the same300×10×top30300\\\!\\times\\\!10\\\!\\times\\\!\\mathrm\{top\}30CEM budget as Reacher\-W; success is declared when the agent reaches the target region within the total budget\.

#### Ball\-in\-Cup \(real\-world setting\)\.

The DMC Ball\-in\-Cup environment\(Tassaet al\.,[2018](https://arxiv.org/html/2607.18715#bib.bib39); Tunyasuvunakoolet al\.,[2020](https://arxiv.org/html/2607.18715#bib.bib86)\)exhibits pendulum\-like contact dynamics: the ball continues to swing under gravity and string tension even when the cup is held stationary, so the action\-invariant component of the transition is neither a constant drift nor a static gravitational pull but a state\-dependent oscillation\. This makes Ball\-in\-Cup a substantially harder testbed than the three flat/W\-variant pairs above, and we report it separately in Section[5\.6](https://arxiv.org/html/2607.18715#S5.SS6)as an additional evaluation beyond the main benchmark\. Following the data scale of the original LeWM protocol, the training set consists of800800episodes×200\\times\\,200steps\.