/jaxlie

Rigid transforms + Lie groups in JAX

Primary LanguagePythonMIT LicenseMIT

jaxlie

build mypy lint codecov pypi_dowlnoads

[ 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:

Group Description Parameterization
jaxlie.SO2 Rotations in 2D. (real, imaginary): unit complex (∈ S1)
jaxlie.SE2 Proper rigid transforms in 2D. (real, imaginary, x, y): unit complex & translation
jaxlie.SO3 Rotations in 3D. (qw, qx, qy, qz): wxyz quaternion (∈ S3)
jaxlie.SE3 Proper rigid transforms in 3D. (qw, qx, qy, qz, x, y, z): wxyz quaternion & translation

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)
  • Broadcasting for leading axes.
  • (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


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)}
}