Sample axis-angle and translation perturbations around transforms

Author: LemonPi

src/pytorch_kinematics/transforms/__init__.py

@@ -29,5 +29,6 @@
     so3_rotation_angle,
 )
 from .transform3d import Rotate, RotateAxisAngle, Scale, Transform3d, Translate
+from pytorch_kinematics.transforms.perturbation import sample_perturbations
 
 __all__ = [k for k in globals().keys() if not k.startswith("_")]

src/pytorch_kinematics/transforms/perturbation.py

@@ -0,0 +1,28 @@
+import torch
+from pytorch_kinematics.transforms.rotation_conversions import axis_angle_to_matrix
+
+
+def sample_perturbations(T, num_perturbations, radian_sigma, translation_sigma):
+    """
+    Sample perturbations around the given transform. The translation and rotation are sampled independently from
+    0 mean gaussians. The angular perturbations' directions are uniformly sampled from the unit sphere while its
+    magnitude is sampled from a gaussian.
+    :param T: given transform to perturb around
+    :param num_perturbations: number of perturbations to sample
+    :param radian_sigma: standard deviation of the gaussian angular perturbation in radians
+    :param translation_sigma: standard deviation of the gaussian translation perturbation in meters / T units
+    :return: perturbed transforms; may not include the original transform
+    """
+    dtype = T.dtype
+    device = T.device
+    perturbed = torch.eye(4, dtype=dtype, device=device).repeat(num_perturbations, 1, 1)
+
+    delta_R = torch.randn((num_perturbations, 3), dtype=dtype, device=device) * radian_sigma
+    delta_R = axis_angle_to_matrix(delta_R)
+    perturbed[:, :3, :3] = delta_R @ T[..., :3, :3]
+    perturbed[:, :3, 3] = T[..., :3, 3]
+
+    delta_t = torch.randn((num_perturbations, 3), dtype=dtype, device=device) * translation_sigma
+    perturbed[:, :3, 3] += delta_t
+
+    return perturbed

src/pytorch_kinematics/transforms/transform3d.py

@@ -9,6 +9,7 @@
 
 from .rotation_conversions import _axis_angle_rotation, matrix_to_quaternion, quaternion_to_matrix, \
     euler_angles_to_matrix
+from pytorch_kinematics.transforms.perturbation import sample_perturbations
 
 DEFAULT_EULER_CONVENTION = "XYZ"
 
@@ -421,6 +422,14 @@ def rotate(self, *args, **kwargs):
     def rotate_axis_angle(self, *args, **kwargs):
         return self.compose(RotateAxisAngle(device=self.device, *args, **kwargs))
 
+    def sample_perturbations(self, num_perturbations, radian_sigma, translation_sigma):
+        mat = self.get_matrix()
+        if mat.shape[0] == 1:
+            mat = mat[0]
+        all_mats = sample_perturbations(mat, num_perturbations, radian_sigma, translation_sigma)
+        out = Transform3d(matrix=all_mats)
+        return out
+
     def clone(self):
         """
         Deep copy of Transforms object. All internal tensors are cloned