jaxsim 0.2.dev191__py3-none-any.whl → 0.6.1.dev5__py3-none-any.whl

jaxsim/logging.py

@@ -1,6 +1,5 @@
 import enum
 import logging
-from typing import Union
 
 import coloredlogs
 
@@ -20,7 +19,7 @@ def _logger() -> logging.Logger:
     return logging.getLogger(name=LOGGER_NAME)
 
 
-def set_logging_level(level: Union[int, LoggingLevel] = LoggingLevel.WARNING):
+def set_logging_level(level: int | LoggingLevel = LoggingLevel.WARNING):
     if isinstance(level, int):
         level = LoggingLevel(level)
 

jaxsim/math/__init__.py

@@ -0,0 +1,13 @@
+# Define the default standard gravity constant.
+StandardGravity = 9.81
+
+from .adjoint import Adjoint
+from .cross import Cross
+from .inertia import Inertia
+from .quaternion import Quaternion
+from .rotation import Rotation
+from .skew import Skew
+from .transform import Transform
+from .utils import safe_norm
+
+from .joint_model import JointModel, supported_joint_motion  # isort:skip

jaxsim/math/adjoint.py

@@ -1,17 +1,20 @@
 import jax.numpy as jnp
+import jaxlie
 
 import jaxsim.typing as jtp
-from jaxsim.sixd import so3
 
-from .quaternion import Quaternion
 from .skew import Skew
 
 
 class Adjoint:
+    """
+    A utility class for adjoint matrix operations.
+    """
+
     @staticmethod
     def from_quaternion_and_translation(
-        quaternion: jtp.Vector = jnp.array([1.0, 0, 0, 0]),
-        translation: jtp.Vector = jnp.zeros(3),
+        quaternion: jtp.Vector | None = None,
+        translation: jtp.Vector | None = None,
         inverse: bool = False,
         normalize_quaternion: bool = False,
     ) -> jtp.Matrix:
@@ -19,42 +22,67 @@ class Adjoint:
         Create an adjoint matrix from a quaternion and a translation.
 
         Args:
-            quaternion (jtp.Vector): A quaternion vector (4D) representing orientation.
-            translation (jtp.Vector): A translation vector (3D).
-            inverse (bool): Whether to compute the inverse adjoint. Default is False.
-            normalize_quaternion (bool): Whether to normalize the quaternion before creating the adjoint.
-                                         Default is False.
+            quaternion: A quaternion vector (4D) representing orientation.
+            translation: A translation vector (3D).
+            inverse: Whether to compute the inverse adjoint.
+            normalize_quaternion: Whether to normalize the quaternion before creating the adjoint.
 
         Returns:
-            jtp.Matrix: The adjoint matrix.
+            The adjoint matrix.
         """
+        quaternion = quaternion if quaternion is not None else jnp.array([1.0, 0, 0, 0])
+        translation = translation if translation is not None else jnp.zeros(3)
         assert quaternion.size == 4
         assert translation.size == 3
 
-        Q_sixd = so3.SO3.from_quaternion_xyzw(xyzw=Quaternion.to_xyzw(quaternion))
+        Q_sixd = jaxlie.SO3(wxyz=quaternion)
         Q_sixd = Q_sixd if not normalize_quaternion else Q_sixd.normalize()
 
         return Adjoint.from_rotation_and_translation(
             rotation=Q_sixd.as_matrix(), translation=translation, inverse=inverse
         )
 
+    @staticmethod
+    def from_transform(transform: jtp.MatrixLike, inverse: bool = False) -> jtp.Matrix:
+        """
+        Create an adjoint matrix from a transformation matrix.
+
+        Args:
+            transform: A 4x4 transformation matrix.
+            inverse: Whether to compute the inverse adjoint.
+
+        Returns:
+            The 6x6 adjoint matrix.
+        """
+
+        A_H_B = jnp.reshape(transform, (-1, 4, 4))
+
+        return (
+            jaxlie.SE3.from_matrix(matrix=A_H_B).adjoint()
+            if not inverse
+            else jaxlie.SE3.from_matrix(matrix=A_H_B).inverse().adjoint()
+        ).reshape(transform.shape[:-2] + (6, 6))
+
     @staticmethod
     def from_rotation_and_translation(
-        rotation: jtp.Matrix = jnp.eye(3),
-        translation: jtp.Vector = jnp.zeros(3),
+        rotation: jtp.Matrix | None = None,
+        translation: jtp.Vector | None = None,
         inverse: bool = False,
     ) -> jtp.Matrix:
         """
         Create an adjoint matrix from a rotation matrix and a translation vector.
 
         Args:
-            rotation (jtp.Matrix): A 3x3 rotation matrix.
-            translation (jtp.Vector): A translation vector (3D).
-            inverse (bool): Whether to compute the inverse adjoint. Default is False.
+            rotation: A 3x3 rotation matrix.
+            translation: A translation vector (3D).
+            inverse: Whether to compute the inverse adjoint. Default is False.
 
         Returns:
-            jtp.Matrix: The adjoint matrix.
+            The adjoint matrix.
         """
+        rotation = rotation if rotation is not None else jnp.eye(3)
+        translation = translation if translation is not None else jnp.zeros(3)
+
         assert rotation.shape == (3, 3)
         assert translation.size == 3
 
@@ -62,14 +90,14 @@ class Adjoint:
         A_o_B = translation.squeeze()
 
         if not inverse:
-            X = A_X_B = jnp.vstack(
+            X = A_X_B = jnp.vstack(  # noqa: F841
                 [
                     jnp.block([A_R_B, Skew.wedge(A_o_B) @ A_R_B]),
                     jnp.block([jnp.zeros(shape=(3, 3)), A_R_B]),
                 ]
             )
         else:
-            X = B_X_A = jnp.vstack(
+            X = B_X_A = jnp.vstack(  # noqa: F841
                 [
                     jnp.block([A_R_B.T, -A_R_B.T @ Skew.wedge(A_o_B)]),
                     jnp.block([jnp.zeros(shape=(3, 3)), A_R_B.T]),
@@ -84,10 +112,10 @@ class Adjoint:
         Convert an adjoint matrix to a transformation matrix.
 
         Args:
-            adjoint (jtp.Matrix): The adjoint matrix (6x6).
+            adjoint: The adjoint matrix (6x6).
 
         Returns:
-            jtp.Matrix: The transformation matrix (4x4).
+            The transformation matrix (4x4).
         """
         X = adjoint.squeeze()
         assert X.shape == (6, 6)
@@ -110,17 +138,23 @@ class Adjoint:
         Compute the inverse of an adjoint matrix.
 
         Args:
-            adjoint (jtp.Matrix): The adjoint matrix.
+            adjoint: The adjoint matrix.
 
         Returns:
-            jtp.Matrix: The inverse adjoint matrix.
+            The inverse adjoint matrix.
         """
-        A_X_B = adjoint
-        A_H_B = Adjoint.to_transform(adjoint=A_X_B)
+        A_X_B = adjoint.reshape(-1, 6, 6)
 
-        A_R_B = A_H_B[0:3, 0:3]
-        A_o_B = A_H_B[0:3, 3]
+        A_R_B_T = jnp.swapaxes(A_X_B[..., 0:3, 0:3], -2, -1)
+        A_T_B = A_X_B[..., 0:3, 3:6]
 
-        return Adjoint.from_rotation_and_translation(
-            rotation=A_R_B, translation=A_o_B, inverse=True
-        )
+        return jnp.concatenate(
+            [
+                jnp.concatenate(
+                    [A_R_B_T, -A_R_B_T @ A_T_B @ A_R_B_T],
+                    axis=-1,
+                ),
+                jnp.concatenate([jnp.zeros_like(A_R_B_T), A_R_B_T], axis=-1),
+            ],
+            axis=-2,
+        ).reshape(adjoint.shape)

jaxsim/math/cross.py

@@ -6,27 +6,36 @@ from .skew import Skew
 
 
 class Cross:
+    """
+    A utility class for cross product matrix operations.
+    """
+
     @staticmethod
     def vx(velocity_sixd: jtp.Vector) -> jtp.Matrix:
         """
         Compute the cross product matrix for 6D velocities.
 
         Args:
-            velocity_sixd (jtp.Vector): A 6D velocity vector [v, ω].
+            velocity_sixd: A 6D velocity vector [v, ω].
 
         Returns:
-            jtp.Matrix: The cross product matrix (6x6).
+            The cross product matrix (6x6).
 
         Raises:
             ValueError: If the input vector does not have a size of 6.
         """
-        v, ω = jnp.split(velocity_sixd.squeeze(), 2)
+        velocity_sixd = velocity_sixd.reshape(-1, 6)
+
+        v, ω = jnp.split(velocity_sixd, 2, axis=-1)
 
-        v_cross = jnp.vstack(
+        v_cross = jnp.concatenate(
             [
-                jnp.block([Skew.wedge(vector=ω), Skew.wedge(vector=v)]),
-                jnp.block([jnp.zeros(shape=(3, 3)), Skew.wedge(vector=ω)]),
-            ]
+                jnp.concatenate(
+                    [Skew.wedge(ω), jnp.zeros((ω.shape[0], 3, 3)).squeeze()], axis=-2
+                ),
+                jnp.concatenate([Skew.wedge(v), Skew.wedge(ω)], axis=-2),
+            ],
+            axis=-1,
         )
 
         return v_cross
@@ -37,10 +46,10 @@ class Cross:
         Compute the negative transpose of the cross product matrix for 6D velocities.
 
         Args:
-            velocity_sixd (jtp.Vector): A 6D velocity vector [v, ω].
+            velocity_sixd: A 6D velocity vector [v, ω].
 
         Returns:
-            jtp.Matrix: The negative transpose of the cross product matrix (6x6).
+            The negative transpose of the cross product matrix (6x6).
 
         Raises:
             ValueError: If the input vector does not have a size of 6.

jaxsim/math/inertia.py

@@ -1,5 +1,3 @@
-from typing import Tuple
-
 import jax.numpy as jnp
 
 import jaxsim.typing as jtp
@@ -8,18 +6,22 @@ from .skew import Skew
 
 
 class Inertia:
+    """
+    A utility class for inertia matrix operations.
+    """
+
     @staticmethod
     def to_sixd(mass: jtp.Float, com: jtp.Vector, I: jtp.Matrix) -> jtp.Matrix:
         """
         Convert mass, center of mass, and inertia matrix to a 6x6 inertia matrix.
 
         Args:
-            mass (jtp.Float): The mass of the body.
-            com (jtp.Vector): The center of mass position (3D).
-            I (jtp.Matrix): The 3x3 inertia matrix.
+            mass: The mass of the body.
+            com: The center of mass position (3D).
+            I: The 3x3 inertia matrix.
 
         Returns:
-            jtp.Matrix: The 6x6 inertia matrix.
+            The 6x6 inertia matrix.
 
         Raises:
             ValueError: If the shape of the inertia matrix I is not (3, 3).
@@ -39,15 +41,15 @@ class Inertia:
         return M
 
     @staticmethod
-    def to_params(M: jtp.Matrix) -> Tuple[jtp.Float, jtp.Vector, jtp.Matrix]:
+    def to_params(M: jtp.Matrix) -> tuple[jtp.Float, jtp.Vector, jtp.Matrix]:
         """
         Convert a 6x6 inertia matrix to mass, center of mass, and inertia matrix.
 
         Args:
-            M (jtp.Matrix): The 6x6 inertia matrix.
+            M: The 6x6 inertia matrix.
 
         Returns:
-            Tuple[jtp.Float, jtp.Vector, jtp.Matrix]: A tuple containing mass, center of mass (3D), and inertia matrix (3x3).
+            A tuple containing mass, center of mass (3D), and inertia matrix (3x3).
 
         Raises:
             ValueError: If the input matrix M has an unexpected shape.

jaxsim/math/joint_model.py

@@ -0,0 +1,289 @@
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+import jax_dataclasses
+import jaxlie
+from jax_dataclasses import Static
+
+import jaxsim.typing as jtp
+from jaxsim.parsers.descriptions import JointGenericAxis, JointType, ModelDescription
+from jaxsim.parsers.kinematic_graph import KinematicGraphTransforms
+
+from .rotation import Rotation
+from .transform import Transform
+
+
+@jax_dataclasses.pytree_dataclass
+class JointModel:
+    """
+    Class describing the joint kinematics of a robot model.
+
+    Attributes:
+        λ_H_pre:
+            The homogeneous transformation between the parent link and
+            the predecessor frame of each joint.
+        suc_H_i:
+            The homogeneous transformation between the successor frame and
+            the child link of each joint.
+        joint_dofs: The number of DoFs of each joint.
+        joint_names: The names of each joint.
+        joint_types: The types of each joint.
+
+    Note:
+        Due to the presence of the static attributes, this class needs to be created
+        already in a vectorized form. In other words, it cannot be created using vmap.
+    """
+
+    λ_H_pre: jtp.Array
+    suc_H_i: jtp.Array
+
+    joint_dofs: Static[tuple[int, ...]]
+    joint_names: Static[tuple[str, ...]]
+    joint_types: Static[tuple[int, ...]]
+    joint_axis: Static[tuple[JointGenericAxis, ...]]
+
+    @staticmethod
+    def build(description: ModelDescription) -> JointModel:
+        """
+        Build the joint model of a model description.
+
+        Args:
+            description: The model description to consider.
+
+        Returns:
+            The joint model of the considered model description.
+        """
+
+        # The link index is equal to its body index: [0, number_of_bodies - 1].
+        ordered_links = sorted(
+            list(description.links_dict.values()),
+            key=lambda l: l.index,
+        )
+
+        # Note: the joint index is equal to its child link index, therefore it
+        # starts from 1.
+        ordered_joints = sorted(
+            list(description.joints_dict.values()),
+            key=lambda j: j.index,
+        )
+
+        # Allocate the parent-to-predecessor and successor-to-child transforms.
+        λ_H_pre = jnp.zeros(shape=(1 + len(ordered_joints), 4, 4), dtype=float)
+        suc_H_i = jnp.zeros(shape=(1 + len(ordered_joints), 4, 4), dtype=float)
+
+        # Initialize an identical parent-to-predecessor transform for the joint
+        # between the world frame W and the base link B.
+        λ_H_pre = λ_H_pre.at[0].set(jnp.eye(4))
+
+        # Initialize the successor-to-child transform of the joint between the
+        # world frame W and the base link B.
+        # We store here the optional transform between the root frame of the model
+        # and the base link frame (this is needed only if the pose of the link frame
+        # w.r.t. the implicit __model__ SDF frame is not the identity).
+        suc_H_i = suc_H_i.at[0].set(ordered_links[0].pose)
+
+        # Create the object to compute forward kinematics.
+        fk = KinematicGraphTransforms(graph=description)
+
+        # Compute the parent-to-predecessor and successor-to-child transforms for
+        # each joint belonging to the model.
+        # Note that the joint indices starts from i=1 given our joint model,
+        # therefore the entries at index 0 are not updated.
+        for joint in ordered_joints:
+            λ_H_pre = λ_H_pre.at[joint.index].set(
+                fk.relative_transform(relative_to=joint.parent.name, name=joint.name)
+            )
+            suc_H_i = suc_H_i.at[joint.index].set(
+                fk.relative_transform(relative_to=joint.name, name=joint.child.name)
+            )
+
+        # Define the DoFs of the base link.
+        base_dofs = 0 if description.fixed_base else 6
+
+        # We always add a dummy fixed joint between world and base.
+        # TODO: Port floating-base support also at this level, not only in RBDAs.
+        return JointModel(
+            λ_H_pre=λ_H_pre,
+            suc_H_i=suc_H_i,
+            # Static attributes
+            joint_dofs=tuple([base_dofs] + [1 for _ in ordered_joints]),
+            joint_names=tuple(["world_to_base"] + [j.name for j in ordered_joints]),
+            joint_types=tuple([JointType.Fixed] + [j.jtype for j in ordered_joints]),
+            joint_axis=tuple(JointGenericAxis(axis=j.axis) for j in ordered_joints),
+        )
+
+    def parent_H_child(
+        self, joint_index: jtp.IntLike, joint_position: jtp.VectorLike
+    ) -> tuple[jtp.Matrix, jtp.Array]:
+        r"""
+        Compute the homogeneous transformation between the parent link and
+        the child link of a joint, and the corresponding motion subspace.
+
+        Args:
+            joint_index: The index of the joint.
+            joint_position: The position of the joint.
+
+        Returns:
+            A tuple containing the homogeneous transformation
+            :math:`{}^{\lambda(i)} \mathbf{H}_i(s)`
+            and the motion subspace :math:`\mathbf{S}(s)`.
+        """
+
+        i = joint_index
+        s = joint_position
+
+        # Get the components of the joint model.
+        λ_Hi_pre = self.parent_H_predecessor(joint_index=i)
+        pre_Hi_suc, S = self.predecessor_H_successor(joint_index=i, joint_position=s)
+        suc_Hi_i = self.successor_H_child(joint_index=i)
+
+        # Compose all the transforms.
+        return λ_Hi_pre @ pre_Hi_suc @ suc_Hi_i, S
+
+    @jax.jit
+    def child_H_parent(
+        self, joint_index: jtp.IntLike, joint_position: jtp.VectorLike
+    ) -> tuple[jtp.Matrix, jtp.Array]:
+        r"""
+        Compute the homogeneous transformation between the child link and
+        the parent link of a joint, and the corresponding motion subspace.
+
+        Args:
+            joint_index: The index of the joint.
+            joint_position: The position of the joint.
+
+        Returns:
+            A tuple containing the homogeneous transformation
+            :math:`{}^{i} \mathbf{H}_{\lambda(i)}(s)`
+            and the motion subspace :math:`\mathbf{S}(s)`.
+        """
+
+        λ_Hi_i, S = self.parent_H_child(
+            joint_index=joint_index, joint_position=joint_position
+        )
+
+        i_Hi_λ = Transform.inverse(λ_Hi_i)
+
+        return i_Hi_λ, S
+
+    def parent_H_predecessor(self, joint_index: jtp.IntLike) -> jtp.Matrix:
+        r"""
+        Return the homogeneous transformation between the parent link and
+        the predecessor frame of a joint.
+
+        Args:
+            joint_index: The index of the joint.
+
+        Returns:
+            The homogeneous transformation
+            :math:`{}^{\lambda(i)} \mathbf{H}_{\text{pre}(i)}`.
+        """
+
+        return self.λ_H_pre[joint_index]
+
+    def predecessor_H_successor(
+        self, joint_index: jtp.IntLike, joint_position: jtp.VectorLike
+    ) -> tuple[jtp.Matrix, jtp.Array]:
+        r"""
+        Compute the homogeneous transformation between the predecessor and
+        the successor frame of a joint, and the corresponding motion subspace.
+
+        Args:
+            joint_index: The index of the joint.
+            joint_position: The position of the joint.
+
+        Returns:
+            A tuple containing the homogeneous transformation
+            :math:`{}^{\text{pre}(i)} \mathbf{H}_{\text{suc}(i)}(s)`
+            and the motion subspace :math:`\mathbf{S}(s)`.
+        """
+
+        pre_H_suc, S = supported_joint_motion(
+            self.joint_types[joint_index],
+            joint_position,
+            self.joint_axis[joint_index].axis,
+        )
+
+        return pre_H_suc, S
+
+    def successor_H_child(self, joint_index: jtp.IntLike) -> jtp.Matrix:
+        r"""
+        Return the homogeneous transformation between the successor frame and
+        the child link of a joint.
+
+        Args:
+            joint_index: The index of the joint.
+
+        Returns:
+            The homogeneous transformation
+            :math:`{}^{\text{suc}(i)} \mathbf{H}_i`.
+        """
+
+        return self.suc_H_i[joint_index]
+
+
+@jax.jit
+def supported_joint_motion(
+    joint_type: jtp.IntLike,
+    joint_position: jtp.VectorLike,
+    joint_axis: jtp.VectorLike | None = None,
+    /,
+) -> tuple[jtp.Matrix, jtp.Array]:
+    """
+    Compute the homogeneous transformation and motion subspace of a joint.
+
+    Args:
+        joint_type: The type of the joint.
+        joint_position: The position of the joint.
+        joint_axis: The optional 3D axis of rotation or translation of the joint.
+
+    Returns:
+        A tuple containing the homogeneous transformation and the motion subspace.
+    """
+
+    # Prepare the joint position
+    s = jnp.array(joint_position).astype(float)
+
+    def compute_F() -> tuple[jtp.Matrix, jtp.Array]:
+        return jaxlie.SE3.identity(), jnp.zeros(shape=(6, 1))
+
+    def compute_R() -> tuple[jtp.Matrix, jtp.Array]:
+
+        # Get the additional argument specifying the joint axis.
+        # This is a metadata required by only some joint types.
+        axis = jnp.array(joint_axis).astype(float).squeeze()
+
+        pre_H_suc = jaxlie.SE3.from_matrix(
+            matrix=jnp.eye(4).at[:3, :3].set(Rotation.from_axis_angle(vector=s * axis))
+        )
+
+        S = jnp.vstack(jnp.hstack([jnp.zeros(3), axis]))
+
+        return pre_H_suc, S
+
+    def compute_P() -> tuple[jtp.Matrix, jtp.Array]:
+
+        # Get the additional argument specifying the joint axis.
+        # This is a metadata required by only some joint types.
+        axis = jnp.array(joint_axis).astype(float).squeeze()
+
+        pre_H_suc = jaxlie.SE3.from_rotation_and_translation(
+            rotation=jaxlie.SO3.identity(),
+            translation=jnp.array(s * axis),
+        )
+
+        S = jnp.vstack(jnp.hstack([axis, jnp.zeros(3)]))
+
+        return pre_H_suc, S
+
+    pre_H_suc, S = jax.lax.switch(
+        index=joint_type,
+        branches=(
+            compute_F,  # JointType.Fixed
+            compute_R,  # JointType.Revolute
+            compute_P,  # JointType.Prismatic
+        ),
+    )
+
+    return pre_H_suc.as_matrix(), S