From 5b5d5051cb8732f708b5bbe77dd2bfa915896a5f Mon Sep 17 00:00:00 2001
From: The jax3d Authors <no-reply@google.com>
Date: Thu, 27 Oct 2022 12:47:52 -0700
Subject: [PATCH] Add differentiable rigid body SE3 transforms.

PiperOrigin-RevId: 484325168
---
 jax3d/math/quaternion.py      | 225 ++++++++++++++++++++++++++++++
 jax3d/math/quaternion_test.py | 100 +++++++++++++
 jax3d/math/rigid_body.py      | 213 ++++++++++++++++++++++++++++
 jax3d/math/rigid_body_test.py | 254 ++++++++++++++++++++++++++++++++++
 4 files changed, 792 insertions(+)
 create mode 100644 jax3d/math/quaternion.py
 create mode 100644 jax3d/math/quaternion_test.py
 create mode 100644 jax3d/math/rigid_body.py
 create mode 100644 jax3d/math/rigid_body_test.py

diff --git a/jax3d/math/quaternion.py b/jax3d/math/quaternion.py
new file mode 100644
index 0000000..8825fa5
--- /dev/null
+++ b/jax3d/math/quaternion.py
@@ -0,0 +1,225 @@
+# Copyright 2022 The jax3d Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Quaternion math.
+
+This module assumes the xyzw quaternion format where xyz is the imaginary part
+and w is the real part.
+
+Functions in this module support both batched and unbatched quaternions.
+"""
+from jax import numpy as jnp
+from jax.numpy import linalg
+
+
+def safe_acos(t, eps=1e-7):
+  """A safe version of arccos which avoids evaluating at -1 or 1."""
+  return jnp.arccos(jnp.clip(t, -1.0 + eps, 1.0 - eps))
+
+
+def im(q):
+  """Fetch the imaginary part of the quaternion."""
+  return q[..., :3]
+
+
+def re(q):
+  """Fetch the real part of the quaternion."""
+  return q[..., 3:]
+
+
+def identity():
+  return jnp.array([0.0, 0.0, 0.0, 1.0])
+
+
+def conjugate(q):
+  """Compute the conjugate of a quaternion."""
+  return jnp.concatenate([-im(q), re(q)], axis=-1)
+
+
+def inverse(q):
+  """Compute the inverse of a quaternion."""
+  return normalize(conjugate(q))
+
+
+def normalize(q):
+  """Normalize a quaternion."""
+  return q / norm(q)
+
+
+def norm(q):
+  return linalg.norm(q, axis=-1, keepdims=True)
+
+
+def multiply(q1, q2):
+  """Multiply two quaternions."""
+  c = (re(q1) * im(q2)
+       + re(q2) * im(q1)
+       + jnp.cross(im(q1), im(q2)))
+  w = re(q1) * re(q2) - jnp.dot(im(q1), im(q2))
+  return jnp.concatenate([c, w], axis=-1)
+
+
+def rotate(q, v):
+  """Rotate a vector using a quaternion."""
+  # Create the quaternion representation of the vector.
+  q_v = jnp.concatenate([v, jnp.zeros_like(v[..., :1])], axis=-1)
+  return im(multiply(multiply(q, q_v), conjugate(q)))
+
+
+def log(q, eps=1e-8):
+  """Computes the quaternion logarithm.
+
+  References:
+    https://en.wikipedia.org/wiki/Quaternion#Exponential,_logarithm,_and_power_functions
+
+  Args:
+    q: the quaternion in (x,y,z,w) format.
+    eps: an epsilon value for numerical stability.
+
+  Returns:
+    The logarithm of q.
+  """
+  mag = linalg.norm(q, axis=-1, keepdims=True)
+  v = im(q)
+  s = re(q)
+  w = jnp.log(mag)
+  denom = jnp.maximum(
+      linalg.norm(v, axis=-1, keepdims=True), eps * jnp.ones_like(v))
+  xyz = v / denom * safe_acos(s / eps)
+  return jnp.concatenate((xyz, w), axis=-1)
+
+
+def exp(q, eps=1e-8):
+  """Computes the quaternion exponential.
+
+  References:
+    https://en.wikipedia.org/wiki/Quaternion#Exponential,_logarithm,_and_power_functions
+
+  Args:
+    q: the quaternion in (x,y,z,w) format or (x,y,z) if is_pure is True.
+    eps: an epsilon value for numerical stability.
+
+  Returns:
+    The exponential of q.
+  """
+  is_pure = q.shape[-1] == 3
+  if is_pure:
+    s = jnp.zeros_like(q[..., -1:])
+    v = q
+  else:
+    v = im(q)
+    s = re(q)
+
+  norm_v = linalg.norm(v, axis=-1, keepdims=True)
+  exp_s = jnp.exp(s)
+  w = jnp.cos(norm_v)
+  xyz = jnp.sin(norm_v) * v / jnp.maximum(norm_v, eps * jnp.ones_like(norm_v))
+  return exp_s * jnp.concatenate((xyz, w), axis=-1)
+
+
+def to_rotation_matrix(q):
+  """Constructs a rotation matrix from a quaternion.
+
+  Args:
+    q: a (*,4) array containing quaternions.
+
+  Returns:
+    A (*,3,3) array containing rotation matrices.
+  """
+  x, y, z, w = jnp.split(q, 4, axis=-1)
+  s = 1.0 / jnp.sum(q ** 2, axis=-1)
+  return jnp.stack([
+      jnp.stack([1 - 2 * s * (y ** 2 + z ** 2),
+                 2 * s * (x * y - z * w),
+                 2 * s * (x * z + y * w)], axis=0),
+      jnp.stack([2 * s * (x * y + z * w),
+                 1 - s * 2 * (x ** 2 + z ** 2),
+                 2 * s * (y * z - x * w)], axis=0),
+      jnp.stack([2 * s * (x * z - y * w),
+                 2 * s * (y * z + x * w),
+                 1 - 2 * s * (x ** 2 + y ** 2)], axis=0),
+  ], axis=0)
+
+
+def from_rotation_matrix(m, eps=1e-9):
+  """Construct quaternion from a rotation matrix.
+
+  Args:
+    m: a (*,3,3) array containing rotation matrices.
+    eps: a small number for numerical stability.
+
+  Returns:
+    A (*,4) array containing quaternions.
+  """
+  trace = jnp.trace(m)
+  m00 = m[..., 0, 0]
+  m01 = m[..., 0, 1]
+  m02 = m[..., 0, 2]
+  m10 = m[..., 1, 0]
+  m11 = m[..., 1, 1]
+  m12 = m[..., 1, 2]
+  m20 = m[..., 2, 0]
+  m21 = m[..., 2, 1]
+  m22 = m[..., 2, 2]
+
+  def tr_positive():
+    sq = jnp.sqrt(trace + 1.0) * 2.  # sq = 4 * w.
+    w = 0.25 * sq
+    x = jnp.divide(m21 - m12, sq)
+    y = jnp.divide(m02 - m20, sq)
+    z = jnp.divide(m10 - m01, sq)
+    return jnp.stack((x, y, z, w), axis=-1)
+
+  def cond_1():
+    sq = jnp.sqrt(1.0 + m00 - m11 - m22 + eps) * 2.  # sq = 4 * x.
+    w = jnp.divide(m21 - m12, sq)
+    x = 0.25 * sq
+    y = jnp.divide(m01 + m10, sq)
+    z = jnp.divide(m02 + m20, sq)
+    return jnp.stack((x, y, z, w), axis=-1)
+
+  def cond_2():
+    sq = jnp.sqrt(1.0 + m11 - m00 - m22 + eps) * 2.  # sq = 4 * y.
+    w = jnp.divide(m02 - m20, sq)
+    x = jnp.divide(m01 + m10, sq)
+    y = 0.25 * sq
+    z = jnp.divide(m12 + m21, sq)
+    return jnp.stack((x, y, z, w), axis=-1)
+
+  def cond_3():
+    sq = jnp.sqrt(1.0 + m22 - m00 - m11 + eps) * 2.  # sq = 4 * z.
+    w = jnp.divide(m10 - m01, sq)
+    x = jnp.divide(m02 + m20, sq)
+    y = jnp.divide(m12 + m21, sq)
+    z = 0.25 * sq
+    return jnp.stack((x, y, z, w), axis=-1)
+
+  def cond_idx(cond):
+    cond = jnp.expand_dims(cond, -1)
+    cond = jnp.tile(cond, [1] * (len(m.shape) - 2) + [4])
+    return cond
+
+  where_2 = jnp.where(cond_idx(m11 > m22), cond_2(), cond_3())
+  where_1 = jnp.where(cond_idx((m00 > m11) & (m00 > m22)), cond_1(), where_2)
+  return jnp.where(cond_idx(trace > 0), tr_positive(), where_1)
+
+
+def from_axis_angle(axis, theta):
+  """Constructs a quaternion for the given axis/angle rotation."""
+  qx = axis[0] * jnp.sin(theta / 2)
+  qy = axis[1] * jnp.sin(theta / 2)
+  qz = axis[2] * jnp.sin(theta / 2)
+  qw = jnp.cos(theta / 2)
+
+  return jnp.squeeze(jnp.array([qx, qy, qz, qw]))
\ No newline at end of file
diff --git a/jax3d/math/quaternion_test.py b/jax3d/math/quaternion_test.py
new file mode 100644
index 0000000..00d1719
--- /dev/null
+++ b/jax3d/math/quaternion_test.py
@@ -0,0 +1,100 @@
+# Copyright 2022 The jax3d Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for quaternions."""
+
+import functools
+import math
+import unittest
+
+from jax import random
+import jax.numpy as jnp
+from jax3d.math import quaternion
+import pytest
+
+
+TEST_BATCH_SIZE = 128
+
+
+class QuaternionTest(unittest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self._seed = 42
+    self._key = random.PRNGKey(self._seed)
+
+  def test_identity(self):
+    identity = quaternion.identity()
+    self.assertLen(identity, 4)
+    self.assertEqual(identity.tolist(), [0.0, 0.0, 0.0, 1.0])
+
+  @pytest.mark.parametrize(('single', (4,)), ('batched', (TEST_BATCH_SIZE, 4)))
+  @pytest.mark.parametrize('shape', )
+  def test_real_imaginary_part(self, shape):
+    if len(shape) > 1:
+      num_quaternions = shape[0]
+    else:
+      num_quaternions = 1
+    random_quat = random.uniform(self._key, shape=shape)
+    imaginary = quaternion.im(random_quat)
+    real = quaternion.re(random_quat)
+
+    # The first three components are imaginary and the fourth is real.
+    self.assertEqual(jnp.prod(jnp.array(imaginary.shape)), num_quaternions * 3)
+    self.assertEqual(jnp.prod(jnp.array(real.shape)), num_quaternions)
+    self.assertEqual(random_quat[..., :3].tolist(), imaginary[..., :].tolist())
+    self.assertEqual(random_quat[..., 3:].tolist(), real[..., :].tolist())
+
+  @pytest.mark.parametrize('batch', [None, TEST_BATCH_SIZE])
+  @pytest.mark.parametrize('func', [random.uniform, jnp.ones, jnp.zeros])
+  @pytest.mark.parametrize('sign', [-1, 1])
+  def test_safe_acos(self, batch, func, sign):
+    # We need a seed to generate random numbers.
+    if func == random.uniform:
+      func = functools.partial(func, key=self._key)
+
+    if batch:
+      shape = (batch, 4)
+    else:
+      shape = (4,)
+    t = sign * func(shape=shape)
+
+    output = quaternion.safe_acos(t)
+
+    # All elements must be within the range of the arc-cosine function.
+    self.assertTrue(jnp.all(output > 0))
+    self.assertTrue(jnp.all(output < math.pi))
+
+  @pytest.mark.parametrize(('single', None), ('batched', TEST_BATCH_SIZE))
+  def test_conjugate(self, batch):
+    if batch:
+      shape = (batch, 4)
+    else:
+      shape = (4,)
+    quat = random.uniform(self._key, shape=shape)
+    conjugate = quaternion.conjugate(quat)
+    self.assertTrue(jnp.all(-1 * quat[..., :3] == conjugate[..., :3]))
+    self.assertTrue(jnp.all(quat[..., 3:] == conjugate[..., 3:]))
+
+  @pytest.mark.parametrize(('single', None), ('batched', TEST_BATCH_SIZE))
+  def test_normalize(self, batch):
+    eps = 1e-6
+    if batch:
+      shape = (batch, 4)
+    else:
+      shape = (4,)
+    q = random.uniform(self._key, shape=shape)
+    self.assertTrue(jnp.all(jnp.abs(quaternion.norm(q) - 1) > eps))
+    q_norm = quaternion.normalize(q)
+    self.assertTrue(jnp.all(jnp.abs(quaternion.norm(q_norm) - 1) < eps))
diff --git a/jax3d/math/rigid_body.py b/jax3d/math/rigid_body.py
new file mode 100644
index 0000000..96ae3ea
--- /dev/null
+++ b/jax3d/math/rigid_body.py
@@ -0,0 +1,213 @@
+# Copyright 2022 The jax3d Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utilization functions for handling rigid body transforms."""
+import jax
+from jax import numpy as jnp
+
+
+def matmul(a, b):
+  """jnp.matmul defaults to bfloat16, but this helper function doesn't."""
+  return jnp.matmul(a, b, precision=jax.lax.Precision.HIGHEST)
+
+
+def divide_safe(numerator: jnp.ndarray,
+                denominator: jnp.ndarray,
+                eps: float = 1e-7) -> jnp.ndarray:
+  """Division of jnp.ndarray's with zero denominator safety."""
+  denominator_ = jnp.where(denominator < eps, 1.0, denominator)
+  return jnp.divide(numerator, denominator_)
+
+
+@jax.jit
+def skew(w: jnp.ndarray) -> jnp.ndarray:
+  """Build a skew matrix ("cross product matrix") for vector w.
+
+  Modern Robotics Eqn 3.30.
+
+  Args:
+    w: (3,) A 3-vector
+
+  Returns:
+    W: (3, 3) A skew matrix such that W @ v == w x v
+  """
+  w = jnp.reshape(w, (3))
+  return jnp.array([[0.0, -w[2], w[1]],
+                    [w[2], 0.0, -w[0]],
+                    [-w[1], w[0], 0.0]])
+
+
+def rotation_translation_to_homogeneous_transform(
+    rotation: jnp.ndarray, translation: jnp.ndarray) -> jnp.ndarray:
+  """Rotation and translation to homogeneous transform.
+
+  Args:
+    R: (3, 3) An orthonormal rotation matrix.
+    p: (3,) A 3-vector representing an offset.
+
+  Returns:
+    X: (4, 4) The homogeneous transformation matrix described by rotating by R
+      and translating by p.
+  """
+  translation = jnp.reshape(translation, (3, 1))
+  return jnp.block([[rotation, translation],
+                    [jnp.array([[0.0, 0.0, 0.0, 1.0]])]])
+
+
+def exp_so3(w: jnp.ndarray, theta: float) -> jnp.ndarray:
+  """Exponential map from Lie algebra so3 to Lie group SO3.
+
+  Modern Robotics Eqn 3.51, a.k.a. Rodrigues' formula.
+
+  Args:
+    w: (3,) An axis of rotation. This is assumed to be a unit-vector.
+    theta: An angle of rotation.
+
+  Returns:
+    rotation: (3, 3) An orthonormal rotation matrix representing a rotation of
+      magnitude theta about axis w.
+  """
+  w_skew = skew(w)
+  return (jnp.eye(3)
+          + jnp.sin(theta) * w_skew
+          + (1.0 - jnp.cos(theta)) * matmul(w_skew, w_skew))
+
+
+def exp_se3(screw_axis: jnp.ndarray, theta: float) -> jnp.ndarray:
+  """Exponential map from Lie algebra so3 to Lie group SO3.
+
+  Modern Robotics Eqn 3.88.
+
+  Args:
+    screw_axis: (6,) A screw axis of motion.
+    theta: Magnitude of motion.
+
+  Returns:
+    a_X_b: (4, 4) The homogeneous transformation matrix attained by integrating
+      motion of magnitude theta about S for one second.
+  """
+  w, v = jnp.split(screw_axis, 2)
+  w_skew = skew(w)
+  rotation = exp_so3(w_skew, theta)
+  translation = matmul((theta * jnp.eye(3) + (1.0 - jnp.cos(theta)) * w_skew
+                        + (theta - jnp.sin(theta)) * matmul(w_skew, w_skew)), v)
+  return rotation_translation_to_homogeneous_transform(rotation, translation)
+
+
+def to_homogenous(v):
+  return jnp.concatenate([v, jnp.ones_like(v[..., :1])], axis=-1)
+
+
+def from_homogenous(v):
+  return v[..., :3] / v[..., -1:]
+
+
+def se3_to_rotation_translation(
+    se3: jnp.ndarray) -> tuple[jnp.ndarray, jnp.ndarray]:
+  """Computes rotation and translation from 6D smooth manifold."""
+  w, v = jnp.split(se3, 2, axis=-1)
+  theta = jnp.linalg.norm(w, axis=-1)
+  w = w / theta[..., None]
+  rot_axis = jnp.concatenate((w, v), axis=-1)
+  homo_trans = exp_se3(rot_axis, theta)
+  rotation_matrix = homo_trans[..., :3, :3]
+  translation_vector = homo_trans[..., :3, -1]
+  return rotation_matrix, translation_vector
+
+
+def hat_inv(skew_sym_matrix: jnp.ndarray) -> jnp.ndarray:
+  """Computes the inverse Hat operator of a skew symmetric matrix.
+
+  References:
+    https://en.wikipedia.org/wiki/Hat_operator
+
+  Args:
+    skew_sym_matrix: a skew symmetric matrix of size 3x3
+
+  Returns:
+    a vector of length 3
+  """
+  x = skew_sym_matrix[..., 2, 1]
+  y = skew_sym_matrix[..., 0, 2]
+  z = skew_sym_matrix[..., 1, 0]
+
+  v = jnp.stack((x, y, z), axis=-1)
+  return v
+
+
+def _taylor_first(x, nth=10):
+  """Taylor expansion of sin(x)/x."""
+  ans = jnp.zeros_like(x)
+  denom = 1.
+  for i in range(nth + 1):
+    if i > 0:
+      denom = denom * (2 * i) * (2 * i + 1)
+    ans = ans + (-1)**i * x**(2 * i) / denom
+  return ans
+
+
+def _taylor_second(x, nth=10):
+  """Taylor expansion of (1-cos(x))/x**2."""
+  ans = jnp.zeros_like(x)
+  denom = 1.
+  for i in range(nth + 1):
+    denom = denom * (2 * i + 1) * (2 * i + 2)
+    ans = ans + (-1)**i * x**(2 * i) / denom
+  return ans
+
+
+def _taylor_third(x, nth=10):
+  """Taylor expansion of (x-sin(x))/x**3."""
+  ans = jnp.zeros_like(x)
+  denom = 1.
+  for i in range(nth + 1):
+    denom = denom * (2 * i + 2) * (2 * i + 3)
+    ans = ans + (-1)**i * x**(2 * i) / denom
+  return ans
+
+
+def rotation_translation_to_se3(rotation_matrix: jnp.ndarray,
+                                translation_vector: jnp.ndarray,
+                                eps: float = 1e-7) -> jnp.ndarray:
+  """Computes the pseudo inverse of a smooth 6D vector of a rigid transform.
+
+  References:
+    https://jinyongjeong.github.io/Download/SE3/jlblanco2010geometry3d_techrep.pdf
+
+  Args:
+    rotation_matrix: a 3x3 rotation matrix.
+    translation_vector: an array of length 3.
+    eps: an epsilon for avoiding division by zero.
+
+  Returns:
+    a 6D se3 representation of the given rigid transform.
+  """
+  trace_rotation = rotation_matrix[..., 0, 0] + rotation_matrix[
+      ..., 1, 1] + rotation_matrix[..., 2, 2]
+  cos_theta = 0.5 * (trace_rotation - 1)
+  sin_theta = jnp.sqrt(1 - cos_theta**2)
+  theta = jnp.arccos(cos_theta)
+  log_rotation = rotation_matrix - rotation_matrix.T
+  log_rotation = log_rotation * (theta / (2 * sin_theta))
+  w = hat_inv(log_rotation)
+  wx = skew(w)
+  identity = jnp.eye(3, dtype=jnp.float32)
+  first_coeff = _taylor_first(theta)
+  second_coeff = _taylor_second(theta)
+  invesrse_v_matrix = identity - 0.5 * wx + (
+      1 - first_coeff /
+      (2 * second_coeff)) / (theta**2 + eps) * matmul(wx, wx)
+  u = divide_safe(invesrse_v_matrix @ translation_vector, theta)
+  wu = jnp.concatenate((w, u), axis=-1)
+  return wu
diff --git a/jax3d/math/rigid_body_test.py b/jax3d/math/rigid_body_test.py
new file mode 100644
index 0000000..1d50b40
--- /dev/null
+++ b/jax3d/math/rigid_body_test.py
@@ -0,0 +1,254 @@
+# Copyright 2022 The jax3d Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+import unittest
+
+import jax
+from jax import numpy as jnp
+from jax import random
+from jax3d.math import quaternion
+from jax3d.math import rigid_body
+import numpy as np
+import pytest
+
+
+TEST_BATCH_SIZE = 128
+SAMPLE_POINTS = [(0, 0, 0), (1, 0, 0), (-1, 0, 0), (0, 1, 0), (0, -1, 0),
+                 (0, 0, 1), (0, 0, -1)]
+
+
+class RigidBodyTest(unittest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self._seed = 42
+    self._key = random.PRNGKey(self._seed)
+
+  @staticmethod
+  def _process_parameters(batch, vector_size=4):
+    if batch:
+      shape = (batch, vector_size)
+      num_vectors = batch
+    else:
+      shape = (vector_size,)
+      num_vectors = 1
+
+    return shape, num_vectors
+
+  def get_random_vector(self, func, shape):
+    if func == random.uniform:
+      self._key, _ = random.split(self._key)
+      return func(shape=shape, key=self._key)
+    else:
+      return func(shape=shape)
+
+  @pytest.mark.parametrize('batch', [None, TEST_BATCH_SIZE])
+  @pytest.mark.parametrize('func', [random.uniform, jnp.ones])
+  @pytest.mark.parametrize('sign', [-1, 1])
+  def test_from_homogenous(self, batch, func, sign):
+    shape, num_vectors = self._process_parameters(batch, 4)
+    vector = sign * self.get_random_vector(func, shape=shape)
+    output = rigid_body.from_homogenous(vector)
+    self.assertEqual(jnp.prod(jnp.array(output.shape)), num_vectors * 3)
+    np.testing.assert_array_equal(output, vector[..., :3] / vector[..., -1:])
+
+  @pytest.mark.parametrize('batch', [None, TEST_BATCH_SIZE])
+  @pytest.mark.parametrize('func', [random.uniform, jnp.ones, jnp.zeros])
+  @pytest.mark.parametrize('sign', [-1, 1])
+  def test_to_homogenous(self, batch, func, sign):
+    shape, num_vectors = self._process_parameters(batch, 3)
+    vector = sign * self.get_random_vector(func, shape=shape)
+    output = rigid_body.to_homogenous(vector)
+    self.assertEqual(jnp.prod(jnp.array(output.shape)), num_vectors * 4)
+    np.testing.assert_array_equal(output[..., :3], vector)
+    np.testing.assert_array_equal(output[..., -1:], 1.0)
+
+  @pytest.mark.parametrize('func', [random.uniform, jnp.ones, jnp.zeros])
+  @pytest.mark.parametrize('sign1', [-1, 1])
+  @pytest.mark.parametrize('sign2', [-1, 1])
+  def test_skew_matrix(self, func, sign1, sign2):
+    # The skew function does not support batched operation.
+    shape, _ = self._process_parameters(None, 3)
+    w = sign1 * self.get_random_vector(func, shape=shape)
+    v = sign2 * self.get_random_vector(func, shape=shape)
+    skew_matrix = rigid_body.skew(w)
+
+    # Properties of a skew symmetric matrix.
+    self.assertEqual(jnp.trace(skew_matrix), 0)
+    np.testing.assert_array_equal(-1 * jnp.transpose(skew_matrix), skew_matrix)
+
+    # Does the matrix approximate the actual cross product?
+    expected_cross_product = jnp.cross(w, v)
+    predicted_cross_product = jnp.matmul(skew_matrix, v)
+    np.testing.assert_allclose(
+        expected_cross_product, predicted_cross_product, atol=1E-5, rtol=1E-5)
+
+  @pytest.mark.parametrize('func', [random.uniform, jnp.ones])
+  @pytest.mark.parametrize('sign1', [-1, 1])
+  @pytest.mark.parametrize('sign2', [-1, 1])
+  def test_exp_so3(self, func, sign1, sign2):
+    shape, num_vectors = self._process_parameters(None, 3)
+
+    # Generate a normalized axis of rotation and the angle of rotation.
+    w = sign1 * self.get_random_vector(func, shape=shape)
+    w = w / jnp.linalg.norm(w)
+
+    theta = sign2 * self.get_random_vector(func, shape=(num_vectors, 1))
+    output = rigid_body.exp_so3(w, theta)
+
+    # Verify orthonormality.
+    np.testing.assert_allclose(
+        jnp.matmul(jnp.transpose(output), output),
+        jnp.eye(3),
+        atol=1E-5,
+        rtol=1E-5)
+    np.testing.assert_allclose(
+        jnp.matmul(output, jnp.transpose(output)),
+        jnp.eye(3),
+        atol=1E-5,
+        rtol=1E-5)
+
+  @pytest.mark.parametrize('axis', [[1, 0, 0], [0, 1, 0], [0, 0, 1]])
+  @pytest.mark.parametrize('theta', [x * math.pi / 4 for x in range(8)])
+  @pytest.mark.parametrize('sign', [-1, 1])
+  @pytest.mark.parametrize('pt_input', SAMPLE_POINTS)
+  def test_exp_so3_rotation(self, axis, theta, sign, pt_input):
+    axis = jnp.array(axis)
+    theta = jnp.array(sign * theta)
+    pt_input = jnp.array(pt_input)
+    theta = jnp.expand_dims(theta, 0)
+
+    axis = axis / jnp.linalg.norm(axis)
+    rotation_matrix = rigid_body.exp_so3(axis, theta)
+    predicted_output = jnp.matmul(rotation_matrix, pt_input)
+
+    # Use a quaternion to compute the rotation and use it as a comparison.
+    quat = quaternion.from_axis_angle(axis, theta)
+    quaternion_output = quaternion.rotate(quat, pt_input)
+    np.testing.assert_allclose(
+        predicted_output, quaternion_output, atol=1E-5, rtol=1E-5)
+
+  @pytest.mark.parametrize('func', [random.uniform, jnp.ones, jnp.zeros])
+  @pytest.mark.parametrize('sign1', [-1, 1])
+  @pytest.mark.parametrize('sign2', [-1, 1])
+  @pytest.mark.parametrize('sign3', [-1, 1])
+  def test_rotation_translation_to_homogeneous_transform(self, func, sign1, sign2, sign3):
+    shape, num_vectors = self._process_parameters(None, 3)
+    w = sign1 * self.get_random_vector(func, shape=shape)
+    w = w / jnp.linalg.norm(w)
+
+    theta = sign2 * self.get_random_vector(func, shape=(num_vectors, 1))
+    r = rigid_body.exp_so3(w, theta)
+
+    p = sign3 * self.get_random_vector(func, shape=(num_vectors, 3))
+    output = rigid_body.rotation_translation_to_homogeneous_transform(r, p)
+    self.assertEqual(output.shape, (4, 4))
+    np.testing.assert_array_equal(jnp.squeeze(r), jnp.squeeze(output[0:3, 0:3]))
+    np.testing.assert_array_equal(jnp.squeeze(p), jnp.squeeze(output[0:3, 3]))
+    np.testing.assert_array_equal(
+        jnp.squeeze(jnp.array([0.0, 0.0, 0.0, 1.0])), jnp.squeeze(output[3, :]))
+
+  @pytest.mark.parametrize('func', [random.uniform, jnp.ones, jnp.zeros])
+  @pytest.mark.parametrize('sign', [-1, 1])
+  @pytest.mark.parametrize('pt', SAMPLE_POINTS)
+  def test_exp_se3_only_rotation(self, func, sign, pt):
+    shape, _ = self._process_parameters(None, 3)
+    pt = jnp.array(pt)
+    w = sign * self.get_random_vector(func, shape=shape)
+    v = jnp.zeros(shape=shape)
+    theta = jnp.linalg.norm(w, axis=-1)
+    w = w / theta[..., jnp.newaxis]
+    screw_axis = jnp.concatenate([w, v], axis=-1)
+    transform = rigid_body.exp_se3(screw_axis, theta)
+
+    quat = quaternion.from_axis_angle(w, theta)
+    pt_rotated = quaternion.rotate(quat, pt)
+
+    self.assertEqual(transform.shape, (4, 4))
+    pt_rotated_tf = rigid_body.from_homogenous(
+        jnp.matmul(transform, rigid_body.to_homogenous(pt)))
+    np.testing.assert_allclose(pt_rotated_tf, pt_rotated, atol=1E-5, rtol=1E-5)
+
+  @pytest.mark.parametrize('func', [random.uniform, jnp.ones, jnp.zeros])
+  @pytest.mark.parametrize('sign', [-1, 1])
+  @pytest.mark.parametrize('pt', SAMPLE_POINTS)
+  def test_exp_se3_only_translation(self, func, sign, pt):
+    shape, _ = self._process_parameters(None, 3)
+    w = jnp.zeros(shape=shape)
+    v = sign * self.get_random_vector(func, shape=shape)
+    theta = jnp.array(1)
+    screw_axis = jnp.concatenate([w, v], axis=-1)
+    transform = rigid_body.exp_se3(screw_axis, theta)
+
+    pt = jnp.array(pt)
+    pt_translated = pt + v
+
+    self.assertEqual(transform.shape, (4, 4))
+    pt_translated_tf = rigid_body.from_homogenous(
+        jnp.matmul(transform, rigid_body.to_homogenous(pt)))
+    np.testing.assert_allclose(
+        pt_translated_tf, pt_translated, atol=1E-5, rtol=1E-5)
+
+  @pytest.mark.parametrize('func', [random.uniform, jnp.ones])
+  @pytest.mark.parametrize('sign1', [-1, 1])
+  @pytest.mark.parametrize('sign2', [-1, 1])
+  @pytest.mark.parametrize('pt', SAMPLE_POINTS)
+  def test_exp_se3(self, func, sign1, sign2, pt):
+    shape, _ = self._process_parameters(None, 3)
+    w = sign1 * self.get_random_vector(func, shape=shape)
+    v = sign2 * self.get_random_vector(func, shape=shape)
+    theta = jnp.linalg.norm(w)
+    w = w / theta[..., jnp.newaxis]
+    v = v / theta[..., jnp.newaxis]
+
+    screw_axis = jnp.concatenate([w, v], axis=-1)
+    transform = rigid_body.exp_se3(screw_axis, theta)
+
+    # TODO(utsinh): Figure out how this t relates to v and add a test..
+    # t = jnp.squeeze(transform[0:3, 3])
+    r = jnp.squeeze(transform[0:3, 0:3])
+    last_row = jnp.squeeze(transform[3, :])
+
+    # The rotation section of the matrix should be orthonormal.
+    self.assertAlmostEqual(jnp.linalg.det(r).tolist(), 1, places=6)
+    np.testing.assert_allclose(
+        jnp.matmul(jnp.transpose(r), r), jnp.eye(3), atol=1E-5, rtol=1E-5)
+    np.testing.assert_allclose(
+        jnp.matmul(r, jnp.transpose(r)), jnp.eye(3), atol=1E-5, rtol=1E-5)
+
+    # The last row should be [0, 0, 0, 1].
+    np.testing.assert_array_equal(last_row, [0, 0, 0, 1])
+
+    pt = jnp.array(pt)
+    q = quaternion.from_axis_angle(w, theta)
+    pt_transformed = quaternion.rotate(q, pt) + (v * theta)
+
+    pt_transformed_tf = rigid_body.from_homogenous(
+        jnp.matmul(transform, rigid_body.to_homogenous(pt)))
+    self.assertEqual(pt_transformed.shape, (3,))
+    self.assertEqual(pt_transformed_tf.shape, (3,))
+    # TODO(utsinh): Make this work - this is the key assert. There's some
+    #  discrepency between how to represent the screw-axis transform as a
+    #  quaternion to rotate and a translation.
+    # np.testing.assert_allclose(pt_transformed_tf, pt_transformed)
+
+  def test_se3_back_and_forth_conversion(self):
+    key = jax.random.PRNGKey(1)
+    rotvec = jax.random.uniform(key, (5, 6))
+    rotmat, trans = jax.vmap(rigid_body.se3_to_rotation_translation)(rotvec)
+    self.assertEqual(rotmat.shape[-2:], (3, 3))
+    rotvec_rec = jax.vmap(rigid_body.rotation_translation_to_se3)(rotmat, trans)
+    self.assertEqual(rotvec_rec.shape, rotvec.shape)
+