Optimal Control as Probabilistic Inference Review

This note is intended to be a reference for personal learning and implementation. It covers literature on path integral method, optimal control, reinforcement learning as a probabilistic inference.

How to solve optimal deterministic control problems in the absence of noise:

• PMP: Pontryagin Minimization Principle
• HJB: Hamilton-Jacobi-Bellman equation How to solve optimal stochastic control problems in the presence of noise:
  • PMP: difficult to solve
  • Stochastic HJB: the curse of dimensionality

Kappen(2005) studied a restrict class of optimal stochastic control problem.

• Dynamics \[ dx = b(x, t) dt + u dt + dw \] where \(dw\) is a Wiener process with \(\langle dw_i, dw_j \rangle = v_{ij} dt\) and \(\nu_{ij}\) is independent of \(x\), \(u\), and \(t\).

  Note: \(\langle dw_i, dw_j \rangle = v_{ij} dt\) defines the covariance structure of the noise. The matrix \(V = [v_{ij}]\) specifies how noise components correlate: diagonal entries are variances, off-diagonal entries are cross-correlations. The independence of \(V\) from state, control, and time is what makes this problem linearly solvable.

• Minimize the following cost function: \[ C(x, u, t) = \mathbb{E}_{x}[\phi(x(t_f) + \int_{t}^{t_f} \big [d\tau \frac{1}{2} u^{\top} R u + q(x, \tau) \big] \] where \(q(x, t)\) is a state dependent potential function.
• Important conclusions: \[ \psi(x_0, t) = \int [\mathrm{d}\xi]_{x_0} \big (- \frac{1}{\lambda} S(\xi) \big) \] where \(V(x, t) = - \lambda \log \psi(x, t)\), \(\int [\mathrm{d}x]_{x_0}\) means an integral over all paths \(x\) that state at \(x_0\) and \[ S(\xi) = \phi(x_f) + \int_{t}^{t_f} d\tau \big ( \frac{1}{2} u^{\top} R u + q(x, \tau) \big ) \]

Note:

• I have changed the symbols to make them more consistent with other papers.
• The dynamics is fully actuated, which means all state variables can be changed by \(u\) or the noise. Theodorou et al.(2010) made some extension, where the dynamics can be under-actuated and only actuated states have noise.

Todorov(2006) studies a special class of MDP problems, the control \(u\in \mathbb{R}^{|\mathcal{S}|}\) is a real-valued vector with dimensionality equal to the number of discrete states. The elements \(u^j\) of \(u\) have the effect of directly modifying the transition probabilities of an uncontrolled Markov chain.

• The controlled transition probabilities \[ p_{ij}(u) = \bar{p}_{ij} {\color{red} \exp(u_j)} \] where \(\bar{p}_{ij}\) is the transition probabilities of an uncontrolled Markov chain, which can be modified directly by \(\exp(u_j)\).

• The one step cost

  \begin{eqnarray*} l(i, u) &=& q(i) + r(i, u) \\ &=& q(i) + {\color{red} \mathrm{KL}(p_i(u)||\bar{p}_i)} \end{eqnarray*}

  Note that The control cost is defined as the KL divergence between the controlled and uncontrolled transition probabilities.

• Bellman equation \[ v(i) = \min_{u \in \mathcal{U}} \{l(i, u) + \sum_j p_{ij} v(j) \} \]

Path integral method was further developed by Theodorou et al.(2010), which studies a class of stochastic optimal control problem, where

• System dynamics:

  \begin{equation} \label{eq:pi-dynamics} \dot{x} = f(x, t) + G(x)(u + \varepsilon) \end{equation}

  where \(\varepsilon\) is Gaussian noise with variance \(\Sigma_{\varepsilon}\).

  Note:

  • The noise term has to be in the control or the directly controlled state, otherwise, the method doesn't apply.
  • It is sometime referred as 'linear in control'.

• Immediate cost function:

  \begin{equation} \label{eq:pi-cost} r_t(x, u) = q_t(x) + \frac{1}{2} u^{\top} R u \end{equation}

  Note:

  • The immediate cost can be split into a state cost and a control dependent cost.
  • It is sometimes referred as 'quadratic'.
• Finite horizon cost function: \[ R(\tau) = \Phi(t_N) + \int_{t_i}^{t_N} r_t(x, u) dt \] where \(\Phi\) is the terminal cost function.
• Value function \[ V(x_{t_i}) = \min_{u_{t}} \mathbb{E}_{\tau}[R(\tau)] \] the expectation of is taken over all possible trajectories, \(\tau\), starting at \(x_{t_i}\)

• The stochastic HJB equation is:

  \begin{equation} \label{eq:sto-hjb} \partial_t V_t = \min_{u} \big ( r_t + (\nabla_x V_t)^{\top} (f + G u) + \frac{1}{2} Tr (\nabla_{xx} V_t) G_t \Sigma_{\varepsilon}G_t^{\top}) \big ) \end{equation}

  which is a diffusion process.

• The Hamiltonian for the stochastic process: \[ H := r_t + (\nabla_x V_t)^{\top} (f + G u) + \color{red}{\frac{1}{2} Tr (\nabla_{xx} V_t) G_t \Sigma_{\varepsilon}G_t^{\top})} \] compared with deterministic process, the difference is the red term, which is from the noise.

  • The optimal control \(u^*\) is given by (set the gradient of Hamiltonian w.r.t. control to zero): \[ u^* = -R^{-1} G_t^{\top}(\nabla_x V_t) \]

    substitute it into the stochastic HJB Eq. \eqref{eq:sto-hjb}, and then use an exponential transformation: \[ V_t = - \lambda \log \Psi_t \] where \(\lambda\) is a scalar.

    Furthermore, set R to be inverse proportional to the noise variance as \(\lambda R^{-1} = \Sigma_{\varepsilon}\), so that \[ \lambda G_t R^{-1} G^{\top} = G_t \Sigma_{\varepsilon}G_t^{\top} = \Sigma(x_t) := \Sigma_t \] we get \[ -\partial_t \Psi_t = - \frac{1}{\lambda}q_t\Psi_t + f^{\top}(\nabla_x\Psi_t) + \frac{1}{2}Tr\big( (\nabla_{xx}\Psi_t)G_t\Sigma_{\varepsilon}G_t^{\top}\big) \] with boundary condition: \(\Psi_{t_N} = \exp(-\frac{1}{\lambda}\Phi_N)\). This partial differential equation (PDE) corresponds to so called Chapman Kolmogorov PDE. One step further, we can apply Feynman-Kac theorem to get one of the major conclusion:

    \begin{equation} \label{eq:pi-psi} \color{red}{\Psi_{t_i} = \mathbb{E}_{\tau_i} \big[ \exp(-\frac{\Phi_N + \int_{t_i}^{t_N} q_t dt}{\lambda}) \big]} \end{equation}

    where \(\tau_i :=(x_{t_i}, \ldots, x_{t_N})\) is a sample path starting at state \(x_{t_i}\).

    Note:

    1. Since \(V_t = -\lambda \log(\Psi_t)\). If we can get \(\Psi_t\), we can get \(V_t\) and thus solve optimal control problem. To get the value function, we don't need to solve the HJB but We can approximate it using forward path integral!
    2. We have replace the control with optimal control, so we don't need to solve it, explicitly.
    3. It is still hard to solve Eq. \eqref{eq:pi-psi}, but we can get its approximation by sampling.
    4. Regarding the simplification \(\lambda R^{-1} = \Sigma_{\varepsilon}\), it couples the control cost with the system dynamics. This assumption transforms the Gaussian probability for state transitions into a quadratic command cost.
    5. In Sutton & Barto(2018), \(\lambda\) is referred as temperature. High temperatures cause the actions to be all (nearly) equiprobable. Low temperatures cause a greater difference in selection probability for actions that differ in their value estimates.

Kappen(2012) established the link between optimal control and probabilistic inference in a clear way.

The optimal control is to minimize the following KL divergence:

where \(p\) is the probability of controlled trajectory, \(q\) is the probability of uncontrolled trajectory, \(\tau = x^{0:T}\), \(S(\tau) = \sum_{t=0}^T C^x(x, t)\), and \(C^x\) is the state dependent cost.

Because

we can rewrite the cost function as: \[ \hat{R}(x^t, u^t, x^{t+1}, t) = \log(\frac{p^t(x^{t+1}| x^t, u^t)}{q^t(x^{t+1}| x^t)}) + C^x(x^t, t) \quad t=0,\ldots,T-1 \] and \[ \hat{R} (x^T, u^T, x^{T+1}, T) = C^x(x^T, T) \]

The result of this KL minimization yields the "Boltzman distribution"

and the optimal cost: \[ C(x^0, p^*) = - \log(Z(x^0)), \] where \(Z(x^0) = \sum_\tau \psi(\tau)\) is a normalization constant.

The optimal control in the current state \(x^0\) is given by \[ p(x^1| x^0) = \sum_{x^{2:T}} p(x^{1:T}| x^0) \]

If we define

where \(C^u\) is the control dependent cost. Eq. \eqref{eq:kappen-kl} becomes

As we can see that \eqref{eq:kappen-kl} is a special form of \eqref{eq:rl-maxent}, where the reward can be partitioned into a state dependent term and a control dependent term.

Eq \eqref{eq:link-q} shows the probability of uncontrolled trajectory is a function of 'control' cost, which sounds wired. But because the noise term is in the control, you can think it as 'noise' cost.

\(\tilde{S}\) is defined as:

and \[ H_{t_j} := G_{t_j}^c R^{-1}G_{t_j}^{c\top} \]

For many systems, when \(dt \to 0\), \[ \frac{x_{t_{j+1}}^c - x_{t_j}^c}{dt} - f_{t_j}^c \to G^c u \]

For some system, when \({G^c_t}^{\top} = {G^c_t}^{-1}\) and \(dt \to 0\) \[ \tilde{s} = \Phi_{t_N} + \int_{t_i}^{t_N} r_t dt + C \] where \(C\) is a constant and it can then be canceled in Eq. \eqref{eq:pi-opt-control}.