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(), andas_matrix()operations. (see ./examples/se3_example.py) - Taylor approximations near singularities.
- 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 jaxlieMisc
jaxlie was originally written when I was learning about Lie groups for our IROS 2021 paper
(link):
@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)}
}