fix handling of batched inputs for serial chains and add new test

Author: Peter

README.md

@@ -22,7 +22,7 @@ See `tests` for code samples; some are also shown here.
 If you use this package in your research, consider citing
 ```
 @software{Zhong_PyTorch_Kinematics_2023,
-author = {Zhong, Sheng and Power, Thomas and Gupta, Ashwin},
+author = {Zhong, Sheng and Power, Thomas and Gupta, Ashwin and Peter, Mitrano},
 doi = {10.5281/zenodo.7700588},
 month = {3},
 title = {{PyTorch Kinematics}},

src/pytorch_kinematics/chain.py

@@ -1,6 +1,4 @@
 from functools import lru_cache
-from pytorch_kinematics.transforms.rotation_conversions import axis_and_angle_to_matrix
-from pytorch_kinematics.transforms.rotation_conversions import axis_and_d_to_pris_matrix
 from typing import Optional, Sequence
 
 import numpy as np
@@ -9,9 +7,10 @@
 import pytorch_kinematics.transforms as tf
 from pytorch_kinematics import jacobian
 from pytorch_kinematics.frame import Frame, Link, Joint
+from pytorch_kinematics.transforms.rotation_conversions import axis_and_angle_to_matrix_44, axis_and_d_to_pris_matrix
 
 
-def get_th_size(th):
+def get_n_joints(th):
     """
 
     Args:
@@ -28,6 +27,19 @@ def get_th_size(th):
         raise NotImplementedError(f"Unsupported type {type(th)}")
 
 
+def get_batch_size(th):
+    if isinstance(th, torch.Tensor) or isinstance(th, np.ndarray):
+        return th.shape[0]
+    elif isinstance(th, dict):
+        elem_shape = get_dict_elem_shape(th)
+        return elem_shape[0]
+    elif isinstance(th, list):
+        # Lists cannot be batched. We don't allow lists of lists.
+        return 1
+    else:
+        raise NotImplementedError(f"Unsupported type {type(th)}")
+
+
 def ensure_2d_tensor(th, dtype, device):
     if not torch.is_tensor(th):
         th = torch.tensor(th, dtype=dtype, device=device)
@@ -282,6 +294,18 @@ def get_frame_indices(self, *frame_names):
         return torch.tensor([self.frame_to_idx[n] for n in frame_names], dtype=torch.long, device=self.device)
 
     def forward_kinematics(self, th, frame_indices: Optional = None):
+        """
+        Compute forward kinematics for the given joint values.
+
+        Args:
+            th: A dict, list, numpy array, or torch tensor of joints values. Possibly batched.
+            frame_indices: A list of frame indices to compute transforms for. If None, all frames are computed.
+                Use `get_frame_indices` to convert from frame names to frame indices.
+
+        Returns:
+            A dict of Transform3d objects for each frame.
+
+        """
         if frame_indices is None:
             frame_indices = self.get_all_frame_indices()
 
@@ -294,7 +318,7 @@ def forward_kinematics(self, th, frame_indices: Optional = None):
         # compute all joint transforms at once first
         # in order to handle multiple joint types without branching, we create all possible transforms
         # for all joint types and then select the appropriate one for each joint.
-        rev_jnt_transform = axis_and_angle_to_matrix(axes_expanded, th)
+        rev_jnt_transform = axis_and_angle_to_matrix_44(axes_expanded, th)
         pris_jnt_transform = axis_and_d_to_pris_matrix(axes_expanded, th)
 
         frame_transforms = {}
@@ -336,7 +360,7 @@ def forward_kinematics(self, th, frame_indices: Optional = None):
     def ensure_tensor(self, th):
         """
         Converts a number of possible types into a tensor. The order of the tensor is determined by the order
-        of self.get_joint_parameter_names().
+        of self.get_joint_parameter_names(). th must contain all joints in the entire chain.
         """
         if isinstance(th, np.ndarray):
             th = torch.tensor(th, device=self.device, dtype=self.dtype)
@@ -362,32 +386,17 @@ def get_all_frame_indices(self):
         frame_indices = self.get_frame_indices(*self.get_frame_names(exclude_fixed=False))
         return frame_indices
 
-    def ensure_dict_of_2d_tensors(self, th):
-        if not isinstance(th, dict):
-            th, _ = ensure_2d_tensor(th, self.dtype, self.device)
-            assert self.n_joints == th.shape[-1]
-            th_dict = dict((j, th[..., i]) for i, j in enumerate(self.get_joint_parameter_names()))
-        else:
-            th_dict = {k: ensure_2d_tensor(v, self.dtype, self.device)[0] for k, v in th.items()}
-        return th_dict
-
     def clamp(self, th):
-        th_dict = self.ensure_dict_of_2d_tensors(th)
+        """
 
-        out_th_dict = {}
-        for joint_name, joint_position in th_dict.items():
-            joint = self.find_joint(joint_name)
-            joint_position_clamped = joint.clamp(joint_position)
-            out_th_dict[joint_name] = joint_position_clamped
+        Args:
+            th: Joint configuration
 
-        return self.match_input_type(out_th_dict, th)
+        Returns: Always a tensor in the order of self.get_joint_parameter_names(), possibly batched.
 
-    @staticmethod
-    def match_input_type(th_dict, th):
-        if isinstance(th, dict):
-            return th_dict
-        else:
-            return torch.stack([v for v in th_dict.values()], dim=-1)
+        """
+        th = self.ensure_tensor(th)
+        return torch.clamp(th, self.low, self.high)
 
     def get_joint_limits(self):
         low = []
@@ -466,22 +475,33 @@ def forward_kinematics(self, th, end_only: bool = True):
             return mat
 
     def convert_serial_inputs_to_chain_inputs(self, th, end_only: bool):
+        # th = self.ensure_tensor(th)
+        th_b = get_batch_size(th)
+        th_n_joints = get_n_joints(th)
+        if isinstance(th, list):
+            th = torch.tensor(th, device=self.device, dtype=self.dtype)
+
         if end_only:
             frame_indices = self.get_frame_indices(self._serial_frames[-1].name)
         else:
             # pass through default behavior for frame indices being None, which is currently
             # to return all frames.
             frame_indices = None
-        if get_th_size(th) < self.n_joints:
+        if th_n_joints < self.n_joints:
             # if th is only a partial list of joints, assume it's a list of joints for only the serial chain.
             partial_th = th
             nonfixed_serial_frames = list(filter(lambda f: f.joint.joint_type != 'fixed', self._serial_frames))
-            if len(nonfixed_serial_frames) != len(partial_th):
-                raise ValueError(f'Expected {len(nonfixed_serial_frames)} joint values, got {len(partial_th)}.')
-            th = torch.zeros(self.n_joints, device=self.device, dtype=self.dtype)
-            for frame, partial_th_i in zip(nonfixed_serial_frames, partial_th):
+            if th_n_joints != len(nonfixed_serial_frames):
+                raise ValueError(f'Expected {len(nonfixed_serial_frames)} joint values, got {th_n_joints}.')
+            th = torch.zeros([th_b, self.n_joints], device=self.device, dtype=self.dtype)
+            for i, frame in enumerate(nonfixed_serial_frames):
+                joint_name = frame.joint.name
+                if isinstance(partial_th, dict):
+                    partial_th_i = partial_th[joint_name]
+                else:
+                    partial_th_i = partial_th[..., i]
                 k = self.frame_to_idx[frame.name]
                 jnt_idx = self.joint_indices[k]
                 if frame.joint.joint_type != 'fixed':
-                    th[jnt_idx] = partial_th_i
+                    th[..., jnt_idx] = partial_th_i
         return frame_indices, th

src/pytorch_kinematics/frame.py

@@ -1,7 +1,7 @@
 import torch
 
 import pytorch_kinematics.transforms as tf
-from pytorch_kinematics.transforms import axis_and_angle_to_matrix
+from pytorch_kinematics.transforms import axis_and_angle_to_matrix_33
 
 
 class Visual(object):
@@ -115,7 +115,7 @@ def get_transform(self, theta):
         dtype = self.joint.axis.dtype
         d = self.joint.axis.device
         if self.joint.joint_type == 'revolute':
-            rot = axis_and_angle_to_matrix(self.joint.axis, 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)

src/pytorch_kinematics/jacobian.py

@@ -47,7 +47,7 @@ def calc_jacobian(serial_chain, th, tool=None):
         elif f.joint.joint_type == "prismatic":
             cnt += 1
             j_eef[:, :3, -cnt] = f.joint.axis.repeat(N, 1) @ cur_transform[:, :3, :3]
-        cur_frame_transform = f.get_transform(th[:, -cnt].reshape(N, 1)).get_matrix()
+        cur_frame_transform = f.get_transform(th[:, -cnt]).get_matrix()
         cur_transform = cur_frame_transform @ cur_transform
 
     # currently j_eef is Jacobian in end-effector frame, convert to base/world frame

src/pytorch_kinematics/transforms/__init__.py

@@ -19,7 +19,7 @@
     random_rotations,
     rotation_6d_to_matrix,
     standardize_quaternion,
-    axis_and_angle_to_matrix,
+    axis_and_angle_to_matrix_33,
     axis_and_d_to_pris_matrix,
     wxyz_to_xyzw,
     xyzw_to_wxyz,

src/pytorch_kinematics/transforms/rotation_conversions.py

@@ -472,7 +472,7 @@ def axis_and_d_to_pris_matrix(axis, d):
     return mat44
 
 
-def axis_and_angle_to_matrix(axis, theta):
+def axis_and_angle_to_matrix_44(axis, theta):
     """
     Creates a 4x4 matrix that represents a rotation around an axis by an angle theta.
     Works with any number of batch dimensions.
@@ -483,6 +483,25 @@ def axis_and_angle_to_matrix(axis, theta):
 
     Returns: [..., 4, 4]
 
+    """
+    rot = axis_and_angle_to_matrix_33(axis, theta)
+    batch_shape = axis.shape[:-1]
+    mat44 = torch.cat((rot, torch.zeros(*batch_shape, 3, 1).to(axis)), -1)
+    mat44 = torch.cat((mat44, torch.tensor([0.0, 0.0, 0.0, 1.0]).expand(*batch_shape, 1, 4).to(axis)), -2)
+    return mat44
+
+
+def axis_and_angle_to_matrix_33(axis, theta):
+    """
+    Creates a 3x3 matrix that represents a rotation around an axis by an angle theta.
+    Works with any number of batch dimensions.
+
+    Argsaxis.sh:
+        axis: [..., 3]
+        theta: [ ...]
+
+    Returns: [..., 3, 3]
+
     """
     # based on https://ai.stackexchange.com/questions/14041/, and checked against wikipedia
     c = torch.cos(theta)  # NOTE: cos is not that precise for float32, you may want to use float64
@@ -501,10 +520,7 @@ def axis_and_angle_to_matrix(axis, theta):
     rot = torch.stack([torch.stack([r00, r01, r02], -1),
                        torch.stack([r10, r11, r12], -1),
                        torch.stack([r20, r21, r22], -1)], -2)
-    batch_shape = axis.shape[:-1]
-    mat44 = torch.cat((rot, torch.zeros(*batch_shape, 3, 1).to(axis)), -1)
-    mat44 = torch.cat((mat44, torch.tensor([0.0, 0.0, 0.0, 1.0]).expand(*batch_shape, 1, 4).to(axis)), -2)
-    return mat44
+    return rot
 
 
 def axis_angle_to_matrix(axis_angle):

tests/test_kinematics.py

@@ -1,6 +1,6 @@
 import math
-from timeit import timeit
 import os
+from timeit import timeit
 
 import torch
 
@@ -260,7 +260,35 @@ def test_fk_val():
     assert torch.allclose(pos, torch.tensor([-0.225692, 0.259045, 0.262139], dtype=torch.float64))
 
 
+def test_fk_partial_batched_dict():
+    # Test that you can pass in dict of batched joint configs for a subset of the joints
+    chain = pk.build_serial_chain_from_mjcf(open(os.path.join(TEST_DIR, "val.xml")).read(), 'left_tool')
+    th = {
+        'joint56': torch.zeros([1000], dtype=torch.float64),
+        'joint57': torch.zeros([1000], dtype=torch.float64),
+        'joint41': torch.zeros([1000], dtype=torch.float64),
+        'joint42': torch.zeros([1000], dtype=torch.float64),
+        'joint43': torch.zeros([1000], dtype=torch.float64),
+        'joint44': torch.zeros([1000], dtype=torch.float64),
+        'joint45': torch.zeros([1000], dtype=torch.float64),
+        'joint46': torch.zeros([1000], dtype=torch.float64),
+        'joint47': torch.zeros([1000], dtype=torch.float64),
+    }
+    chain = chain.to(dtype=torch.float64)
+    tg = chain.forward_kinematics(th)
+
+
+def test_fk_partial_batched():
+    # Test that you can pass in dict of batched joint configs for a subset of the joints
+    chain = pk.build_serial_chain_from_mjcf(open(os.path.join(TEST_DIR, "val.xml")).read(), 'left_tool')
+    th = torch.zeros([1000, 9], dtype=torch.float64)
+    chain = chain.to(dtype=torch.float64)
+    tg = chain.forward_kinematics(th)
+
+
 if __name__ == "__main__":
+    test_fk_partial_batched()
+    test_fk_partial_batched_dict()
     test_fk_val()
     test_sdf_serial_chain()
     test_urdf_serial()

tests/test_rotation_conversions.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from pytorch_kinematics.transforms.rotation_conversions import axis_and_angle_to_matrix, axis_angle_to_matrix, \
+from pytorch_kinematics.transforms.rotation_conversions import axis_and_angle_to_matrix_33, axis_angle_to_matrix, \
     pos_rot_to_matrix, matrix_to_pos_rot, random_rotations
 
 
@@ -18,7 +18,7 @@ def test_axis_angle_to_matrix_perf():
     dt1 = timeit.timeit(lambda: axis_angle_to_matrix(axis_angle), number=number)
     print(f'Old method: {dt1:.5f}')
 
-    dt2 = timeit.timeit(lambda: axis_and_angle_to_matrix(axis=axis_1d, theta=theta), number=number)
+    dt2 = timeit.timeit(lambda: axis_and_angle_to_matrix_33(axis=axis_1d, theta=theta), number=number)
     print(f'New method: {dt2:.5f}')