LQRax

Last updated on Apr 24, 2025
GitHub

GitHub link: https://github.com/MaxMSun/lqrax

2D drone landing example with ilqr

LQRax is JAX-enabled continuous-time LQR solver. It is a Riccati equation solver natively written in JAX:

  • It accelerates numerical simulation through JAX’s scan mechanism;

  • It allows rapid prototyping of iterative LQR (iLQR) for nonlinear control by supporting auto-differentiation on the loss function and dynamics;

  • It also enables batch-based large-scale optimal control using JAX’s vmap mechanism.

Code
Max Muchen Sun
Max Muchen Sun
Ph.D. Candidate in Robotics

I develop algorithms to make robots self-sufficient.