# Continuity Equation Based Convexification of Control Problems

In this note, I introduce an approach to convexify control problems for nonlinear systems. This method is due to [optimal transport theory](https://en.wikipedia.org/wiki/Transportation_theory_(mathematics)). However, one does not need optimal transport to understand or appreciate this convexification strategy.

**Continuity Equation**
Before we get to the control problems, lets review the continuity equation and its relation to ODEs and dynamical systems.

Consider an ordinary differential equation
\begin{align}
\dot{x} = f(x) \\
x(0) = x_0
\end{align}For general vector fields $f:\mathbb{R}^d \rightarrow \mathbb{R}^d$ this is a nonlinear differential equation. However, instead of looking at how the solution flows from an initial condition along the vector field, one can consider an alternative quantity, $\rho(y,t)$, the probability of finding a state $x(t)$ at a point $y$ at time $t$, $\mathbb{P}(x(t) \in dy) = \rho(y,t)dy$ of finding a system in a state $y$ at time $t$. This quantity is known to evolve according to a linear partial differential equation, the [continuity equation](https://):
\begin{align*}
\partial_t \rho +\nabla_x \cdot (f(x,u)\rho) = 0 \\
\rho(x,0) = \rho_0
\end{align*} where $\nabla_x : \sum_{i=1}^d \partial_{x_i}$ is the divergence operation. Even though original ODE is nonlinear the continuity equation is linear. This perspective has been very fruitful in studying problems with chaos, where it is more tractable to study the evolution of a colleciton of ensembles instead of a single initial condition.

Following plot shows the 2D projection of a [3D Lorentz system](https://en.wikipedia.org/wiki/Lorenz_system) and its corresponding density evolution.

|  |
|:-----------------------------------------------------------------------:|
| Evolution of particles in a flow |

## Optimal Controller Problems

Next, we consider how the continuity equation can be used to reformulate control problems. Consider the optimal control problem,
\begin{align}
&\inf_{x,u} \int_0^T L(x(t))dt + \int_0^T M(u(t))dt \\
&\dot{x} = f_0(x)+\sum_{i}u_i(t)f_i(x) := f(x,u) \\
& x(0) = x_0
\end{align}For general vector fields $f_i(x)$ this problem is non-convex even if $L$ and $M$ are convex.

The evolution of the probability density is given by the continuity equation as mentioned above.
\begin{align*}
\partial_t \rho +\nabla_x \cdot (f(x,u)\rho) = 0 \\
\rho(x,0) = \rho_0
\end{align*}We can now treat this continuity equation as our state equation instead of the ODE above.

Suppose, additionally instead of *open loop* controls $u_i(t)$ one looks for *feedback controls* $u_i(x,t)$, a natural extension of the optimal cost is the following
\begin{equation}
\inf_{\rho, u(x,t)} \int_0^T \int_{\mathbb{R}^d} L(x)\rho(x,t)dxdt + \int_0^T \int_{\mathbb{R}^d} M(u(x,t))\rho(x,t)dxdt
\end{equation} This cost function can be interpreted as the *expected cost* of states and controls as one is averaging the cost against the distribution $\rho(x,t)$ of states.

As such, this problem is not convex in the decision variables $u_i(x,t)$ and $\rho(x,t)$. However, there is a way to convexify it. Let us define the variables,
\begin{equation}
m_i(x,t) = u_i(x,t) \rho(x,t)
\end{equation}Then the optimization problem can be rewritten as following **convex problem**.
\begin{align*}
\inf_{\rho, u(x,t)} \int_0^T \int_{\mathbb{R}^d} L(x)\rho(x,t)dxdt + \int_0^T \int_{\mathbb{R}^d} M(\frac{m(x,t)}{\rho(x,t)})\rho(x,t)dxdt \\
\partial_t \rho +\nabla_x \cdot \Big (f_0(x) \rho(x,t)+\sum_i f_i(x) m_i(x,t) \Big ) = 0 \\
\rho(x,0) = \rho_0 \\
\rho(x,t) \geq 0
\end{align*} It can be checked that the functional $\int_0^T \int_{\mathbb{R}^d} M(\frac{m(x,t)}{\rho(x,t)})\rho(x,t)$ is convex in $m_i$ and $\rho$, whenever $M$ is convex. Moreover, the continuity equation constraint is linear in the decision variables. This was a groundbreaking observation to due [Benamou-Brenier](https://link.springer.com/article/10.1007/s002110050002) in optimal transport theory.

Hence, the control problem has been convexified.

## Optimal Observer Problem

Can we use a similar approach to filtering or observability?

The goal is construct the state of the system by observing the outputs of the state. Consider an optimization formulation of this observability problem. Given observations $y(t)$ solve
\begin{align}
\inf_{\tilde{x}(t)} \int_0^T |y(t)-\tilde{y}(t)|^2dt + \frac{\lambda}{2}|\tilde{x}(0)|^2\\
\dot{\tilde{x}} = f(\tilde{x},u) \\
\tilde{x}(0) = \tilde{x}_0 \\
\tilde{y}(t) = h(\tilde{x}(t))
\end{align}where $h:\mathbb{R}^d \rightarrow \mathbb{R}^m$ is the output map. Clearly, for general $f$ and $h$ this is a non-convex problem. Now, let's formulate this problem in terms of the continuity equation, by replacing $x(t)$ with $\rho(t,x)$ as the state variable.

In this case, similar to what we did for optimal control above, one can formulate the expected mean-square error cost.
\begin{align}
\inf_{\tilde{\rho}_0(x)} \int_0^T |y(t) - \int_{\mathbb{R}^d} h(x)\tilde{\rho}(x,t)dx|^2 dt + \frac{\lambda}{2} \int |x|^2\rho_0(x)dx
\end{align}subject to the continuity equation constraint.
\begin{align*}
\partial_t \tilde{\rho} +\nabla_x \cdot (f(x,u)\tilde{\rho}) = 0 \\
\rho(x,0) = \tilde{\rho}_0
\end{align*} The second term in the cost function is clearly convex. The first term is convex because it is composition of a convex functional with a linear one. Since the control is fixed, the continuity equation constraint remains convex without needing to define the variable $m_i= u_i\tilde{\rho}$. This convexifies the optimal observation problem.