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

2026-07-16

AI Summary

World Action Models (WAMs) exhibit insufficient robustness under distributional shifts. This work proposes a novel approach that, for the first time, integrates mechanistic interpretability with optimal control to enhance robustness without additional training. By analyzing perturbation representations in the WAM activation space that correlate with robustness, the method introduces a contrastive activation steering technique guided by interpretable directions. Building upon locally linear dynamics, it further develops a lightweight feedback controller, WA-LQR. Evaluated on Cosmos-Policy and DiT4DiT, the approach significantly improves robustness against camera, gripper, and visual noise perturbations, outperforming both unguided and prompt-guided baselines, and demonstrates strong generalization to unseen tasks.

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.

Problem

• World Action Models
• distribution shift
• robustness
• activation space
• steerability

Innovation

• mechanistic interpretability
• optimal control
• World Action Models
• activation steering
• WA-LQR