Diffilqrax: Differentiable optimal control

What is it?

Diffilqrax is a JAX library for LQR and iterative LQR (iLQR) optimal control. It provides:

• Sequential and parallel solvers. The sequential solvers use a Riccati recursion (lqr.solve_lqr, ilqr.ilqr_solver). The parallel solvers use an associative scan with O(log T) depth (plqr.solve_plqr, parallel_ilqr.pilqr_solver).
• Robust iLQR. It uses adaptive Levenberg–Marquardt regularisation and an Armijo line search, and stops based on the KKT conditions.
• Exact implicit gradients. diff_lqr.dlqr, diff_ilqr.dilqr and their parallel versions pdlqr and pdilqr define custom VJPs, so jax.grad differentiates the optimal trajectory with respect to the dynamics and cost parameters without backpropagating through the solver iterations.
• Batched solves. batch.batched_ilqr_solver and batch.batched_dilqr solve many problems that share a model, with vmap, chunked lax.map or shard_map across devices.

Note on reverse mode. The solvers run as a lax.while_loop by default (loop="while"), so jax.grad cannot backpropagate through ilqr_solver or pilqr_solver directly. To get gradients of the solution, use dilqr, pdilqr or batched_dilqr. To backpropagate through the iterations themselves, pass loop="scan". This stores every iteration and always runs all max_iter of them under vmap.

Installation

From PyPI, with pip or uv:

pip install diffilqrax
uv pip install diffilqrax       # into the active environment
uv add diffilqrax               # as a dependency of a uv project

The latest version straight from GitHub:

uv pip install "git+https://github.com/ThomasMullen/diffilqrax"

For development, from a local clone:

git clone git@github.com:ThomasMullen/diffilqrax.git
cd diffilqrax
uv pip install -e ".[dev]"      # or: pip install -e ".[dev]"   (tests and linting)
uv pip install -e ".[doc]"      # documentation
uv build                        # sdist + wheel in dist/

Importing diffilqrax enables JAX double precision (jax_enable_x64).

Quick example

This example swings a pendulum up from hanging down to upright:

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

dims = ModelDims(n=2, m=1, horizon=100, dt=0.05)

def dynamics(t, x, u, theta):          # x = [angle, angular velocity]
    acc = -theta.g / theta.l * jnp.sin(x[0]) + u[0] / (theta.m * theta.l**2)
    return x + dims.dt * jnp.stack([x[1], acc])

def cost(t, x, u, theta):
    return 0.01 * jnp.sum(u**2)

def costf(x, theta):
    return 10.0 * ((x[0] - jnp.pi) ** 2 + x[1] ** 2)

model = System(cost, costf, dynamics, dims)
params = iLQRParams(x0=jnp.zeros(2), theta=PendulumParams(m=1.0, l=1.0, g=9.81))
Us_init = jnp.zeros((dims.horizon, dims.m))

(Xs, Us, lambdas), total_cost, cost_history = ilqr_solver(model, params, Us_init, max_iter=200)

# differentiable version: tau = [x_t, u_t], with implicit gradients w.r.t. params
tau = dilqr(model, params, Us_init, max_iter=200)

More examples are in the documentation notebooks (docs/source/examples/) and in examples/.

Documentation

The documentation is hosted at https://thomasmullen.github.io/diffilqrax/.

Build the docs locally (this also needs pandoc, for example brew install pandoc):

pip install -e ".[doc]"
make -C docs html
open docs/build/html/index.html

Tests

pytest tests/ -m "not slow and not heavy"   # quick tier, as run in CI
pytest tests/ -m "not slow"                  # fast suite: adds the heavy cases
pytest tests/ -m slow                        # large-dimension tests

The heavy cases (listed in tests/conftest.py) are too slow or memory-hungry for a GitHub runner and re-check properties that a cheaper test already covers. CI runs the quick tier; the full tiers can be started by hand from the "Python package" workflow.

Benchmarks import the shared test problems from tests/, so run them as modules from the repository root, for example python -m benchmarks.timing --quick.

The tests run on the CPU by default. Set DIFFILQRAX_TEST_PLATFORM=gpu to run them on a GPU instead.

Design principle

License

This project is licensed under the Apache License 2.0. See the LICENSE file for details.

Acknowledgments

This project has received funding from the European Union’s Horizon 2020 research and innovation programme under the Marie Skłodowska-Curie grant agreement #813457.