CausalWM: Causal Chain-of-Thought Reasoning for Embodied World Model

Method: Causal CoT

CausalWM uses observable physical variables to decompose future prediction. Optical flow represents motion, depth captures geometry, and pointmaps represent three-dimensional spatial relationships. The model generates intermediate variables sequentially according to predefined causal dependencies:

Optical Flow → Geometry / Pointmap → Future RGB

Each prediction is inserted back into the context to condition the next step. Intermediate variables therefore serve as both prediction targets and inputs to subsequent reasoning, rather than only auxiliary supervision or independent outputs.

The analogy to language-model chain-of-thought is the decomposition into intermediate steps. Here, those steps are visual physical representations: motion informs the prediction of geometry, and the resulting context informs future appearance. The objective is to make the intermediate structure useful for downstream generation, rather than merely reconstructing several modalities in parallel.

A causal mask prevents later variables from leaking into earlier reasoning steps during training. Training and inference use the same order, so intermediate predictions actively participate in future video generation at inference time.

For example, an earlier motion-prediction step must not obtain its answer from a later geometry or RGB target. Blocking that information preserves the intended direction of conditioning and keeps the training setup consistent with the sequential prediction process used at inference.

Data and Training

To train CausalWM, Aether AI collected approximately 31,000 hours of embodied video across 20 dataset families, spanning human egocentric activity, real-robot demonstrations, and simulated manipulation.

Data preparation combines length-based and motion-based filtering with action-quality checks. Accepted videos are segmented into events and recaptioned to align each clip with its instruction. The diagram below summarizes the source mixture and processing pipeline.

Training proceeds from a broad video-generation prior to explicit physical reasoning, then optimizes the reasoning process through the quality of its final outcomes.

Stage 1: Pixel-Level Pre-training

Built on LTX-2.3-22B, the model first learns language-conditioned future prediction. CD-LAM extracts a unified latent action representation from video, enabling action-conditioned dynamics learning when action labels are scarce or robot action spaces differ. A multi-view version supports prediction across multiple camera views.

This stage establishes a pixel-level world generation prior before introducing explicit intermediate reasoning. It combines the ability to generate plausible future observations with conditioning on language or inferred actions, providing the foundation for the subsequent CoT training stage.

Stage 2: Causal CoT Mid-training

Optical flow, depth, and pointmaps are extracted automatically from raw video, allowing intermediate supervision to scale with the data. The model learns the fixed sequence Flow → Pointmap → Future RGB, reusing earlier predictions as context under the causal mask described above.

These variables expose complementary aspects of the same interaction: flow captures movement, while depth and pointmaps describe spatial structure. Automatic extraction makes it possible to train on intermediate physical information without manually annotating a reasoning trajectory for every video.

Stage 3: Multi-Objective RL Post-training

Video generation admits many possible futures. Supervised learning alone does not directly optimize every property that matters for a world model, including physical consistency, temporal quality, and task completion. Reinforcement learning uses the quality of generated outcomes to refine the prediction process.

The model samples multiple possible futures from the same context and applies group-based optimization using rewards for physical consistency, visual quality, and task completion. Final video quality also guides exploration of better Causal CoT reasoning trajectories: effective intermediate predictions are reinforced, while implausible paths are discouraged. This makes the reasoning process itself an optimization target.

In-Context Control

Repeatedly conditioning generation on intermediate predictions teaches CausalWM to use in-context visual features. This interface can also accept control variables outside the original CoT design, provided they can be represented visually.

For robotic trajectory control:

1. Obtain robot joint trajectories from a simulator.
2. Use URDF and forward kinematics to render them as image-space action trajectories.
3. Encode the rendered trajectories into visual tokens and insert them into the model’s context.

With a small amount of fine-tuning, the model learns to generate future video consistent with these controls. Other structured visual signals, such as manually drawn object trajectories, can use the same interface. Intermediate physical variables can thus be predicted by the model or supplied to control generation.

Fine-tuning adapts the model’s learned use of additional visual context to a specified control signal, allowing reasoning and control to share a common conditioning interface.

Evaluation

Evaluation covers language-conditioned and action-conditioned prediction, with both single-view and multi-view settings. Language conditions specify behavior through instructions, while action conditions provide control signals. Multi-view prediction extends the task to observations from multiple camera viewpoints. The reported results include:

• TriWorldBench: No. 1 on the leaderboard, with a TWB-Score of 66.04 in the snapshot shown above.
• PAI-Bench: CausalWM ranks #1 in the robot domain, outperforming Cosmos 3 Super while significantly outperforming other models such as Veo 3 and Wan 2.2.

These results support explicitly modeling intermediate physical structure as a way to improve world models alongside scaling. The qualitative comparison below shows intermediate predictions and generated video for a bottle-to-drawer manipulation task.

Research Direction

CausalWM is Aether AI’s first embodied causal world model. Future work focuses on better causal representation learning and causal reasoning, identifying the variables that determine outcomes and the dependencies between them. The longer-term goal is to connect causal world understanding with action generation, so models can reason about how different actions and interventions change the future.