Introduce batch fk capabilities for floating and planar joints

urchin/urdf.py

@@ -2388,37 +2388,23 @@ def get_child_pose(self, cfg=None):
         elif self.joint_type == "fixed":
             return self.origin
         elif self.joint_type in ["revolute", "continuous"]:
-            if cfg is None:
-                cfg = 0.0
-            else:
-                cfg = float(cfg)
+            cfg = float(cfg)
             R = trimesh.transformations.rotation_matrix(cfg, self.axis)
             return self.origin.dot(R)
         elif self.joint_type == "prismatic":
-            if cfg is None:
-                cfg = 0.0
-            else:
-                cfg = float(cfg)
+            cfg = float(cfg)
             translation = np.eye(4, dtype=np.float64)
             translation[:3, 3] = self.axis * cfg
             return self.origin.dot(translation)
         elif self.joint_type == "planar":
-            if cfg is None:
-                cfg = np.zeros(2, dtype=np.float64)
-            else:
-                cfg = np.asanyarray(cfg, dtype=np.float64)
+            cfg = np.asanyarray(cfg, dtype=np.float64)
             if cfg.shape != (2,):
                 raise ValueError("(2,) float configuration required for planar joints")
             translation = np.eye(4, dtype=np.float64)
             translation[:3, 3] = self.origin[:3, :2].dot(cfg)
             return self.origin.dot(translation)
         elif self.joint_type == "floating":
-            if cfg is None:
-                cfg = np.zeros(6, dtype=np.float64)
-            else:
-                cfg = configure_origin(cfg)
-            if cfg is None:
-                raise ValueError("Invalid configuration for floating joint")
+            cfg = configure_origin(cfg)
             return self.origin.dot(cfg)
         else:
             raise ValueError("Invalid configuration")
@@ -2437,8 +2423,10 @@ def get_child_poses(self, cfg, n_cfgs):
             - ``prismatic`` - a translation along the axis in meters.
             - ``revolute`` - a rotation about the axis in radians.
             - ``continuous`` - a rotation about the axis in radians.
-            - ``planar`` - Not implemented.
-            - ``floating`` - Not implemented.
+            - ``planar`` - the x and y translation values in the plane.
+              as an (n,2) matrix.
+            - ``floating`` - the xyz values followed by the rpy values,
+              as (n,6) or a (n,4,4) matrix.
 
             If ``cfg`` is ``None``, then this just returns the joint pose.
 
@@ -2452,19 +2440,21 @@ def get_child_poses(self, cfg, n_cfgs):
         elif self.joint_type == "fixed":
             return np.tile(self.origin, (n_cfgs, 1, 1))
         elif self.joint_type in ["revolute", "continuous"]:
-            if cfg is None:
-                cfg = np.zeros(n_cfgs)
             return np.matmul(self.origin, self._rotation_matrices(cfg, self.axis))
         elif self.joint_type == "prismatic":
-            if cfg is None:
-                cfg = np.zeros(n_cfgs)
             translation = np.tile(np.eye(4), (n_cfgs, 1, 1))
             translation[:, :3, 3] = self.axis * cfg[:, np.newaxis]
             return np.matmul(self.origin, translation)
         elif self.joint_type == "planar":
-            raise NotImplementedError()
+            cfg = np.asanyarray(cfg, dtype=np.float64)
+            if cfg.shape[-1] != 2:
+                raise ValueError("(...,2) float configuration required for planar joints")
+            translation = np.tile(np.eye(4), (n_cfgs, 1, 1))
+            translation[:, :3, 3] = np.matmul(self.origin[np.newaxis, :3, :2], cfg[..., np.newaxis]).squeeze()
+            return np.matmul(self.origin, translation)
         elif self.joint_type == "floating":
-            raise NotImplementedError()
+            cfg = configure_origin(cfg)
+            return np.matmul(self.origin, cfg)
         else:
             raise ValueError("Invalid configuration")
 
@@ -4029,12 +4019,27 @@ def _process_cfg(self, cfg):
                     joint_cfg[joint] = cfg[joint]
         elif isinstance(cfg, (list, tuple, np.ndarray)):
             if len(cfg) != len(self.actuated_joints):
-                raise ValueError(
-                    "Cfg must have same length as actuated joints "
-                    "if specified as a numerical array"
-                )
-            for joint, value in zip(self.actuated_joints, cfg):
-                joint_cfg[joint] = value
+                try:
+                    cfg = np.asanyarray(cfg, dtype=np.float64)
+                except ValueError:
+                    raise ValueError(
+                        "Cfg must have same length as actuated joints or dof "
+                        "if specified as a numerical array"
+                    )
+            if not isinstance(cfg, np.ndarray):
+                for joint, value in zip(self.actuated_joints, cfg):
+                    joint_cfg[joint] = value
+            else:
+                sidx = 0
+                for j in self.actuated_joints:
+                    if j.joint_type == "planar":
+                        eidx = sidx + 2
+                    elif j.joint_type == "floating":
+                        eidx = sidx + 6
+                    else:
+                        eidx = sidx + 1
+                    joint_cfg[j] = cfg[sidx:eidx].squeeze()
+                    sidx = eidx
         else:
             raise TypeError("Invalid type for config")
         return joint_cfg
@@ -4068,9 +4073,25 @@ def _process_cfgs(self, cfgs):
             elif cfgs[0] is None:
                 pass
             else:
-                cfgs = np.asanyarray(cfgs, dtype=np.float64)
-                for i, j in enumerate(self.actuated_joints):
-                    joint_cfg[j] = cfgs[:, i]
+                if isinstance(cfgs, (list, tuple)):
+                    if len(cfgs) == len(self.actuated_joints):
+                        for i, j in enumerate(self.actuated_joints):
+                            joint_cfg[j].append(cfgs[i])
+                    else:
+                        cfgs = np.asanyarray(cfgs, dtype=np.float64)
+                if isinstance(cfgs, np.ndarray):
+                    sidx = 0
+                    for j in self.actuated_joints:
+                        if j.joint_type == "planar":
+                            eidx = sidx + 2
+                        elif j.joint_type == "floating":
+                            eidx = sidx + 6
+                        else:
+                            eidx = sidx + 1
+                        joint_cfg[j] = cfgs[:, sidx:eidx]
+                        if joint_cfg[j].shape[-1] == 1:
+                            joint_cfg[j] = joint_cfg[j].squeeze(-1)
+                        sidx = eidx
         else:
             raise ValueError("Incorrectly formatted config array")
 

urchin/utils.py

@@ -20,23 +20,30 @@ def rpy_to_matrix(coords):
 
     Parameters
     ----------
-    coords : (3,) float
+    coords : (..., 3) float
         The roll-pitch-yaw coordinates in order (x-rot, y-rot, z-rot).
 
     Returns
     -------
-    R : (3,3) float
+    R : (..., 3,3) float
         The corresponding homogenous 3x3 rotation matrix.
     """
     coords = np.asanyarray(coords, dtype=np.float64)
-    c3, c2, c1 = np.cos(coords)
-    s3, s2, s1 = np.sin(coords)
-
-    return np.array([
+    coords_cos = np.cos(coords)
+    coords_sin = np.sin(coords)
+    c1 = coords_cos[..., 2]
+    c2 = coords_cos[..., 1]
+    c3 = coords_cos[..., 0]
+    s1 = coords_sin[..., 2]
+    s2 = coords_sin[..., 1]
+    s3 = coords_sin[..., 0]
+
+    out =  np.array([
         [c1 * c2, (c1 * s2 * s3) - (c3 * s1), (s1 * s3) + (c1 * c3 * s2)],
         [c2 * s1, (c1 * c3) + (s1 * s2 * s3), (c3 * s1 * s2) - (c1 * s3)],
         [-s2, c2 * s3, c2 * c3]
     ], dtype=np.float64)
+    return np.moveaxis(out, (0, 1), (-2, -1))
 
 
 def matrix_to_rpy(R, solution=1):
@@ -113,17 +120,19 @@ def xyz_rpy_to_matrix(xyz_rpy):
 
     Parameters
     ----------
-    xyz_rpy : (6,) float
+    xyz_rpy : (...,6) float
         The xyz_rpy vector.
 
     Returns
     -------
-    matrix : (4,4) float
+    matrix : (...,4,4) float
         The homogenous transform matrix.
     """
-    matrix = np.eye(4, dtype=np.float64)
-    matrix[:3,3] = xyz_rpy[:3]
-    matrix[:3,:3] = rpy_to_matrix(xyz_rpy[3:])
+    xyz_rpy = np.asanyarray(xyz_rpy, dtype=np.float64)
+    matrix = np.zeros(xyz_rpy.shape[:-1] + (4,4), dtype=np.float64)
+    matrix[...,3,3] = 1.0
+    matrix[...,:3,3] = xyz_rpy[...,:3]
+    matrix[...,:3,:3] = rpy_to_matrix(xyz_rpy[...,3:])
     return matrix
 
 
@@ -262,9 +271,9 @@ def configure_origin(value):
         value = np.eye(4, dtype=np.float64)
     elif isinstance(value, (list, tuple, np.ndarray)):
         value = np.asanyarray(value, dtype=np.float64)
-        if value.shape == (6,):
+        if value.shape[-1] == 6:
             value = xyz_rpy_to_matrix(value)
-        elif value.shape != (4,4):
+        elif value.shape[-2:] != (4,4):
             raise ValueError('Origin must be specified as a 4x4 '
                              'homogenous transformation matrix')
     else: