CST-WM

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.

The idea

The hardest moments in tracking are not when the target is clearly visible, but when it is temporarily occluded, leaves the field of view, or becomes confusable with distractors. The value of an action then depends not only on how it changes the robot state, but also on how it changes future target visibility and apparent scale, so embodied visual tracking is a predictive decision problem rather than a reactive mapping from the current frame to the next control.

Learning the predictor from logged tracking data has a pitfall: the recorded robot turns toward and slows down for the target, so its actions are correlated with target evidence. A predictor can exploit that correlation by attributing to the current action a direct effect on target evidence. The error is invisible on logged action sequences, where both explanations predict similar futures, and surfaces where planning relies on the model, under candidate actions that the recording robot never took.

Method

GroundingDINO returns a confidence H^vis and a box whose relative area gives the apparent scale H^area, so the target evidence H_ℓ = [H^vis_ℓ, H^area_ℓ] says whether the target is visible and how far away it is. The rest of the scene enters through the latent Z_ℓ of a pre-trained VAE, and the robot pose x_ℓ is integrated from past actions.

Physically, the action can change target evidence only by moving the camera: a_ℓ → x_ℓ+1 → Z_ℓ+1 → H_ℓ+1. Each transition layer updates three groups of tokens in a fixed order, and attention masks decide what each group may read. The target-evidence tokens E_ℓ+1 are updated first and read nothing that carries a_ℓ; the robot tokens R_ℓ+1 are the only representation that carries the current action; the observation tokens come last and read both. There is no path from a_ℓ to E_ℓ+1, so ∂Eℓ+1/∂aℓ = 0 holds by construction, while the action still shapes all later target evidence through the view.

An evidence head g_ψ, distilled from the detector, reads the target evidence of an imagined state from the observation-branch tokens, so imagined states are scored without decoding. Planning uses the Cross-Entropy Method (horizon T = 10, N = 128 candidates, K = 16 elites, M = 4 iterations): each candidate is scored by the visibility and apparent scale read from its imagined views, with penalties for invalid actions and unsafe predicted motion, and the robot executes the first action of the best sequence and replans. While the target is hidden, a candidate whose imagined views show the target again collects the visibility and scale terms from that step on.

Is apparent scale a usable stand-in for distance?

The detector's score tracks the relative distance over most of the following range and saturates only very close to the target, which is why the planner can use it as a distance signal. As the planner's target evidence it comes close to the simulator's true distance (0.76 against 0.78 re-acquisition), and the distance-aware loss, used only in training, raises its rank correlation with true distance from 0.71 to 0.83 across humanoids.

Demo

Real-world tracking on a Unitree Go2 with an Intel RealSense D435i RGB camera, recorded across four environments that differ in lighting and layout, plus one clip isolating a target-loss and re-acquisition episode. Pick a segment:

Bright indoor. Office corridors under normal lighting. The robot holds a steady following distance down a straight corridor while the target walks away from it.

All four environment segments, uncut, in one file: full video (2:58).

Real-world trials

We follow the real-robot protocol of TrackVLA: occlusion, distractor crossing, and fast motion, with 10 trials per scenario and method. As in TrackVLA and OA-VAT, planning runs on an offboard RTX 6000 Ada GPU with a reduced budget, one planning step every 1.4 s. The model keeps no target identity, so a trial in which the robot follows the distractor counts as a failure. CST-WM succeeds in 8, 6, and 6 of 10 trials, against 4, 6, and 4 for TrackVLA.

Results

EVT-Bench and Habitat 3.0 with egocentric observations and scene-disjoint splits; means over three seeds. The adapted NWM keeps NWM's non-factorized CDiT-XL/2 transition, in which the action conditions every token, and shares data, optimizer, evidence read-out, objective, CEM budget, and action bounds with CST-WM, at 6.6× our parameters and 3.6× our compute per pass, so any difference comes from the transition.

Tracking

On EVT-Bench single-target tracking, CST-WM reaches 88.7 SR and 83.4 TR, 2.7 and 2.4 points above TrackVLA++, with the lowest collision rate. On Habitat 3.0 it leads on all four metrics; against the adapted NWM, which differs only in the transition, it improves F and ES by 0.12 each and lowers CR by 0.06, although the adapted NWM is the larger model. The advantage carries over to cross-dataset transfer.

Main results: EVT-Bench single-target tracking; Habitat 3.0 standard tracking from a visible target; transfer from EVT-Bench to Habitat 3.0. ^†Released pipeline without the diffusion planner.

Recovery after temporary target loss

Under all four causes of target loss, CST-WM has the highest Re-acq., Post-F, and Rec-ES and the shortest TTR. Its margin over the adapted NWM is 0.09 in Re-acq. under short occlusion and 0.15–0.16 under long occlusion, out-of-view drift, and distractor crossing, where the target stays hidden longest.

Do imagined futures rank candidates like the simulator?

From each of 500 saved simulator states, the same 32 candidate action sequences are executed in the simulator and imagined by every model. CST-WM has lower latent prediction error and higher target-visibility AUROC at every horizon, and the gap grows with the horizon. Its ranking of the candidates agrees best with the simulator's: Spearman ρ 0.68, against 0.47 for the adapted NWM and 0.52 for a TrackVLA-style predictor.

Does the structural constraint hold?

The Jacobian |∂Eℓ+1/∂aℓ| and the change of E_ℓ+1 under action swaps are near zero for CST-WM, as the mask requires, whereas the adapted NWM absorbs action information. The mask does not guarantee a third property: with the humanoid held still and the robot moving, E_ℓ+1 can change through Z_ℓ and x_ℓ. CST-WM keeps it stable (0.98 against 0.61) while its predicted views follow the viewpoint (0.84 against 0.78). The action moves the view but not the target.

Ablations

Each ablation changes one component. Removing the target-evidence mask costs the most re-acquisition (0.76 to 0.61) and planning-value consistency (0.68 to 0.55). Without the mask the evidence head predicts the logged future better after 8 steps, yet ranks the counterfactual candidates worse and recovers worse in closed loop: accuracy on the actions the behavior policy took does not carry over to the actions it did not.

Planning cost

A planning step runs its denoising passes over all candidates as one batch; on one RTX A6000 a pass over 128 candidates takes 0.36 s. The simulator waits for the planner, so planning time does not enter the simulated results, and a 4.7× larger budget adds only 0.02 re-acquisition.

Citation

@article{hu2026cstwm,
  title   = {CST-WM: A Causally Structured World Model for Embodied Visual Tracking},
  author  = {Hu, Junyi and Yuan, Shuaihang and Liang, Jiazhao and Fang, Yi},
  journal = {arXiv preprint arXiv:2609.06302},
  year    = {2026}
}