CST-WM: A Causally Structured World Model for Embodied Visual Tracking

Junyi Hu; Shuaihang Yuan; Jiazhao Liang; Yi Fang

cs.CVSep 5, 2026

CST-WM: A Causally Structured World Model for Embodied Visual Tracking

Authors: Junyi Hu, Shuaihang Yuan, Jiazhao Liang, Yi Fang

Organizations: New York University Abu Dhabi

Abstract

Embodied visual tracking requires a robot to choose actions that keep a moving target observable at a suitable distance, and to recover it after occlusion, out-of-view drift, or distractor crossings. We cast the task as planning over future target evidence with an action-conditioned world model. In logged tracking data, however, the behavior policy's actions are correlated with where the target is, so a generic predictor can learn a shortcut: it writes the current action directly into its prediction of target evidence, instead of letting the action affect that evidence only by moving the robot and changing what it observes. We call this failure causal hallucination; the resulting rollouts look plausible but rank candidate actions for the wrong reason. We propose CST-WM, a causally structured world model whose state is split into target-evidence, robot, and observation branches. Its transition removes the same-step edge from action to target evidence but keeps the path through robot motion and the resulting views, so candidate actions are still distinguished by their predicted ego-motion. With rollout-based model-predictive control, a single model handles both steady following and re-acquisition after target loss. On EVT-Bench and Habitat 3.0, covering standard tracking, target-loss recovery, and cross-dataset transfer, CST-WM improves following, distance-range control, safety, and re-acquisition over reactive trackers and world-model baselines, and removing the action mask causes the largest drop in re-acquisition among our ablations. Offline, CST-WM has lower multi-step rollout error, and its ranking of candidate actions agrees better with the simulator's. On a Unitree Go2 quadruped, CST-WM succeeds in 20 of 30 real-world trials under occlusion, distractor crossing, and fast motion, against 14 for TrackVLA.

Figures & tables

Figure 1: Causal structure of action-conditioned prediction in embodied visual tracking. Given an egocentric observation and candidate action sequences (A: forward, turn-left; B: forward, turn-right), the model must imagine which sequence preserves target observability and following distance. Conventional world models route the current action directly into the target-evidence update, producing plausible rollouts that fail to discriminate A from B. CST-WM removes this direct edge: the target-evidence branch $E_{\ell+1}$ is updated without action injection, and the action only reaches future observations through the robot branch $R_{\ell+1}$ . The resulting predictions distinguish the two sequences and agree with the simulator.

Figure 2: Overview of CST-WM. (a) State. GroundingDINO reads the target evidence $H_{\ell}$ (confidence, box area) from $O_{\ell}$ , a frozen VAE encodes $O_{\ell}$ into $Z_{\ell}$ , and the pose $x_{\ell}$ is integrated from past actions. (b) Transition. (b1) The edge $a_{\ell}\rightarrow E_{\ell+1}$ is removed, and the action reaches the next target evidence only along the highlighted path through $R_{\ell+1}$ and $Z_{\ell+1}$ . (b2) Each layer updates the target-evidence, robot, and observation branches in this order, and the evidence head $g_{\psi}$ , distilled from the detector, reads $\widehat{H}_{\ell+1}$ from the observation branch. (c) Planning. CEM rolls out $N$ sampled sequences, scores the evidence read from the imagined views, refits to the top $K$ , and executes the best first action.

EVT-Bench (STT)				Habitat 3.0, standard tracking					Cross-dataset transfer to Habitat 3.0
Method	SR $\uparrow$	TR $\uparrow$	CR $\downarrow$	Method	F $\uparrow$	DRS $\uparrow$	CR $\downarrow$	ES $\uparrow$	Method	F $\uparrow$	DRS $\uparrow$	CR $\downarrow$	ES $\uparrow$
Uni-NaVid	25.7	39.5	41.9	Habitat 3.0 baseline	0.29	0.47	0.48	0.40	Uni-NaVid	0.40	0.56	0.39	0.45
OA-VAT †	32.4	63.9	3.70	Uni-NaVid	0.34	0.52	0.43	0.47	TrackVLA	0.38	0.54	0.41	0.43
TrackVLA	85.1	78.6	1.65	Adapted NWM	0.41	0.61	0.33	0.49	Adapted NWM	0.43	0.58	0.35	0.47
TrackVLA++	86.0	81.0	2.10	SDA-S2	0.39	0.63	0.57	0.43	TrackVLA++	0.42	0.61	0.38	0.43
Ours	88.7	83.4	1.41	Ours	0.53	0.70	0.27	0.61	Ours	0.48	0.65	0.30	0.54

Table 1: Main results. Left: EVT-Bench, single-target tracking (STT). Middle: Habitat 3.0, standard tracking from a visible target. Right: transfer from EVT-Bench to Habitat 3.0. † Released pipeline without the diffusion planner (Appendix C ). Best in bold , second underlined .

	Short occlusion		Long occlusion		Out-of-view drift		Distractor crossing
Method	Re-acq. $\uparrow$	TTR $\downarrow$	Re-acq. $\uparrow$	TTR $\downarrow$	Re-acq. $\uparrow$	TTR $\downarrow$	Re-acq. $\uparrow$	TTR $\downarrow$
TrackVLA	0.71	8.9	0.48	14.6	0.52	12.8	0.43	15.2
Adapted NWM	0.75	8.1	0.54	13.2	0.57	11.7	0.49	14.0
Ours	0.84	6.4	0.69	9.8	0.73	8.7	0.65	10.6
Method	Post-F $\uparrow$	Rec-ES $\uparrow$	Post-F $\uparrow$	Rec-ES $\uparrow$	Post-F $\uparrow$	Rec-ES $\uparrow$	Post-F $\uparrow$	Rec-ES $\uparrow$
TrackVLA	0.58	0.47	0.41	0.29	0.45	0.33	0.37	0.24

Table 2: Recovery after temporary target loss (Habitat 3.0 target-loss protocol). Re-acq.: re-acquisition success; TTR: time to re-acquire (steps); Post-F: following rate after re-acquisition; Rec-ES: recovery episode success. Best in bold , second underlined .

Table 6

Figure 5: Real-world tracking with CST-WM in a library, office corridors, and a parking lot at night. Each row is one episode, left to right in time; the blue curve is the planned path.

Table 9

	Target evidence	F $\uparrow$	DRS $\uparrow$	Re-acq. $\uparrow$	CR $\downarrow$	Plan. $\rho\uparrow$	Time (s) $\downarrow$
Privileged	Simulator true distance	0.54	0.72	0.78	0.26	–	–
Privileged	Detector on rendered future views	0.48	0.67	0.64	0.24	0.74	5.2
Read-out	Detector on decoded views	0.38	0.46	0.46	0.40	0.58	5.8
	Separate latent head on $\widehat{Z}$	0.45	0.60	0.57	0.33	0.54	0.030
	Head, frozen transition	0.34	0.37	0.43	0.39	0.62	0.021
Signal	Raw detector score	0.46	0.60	0.66	0.32	–	–

Table 8: Target evidence for planning on Habitat 3.0: privileged sources, imagined-step read-outs, and signal variants; transition, objective, and CEM budget fixed. Time: read-out per imagined step.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Table 11

Item	Details
Data	EVT-Bench / Habitat 3.0; scene-disjoint splits. Train scenes: 64 / 72; val: 8 / 9; test: 8 / 9. Trajectories: 48,000 / 61,000. One-step clips: 1.92M / 2.44M. Filtering removes corrupted frames, invalid robot states, reset fragments, terminal-only collision fragments, and action-invalid transitions.
Observation and action	Observations sampled at the simulator control frequency; consecutive frames without temporal skipping; actions at the same rate.
Model	VAE ( Blattmann et al., 2023 ) (sd-vae-ft-ema; 34.2M-parameter encoder, 49.5M-parameter decoder), $224{\times}224$ views, latent $4{\times}28{\times}28$ . Target-evidence branch input: the scalar $0.4\,H^{\mathrm{vis}}+0.6\,(H^{\mathrm{area}})^{1/2}$ (the evidence head outputs the two components); conditioning on the current step only: one observation latent, no history; robot state: planar position and yaw; action $(v,\omega)$ . Transition: 12 causal layers, each with target-evidence, robot (a pose sub-layer and an action sub-layer, giving $R_{\ell+1}=[x_{\ell+1},W_{a}(a_{\ell})]$ ), and observation sub-layers (attention + FFN), hidden 512, 8 heads, FFN 2048; 153.6M parameters, 89 GFLOPs per denoising pass. Evidence: GroundingDINO Swin-T (172.8M parameters), prompt “humanoid”, views resized to $800{\times}800$ , on the current observation and for training labels; evidence head $g_{\psi}$ (0.26M parameters) on imagined steps.
Diffusion training	Linear DDPM with fixed variance, 1000 train steps, $\varepsilon$ -prediction on the next observation latent, 20 DDIM inference steps; AdamW, lr $8{\times}10^{-5}$ , no weight decay, batch 16, 150 epochs, grad-norm clip 10, EMA 0.9999, bf16 mixed precision; 3 seeds; evidence-head loss weight $\lambda_{\mathrm{ev}}=1$ (joint training from the 12k-step transition); auxiliary distance-aware loss weight $\lambda_{\mathrm{dist}}=0.2$ .
Planning	$T{=}10$ , $N{=}128$ , $K{=}16$ , $M{=}4$ , 20 DDIM steps; $\beta_{\mathrm{vis}}{=}\beta_{\mathrm{dist}}{=}\lambda_{\mathrm{valid}}{=}\lambda_{\mathrm{safe}}{=}1.0$ ; $\alpha{=}0.5$ ; $d_{\mathrm{safe}}{=}0.20$ m; one stochastic rollout per candidate; target evidence read by $g_{\psi}$ at every imagined step.
Compute	Training: a single GPU per model, bf16 mixed precision, deterministic evaluation seeds. Runtime: measured on one RTX A6000, and on an RTX 6000 Ada for the NWM-protocol trajectory times and the robot-budget planning step; per-component times, memory, and the NWM-protocol trajectory times are in Appendix D.4 and Table 17 . Random seeds per main model: 3. Evaluation episodes per model: 500 starting states $\times$ 32 candidate sequences for offline diagnostics, plus standard EVT-Bench / Habitat 3.0 splits for downstream metrics.

Appendix

Table 9: Implementation details: dataset scale, architecture, optimization, planner, and compute.

Method	Type	Target evidence	Action selection	Model and compute
Uni-NaVid, TrackVLA, TrackVLA++	reactive VLA	learned, from RGB and instruction	one policy pass per step	own architecture and training data
OA-VAT †	reactive tracker	own detector, tracker, and re-identification	PID control law	released weights, no planner
Habitat 3.0 baseline, SDA-S2	reactive	own perception	one policy pass per step	own architecture and training data
Adapted NWM	world model, CEM	GroundingDINO on current view, evidence head on imagined views	Eq. ( 5 ), same budget	CDiT-XL/2, 1.01B, 321 GFLOPs per pass; same as ours otherwise
Ablations (Table 7 )	world model, CEM	GroundingDINO on current view, evidence head on imagined views	Eq. ( 5 ), same budget	full model with one component changed
CST-WM (ours)	world model, CEM	GroundingDINO on current view, evidence head on imagined views	Eq. ( 5 )	causal transition, 0.15B, 89 GFLOPs per pass

Appendix

Table 10: What each compared method receives at test time. “Same as ours” for the adapted NWM and the ablations means identical training data, optimizer settings, evidence read-out, planning objective (Eq. ( 5 )), and CEM budget (

T{=}10

N{=}128

K{=}16

M{=}4

, 20 DDIM steps); ablations differ from the full model only in the named component. † Run by us on EVT-Bench STT with the released pipeline, without the diffusion planner.

Method	All states $\downarrow$	In-view $\downarrow$	Out-of-view $\downarrow$	Corr. with performance
Adapted NWM	0.112	0.096	0.131	$-0.41$
Leaky variant B	0.331	0.304	0.357	$-0.53$
Leaky variant C	0.086	0.071	0.102	$-0.36$
Ours	0.001	0.001	0.002	$-0.04$

Appendix

Table 11: Jacobian

|\partial E_{\ell+1}/\partial a_{\ell}|

of the target-evidence representation over all rollout states and by target visibility, and its correlation with per-episode tracking performance. Best in bold , second underlined (Jacobian columns).

Table 15

Figure 7: GroundingDINO target-evidence score against relative robot–human distance on EVT-Bench. These readings are the current evidence and the training labels of the evidence head.

Table 17

Stage	Default	Fast	Robot	Robot, Ada
Budget $(T,N,M,S)$	(10,128,4,20)	(6,64,2,10)	(4,24,2,10)	(4,24,2,10)
GroundingDINO, current view	0.09 s	0.09 s	0.09 s	–
VAE encoder, current frame	7 ms	7 ms	7 ms	–
Denoising passes (sequential, batch $N$ )	800	120	80	80
Evidence-head reads (batch $N$ )	40	12	8	8
Time per action (whole step)	267 s	22.5 s	8.2 s	1.4 s

Appendix

Table 16: Wall-clock time of one planning step from the received observation to the issued action, for the three budgets, on one RTX A6000 (median of full steps after warm-up); the robot budget is also given on the offboard RTX 6000 Ada of the real-world trials. Denoising passes run in sequence, each over the

N

candidates as one batch; per-pass times are in Table 18 .

Model, GPU	Base	+Skip	+6 steps †	+4-bit ∗
NWM, Ada (reported)	30.3	14.7	0.4	0.1
NWM, A6000	166.2	84.0	2.30	0.57
Ours, A6000	135.6	67.7	1.88	0.47
Ours, Ada	23.4	12.3	0.34	0.08

Appendix

Table 17: Time (s) to simulate one trajectory at batch size 1 (NWM protocol): 16 states with 250 steps (Base), 8 states (+Skip), 6 sampling steps, and 4-bit quantization. Ada: RTX 6000 Ada. ∗ Estimated as in NWM: 6-step time divided by 4. † NWM distills to 6 steps; we time a 6-step sampler.

Component	Size	Batch	Time	Peak memory
VAE encoder, 1 / 4 frames	34.2M	1	7 / 23 ms	–
Ours, one denoising pass	153.6M, 89 GFLOPs	1 / 24 / 64	31 / 106 / 198 ms	1.2 / 1.6 / 2.4 GB
		128 / 192	362 / 572 ms	3.6 / 4.8 GB
NWM CDiT-XL/2, one denoising pass	1.01B, 321 GFLOPs	1 / 128	39 / 1017 ms	–
Evidence head $g_{\psi}$	0.26M	24 / 64 / 128 / 192	0.21 / 0.20 / 0.20 / 0.27 ms	–
VAE decoder, per view	49.5M	1 / 64	15.6 / 10.6 ms	–

Appendix

Table 18: Per-component runtime and peak allocated GPU memory on one RTX A6000. Denoising passes are timed at the candidate batch sizes of the robot budget and of the fast, default, and large budgets of Figure 6 ; the NWM backbone is listed for reference. The evidence head reads imagined views during planning, GroundingDINO reads the current observation, and the decoder serves visualization and the decode-and-detect read-out of Table 8 . ∗ Batch size 1 includes CPU resizing in the reference code.

Explore similar work

CausalWM: Causal Chain-of-Thought Reasoning for Embodied World Model
Sep 19, 2026Ziming Xu, Shuang Liang, Ruobing Han +10Efficient World-Action Model Future Video Prediction
Aether AI · University of California, San Diego · Vanderbilt University
From World Models to World Action Models: A Concise Tutorial for Robotics
Jul 1, 2026Xiaoxiong Zhang, Xiong Zeng, Wei ZhangWorld Models Action Prediction
School of Automation and Intelligent Manufacturing, Southern University of Science and Technology, Shenzhen, China · LimX Dynamics, Shenzhen, China
CSWAM: Better Causal Semantic Representations for Out-of-Distribution Generalization in World Action Models
Sep 16, 2026Tianbin Liu, Jian Zhu, Taiyi Su +5Efficient World-Action Model World Models
AIRC, Midea Group · Tongji University

Sep 19, 2026, cs.CVJ/K move · Enter open · S save

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

Ziming Xu, Shuang Liang, Ruobing Han, Ziqiao Xi+9

Aether AI · University of California, San Diego · Vanderbilt University

Efficient World-Action Model Future Video Prediction Causal Chain-of-Thought Reasoning Reinforcement Learning Post-Training

Open paper AskPDF arXiv

Embodied world models learn to predict future physical dynamics from visual observations and control signals, where physical knowledge is implicitly entangled within latent representations. We introduce CausalWM, a 16B embodied world model that performs explicit causal chain-of-thought reasoning before future video prediction. CausalWM organizes useful variables into a reasoning trajectory, allowing the model to progressively capture causal dependencies underlying physical evolution. To train CausalWM, we collect 31K hours embodied data and develop a three-stage paradigm consisting of large-scale video pre-training, causal CoT mid-training, and multi-objective RL post-training. Despite using only a limited set of supervised CoT variables, CausalWM exhibits emergent in-context learning capabilities, enabling contextual visual feature guidance and efficient few-step generation. CausalWM achieves state-of-the-art performance across language-conditioned, action-conditioned, single-view and multi-view benchmarks, including Top-1 performance on TriWorldBench leaderboard.

Abstract

Explore similar work

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

From World Models to World Action Models: A Concise Tutorial for Robotics

CSWAM: Better Causal Semantic Representations for Out-of-Distribution Generalization in World Action Models