Lie algebra for SO(3) and SE(3) operations in PyTorch
| 70 | |
| 71 | |
| 72 | class Lie: |
| 73 | """ |
| 74 | Lie algebra for SO(3) and SE(3) operations in PyTorch |
| 75 | """ |
| 76 | |
| 77 | def so3_to_SO3(self, w): # [...,3] |
| 78 | wx = self.skew_symmetric(w) |
| 79 | theta = w.norm(dim=-1)[..., None, None] |
| 80 | I = torch.eye(3, device=w.device, dtype=torch.float) |
| 81 | A = self.taylor_A(theta) |
| 82 | B = self.taylor_B(theta) |
| 83 | R = I + A * wx + B * wx @ wx |
| 84 | return R |
| 85 | |
| 86 | def SO3_to_so3(self, R, eps=1e-7): # [...,3,3] |
| 87 | trace = R[..., 0, 0] + R[..., 1, 1] + R[..., 2, 2] |
| 88 | theta = ((trace - 1) / 2).clamp(-1 + eps, 1 - eps).acos_()[ |
| 89 | ..., None, None |
| 90 | ] % np.pi # ln(R) will explode if theta==pi |
| 91 | lnR = ( |
| 92 | 1 / (2 * self.taylor_A(theta) + 1e-8) * (R - R.transpose(-2, -1)) |
| 93 | ) # FIXME: wei-chiu finds it weird |
| 94 | w0, w1, w2 = lnR[..., 2, 1], lnR[..., 0, 2], lnR[..., 1, 0] |
| 95 | w = torch.stack([w0, w1, w2], dim=-1) |
| 96 | return w |
| 97 | |
| 98 | def se3_to_SE3(self, wu): # [...,3] |
| 99 | w, u = wu.split([3, 3], dim=-1) |
| 100 | wx = self.skew_symmetric(w) |
| 101 | theta = w.norm(dim=-1)[..., None, None] |
| 102 | I = torch.eye(3, device=w.device, dtype=torch.float) |
| 103 | A = self.taylor_A(theta) |
| 104 | B = self.taylor_B(theta) |
| 105 | C = self.taylor_C(theta) |
| 106 | R = I + A * wx + B * wx @ wx |
| 107 | V = I + B * wx + C * wx @ wx |
| 108 | Rt = torch.cat([R, (V @ u[..., None])], dim=-1) |
| 109 | return Rt |
| 110 | |
| 111 | def SE3_to_se3(self, Rt, eps=1e-8): # [...,3,4] |
| 112 | R, t = Rt.split([3, 1], dim=-1) |
| 113 | w = self.SO3_to_so3(R) |
| 114 | wx = self.skew_symmetric(w) |
| 115 | theta = w.norm(dim=-1)[..., None, None] |
| 116 | I = torch.eye(3, device=w.device, dtype=torch.float) |
| 117 | A = self.taylor_A(theta) |
| 118 | B = self.taylor_B(theta) |
| 119 | invV = I - 0.5 * wx + (1 - A / (2 * B)) / (theta ** 2 + eps) * wx @ wx |
| 120 | u = (invV @ t)[..., 0] |
| 121 | wu = torch.cat([w, u], dim=-1) |
| 122 | return wu |
| 123 | |
| 124 | def skew_symmetric(self, w): |
| 125 | w0, w1, w2 = w.unbind(dim=-1) |
| 126 | O = torch.zeros_like(w0) |
| 127 | wx = torch.stack( |
| 128 | [ |
| 129 | torch.stack([O, -w2, w1], dim=-1), |