Add spherical joints

src/adam/casadi/casadi_like.py

@@ -380,12 +380,15 @@ def mxv(m: CasadiLike, v: CasadiLike) -> CasadiLike:
     def vxs(v: CasadiLike, c: CasadiLike) -> CasadiLike:
         """
         Vector times scalar multiplication for CasADi.
+        For multi-DOF joints (v has more than 1 column), performs matrix multiply.
 
         Args:
-            v: Vector (CasadiLike)
-            c: Scalar (CasadiLike)
+            v: Vector or matrix (CasadiLike)
+            c: Scalar or vector (CasadiLike)
 
         Returns:
-            CasadiLike: Result of vector times scalar
+            CasadiLike: Result of v * c (1-DOF) or v @ c (multi-DOF)
         """
+        if v.array.shape[1] > 1:
+            return CasadiLike(cs.mtimes(v.array, c.array))
         return CasadiLike(v.array * c.array)

src/adam/casadi/computations.py

@@ -4,12 +4,14 @@
 
 import casadi as cs
 import numpy as np
+import idyntree.bindings as iDynTree
 
 from adam.casadi.casadi_like import SpatialMath
 from adam.core import RBDAlgorithms
 from adam.core.constants import Representations
 from adam.model import Model, build_model_factory
 from adam.model.kindyn_mixin import KinDynFactoryMixin
+from adam.model.conversions import from_idyntree_model
 
 
 class KinDynComputations(KinDynFactoryMixin):
@@ -22,17 +24,41 @@ def __init__(
         root_link: str = None,
         gravity: np.array = np.array([0.0, 0.0, -9.80665, 0.0, 0.0, 0.0]),
         f_opts: dict = dict(jit=False, jit_options=dict(flags="-Ofast"), cse=True),
+        convert_revolute_to_spherical: bool = False,
     ) -> None:
         """
         Args:
             urdfstring (str): path/string of a URDF or a MuJoCo MjModel.
                 NOTE: The parameter name `urdfstring` is deprecated and will be renamed to `model` in a future release.
             joints_name_list (list): list of the actuated joints
             root_link (str, optional): Deprecated. The root link is automatically chosen as the link with no parent in the URDF. Defaults to None.
+            convert_revolute_to_spherical (bool, optional): If True, converts revolute joints to spherical joints. Defaults to False.
         """
         math = SpatialMath()
-        factory = build_model_factory(description=urdfstring, math=math)
-        model = Model.build(factory=factory, joints_name_list=joints_name_list)
+        # if convert_revolute_to_spherical:
+        #     model_loader = iDynTree.ModelLoader()
+        #     options = model_loader.parsingOptions()
+        #     options.convertThreeRevoluteJointsToSphericalJoint = (
+        #         convert_revolute_to_spherical
+        #     )
+        #     model_loader.setParsingOptions(options)
+        #     model_loader.loadReducedModelFromFile(str(urdfstring), joints_name_list)
+        #     model_idy = model_loader.model()
+        #     model = from_idyntree_model(model_idy, math=math)
+        # else:
+        #     factory = build_model_factory(description=urdfstring, math=math)
+        #     model = Model.build(factory=factory, joints_name_list=joints_name_list)
+
+        model_loader = iDynTree.ModelLoader()
+        options = model_loader.parsingOptions()
+        options.convertThreeRevoluteJointsToSphericalJoint = (
+            convert_revolute_to_spherical
+        )
+        model_loader.setParsingOptions(options)
+        model_loader.loadReducedModelFromFile(str(urdfstring), joints_name_list)
+        model_idyntree = model_loader.model()
+        model = from_idyntree_model(model_idyntree, math=math)
+
         self.rbdalgos = RBDAlgorithms(model=model, math=math)
         self.NDoF = self.rbdalgos.NDoF
         self.g = gravity

src/adam/casadi/inverse_kinematics.py

@@ -43,6 +43,7 @@ def __init__(
         joints_list: list[str],
         joint_limits_active: bool = True,
         solver_settings: dict[str, Any] = None,
+        convert_revolute_to_spherical: bool = False,
     ):
         """Initialize the InverseKinematics solver.
 
@@ -52,9 +53,13 @@ def __init__(
             joint_limits_active (bool, optional): If True, enforces joint limits. Defaults to True.
             solver_settings (dict[str, ], optional): Settings for the solver. Defaults to None.
         """
-        self.kd = KinDynComputations(urdf_path, joints_list)
-        self.ndof = len(joints_list)
-        self.joints_list = joints_list
+        self.kd = KinDynComputations(
+            urdf_path,
+            joints_list,
+            convert_revolute_to_spherical=convert_revolute_to_spherical,
+        )
+        self.ndof = self.kd.NDoF
+        self.joints_list = self.kd.rbdalgos.model.actuated_joints
 
         # Opti variables --------------------------------------------------
         self.opti = cs.Opti()

src/adam/core/rbd_algorithms.py

@@ -131,13 +131,16 @@ def tile_batch(arr):
                     Xup[idx],
                 )
 
-        def block_index(node_idx: int) -> int | None:
+        def block_info(node_idx: int) -> list[tuple[int, int]] | None:
+            """Return list of (block_index, local_dof_offset) for a node, or None."""
             if node_idx == root_idx:
-                return 0
+                return [(0, 0)]
             joint_idx = joint_indices[node_idx]
             if joint_idx is None:
                 return None
-            return 1 + int(joint_idx)
+            if isinstance(joint_idx, tuple):
+                return [(1 + j, k) for k, j in enumerate(joint_idx)]
+            return [(1 + int(joint_idx), 0)]
 
         sizes = [6] + [1] * n
         blocks: list[list[npt.ArrayLike | None]] = [
@@ -148,30 +151,38 @@ def block_index(node_idx: int) -> int | None:
             Phi_i = Phi[idx]
             Phi_T = math.swapaxes(Phi_i, -2, -1)
             F = math.mtimes(Ic_comp[idx], Phi_i)
-            ri = block_index(idx)
+            ri_info = block_info(idx)
 
             if idx == root_idx:
                 blocks[0][0] = math.mtimes(Phi_T, F)
                 continue
 
-            if ri is not None:
-                blocks[ri][ri] = math.mtimes(Phi_T, F)
+            if ri_info is not None:
+                diag = math.mtimes(Phi_T, F)
+                for ri, oi in ri_info:
+                    for rj, oj in ri_info:
+                        blocks[ri][rj] = diag[
+                            ..., oi : oi + sizes[ri], oj : oj + sizes[rj]
+                        ]
 
             current = idx
             F_path = F
             while current != root_idx:
                 parent = parent_indices[current]
                 F_path = math.mtimes(Xup_T[current], F_path)
-                rj = block_index(parent)
-                if rj is None:
+                rj_info = block_info(parent)
+                if rj_info is None:
                     current = parent
                     continue
-                Bij = math.mtimes(math.swapaxes(F_path, -2, -1), Phi[parent])
-                if ri is None:
+                Bij_full = math.mtimes(math.swapaxes(F_path, -2, -1), Phi[parent])
+                if ri_info is None:
                     current = parent
                     continue
-                blocks[ri][rj] = Bij
-                blocks[rj][ri] = math.swapaxes(Bij, -2, -1)
+                for ri, oi in ri_info:
+                    for rj, oj in rj_info:
+                        sub = Bij_full[..., oi : oi + sizes[ri], oj : oj + sizes[rj]]
+                        blocks[ri][rj] = sub
+                        blocks[rj][ri] = math.swapaxes(sub, -2, -1)
                 current = parent
 
         row_tensors = []
@@ -206,14 +217,17 @@ def block_index(node_idx: int) -> int | None:
         )
 
         zero_col = math.factory.zeros(batch_shape + (6, 1))
+        joint_local_dof = self._joint_index_to_local_dof
         joint_cols = []
         for jidx in range(n):
             node_idx = joint_to_node.get(jidx)
             if node_idx is not None:
-                col = math.mtimes(
+                full_col = math.mtimes(
                     math.swapaxes(X_G[node_idx], -2, -1),
                     math.mtimes(Ic_comp[node_idx], Phi[node_idx]),
                 )
+                local = joint_local_dof.get(jidx, 0)
+                col = full_col[..., :, local : local + 1]
             else:
                 col = zero_col
             joint_cols.append(col)
@@ -547,8 +561,8 @@ def jacobian_dot(
             a = a + self.math.vxs(J_dot_j, q_dot)
 
             if joint.idx is not None:
-                cols[joint.idx] = J_j
-                cols_dot[joint.idx] = J_dot_j
+                self._scatter_cols(cols, joint.idx, J_j)
+                self._scatter_cols(cols_dot, joint.idx, J_dot_j)
 
         zero_col = self.math.factory.zeros(batch_size + (6, 1))
         cols = [zero_col if c is None else c for c in cols]
@@ -829,7 +843,8 @@ def rnea(
             else:
                 joint_idx = joint_indices[idx]
                 if joint_idx is not None:
-                    tau_joint_cols[joint_idx] = math.mxv(Phi_T, Fi)
+                    tau_val = math.mxv(Phi_T, Fi)
+                    self._scatter_entries(tau_joint_cols, joint_idx, tau_val)
                 parent = parent_indices[idx]
                 if parent >= 0:
                     f[parent] = f[parent] + math.mxv(
@@ -1035,8 +1050,12 @@ def tile_batch(arr):
                 U_i = math.mtimes(IA[idx], Si)
                 d_i = math.mtimes(math.swapaxes(Si, -2, -1), U_i)
                 joint_idx = joint_indices[idx]
-                tau_vec = joint_torques_eff[..., joint_idx]
-                Si_T_pA = math.mxv(math.swapaxes(Si, -2, -1), pA[idx])[..., 0]
+                if isinstance(joint_idx, tuple):
+                    tau_slice = slice(joint_idx[0], joint_idx[-1] + 1)
+                else:
+                    tau_slice = slice(joint_idx, joint_idx + 1)
+                tau_vec = joint_torques_eff[..., tau_slice]
+                Si_T_pA = math.mxv(math.swapaxes(Si, -2, -1), pA[idx])
                 u_i = tau_vec - Si_T_pA
 
                 d_list[idx] = d_i
@@ -1080,14 +1099,15 @@ def tile_batch(arr):
 
             if Si is not None and joint_idx is not None:
                 U_i = U_list[idx]
-                U_T_rel_acc = math.mxv(math.swapaxes(U_i, -2, -1), rel_acc)[..., 0]
+                U_T_rel_acc = math.mxv(math.swapaxes(U_i, -2, -1), rel_acc)
                 num = u_list[idx] - U_T_rel_acc
                 inv_d = inv_d_list[idx]
                 num_expanded = expand_to_match(num, inv_d)
                 gain_qdd = math.mtimes(inv_d, num_expanded)
                 qdd_col = gain_qdd[..., 0]
-                if joint_idx < n:
-                    qdd_entries[joint_idx] = qdd_col
+                max_idx = max(joint_idx) if isinstance(joint_idx, tuple) else joint_idx
+                if max_idx < n:
+                    self._scatter_entries(qdd_entries, joint_idx, qdd_col)
                 a_correction_vec = math.mxv(Si, qdd_col)
                 a[idx] = a_pre + a_correction_vec
             else:
@@ -1223,12 +1243,35 @@ def _joints_jacobian_from_traversal(
                 continue
             L_H_j = L_H_B @ B_H_j
             S = joint.motion_subspace()
-            cols[joint.idx] = self.math.adjoint(L_H_j) @ S
+            J_block = self.math.adjoint(L_H_j) @ S
+            self._scatter_cols(cols, joint.idx, J_block)
 
         zero_col = self.math.factory.zeros(batch_size + (6, 1))
         cols = [zero_col if col is None else col for col in cols]
         return self.math.concatenate(cols, axis=-1)
 
+    @staticmethod
+    def _scatter_cols(
+        cols: list, idx: int | tuple[int, ...], block: npt.ArrayLike
+    ) -> None:
+        """Assign columns of *block* to *cols* at the DOF positions given by *idx*."""
+        if isinstance(idx, tuple):
+            for i, j in enumerate(idx):
+                cols[j] = block[..., :, i : i + 1]
+        else:
+            cols[idx] = block
+
+    @staticmethod
+    def _scatter_entries(
+        entries: list, idx: int | tuple[int, ...], vec: npt.ArrayLike
+    ) -> None:
+        """Assign elements of *vec* to *entries* at the DOF positions given by *idx*."""
+        if isinstance(idx, tuple):
+            for i, j in enumerate(idx):
+                entries[j] = vec[..., i : i + 1]
+        else:
+            entries[idx] = vec
+
     def _default_joint_value(
         self, joint_positions: npt.ArrayLike, fallback: npt.ArrayLike | None = None
     ) -> npt.ArrayLike:
@@ -1266,11 +1309,14 @@ def _prepare_tree_cache(self) -> None:
         self._node_indices = tuple(range(node_count))
         self._rev_node_indices = tuple(reversed(self._node_indices))
         self._parent_indices = [-1] * node_count
-        self._joint_indices_per_node: list[int | None] = [None] * node_count
+        self._joint_indices_per_node: list[int | tuple[int, ...] | None] = [
+            None
+        ] * node_count
         self._motion_subspaces: list[npt.ArrayLike] = [None] * node_count
         self._spatial_inertias: list[npt.ArrayLike] = [None] * node_count
         self._joints_per_node: list[Joint | None] = [None] * node_count
         self._joint_index_to_node: dict[int, int] = {}
+        self._joint_index_to_local_dof: dict[int, int] = {}
         self._joint_by_child: dict[str, Joint] = {}
         self._joint_chain_cache: dict[str, tuple[Joint, ...]] = {
             self.root_link: tuple()
@@ -1293,7 +1339,13 @@ def _prepare_tree_cache(self) -> None:
                 self._motion_subspaces[idx] = joint.motion_subspace()
                 self._joint_indices_per_node[idx] = joint.idx
                 if joint.idx is not None:
-                    self._joint_index_to_node[int(joint.idx)] = idx
+                    if isinstance(joint.idx, tuple):
+                        for local, global_idx in enumerate(joint.idx):
+                            self._joint_index_to_node[global_idx] = idx
+                            self._joint_index_to_local_dof[global_idx] = local
+                    else:
+                        self._joint_index_to_node[int(joint.idx)] = idx
+                        self._joint_index_to_local_dof[int(joint.idx)] = 0
 
         for joint in self.model.joints.values():
             self._joint_by_child[joint.child] = joint