Transformers As Generalizable Optimal Controllers

TL;DR

This paper demonstrates that transformer-based policies can effectively learn near-optimal state-feedback controllers for a broad class of linear systems, generalizing across different system dimensions and parameters.

Contribution

It introduces a transformer-based approach to approximate optimal controllers for heterogeneous LTI systems, enabling generalization and adaptability without explicit plant models.

Findings

Achieves small sub-optimality compared to LQR

Remains stabilizing under moderate perturbations

Benefits from lightweight fine-tuning on unseen systems

Abstract

We study whether optimal state-feedback laws for a family of heterogeneous Multiple-Input, Multiple-Output (MIMO) Linear Time-Invariant (LTI) systems can be captured by a single learned controller. We train one transformer policy on LQR-generated trajectories from systems with different state and input dimensions, using a shared representation with standardization, padding, dimension encoding, and masked loss. The policy maps recent state history to control actions without requiring plant matrices at inference time. Across a broad set of systems, it achieves empirically small sub-optimality relative to Linear Quadratic Regulator (LQR), remains stabilizing under moderate parameter perturbations, and benefits from lightweight fine-tuning on unseen systems. These results support transformer policies as practical approximators of near-optimal feedback laws over structured linear-system families.