jaxfg

Factor graph-based nonlinear optimization library for JAX.

Applications include sensor fusion, control, planning, SLAM. Borrows heavily from a wide set of existing libraries, including: Ceres Solver, g2o, GTSAM, minisam, SwiftFusion.

Features:

• Autodiff-powered (sparse) Jacobians.
• Automatic batching of factor computations.
• Out-of-the-box support for optimization on SO(2), SO(3), SE(2), and SE(3).
• 100% implemented in Python!

Current limitations:

• JIT compilation adds significant startup overhead. This could likely be optimized (for example, by specifying more analytical Jacobians) but is mostly unavoidable with JAX/XLA. Limits applications for systems that are online or require dynamic graph alterations.
• Python >=3.7 only, due to features needed for generic types.

Installation

scikit-sparse require SuiteSparse:

sudo apt update
sudo apt install -y libsuitesparse-dev

Then, from your environment of choice:

git clone https://github.com/brentyi/jaxfg.git
cd jaxfg
pip install -e .

Example scripts

Toy pose graph optimization:

python scripts/pose_graph_simple.py

Pose graph optimization from .g2o files:

python scripts/pose_graph_g2o.py --help

To-do

• Preliminary graph, variable, factor interfaces
• Real vector variable types
• Refactor into package
• Nonlinear optimization for MAP inference
  • Conjugate gradient linear solver
  • CHOLMOD linear solver
    • Basic implementation. JIT-able, but no vmap, pmap, or autodiff support.
  • Gauss-Newton implementation
  • Termination criteria
  • Damped least squares
  • Dogleg
  • Inexact Newton steps
  • Revisit termination criteria
  • Reduce redundant code
  • Robust losses
• Marginalization
  • Working prototype using sksparse/CHOLMOD
  • JAX implementation?
• Validate g2o example
• Performance
  • More intentional JIT compilation
  • Re-implement parallel factor computation
  • Vectorized linearization
  • Basic (Jacobi) CGLS preconditioning
• Manifold optimization (mostly offloaded to jaxlie)
  • Basic interface
  • Manifold optimization on SO2
  • Manifold optimization on SE2
  • Manifold optimization on SO3
  • Manifold optimization on SE3
• Usability + code health (low priority)
  • Basic cleanup/refactor
    • Better parallel factor interface
    • Separate out utils, lie group helpers
    • Put things in folders
  • Resolve typing errors
  • Cleanup/refactor (more)
  • Package cleanup: dependencies, etc
  • Add CI:
    • mypy
    • lint
    • build
    • coverage
  • More comprehensive tests
  • Clean up docstrings