Awesome

jaxlie

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:

<table> <thead> <tr> <th>Group</th> <th>Description</th> <th>Parameterization</th> </tr> </thead> <tbody valign="top"> <tr> <td><code>jaxlie.<strong>SO2</strong></code></td> <td>Rotations in 2D.</td> <td><em>(real, imaginary):</em> unit complex (∈ S<sup>1</sup>)</td> </tr> <tr> <td><code>jaxlie.<strong>SE2</strong></code></td> <td>Proper rigid transforms in 2D.</td> <td><em>(real, imaginary, x, y):</em> unit complex & translation</td> </tr> <tr> <td><code>jaxlie.<strong>SO3</strong></code></td> <td>Rotations in 3D.</td> <td><em>(qw, qx, qy, qz):</em> wxyz quaternion (∈ S<sup>3</sup>)</td> </tr> <tr> <td><code>jaxlie.<strong>SE3</strong></code></td> <td>Proper rigid transforms in 3D.</td> <td><em>(qw, qx, qy, qz, x, y, z):</em> wxyz quaternion & translation</td> </tr> </tbody> </table>

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)
Taylor approximations near singularities.
Helpers for optimization on manifolds (see ./examples/se3_optimization.py, <code>jaxlie.<strong>manifold.*</strong></code>).
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

Misc

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