Home

Awesome

jaxfg

build lint mypy codecov

<!-- vim-markdown-toc GFM --> <!-- vim-markdown-toc -->

jaxfg is a factor graph-based nonlinear least squares library for JAX. Typical applications include sensor fusion, SLAM, bundle adjustment, optimal control.

The premise: we provide a high-level interface for defining probability densities as factor graphs. MAP inference reduces to nonlinear optimization, which we accelerate by analyzing the structure of the graph. Repeated factor and variable types have operations vectorized, and the sparsity of graph connections is translated into sparse matrix operations.

Features:

This library is released as part of our IROS 2021 paper (more info in our core experiment repository here) and borrows heavily from a wide set of existing libraries, including GTSAM, Ceres Solver, minisam, SwiftFusion, and g2o. For technical background and concepts, GTSAM has a great set of tutorials.

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  # For options, pass in a --help flag

Development

If you're interested in extending this library to define your own factor graphs, we'd recommend first familiarizing yourself with:

  1. Pytrees in JAX: https://jax.readthedocs.io/en/latest/jax-101/05.1-pytrees.html
  2. Python dataclasses: https://docs.python.org/3/library/dataclasses.html
    • We currently take a "make everything a dataclass" philosophy for software engineering in this library. This is convenient for several reasons, but notably makes it easy for objects to be registered as pytree nodes. See jax_dataclasses for details on this.
  3. Type annotations: https://docs.python.org/3/library/typing.html
    • We rely on generics (typing.Generic and typing.TypeVar) particularly heavily. If you're familiar with C++ this should come very naturally (~templates).
  4. Explicit decorators for overrides/inheritance: https://github.com/mkorpela/overrides
    • The @overrides and @final decorators signal which methods are being and/or shouldn't be overridden. The same goes for @abc.abstractmethod.

From there, we have a few references for defining your own factor graphs, factors, and manifolds:

Current limitations

  1. In XLA, JIT compilation needs to happen for each unique set of input shapes. Modifying graph structures can thus introduce significant re-compilation overheads; this can restrict applications that are dynamic or online.
  2. Our marginalization implementation is not very good.

To-do

This library's still in development mode! Here's our TODO list: