Rotation matrix to axis–angle v.s. Log map from SO(3) to so(3)

1
2
3
4
5
6
7
8
9
10
11
12
from pytorch3d.transforms.rotation_conversions import matrix_to_axis_angle, random_rotations
from pytorch3d.transforms.so3 import so3_log_map
import torch

R = random_rotations(500)
axis_angle = matrix_to_axis_angle(R)
angle_in_2pi = axis_angle.norm(dim=1)
mask = angle_in_2pi > torch.pi
axis_angle[mask] = torch.einsum('km,k->km',
    axis_angle[mask],
    1-2*torch.pi/angle_in_2pi[mask])
so3_axis_angle = so3_log_map(R)
\[\begin{aligned} \theta &= \left\{ \begin{array}{ll} \theta'-2\pi, & \theta' > \pi \\ \theta', & \text{elsewise} \end{array} \right. \\ \Downarrow\\ \boldsymbol{\phi} &= \left\{ \begin{array}{ll} \frac{\boldsymbol{\phi}'}{\theta'}(\theta'-2\pi), & \theta' > \pi \\ \boldsymbol{\phi}', & \text{elsewise} \end{array} \right. \end{aligned}\]

Noted that when the rotation angle is close to pi, error would be raised:

1
2
3
4
5
try:
    assert torch.allclose(axis_angle, so3_axis_angle, atol=1e-3)
except AssertionError:
    idx = torch.where(torch.abs(axis_angle-so3_axis_angle)>1e-3)[0].unique()
    print(angle_in_2pi[idx], so3_axis_angle[idx].norm(dim=1))

There had been discussions on the numerical instabilities of so3_log_map (e.g. https://github.com/facebookresearch/pytorch3d/issues/188).

The function matrix_to_axis_angle is more stable since it is a wrapper of quaternion_to_axis_angle(matrix_to_quaternion(matrix)) utilizing the advantages of quaternions (see here).

NOTE

  • matrix_to_axis_angle: \(\theta \in [0, 2\pi)\), for the same rotation axis, the direction is fixed
  • so3_log_map: \(\theta \in [0, \pi)\), the direction of the rotation axis is used.