Generalizable Optimal Control with Transformers: One Policy Across Diverse Systems

Abstract:Classical optimal control designs a separate controller for each plant. Even for the Linear Quadratic Regulator (LQR), every new model must be identified and its Riccati equation re-solved. We ask whether a single learned policy can instead serve an entire family of systems, and we show that one transformer can. We train the policy to imitate optimal LQR state feedback across a collection of heterogeneous Multiple-Input, Multiple-Output (MIMO) Linear Time-Invariant (LTI) systems that differ in their state and input dimensions and in their cost objectives. A shared representation lets the same parameters control every member of the family. It combines system-wise standardization, zero-padding and masking across dimensions, and an explicit encoding of the cost matrices. At run time, the policy maps a short window of recent states and the specified cost to a control action. It uses no plant matrices and identifies the dynamics implicitly from the state history. We evaluate on $28$ simulated systems over $9{,}675$ closed-loop rollouts, and no unstable rollout was observed in any of them. On the systems seen during training, it attains a median relative sub-optimality of $0.022\%$, even under parameter perturbations of up to $\pm10\%$. It transfers to unseen systems with lightweight fine-tuning, reaching a median sub-optimality of $0.19\%$. These results support transformers as generalizable near-optimal controllers for structured families of linear systems.

Systems and Control (eess.SY); Robotics (cs.RO)

arXiv:2603.14910 [eess.SY]

(or arXiv:2603.14910v2 [eess.SY] for this version)

https://doi.org/10.48550/arXiv.2603.14910 Focus to learn more arXiv-issued DOI via DataCite

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.