PaperlayerView exact revision on arXiv

Source-linked AI summary

CausalWAM: Learning the Causes of Actions in World Action Models

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

arXiv:2610.42096v1cs.AIcs.ROstat.ML

TL;DR

World action models must predict an entire scene, but only some latent variables cause actions, and demonstrations do not by themselves guarantee that action generation excludes nuisance variables. The paper imposes sparsity on the action pathway and proves identification of the action-relevant latent block under its conditions. CWAM applies this constraint to a pretrained model, and experiments report improved aggregate closed-loop success over Cosmos 3 without the routing path.

  • Problem

    Demonstration-fitted world action models need not generate actions solely from their true latent causes, because observationally equivalent models can mix action-relevant and nuisance variables.

    Source: §1 · p. 1

  • Method

    CWAM adds a sparse routing path to the action head and applies a sparsity penalty there while leaving the world-modeling objective unchanged.

    Source: §4 · p. 5

  • Results

    The routing path improves aggregate closed-loop success over Cosmos 3 without it on RoboCasa365, LIBERO, and a real Franka R3 robot.

    Source: §1 · p. 2

  • Takeaways & Limitations

    CWAM assists task completion in simulation and on a real robot, while identifying the action-relevant block without component-wise identification of individual latents.

    Source: §6 · p. 10

    The theory does not provide component-wise identifiability of individual latents, which may be useful for precise interventions.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

World action models must predict the whole scene, although only some latent variables cause actions. This work asks whether demonstrations alone can separate those action causes from nuisance variables and shows that sparsity on the action pathway identifies them under stated conditions.

1 INTRODUCTION

  • World action models learn latent variables for both scene prediction and action generation, so their latents can include nuisance factors that change images without affecting motion.The stacking example’s action depends on block positions and whether the gripper holds the upper block, while tablecloth and background clutter are nuisance variables.Source: §1 · p. 1
  • Fitting demonstrations alone does not ensure that actions depend only on their true latent causes, because equivalent models can mix action-relevant and nuisance variables.A flexible action head may still fit demonstrations while using variables that contain information about scene elements, such as the tablecloth, that do not cause the motion.Source: §1 · p. 1
  • Existing identifiability approaches use auxiliary variables, temporal dependence, latent-dynamics conditions, or known-target interventions, which may not hold or be feasible in robotics demonstrations.The paper contrasts these requirements with identifying only the subset of latents relevant to actions, using actions themselves as the indicator of relevance.Source: §1 · pp. 1–2
  • The paper frames its central question as whether demonstrations alone can separate what causes actions from what prediction needs.Source: §1 · p. 2
  • Sparsity on the action pathway identifies the action-to-latent dependency structure up to permutation and makes the action-generating variables an invertible function of the true action causes alone.The prediction pathway remains free to use every latent variable; the guarantee requires neither interventions nor sufficient temporal dependence and allows dependence between relevant and nuisance variables.Source: §1 · p. 2

2 PROBLEM SETUP

The setup models demonstrations as observations and actions generated from latent states through an invertible observation map and an action mechanism. It focuses identification on the latent coordinates used for actions, while the full state remains available for prediction.

2 PROBLEM SETUP

  • Demonstration data are observation-action sequences, with relevance defined by executed actions and no task, skill, or segment labels assumed.Source: §2 · p. 2
  • The observational process maps latent state st to observation ot through f and maps st, with exogenous noise ηt, to action at through π.The observation map is a diffeomorphism onto its image; the action mechanism also applies to an H-step chunk when the policy emits one.Source: §2 · p. 2
  • Action-relevant variables are those on which π functionally depends, and consecutive steps may be disconnected or i.i.d.Because f is invertible, observations depend on every latent variable, whereas actions depend only on action-relevant ones; a fitted action mechanism may include additional variables.Source: §2 · p. 3
  • The action-relevant subvector is defined by the latent coordinates in the functional support of π’s Jacobian, distinguishing functional relevance from correlation.The support records which latent coordinates can change each action dimension; IA is the union of those coordinates across action dimensions.Source: §2 · p. 3
  • Under observational equivalence, the estimated state determines every latent variable for prediction, while identification asks whether the compact action state isolates the action-relevant block.The compact state is formed from estimated latent coordinates selected by the action mechanism; actions are not inputs to the encoder at deployment.Source: §2 · p. 3
  • Observational equivalence alone permits an arbitrary compatible diffeomorphic reparameterization, so it does not identify the full latent state or a unique encoder.The stated analysis assumes continuously differentiable maps, matching observation-map images, and a latent-state density positive over R^n, and uses a pathwise coupling with shared noise.Source: §2 · p. 3

3 THEORY

Without support sparsity, observationally equivalent fitted mechanisms may use extra latent coordinates, leaving nuisance variables entangled with the action pathway. Under the stated conditions, sparsity identifies the action-dependency pattern up to permutation and makes the compact state an invertible re-encoding of the action-relevant block.

3 THEORY

  • The guarantees apply within each observation-action step and require no cross-time structure; consecutive steps may be disconnected.They concern one observation-action pair at a time.Source: §3 · p. 3

3.1 IDENTIFIABILITY WITHOUT SPARSITY

  • Without support sparsity, fitting may over-cover action dependencies, so nuisance coordinates can enter the action pathway and neither the true support pattern nor the action-relevant block is guaranteed to be recovered.The unconstrained result gives only a one-sided support bound.Source: §3.1 · p. 4
  • The no-sparsity bound relies on a span condition: each action dimension’s Jacobian row must span the coordinate subspace of its true parents.Smoothness and positive state density alone do not imply this separate assumption on the action mechanism.Source: §3.1 · p. 4

3.2 IDENTIFIABILITY UNDER SUPPORT SPARSITY

  • With support sparsity, the fitted mechanism’s dependency pattern matches the true pattern up to a common permutation of latent coordinates.Sparsity turns the lower bound on fitted support into an exact support match.Source: §3.2 · p. 4
  • The compact state is an invertible re-encoding of the true action-relevant block and does not mix nuisance coordinates into its estimated coordinates.This functional guarantee holds regardless of statistical dependence between the action-relevant block and nuisance coordinates.Source: §3.2 · p. 5

4 ESTIMATION METHOD

CWAM applies support sparsity to the action mechanism while retaining the pretrained world-modeling objective. It routes selected clean context tokens to the action head, with the default policy retaining its dense path.

4 ESTIMATION METHOD

  • CWAM applies a sparsity penalty to the action mechanism, using the theory’s support-sparsity principle to identify action-relevant latents and their dependency pattern.The deployed interface can emit an H-step action chunk.Source: §4 · p. 5
  • CWAM keeps Cosmos 3’s joint rectified-flow video-and-action training objective unchanged and applies the added sparse loss only to the router readout.The matched Cosmos 3 baseline uses the same joint post-training loss on the same demonstration windows.Source: §4 · pp. 5–6
  • Each action hidden state routes over clean task and vision context, and the top-k routed tokens form the compact state.A learned gate adds the routed signal to the hidden state, and training regularizes the router’s top-k attention mass.Source: §4 · pp. 5–6, Figure 1 · p. 6
  • The default policy retains the dense visual-to-action path alongside the added sparse routing path.The figure depicts the sparse path as routing an estimate of the action-relevant latent block to the action head.Source: Figure 1 · p. 6, §4 · p. 6

5 EXPERIMENTS

Across synthetic and robotic evaluations, CWAM’s sparse action path recovers action-relevant information and improves closed-loop success over a matched baseline. Routed tokens retain much of CWAM’s success on its strongest tasks, while equal-count random tokens cause a larger drop.

5.1 THE COMPACT STATE RECOVERS THE ACTION-RELEVANT BLOCK

  • On synthetic data, CWAM’s selected compact state decodes action-relevant variables about as well as the dense readout while showing near-zero excess decoding of nuisance variables.Dense and VAE-selected readouts decode nuisance variables beyond the true action-relevant block, while four random coordinates recover relevant variables poorly.Source: §5.1 · p. 7
  • Zeroing CWAM’s compact state substantially raises action error, whereas zeroing its complement costs much less, indicating that the fitted head relies on the selected coordinates.Random-matched coordinates show the opposite pattern, so coordinate selection matters beyond the number of retained coordinates.Source: §5.1 · p. 7
  • As the number of hidden, step-resampled nuisance regimes grows, action-relevant decoding remains accurate, while only the dense readout retains large nuisance excess.The action mechanism is held fixed in this test.Source: §5.1 · p. 7

5.2 THE SPARSE ROUTING PATH IMPROVES CLOSED-LOOP SUCCESS

  • CWAM improves success over matched Cosmos 3 in every RoboCasa365 skill family, every LIBERO suite, and the real-robot task, with nearly twice the baseline’s average success across RoboCasa365 tasks.On the real robot, CWAM succeeds in a majority of trials, while the baseline succeeds in fewer than half.Source: §5.2 · pp. 7–8, Figure 2 · p. 8

5.3 ROUTED TOKENS ALONE KEEP MOST OF CWAM’S SUCCESS ON ITS BEST TASKS

  • On four RoboCasa365 tasks where CWAM performs best, routed tokens retain most default-policy success, while an equal number of random tokens causes a substantially larger success drop.The checkpoint is fixed during both visual-context interventions.Source: Table 2 · p. 8, §5.3 · pp. 8–9

6 CONCLUSION

The paper formalizes learning action causes from observational decompositions and uses sparsity to separate the action pathway from the full state needed for prediction. Experiments support the theory and show task-completion benefits, while individual latent variables are not identified component by component.

6 CONCLUSION

  • Sparsity regularization identifies the action causes and latent-to-action structure up to permutation, while the full estimated state remains available to describe the scene for prediction.Without sparsity, a fitted action mechanism can include action-irrelevant latent variables.Source: §6 · p. 9
  • The identification results require neither interventions, temporal structure, nor supervision, and synthetic experiments agree with the theory while CWAM assists task completion in simulation and on a real robot.The conclusion presents the work as an early step toward a causal understanding of actions in WAMs.Source: §6 · p. 10
  • The theory does not identify individual latent variables component by component, which could be useful for precise interventions.The authors note that component-wise identification may not be required to separate actions from nuisance variables.Source: §6 · p. 10

AI STATEMENT

The authors report AI assistance for coding, language editing, and figure plotting, and provide a proof of Proposition 1 based on equivalent observational models and support arguments.

AI STATEMENT

  • Generative AI tools assisted with coding, language editing, and figure plotting; the authors reviewed the work and take responsibility for the final content.Source: AI STATEMENT · p. 10

A.1 PROOF OF PROPOSITION 1

  • The Proposition 1 proof begins by reparameterizing observationally equivalent models and differentiating their coupled action mechanisms.The proof assumes linearly independent action-Jacobian rows at selected states and regularity conditions.Source: §A.1 · p. 12
  • Using row independence and support constraints, the proof derives a common permutation that relates the latent-coordinate supports across action dimensions.The same permutation is chosen for all action dimensions.Source: §A.1 · p. 13

A.2 PROOF OF THEOREM 1

Theorem 1 shows that under support sparsity, the estimated latent-to-action dependency pattern matches the true pattern up to one shared permutation of latent coordinates.

A.2 PROOF OF THEOREM 1

  • Under the theorem’s hypotheses and support sparsity, one permutation aligns the estimated latent-to-action support with the true support for every action dimension.Source: §A.2 · p. 14
  • The proof applies the earlier support argument to each action dimension using the same permutation, yielding a shared coordinate correspondence.Source: §A.2 · p. 14
  • Sparsity forces the per-dimension support bounds to be equalities, completing the support-level equivalence.Source: §A.2 · p. 14

A.3 PROOF OF THEOREM 2

Under the stated hypotheses, the estimated action-relevant variables recover the true action-relevant block through an invertible map, with variables matched up to a permutation. The proof uses derivative-support constraints, while the conditions require sufficient variation in action-mechanism derivatives and exclude simultaneous support cancellation.

A.3 PROOF OF THEOREM 2

  • Theorem 2 identifies the estimated action-relevant variables as an invertible function of the true action-relevant block, with the same permutation as Theorem 1.The proof establishes a zero upper-right block in the reparameterization Jacobian, yielding block-wise functional dependence and an inverse on the image.Source: §A.3 · pp. 14–15

B.1 DETAILED DISCUSSION ON MAIN CONDITIONS

  • The spanning condition requires action-mechanism Jacobian rows to vary enough to span each action dimension’s true parent set; smoothness and positive state density alone do not ensure it.For π_i(s) = s_1 + s_2, the rows are all (1, 1, 0, …), so the condition fails regardless of the state distribution.Source: §B.1 · p. 16
  • The support-inclusion condition must hold at the spanning states, excluding measure-zero cases where partial derivatives cancel at all of them.The stated assumptions also require that reparameterization cannot hide a true dependency through such cancellation.Source: §B.1 · p. 16

C EVALUATION PROTOCOL

The protocol tests block recovery in a synthetic process with known action-relevant variables and nuisance coupling, then varies hidden environments that alter nuisance transitions but not the action mechanism. It evaluates both action-relevant and nuisance-variable decoding with held-out probes.

C EVALUATION PROTOCOL

  • The synthetic generator uses 20 latent dimensions, with the first four action-relevant, and hidden regimes that modulate the action mechanism and transition while nuisance variables couple to relevant ones.A second environment index rescales nuisance transition coordinates, leaves the action mechanism fixed, and is not provided to the encoder.Source: C EVALUATION PROTOCOL · p. 16

D HIDDEN ENVIRONMENTS (SYNTHETIC)

  • As hidden environments are added, CWAM’s action-relevant decoding does not degrade relative to the dense readout, while the dense readout alone retains excess nuisance decoding.Each environment rescales nuisance transition variables; the action mechanism remains fixed and the environment index is never an encoder input.Source: Table 4 · p. 17, D HIDDEN ENVIRONMENTS (SYNTHETIC) · p. 17

C EVALUATION PROTOCOL

  • Block recovery is measured with held-out ridge-probe R2 for action-relevant and nuisance variables, with the true action-relevant block serving as the nuisance-decoding oracle.The protocol uses disjoint calibration and reporting episodes; the oracle reference measures nuisance predictability already present in the true relevant block.Source: C EVALUATION PROTOCOL · pp. 16–17

E ROBOCASA365 PER-TASK SUCCESS

On RoboCasa365, CWAM improves success across the task set, with the table reporting per-task rates and the macro average. The robot comparison uses a matched Cosmos 3 baseline trained with the same recipe and optimization, differing by the router and sparsity loss.

E ROBOCASA365 PER-TASK SUCCESS

  • CWAM raises success on 13 of 17 RoboCasa365 tasks and increases macro-average success from 0.24 to 0.43.Table 5 reports success rates over 100 episodes per task, grouped by skill family; Figure 2(a) averages task success within each family.Source: E ROBOCASA365 PER-TASK SUCCESS · p. 17, Table 5 · p. 18

F COSMOS 3 POST-TRAINING SETUP

F.1 TRAINING STAGES AND SERVED CHECKPOINT

The appendix documents the matched Cosmos 3 post-training setup and evaluation protocols, then specifies the shared model objective and CWAM’s added router and sparsity loss.

F.1 TRAINING STAGES AND SERVED CHECKPOINT

  • CWAM starts from the public Cosmos 3-Nano checkpoint, trains with a freshly initialized router for 10,000 steps, then continues to 50,000 steps; the matched baseline follows the same stages without the router.The second stage uses a flat learning-rate multiplier of 0.3871, and training runs on one 8×H200 node.Source: §F.1 · pp. 18–19

F.2 DEMONSTRATION DATA RECIPE (ROBOCASA365)

  • RoboCasa365 training uses 20 Hz policy windows of 16 frames, concatenated third-person and wrist views, language prompts, and frame-wise relative-rotation actions.The leading frame is observed and later video tokens are noised under the rectified-flow objective.Source: §F.2 · p. 19

F.3 CLOSED-LOOP TASK PROTOCOLS

  • Closed-loop evaluation covers 17 RoboCasa365 tasks (100 episodes each, 400-step cap, 12-step open-loop horizon) and 40 LIBERO tasks across four suites, with per-suite success rates reported.Source: §F.3 · p. 19

F.4 SHARED OPTIMIZATION HYPERPARAMETERS

  • Table 6 lists optimization hyperparameters shared by the baseline and CWAM, with router-specific settings documented separately.Source: §F.4 · p. 19, Table 6 · p. 19

G COSMOS 3 BASE MODEL AND RECTIFIED-FLOW OBJECTIVE

  • The Cosmos 3-Nano model jointly trains on noised future video tokens and action chunks conditioned on clean history and language, using the same rectified-flow loss for baseline and CWAM.CWAM adds a sparsity loss only on the router readout.Source: G COSMOS 3 BASE MODEL AND… · p. 20

H ACTION-RELEVANCE ROUTER

  • The residual action-relevance router attends over clean task and vision context, selects a top-32 support, and receives a sparsity penalty that ramps during training.The dense visual-to-action path remains active during training and default evaluation; C-only and random-matched readouts are evaluation-only interventions without refitting.Source: H ACTION-RELEVANCE ROUTER · p. 20
Loading 2610.42096v1…