brentyi / jaxls

Sparse nonlinear least squares for JAX

jaxls

status: work-in-progress

jaxls is a library for nonlinear least squares in JAX.

We provide a factor graph interface for specifying and solving least squares problems. We accelerate optimization by analyzing the structure of graphs: repeated factor and variable types are vectorized, and the sparsity of adjacency in the graph is translated into sparse matrix operations.

Features:

• Automatic sparse Jacobians.
• Optimization on manifolds; SO(2), SO(3), SE(2), and SE(3) implementations included.
• Nonlinear solvers: Levenberg-Marquardt and Gauss-Newton.
• Linear solvers: both direct (sparse Cholesky via CHOLMOD, on CPU) and iterative (Jacobi-preconditioned Conjugate Gradient).

Use cases are primarily in least squares problems that are inherently (1) sparse and (2) inefficient to solve with gradient-based methods. In robotics, these are ubiquitous across classical approaches to perception, planning, and control.

For the first iteration of this library, written for IROS 2021, see jaxfg. jaxls is a rewrite that aims to be faster and easier to use. For additional references, see inspirations like GTSAM, Ceres Solver, minisam, SwiftFusion, g2o.

Installation

jaxls supports python>=3.12.

For Cholesky factorization via CHOLMOD, scikit-sparse requires SuiteSparse:

# Via conda.
conda install conda-forge::suitesparse
# Via apt.
sudo apt update
sudo apt install -y libsuitesparse-dev
# Via brew.
brew install suite-sparse

Then, from your environment of choice:

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

Pose graph example

import jaxls
import jaxlie

Defining variables. Each variable is given an integer ID. They don't need to be contiguous.

pose_vars = [jaxls.SE2Var(0), jaxls.SE2Var(1)]

Defining factors. Factors are defined using a callable cost function and a set of arguments.

# Factors take two arguments:
# - A callable with signature `(jaxls.VarValues, *Args) -> jax.Array`.
# - A tuple of arguments: the type should be `tuple[*Args]`.
#
# All arguments should be PyTree structures. Variable types within the PyTree
# will be automatically detected.
factors = [
    # Cost on pose 0.
    jaxls.Factor.make(
        lambda vals, var, init: (vals[var] @ init.inverse()).log(),
        (pose_vars[0], jaxlie.SE2.from_xy_theta(0.0, 0.0, 0.0)),
    ),
    # Cost on pose 1.
    jaxls.Factor.make(
        lambda vals, var, init: (vals[var] @ init.inverse()).log(),
        (pose_vars[1], jaxlie.SE2.from_xy_theta(2.0, 0.0, 0.0)),
    ),
    # Cost between poses.
    jaxls.Factor.make(
        lambda vals, var0, var1, delta: (
            (vals[var0].inverse() @ vals[var1]) @ delta.inverse()
        ).log(),
        (pose_vars[0], pose_vars[1], jaxlie.SE2.from_xy_theta(1.0, 0.0, 0.0)),
    ),
]

Factors with similar structure, like the first two in this example, will be vectorized under-the-hood.

Solving optimization problems. We can set up the optimization problem, solve it, and print the solutions:

graph = jaxfg2.FactorGraph.make(factors, pose_vars)
solution = graph.solve()
print("All solutions", solution)
print("Pose 0", solution[pose_vars[0]])
print("Pose 1", solution[pose_vars[1]])

jaxls's Issues

Bug: cost depends on initial_assignment variable order

jaxfg_bug.py.txt

The issue is that the cost (sum of residual **2) depends on how the variables are ordered in initial_assignments.

RealVectorVariable Class Issue

https://github.com/brentyi/jaxfg/blob/master/jaxfg/core/_variables.py#L151

I tried to use the RealVectorVariable[dim] to instantiate a real vector variable. But I got an TypeError: Parameter list is too short. I fixed error by changing the code from

        class _RealVectorVariable(VariableBase[hints.Array]):
            @staticmethod
            @overrides
            @final
            def get_default_value() -> hints.Array:
                return jnp.zeros(dim)

to

        class _RealVectorVariable(VariableBase[hints.Array]):
            @classmethod
            @overrides
            @final
            def get_default_value(cls) -> hints.Array:
                return jnp.zeros(dim)

A possible bug with the order in solution_assignments?

Hi Brent,

It seems that the solution_assignments is not consistent with the order of pose_variables in pose_graph_simple.py(https://github.com/brentyi/jaxfg/blob/f5204945bb6afa444810e6163e9a913bdbdd636b/scripts/pose_graph_simple.py). For the original code, the outputs of

# Grab and print a single variable value at a time.
print("First pose (jaxlie.SE2 object):")
print(solution_assignments.get_value(pose_variables[0]))
print()

print("Second pose (jaxlie.SE2 object):")
print(solution_assignments.get_value(pose_variables[1]))
print()

was

First pose (jaxlie.SE2 object): 
SE2(unit_complex=[1. 0.], xy=[0.33333 0.     ])

Second pose (jaxlie.SE2 object):
SE2(unit_complex=[1. 0.], xy=[1.66667 0.     ])

However, if I change the code of creating the factor graph to (switch the order of the two PriorFactors)

factors: List[jaxfg.core.FactorBase] = [
jaxfg.geometry.PriorFactor.make(
        variable=pose_variables[1],
        mu=jaxlie.SE2.from_xy_theta(2.0, 0.0, 0.0),
        noise_model=jaxfg.noises.DiagonalGaussian(jnp.ones(3)),
    ),
    jaxfg.geometry.PriorFactor.make(
        variable=pose_variables[0],
        mu=jaxlie.SE2.from_xy_theta(0.0, 0.0, 0.0),
        noise_model=jaxfg.noises.DiagonalGaussian(jnp.ones(3)),
    ),
    jaxfg.geometry.BetweenFactor.make(
        variable_T_world_a=pose_variables[0],
        variable_T_world_b=pose_variables[1],
        T_a_b=jaxlie.SE2.from_xy_theta(1.0, 0.0, 0.0),
        noise_model=jaxfg.noises.DiagonalGaussian(jnp.ones(3)),
    ),
]

Then, the outputs become

First pose (jaxlie.SE2 object):
SE2(unit_complex=[1. 0.], xy=[1.66667 0.     ])

Second pose (jaxlie.SE2 object):
SE2(unit_complex=[1. 0.], xy=[0.33333 0.     ])

It seems that the order in the solution_assignments is determined by the order of the factors?

Changing and indexing prior factors after making graph

Thanks for the library! I'm trying to update the mu of a prior factor based on the optimized pose of another graph's node. However I'm stuck figuring out how the rows of graph.factor_stacks[1].factor.mu.unit_complex_xy correspond to which nodes. Any help? Thanks

SLAM example

Hello

This project seems very interesting, thanks for sharing it.

Could you please provide an example of a SLAM implementation using in the examples section using a popular dataset (KITTI or oxford dataset for example)?

This would be a very helpful example to learn how to use the library.

Thanks in advance

About the performance compare between other graph optimization framework.

Hi, is there any comparison of performance between this library and others( such as ceres solver, g2o, gtsam)?
And, since the jaxfg is based on Jax, can the jaxfg use cuda to significantly speed up optimization speed.?

Open source license?

Hi, this projects looks awesome! I definitely see some potential synergy/overlap with JAXopt: https://github.com/google/jaxopt

Would you consider adding an open source license so others can use your code? GitHub has some good guidance here:
https://docs.github.com/en/github/creating-cloning-and-archiving-repositories/creating-a-repository-on-github/licensing-a-repository

Thanks!