Merge pull request #45 from fangnan99/fix/prismatic

fix mistake in prismatic joint calculation

Author: LemonPi

src/pytorch_kinematics/frame.py

@@ -133,7 +133,8 @@ def get_transform(self, theta):
             rot = axis_and_angle_to_matrix_33(self.joint.axis, theta)
             t = tf.Transform3d(rot=rot, dtype=dtype, device=d)
         elif self.joint.joint_type == 'prismatic':
-            t = tf.Transform3d(pos=theta * self.joint.axis, dtype=dtype, device=d)
+            pos = theta.unsqueeze(1) * self.joint.axis
+            t = tf.Transform3d(pos=pos, dtype=dtype, device=d)
         elif self.joint.joint_type == 'fixed':
             t = tf.Transform3d(default_batch_size=theta.shape[0], dtype=dtype, device=d)
         else:

src/pytorch_kinematics/jacobian.py

@@ -46,7 +46,7 @@ def calc_jacobian(serial_chain, th, tool=None, ret_eef_pose=False):
             j_eef[:, :, -cnt] = torch.cat((position_jacobian, axis_in_eef), dim=-1)
         elif f.joint.joint_type == "prismatic":
             cnt += 1
-            j_eef[:, :3, -cnt] = f.joint.axis.repeat(N, 1) @ cur_transform[:, :3, :3]
+            j_eef[:, :3, -cnt] = (f.joint.axis.repeat(N, 1, 1) @ cur_transform[:, :3, :3])[:, 0, :]
         cur_frame_transform = f.get_transform(th[:, -cnt]).get_matrix()
         cur_transform = cur_frame_transform @ cur_transform