Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

third_party/Dyn-HaMR/dyn-hamr/HMP/nemf/losses.py

@@ -0,0 +1,88 @@
+import torch
+import torch.nn as nn
+
+
+class GeodesicLoss(nn.Module):
+    def __init__(self):
+        super(GeodesicLoss, self).__init__()
+
+    def compute_geodesic_distance(self, m1, m2):
+        """ Compute the geodesic distance between two rotation matrices.
+
+        Args:
+            m1, m2: Two rotation matrices with the shape (batch x 3 x 3).
+
+        Returns:
+            The minimal angular difference between two rotation matrices in radian form [0, pi].
+        """
+        batch = m1.shape[0]
+        m = torch.bmm(m1, m2.transpose(1, 2))  # batch*3*3
+
+        cos = (m[:, 0, 0] + m[:, 1, 1] + m[:, 2, 2] - 1) / 2
+        # cos = torch.min(cos, torch.autograd.Variable(torch.ones(batch).cuda()))
+        # cos = torch.max(cos, torch.autograd.Variable(torch.ones(batch).cuda()) * -1)
+        eps = 1e-7
+        cos = torch.clamp(cos, -1 + eps, 1 - eps)
+
+        theta = torch.acos(cos)
+
+        return theta
+    
+    def rot_smooth_loss(self, rotmat):
+    
+        diff = self.compute_geodesic_distance(rotmat[:, 1:], rotmat[:, :-1])
+        vel = diff * 30
+        loss = vel ** 2
+        loss = 0.5 * torch.mean(loss)
+        return loss
+
+    def __call__(self, m1, m2, reduction='mean'):
+        loss = self.compute_geodesic_distance(m1, m2)
+
+        if reduction == 'mean':
+            return loss.mean()
+        elif reduction == 'none':
+            return loss
+        else:
+            raise RuntimeError(f'unsupported reduction: {reduction}')
+
+
+def pos_smooth_loss(pos):
+    # minimize delta steps
+    diff = (pos[:, 1:] - pos[:, :-1])
+    vel = diff * 30
+    loss = vel ** 2
+    loss = 0.5 * torch.mean(loss)
+    return loss
+
+def compute_geodesic_distance(m1, m2):
+    """ Compute the geodesic distance between two rotation matrices.
+
+    Args:
+        m1, m2: Two rotation matrices with the shape (batch x 3 x 3).
+
+    Returns:
+        The minimal angular difference between two rotation matrices in radian form [0, pi].
+    """
+    batch = m1.shape[:-2]
+    m = torch.matmul(m1, m2.transpose(-2, -1))  # batch*3*3
+
+    cos = (m[..., 0, 0] + m[..., 1, 1] + m[..., 2, 2] - 1) / 2
+    cos = torch.min(cos, torch.ones(*batch).cuda())
+    cos = torch.max(cos, torch.ones(*batch).cuda() * -1)
+    
+    eps = 1e-7
+    cos = torch.clamp(cos, -1 + eps, 1 - eps)
+        
+    theta = torch.acos(cos)
+    return theta
+
+
+def rot_smooth_loss(rotmat):
+    
+    diff = compute_geodesic_distance(rotmat[:, 1:], rotmat[:, :-1])
+    vel = diff * 30
+    loss = vel ** 2
+    loss = 0.5 * torch.mean(loss)
+    return loss
+