Steering Robustness into World Action Models via Mechanistic Interpretability and Optimal Control
RoboticsArtificial IntelligenceMachine Learning
Summary
The authors study how World Action Models (WAMs), which help control systems by understanding the environment, respond to changes that can cause failures. They find that some WAMs represent important features in a way that can be separated linearly, which helps guide the model's responses without extra training. Using this insight, the authors develop a new feedback controller called WA-LQR that steers these models efficiently by leveraging their internal linear behavior. They test their approach on different WAMs, showing that WA-LQR improves robustness to various disturbances in certain models but not all. Their work helps understand and improve how control systems stay reliable when conditions change.
Authors
Jihoon Hong, Julian Skifstad, Qiyue Dai, Alice Chan, Glen Chou
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.