CausalWAM: Learning the Causes of Actions in World Action Models

AI summary

Overview

Research area: Causal representation learning applied to robot learning — specifically world action models (WAMs) that jointly predict video and generate actions from demonstrations.

Technical level: Advanced. The paper's core is an identifiability theory (nonlinear ICA-style arguments over Jacobian supports, diffeomorphisms, and permutation equivalence), paired with a concrete architectural modification to a pretrained video/action foundation model.

Scope (one sentence): The paper proves that a sparsity constraint on the latent-to-action dependency structure isolates the latent variables that actually cause demonstrated actions, and instantiates that constraint as a sparse routing path added to Cosmos 3, evaluated on synthetic data, RoboCasa365, LIBERO, and a real Franka R3 robot.

What This Paper Is About

World action models learn latent variables from images and use them both to predict how a scene will evolve and to generate robot actions. Because prediction must account for the entire scene, those latents necessarily encode irrelevant background details (a tablecloth, clutter), and nothing in standard demonstration-fitting forces the action head to ignore them. The paper asks whether, from demonstrations alone, one can separate what causes actions from what prediction merely needs, and answers yes — provided the number of latent-to-action dependencies is constrained to be minimal.

Key Contributions

1. A negative result without sparsity (Proposition 1). For any mechanism fitted to demonstrations under observational equivalence plus nondegeneracy conditions, each action dimension can depend on at least as many latent variables as the true mechanism — an "over-coverage" guarantee that permits nuisance coordinates to leak into the action pathway.
2. Identifiability of the action structure under support sparsity (Theorem 1). Adding a constraint that the estimated mechanism activate no more latent coordinates than the true one (Equation 3) turns the lower bound into an equality, identifying the latent-to-action dependency pattern up to a permutation of the latent variables.
3. Identifiability of the compact state itself (Theorem 2). The selected action-relevant coordinates form a state C_t that is an invertible function of the true action-relevant block s_{t,I_A} alone — not merely correlated with it — so nuisance variables cannot mix in. The guarantee requires no interventions, no sufficient temporal dependence, no assumptions on how unobserved regimes evolve, and permits statistical dependence between action-relevant and nuisance variables.
4. CWAM (Causal WAM), a practical instantiation. A sparse routing path is added to the action head of a pretrained Cosmos 3 Nano backbone, with a token-level surrogate of the sparsity constraint, while the world-modeling loss is left completely unchanged.

Main Findings

• Without sparsity, action mechanisms over-cover. Proposition 1 shows fitting (o, a) guarantees only |I(J_π̂)_{i,·}|_0 ≥ |I(J_π)_{i,·}|_0 per action dimension, so spurious edges from nuisance coordinates to actions are not excluded.
• Sparsity yields exact structure recovery. Theorem 1 gives I(J_π̂) = I(J_π) P_ρ and I((J_π̂)_{i,·}) = ρ(I_i) for every action dimension i, for a single permutation ρ shared across dimensions.
• The compact state carries only causal information. Theorem 2 gives C_t = ŝ_{t,ρ(I_A)} = h(s_{t,I_A}) with h invertible, so C_t depends only on the action-relevant block and determines it.
• Synthetic recovery of the action-relevant block (n = 20 latents, |I_A| = 4, five seeds). CWAM with sparsity reaches a relevant R² of 0.9957 ± 0.0014 with a nuisance excess Δ_nuis of 0.039 ± 0.093, versus 0.9961 ± 0.0015 / 0.646 ± 0.115 for the dense 20-dimensional readout and 0.9440 ± 0.0448 / 0.316 ± 0.186 for a standard MLP VAE. Four random coordinates recover the action-relevant variables poorly (0.5493 ± 0.2964).
• The selected coordinates are what the head actually uses. Zeroing C_t raises action error by a factor of 1.704 ± 1.398 for CWAM with sparsity and 3.456 ± 3.180 for CWAM without sparsity, while zeroing the complement costs far more (54.65 ± 25.78 and 50.58 ± 28.81 respectively). Random matched shows the opposite pattern (48.067 ± 31.855 selected/full, 6.36 ± 6.42 complement/full), indicating that which coordinates are kept matters more than how many.
• Hidden regime variation does not break recovery. A second hidden index e_t, resampled each step and never shown to the encoder, rescales only nuisance variables in the transition; as its number of values E grows, neither readout loses accuracy on the action-relevant latents, and only the dense readout retains a large nuisance excess (Appendix D, Table 4).
• Closed-loop gains on RoboCasa365. CWAM beats the matched Cosmos 3 baseline in every foundational skill family over 17 tasks with 100 episodes per task; gains are smallest for button pressing and largest for drawer opening and closing, rack sliding, insertion, and pick-and-place, with CWAM's averaged success rate nearly twice the baseline's (Table 5).
• Closed-loop gains on LIBERO. Both policies succeed often over 40 tasks, with CWAM outperforming across all suites and the largest gain on LIBERO-Goal.
• Real-robot gains. On a Franka R3 cube-stacking task with 50 trials per policy, CWAM succeeds in a majority of trials whereas the baseline succeeds in fewer than half.
• Routed tokens alone retain most of the success (Table 2, four best RoboCasa365 tasks, checkpoint fixed at 50k steps). Averages: CWAM 71.3%, C-only 60.3%, Random matched 24.5% — with the same ordering on every task (e.g., pick place sink to counter: 74% / 63% / 28%; open stand mixer head: 69% / 58% / 16%).
• Router attention is object-centric in most frames. Routing mass concentrates on the object being moved and its target in the RoboCasa365 rollout and in most real-robot frames; in some frames it disperses onto the robot arm and the wall, as in the fourth frame of the real-robot sequence.

Methodology in Plain English

The authors first write down a minimal generative story for demonstrations: an image o_t = f(s_t) is an invertible (diffeomorphic) view of an unknown latent state, and the action a_t = π(s_t, η_t) is produced by an unknown mechanism from that state plus noise. Only a subset of the latent coordinates — those π functionally depends on — actually causes the action; everything else is nuisance that only prediction needs.

They then analyze when fitting this model to demonstrations pins down the causal part. Without extra bias, many reparameterizations of the latent space fit equally well, and a flexible action head can route nuisance information into the action while still matching the data — this is Proposition 1's one-sided bound. The fix is to require that the estimated mechanism use no more latent coordinates than necessary (a sparsity constraint on the Jacobian support). Combined with the lower bound, this squeezes the estimated support to exactly the true one up to relabeling (Theorem 1), and a block-triangularity argument on the reparameterization Jacobian shows the selected coordinates cannot mix in nuisance variables (Theorem 2).

For the practical model, the authors take Cosmos 3 Nano, which co-trains video prediction and action generation with a mixture-of-transformers generator under a joint rectified-flow objective, and post-train it on demonstrations with the standard recipe. They add a router after the generator tower: each action hidden state attends over clean context tokens (task tokens and clean vision, excluding future noised tokens) with relation-aware biases on position, modality, and chunk progress (d_r = 512, 8 heads), and the routed signal enters through a small, learned gate. A sparsity loss rewards putting high attention mass on the top-k context tokens, with λ = 0.01 ramped in after 10^3 steps; the total objective is L_WAM + L_sparse, where L_WAM is identical to the baseline's. The compact state C_t is defined as the clean tokens appearing in the top-k routing mass of some action query.

Synthetic experiments test the theory with known latents (n = 20, |I_A| = 4), using an observation-only MLP encoder plus a linear action head with a group penalty on head columns, and held-out linear probes to measure how much C_t reveals about the action-relevant versus nuisance variables. Real experiments compare against an identically post-trained Cosmos 3 without the router and without the sparsity loss, then ablate visual context at inference (routed tokens only vs. an equal number of random tokens) without refitting.

Why This Matters

The work gives a formal reason why adding a bottleneck on the action side of a world model — rather than compressing the whole latent state — should help, and shows that the guarantee survives conditions that break most identifiability results: no interventions, no temporal dependence, arbitrary unobserved regime shifts, and correlation between causal and nuisance factors.

Real-world applications:

• Tabletop manipulation such as cube stacking, pick-and-place, and insertion, where background clutter and table coverings change without affecting the correct motion.
• Articulated-object tasks such as opening drawers, sliding racks, and opening a stand mixer head — the categories where the reported gains are largest.
• Post-training of pretrained robot foundation models, where the modification is a cheap add-on that leaves the existing world-modeling objective untouched.
• Deployment under visual distribution shift, since the router's selection is learned from demonstrations rather than from task or segment labels.

Industry relevance: teams building or fine-tuning video-pretrained robot policies can adopt the router and sparsity loss as a post-training change without redesigning the backbone, and the theory offers a principled justification for sparse action heads rather than an empirical heuristic.

Future Directions

• Component-wise identifiability. The authors name this as the main limitation: the theory identifies the action-relevant block up to an invertible transform and the structure up to a permutation, but does not pin down individual latents — which they note may be unnecessary for separating actions but useful for precise interventions.
• Relaxing the nondegeneracy condition. The span condition in Proposition 1 is an explicit extra assumption on the action mechanism that smoothness and positive state density do not imply (discussed in Appendix B.1), leaving room to weaken it.
• Explaining dispersed routing. Routing mass falls on the manipulated object and target in most frames but drifts to the robot arm and wall in others; understanding and correcting those frames is an open question.
• Extending causal analysis to other WAM problems. The conclusion explicitly frames this as "one of the first steps towards a causal understanding of actions," inviting further work on what else causality can contribute to world action models.

Target Audience

Researchers in causal representation learning and nonlinear ICA who want identifiability results that hold without interventions or temporal dependence; robot learning researchers working on imitation learning and vision-language-action models; and practitioners post-training video-pretrained world models who want a lightweight, theoretically motivated architectural change. Readers need comfort with Jacobians, diffeomorphisms, and support-based identifiability arguments to follow the theory, though the method section and experiments are accessible to applied robot-learning engineers.

Authors’ abstract

World action models (WAMs) learn latent variables from observations and use these to both predict the evolution of the scene and to generate actions. The prediction has to be of the whole scene, whereas only a portion of the scene causes actions. WAMs fitted to demonstrations do not, in general, guarantee that they will generate actions solely from the latent causes of the demonstrated actions (the action-relevant latent variables), as the model could just as well learn to imitate in this way whilst entangling the action-relevant latent variables with the nuisance variables. We therefore separate out what causes action from what prediction depends on. We prove that under a sparsity constraint the latent variables used to generate actions are an invertible function of the action-relevant latent variables alone and the true structure is identified up to a permutation of the latent variables. This requires neither sufficient temporal dependence nor any assumptions on how the unobserved regimes evolve and holds even when the action-relevant latent variables are statistically dependent on the rest. We therefore propose CWAM (Causal WAM), which incorporates the constraint into a pretrained WAM by adding a sparse routing path to its action head, with the world modeling objective remaining unchanged. Experiments on both synthetic data and real-world robots validate the theory as well as showcase the improved efficiency in closed-loop success.