Add batched FK for planar and floating joints

tests/unit/test_joint.py

@@ -0,0 +1,38 @@
+import numpy as np
+import pytest
+
+from urchin import Joint, xyz_rpy_to_matrix
+
+
+@pytest.mark.parametrize("joint_type", ["continuous", "planar", "floating"])
+def test_none_configuration_uses_joint_origin(joint_type):
+    origin = xyz_rpy_to_matrix([1.0, 2.0, 3.0, 0.1, 0.2, 0.3])
+    joint = Joint("joint", joint_type, "parent", "child", origin=origin)
+
+    assert np.allclose(joint.get_child_pose(None), origin)
+    assert np.allclose(joint.get_child_poses(None, 2), np.tile(origin, (2, 1, 1)))
+
+
+def test_planar_get_child_poses_matches_scalar_poses():
+    origin = xyz_rpy_to_matrix([1.0, 2.0, 3.0, 0.1, 0.2, 0.3])
+    joint = Joint("joint", "planar", "parent", "child", origin=origin)
+    cfgs = np.array([[0.5, 1.5], [-1.0, 2.0]])
+
+    expected = np.stack([joint.get_child_pose(cfg) for cfg in cfgs])
+
+    assert np.allclose(joint.get_child_poses(cfgs, len(cfgs)), expected)
+
+
+def test_floating_get_child_poses_matches_scalar_poses():
+    origin = xyz_rpy_to_matrix([1.0, 2.0, 3.0, 0.1, 0.2, 0.3])
+    joint = Joint("joint", "floating", "parent", "child", origin=origin)
+    cfgs = np.array(
+        [
+            [0.5, 1.5, 2.5, 0.1, 0.2, 0.3],
+            [-1.0, 2.0, 0.0, -0.3, 0.0, 0.2],
+        ]
+    )
+    expected = np.stack([joint.get_child_pose(cfg) for cfg in cfgs])
+
+    assert np.allclose(joint.get_child_poses(cfgs, len(cfgs)), expected)
+    assert np.allclose(joint.get_child_poses(xyz_rpy_to_matrix(cfgs), len(cfgs)), expected)

tests/unit/test_urdf.py

@@ -2,7 +2,22 @@
 import pytest
 import trimesh
 
-from urchin import URDF, Joint, Link, Material, Transmission
+from urchin import URDF, Joint, Link, Material, Transmission, xyz_rpy_to_matrix
+
+
+def _make_multi_dof_urdf():
+    links = [
+        Link("base", None, None, None),
+        Link("floating_link", None, None, None),
+        Link("planar_link", None, None, None),
+        Link("tip", None, None, None),
+    ]
+    joints = [
+        Joint("floating_joint", "floating", "base", "floating_link"),
+        Joint("planar_joint", "planar", "floating_link", "planar_link"),
+        Joint("continuous_joint", "continuous", "planar_link", "tip"),
+    ]
+    return URDF("multi_dof", links, joints)
 
 
 def test_urchin(tmpdir):
@@ -109,3 +124,47 @@ def test_urchin(tmpdir):
     assert isinstance(x, URDF)
     x = x.copy(scale=[1, 1, 3])
     assert isinstance(x, URDF)
+
+
+def test_multi_dof_flat_configurations():
+    robot = _make_multi_dof_urdf()
+    cfgs = np.array(
+        [
+            [0.5, 1.5, 2.5, 0.1, 0.2, 0.3, 3.5, 4.5, 0.4],
+            [-1.0, 2.0, 0.0, -0.3, 0.0, 0.2, 1.0, -2.0, -0.5],
+        ]
+    )
+
+    expected = []
+    for cfg in cfgs:
+        expected.append(
+            robot.link_fk(
+                {
+                    "floating_joint": cfg[:6],
+                    "planar_joint": cfg[6:8],
+                    "continuous_joint": cfg[8],
+                },
+                link="tip",
+            )
+        )
+        assert np.allclose(robot.link_fk(cfg.tolist(), link="tip"), expected[-1])
+        assert np.allclose(robot.link_fk([cfg[:6], cfg[6:8], cfg[8]], link="tip"), expected[-1])
+
+    expected_batch = np.stack(expected)
+    assert np.allclose(robot.link_fk_batch(cfgs, link="tip"), expected_batch)
+    assert np.allclose(
+        robot.link_fk_batch(
+            {
+                "floating_joint": xyz_rpy_to_matrix(cfgs[:, :6]),
+                "planar_joint": cfgs[:, 6:8],
+                "continuous_joint": cfgs[:, 8],
+            },
+            link="tip",
+        ),
+        expected_batch,
+    )
+
+    with pytest.raises(ValueError, match="degree of freedom"):
+        robot.link_fk(np.zeros(3))
+    with pytest.raises(ValueError, match="degree of freedom"):
+        robot.link_fk_batch(np.zeros((2, 3)))

tests/unit/test_utils.py

@@ -14,6 +14,16 @@ def test_rpy_to_matrix():
     assert np.allclose(rpy_to_matrix([0, np.pi / 2, 0]), yr45)
     assert np.allclose(rpy_to_matrix([0, 0, np.pi / 2]), zr45)
     assert np.allclose(rpy_to_matrix(np.pi / 2 * np.ones(3)), c)
+    assert np.allclose(
+        rpy_to_matrix(
+            [
+                [np.pi / 2, 0, 0],
+                [0, np.pi / 2, 0],
+                [0, 0, np.pi / 2],
+            ]
+        ),
+        np.stack([xr45, yr45, zr45]),
+    )
 
 
 def test_matrix_to_rpy():
@@ -51,3 +61,13 @@ def test_xyz_rpy_to_matrix():
     assert np.allclose(xyz_rpy_to_matrix([1, 2, 3, np.pi / 2, 0, 0]), xr45)
     assert np.allclose(xyz_rpy_to_matrix([2, 3, 1, 0, np.pi / 2, 0]), yr45)
     assert np.allclose(xyz_rpy_to_matrix([3, 1, 2, 0, 0, np.pi / 2]), zr45)
+    assert np.allclose(
+        xyz_rpy_to_matrix(
+            [
+                [1, 2, 3, np.pi / 2, 0, 0],
+                [2, 3, 1, 0, np.pi / 2, 0],
+                [3, 1, 2, 0, 0, np.pi / 2],
+            ]
+        ),
+        np.stack([xr45, yr45, zr45]),
+    )

urchin/joint.py

@@ -748,16 +748,18 @@ def get_child_poses(self, cfg: Optional[npt.ArrayLike], n_cfgs: int) -> np.ndarr
 
         Parameters
         ----------
-        cfg : (n,) float or None
+        cfg : (n,) float, (n,2) float, (n,6) float, (n,4,4) float, or None
             The configuration values for this joint. They are interpreted
             based on the joint type as follows:
 
             - ``fixed`` - not used.
             - ``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 an (n,6) matrix, or an (n,4,4) matrix.
 
             If ``cfg`` is ``None``, then this just returns the joint pose.
 
@@ -783,9 +785,19 @@ def get_child_poses(self, cfg: Optional[npt.ArrayLike], n_cfgs: int) -> np.ndarr
             translation[:, :3, 3] = self.axis * cfg_arr[:, np.newaxis]
             return np.matmul(self.origin, translation)
         elif self.joint_type == "planar":
-            raise NotImplementedError()
+            cfg_arr = np.asanyarray(cfg, dtype=np.float64)
+            if cfg_arr.shape != (n_cfgs, 2):
+                raise ValueError("(n,2) float configuration required for planar joints")
+            translation = np.tile(np.eye(4, dtype=np.float64), (n_cfgs, 1, 1))
+            translation[:, :3, 3] = cfg_arr.dot(self.origin[:3, :2].T)
+            return np.matmul(self.origin, translation)
         elif self.joint_type == "floating":
-            raise NotImplementedError()
+            cfg_arr = configure_origin(cfg)
+            if cfg_arr.shape != (n_cfgs, 4, 4):
+                raise ValueError(
+                    "(n,6) or (n,4,4) float configuration required for floating joints"
+                )
+            return np.matmul(self.origin, cfg_arr)
         else:
             raise ValueError("Invalid configuration")
 

urchin/urdf.py

@@ -364,8 +364,9 @@ def link_fk(
         ----------
         cfg : dict or (n), float
             A map from joints or joint names to configuration values for
-            each joint, or a list containing a value for each actuated joint
-            in sorted order from the base link.
+            each joint, a list containing a value for each actuated joint,
+            or a flat list containing a value for each degree of freedom in
+            sorted order from the base link.
             If not specified, all joints are assumed to be in their default
             configurations.
         link : str or :class:`.Link`
@@ -452,7 +453,7 @@ def link_fk_batch(self, cfgs=None, link=None, links=None, use_names=False):
             One of the following: (A) a map from joints or joint names to vectors
             of joint configuration values, (B) a list of maps from joints or joint names
             to single configuration values, or (C) a list of ``n`` configuration vectors,
-            each of which has a vector with an entry for each actuated joint.
+            each of which has an entry for each degree of freedom.
         link : str or :class:`.Link`
             A single link or link name to return a pose for.
         links : list of str or list of :class:`.Link`
@@ -1350,6 +1351,14 @@ def _validate_graph(self) -> tuple[Link, list[Link]]:
             raise ValueError("URDF has no base link")
         return base_link, end_links
 
+    @classmethod
+    def _joint_dof(cls, joint: Joint) -> int:
+        if joint.joint_type == "planar":
+            return 2
+        elif joint.joint_type == "floating":
+            return 6
+        return 1
+
     def _process_cfg(
         self,
         cfg: Union[
@@ -1358,11 +1367,11 @@ def _process_cfg(
             npt.ArrayLike,
             None,
         ],
-    ) -> dict[Joint, float]:
+    ) -> dict[Joint, Union[float, npt.NDArray[np.float64]]]:
         """Process a joint configuration spec into a dictionary mapping
         joints to configuration values.
         """
-        joint_cfg: dict[Joint, float] = {}
+        joint_cfg: dict[Joint, Union[float, npt.NDArray[np.float64]]] = {}
         if cfg is None:
             return joint_cfg
         if isinstance(cfg, dict):
@@ -1372,12 +1381,36 @@ def _process_cfg(
                 elif isinstance(joint, Joint):
                     joint_cfg[joint] = cfg[joint]
         elif isinstance(cfg, (list, tuple, np.ndarray)):
-            if len(cfg) != len(self.actuated_joints):
+            # Preserve the existing per-joint nested form before treating scalar
+            # sequences as flattened values for each degree of freedom.
+            if len(cfg) == len(self.actuated_joints) and any(np.ndim(value) > 0 for value in cfg):
+                for joint, value in zip(self.actuated_joints, cfg):
+                    joint_cfg[joint] = value
+                return joint_cfg
+
+            total_dofs = sum(self._joint_dof(joint) for joint in self.actuated_joints)
+            try:
+                cfg_arr = np.asanyarray(cfg, dtype=np.float64)
+            except (TypeError, ValueError) as exc:
+                raise ValueError(
+                    "Cfg must contain one value per actuated joint or one value "
+                    "per degree of freedom"
+                ) from exc
+            if cfg_arr.shape != (total_dofs,):
                 raise ValueError(
-                    "Cfg must have same length as actuated joints if specified as a numerical array"
+                    "Cfg must contain one value per actuated joint or one value "
+                    "per degree of freedom"
                 )
-            for joint, value in zip(self.actuated_joints, cfg):
-                joint_cfg[joint] = value
+
+            start = 0
+            for joint in self.actuated_joints:
+                dof = self._joint_dof(joint)
+                end = start + dof
+                if dof == 1:
+                    joint_cfg[joint] = float(cfg_arr[start])
+                else:
+                    joint_cfg[joint] = cfg_arr[start:end]
+                start = end
         else:
             raise TypeError("Invalid type for config")
         return joint_cfg
@@ -1415,7 +1448,9 @@ def _process_cfgs(
                     n_cfgs = len(cfgs[joint])
         elif isinstance(cfgs, (list, tuple, np.ndarray)):
             n_cfgs = len(cfgs)
-            if isinstance(cfgs[0], dict):
+            if n_cfgs == 0:
+                pass
+            elif isinstance(cfgs[0], dict):
                 for cfg in cfgs:
                     for joint in cfg:
                         if isinstance(joint, str):
@@ -1431,20 +1466,39 @@ def _process_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] = cast(npt.NDArray[np.float64], cfgs[:, i])
+                cfgs_arr = np.asanyarray(cfgs, dtype=np.float64)
+                total_dofs = sum(self._joint_dof(joint) for joint in self.actuated_joints)
+                if cfgs_arr.ndim != 2 or cfgs_arr.shape[1] != total_dofs:
+                    raise ValueError(
+                        "Cfgs must contain one column per degree of freedom "
+                        f"(expected {total_dofs})"
+                    )
+
+                start = 0
+                for joint in self.actuated_joints:
+                    dof = self._joint_dof(joint)
+                    end = start + dof
+                    batch_joint_values = cfgs_arr[:, start:end]
+                    joint_cfg[joint] = batch_joint_values[:, 0] if dof == 1 else batch_joint_values
+                    start = end
         else:
             raise ValueError("Incorrectly formatted config array")
 
         for j in joint_cfg:
-            if isinstance(joint_cfg[j], list):
+            configured_values = joint_cfg[j]
+            if isinstance(configured_values, list):
                 from typing import cast as _cast
 
-                if len(_cast(list[float], joint_cfg[j])) == 0:
+                if len(_cast(list[float], configured_values)) == 0:
                     joint_cfg[j] = None
-                elif n_cfgs is not None and len(_cast(list[float], joint_cfg[j])) != n_cfgs:
+                elif n_cfgs is not None and len(_cast(list[float], configured_values)) != n_cfgs:
                     raise ValueError("Inconsistent number of configurations for joints")
+            elif (
+                configured_values is not None
+                and n_cfgs is not None
+                and len(configured_values) != n_cfgs
+            ):
+                raise ValueError("Inconsistent number of configurations for joints")
 
         from typing import cast as _cast