CausalWAM: A Causal World Action Model for Generalizable Robot Control

Deploying robots in complex real-world environments requires more than predicting future observations or directly mapping visual inputs to actions. Robots must understand how their actions influence the physical world, identify task-relevant causal variables, and reason about how these variables evolve through interaction. Achieving robust and generalizable robotic intelligence therefore requires a unified model that connects causal understanding, future prediction, and action generation within a coherent framework.

We introduce the Causal World Action Model (CausalWAM), a unified framework that integrates causal world modeling with robot policy learning. Built upon CausalWM[1], CausalWAM explicitly models the structural knowledge underlying physical interactions and extends future prediction to action generation. Rather than directly predicting future visual observations or robot actions, CausalWAM organizes its reasoning around Causal Chain-of-Thought, progressively modeling intermediate physical variables and task-relevant state transitions. Through a unified world-action architecture, CausalWAM leverages these representations to predict both how the physical world will evolve and which actions should be taken to achieve desired outcomes. This allows world prediction and action generation to mutually benefit from a shared understanding of physical dynamics.

This design enables three key capabilities. Causal Reasoning: CausalWAM explicitly models intermediate physical variables and their dependencies, enabling causal reasoning about how actions influence future world states. Predictive Control: by integrating future prediction with action generation, CausalWAM can anticipate physical consequences and generate actions informed by expected state transitions. Generalization: its shared causal representations facilitate knowledge transfer across tasks, environments, and interaction scenarios, reducing reliance on purely reactive behavior and task-specific action patterns.

01One Model for Physical Prediction and Control

CausalWAM is built on CausalWM using a Mixture-of-Transformers (MoT) architecture[2], where we introduce an additional Action Expert alongside the original world-modeling expert to jointly model future states and robot actions. Based on this architecture, we make three key modifications for embodied control: Unified State Encoding, which integrates visual observations, proprioception, and CoT context into a shared DiT context; Cross-Embodiment Unified Action, which standardizes action representations and sampling frequencies across heterogeneous robots; and Long-Horizon Memory, which maintains a bounded interaction history to support iterative prediction and replanning over extended horizons.

Unified State Encoding

CausalWAM encodes visual observations, proprioceptive signals, and chain-of-thought (CoT) tokens into a unified representation, which is jointly placed in the context of the DiT backbone. Visual inputs describe scene and object configurations, proprioception captures robot configuration and motion, and CoT tokens provide intermediate reasoning about task progress and relevant physical factors. Built on the CausalWM [1], this unified context provides the shared conditioning space for subsequent state prediction and action generation, enabling the model to jointly model environment dynamics and robot behavior. During execution, updated observations and recent interaction history are continuously re-encoded to condition subsequent predictions, allowing the model to adapt its actions through iterative prediction and replanning.

Cross-Embodiment Unified Action

CausalWAM unifies robot actions using two complementary representations: an end-effector representation and a camera-grounded representation. The former provides an embodiment-agnostic action space across robots with different kinematics, while the latter expresses motion in the camera frame to align actions with visual observations. The camera-grounded representation is used as an input condition, whereas the end-effector representation serves as the unified action prediction objective. To further unify heterogeneous datasets, we normalize the sampling frequency of action trajectories by resampling them onto a shared temporal grid. This ensures that action targets correspond to consistent time intervals and prediction horizons across datasets with different control rates, enabling stable joint training across embodiments.

Long-Horizon Memory

CausalWAM maintains a bounded rolling memory of recent interaction history within each episode. The memory stores preceding observations, state transitions, and interaction outcomes, providing temporal context beyond the current observation for subsequent prediction and control. At each replanning step, newly acquired observations are appended to the memory while older context is discarded once the window limit is reached. The memory is initialized empty and reset between episodes, allowing execution feedback to continuously influence future predictions without updating model parameters.

02Training Across Robots and Views

CausalWAM is trained on a diverse mixture of robot and egocentric videos, together with action-annotated trajectories collected from real-world robots, UMI demonstrations, and simulation. These sources provide broad coverage of robot embodiments, manipulation tasks, environments, and visual conditions.

Preparing the Training Data

To ensure high-quality supervision, we apply a comprehensive data curation pipeline. We identify recording failures, anomalous motion, temporal misalignment, and inconsistencies across visual observations, proprioceptive states, actions, and task annotations. Recoverable issues are corrected, while unreliable segments are discarded. We retain only clean, temporally aligned training windows with complete prediction horizons, resulting in 15.8k trainable hours from 19.5k raw hours. [4-13] Additional load-time validation ensures valid wrist-camera action representations and frame rates compatible with the target video sampling rate.

Training Pipeline

CausalWAM is pre-trained in three stages: single-view video learning, multi-view video learning, and joint video–action learning. The first two stages adapt the video backbone to robot-centric scenes and synchronized multi-view observations, enabling the model to capture scene dynamics from visual context. The third stage introduces the Action Expert and jointly optimizes future video and action prediction with flow matching, explicitly coupling visual dynamics with robot control. During supervised fine-tuning, CausalWAM learns to exploit recent execution history as additional conditioning without introducing extra modules or objectives. We vary the amount of historical context and randomly mask historical states and actions while retaining visual observations, encouraging robust use of interaction history.

03Testing Generalization on Robot Benchmarks

LIBERO-Plus

We fine-tune CausalWAM on demonstrations from all four standard LIBERO suites[14] and evaluate it on LIBERO-Plus[15] without additional fine-tuning, following the official protocol. The results below show the success rate for each perturbation dimension and the overall performance, together with representative VLA and WAM baselines.

RoboTwin 2.0

We evaluate CausalWAM on RoboTwin 2.0[16] under the Clean2Random protocol, fine-tuning exclusively on clean demonstrations and evaluating separately in clean and randomized environments, following the official task definitions. These two settings assess in-distribution performance and out-of-distribution generalization, respectively. The results below show success rates under both conditions and their average, alongside representative VLA and WAM baselines.

04From Policy to Robot Agent

Agent Framework

We integrate CausalWAM into an agent harness with a VLM planner and four tools: Observe for querying object states and task progress, Move for cuRobo[19]-based repositioning, and Skill/Delegate for invoking CausalWAM with short commands or full task instructions. At each turn, the planner receives camera observations and robot state, while CausalWAM executes manipulation actions with visual feedback. Although fine-tuned only on full task instructions, CausalWAM follows short skill commands zero-shot, allowing the planner to compose learned behaviors into longer-horizon tasks without agent-specific training.

RoboDojo Performance

We evaluate all 54 RoboDojo[17] tasks on layouts 0–4 with seed 0, for 270 cells in total, using the official per-task step limits and the official simulator image. The planner is capped at 240 turns per episode. We mount this harness on a CausalWAM fine-tuned RoboDojo checkpoint.

On make_toast, delegating execution to CausalWAM reduced mean planner API cost per attempt from $11.04 to $0.03, while success increased from 20% to 50%. These costs exclude local CausalWAM inference. Handing execution to CausalWAM cut the cost of the task by more than 99% while raising its success rate.

The shared-cell comparison uses RoboDawn[18] as its baseline.

05Emergent Zero-Shot Generalization

We evaluate CausalWAM’s zero-shot generalization on RoboLab[20] by directly deploying the pre-trained checkpoint without any additional fine-tuning. This evaluation probes whether the manipulation capabilities acquired during pre-training transfer to novel tasks and scenes in high-fidelity simulation.

The accompanying videos showcase CausalWAM’s policy execution and future-video predictions across different scenarios, providing a qualitative view of its ability to generate actions and anticipate the resulting scene changes.

@misc{shao2026causalwam,
  title = {CausalWAM: A Causal World Action Model for
           Generalizable Robot Control},
  author = {Yifei Shao and Junbo Huang and Ziqiao Xi and Fang Nan and
            Ruobing Han and Yuchen Yan and Lingjun Mao and Jinglin Cao and
            Wenpeng Xu and Zijun Zhang and Lukun He and Ziming Xu and
            Shuang Liang and Kun Zhou and Biwei Huang},
  year = {2026}
}

References