fix mistake in prismatic joint calculation

UM-ARM-Lab · Oct 2, 2024 · 2bbb216 · 2bbb216
1 parent fcadfde
commit 2bbb216
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/src/pytorch_kinematics/frame.py b/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:

diff --git a/src/pytorch_kinematics/jacobian.py b/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