Generalizable Optimal Control with Transformers: One Policy Across Diverse Systems

arXiv:2603.14910v2 Announce Type: replace-cross 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.

Source: arXiv cs.RO | 2026-09-22