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

As robotic foundation models move toward World Action Models (WAMs)—which couple action prediction with future-state modeling—they remain brittle under distribution shifts like camera changes or sensor noise. The authors investigate whether the internal activation space of these models contains interpretable, linear representations of robustness-critical features that can be used to improve performance without costly retraining or fine-tuning.

Using mechanistic interpretability, the researchers analyze the activation space of three WAM architectures: Cosmos-Policy, DiT4DiT, and LingBot-VA. They compare activations across successful and unsuccessful rollouts to test for low-dimensional linear separability. Based on these findings, they introduce World-Action Linear Quadratic Regulator (WA-LQR), a training-free steering method. WA-LQR projects activations into a low-dimensional contrastive subspace and uses local linear dynamics to synthesize a closed-loop controller that steers the model toward robust behavior during inference.

The study reveals that steerability is not universal across WAMs; Cosmos-Policy and DiT4DiT exhibit clear linear structure in their activations, whereas LingBot-VA does not. This linear separability is a strong predictor of how well a model can be steered. On models with high steerability, WA-LQR effectively generalizes contrastive directions to new tasks and significantly improves robustness against camera, gripper, and visual-noise perturbations compared to standard prompt-steering or open-loop activation addition techniques.

This work provides a diagnostic framework for evaluating the robustness of robotic foundation models before deployment. By showing that mechanistic interpretability can guide inference-time interventions, the authors offer a path to improving robot reliability in dynamic environments without the need for extensive data collection or model retraining.