add SE3.transform_from and SE3.transform_to

theseus/geometry/se3.py

@@ -284,26 +284,88 @@ def vee(matrix: torch.Tensor) -> torch.Tensor:
         SE3._hat_matrix_check(matrix)
         return torch.cat((matrix[:, :3, 3], SO3.vee(matrix[:, :3, :3])), dim=1)
 
-    def _transform_shape_check(self, point: Union[Point3, torch.Tensor]):
-        raise NotImplementedError
-
     def _copy_impl(self, new_name: Optional[str] = None) -> "SE3":
         return SE3(data=self.data.clone(), name=new_name)
 
     # only added to avoid casting downstream
     def copy(self, new_name: Optional[str] = None) -> "SE3":
         return cast(SE3, super().copy(new_name=new_name))
 
+    def _transform_shape_check(self, point: Union[Point3, torch.Tensor]):
+        err_msg = "SE3 can only transform 3-D vectors."
+        if isinstance(point, torch.Tensor):
+            if not point.ndim == 2 or point.shape[1] != 3:
+                raise ValueError(err_msg)
+        elif point.dof() != 3:
+            raise ValueError(err_msg)
+        if (
+            point.shape[0] != self.shape[0]
+            and point.shape[0] != 1
+            and self.shape[0] != 1
+        ):
+            raise ValueError(
+                "Input point batch size is not broadcastable with group batch size."
+            )
+
     def transform_to(
         self,
         point: Union[Point3, torch.Tensor],
         jacobians: Optional[List[torch.Tensor]] = None,
     ) -> Point3:
-        raise NotImplementedError
+        self._transform_shape_check(point)
+        batch_size = max(self.shape[0], point.shape[0])
+        if isinstance(point, torch.Tensor):
+            p = point.view(-1, 3, 1)
+        else:
+            p = point.data.view(-1, 3, 1)
+
+        ret = Point3(data=(self[:, :, :3] @ p).view(-1, 3))
+        ret[:, 3:] += self[:, :, 3:]
+
+        if jacobians is not None:
+            self._check_jacobians_list(jacobians)
+            # Right jacobians for SE(3) are computed
+            Jg = torch.zeros(batch_size, 3, 6)
+            Jg[:, :, :3] = self[:, :, :3]
+            Jg[:, :, 3:] = -self[:, :, :3] @ SO3.hat(p)
+            # Jacobians for point
+            Jpnt = Jg[:, :, :3]
+
+            jacobians.extend([Jg, Jpnt])
+
+        return ret
 
     def transform_from(
         self,
         point: Union[Point3, torch.Tensor],
         jacobians: Optional[List[torch.Tensor]] = None,
     ) -> Point3:
-        raise NotImplementedError
+        self._transform_shape_check(point)
+        batch_size = max(self.shape[0], point.shape[0])
+        if isinstance(point, torch.Tensor):
+            p = point.view(-1, 3, 1)
+        else:
+            p = point.data.view(-1, 3, 1)
+
+        temp = p - self[:, :, 3:]
+        ret = Point3(data=(self[:, :, :3].transpose(1, 2) @ temp).view(-1, 3))
+
+        if jacobians is not None:
+            self._check_jacobians_list(jacobians)
+            # Right jacobians for SE(3) are computed
+            Jg = torch.zeros(batch_size, 3, 6)
+            Jg[:, 0, 0] = -1
+            Jg[:, 1, 1] = -1
+            Jg[:, 2, 2] = -1
+            Jg[:, 0, 4] = -ret[:, 2]
+            Jg[:, 1, 3] = ret[:, 2]
+            Jg[:, 0, 5] = ret[:, 1]
+            Jg[:, 2, 3] = -ret[:, 1]
+            Jg[:, 1, 5] = -ret[:, 0]
+            Jg[:, 2, 4] = ret[:, 0]
+            # Jacobians for point
+            Jpnt = self[:, :, :3].transpose(1, 2).expand(batch_size, 3, 3)
+
+            jacobians.extend([Jg, Jpnt])
+
+        return ret

theseus/geometry/so3.py

@@ -299,7 +299,7 @@ def rotate(
         ret = Point3(data=(self.data @ p).view(-1, 3))
         if jacobians is not None:
             self._check_jacobians_list(jacobians)
-            # Jacobians for SO3: left-invariant jacobians are computed
+            # Right jacobians for SO(3) are computed
             Jrot = -self.data @ SO3.hat(p)
             # Jacobians for point
             Jpnt = self.to_matrix().expand(batch_size, 3, 3)
@@ -323,7 +323,7 @@ def unrotate(
         ret = Point3(data=(self.data.transpose(1, 2) @ p).view(-1, 3))
         if jacobians is not None:
             self._check_jacobians_list(jacobians)
-            # Jacobians for SO3: left-invariant jacobians are computed
+            # Left jacobians for SO3 are computed
             Jrot = torch.zeros(batch_size, 3, 3, dtype=self.dtype, device=self.device)
             Jrot[:, 0, 1] = -ret[:, 2]
             Jrot[:, 1, 0] = ret[:, 2]