GTSAM Posts

Authors: Gerry Chen, Yetong Zhang, and Frank Dellaert

$\DeclareMathOperator*{\argmin}{argmin} \newcommand{\coloneqq}{\mathrel{:=}}$

Introduction

Factor graph structure. The objective factors are marked with dashed lines, and the constrain factors are marked with solid lines. — **Figure 1** Factor graph structure for an LQR problem with 3 time steps. The cost factors are marked with dashed lines and the dynamics constraint factors are marked with solid lines.

In this post we explain how optimal control problems can be formulated as factor graphs and solved by performing variable elimination on the factor graph.

Specifically, we will show the factor graph formulation and solution for the Linear Quadratic Regulator (LQR). LQR is a state feedback controller which derives the optimal gains for a linear system with quadratic costs on control effort and state error.

We consider the finite-horizon, discrete LQR problem. The task is to find the optimal controls $u_k$ at time instances $t_k$ so that a total cost is minimized. Note that we will later see the optimal controls can be represented in the form $u^*_k = K_kx_k$ for some optimal gain matrices $K_k$ . The LQR problem can be represented as a constrained optimization problem where the costs of control and state error are represented by the minimization objective $\eqref{eq:cost}$ , and the system dynamics are represented by the constraints $\eqref{eq:dyn_model}$ .

$\begin{equation} \argmin\limits_{u_{1\sim N}} ~~x_N^T Q x_N + \sum\limits_{k=1}^{N-1} x_k^T Q x_k + u_k^T R u_k \label{eq:cost} \end{equation}$ $\begin{equation} s.t. ~~ x_{k+1}=Ax_k+Bu_k ~~\text{for } k=1 \text{ to } N-1 \label{eq:dyn_model} \end{equation}$

We can visualize the objective function and constraints in the form of a factor graph as shown in Figure 1. This is a simple Markov chain, with the oldest states and controls on the left, and the newest states and controls on the right. The ternary factors represent the dynamics model constraints and the unary factors represent the state and control costs we seek to minimize via least-squares.

Variable Elimination

To optimize the factor graph, which represents minimizing the least squares objectives above, we can simply eliminate the factors from right to left. In this section we demonstrate the variable elimination graphically and algebraically, but the matrix elimination is also provided in the Appendix.

**Figure 2a** Elimination of state $x_2$

**Figure 3a** Elimination of state $u_1$

❮ ❯

Eliminate a State

Let us start at the last state, $x_2$ . Gathering the two factors (marked in red Figure 2a), we have $\eqref{eq:potential}$ the objective function $\phi_1$ , and $\eqref{eq:constrain}$ the constraint equation on $x_2$ , $u_1$ and $x_1$ :

$\begin{equation} \phi_1(x_2) = x_2^T Q x_2 \label{eq:potential} \end{equation}$

$\begin{equation} x_2 = Ax_1 + Bu_1 \label{eq:constrain} \end{equation}$

By substituting $x_2$ from the dynamics constraint $\eqref{eq:constrain}$ into the objective function $\eqref{eq:potential}$ , we create a new factor representing the cost of state $x_2$ as a function of $x_1$ and $u_1$ :

$\begin{equation} \phi_2(x_1, u_1) = (Ax_1 + Bu_1)^T Q (Ax_1 + Bu_1) \label{eq:potential_simplified} \end{equation}$

The resulting factor graph is illustrated in Figure 2b. Note that the dynamics constraint is now represented by the bayes net factors shown as gray arrows.

To summarize, we used the dynamics constraint to eliminate variable $x_2$ and the two factors marked in red, and we replaced them with a new binary cost factor on $x_1$ and $u_1$ , marked in blue.

next >>

Eliminate a Control

To eliminate $u_1$ , we seek to replace the two factors marked red in Figure 3a with a new cost factor on $x_1$ and an equation for the optimal control $u_1^*(x_1)$ .

Adding the control cost to $\eqref{eq:potential_simplified}$ , the combined cost of the two red factors in Figure 3a is given by:

$\begin{equation} \phi_3(x_1, u_1) = u_1^TRu_1 + (Ax_1 + Bu_1)^T Q (Ax_1 + Bu_1) \label{eq:potential_u1} \end{equation}$

$\phi_3$ is sometimes referred to as the optimal action value function and we seek to minimize it over $u_1$ . We do so by setting the derivative of $\eqref{eq:potential_u1}$ wrt $u_1$ to zero yielding the expression for the optimal control input $u_1^*$ as

$\begin{align} u_1^*(x_1) &= \argmin\limits_{u_1}\phi_3(x_1, u_1) \nonumber \\ &= -(R+B^TQB)^{-1}B^TQAx_1 \label{eq:control_law} \\ &= K_1x_1 \nonumber \end{align}$

where $K_1\coloneqq -(R+B^TQB)^{-1}B^TQA$ .

Finally, we substitute the expression of our optimal control, $u_1^* = K_1x_1$ , into our potential $\eqref{eq:potential_u1}$ to obtain a new unary cost factor on $x_1$ :

$\begin{align} \phi_4(x_1) &= \phi_3(x_1, u_1^*(x_1)) \nonumber \\ &= (K_1x_1)^T RK_1x_1 + (Ax_1 + BKx_1)^T Q (Ax_1 + BKx_1) \nonumber \\ &= x_1^T(A^TQA-K_1^TB^TQA)x_1 \label{eq:potential_x1} \end{align}$ Note that we simplified $K_1^TRK_1 + K_1^TB^TQBK_1 = -K_1^TB^TQA$ by substituting in for $K_1$ using $\eqref{eq:control_law}$ .

The resulting factor graph is illustrated in Figure 3b.

For convenience, we will also define $P_k$ where $x_k^TP_kx_k$ represents the aggregate of the two unary costs on $x_k$ . In the case of $P_1$ , $\begin{align} x_1^TP_1x_1 &= x_1^TQx_1 + \phi_4(x_1) \nonumber \end{align}$ is the aggregation of the two unary factors labeled in green in Figure 3c.

<< prevnext >>

Turning into a Bayes Network

By eliminating all the variables from right to left, we can get a Bayes network as shown in Figure 4d. Each time we eliminate a state and control, we simply repeat the steps in Eliminate a state and Eliminate a control: we express the state $x_{k+1}$ with the dynamics model, then find the optimal control $u_k$ as a function of state $x_k$ .

Eliminating a general state, $x_{k+1}$ , and control $u_k$ , we obtain the recurrence relations:

$\begin{equation} \boxed{K_k = -(R+B^TP_{k+1}B)^{-1}B^TP_{k+1}A} \label{eq:control_update_k} \end{equation}$

$\begin{equation} \boxed{P_k = Q+A^TP_{k+1}A - K_k^TB^TP_{k+1}A} \label{eq:cost_update_k} \end{equation}$

with $P_{N}=Q$ is the cost at the last time step.

The final Bayes net in Figure 4d shows graphically the optimal control law: $\begin{equation} \boxed{u^*_k = K_k x_k} \end{equation}$

<< prev

Intuition

Comparison between LQR control as solved by factor graphs and by the Ricatti Equation. (they are the same) — **Figure 5** Example LQR control solutions as solved by factor graphs (middle) and the traditional Discrete Algebraic Ricatti Equations (right). The optimal control gains and cost-to-go factors are compared (left). All plots show exact agreement between factor graph and Ricatti equation solutions.

We introduce the cost-to-go (also known as return cost, optimal state value function, or simply value function) as $V_k(x) \coloneqq x^TP_kx$ which intuitively represents the total cost that will be accrued from here on out, assuming optimal control.

In our factor graph representation, it is becomes obvious that $V_k(x)$ corresponds to the total cost at and after the state $x_k$ assuming optimal control because we eliminate variables backwards in time with the objective of minimizing cost. Eliminating a state just re-expresses the future cost in terms of prior states/controls. Each time we eliminate a control, $u$ , the future cost is recalculated assuming optimal control (i.e. $\phi_4(x) = \phi_3(x, u^*)$ ).

This “cost-to-go” is depicted as a heatmap in Figure 5. The heat maps depict the $V_k$ showing that the cost is high when $x$ is far from 0, but also showing that after iterating sufficient far backwards in time, $V_k(x)$ begins to converge. That is to say, the $V_0(x)$ is very similar for $N=30$ and $N=100$ . Similarly, the leftmost plot of Figure 5 depicts $K_k$ and $P_k$ and shows that they (predictably) converge as well.

This convergence allows us to see that we can extend to the infinite horizon LQR problem (continued in the next section).

Equivalence to the Ricatti Equation

In traditional descriptions of discrete, finite-horizon LQR (i.e. Chow, Kirk, Stanford), the control law and cost function are given by

$u_k = K_kx_k$

$\begin{equation} K_k = -(R+B^TP_{k+1}B)^{-1}B^TP_{k+1}A \label{eq:control_update_k_ricatti} \end{equation}$

$\begin{equation} P_k = Q+A^TP_{k+1}A - K_k^TB^TP_{k+1}A \label{eq:cost_update_k_ricatti} \end{equation}$

with $P_k$ commonly referred to as the solution to the dynamic Ricatti equation and $P_N=Q$ is the value of the Ricatti function at the final time step. $\eqref{eq:control_update_k_ricatti}$ and $\eqref{eq:cost_update_k_ricatti}$ correspond to the same results as we derived in $\eqref{eq:control_update_k}$ and $\eqref{eq:cost_update_k}$ respectively.

Recall that $P_0$ and $K_0$ appear to converge as the number of time steps grows. They will approach a stationary solution to the equations

$\begin{align} K &= -(R+B^TPB)^{-1}B^TPA \nonumber \\ P &= Q+A^TPA - K^TB^TPA \nonumber \end{align}$

as $N\to\infty$ . This is the Discrete Algebraic Ricatti Equations (DARE) and $\lim_{N\to\infty}V_0(x)$ and $\lim_{N\to\infty}K_0$ are the cost-to-go and optimal control gain respectively for the infinite horizon LQR problem. Indeed, one way to calculate the solution to the DARE is to iterate on the dynamic Ricatti equation.

Implementation using GTSAM

**Edit (Apr 17, 2021): Code updated to new Python wrapper as of GTSAM 4.1.0.

You can view an example Jupyter notebook on google colab or download the modules/examples that you can use in your projects to:

Calculate the closed loop gain matrix, K, using GTSAM
Calculate the “cost-to-go” matrix, P (which is equivalent to the solutions to the dynamic Ricatti equation), using GTSAM
Calculate the LQR solution for a non-zero, non-constant goal position, using GTSAM
Visualize the cost-to-go and how it relates to factor graphs and the Ricatti equation
and more!

A brief example of the open-loop finite horizon LQR problem using factor graphs is shown below:

def solve_lqr(A, B, Q, R, X0=np.array([0., 0.]), num_time_steps=500):
    '''Solves a discrete, finite horizon LQR problem given system dynamics in
    state space representation.
    Arguments:
        A, B: nxn state transition matrix and nxp control input matrix
        Q, R: nxn state cost matrix and pxp control cost matrix
        X0: initial state (n-vector)
        num_time_steps: number of time steps, T
    Returns:
        x_sol, u_sol: Txn array of states and Txp array of controls
    '''
    n = np.size(A, 0)
    p = np.size(B, 1)

    # noise models
    prior_noise = gtsam.noiseModel.Constrained.All(n)
    dynamics_noise = gtsam.noiseModel.Constrained.All(n)
    q_noise = gtsam.noiseModel.Gaussian.Information(Q)
    r_noise = gtsam.noiseModel.Gaussian.Information(R)

    # Create an empty Gaussian factor graph
    graph = gtsam.GaussianFactorGraph()

    # Create the keys corresponding to unknown variables in the factor graph
    X = []
    U = []
    for k in range(num_time_steps):
        X.append(gtsam.symbol('x', k))
        U.append(gtsam.symbol('u', k))

    # set initial state as prior
    graph.add(X[0], np.eye(n), X0, prior_noise)

    # Add dynamics constraint as ternary factor
    #   A.x1 + B.u1 - I.x2 = 0
    for k in range(num_time_steps-1):
        graph.add(X[k], A, U[k], B, X[k+1], -np.eye(n),
                  np.zeros((n)), dynamics_noise)

    # Add cost functions as unary factors
    for x in X:
        graph.add(x, np.eye(n), np.zeros(n), q_noise)
    for u in U:
        graph.add(u, np.eye(p), np.zeros(p), r_noise)

    # Solve
    result = graph.optimize()
    x_sol = np.zeros((num_time_steps, n))
    u_sol = np.zeros((num_time_steps, p))
    for k in range(num_time_steps):
        x_sol[k, :] = result.at(X[k])
        u_sol[k] = result.at(U[k])
    
    return x_sol, u_sol

Future Work

The factor graph (Figure 1) for our finite horizon discrete LQR problem can be readily extended to LQG, iLQR, DDP, and reinforcement learning using non-deterministic dynamics factors, nonlinear factors, discrete factor graphs, and other features of GTSAM (stay tuned for future posts).

Appendix

Least Squares Implementation in GTSAM

GTSAM can be specified to use either of two methods for solving the least squares problems that appear in eliminating factor graphs: Cholesky Factorization or QR Factorization. Both arrive at the same result, but we will take a look at QR since it more immediately illustrates the elimination algorithm at work.

QR Factorization

factor graph partially eliminated — **Figure 6a** Initial factor graph and elimination matrix

❮ ❯

$P_{k}$	$=$	$Q + A^TP_{k+1}A - K_k^TB^TP_{k+1}A$
$D_{k}$	$=$	$R + B^TP_{k+1}B$
$K_k$	$=$	$-D_{k}^{-1/2}(R + B^TP_{k+1}B)^{-T/2}B^TP_{k+1}A$
	$=$	$-(R + B^TP_{k+1}B)^{-1}B^TP_{k+1}A$

The factorization process is illustrated in Figure 6 for a 3-time step factor graph, where the noise matrices and elimination matrices are shown with the corresponding states of the graph. The noise matrix (NM) is $0$ for a hard constraint and $I$ for a minimization objective. The elimination matrix is formatted as an augmented matrix $[A|b]$ for the linear least squares problem $\argmin\limits_x||Ax-b||_2^2$ with ${x=[x_2;u_1;x_1;u_0;x_0]}$ is the vertical concatenation of all state and control vectors. The recursive expressions for $P$ , $D$ , and $K$ when eliminating control variables (i.e. $u_1$ in Figure 6e) are derived from block QR Factorization.

Note that all $b_i=0$ in the augmented matrix for the LQR problem of finding minimal control to reach state $0$ , but simply changing values of $b_i$ intuitively extends GTSAM to solve LQR problems whose objectives are to reach different states or even follow trajectories.