LQRax

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

LQRax is a GPU-friendly, auto-differentiable solver for continuous-time LQR problems based on Riccati equations, enabled by JAX.

• It accelerates numerical simulation through JAX’s scan mechanism;
• It enables rapid prototyping of single-agent and multi-agent nonlinear control algorithms, with auto-differentiation support on the loss function and dynamics;
• It enables batch-based large-scale optimal control on GPUs using JAX’s vmap mechanism.
• All the operations, including trajectory simulation and control synthesis, are backward-differentiable.