SO3 improvements for stable gradients.

Summary:
Improves so3 functions to make gradient computation stable:
- Instead of `torch.acos`, uses `acos_linear_extrapolation` which has finite gradients of reasonable magnitude for all inputs.
- Adds tests for the latter.

The tests of the finiteness of the gradient in `test_so3_exp_singularity`, `test_so3_exp_singularity`, `test_so3_cos_bound` would fail if the `so3` functions would call `torch.acos` instead of `acos_linear_extrapolation`.

Reviewed By: bottler

Differential Revision: D23326429

fbshipit-source-id: dc296abf2ae3ddfb3942c8146621491a9cb740ee

Author: davnov134

pytorch3d/transforms/__init__.py

@@ -22,6 +22,7 @@
 )
 from .so3 import (
     so3_exponential_map,
+    so3_exp_map,
     so3_log_map,
     so3_relative_angle,
     so3_rotation_angle,

pytorch3d/transforms/so3.py

@@ -1,13 +1,20 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
 
+from typing import Tuple
 
 import torch
 
+from ..transforms import acos_linear_extrapolation
 
 HAT_INV_SKEW_SYMMETRIC_TOL = 1e-5
 
 
-def so3_relative_angle(R1, R2, cos_angle: bool = False):
+def so3_relative_angle(
+    R1: torch.Tensor,
+    R2: torch.Tensor,
+    cos_angle: bool = False,
+    cos_bound: float = 1e-4,
+) -> torch.Tensor:
     """
     Calculates the relative angle (in radians) between pairs of
     rotation matrices `R1` and `R2` with `angle = acos(0.5 * (Trace(R1 R2^T)-1))`
@@ -20,8 +27,12 @@ def so3_relative_angle(R1, R2, cos_angle: bool = False):
         R1: Batch of rotation matrices of shape `(minibatch, 3, 3)`.
         R2: Batch of rotation matrices of shape `(minibatch, 3, 3)`.
         cos_angle: If==True return cosine of the relative angle rather than
-                   the angle itself. This can avoid the unstable
-                   calculation of `acos`.
+            the angle itself. This can avoid the unstable calculation of `acos`.
+        cos_bound: Clamps the cosine of the relative rotation angle to
+            [-1 + cos_bound, 1 - cos_bound] to avoid non-finite outputs/gradients
+            of the `acos` call. Note that the non-finite outputs/gradients
+            are returned when the angle is requested (i.e. `cos_angle==False`)
+            and the rotation angle is close to 0 or π.
 
     Returns:
         Corresponding rotation angles of shape `(minibatch,)`.
@@ -32,10 +43,15 @@ def so3_relative_angle(R1, R2, cos_angle: bool = False):
         ValueError if `R1` or `R2` has an unexpected trace.
     """
     R12 = torch.bmm(R1, R2.permute(0, 2, 1))
-    return so3_rotation_angle(R12, cos_angle=cos_angle)
+    return so3_rotation_angle(R12, cos_angle=cos_angle, cos_bound=cos_bound)
 
 
-def so3_rotation_angle(R, eps: float = 1e-4, cos_angle: bool = False):
+def so3_rotation_angle(
+    R: torch.Tensor,
+    eps: float = 1e-4,
+    cos_angle: bool = False,
+    cos_bound: float = 1e-4,
+) -> torch.Tensor:
     """
     Calculates angles (in radians) of a batch of rotation matrices `R` with
     `angle = acos(0.5 * (Trace(R)-1))`. The trace of the
@@ -47,8 +63,13 @@ def so3_rotation_angle(R, eps: float = 1e-4, cos_angle: bool = False):
         R: Batch of rotation matrices of shape `(minibatch, 3, 3)`.
         eps: Tolerance for the valid trace check.
         cos_angle: If==True return cosine of the rotation angles rather than
-                   the angle itself. This can avoid the unstable
-                   calculation of `acos`.
+            the angle itself. This can avoid the unstable
+            calculation of `acos`.
+        cos_bound: Clamps the cosine of the rotation angle to
+            [-1 + cos_bound, 1 - cos_bound] to avoid non-finite outputs/gradients
+            of the `acos` call. Note that the non-finite outputs/gradients
+            are returned when the angle is requested (i.e. `cos_angle==False`)
+            and the rotation angle is close to 0 or π.
 
     Returns:
         Corresponding rotation angles of shape `(minibatch,)`.
@@ -68,20 +89,19 @@ def so3_rotation_angle(R, eps: float = 1e-4, cos_angle: bool = False):
     if ((rot_trace < -1.0 - eps) + (rot_trace > 3.0 + eps)).any():
         raise ValueError("A matrix has trace outside valid range [-1-eps,3+eps].")
 
-    # clamp to valid range
-    rot_trace = torch.clamp(rot_trace, -1.0, 3.0)
-
     # phi ... rotation angle
-    phi = 0.5 * (rot_trace - 1.0)
+    phi_cos = (rot_trace - 1.0) * 0.5
 
     if cos_angle:
-        return phi
+        return phi_cos
     else:
-        # pyre-fixme[16]: `float` has no attribute `acos`.
-        return phi.acos()
+        if cos_bound > 0.0:
+            return acos_linear_extrapolation(phi_cos, 1.0 - cos_bound)
+        else:
+            return torch.acos(phi_cos)
 
 
-def so3_exponential_map(log_rot, eps: float = 0.0001):
+def so3_exp_map(log_rot: torch.Tensor, eps: float = 0.0001) -> torch.Tensor:
     """
     Convert a batch of logarithmic representations of rotation matrices `log_rot`
     to a batch of 3x3 rotation matrices using Rodrigues formula [1].
@@ -94,18 +114,31 @@ def so3_exponential_map(log_rot, eps: float = 0.0001):
     which is handled by clamping controlled with the `eps` argument.
 
     Args:
-        log_rot: Batch of vectors of shape `(minibatch , 3)`.
+        log_rot: Batch of vectors of shape `(minibatch, 3)`.
         eps: A float constant handling the conversion singularity.
 
     Returns:
-        Batch of rotation matrices of shape `(minibatch , 3 , 3)`.
+        Batch of rotation matrices of shape `(minibatch, 3, 3)`.
 
     Raises:
         ValueError if `log_rot` is of incorrect shape.
 
     [1] https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula
     """
+    return _so3_exp_map(log_rot, eps=eps)[0]
+
+
+so3_exponential_map = so3_exp_map
+
 
+def _so3_exp_map(
+    log_rot: torch.Tensor, eps: float = 0.0001
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    A helper function that computes the so3 exponential map and,
+    apart from the rotation matrix, also returns intermediate variables
+    that can be re-used in other functions.
+    """
     _, dim = log_rot.shape
     if dim != 3:
         raise ValueError("Input tensor shape has to be Nx3.")
@@ -117,27 +150,35 @@ def so3_exponential_map(log_rot, eps: float = 0.0001):
     fac1 = rot_angles_inv * rot_angles.sin()
     fac2 = rot_angles_inv * rot_angles_inv * (1.0 - rot_angles.cos())
     skews = hat(log_rot)
+    skews_square = torch.bmm(skews, skews)
 
     R = (
         # pyre-fixme[16]: `float` has no attribute `__getitem__`.
         fac1[:, None, None] * skews
-        + fac2[:, None, None] * torch.bmm(skews, skews)
+        + fac2[:, None, None] * skews_square
         + torch.eye(3, dtype=log_rot.dtype, device=log_rot.device)[None]
     )
 
-    return R
+    return R, rot_angles, skews, skews_square
 
 
-def so3_log_map(R, eps: float = 0.0001):
+def so3_log_map(
+    R: torch.Tensor, eps: float = 0.0001, cos_bound: float = 1e-4
+) -> torch.Tensor:
     """
     Convert a batch of 3x3 rotation matrices `R`
     to a batch of 3-dimensional matrix logarithms of rotation matrices
     The conversion has a singularity around `(R=I)` which is handled
-    by clamping controlled with the `eps` argument.
+    by clamping controlled with the `eps` and `cos_bound` arguments.
 
     Args:
         R: batch of rotation matrices of shape `(minibatch, 3, 3)`.
         eps: A float constant handling the conversion singularity.
+        cos_bound: Clamps the cosine of the rotation angle to
+            [-1 + cos_bound, 1 - cos_bound] to avoid non-finite outputs/gradients
+            of the `acos` call when computing `so3_rotation_angle`.
+            Note that the non-finite outputs/gradients are returned when
+            the rotation angle is close to 0 or π.
 
     Returns:
         Batch of logarithms of input rotation matrices
@@ -152,22 +193,26 @@ def so3_log_map(R, eps: float = 0.0001):
     if dim1 != 3 or dim2 != 3:
         raise ValueError("Input has to be a batch of 3x3 Tensors.")
 
-    phi = so3_rotation_angle(R)
+    phi = so3_rotation_angle(R, cos_bound=cos_bound, eps=eps)
 
-    phi_sin = phi.sin()
+    phi_sin = torch.sin(phi)
 
-    phi_denom = (
-        torch.clamp(phi_sin.abs(), eps) * phi_sin.sign()
-        + (phi_sin == 0).type_as(phi) * eps
-    )
+    # We want to avoid a tiny denominator of phi_factor = phi / (2.0 * phi_sin).
+    # Hence, for phi_sin.abs() <= 0.5 * eps, we approximate phi_factor with
+    # 2nd order Taylor expansion: phi_factor = 0.5 + (1.0 / 12) * phi**2
+    phi_factor = torch.empty_like(phi)
+    ok_denom = phi_sin.abs() > (0.5 * eps)
+    phi_factor[~ok_denom] = 0.5 + (phi[~ok_denom] ** 2) * (1.0 / 12)
+    phi_factor[ok_denom] = phi[ok_denom] / (2.0 * phi_sin[ok_denom])
+
+    log_rot_hat = phi_factor[:, None, None] * (R - R.permute(0, 2, 1))
 
-    log_rot_hat = (phi / (2.0 * phi_denom))[:, None, None] * (R - R.permute(0, 2, 1))
     log_rot = hat_inv(log_rot_hat)
 
     return log_rot
 
 
-def hat_inv(h):
+def hat_inv(h: torch.Tensor) -> torch.Tensor:
     """
     Compute the inverse Hat operator [1] of a batch of 3x3 matrices.
 
@@ -188,9 +233,9 @@ def hat_inv(h):
     if dim1 != 3 or dim2 != 3:
         raise ValueError("Input has to be a batch of 3x3 Tensors.")
 
-    ss_diff = (h + h.permute(0, 2, 1)).abs().max()
+    ss_diff = torch.abs(h + h.permute(0, 2, 1)).max()
     if float(ss_diff) > HAT_INV_SKEW_SYMMETRIC_TOL:
-        raise ValueError("One of input matrices not skew-symmetric.")
+        raise ValueError("One of input matrices is not skew-symmetric.")
 
     x = h[:, 2, 1]
     y = h[:, 0, 2]
@@ -201,7 +246,7 @@ def hat_inv(h):
     return v
 
 
-def hat(v):
+def hat(v: torch.Tensor) -> torch.Tensor:
     """
     Compute the Hat operator [1] of a batch of 3D vectors.
 
@@ -225,7 +270,7 @@ def hat(v):
     if dim != 3:
         raise ValueError("Input vectors have to be 3-dimensional.")
 
-    h = v.new_zeros(N, 3, 3)
+    h = torch.zeros((N, 3, 3), dtype=v.dtype, device=v.device)
 
     x, y, z = v.unbind(1)