Generalizable Optimal Control with Transformers: One Policy Across Diverse Systems

Turki Bin Mohaya; Maitham F. AL-Sunni; John M. Dolan; Peter Seiler

eess.SYMar 16, 2026

Generalizable Optimal Control with Transformers: One Policy Across Diverse Systems

Authors: Turki Bin Mohaya, Maitham F. AL-Sunni, John M. Dolan, Peter Seiler

Organizations: Robotics Institute, Carnegie Mellon University, Pittsburgh, PA, USA

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.

Figures & tables

Fig. 1: Unlike the classical LQR approach, we train a single transformer policy

\pi_{\phi}

that maps a window of recent states and the cost matrices

(Q,R)

to a near-optimal control

u_{t}

for every member of the family, with unseen systems handled by lightweight fine-tuning.

Seen Systems: 1. Simple Pendulum, 2. Two Link Arm, 3. Spring Damper, 4. Suspension System, 5. DC Motor, 6. Three Link Manipulator, 7. Differential Drive Robot, 8. SCARA Robot, 9. Omnidirectional Robot, 10. Cable Driven Robot, 11. Flexible Joint Robot, 12. Six DOF Manipulator, 13. Dual Arm Robot, 14. Double Integrator, 15. Lotka Volterra.
Unseen Systems: 16. Inverted Pendulum, 17. Segway Robot, 18. Asymmetric Oscillator, 19. Active Mass Damper, 20. Coupled Oscillators, 21. Damped Oscillator, 22. Triple Mass Spring, 23. Electromechanical Actuator, 24. Thermal System, 25. Fluid Tank,
26. Vibrating Beam, 27. Motor Generator, 28. Mechanical Linkage.

TABLE I: Benchmark systems

Data:

time step

\Delta t{=}0.02

s , rollout length

T{=}1250

1.5{\times}10^{6}

sequences/system , perturbation

\delta{=}{\pm}10\%

Architecture:

window

w{=}12

, max. dimensions

n_{x}^{\mathrm{max}}{=}12

n_{u}^{\mathrm{max}}{=}6

, embedding width

d_{\mathrm{m}}{=}64

, heads

h{=}16

, blocks

L{=}4

, feed-forward width

d_{\mathrm{ff}}{=}256

Training:

epochs

N_{\mathrm{epochs}}{=}500

, batch size

|\mathcal{B}|{=}4096

, learning rate

\eta{=}5{\times}10^{-5}

, Cauchy scale

\xi{=}1

TABLE II: Hyperparameters

Fig. 2: Relative sub-optimality

\Delta^{(i)}

for each system in Table I .

Fig. 3: Performance on the dual arm robot (system #13, seen).

Fig. 4: Performance on the vibrating beam (system #26, unseen).

Fig. 5: Validation loss versus window length

w

Optimal Control

Optimal control studies how to choose actions over time to optimize performance while satisfying system dynamics and constraints.

See all 135 papers

An Optimal Control Approach To Transformer Training
Mar 10, 2026Kağan Akman, Naci Saldi, Serdar YükselTransformer Transformer Training Dynamics
Department of Mathematics, Bilkent University, Ankara, 06800, Turkey · Department of Mathematics and Statistics, Queen’s University, Kingston K7L 3N6, ON, Canada
Transformer-like Inference from Optimal Control
May 15, 2026Aditya Kudre, Heng-Sheng Chang, Prashant G. MehtaTransformer Inference Decoder-Only Language Models
Coordinated Science Laboratory · Electrical and Computer Engineering · University of Illinois Urbana-Champaign +1
From Policy Evaluation to Off-Policy Control: Provable Emergence of In-Context Q-Learning
Oct 13, 2026Yingjun Mo, Yajie Zhao, Fang KongIn-Context RL Transformer-Based RL
Southern University of Science and Technology (SUSTech) · University of Electronic Science and Technology of China (UESTC)

Mar 10, 2026J/K move · Enter open · S save

An Optimal Control Approach To Transformer Training

Kağan Akman, Naci Saldi, Serdar Yüksel

Department of Mathematics, Bilkent University, Ankara, 06800, Turkey · Department of Mathematics and Statistics, Queen’s University, Kingston K7L 3N6, ON, Canada

Transformer Transformer Training Dynamics Optimal Control

Open paper AskPDF arXiv

In this paper, we develop a rigorous optimal control-theoretic approach to Transformer training that respects key structural constraints such as (i) realized-input-independence during execution, (ii) the ensemble control nature of the problem, and (iii) positional dependence. We model the Transformer architecture as a discrete-time controlled particle system with shared actions, exhibiting noise-free controlled McKean--Vlasov dynamics, and we show that lifting the dynamics to probability measures produces a fully-observed Markov decision process (MDP). Positional encodings are incorporated into the state space to preserve the sequence order under lifting. Using the dynamic programming principle, we establish the existence of globally optimal policies under mild assumptions of compactness. We further prove that closed-loop policies in the lifted problem is equivalent to an initial-distribution dependent open-loop policy, which are realized-input-independent and compatible with standard Transformer training. To train a Transformer, we propose a triply quantized training procedure for the lifted MDP by quantizing the state space, the space of probability measures, and the action space, and show that any optimal policy for the triply quantized model is near-optimal for the original training problem. Finally, we establish stability and empirical consistency properties of the lifted model by showing that the value function is continuous with respect to the perturbations of the initial empirical measures and convergence of policies as the data size increases. This approach provides a globally optimal and robust theoretical alternative to gradient-based training without requiring smoothness or convexity.

Abstract

Optimal Control

Explore similar work

An Optimal Control Approach To Transformer Training

Transformer-like Inference from Optimal Control

From Policy Evaluation to Off-Policy Control: Provable Emergence of In-Context Q-Learning