DeSKO: Stability-Assured Robust Control of Nonlinear Systems with a Deep Stochastic Koopman Operator

Soft Robotics Lab, ETH Zürich  ·  Harbin Institute of Technology
ICLR 2022
DeSKO architecture: a state is encoded as a distribution of observables, propagated by Koopman and control matrices, and mapped back to state space.

DeSKO models nonlinear uncertainty as a distribution of Koopman observables, then uses that distribution for efficient robust control with a closed-loop stability guarantee.

Abstract

The Koopman operator theory linearly describes nonlinear dynamical systems in a high-dimensional functional space and allows linear control methods to be applied to highly nonlinear systems. However, the Koopman operator does not account for uncertainty, which limits its performance in real-world applications.

We propose a deep stochastic Koopman operator (DeSKO) in a robust learning-control framework for nonlinear stochastic systems. DeSKO captures uncertainty by inferring a distribution of observables, which is then used to design a robust, stabilizing closed-loop controller. Across advanced control benchmarks—including a soft robotic arm, a legged robot, and a biological gene regulatory network—DeSKO is more robust and scalable than prior deep Koopman and reinforcement-learning methods. It also resists previously unseen external disturbances with magnitudes up to five times the maximum control input.

Method

1

Encode uncertainty

Two neural networks map the current state to the mean and diagonal variance of a distribution over observables.

2

Propagate linearly

Learned Koopman and control matrices recursively predict the observable distribution over the control horizon.

3

Control robustly

Robust MPC exploits the linear representation while providing stability guarantees under noise and residual error.

Essential Equations

From a noisy state to a stability-assured control input.

1Encode
\[p_\theta(\psi_t\mid x_t)=\mathcal{N}\!\left(\mu_\theta(x_t),\,\sigma_\theta(x_t)\right),\qquad \psi_t^{,i}=\mu_\theta(x_t)+\epsilon_i\sigma_\theta(x_t),\ \epsilon_i\sim\mathcal{N}(0,I)\]

The encoder maps the state to a distribution in the lifted observable space; reparameterized samples make the model trainable end to end.

2Propagate & decode
\[\psi_{t+k}^{,i}=A\psi_{t+k-1}^{,i}+Bu_{t+k-1},\qquad \hat{x}_{t+k}=C\psi_{t+k}^{,i}\]

The learned Koopman matrix \(A\) and control matrix \(B\) propagate each observable sample linearly; \(C\) maps it back to the original state.

3Learn
\[\min_{A,B,C,\theta}\ \mathbb{E}_{\mathcal D}\sum_{k=1}^{H}\sum_{i=1}^{N} \left\|x_{t+k}-C\psi_{t+k}^{,i}\right\|\]

A multi-step prediction objective trains the encoder, Koopman dynamics, and decoder jointly over horizon \(H\), reducing recursive rollout error.

4Control
\[\hat\mu_{t+1}=A\hat\mu_t+Bc_t,\qquad u_t=c_t^*+K\!\left(\mu_\theta(x_t)-\hat\mu_t\right)\]

MPC plans the nominal input \(c_t^*\) in the linear latent dynamics, while LQR feedback \(K\) corrects deviations caused by uncertainty.

Stability guarantee. With bounded noise and bounded Koopman approximation residuals, the closed loop is uniformly ultimately bounded by \[\frac{\beta\sigma(Lb+\gamma)}{1-\beta}+\eta.\]

Key Results

Evaluated on eight settings spanning CartPole, gene regulatory networks, HalfCheetah, and the SoPrA soft arm.

Robust to 5× control authority

DeSKO achieved the best robustness in every disturbance test. In key CartPole settings, it resisted persistent 100 N disturbances without failure—even though the maximum control input was only 20 N.

100 trials at every disturbance magnitude.

CartPole death rate under increasing disturbances for DeSKO, DKO, and SAC.
CartPole with observation noise. Lower death rate is better.
HalfCheetah velocity tracking for DeSKO, DKO, SAC, and the reference trajectory.
Scalable control. DeSKO remains competitive on high-dimensional HalfCheetah locomotion.
Real SoPrA soft-arm endpoint tracking a circular trajectory.
Real-world tracking. DeSKO accurately controls the physical SoPrA soft robotic arm.

Contributions

Stochastic Koopman learning. DeSKO captures uncertainty directly in the lifted observable space while preserving linear propagation.

Efficient robust MPC. The controller uses distributional predictions without giving up the computational benefits of Koopman linearity.

Closed-loop stability. The analysis proves uniform ultimate boundedness under stochastic noise and Koopman approximation error.

BibTeX

@inproceedings{han2022desko,
  title     = {DeSKO: Stability-Assured Robust Control of Nonlinear Systems with a Deep Stochastic Koopman Operator},
  author    = {Han, Minghao and Euler-Rolle, Jacob and Katzschmann, Robert K.},
  booktitle = {International Conference on Learning Representations},
  year      = {2022}
}