brentyi/jaxlie

Rigid transforms + Lie groups in JAX

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:

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

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