PyTorch Geodesic Loss

A PyTorch criterion for computing the distance between rotation matrices (as described here and here). Specifically, the loss is calculated as:

import torch

from geodesic_loss import GeodesicLoss
from scipy.spatial.transform import Rotation

N = 100
Rs = torch.Tensor(Rotation.random(2 * N).as_matrix())
R_Ss = Rs[:N]
R_Ts = Rs[N:]

criterion = GeodesicLoss(reduction="none")
dists = criterion(R_Ss, R_Ts)