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 $\pm 10\%$. 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.

This work has been accepted for publication in the proceedings of the 2026 IEEE Conference on Decision and Control (CDC), Honolulu, Hawaii, USA