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

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.

• The paper introduces WA-LQR, a training-free closed-loop controller that steers robustness-relevant activation features in a reduced rank-64 subspace using locally linear dynamics and LQR feedback.
• The paper finds that activation separability predicts steering success, with WA-LQR improving Cosmos-Policy success from 46.0% to 59.3% under camera changes and from 61.3% to 72.7% under gripper perturbations on LIBERO-10.
• The paper shows that steering depends on architecture and task clusters: it works well for Cosmos-Policy and DiT4DiT but offers limited gains for LingBot-VA, while real-robot transfer and automatic cluster discovery remain open problems.

Overview

This paper, by Hong, Skifstad, Dai, Chan, and Chou at Georgia Institute of Technology (2607.14943), addresses the brittleness of World Action Models (WAMs) under out-of-distribution (OOD) perturbations. WAMs couple action prediction with future-state modeling through video-generative backbones, but remain fragile to camera viewpoint changes, robot initial-state perturbations, and visual corruption. The authors pursue two questions: whether robustness-relevant nuisance features are linearly represented in WAM activation space, and whether such representations can be exploited for training-free robustness improvement at inference time. Their answer combines a mechanistic interpretability (MI) study with a control-theoretic steering method, the World-Action Linear Quadratic Regulator (WA-LQR), evaluated on Cosmos-Policy 2B, DiT4DiT, and LingBot-VA on LIBERO-10.

Mechanistic analysis of robustness features

The mechanistic study compares activations from nominal and perturbed rollouts across three perturbation types: initial gripper position, camera orientation, and Gaussian image noise. Because DiT activations have shape $(F,H,W,D)$ with $D$ large relative to sample counts, the authors mean-pool over token positions and action-chunk timesteps, compute contrastive directions $d_k = \bar{x}^+_k - \bar{x}^-_k$, project onto the top three principal components of these directions, and fit a linear SVM whose average hinge loss serves as a separability metric (0 for perfect separation, 1 for random).

The central finding is that linear separability is strongly architecture- and task-dependent. Cosmos-Policy 2B and DiT4DiT exhibit clear low-dimensional linear structure across multiple perturbations, whereas LingBot-VA shows little to no separability in either its action or video modules. Notably, separable features are not shared universally across tasks: certain clusters of tasks share representations of the same perturbation while others do not. This task clustering is what makes cross-task steering possible — contrastive directions computed on one task transfer to other tasks within the same cluster — but it also means no single direction applies to all scenes, an assumption the paper explicitly concedes is infeasible.

A further observation is that hinge loss correlates negatively with steering gains, making separability loss a practical diagnostic: models predicted steerable from MI alone are indeed the ones where interventions succeed. This predictive validity is arguably the paper's most useful contribution beyond the steering method itself, since it allows practitioners to determine a priori whether a given WAM is amenable to activation steering.

Open-loop steering via activation addition

As a baseline intervention, the authors adapt ActAdd from the LLM literature: contrastive directions are averaged over prompt pairs and action-chunk timesteps into a steering vector $a_{l,t}$, added at inference time with strength $\gamma$. The method is training-free and weight-preserving but open-loop — it ignores the current activation state and applies a fixed perturbation regardless of whether the feature is already at its desired strength. The appendix sensitivity analysis quantifies this weakness sharply: on Cosmos-Policy under Gaussian noise, ActAdd peaks at 67.4% mean success at $\gamma=0.1$ but collapses to 0% success at $\gamma \geq 0.2$, demonstrating extreme hyperparameter fragility that motivates closed-loop control.

WA-LQR: closed-loop steering via reduced-order optimal control

WA-LQR treats WAM inference as a controlled dynamical system. For each layer–denoising-timestep pair $(l,t)$, contrastive directions are pooled and compressed via streaming randomized SVD into a rank-64 orthonormal basis $V_{l,t}$, defining latent states $z = P x$. Two empirical assumptions justify this reduction, both validated experimentally:

1. Variance concentration: contrastive-vector variance is mostly captured by the top few singular directions, so the perturbation-relevant information survives projection.
2. Local linearity: cosine similarity and magnitude ratios between first-order Jacobian predictions and actual block outputs remain close to 1 throughout inference, and the reduced Jacobians $D$0 show high subspace overlap across inputs and tasks, allowing matrices computed on one input to be reused elsewhere.

Within this latent space, WA-LQR defines a feature setpoint along the normalized contrastive direction $D$1, computes the online tracking error $D$2 per action-chunk index $D$3, and solves an independent $D$4-step LQR over transformer blocks for each $D$5. A key design element is an action-decay schedule on the control penalty $D$6, which grows exponentially with chunk index and saturates at a large value, so steering is strongest early in the rollout and effectively vanishes later. The resulting controller $D$7 intervenes proportionally to realized feature error rather than at fixed magnitude. Solving each subproblem costs $D$8 via Riccati recursion, with Jacobian products handled by JVPs/VJPs without materializing full Jacobians. To the authors' knowledge, these are the first open- and closed-loop activation steering methods for WAMs.

Empirical results

On LIBERO-10 with 30 trials per task, WA-LQR improves Cosmos-Policy robustness substantially. Representative averages:

Perturbation	No steering	Prompt steering	ActAdd	WA-LQR
Camera orientation	46.0%	46.7%	49.3%	59.3%
Initial gripper position	61.3%	60.7%	63.3%	72.7%
Camera Gaussian noise	26.7%	0.7%	67.3%	58.7%

Individual results reach up to a 41-percentage-point improvement (e.g., gripper Task 1→Task 2 rising from 83.3% to 100%, noise Task 6→Task 1 from 53.3% to 73.3%). Prompt steering is notably counterproductive under Gaussian noise, dropping average success to 0.7%. One exception to WA-LQR's dominance is Gaussian noise corruption, where ActAdd outperforms it; the authors attribute this to the very strong separability of that feature (hinge loss near 0.04), which makes simple open-loop addition viable. On DiT4DiT, WA-LQR again leads on both gripper position (71.7% vs. 65.7% unsteered) and Gaussian noise (48.9% vs. 15.6%).

LingBot-VA results confirm the mechanistic prediction of weak steerability: neither ActAdd nor WA-LQR produces consistent improvements, and ActAdd frequently degrades performance through oversteering, while WA-LQR's feedback modulation limits damage. Across all models, the negative correlation between hinge loss and steering benefit supports using separability as a screening tool before deploying any steering intervention.

Limitations and open questions

The paper is candid about scope constraints. First, there is currently no interpretable method for predicting when tasks or environments share transferable representations; cluster discovery requires per-setting mechanistic analysis, limiting plug-and-play applicability. Second, steerability depends heavily on architecture — LingBot-VA's mixture-of-transformers design resists the linear-feature assumption — and the paper does not explain why some architectures develop separable robustness features while others do not. Third, the rank-64 subspace, the exponential decay schedule parameters ($D$9, $d_k = \bar{x}^+_k - \bar{x}^-_k$0), and the strength parameter $d_k = \bar{x}^+_k - \bar{x}^-_k$1 were tuned via search, and the sensitivity of WA-LQR itself to these choices is less thoroughly characterized than ActAdd's. Finally, evaluations are confined to LIBERO-10 simulation; real-robot transfer remains untested. These observations motivate the authors' stated goal of designing WAMs that are inherently steerable while preserving base-model representations.

Conclusion

This work connects mechanistic interpretability to optimal control for robotic foundation models. It establishes that robustness-critical features in WAMs can be linearly separable in a low-dimensional activation subspace — architecture permitting — that separability loss predicts steering success, and that locally linear latent dynamics support efficient closed-form LQR feedback controllers that outperform open-loop steering while avoiding its oversteering failure mode. Success-rate improvements of up to 41 points without any finetuning demonstrate that inference-time representation control is a viable complement to data-centric approaches for WAM robustness, though its applicability remains contingent on architectural properties that are not yet predictable in advance.