DiffiLQRax v1.0: a more stable, faster and batched iLQR solver

Last year I didn’t spend much time maintaining DiffiLQRax. I knew some bugs made the solver unstable, especially with very non-convex cost functions or state spaces larger than about 20 dimensions. Over the past few weeks I have gone back to the code and rewritten the core so the algorithm reliably converges to a local optimum. I have also added a few features I have wanted for a while. This post summarises the changes.

Stability fixes

Adaptive Levenberg–Marquardt regularisation. In the backward pass, the input Hessian \(Q_{uu}\) must be positive definite for the step to be a descent direction. With non-convex costs this often fails. The solver now adds an adaptive Levenberg–Marquardt term \(\mu I\) until \(Q_{uu}\) is positive definite. If a step is rejected, \(\mu\) increases and the step is retried; if it is accepted, \(\mu\) decreases again. Previously, this was a fixed tolerance hyperparameter.

Correct trajectory updates. I fixed how the new state and input trajectories are rolled out during each iLQR iteration. The forward pass now applies the feedback gains around the updated trajectory, and the step size comes from an Armijo line search. The solver stops when the KKT conditions are met (kkt_tol), not after a fixed number of iterations.

Implicit differentiation instead of unrolling. JAX no longer backpropagates through the iLQR iterations. At the optimal state–input trajectory, the solver builds a local LQR problem and differentiates that instead. The result is exact implicit gradients with respect to the initial state and every leaf of the parameter pytree, at a memory cost that does not depend on the number of iterations.

from diffilqrax.diff_ilqr import dilqr

def loss(theta):
    tau = dilqr(model, iLQRParams(x0=x0, theta=theta), Us_init)
    return jnp.sum(tau[:, dims.n:] ** 2)

grads = jax.grad(loss)(theta)

New features

Batched solves on CPU, GPU or several devices. batched_ilqr_solver and batched_dilqr solve many problems that share one model in a single call, for example several target trajectories or initial conditions. The strategy argument selects "vmap" (vectorised), "map" (lax.map, chunked by batch_size, lower memory) or "shard" (shard_map across jax.devices(), which includes multi-host setups after jax.distributed.initialize()). Gradients of a loss summed over the batch work with every strategy.

from diffilqrax.batch import batched_ilqr_solver

(Xs, Us, _), J, _ = batched_ilqr_solver(model, params, Us_init, strategy="shard")

DDP or Gauss–Newton. Set exact_hessian=True (the default) to include the second-order dynamics term \(\lambda^\top \nabla^2 f\), as in DDP, so the gradients match direct differentiation. Set exact_hessian=False to keep only the cost Hessians (Gauss–Newton). This is cheaper, and exact only when the dynamics are linear.

Parallel LQR and iLQR. The Riccati recursion can now be solved as an associative scan, following Särkkä & García-Fernández. This reduces the time-depth from \(O(T)\) to \(O(\log T)\), which helps on GPUs with long horizons. solve_plqr, pilqr_solver and pdilqr are drop-in replacements for their sequential versions and return the same solution. On CPU or short horizons, the sequential solvers are usually still faster.

Examples

The documentation has been rewritten and includes notebook tutorials:

• LQR: bringing a double integrator to rest
• LQR tracking of a reference trajectory
• iLQR: swinging up a torque-limited pendulum
• iLQR compared with gradient descent
• Implicit gradients through the iLQR solver
• Parallel (associative-scan) solvers
• iLQR inference: reconstructing a zebrafish tail bout

Getting started

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

Importing a solver turns on JAX double precision (jax_enable_x64), because the default tolerances assume float64. For GPU support, install the matching JAX build first.

If you find DiffiLQRax useful, please star or watch the repo. Issues and pull requests are welcome.