API Reference

Every solver comes in a sequential form, which uses a Riccati recursion and suits CPUs, and a parallel form, which uses an associative scan with \(O(\log T)\) depth and suits GPUs on long horizons. Each also has a differentiable wrapper. The wrapper returns \(\tau = [x_t, u_t]\) and has implicit (adjoint-LQR) gradients with respect to the problem parameters.

Solvers

lqr

Sequential LQR: solve_lqr().

LQR (sequential)

ilqr

Sequential iLQR: ilqr_solver().

iLQR (sequential)

plqr

Parallel LQR: solve_plqr().

LQR (parallel)

parallel_ilqr

Parallel iLQR: pilqr_solver().

iLQR (parallel)

Differentiable solvers

diff_lqr

dlqr()

Differentiable LQR

diff_ilqr

dilqr()

Differentiable iLQR

parallel_dlqr

pdlqr()

Differentiable LQR (parallel)

parallel_dilqr

pdilqr()

Differentiable iLQR (parallel)

Batched solves

batch

Many problems sharing one model: batched_ilqr_solver() and batched_dilqr(), by vmap, lax.map or shard_map.

Batched solves

Types and utilities

typs

ModelDims, System, LQR, LQRParams, iLQRParams and more.

Data types

utils

Key splitting, random stable dynamics and linearisation helpers.

Utilities