Markov Chain Monte Carlo for Koopman-based Optimal Control: Technical Report

This paper address the optimal control of nonlinear systems that have reasonably accurate finite-dimensional representations in terms of their Koopman operator [20]. While Koopman’s pioneering work is almost a century old, its use as a practical tool to model complex dynamics is much more recent and only became practical when computational tools became available for the analysis of systems with hundreds to thousands of dimensions. The use of Koopman models for control is even more recent, but has attracted significant attention in the last few years [9, 21, 28, 23, 6, 27].

The linear structure of the Koopman representation permits very efficient solutions for optimal control when the lifted dynamics are linear on the control input, as in [9, 21, 28]. However, linearity in the control severely limits the class of applicable dynamics. Bilinear representations are more widely applicable [23, 6, 27], but are also harder to control.

When the set of admissible control inputs is finite, the Koopman representation can be viewed as a switched linear system, where the optimal control selects, at each time step, one out of several admissible dynamics [7]. The optimal control of switched system it typically difficult [5, 31], but when the optimization criterion is linear in the lifted state, the dynamic programming cost-to-go is concave and piecewise linear, with a simple representation in terms of a minimum over a finite set of linear functions. This observation enabled the design of efficient algorithms that combine dynamic programming with dynamic pruning [7].

This paper exploits the structure of the Koopman representation to develop efficient Markov Chain Monte Carlo (MCMC) sampling method for optimal control. Following the pioneering work of [19, 10, 14], we draw samples from a Boltzmann distribution with energy proportional to the criterion to minimize. The original use of MCMC methods for combinatorial optimization relied on a gradual decrease in temperature, now commonly known as simulated annealing, to prevent the chain from getting trapped into states that do not minimize energy. An alternative approach relies on the use of multiple replicas of a Markov chain, each generating samples for a different temperature. The introduction of multiple replicas of a Markov chain to improve the mixing time can be traced back to [30]. The more recent form of parallel tempering (also known as Metropolis–coupled MCMC, or exchange Monte Carlo) is due to [15, 17].

Our key contribution is an MCMC algorithm that combines Gibbs sampling [14, 13] with parallel tempering to solve the switched linear optimizations that arise from the Koopman representation of nonlinear optimal control problems. This algorithm is computationally efficient due to the combination of two factors: (i) the linear structure of the cost function enables the full variable sweep need for Gibbs sampling to be performed with computation that scales linearly with the horizon length $T$ and (ii) parallel tempering can be fully parallelized across computation cores, as noted in [11]. With regard to (i), the computational complexity of evaluation Gibbs’ conditional distribution for each optimization variable is of order $Tn_{\psi}^{2}$, where $n_{\psi}$ denotes the size of the lifted state. Our algorithm computes the conditional distributions for all $T$ variables in a full Gibbs’ sweep with computation still just of order $Tn_{\psi}^{2}$. With regard to (ii), while here we only explore parallelization across CPU cores, in the last few years hardware parallelization using GPUs and FPGAs has achieved orders of magnitude increase in the number of samples generated with MCMC sampling [25, 3].

The remaining of this paper is organized as follows: Section II shows how a nonlinear control problem can be converted into a switching linear systems optimization, using the Koopman operator. Section III provides basic background on MCMC, Gibbs sampling, and parallel tempering. Our optimization algorithm is described in Section IV and its use is illustrated in Section V in the context of a video game.

Given a discrete-time nonlinear system of the form

with the time $t$ taking values over $\mathcal{T}\coloneq\{1,\dots,T\}$, our goal is to solve a final-state optimal control problem of the form

where $u\coloneq(u_{1},\dots,u_{T})\in\mathcal{U}\coloneq\mathcal{U}_{1}\times\cdots\times\mathcal{U}_{T}$ denotes the control sequence to be optimized, which we assume finite but potentially with a large number of elements.

For each input $u\in\mathcal{U}_{t}$, $t\in\mathcal{T}$, the Koopman operator $K_{u}$ for the system (1) operates on the linear space of functions $\mathcal{F}$ from $\mathcal{X}$ to $\mathbb{R}$ and is defined by

Assuming there is a finite dimensional linear subspace $\mathcal{F}_{\mathrm{inv}}$ of $\mathcal{F}$ that is invariant for every Koopman operator in the family $\{K_{u}:\forall u\in\mathcal{U}_{t},\;t\in\mathcal{T}\}$, there is an associated family of matrices $\{A(u)\in\mathbb{R}^{n_{\psi}\times n_{\psi}}:\forall u\in\mathcal{U}_{t},\;t\in\mathcal{T}\}$ such that

where $\psi_{t}\coloneq\begin{bmatrix}\varphi_{1}(x_{t})&\cdots&\varphi_{n_{\psi}}(x_{t})\end{bmatrix}^{\prime}\in\mathbb{R}^{n_{\psi}}$ and the functions $\{\varphi_{1}(\cdot),\dots,\varphi_{n_{\psi}}(\cdot)\}$ form a basis for $\mathcal{F}_{\mathrm{inv}}$ [12, 7, 16]. If in addition, the function $g(\cdot)$ in (2) also belongs to $\mathcal{F}_{\mathrm{inv}}$, there also exists a row vector $c\in\mathbb{R}^{1\times n_{\psi}}$ such that

which enables re-writing the optimization criterion (2) as

We can thus view the original optimal control problem for the nonlinear system (1) as a switched linear control problem [7]. Note that it is always possible to make sure that (4) hold for some vector $c$ by including $g(x_{t})$ as one of the entries of $\psi_{t}$.

It is generally not possible to find a finite-dimensional subspace $\mathcal{F}_{\mathrm{inv}}$ that contains $g(\cdot)$ and is invariant for every $\{K_{u}:u\in\mathcal{U}_{t},\,t\in\mathcal{T}\}$. Instead, we typically work with a finite dimensional space for which (3)–(4) hold up to some error. However, to make this error small, we typically need to work with high-dimensional subspaces, i.e., large values of $n_{\psi}$. Quantifying the impact on the error of the lifted space dimension $n_{\psi}$ and of the size of the dataset used to learn the Koopman dynamics is challenging, but notable progress has been reported in [16, 26]. Much work remains, e.g., in quantifying errors for systems with isolated limit sets [4, 22].

To solve a combinatorial minimization of the form

over a finite set $\mathcal{U}$, it is convenient to consider the following Boltzmann distribution with energies $J(u)$, $u\in\mathcal{U}$:

for some constant $\beta\geq 0$. For consistency with statistical mechanics, we say that values of $\beta$ close to zero correspond to high temperatures, whereas large values correspond to low temperatures. The normalization function $Q(\beta)$ is called the canonical partition function.

For $\beta=0$ (high temperature), the Boltzmann distribution becomes uniform and all $u\in\mathcal{U}$ are equally probable. However, as we increase $\beta$ (lower the temperature) all probability mass is concentrated on the subset of $\mathcal{U}$ that minimizes the energy $J(u)$. This can be seen by noting that the ratio of the probability between two $u,\bar{u}\in\mathcal{U}$ is given by

which shows that the value with higher energy (and higher cost) will eventually never be selected.

This motivates a procedure to solve (6): draw samples from a random variable ${\bm{u}}_{\beta}$ with Boltzmann distribution given by (7) with $\beta$ sufficiently high (temperature sufficiently low) so that all samples correspond to states with the minimum energy/cost.

The following result confirms that the probability $\delta$ that we can extract $N$ samples from the Boltzmann distribution with energies much larger than the minimum $J^{*}$ can be made arbitrarily small by increasing $N$ in proportion to just the logarithm of $\delta$, provided that $\beta$ is sufficiently high (temperature is sufficiently low).

The number of samples $N$ required to achieve a small probability $\delta$ is a function of the fraction $|\mathcal{U}^{*}|/|\mathcal{U}|$ of elements of $\mathcal{U}$ that minimize $J(u)$. The result is particularly interesting in that the magnitude of $\beta$ and number of samples $N$ depends on the logarithm of this fraction, meaning that the size of $\mathcal{U}$ could grow exponentially with respect to some scaling variable (like the time horizon in optimal control) and yet $\beta$ and $N$ would only need to grow linearly. However, the optimism arising from this bound needs to be tempered by the observation that this result assumes that the samples are independent and this is typically hard to achieve for large values of $\beta$, as discussed below.

Our goal is to optimize a criterion of the form (5). In the context of MCMC sampling from the Boltzmann distribution (7), this corresponds to a Gibbs update for the control ${\bm{u}}^{\beta}_{t}[k+1]$ with a distribution (14) that can be computed using

and a tempering flip probability in (18) computed using

where

The following algorithm implements Gibbs sampling with tempering using recursive formulas to evaluate (19)–(20).