CausalWAM: Learning the Causes of Actions in World Action Models

Juntao Cheng; Jinzhou Tang; Fan Feng; Henrik I. Christensen; Yujia Zheng

cs.AIOct 30, 2026

CausalWAM: Learning the Causes of Actions in World Action Models

Authors: Juntao Cheng, Jinzhou Tang, Fan Feng, Henrik I. Christensen, Yujia Zheng

Organizations: UIUC · UCSD

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.

Figures & tables

Figure 1: The action depends only on a block of the latent state, and CWAM routes an estimate of this block to the action head through a sparse path. (a) In the world action process equation 1 this block is

s_{t,I_{\mathrm{A}}}

, and the learner

(\hat{f},\hat{\pi})

estimates it from observations as the compact state

C_{t}

. (b) CWAM keeps the joint rectified-flow training of Cosmos 3 and adds a sparse router over clean context tokens, whose top-

k

routed tokens form

C_{t}

. Gray dashed arrows mark the dense path to the action head, which the default policy keeps.

Method	$\|C_{t}\|$	Relevant $R^{2}\uparrow$	$\Delta_{\mathrm{nuis}}\downarrow$	Selected/full	Complement/full
Dense (20D readout)	20	0.9961 $\pm$ 0.0015	0.646 $\pm$ 0.115	n/a	n/a
Random matched (4D)	4	0.5493 $\pm$ 0.2964	0.392 $\pm$ 0.144	48.067 $\pm$ 31.855	6.36 $\pm$ 6.42
CWAM without sparsity	4	0.9876 $\pm$ 0.0164	-0.008 $\pm$ 0.006	3.456 $\pm$ 3.180	50.58 $\pm$ 28.81
Standard MLP VAE	4	0.9440 $\pm$ 0.0448	0.316 $\pm$ 0.186	3.259 $\pm$ 2.211	51.13 $\pm$ 28.54
CWAM with sparsity	4	0.9957 $\pm$ 0.0014	0.039 $\pm$ 0.093	1.704 $\pm$ 1.398	54.65 $\pm$ 25.78

Table 1: The coordinates selected by CWAM with sparsity decode the action-relevant latent variables as well as the dense readout, and the nuisance variables only about as well as the true block. Zeroing them breaks the fitted head. Held-out linear probes give the relevant

R^{2}

and

\Delta_{\mathrm{nuis}}

, the nuisance

R^{2}

C_{t}

minus that of

s_{t,I_{\mathrm{A}}}

. Selected/full (Complement/full) is the held-out action MSE of the fitted head, not refitted, after zeroing the coordinates outside (inside)

C_{t}

, relative to the full

\hat{s}_{t}

. Mean

\pm

population std over five seeds, with ratios computed per seed.

Figure 2: CWAM improves closed-loop success over the matched Cosmos 3 baseline in every RoboCasa365 skill family, every LIBERO suite, and on the real robot. (a) RoboCasa365, mean success per skill family over $17$ tasks (tasks in Table 5 ). (b) LIBERO, success per suite over $40$ tasks. (c) Franka stack cube, $50$ trials per policy. Panels (a) and (b) use $100$ episodes per task. The baseline (Cosmos 3 in the legend) is post-trained on the same data and recipe, without the router and $\mathcal{L}_{\mathrm{sparse}}$ .

Figure 3: In most frames, the routing mass falls on the object being moved and its target. (Top) A RoboCasa365 rollout. (Bottom) The real robot stacking the red cube on the green cube. Rows show the wrist and agent (third-person) views, each followed by the same view with the routing mass overlaid. Warmer colors mark more of the top- $k$ routing mass over clean visual tokens.

Task	CWAM	C-only	Random matched
Pick place sink to counter	74%	63%	28%
Pick place toaster to counter	72%	61%	21%
Open drawer	70%	59%	33%
Open stand mixer head	69%	58%	16%
Average (4 tasks)	71.3%	60.3%	24.5%

Table 2: On the four RoboCasa365 tasks where CWAM succeeds most, keeping only the routed tokens costs far less success than keeping as many random tokens. The CWAM checkpoint is fixed at

50\mathrm{k}

steps and run on the episodes of the main evaluation. The action queries see the full visual context (CWAM), only the routed tokens that form

C_{t}

(C-only), or as many visual tokens drawn uniformly at random from the clean context (random matched). Neither intervention refits the policy.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Meaning
$o\in\mathbb{R}^{d_{o}}$	Observation
$s\in\mathbb{R}^{n}$	Latent state
$a\in\mathbb{R}^{d_{a}}$	Action
$f$	Observation map $s\mapsto o$ , diffeomorphism onto its image
$\hat{f}$	Learner observation map $\hat{s}\mapsto o$ , diffeomorphism onto its image; its inverse $\hat{f}^{-1}$ is the encoder
$\pi$	Action mechanism $(s,\eta)\mapsto a$ in equation 1b

Appendix

Table 3: Summary of notation for the identifiability proofs.

	Dense (20D readout)		CWAM with sparsity (4D)
$E$	Relevant $R^{2}\uparrow$	$\Delta_{\mathrm{nuis}}\downarrow$	Relevant $R^{2}\uparrow$	$\Delta_{\mathrm{nuis}}\downarrow$
1	0.9961 $\pm$ 0.0015	0.646 $\pm$ 0.115	0.9957 $\pm$ 0.0014	0.039 $\pm$ 0.093
2	0.9961 $\pm$ 0.0016	0.640 $\pm$ 0.110	0.9961 $\pm$ 0.0016	-0.004 $\pm$ 0.003
3	0.9961 $\pm$ 0.0017	0.677 $\pm$ 0.109	0.9961 $\pm$ 0.0017	-0.004 $\pm$ 0.003
5	0.9961 $\pm$ 0.0017	0.683 $\pm$ 0.114	0.9961 $\pm$ 0.0017	-0.003 $\pm$ 0.002

Appendix

Table 4: Block recovery by CWAM with sparsity does not degrade as hidden environments are added. At every

E

it decodes the action-relevant block as well as the dense readout, which alone decodes nuisance variables beyond the true block. Each hidden environment rescales the nuisance variables in the transition, with the action mechanism fixed across

E

, and is never an encoder input. Metrics as in Table 1 , whose fits give the row

E=1

. Mean

\pm

population std over five seeds.

Foundational skill family	Task	Baseline	CWAM
Pick and place	PickPlaceCtrToCab	.31	.43
	PickPlaceCtrToStove	.22	.45
	PickPlaceDrawerToCtr	.15	.07
	PickPlaceSinkToCtr	.24	.74
	PickPlaceToasterToCtr	.28	.72
Opening / closing doors	CloseFridge	.14	.61

Appendix

Table 5: CWAM raises success on

13

of the

17

RoboCasa365 tasks and the macro average from

0.24

0.43

. Each cell is the success rate over

100

episodes. Tasks are grouped by skill family, and each point in Figure 2 (a) averages the tasks of one family. Baseline and CWAM are as in Figure 2 .

Setting	Value
Optimizer	FusedAdamW ( $\beta_{1}{=}0.9$ , $\beta_{2}{=}0.99$ , ${\epsilon}{=}10^{-8}$ )
Base learning rate	$5\times 10^{-5}$
LR multipliers	$5.0$ on action2llm , llm2action , action_modality_embed
Weight decay	$0.05$
Gradient clipping	$1.0$ (force finite)
Precision	bf16

Appendix

Table 6: Shared Cosmos 3 RoboCasa365 post-training hyperparameters (baseline and CWAM).

World Action Models

World Action Models connect predictions of how scenes change under action to executable policies for robots, vehicles, and other embodied agents.

See all 226 papers

What Matters in Designing World Action Models: An Empirical Study
Sep 21, 2026Chao Tang, Haoqing Wang, Zilang Cen +9World Model Learning Causal World Models
Samsung Robotics eXperience · Samsung R&D Institute China–Beijing · Zhongguancun Academy +2
World-Action Models for Robot Learning and Control: A Survey
Sep 13, 2026Zuxing Lu, Hongjia Zhai, Guanzhi Wang +13World Action Models Action-Conditioned World Models
MBZUAI · Caltech · Amazon FAR +4
CAG-WAM: Causal Action-Gain World Action Models for Robot Manipulation
Oct 10, 2026Peizheng Guo, Cheng Zhang, Xingyu Zhang +2Action-Conditioned World Models Robotic Manipulation
National Key Laboratory of Space Integrated Information System, Institute of Software Chinese Academy of Sciences, Beijing, China · University of Chinese Academy of Sciences, Beijing, China · Roban

Sep 21, 2026J/K move · Enter open · S save

What Matters in Designing World Action Models: An Empirical Study

Chao Tang, Haoqing Wang, Zilang Cen, Weishi Mi+8

Samsung Robotics eXperience · Samsung R&D Institute China–Beijing · Zhongguancun Academy+2

World Model Learning Causal World Models World Action Models Latent World Models World Model Evaluation

Open paper AskPDF arXiv

World Action Models (WAMs) have emerged as a promising paradigm for generalizable robot control. Despite the growing number of WAM systems, existing works often introduce unified systems that bundle together multiple design choices, such as architecture and training strategy, making it difficult to isolate individual contributions and systematically compare alternative designs. In this work, we present a controlled study that disentangles these design choices and analyzes not only their empirical effects, but also how and why they shape WAMs. More specifically, we focus on three fundamental questions in building WAMs: (1) what causal structure should govern the interaction between world modeling and action generation? (2) in which latent space should world modeling be performed? and (3) how do different world-action modeling objectives affect model behavior and performance? Through structurally controlled experiments on three representative benchmarks, RoboCasa-GR1, LIBERO, and LIBERO-Plus, we systematically compare six causal structures, eight latent representations, and four training objectives, covering popular design choices in existing WAMs. We further validate our key findings on real-robot data from the DROID dataset. We hope to provide a systematic understanding of how core design choices affect world-action modeling and what principles can guide the development of future WAM systems.

Abstract

World Action Models

Explore similar work

What Matters in Designing World Action Models: An Empirical Study

World-Action Models for Robot Learning and Control: A Survey

CAG-WAM: Causal Action-Gain World Action Models for Robot Manipulation