Pick-place the ball and pull the puck teach you four things, not two tasks. If a model stores them as four things — two verbs, two objects — then pick-place the puck is not a new task at all. It is a new sentence made of old words. This work proves when a network can be forced to store the world that way, and builds one that does.
The problem
An agent that has mastered push ball to place A can still fail at push ball to place B. The paper’s own framing is blunt: such a policy “might perform notoriously poorly” on what looks like a trivial variation.
The standard remedies each miss for a different reason.
- Augmentation and pretrained visual encoders harden an agent against changes in appearance. They do nothing about changes in dynamics.
- Invariant representation learning keeps what is stable across training tasks — which discards exactly the parts that vary and would need to be recombined.
- Meta-RL is expensive and offers no guarantee of fast convergence when environments vary a lot.
- World models such as DreamerV3 and TD-MPC2 reconstruct observations well, but optimizing for reconstruction yields an entangled latent space. An entangled latent cannot be recombined, because there is no part to swap.
Why language is the seam
Humans do not learn tasks holistically; we learn push, ball, place A and their relationships. That decomposition is what lets us attempt push puck to place A without starting over.
Causal systems have exactly the two properties that make this work. They are modular — each mechanism can be learned on its own and recombined — and sparse, so a distribution shift perturbs a few components and leaves the rest intact. Carve the latent state along causal lines and compositional generalization becomes nearly free: to reach a new task you only repair what actually changed.
The obstacle is that you cannot see the carving. From pixels alone, “verb-ness” and “object-ness” are hopelessly mixed, and you need a supervisory signal that is already compositional to reveal where the seams are.
Language is such a signal. Task instructions arrive pre-factored into verb and object, and each component acts as a separate control signal touching a disjoint set of state dimensions. The structural point is subtle but decisive: this supplies many auxiliary variables instead of one. Prior nonlinear ICA and temporal causal representation learning rely on a single auxiliary variable wired to every latent — which can never express “this word governs these dimensions, that word governs those.”
The environment model. A causal graphical view of the POMDP: each language component controls its own block of the latent state, the blocks interact through the transition dynamics, and only the pixels and reward are observed.
What we can prove
Observations and reward are generated from a latent state that partitions into disjoint, language-controlled components:
Each component is conditioned on the entire previous state, not just its own past — so components interact over time while remaining instantaneously separated by their controlling word.
- A1 — Invertible mixing. The map from latent state to observation and reward is invertible and smooth.
- A2 — Positive density. The set of latent values with zero density has measure zero.
- A3 — Smoothness. The conditional density is at least first-order differentiable.
- A4 — Conditional independence. Given the word, the previous state and the previous action, latent dimensions are mutually independent.
- A5 — Sufficient variation. For each component there exist values of its controlling word whose log-density-gradient differences form an invertible matrix.
Theorem · block-wise identifiability. Under A1–A5, each language-controlled component is identifiable up to an invertible transform within its own block: the estimated block carries exactly the information of the true block and nothing from the others.
The proof matches observation distributions, pushes them through a change of variables, and then constructs differences across values of the controlling word — which cancels an otherwise intractable Jacobian determinant and forces every cross-block derivative to vanish.
Two things are worth being precise about. Identifiability is block-wise, not dimension-wise: individual latent dimensions inside a block are not claimed to be recovered. That weakening is deliberate, and it buys the striking part of the result — to isolate a component you need only distinct values of the word that controls it. To separate the object block, in other words, you need one more object than the block has dimensions.
Building it
The theory says identification is possible; a prototype built on DreamerV3 shows it is reachable. The framework is backbone-agnostic — DreamerV3 is a vehicle, not the contribution.
Factorized dynamics. Both the representation model and the transition model are split into per-component modules, and Dreamer’s single KL term becomes a sum of component-wise KLs — a soft independence constraint that pushes the components apart.
Mutual information constraints. In practice a word’s influence can be too weak for variation alone to separate anything. So for each component, one estimator maximizes mutual information with its own word while the others minimize it against the remaining words. These are MINE estimators trained adversarially inside the world-model loop — fragile enough that the coefficient has to be annealed from zero to avoid destabilizing early training.
Adaptive sparsity. Learnable gated masks select which latents feed the observation, reward and continuation heads. The gating design is borrowed from gated sparse autoencoders in mechanistic interpretability, chosen to avoid plain L1’s shrinkage — which would otherwise destroy information exactly when the agent is still exploring.
Composable world model learning. Per-component posteriors and priors, mutual-information blocks between components, and learnable causal masks deciding which latents reach the observation, reward and continuation heads.
The policy never sees the full state. Only the latents the reward mask selects are passed to the actor and critic, so the agent plans over a compact, reward-relevant slice of what the world model knows.
Imagination and policy learning. Prompted by the language components and an initial observation, the world model composes rollouts from its per-component dynamics; only the reward-relevant latents reach the actor and critic.
Results
The components separate. On a synthetic system built to satisfy the theory — three language components, three values each, twenty-four training combinations and three held out — kernel ridge regression from estimated to true components gives an aggregate R² of 0.9286, with a diagonal of 0.91, 0.92 and 0.95 and off-diagonal terms around 0.1.
Recovered versus true components. Each estimated block explains its own true block and almost none of the others — which is what block-wise identifiability looks like when it works.
That separation is not something the baselines approach. Against iVAE, TCL, TDRL, NCTRL, DreamerV3 and TD-MPC2, WM3C is the only method that gets off the ground: every baseline plateaus below 0.4, including two state-of-the-art world models whose reconstruction quality is not in question. Fitting the observations well and carving the state correctly turn out to be different objectives.
Average R² during training. Shaded bands are the standard deviation over three seeds. Every baseline saturates early and far below WM3C.
Imagining combinations it has never seen. A harder test is to roll the latent state forward in an unseen environment with no observations to correct it. Errors compound, so a model that has merely memorized the training tasks degrades quickly. WM3C holds its accuracy across the horizon.
Imagination on unseen combinations. Average R² of latent rollouts on verb × object pairings the model never trained on, with no observations available during the rollout.
Simulated robot manipulation. In Meta-World — a MuJoCo suite, not physical robots — 18 tasks are used for training and 9 held out. Seven of the nine are pure recombinations of a known verb and a known object; the remaining two introduce a genuinely novel component. Trained on Close Door and Open Faucet, the agent is asked for Close Faucet and Open Door.
Even before adaptation, the decomposition pays off during training itself: on tasks such as box-close and door-unlock, where DreamerV3 struggles to learn a useful representation at all, WM3C still gets off the ground.
Learning curves across the 18 training tasks. Average success rate plus nine individual tasks, three seeds each.
The framing that makes the adaptation result meaningful is the asymmetry of the comparison: WM3C adapts by tuning only its dynamics module, while DreamerV3 and MT-SAC are fine-tuned on all parameters — and WM3C still adapts faster on the seven recombination tasks. That is the theory’s prediction made operational. Only the interactions changed, so only the interactions need retraining.
Adapting to nine unseen tasks. WM3C tunes only its dynamics module; the baselines get full-parameter fine-tuning. On coffee-push, the strongest baseline essentially fails while WM3C recovers the task.
Delete the verb and the gripper disappears
The most direct evidence that the decomposition is real is not a metric. Intervene on the verb component and the robot’s end effector and its adjacent joint vanish from the reconstruction — while the object sits there untouched. Intervene on the object component and the object mutates, arm intact.
Nobody labelled which pixels were “the arm.” With only task instructions as supervision, the model put the thing that does the doing and the thing being done to in different boxes.
Intervening on the language-controlled components. Top row: what the agent sees. Middle: reconstruction after intervening on the object component. Bottom: after intervening on the verb component.
What it cost, and where it breaks
The 18-task Meta-World run takes 8 days on 4× RTX 3090, plus 8 hours for each adaptation task — worth stating plainly, since compositional world models are often discussed as though they were cheap.
The ablation removes the masks and the mutual-information constraints together, so it cannot attribute credit between them. It does surface something counterintuitive: the benefit of those constraints grows as the number of tasks grows. More language variation might have been expected to make the extra structure redundant. The opposite happened.
Ablation at two scales. Removing the causal masks and the mutual-information constraints costs performance at both 5 tasks / 1M steps and 18 tasks / 2M steps — and costs more at the larger scale. The two panels use different y-axis maxima, so read each on its own.
The honest limits are structural. The theory rests on a disjoint partition, so environments with highly overlapping components fall outside it. Long or ambiguous instructions need to be pre-structured — plausibly by an LLM — before they can act as control signals. And on the two test tasks containing a genuinely novel verb or object, adapting only the dynamics module is no longer uniformly sufficient. The authors are also candid that not every identifiability condition is met in practice, even where the recovered latents remain meaningful.
Why it matters
For causal representation learning, this is the first identifiability result for components separately controlled by multiple intermittent signals in the general nonlinear case — moving past the single-auxiliary-variable setting that earlier work inherited from nonlinear ICA.
For reinforcement learning, it suggests a different way to price generalization. If a model’s state is carved along causal lines, the cost of reaching a new task scales with what actually changed in the world, not with the size of the network. Adaptation stops being retraining and starts being repair.