DiffiLQRax in practice: inference, motor control and learning dynamics

DiffiLQRax is a differentiable iLQR solver in JAX. You describe a problem as a cost and dynamics, and it returns the optimal trajectory along with exact gradients through that optimum. This post shows the same solver in three roles: as an inference engine, as a controller, and as an encoder inside a model you train. For what changed in v1.0, see the release post.

1. Inference: which inputs produced this behaviour?

A trained iLQR-VAE models a larval zebrafish tail bout as a 120-unit recurrent network driven by 10 sparse inputs, read out as 7 tail angles. Inferring the inputs behind an observed bout is an optimal control problem. The cost is the Student-t prior on the inputs plus the reconstruction error, and iLQR finds its minimum. DiffiLQRax matches the objective of our original OCaml implementation to round-off.

model = System(cost, costf, dynamics, ModelDims(n=120, m=10, horizon=T, dt=1.0))
theta = (network_params, observed_bout)           # cost = prior(u) + likelihood(y | x)
(xs, us, _), J, _ = ilqr_solver(model, iLQRParams(x0=jnp.zeros(120), theta=theta), jnp.zeros((T, 10)))

2. Motor control: optimal commands for a muscle-driven arm

The plant is a two-link arm driven by six muscles, assembled from feedbax components. The task is to reach to 8 targets around a start point. iLQR finds the open-loop muscle excitations that minimise effort plus endpoint error. It converges in 19–35 iterations and ends within 0.25 cm of each target. This optimum serves as a benchmark for learned controllers. A GRU feedback controller trained by backpropagation through the arm reaches a cost 0.2–3.6% above it.

model = System(running_cost, final_cost, arm_dynamics, ModelDims(n=n_state, m=6, horizon=T, dt=0.01))
solve = lambda target: ilqr_solver(model, iLQRParams(x0=x0, theta=target), us_init)
(xs, us, _), J, _ = jax.vmap(solve)(targets)      # all 8 reaches in one call

Because the target is just a parameter (theta), solving every reach at once is a single vmap. The batched solvers in v1.0 can also shard such batches across devices (details).

3. Learning dynamics: iLQR as a differentiable encoder

In an iLQR-VAE the recognition network is not a neural network. The posterior mean of the inputs is the solution of an iLQR problem defined by the generative model itself. Training the model needs gradients through that solution. dilqr provides them by implicit differentiation at the optimum, so it doesn’t backpropagate through the solver iterations. Here the generative model is fit to noisy 10-D observations of a Lorenz system that receives occasional kicks. On held-out trials, the inferred trajectories recover the true 3-D state with R² = 0.996.

def elbo_loss(params, ys):
    tau = dilqr(model, iLQRParams(x0=x0, theta=(params, ys)), us_init)   # posterior mean (x, u)
    return -elbo(params, tau, ys)

grads = jax.grad(elbo_loss)(params, ys)            # implicit gradients through the iLQR optimum

The encoder works well. Getting the learned latent dynamics to reproduce the attractor in long autonomous rollouts is work in progress.

Further reading

• DiffiLQRax v1.0 release notes: stability fixes, implicit gradients, batched and parallel solvers
• The original DiffiLQRax post
• Documentation and tutorials, including the zebrafish tail reconstruction
• Code on GitHub