_images/diffilqrax_logo_dm.png _images/diffilqrax_logo.png

Welcome to DiffiLQRax!#

DiffiLQRax is an open source JAX library for LQR and iterative LQR (iLQR) optimal control. It provides:

  • sequential (Riccati) and parallel (associative-scan) solvers for LQR and iLQR,

  • adaptive regularisation and Armijo line search, with stopping based on the KKT conditions,

  • exact implicit gradients through every solver via jax.grad, so a controller can be trained inside a larger model.

Please support the development of DiffiLQRax by starring or watching the project on GitHub.

Quick Examples#

Linear Quadratic Regulator (LQR)#

LQR: bringing a double integrator to rest
LQR regulation
LQR: bringing a double integrator to rest
LQR tracking of a reference trajectory
LQR tracking
LQR tracking: following a reference trajectory

Iterative LQR (iLQR)#

iLQR: pendulum swing-up
pendulum swing-up
iLQR: swinging up a torque-limited pendulum
iLQR compared with gradient descent
iLQR vs gradient descent
iLQR compared with gradient descent
Implicit gradients through iLQR
implicit vs direct gradients
Implicit gradients through the iLQR solver
Parallel (associative-scan) solvers
parallel solvers
Parallel (associative-scan) solvers
iLQR inference: reconstructing a zebrafish tail bout
zebrafish tail reconstruction
iLQR inference: reconstructing a zebrafish tail bout

Installation#

git clone git@github.com:ThomasMullen/diffilqrax.git
cd diffilqrax
pip install -e .        # or: uv sync

See Getting started for optional extras, GPU notes and a first example.

Contents#

Contents

Acknowledgements#

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.