Steering Robustness into World Action Models via Mechanistic Interpretability and Optimal Control

Jihoon Hong*  ·  Julian Skifstad*  ·  Qiyue Dai  ·  Alice Chan  ·  Glen Chou

*Equal contribution  ·  Georgia Institute of Technology

TL;DR

World Action Models (WAMs) fail under camera shifts, gripper-position changes, and sensor noise. We find that some WAMs (Cosmos-Policy, DiT4DiT) encode these failure modes as simple, low-dimensional linear directions in activation space — while others (LingBot-VA) do not. We use this mechanistic signal to predict which models are steerable, then exploit it with WA-LQR, a training-free, closed-loop LQR controller that nudges WAM activations toward robust behavior at inference time — improving task success rates by up to +41% with no finetuning.

Teaser figure: WA-LQR steers a Cosmos-Policy World Action Model to succeed under camera, gripper, and noise perturbations where the unsteered model fails.
WA-LQR makes World Action Models more robust to perturbations including gripper-position changes, camera-orientation shifts, and Gaussian sensor noise. In these Cosmos-Policy examples from LIBERO-10, the yellow and green boxes mark the two objects that must be placed in the basket. Without steering, the WAM fails; with WA-LQR, it succeeds.
Abstract

World Action Models (WAMs) enable semantically- and physically-informed control but are brittle under distribution shift. In this work, we use mechanistic interpretability to study how robustness-relevant perturbations are represented in WAM activation space. Comparing activations across successful and unsuccessful rollouts, we find some WAM architectures exhibit low-dimensional linear separability for robustness-critical features, while others do not. This motivates the use of contrastive activation directions for training-free WAM steering. We also show that local linearity in WAM activation dynamics enables efficient feedback steering via model-based optimal control, yielding World-Action Linear Quadratic Regulator (WA-LQR), a minimally-invasive reduced-order LQR controller.

Via mechanistic evaluations, we predict strong steerability in the Cosmos-Policy and DiT4DiT models but weak steerability in LingBot-VA, consistent with steering intervention results. On Cosmos-Policy and DiT4DiT, WA-LQR generalizes contrastive directions to new tasks and improves robustness to camera, gripper, and visual-noise perturbations over unsteered and prompt steering baselines.

Mechanistic Study

Steerability is architecture-dependent

We compare WAM activations from nominal vs. perturbed rollouts (camera pose, gripper position, image noise), project the contrastive directions onto their top-3 principal components, and fit a linear SVM. The hinge loss quantifies separability: 0 means perfectly linearly separable, 1 means random.

PCA visualization showing clean linear separability between clean and perturbed activations for Cosmos-Policy and DiT4DiT.
Cosmos-Policy & DiT4DiT exhibit clean linear separability between clean and perturbed activations across denoising blocks, under both noise (a) and camera (b) perturbations.
PCA visualization showing weak separability between clean and perturbed activations for LingBot-VA.
LingBot-VA shows substantially weaker separability — clean and camera-perturbed activations overlap heavily at the first, best intermediate, and final transformer blocks.
Pairwise separability matrix for Cosmos-Policy under Gaussian noise corruption across LIBERO-10 tasks.
Some task pairs share feature representations. Pairwise separability for Cosmos-Policy under Gaussian noise corruption across LIBERO-10 tasks. High linear separability and shared feature clusters across certain task pairs suggest reusable representations that WA-LQR can exploit for cross-task steering.
Method

Training-free activation steering for WAMs

Given paired rollouts that differ mainly in a robustness-relevant feature (e.g. nominal vs. perturbed camera pose), we compute contrastive activation directions and use them to steer WAM inference — with no finetuning and no weight updates.

Open-loop

Activation Addition (ActAdd)

Average contrastive directions into a single steering vector, then add it at every denoising step with a fixed strength γ. Simple and training-free, but ignores the model's current activation state — it can oversteer or degrade the action.

xl,t,τ ← xl,t,τ + γ al,t
Closed-loop · ours

World-Action LQR (WA-LQR)

Project activations into a low-dimensional contrastive subspace, treat the DiT block sequence as a locally-linear dynamical system, and solve a per-timestep LQR problem to track a target feature setpoint — steering only as much as the realized activation deviates from it.

  1. Project — zl,t,τ = Pl,t xl,t,τ ∈ ℝdz, dz ≪ dx
  2. Linearize across transformer blocks via Jacobian-vector products
  3. Track a feature setpoint with online error αl,t,τ
  4. Solve LQR offline via Riccati recursion → gains Kl,t,τ
  5. Apply online: u*l,t,τ = ūl,t,τ + Kl,t,τ αl,t,τ vzl,t
Results

Steering improves LIBERO-10 robustness — when the feature is separable

We evaluate WA-LQR and ActAdd against unsteered and prompt-steering baselines on LIBERO-10 tasks under camera-orientation, initial-gripper-position, and Gaussian-noise perturbations (30 trials/task).

Best-case robustness gain
+41%
Success-rate improvement from WA-LQR steering, no finetuning
Effective feature dimensionality
3 PCs
Contrastive directions are well summarized by as few as 3 principal components
Architecture-dependence
2 of 3
Models (Cosmos-Policy, DiT4DiT) show strong linear separability — LingBot-VA does not
No Steering ActAdd (open-loop) WA-LQR (closed-loop, ours)

Cosmos-Policy 2B

LIBERO-10 success rate by perturbation type, averaged across task-transfer pairs
46.0%
49.3%
59.3%
Camera
Orientation
61.3%
63.3%
72.7%
Initial Gripper
Position
26.7%
67.3%
58.7%
Camera
Gaussian Noise

DiT4DiT

LIBERO-10 success rate by perturbation type, averaged across task-transfer pairs
65.7%
65.8%
71.7%
Initial Gripper
Position
15.6%
43.3%
48.9%
Camera
Gaussian Noise

LingBot-VA

LIBERO-10 success rate by perturbation type — weak separability, weak steering gains
48.0%
41.0%
51.0%
Camera
Orientation
72.0%
72.0%
75.0%
Initial Gripper
Position
59.0%
38.0%
56.0%
Camera
Gaussian Noise
Scatter plot showing negative correlation between SVM hinge loss and steering performance gain across tasks and models.
Separability predicts steerability. Hinge loss (lower = more separable) correlates negatively with steering performance gains across tasks and models — validating hinge loss as a cheap, training-free diagnostic for whether a WAM is worth steering.
Cumulative variance explained by top-k singular vectors of contrastive activation directions, showing rapid decay.
Low-dimensional structure. Variance of latent contrastive vectors for camera-orientation perturbations on Cosmos-Policy concentrates in a handful of singular directions — enabling the reduced-order LQR design behind WA-LQR.
Video Rollouts

Unsteered failures vs. WA-LQR steered successes

Paired rollouts from the same task and perturbation, with and without WA-LQR steering. In every pair, the unsteered policy fails and the WA-LQR-steered policy succeeds.

Cosmos-Policy 2B

Camera Orientation
Unsteered — fails
WA-LQR steered — succeeds
Initial Gripper Position
Unsteered — fails
WA-LQR steered — succeeds

DiT4DiT

Initial Gripper Position
Unsteered — fails
WA-LQR steered — succeeds

LingBot-VA

Initial Gripper Position
Unsteered — fails
WA-LQR steered — succeeds

Note: LingBot-VA showed weak activation separability, so steering gains are small and inconsistent overall — this pair illustrates a favorable case rather than the typical outcome.

Citation

BibTeX

arXiv

@article{hong2026steering,
  title   = {Steering Robustness into World Action Models via Mechanistic Interpretability and Optimal Control},
  author  = {Hong, Jihoon and Skifstad, Julian and Dai, Qiyue and Chan, Alice and Chou, Glen},
  journal = {arXiv preprint arXiv:2607.14943},
  year    = {2026}
}

RSS Robot World Models Workshop

@inproceedings{skifstad2026steering,
  title     = {Steering Robustness into World Action Models via Mechanistic Interpretability and Optimal Control},
  author    = {Skifstad, Julian and Hong, Jihoon and Dai, Qiyue and Chan, Alice and Chou, Glen},
  booktitle = {Robot World Models},
  year      = {2026},
  url       = {https://openreview.net/forum?id=K691Pk5J1f}
}