all tests pass!

Author: Peter

src/pytorch_kinematics/chain.py

@@ -38,6 +38,16 @@ def ensure_2d_tensor(th, dtype, device):
     return th, N
 
 
+def get_dict_elem_shape(th_dict):
+    elem = th_dict[list(th_dict.keys())[0]]
+    if isinstance(elem, np.ndarray):
+        return elem.shape
+    elif isinstance(elem, torch.Tensor):
+        return elem.shape
+    else:
+        return ()
+
+
 class Chain:
     """
     Robot model that may be constructed from different descriptions via their respective parsers.
@@ -174,6 +184,7 @@ def get_joints(self, exclude_fixed=True):
         joints = self._get_joints(self._root, exclude_fixed=exclude_fixed)
         return joints
 
+    @lru_cache()
     def get_joint_parameter_names(self, exclude_fixed=True):
         names = []
         for f in self.get_joints(exclude_fixed=exclude_fixed):
@@ -270,11 +281,16 @@ def get_frame_indices(self, *frame_names):
 
     def forward_kinematics(self, th, frame_indices: Optional = None):
         """
-        Instead of a tree, we can use a flat data structure with indexes to represent the parent
-        then instead of recursion we can just iterate in order and use parent pointers. This
-        reduces function call overhead and moves some of the indexing work to the constructor.
+        Compute forward kinematics.
+
+        Args:
+            th: A dict, list, numpy array, or torch tensor of ALL joints values. Possibly batched.
+               the fastest thing to use is a torch tensor, all other types get converted to that.
+               If any joint values are missing, an exception will be thrown.
+            frame_indices: A list of frame indices to compute forward kinematics for. If None, all frames are computed.
 
-        use `get_frame_indices` to get the indices for the frames you want to compute the pose for.
+        Returns:
+            A dict of frame names and their corresponding Transform3d objects.
         """
         if frame_indices is None:
             frame_indices = self.get_all_frame_indices()
@@ -343,18 +359,28 @@ def forward_kinematics(self, th, frame_indices: Optional = None):
         return frame_names_and_transform3ds
 
     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().
+        """
         if isinstance(th, np.ndarray):
             th = torch.tensor(th, device=self.device, dtype=self.dtype)
-        if isinstance(th, list):
+        elif isinstance(th, list):
             th = torch.tensor(th, device=self.device, dtype=self.dtype)
-        if isinstance(th, dict):
+        elif isinstance(th, dict):
             # convert dict to a flat, complete, tensor of all joints values. Missing joints are filled with zeros.
             th_dict = th
-            th = torch.zeros(self.n_joints, device=self.device, dtype=self.dtype)
+            elem_shape = get_dict_elem_shape(th_dict)
+            th = torch.ones([*elem_shape, self.n_joints], device=self.device, dtype=self.dtype) * torch.nan
             joint_names = self.get_joint_parameter_names()
             for joint_name, joint_position in th_dict.items():
                 jnt_idx = joint_names.index(joint_name)
-                th[jnt_idx] = joint_position
+                th[..., jnt_idx] = joint_position
+            if torch.any(torch.isnan(th)):
+                msg = "Missing values for the following joints:\n"
+                for joint_name, th_i in zip(self.get_joint_parameter_names(), th):
+                    msg += joint_name + "\n"
+                raise ValueError(msg)
         return th
 
     def get_all_frame_indices(self):
@@ -454,6 +480,7 @@ def jacobian(self, th, locations=None):
         return jacobian.calc_jacobian(self, th, tool=locations)
 
     def forward_kinematics(self, th, end_only: bool = True):
+        """ Like the base class, except `th` only needs to contain the joints in the SerialChain, not all joints. """
         if end_only:
             frame_indices = self.get_frame_indices(self._serial_frames[-1].name)
         else:

src/pytorch_kinematics/sdf.py

@@ -100,3 +100,9 @@ def build_chain_from_sdf(data):
                                  _convert_visuals(root_link.visuals))
     root_frame.children = _build_chain_recurse(root_frame, lmap, joints)
     return chain.Chain(root_frame)
+
+
+def build_serial_chain_from_sdf(data, end_link_name, root_link_name=""):
+    mjcf_chain = build_chain_from_sdf(data)
+    serial_chain = chain.SerialChain(mjcf_chain, end_link_name, "" if root_link_name == "" else root_link_name)
+    return serial_chain

tests/test_kinematics.py

@@ -36,7 +36,16 @@ def test_fk_mjcf():
     chain = chain.to(dtype=torch.float64)
     print(chain)
     print(chain.get_joint_parameter_names())
-    th = {'hip_1': 1.0, 'ankle_1': 1}
+    th = {
+        'hip_1':   1.0,
+        'ankle_1': 1,
+        'hip_2':   0.0,
+        'ankle_2': 0.0,
+        'hip_3':   0.0,
+        'ankle_3': 0.0,
+        'hip_4':   0.0,
+        'ankle_4': 0.0,
+    }
     ret = chain.forward_kinematics(th)
     tg = ret['aux_1']
     pos, rot = quat_pos_from_transform3d(tg)
@@ -151,14 +160,24 @@ def test_fk_simple_arm():
     chain = chain.to(dtype=torch.float64)
     # print(chain)
     # print(chain.get_joint_parameter_names())
-    ret = chain.forward_kinematics({'arm_elbow_pan_joint': math.pi / 2.0, 'arm_wrist_lift_joint': -0.5})
+    ret = chain.forward_kinematics({
+        'arm_shoulder_pan_joint': 0.,
+        'arm_elbow_pan_joint':    math.pi / 2.0,
+        'arm_wrist_lift_joint':   -0.5,
+        'arm_wrist_roll_joint':   0.,
+    })
     tg = ret['arm_wrist_roll']
     pos, rot = quat_pos_from_transform3d(tg)
     assert quaternion_equality(rot, torch.tensor([0.70710678, 0., 0., 0.70710678], dtype=torch.float64))
     assert torch.allclose(pos, torch.tensor([1.05, 0.55, 0.5], dtype=torch.float64))
 
     N = 100
-    ret = chain.forward_kinematics({'arm_elbow_pan_joint': torch.rand(N, 1), 'arm_wrist_lift_joint': torch.rand(N, 1)})
+    ret = chain.forward_kinematics({
+        'arm_shoulder_pan_joint': torch.rand(N),
+        'arm_elbow_pan_joint':    torch.rand(N),
+        'arm_wrist_lift_joint':   torch.rand(N),
+        'arm_wrist_roll_joint':   torch.rand(N),
+    })
     tg = ret['arm_wrist_roll']
     assert list(tg.get_matrix().shape) == [N, 4, 4]
 
@@ -176,7 +195,12 @@ def test_cuda():
         chain = pk.build_chain_from_sdf(open(os.path.join(TEST_DIR, "simple_arm.sdf")).read())
         chain = chain.to(dtype=dtype, device=d)
 
-        ret = chain.forward_kinematics({'arm_elbow_pan_joint': math.pi / 2.0, 'arm_wrist_lift_joint': -0.5})
+        ret = chain.forward_kinematics({
+            'arm_shoulder_pan_joint': 0,
+            'arm_elbow_pan_joint': math.pi / 2.0,
+            'arm_wrist_lift_joint': -0.5,
+            'arm_wrist_roll_joint': 0,
+        })
         tg = ret['arm_wrist_roll']
         pos, rot = quat_pos_from_transform3d(tg)
         assert quaternion_equality(rot, torch.tensor([0.70710678, 0., 0., 0.70710678], dtype=dtype, device=d))