Let \(x_t\in\mathbf{C}^{n_x}\) denote the state vector at time \(t\), and \(u_t\in\mathbf{C}^{n_u}\) denote the control signal at time \(t\). We describe the dynamics of the system according to

Moreover, we define a cost-to-go function as follows,

where \(V_{t:T}\) designates the cumulative cost in the trajectory starting at time \(t\) and ending at time \(T\), and \(\ell\) is a per-step cost function. As we mentioned previously, iLQR approximates the cost function with a second degree Taylor approximation. So, rather than requiring this approximation, we'll define a second-order cost function as follows:

where \(Q\in\mathbf{R}^{n_x\times n_x}\) and \(R\in\mathbf{R}^{n_u\times n_u}\) decide how to penalize trajectories based on their states and controls. With this quadratic parameterization of the cost function, we don't need to worry about the Taylor series approximation.

In order to optimize the cost of trajectories, iLQR makes local linear approximations of the dynamics of the system. To compute this, we analyze the perturbation of the next state \(\delta x_{t+1}\) due to a small perturbation in state \(\delta x_t\) and control \(\delta u_t\):

where \(\delta x_t\triangleq x - x_t\) and \(\delta u_t\triangleq u - u_t\). Here, the \(x,u\) refer to states and controls in a proposed trajectory, while \(x_t, u_t\) refer to states and controls in the trajectory that is being improved (referring to the pseudocode above, \(x, u\) would be in states and controls, while \(x_t, u_t\) would be in state_trajectory and control_trajectory). Also, noting that \(x_{t+1}\equiv f(x_t, u_t)\), we have

where \(A_t\in\mathbf{C}^{n_x\times n_x},B_t\in\mathbf{C}^{n_x,n_u}\) are the Jacobians of \(f\) with respect to \(x\) and \(u\) respectively, evaluated at \((x_t, u_t)\).

With that out of the way, we can think about how to optimize the cost-to-go of our trajectory. Firstly, note the recursive property of the cost to go:

This is called the Hamilton-Jacobi-Bellman equation, and it is beautiful (reinforcement learning folk are obligated to agree). The minimum over controls is taken here since the cost-to-go should represent the optimal cost of completing the trajectory from state \(x_t\). Keeping with reinforcement learning style notation, we'll define \(Q_{t:T}(x_t, u_t) = \ell(x_t, u_t) + V_{t+1:T}(f(x_t, u_t))\), so \(V_{t:T}(x_t) = \min_{u_t}Q_{t:T}(x_t, u_t)\). In order to minimize the cost-to-go with respect to controls, we approximate the loss function by a quadratic, as mentioned previously. Now we'll analyze how the cost-to-go is perturbed due to a small change in \(x_t\), using a second order Taylor expansion (note, the Jacobians will be very simple due to the quadratic costs):

To compute the Jacobians of \(V_{t:T}\), we must compute the Jacobians of \(Q_{t:T}\). In a similar fashion, we have

Phew! Believe me, that was not any more complicated than a second order Taylor series expansion. Thankfully, we can write this in matrix form, where the Jacobians will be replaced with cleaner symbols:

The matrices \(Q^{(t)}_{xx}, Q^{(t)}_{uu}, Q^{(t)}_{ux}, Q^{(t)}_{xu}, Q^{(t)}_x, Q^{(t)}_u\) are derived below:

You may be concerned about the dependence of these matrices on the \(M\) and \(N\) matrices "from the future", but fortunately we can solve for these matrices backwards in time (hence, the backwards update step). Note that at \(t=T\) (the end of the horizon), the cost to go is \(V_{T:T}(x_T) = \ell(x_T) = x_T^\top Q_f x_T\), where we optionally specify another cost matrix \(Q_f\) to distinguish "transient" penalties from the penalty on the final state. Therefore, by \eqref{eq:dV}, we have \(M_{T} = Q_f(x_T - x^*)\) for target state \(x^*\), and \(N_T = Q_f\). With this edge case, we can compute the Jacobian matrices backward from \(t=T\) to \(t=0\).

To minimize the cost-to-go \(V_{t:T}(x_t)\), we must solve for \(\delta u_t = \arg\min_{\delta u_t}\delta Q_{t:T}\). We proceed to do this using standard calculus optimization techniques. Due to Bellman's principle of optimality, we can simply optimize \(\delta u_t\) individually for each timestep. Referring to \eqref{eq:dQ}, have

Notably, we have

where \(K_t = -(Q_{uu}^{(t)})^{-1}Q^{(t)}_{ux}\) and \(d_t = -(Q_{uu}^{(t)})^{-1}Q^{(t)}_u\), so optimal changes to controls at each step are affine transformations of the measured error in state (we will discuss how this state error is calculated soon).

Recall that in order to compute these gains, we need the \(M\) and \(N\) matrices for each timestep. First, substitute the \(\delta u_t\) computed in \eqref{eq:du} into \eqref{eq:dQ}:

Recall that \(M_t\) is the term that is linear in \(\delta x_t\) in \eqref{eq:dV}. Since \(\frac{\partial V_{t:T}}{\partial\delta x_t} = \frac{\partial}{\partial\delta x_t}\min_{\delta u_t}Q_{t:T}(\delta x_t,\delta u_t) = \delta Q_{t:T}\), \(M_T\) is obtained by extracting the linear terms in \(\delta x_t\) from the matrix equation above:

Similarly, since \(N_t\) is the term that is quadratic in \(\delta x_t\), we have

By starting at \(t=T\) and working backwards, we can compute all of the \(K_t\) and \(d_t\) vectors. This is the job of BackwardUpdate in the pseudocode above.

Before iLQR even begins, we must have some initial trajectory of states and controls. It is natural to wonder how one concocts such an initial trajectory. Unfortunately, in the general case, there is no single "right way" to do this, unless somehow you know of a trajectory that is somewhat close to optimal. Otherwise, to the best of my knowledge, we can use some heuristics to generate a trajectory of controls and initialize the state trajectory according to the dynamics following those controls. Here's some ideas for generating the control trajectory that I played with:

• \(u_t = 0\): This one was suggested by Li and Todorov (the inventors of iLQR). Intuitively, this might make training a little slow if the starting state is in a stable equilibrium.
• \(u_t = \alpha\sin(\omega t), \alpha,\omega\in\mathbf{R}\): This is likely to make more sense for scalar controls, but it ensures that a nice range of controls will be attempted in the initial trajectory. I found this worked slightly better than the "zero initialization" for the swingup.
• \(u_t\sim\mathcal{N}(0, \Sigma), \Sigma=\Sigma^\top\in\mathbf{R}^{n_u\times n_u}_{\geq 0}\succ 0\): Zero-centered random controls. This one was the best during my experiments.

Since iLQR is an iterative nonlinear optimization algorithm, initialization can definitely have an affect on the final controller. Playing with various initialization schemas would be wise.

I cannot emphasize enough how essential this trick was. When tuning my iLQR implementation, I noticed that the controller would fail to swing up the pole because it wouldn't accelerate fast enough, which I attributed to penalizing velocities and control magnitudes too much. But when I would try reducing these penalties, iLQR would outright explode: velocities and angular velocities would shoot up, and the loss would skyrocket, leading to a completely abysmal final trajectory. After further inspection, I noticed that \(|\delta u_t|\) was blowing up. To prevent this from happening, I implemented the following adjustment

\( \delta u_t = \text{sgn}(K_t\delta x_t + d_t)\min(K_t\delta x_t + d_t, \Delta) \)

where \(\text{sgn}(x) = |x|/x\), which essentially clips the control updates to some magnitude \(\Delta\). Note that while this transformation is defined for scalar controls, an equivalent transformation for vector controls can be derived without much trouble. Upon employing this trick, these divergence issues ceased immediately. This allowed me to be much more flexible with my cost functions, and ultimately allowed me to tune my iLQR to a competent solution in a finite amount of time.

This trick does impose an additional hyperparameter, however in my experiments I found that this hyperparameter was very easy to tune. Moreover, when \(\Delta\) is "too low", the only consequence is that iLQR won't make progress as quickly, which normally isn't such a big deal.

Alternatively, line search methods or adjustable (perhaps even adaptive) \(\delta u_t\) step sizes could potential help solve this problem. However, these methods involve hyperparameters of their own, and I still believe gradient clipping could be necessary – for instance, say you update \(u_t\leftarrow\bar{u}_t + \alpha\delta u_t\) for some small step size \(\alpha\), but \(||\delta u_t||\) is enormous. The small step size would dampen the effect of this, but without forcing very small step sizes, this can't eliminate the divergence issue I was experiencing. Combining gradient clipping with line search or adaptive step sizes would likely be the most robust strategy, however gradient clipping on its own was more than sufficient enough for my purposes.

This trick allowed me to design my cost functions with much more intuition, which proved to be very helpful. Firstly, if you distinguish \(Q_f\) from \(Q\), the influence of the penalties governed by \(Q\) to the total trajectory cost will devastatingly outweigh that of \(Q_f\) for long time horizons. It might even be wise to let \(Q = 0\) if you mostly care about penalizing state at the end of the trajectory. Otherwise, I found that scaling \(Q\) by a factor of \(1/T\) helped design \(Q\) and \(Q_f\) quite a bit.

Moreover, if the state variables are measured on different scales, this should be compensated for. Let's say you want to penalize deviations in angle twice as much as deviations in velocity. Angles are confined to \([0,2\pi]\), while velocities could be much larger, say in the range \([-v_{\max},v_{\max}]\). Therefore it doesn't suffice to assign twice the weight to deviations in angle, since a deviation of \(\pi\) radians is most likely much more significant than a deviation of \(\pi\) in velocity. Let \(r_i\) denote the range of state \(i\), which would be \(2\pi\) for angles and \(2v_{\max}\) for the velocity example. To compensate for changes of scale, let \(\tilde{Q}\) be the "intuitive" cost matrix (that is, one where you assign twice the weight to a state deviation you want to penalize twice as much), and compute \(Q\) according to

In the derivation of the iLQR optimization math, we needed to compute a bunch of Jacobians. This doesn't necessarily cause too much trouble, since given the dynamics model we can simply write out the Taylor expansions on paper and then code them up. I can think of two reasons why we might prefer a workaround for this:

1. The state-control space is large, so deriving the Jacobians could be very time consuming and/or error-prone (or you're simply too lazy to derive the Jacobians)
2. You don't actually have an explicit dynamics model (for instance, perhaps you have a powerful simulator instead)

Fortunately, we can automatically approximate Jacobians using the method of finite differences, which is so easy to implement that it's almost irresistible. The method essentially works by perturbing each variable to be differentiated by some small \(\epsilon>0\) in each direction and examining how that changes the output, like so:

If the function \(f\) is "smooth enough", we can approximate the Jacobian arbitrarily well by shrinking \(\epsilon\). For example, approximating the Jacobian with respect to states for a dynamics function \(f:\mathbf{R}^n\times \mathbf{R}^{n_u}\to\mathbf{R}^n\) can be achieved as follows:

def finite_differences(dynamics, x, u, eps):
  n = x.shape[0]
  jacobian = np.zeros((n, n))
  for i in range(n):
    dxp = x.copy()
    dxp[i] += eps
    x_inc = dynamics(dxp, u)
    dxm = x.copy()
    dxm[i] -= eps
    x_dec = dynamics(dxm, u)
    jacobian[:, i] = (x_inc - x_dec) / (2 * eps)
  return jacobian