Getting started#

Installation#

diffilqrax needs Python 3.10 or later and depends only on jax, numpy and jaxopt. To install from a local clone with pip:

git clone git@github.com:ThomasMullen/diffilqrax.git
cd diffilqrax
pip install -e .              # library only
pip install -e ".[dev]"       # + test and lint tools
pip install -e ".[doc]"       # + documentation tools

Or with uv, which uses the committed uv.lock:

uv sync --extra dev            # creates .venv with the locked dependencies
uv run pytest tests/ -m "not slow"

For GPU support, install the matching jax build first, as described in the JAX installation guide.

Note

Importing a diffilqrax solver turns on JAX double precision (jax_enable_x64). The default tolerances, for example kkt_tol=1e-10, assume float64.

Core types#

All problem data is held in the small NamedTuple containers defined in diffilqrax.typs:

ModelDims(n, m, horizon, dt)

State dimension, input dimension, number of time steps \(T\), and the time step.

LQR(A, B, a, Q, q, R, r, S, Qf, qf)

A time-varying LQR problem. The time-varying terms are stacked along a leading axis of length \(T\): A has shape (T, n, n), B has shape (T, n, m) and S has shape (T, n, m). The terminal terms Qf and qf have no time axis.

LQRParams(x0, lqr)

Initial state and the LQR problem.

System(cost, costf, dynamics, dims)

A non-linear problem. Every function takes the parameters theta as its last argument: cost(t, x, u, theta), costf(x, theta) and dynamics(t, x, u, theta), which returns \(x_{t+1}\). Derivatives are taken with JAX automatic differentiation.

iLQRParams(x0, theta)

Initial state and an arbitrary pytree of parameters. Gradients are taken with respect to this pytree.

Solving an LQR problem#

Consider optimal control for a double integrator driving point mass to rest state. The linear state dynamics are

(1)#\[\begin{split}x_{t+1} = A x_t + B u_t, \qquad A = \begin{pmatrix} 1 & \delta t \\ 0 & 1 \end{pmatrix}, \quad B = \begin{pmatrix} \tfrac12 \delta t^2 \\ \delta t \end{pmatrix}.\end{split}\]

The solver is searching for a control sequence that minimises quadratic cost

\[\begin{split}\min_{u_{0:T-1}} \; J = \tfrac{1}{2} x_T^\top Q_f x_T + \tfrac{1}{2} \sum_{t=0}^{T-1} \left( x_t^\top Q x_t + r\, u_t^2 \right), \\ Q = \begin{pmatrix} 1 & \delta t \\ 0 & 1 \end{pmatrix}, \quad Q_f = \begin{pmatrix} 10 & \delta t \\ 0 & 10 \end{pmatrix}, \quad R = \begin{pmatrix} 0.1 & \delta t \\ 0 & 0.1 \end{pmatrix}, \quad\end{split}\]

subject to (1), with \(Q, Q_f \succeq 0\) and \(r > 0\).

The solution is a linear feedback law \(u_t = -K_t x_t\), where the gains follow from the backward Riccati recursion, starting at \(V_T = Q_f\):

\[\begin{split}K_t &= \left( r + B_d^\top V_{t+1} B_d \right)^{-1} B_d^\top V_{t+1} A_d \\ V_t &= Q + A_d^\top V_{t+1} \left( A_d - B_d K_t \right)\end{split}\]
import jax.numpy as jnp
from diffilqrax.typs import LQR, LQRParams, ModelDims
from diffilqrax.lqr import solve_lqr

dims = ModelDims(n=2, m=1, horizon=50, dt=0.1)
T, n, m = dims.horizon, dims.n, dims.m

A = jnp.array([[1.0, dims.dt], [0.0, 1.0]])
B = jnp.array([[0.0], [dims.dt]])
lqr = LQR(
    A=jnp.tile(A, (T, 1, 1)), a=jnp.zeros((T, n)),
    B=jnp.tile(B, (T, 1, 1)),
    Q=jnp.tile(jnp.eye(n), (T, 1, 1)), q=jnp.zeros((T, n)),
    S=jnp.zeros((T, n, m)),
    R=jnp.tile(0.1 * jnp.eye(m), (T, 1, 1)), r=jnp.zeros((T, m)),
    Qf=10.0 * jnp.eye(n), qf=jnp.zeros(n),
)
params = LQRParams(x0=jnp.array([1.0, 0.0]), lqr=lqr)

Xs, Us, lambdas = solve_lqr(params) # shapes (T+1, n), (T, m), (T+1, n)

diffilqrax.plqr.solve_plqr() takes the same arguments and returns the same solution, computed with an associative scan.

Solving a non-linear problem with iLQR#

This example swings a pendulum up to \(\theta = \pi\):

import jax
import jax.numpy as jnp
from diffilqrax.typs import System, iLQRParams, PendulumParams
from diffilqrax.ilqr import ilqr_solver

DAMPING, U_MAX = 0.1, 3.0
dims = ModelDims(n=2, m=1, horizon=100, dt=0.05)

def dynamics(t, x, u, theta):
    """Euler step; x = [angle, angular velocity], angle 0 is hanging down"""
    acc = -theta.g / theta.l * jnp.sin(x[0]) - DAMPING * x[1] + u[0] / (theta.m * theta.l**2)
    return x + DT * jnp.stack([x[1], acc])

def cost(t, x, u, theta):
    """Small running cost on angle, velocity and torque, plus a smooth barrier at |u| = U_MAX"""
    barrier = jax.nn.softplus(5.0 * (u - U_MAX)) ** 2 + jax.nn.softplus(5.0 * (-u - U_MAX)) ** 2
    return 0.1 * (1.0 + jnp.cos(x[0])) + 0.01 * x[1] ** 2 + 0.01 * jnp.sum(u**2) + jnp.sum(barrier)

def costf(x, theta):
    """Be upright (angle ±pi) and still at the final time"""
    return 50.0 * (1.0 + jnp.cos(x[0])) + 5.0 * x[1] ** 2

# Construct problem
model = System(cost, costf, dynamics, dims)
params = iLQRParams(x0=jnp.array([0.1, 0.0]), theta=PendulumParams(m=1.0, l=1.0, g=9.8))
Us_init = 0.5 * jr.normal(jr.PRNGKey(0), (dims.horizon, dims.u))

# Solve problem
(Xs, Us, lambdas), total_cost, cost_log = ilqr_solver(model, params, Us_init, max_iter=300)

The main options are:

  • kkt_tol: the solver stops when \(\max_t |\partial L / \partial u_t|\) is below this value.

  • max_iter: the maximum number of outer iterations.

  • line-search keyword arguments for armijo_linesearch(): tau, beta, alpha_min and max_iter_linesearch.

Regularisation of the input Hessian is adaptive, so it needs no tuning.

Differentiating through the solver#

diffilqrax.diff_ilqr.dilqr() solves the same problem and returns \(\tau = [x_t, u_t]\) as an array of shape (T+1, n+m). It defines a custom VJP, so jax.grad gives exact implicit gradients with respect to x0 and every leaf of theta. JAX does not backpropagate through the solver iterations.

import jax
from diffilqrax.diff_ilqr import dilqr

def loss(theta):
    tau = dilqr(model, iLQRParams(x0=jnp.zeros(2), theta=theta), Us_init, max_iter=200, kkt_tol=1e-6)
    return jnp.sum(tau[:, dims.n:] ** 2)   # control effort of the optimal solution

grads = jax.grad(loss)(params.theta)       # a PendulumParams of gradients

For LQR problems, diffilqrax.diff_lqr.dlqr() does the same job: dlqr(dims, params, None).

Parallel solvers#

diffilqrax.parallel_ilqr.pilqr_solver() and diffilqrax.parallel_dilqr.pdilqr() are drop-in replacements for ilqr_solver and dilqr. They give the same solution, but the Riccati recursion becomes an associative scan:

from diffilqrax.parallel_ilqr import pilqr_solver

(Xs, Us, _), J, _ = pilqr_solver(model, params, Us_init, max_iter=200, kkt_tol=1e-6)

With rollout="parallel", the non-linear rollout is also replaced, by multiple shooting, so each iteration has \(O(\log T)\) depth. This is worthwhile on a GPU with long horizons. On a CPU the sequential solvers are usually faster. See Parallel (associative-scan) solvers for a comparison.

Next steps#

  • Tutorials has worked notebooks, each ending with a How to adapt section.

  • Design Principles gives an overview of the modules.

  • examples/ in the repository has standalone scripts, such as the pendulum swing-up and a 4000-step tracking problem.

Building the documentation locally#

pip install -e ".[doc]"        # also needs the pandoc executable, e.g. brew install pandoc
make -C docs html              # or: sphinx-build -b html docs/source docs/build/html
open docs/build/html/index.html