Introduction to jaxlie
Jaxlie is a specialized library focused on implementations of Lie groups used predominantly in the domain of computer vision and robotics, especially for tasks involving rigid body transformations. It is built using JAX, a high-performance numerical computation library for machine learning research, and draws inspiration from the well-known C++ library called Sophus.
Core Components
Jaxlie provides high-level data classes to represent various Lie groups, primarily focusing on transformations in 2D and 3D spaces. These are:
-
SO2: This class handles rotations in a two-dimensional plane. It uses unit complex numbers for parameterization, representing angles on a circle (S¹).
-
SE2: This class is for proper rigid transformations in 2D, combining rotation and translation using a unit complex number along with x and y coordinates.
-
SO3: This class deals with three-dimensional rotations, parameterized using quaternions, which are mathematical entities that provide an efficient way to handle 3D rotations.
-
SE3: This class extends the concept of rigid transformations to 3D spaces, using quaternions for rotation and adding x, y, and z coordinates for translation.
Functionality
Each of these classes within jaxlie comes equipped with a variety of operations, which are essential for computational tasks in robotics and computer vision. Some of the key operations include:
-
exp() and log(): These operations convert between exponential coordinates and Lie groups, which are crucial for smooth interpolation and movement along a curve.
-
adjoint(), inverse(), and identity(): These functions perform transformations such as finding the adjoint matrix, evaluating inverse transformations, and returning the identity element of a group.
-
apply() and multiply(): These operations allow for applying transformations to points and multiplying group elements, respectively.
-
from_matrix() and as_matrix(): These functions convert transformations to and from matrix representations, aiding in compatibility with various matrix-based computations.
In addition, jaxlie supports forward- and reverse-mode automatic differentiation, which is beneficial for optimization on manifolds. Helpers for manifold optimization, serialization through the Flax library, and compatibility with standard JAX function transformations enhance its utility.
Features
Jaxlie includes several useful features to facilitate unique computational needs:
-
Sampling and Conversion Utilities: It offers utilities like uniform random sampling and conversion between quaternion representations and Euler angles in its SO3 class.
-
Optimization: There are tools designed for optimization on manifolds, making tasks like pose graph optimization and bundle adjustment more efficient.
-
Compatibility and Serialization: Standard JAX transformation compatibility ensures smooth integration with existing JAX workflows. Serialization is supported via the Flax library, improving storage and retrieval of data models.
Installation
Jaxlie can be easily installed for Python versions 3.7 and above:
pip install jaxlie
(Although releases are available for Python 3.6, they are no longer updated.)
Applications in the Community
Jaxlie has found applications in several influential projects:
-
jaxfg: This project applies jaxlie to tackle nonlinear least squares problems with block-sparse structures, aiding in complex tasks like pose graph optimization and bundle adjustment.
-
tensorf-jax: An unofficial implementation of Tensorial Radiance Fields, a technique for 3D representation, uses jaxlie. This demonstrates jaxlie's efficacy in rendering applications.
Academic Contribution
Jaxlie was initially developed for a research paper presented at IROS 2021. The library has played a role in advancing research in differentiable factor graph optimization. It is an integral part of efforts aimed at improving smoothers in learning contexts, with its academic background further validating its utility in scientific research.
Anyone benefiting from jaxlie is encouraged to cite the original paper to acknowledge its contribution to the field.