lqrax

GPU-friendly, auto-differentiable LQR solver with JAX.

Science Score: 44.0%

This score indicates how likely this project is to be science-related based on various indicators:

✓
CITATION.cff file
Found CITATION.cff file
✓
codemeta.json file
Found codemeta.json file
✓
.zenodo.json file
Found .zenodo.json file
○
DOI references
○
Academic publication links
○
Academic email domains
○
Institutional organization owner
○
JOSS paper metadata
○
Scientific vocabulary similarity
Low similarity (12.2%) to scientific vocabulary

Keywords

differential-games dynamic-games ilqr iterative-linear-quadratic jax linear-quadratic-regularator lqr lqr-controller optimal-control riccati-equations riccati-solver trajectory-optimization

Last synced: 6 months ago · JSON representation ·

Repository

GPU-friendly, auto-differentiable LQR solver with JAX.

Basic Info

Host: GitHub
Owner: MaxMSun
License: gpl-3.0
Language: Python
Default Branch: main
Homepage:
Size: 7.96 MB

Statistics

Stars: 172
Watchers: 3
Forks: 6
Open Issues: 0
Releases: 5

Topics

differential-games dynamic-games ilqr iterative-linear-quadratic jax linear-quadratic-regularator lqr lqr-controller optimal-control riccati-equations riccati-solver trajectory-optimization

Created 11 months ago · Last pushed 7 months ago

Metadata Files

Readme License Citation

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.

This repo is currently under active development.

Examples

Please also checkout Linear Quadratic Flow Matching that uses this package.

Install

Follow the instructions to install JAX before installing this package.

To install: pip install lqrax

Usage

There are two modules: LQR and iLQR,

The LQR module solves the following time-varying LQR problem:

$$ \arg\min{u(t)} \int0^T \Big[ (x(t)-x{ref}(t))^\top Q (x(t)-x{ref}(t)) + u(t)^\top R u(t) \Big] dt $$ $$ \text{s.t. } \dot{x}(t) = A(t) x(t) + B(t) u(t), \quad x(0) = x_0 $$

The iLQR module solves a different time-varying LQR problem:

$$ \arg\min{v(t)} \int0^T \Big[ z(t)^\top Q z(t) + v(t)^\top R v(t) + z(t)^\top a(t) + v(t)^\top b(t) \Big] dt $$ $$ \text{s.t. } \dot{z}(t) = A(t) z(t) + B(t) v(t), \quad z(0) = 0. $$

This formulation is often used as the sub-problem for iterative linear quadratic regulator (iLQR) to calculate the steepest descent direction on the control for a general nonlinear control problem:

$$ \arg\min{u(t)} \int0^T l(x(t), u(t)) dt, \text{ s.t. } \dot{x}(t) = f(x(t), u(t)), $$

where the $z(t)$ and $v(t)$ are perturbations on the system's state $x(t)$ and control $u(t)$, and $A(t)$ and $B(t)$ are the linearized system dynamics $f(x(t), u(t))$ on the current system trajectory with respect to the state and control.

Copyright and License

The implementations contained herein are copyright (C) 2024 - 2025 by Max Muchen Sun, and are distributed under the terms of the GNU General Public License (GPL) version 3 (or later). Please see the LICENSE for more information.

If you use the package in your research, please cite this repository. You can see the citation information at the right side panel under "About". The BibTeX file is attached below: @software{sun_lqrax_2025, author = {["Sun"], Max Muchen}, license = {GPL-3.0}, month = march, title = {{LQRax: JAX-enabled continuous-time LQR solver}}, url = {https://github.com/MaxMSun/lqrax}, version = {0.0.6}, year = {2025} }

Contact: msun@u.northwestern.edu

Owner

Name: Max M. Sun
Login: MaxMSun
Kind: user
Location: Chicago, IL
Company: Northwestern University

Website: https://maxsun.io/
Twitter: max_msun
Repositories: 4
Profile: https://github.com/MaxMSun

Robotics Ph.D. candidate at Northwestern University

Citation (CITATION.cff)

cff-version: 1.2.0
message: "If you use this software, please cite it as below."
authors:
  - family-names: [Sun]
    given-names: [Max Muchen]
    orcid: [0000-0002-3704-6315]
title: "LQRax: JAX-Enabled continuous-time Riccati equation solver"
version: "0.0.6"
date-released: 2025-03-28
url: "https://github.com/MaxMSun/lqrax"
repository-code: "https://github.com/MaxMSun/lqrax"
license: "GPL-3.0"
keywords:
  - optimal control

GitHub Events

Total

Create event: 5
Issues event: 1
Release event: 4
Watch event: 117
Issue comment event: 2
Public event: 1
Push event: 26
Fork event: 6

Last Year

Create event: 5
Issues event: 1
Release event: 4
Watch event: 117
Issue comment event: 2
Public event: 1
Push event: 26
Fork event: 6

Issues and Pull Requests

Last synced: 7 months ago

All Time

Total issues: 1
Total pull requests: 0
Average time to close issues: 25 days
Average time to close pull requests: N/A
Total issue authors: 1
Total pull request authors: 0
Average comments per issue: 2.0
Average comments per pull request: 0
Merged pull requests: 0
Bot issues: 0
Bot pull requests: 0

Past Year

Issues: 1
Pull requests: 0
Average time to close issues: 25 days
Average time to close pull requests: N/A
Issue authors: 1
Pull request authors: 0
Average comments per issue: 2.0
Average comments per pull request: 0
Merged pull requests: 0
Bot issues: 0
Bot pull requests: 0

ecosyste.ms

Data

Tools

Indexes

Applications

Experiments

Open Source Science