jaxlie

[ API reference ] [ PyPI ]

jaxlie is a library containing implementations of Lie groups commonly used for rigid body transformations, targeted at computer vision & robotics applications written in JAX. Heavily inspired by the C++ library Sophus.

We implement Lie groups as high-level (data)classes:

Where each group supports:

• Forward- and reverse-mode AD-friendly exp(), log(), adjoint(), apply(), multiply(), inverse(), identity(), from_matrix(), and as_matrix() operations. (see ./examples/se3_example.py)
• Helpers for optimization on manifolds (see ./examples/se3_optimization.py, jaxlie.manifold.*).
• Compatibility with standard JAX function transformations. (see ./examples/vmap_example.py)
• (Un)flattening as pytree nodes.
• Serialization using flax.

We also implement various common utilities for things like uniform random sampling (sample_uniform()) and converting from/to Euler angles (in the SO3 class).

Install (Python >=3.7)

# Python 3.6 releases also exist, but are no longer being updated.
pip install jaxlie

In the wild

• jaxfg applies jaxlie to nonlinear least squares problems with block-sparse structure. (for pose graph optimization, bundle adjustment, etc)
• tensorf-jax is an unofficial implementation of Tensorial Radiance Fields (Chen et al, ECCV 2022) using jaxlie.

Misc

jaxlie was originally written for our IROS 2021 paper (link). If it's useful for you, you're welcome to cite:

@inproceedings{yi2021iros,
    author={Brent Yi and Michelle Lee and Alina Kloss and Roberto Mart\'in-Mart\'in and Jeannette Bohg},
    title = {Differentiable Factor Graph Optimization for Learning Smoothers},
    year = 2021,
    BOOKTITLE = {2021 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)}
}