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

2026/07/16 by Jihoon Hong, Julian Skifstad, Qiyue Dai +2
#cs.RO #cs.AI #cs.LG #cs.SY #eess.SY #math.OC

paper · pdf

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.

Citations

• Do World Action Models Generalize Better than VLAs? A Robustness Study
• Controlling Language and Diffusion Models by Transporting Activations
• External Invariants: A Cryptographic Trust Architecture for Institutional AI Inference
• The Geometry of Truth: Emergent Linear Structure in Large Language Model Representations of True/False Datasets
• Finding Structure with Randomness: Probabilistic Algorithms for Constructing Approximate Matrix Decompositions

Related

• DWM: Separating World Effects from Actions in Latent World Models
• Mechanistic interpretability for steering vision-language-action models
• WMPO: World Model-based Policy Optimization for Vision-Language-Action Models
• World Models as Reference Trajectories for Rapid Motor Adaptation
• RENEW: Towards Learning World Models and Repairing Model Exploitation from Preferences