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

jaxsim/api/kin_dyn_parameters.py

@@ -0,0 +1,901 @@
+from __future__ import annotations
+
+import dataclasses
+
+import jax.lax
+import jax.numpy as jnp
+import jax_dataclasses
+import numpy as np
+import numpy.typing as npt
+from jax_dataclasses import Static
+
+import jaxsim.typing as jtp
+from jaxsim.math import Adjoint, Inertia, JointModel, supported_joint_motion
+from jaxsim.parsers.descriptions import JointDescription, ModelDescription
+from jaxsim.utils import HashedNumpyArray, JaxsimDataclass
+
+
+@jax_dataclasses.pytree_dataclass(eq=False, unsafe_hash=False)
+class KinDynParameters(JaxsimDataclass):
+    r"""
+    Class storing the kinematic and dynamic parameters of a model.
+
+    Attributes:
+        link_names: The names of the links.
+        parent_array: The parent array :math:`\lambda(i)` of the model.
+        support_body_array_bool:
+            The boolean support parent array :math:`\kappa_{b}(i)` of the model.
+        link_parameters: The parameters of the links.
+        frame_parameters: The parameters of the frames.
+        contact_parameters: The parameters of the collidable points.
+        joint_model: The joint model of the model.
+        joint_parameters: The parameters of the joints.
+    """
+
+    # Static
+    link_names: Static[tuple[str]]
+    _parent_array: Static[HashedNumpyArray]
+    _support_body_array_bool: Static[HashedNumpyArray]
+
+    # Links
+    link_parameters: LinkParameters
+
+    # Contacts
+    contact_parameters: ContactParameters
+
+    # Frames
+    frame_parameters: FrameParameters
+
+    # Joints
+    joint_model: JointModel
+    joint_parameters: JointParameters | None
+
+    @property
+    def parent_array(self) -> jtp.Vector:
+        r"""
+        Return the parent array :math:`\lambda(i)` of the model.
+        """
+        return self._parent_array.get()
+
+    @property
+    def support_body_array_bool(self) -> jtp.Matrix:
+        r"""
+        Return the boolean support parent array :math:`\kappa_{b}(i)` of the model.
+        """
+        return self._support_body_array_bool.get()
+
+    @staticmethod
+    def build(model_description: ModelDescription) -> KinDynParameters:
+        """
+        Construct the kinematic and dynamic parameters of the model.
+
+        Args:
+            model_description: The parsed model description to consider.
+
+        Returns:
+            The kinematic and dynamic parameters of the model.
+
+        Note:
+            This class is meant to ease the management of parametric models in
+            an automatic differentiation context.
+        """
+
+        # Extract the links ordered by their index.
+        # The link index corresponds to the body index ∈ [0, num_bodies - 1].
+        ordered_links = sorted(
+            list(model_description.links_dict.values()),
+            key=lambda l: l.index,
+        )
+
+        # Extract the joints ordered by their index.
+        # The joint index matches the index of its child link, therefore it starts
+        # from 1. Keep this in mind since this 1-indexing might introduce bugs.
+        ordered_joints = sorted(
+            list(model_description.joints_dict.values()),
+            key=lambda j: j.index,
+        )
+
+        # ================
+        # Links properties
+        # ================
+
+        # Create a list of link parameters objects.
+        link_parameters_list = [
+            LinkParameters.build_from_spatial_inertia(index=link.index, M=link.inertia)
+            for link in ordered_links
+        ]
+
+        # Create a vectorized object of link parameters.
+        link_parameters = jax.tree.map(lambda *l: jnp.stack(l), *link_parameters_list)
+
+        # =================
+        # Joints properties
+        # =================
+
+        # Create a list of joint parameters objects.
+        joint_parameters_list = [
+            JointParameters.build_from_joint_description(joint_description=joint)
+            for joint in ordered_joints
+        ]
+
+        # Create a vectorized object of joint parameters.
+        joint_parameters = (
+            jax.tree.map(lambda *l: jnp.stack(l), *joint_parameters_list)
+            if len(ordered_joints) > 0
+            else JointParameters(
+                index=jnp.array([], dtype=int),
+                friction_static=jnp.array([], dtype=float),
+                friction_viscous=jnp.array([], dtype=float),
+                position_limits_min=jnp.array([], dtype=float),
+                position_limits_max=jnp.array([], dtype=float),
+                position_limit_spring=jnp.array([], dtype=float),
+                position_limit_damper=jnp.array([], dtype=float),
+            )
+        )
+
+        # Create an object that defines the joint model (parent-to-child transforms).
+        joint_model = JointModel.build(description=model_description)
+
+        # ===================
+        # Contacts properties
+        # ===================
+
+        # Create the object storing the parameters of collidable points.
+        # Note that, contrarily to LinkParameters and JointsParameters, this object
+        # is not created with vmap. This is because the "body" attribute of the object
+        # must be Static for JIT-related reasons, and tree_map would not consider it
+        # as a leaf.
+        contact_parameters = ContactParameters.build_from(
+            model_description=model_description
+        )
+
+        # =================
+        # Frames properties
+        # =================
+
+        # Create the object storing the parameters of frames.
+        # Note that, contrarily to LinkParameters and JointsParameters, this object
+        # is not created with vmap. This is because the "name" attribute of the object
+        # must be Static for JIT-related reasons, and tree_map would not consider it
+        # as a leaf.
+        frame_parameters = FrameParameters.build_from(
+            model_description=model_description
+        )
+
+        # ===============
+        # Tree properties
+        # ===============
+
+        # Build the parent array λ(i) of the model.
+        # Note: the parent of the base link is not set since it's not defined.
+        parent_array_dict = {
+            link.index: link.parent.index
+            for link in ordered_links
+            if link.parent is not None
+        }
+        parent_array = jnp.array([-1, *list(parent_array_dict.values())], dtype=int)
+
+        # Instead of building the support parent array κ(i) for each link of the model,
+        # that has a variable length depending on the number of links connecting the
+        # root to the i-th link, we build the corresponding boolean version.
+        # Given a link index i, the boolean support parent array κb(i) is an array
+        # with the same number of elements of λ(i) having the i-th element set to True
+        # if the i-th link is in the support parent array κ(i), False otherwise.
+        # We store the boolean κb(i) as static attribute of the PyTree so that
+        # algorithms that need to access it can be jit-compiled.
+        def κb(link_index: jtp.IntLike) -> jtp.Vector:
+            κb = jnp.zeros(len(ordered_links), dtype=bool)
+
+            carry0 = κb, link_index
+
+            def scan_body(carry: tuple, i: jtp.Int) -> tuple[tuple, None]:
+
+                κb, active_link_index = carry
+
+                κb, active_link_index = jax.lax.cond(
+                    pred=(i == active_link_index),
+                    false_fun=lambda: (κb, active_link_index),
+                    true_fun=lambda: (
+                        κb.at[active_link_index].set(True),
+                        parent_array[active_link_index],
+                    ),
+                )
+
+                return (κb, active_link_index), None
+
+            (κb, _), _ = jax.lax.scan(
+                f=scan_body,
+                init=carry0,
+                xs=jnp.flip(jnp.arange(start=0, stop=len(ordered_links))),
+            )
+
+            return κb
+
+        support_body_array_bool = jax.vmap(κb)(
+            jnp.arange(start=0, stop=len(ordered_links))
+        )
+
+        # =================================
+        # Build and return KinDynParameters
+        # =================================
+
+        return KinDynParameters(
+            link_names=tuple(l.name for l in ordered_links),
+            _parent_array=HashedNumpyArray(array=parent_array),
+            _support_body_array_bool=HashedNumpyArray(array=support_body_array_bool),
+            link_parameters=link_parameters,
+            joint_model=joint_model,
+            joint_parameters=joint_parameters,
+            contact_parameters=contact_parameters,
+            frame_parameters=frame_parameters,
+        )
+
+    def __eq__(self, other: KinDynParameters) -> bool:
+
+        if not isinstance(other, KinDynParameters):
+            return False
+
+        return hash(self) == hash(other)
+
+    def __hash__(self) -> int:
+
+        return hash(
+            (
+                hash(self.number_of_links()),
+                hash(self.number_of_joints()),
+                hash(self.frame_parameters.name),
+                hash(self.frame_parameters.body),
+                hash(self._parent_array),
+                hash(self._support_body_array_bool),
+            )
+        )
+
+    # =============================
+    # Helpers to extract parameters
+    # =============================
+
+    def number_of_links(self) -> int:
+        """
+        Return the number of links of the model.
+
+        Returns:
+            The number of links of the model.
+        """
+
+        return len(self.link_names)
+
+    def number_of_joints(self) -> int:
+        """
+        Return the number of joints of the model.
+
+        Returns:
+            The number of joints of the model.
+        """
+
+        return len(self.joint_model.joint_names) - 1
+
+    def number_of_frames(self) -> int:
+        """
+        Return the number of frames of the model.
+
+        Returns:
+            The number of frames of the model.
+        """
+
+        return len(self.frame_parameters.name)
+
+    def support_body_array(self, link_index: jtp.IntLike) -> jtp.Vector:
+        r"""
+        Return the support parent array :math:`\kappa(i)` of a link.
+
+        Args:
+            link_index: The index of the link.
+
+        Returns:
+            The support parent array :math:`\kappa(i)` of the link.
+
+        Note:
+            This method returns a variable-length vector. In jit-compiled functions,
+            it's better to use the (static) boolean version `support_body_array_bool`.
+        """
+
+        return jnp.array(
+            jnp.where(self.support_body_array_bool[link_index])[0], dtype=int
+        )
+
+    # ========================
+    # Quantities used by RBDAs
+    # ========================
+
+    @jax.jit
+    def links_spatial_inertia(self) -> jtp.Array:
+        """
+        Return the spatial inertia of all links of the model.
+
+        Returns:
+            The spatial inertia of all links of the model.
+        """
+
+        return jax.vmap(LinkParameters.spatial_inertia)(self.link_parameters)
+
+    @jax.jit
+    def tree_transforms(self) -> jtp.Array:
+        r"""
+        Return the tree transforms of the model.
+
+        Returns:
+            The transforms
+            :math:`{}^{\text{pre}(i)} H_{\lambda(i)}`
+            of all joints of the model.
+        """
+
+        pre_Xi_λ = jax.vmap(
+            lambda i: self.joint_model.parent_H_predecessor(joint_index=i)
+            .inverse()
+            .adjoint()
+        )(jnp.arange(1, self.number_of_joints() + 1))
+
+        return jnp.vstack(
+            [
+                jnp.zeros(shape=(1, 6, 6), dtype=float),
+                pre_Xi_λ,
+            ]
+        )
+
+    @jax.jit
+    def joint_transforms(
+        self, joint_positions: jtp.VectorLike, base_transform: jtp.MatrixLike
+    ) -> jtp.Array:
+        r"""
+        Return the transforms of the joints.
+
+        Args:
+            joint_positions: The joint positions.
+            base_transform: The homogeneous matrix defining the base pose.
+
+        Returns:
+            The stacked transforms
+            :math:`{}^{i} \mathbf{H}_{\lambda(i)}(s)`
+            of each joint.
+        """
+
+        return self.joint_transforms_and_motion_subspaces(
+            joint_positions=joint_positions,
+            base_transform=base_transform,
+        )[0]
+
+    @jax.jit
+    def joint_motion_subspaces(
+        self, joint_positions: jtp.VectorLike, base_transform: jtp.MatrixLike
+    ) -> jtp.Array:
+        r"""
+        Return the motion subspaces of the joints.
+
+        Args:
+            joint_positions: The joint positions.
+            base_transform: The homogeneous matrix defining the base pose.
+
+        Returns:
+            The stacked motion subspaces :math:`\mathbf{S}(s)` of each joint.
+        """
+
+        return self.joint_transforms_and_motion_subspaces(
+            joint_positions=joint_positions,
+            base_transform=base_transform,
+        )[1]
+
+    @jax.jit
+    def joint_transforms_and_motion_subspaces(
+        self, joint_positions: jtp.VectorLike, base_transform: jtp.MatrixLike
+    ) -> tuple[jtp.Array, jtp.Array]:
+        r"""
+        Return the transforms and the motion subspaces of the joints.
+
+        Args:
+            joint_positions: The joint positions.
+            base_transform: The homogeneous matrix defining the base pose.
+
+        Returns:
+            A tuple containing the stacked transforms
+            :math:`{}^{i} \mathbf{H}_{\lambda(i)}(s)`
+            and the stacked motion subspaces :math:`\mathbf{S}(s)` of each joint.
+
+        Note:
+            The first transform, at index 0, provides the pose of the base link
+            w.r.t. the world frame. For both floating-base and fixed-base systems,
+            it takes into account the base pose and the optional transform
+            between the root frame of the model and the base link.
+        """
+
+        # Rename the base transform.
+        W_H_B = base_transform
+
+        # Extract the parent-to-predecessor fixed transforms of the joints.
+        λ_H_pre = jnp.vstack(
+            [
+                jnp.eye(4)[jnp.newaxis],
+                self.joint_model.λ_H_pre[1 : 1 + self.number_of_joints()],
+            ]
+        )
+
+        # Compute the transforms and motion subspaces of the joints.
+        if self.number_of_joints() == 0:
+            pre_H_suc_J, S_J = jnp.empty((0, 4, 4)), jnp.empty((0, 6, 1))
+        else:
+            pre_H_suc_J, S_J = jax.vmap(supported_joint_motion)(
+                jnp.array(self.joint_model.joint_types[1:]).astype(int),
+                jnp.array(joint_positions),
+                jnp.array([j.axis for j in self.joint_model.joint_axis]),
+            )
+
+        # Extract the transforms and motion subspaces of the joints.
+        # We stack the base transform W_H_B at index 0, and a dummy motion subspace
+        # for either the fixed or free-floating joint connecting the world to the base.
+        pre_H_suc = jnp.vstack([W_H_B[jnp.newaxis, ...], pre_H_suc_J])
+        S = jnp.vstack([jnp.zeros((6, 1))[jnp.newaxis, ...], S_J])
+
+        # Extract the successor-to-child fixed transforms.
+        # Note that here we include also the index 0 since suc_H_child[0] stores the
+        # optional pose of the base link w.r.t. the root frame of the model.
+        # This is supported by SDF when the base link <pose> element is defined.
+        suc_H_i = self.joint_model.suc_H_i[jnp.arange(0, 1 + self.number_of_joints())]
+
+        # Compute the overall transforms from the parent to the child of each joint by
+        # composing all the components of our joint model.
+        i_X_λ = jax.vmap(
+            lambda λ_Hi_pre, pre_Hi_suc, suc_Hi_i: Adjoint.from_transform(
+                transform=λ_Hi_pre @ pre_Hi_suc @ suc_Hi_i, inverse=True
+            )
+        )(λ_H_pre, pre_H_suc, suc_H_i)
+
+        return i_X_λ, S
+
+    # ============================
+    # Helpers to update parameters
+    # ============================
+
+    def set_link_mass(
+        self, link_index: jtp.IntLike, mass: jtp.FloatLike
+    ) -> KinDynParameters:
+        """
+        Set the mass of a link.
+
+        Args:
+            link_index: The index of the link.
+            mass: The mass of the link.
+
+        Returns:
+            The updated kinematic and dynamic parameters of the model.
+        """
+
+        link_parameters = self.link_parameters.replace(
+            mass=self.link_parameters.mass.at[link_index].set(mass)
+        )
+
+        return self.replace(link_parameters=link_parameters)
+
+    def set_link_inertia(
+        self, link_index: jtp.IntLike, inertia: jtp.MatrixLike
+    ) -> KinDynParameters:
+        r"""
+        Set the inertia tensor of a link.
+
+        Args:
+            link_index: The index of the link.
+            inertia: The :math:`3 \times 3` inertia tensor of the link.
+
+        Returns:
+            The updated kinematic and dynamic parameters of the model.
+        """
+
+        inertia_elements = LinkParameters.flatten_inertia_tensor(I=inertia)
+
+        link_parameters = self.link_parameters.replace(
+            inertia_elements=self.link_parameters.inertia_elements.at[link_index].set(
+                inertia_elements
+            )
+        )
+
+        return self.replace(link_parameters=link_parameters)
+
+
+@jax_dataclasses.pytree_dataclass
+class JointParameters(JaxsimDataclass):
+    """
+    Class storing the parameters of a joint.
+
+    Attributes:
+        index: The index of the joint.
+        friction_static: The static friction of the joint.
+        friction_viscous: The viscous friction of the joint.
+        position_limits_min: The lower position limit of the joint.
+        position_limits_max: The upper position limit of the joint.
+        position_limit_spring: The spring constant of the position limit.
+        position_limit_damper: The damper constant of the position limit.
+
+    Note:
+        This class is used inside KinDynParameters to store the vectorized set
+        of joint parameters.
+    """
+
+    index: jtp.Int
+
+    friction_static: jtp.Float
+    friction_viscous: jtp.Float
+
+    position_limits_min: jtp.Float
+    position_limits_max: jtp.Float
+
+    position_limit_spring: jtp.Float
+    position_limit_damper: jtp.Float
+
+    @staticmethod
+    def build_from_joint_description(
+        joint_description: JointDescription,
+    ) -> JointParameters:
+        """
+        Build a JointParameters object from a joint description.
+
+        Args:
+            joint_description: The joint description to consider.
+
+        Returns:
+            The JointParameters object.
+        """
+
+        s_min = joint_description.position_limit[0]
+        s_max = joint_description.position_limit[1]
+
+        position_limits_min = jnp.minimum(s_min, s_max)
+        position_limits_max = jnp.maximum(s_min, s_max)
+
+        friction_static = jnp.array(joint_description.friction_static).squeeze()
+        friction_viscous = jnp.array(joint_description.friction_viscous).squeeze()
+
+        position_limit_spring = jnp.array(
+            joint_description.position_limit_spring
+        ).squeeze()
+
+        position_limit_damper = jnp.array(
+            joint_description.position_limit_damper
+        ).squeeze()
+
+        return JointParameters(
+            index=jnp.array(joint_description.index).squeeze().astype(int),
+            friction_static=friction_static.astype(float),
+            friction_viscous=friction_viscous.astype(float),
+            position_limits_min=position_limits_min.astype(float),
+            position_limits_max=position_limits_max.astype(float),
+            position_limit_spring=position_limit_spring.astype(float),
+            position_limit_damper=position_limit_damper.astype(float),
+        )
+
+
+@jax_dataclasses.pytree_dataclass
+class LinkParameters(JaxsimDataclass):
+    r"""
+    Class storing the parameters of a link.
+
+    Attributes:
+        index: The index of the link.
+        mass: The mass of the link.
+        inertia_elements:
+            The unique elements of the :math:`3 \times 3` inertia tensor of the link.
+        center_of_mass:
+            The translation :math:`{}^L \mathbf{p}_{\text{CoM}}` between the origin
+            of the link frame and the link's center of mass, expressed in the
+            coordinates of the link frame.
+
+    Note:
+        This class is used inside KinDynParameters to store the vectorized set
+        of link parameters.
+    """
+
+    index: jtp.Int
+
+    mass: jtp.Float
+    center_of_mass: jtp.Vector
+    inertia_elements: jtp.Vector
+
+    @staticmethod
+    def build_from_spatial_inertia(index: jtp.IntLike, M: jtp.Matrix) -> LinkParameters:
+        r"""
+        Build a LinkParameters object from a :math:`6 \times 6` spatial inertia matrix.
+
+        Args:
+            index: The index of the link.
+            M: The :math:`6 \times 6` spatial inertia matrix of the link.
+
+        Returns:
+            The LinkParameters object.
+        """
+
+        # Extract the link parameters from the 6D spatial inertia.
+        m, L_p_CoM, I_CoM = Inertia.to_params(M=M)
+
+        # Extract only the necessary elements of the inertia tensor.
+        inertia_elements = I_CoM[jnp.triu_indices(3)]
+
+        return LinkParameters(
+            index=jnp.array(index).squeeze().astype(int),
+            mass=jnp.array(m).squeeze().astype(float),
+            center_of_mass=jnp.atleast_1d(jnp.array(L_p_CoM).squeeze()).astype(float),
+            inertia_elements=jnp.atleast_1d(inertia_elements.squeeze()).astype(float),
+        )
+
+    @staticmethod
+    def build_from_inertial_parameters(
+        index: jtp.IntLike, m: jtp.FloatLike, I: jtp.MatrixLike, c: jtp.VectorLike
+    ) -> LinkParameters:
+        r"""
+        Build a LinkParameters object from the inertial parameters of a link.
+
+        Args:
+            index: The index of the link.
+            m: The mass of the link.
+            I: The :math:`3 \times 3` inertia tensor of the link.
+            c: The translation between the link frame and the link's center of mass.
+
+        Returns:
+            The LinkParameters object.
+        """
+
+        # Extract only the necessary elements of the inertia tensor.
+        inertia_elements = I[jnp.triu_indices(3)]
+
+        return LinkParameters(
+            index=jnp.array(index).squeeze().astype(int),
+            mass=jnp.array(m).squeeze().astype(float),
+            center_of_mass=jnp.atleast_1d(c.squeeze()).astype(float),
+            inertia_elements=jnp.atleast_1d(inertia_elements.squeeze()).astype(float),
+        )
+
+    @staticmethod
+    def build_from_flat_parameters(
+        index: jtp.IntLike, parameters: jtp.VectorLike
+    ) -> LinkParameters:
+        """
+        Build a LinkParameters object from a flat vector of parameters.
+
+        Args:
+            index: The index of the link.
+            parameters: The flat vector of parameters.
+
+        Returns:
+            The LinkParameters object.
+        """
+        index = jnp.array(index).squeeze().astype(int)
+
+        m = jnp.array(parameters[0]).squeeze().astype(float)
+        c = jnp.atleast_1d(parameters[1:4].squeeze()).astype(float)
+        inertia_elements = jnp.atleast_1d(parameters[4:].squeeze()).astype(float)
+
+        return LinkParameters(
+            index=index, mass=m, inertia_elements=inertia_elements, center_of_mass=c
+        )
+
+    @staticmethod
+    def flat_parameters(params: LinkParameters) -> jtp.Vector:
+        """
+        Return the parameters of a link as a flat vector.
+
+        Args:
+            params: The link parameters.
+
+        Returns:
+            The parameters of the link as a flat vector.
+        """
+
+        return (
+            jnp.hstack(
+                [params.mass, params.center_of_mass.squeeze(), params.inertia_elements]
+            )
+            .squeeze()
+            .astype(float)
+        )
+
+    @staticmethod
+    def inertia_tensor(params: LinkParameters) -> jtp.Matrix:
+        r"""
+        Return the :math:`3 \times 3` inertia tensor of a link.
+
+        Args:
+            params: The link parameters.
+
+        Returns:
+            The :math:`3 \times 3` inertia tensor of the link.
+        """
+
+        return LinkParameters.unflatten_inertia_tensor(
+            inertia_elements=params.inertia_elements
+        )
+
+    @staticmethod
+    def spatial_inertia(params: LinkParameters) -> jtp.Matrix:
+        r"""
+        Return the :math:`6 \times 6` spatial inertia matrix of a link.
+
+        Args:
+            params: The link parameters.
+
+        Returns:
+            The :math:`6 \times 6` spatial inertia matrix of the link.
+        """
+
+        return Inertia.to_sixd(
+            mass=params.mass,
+            I=LinkParameters.inertia_tensor(params),
+            com=params.center_of_mass,
+        )
+
+    @staticmethod
+    def flatten_inertia_tensor(I: jtp.Matrix) -> jtp.Vector:
+        r"""
+        Flatten a :math:`3 \times 3` inertia tensor into a vector of unique elements.
+
+        Args:
+            I: The :math:`3 \times 3` inertia tensor.
+
+        Returns:
+            The vector of unique elements of the inertia tensor.
+        """
+
+        return jnp.atleast_1d(I[jnp.triu_indices(3)].squeeze())
+
+    @staticmethod
+    def unflatten_inertia_tensor(inertia_elements: jtp.Vector) -> jtp.Matrix:
+        r"""
+        Unflatten a vector of unique elements into a :math:`3 \times 3` inertia tensor.
+
+        Args:
+            inertia_elements: The vector of unique elements of the inertia tensor.
+
+        Returns:
+            The :math:`3 \times 3` inertia tensor.
+        """
+
+        I = jnp.zeros([3, 3]).at[jnp.triu_indices(3)].set(inertia_elements.squeeze())
+        return jnp.atleast_2d(jnp.where(I, I, I.T)).astype(float)
+
+
+@jax_dataclasses.pytree_dataclass
+class ContactParameters(JaxsimDataclass):
+    """
+    Class storing the contact parameters of a model.
+
+    Attributes:
+        body:
+            A tuple of integers representing, for each collidable point, the index of
+            the body (link) to which it is rigidly attached to.
+        point:
+            The translations between the link frame and the collidable point, expressed
+            in the coordinates of the parent link frame.
+        enabled:
+            A tuple of booleans representing, for each collidable point, whether it is
+            enabled or not in contact models.
+
+    Note:
+        Contrarily to LinkParameters and JointParameters, this class is not meant
+        to be created with vmap. This is because the `body` attribute must be `Static`.
+    """
+
+    body: Static[tuple[int, ...]] = dataclasses.field(default_factory=tuple)
+
+    point: jtp.Matrix = dataclasses.field(default_factory=lambda: jnp.array([]))
+
+    enabled: Static[tuple[bool, ...]] = dataclasses.field(default_factory=tuple)
+
+    @property
+    def indices_of_enabled_collidable_points(self) -> npt.NDArray:
+        """
+        Return the indices of the enabled collidable points.
+        """
+        return np.where(np.array(self.enabled))[0]
+
+    @staticmethod
+    def build_from(model_description: ModelDescription) -> ContactParameters:
+        """
+        Build a ContactParameters object from a model description.
+
+        Args:
+            model_description: The model description to consider.
+
+        Returns:
+            The ContactParameters object.
+        """
+
+        if len(model_description.collision_shapes) == 0:
+            return ContactParameters()
+
+        # Get all the links so that we can take their updated index.
+        links_dict = {link.name: link for link in model_description}
+
+        # Get all the enabled collidable points of the model.
+        collidable_points = model_description.all_enabled_collidable_points()
+
+        # Extract the positions L_p_C of the collidable points w.r.t. the link frames
+        # they are rigidly attached to.
+        points = jnp.vstack([cp.position for cp in collidable_points])
+
+        # Extract the indices of the links to which the collidable points are rigidly
+        # attached to.
+        link_index_of_points = tuple(
+            links_dict[cp.parent_link.name].index for cp in collidable_points
+        )
+
+        # Build the ContactParameters object.
+        cp = ContactParameters(
+            point=points,
+            body=link_index_of_points,
+            enabled=tuple(True for _ in link_index_of_points),
+        )
+
+        assert cp.point.shape[1] == 3, cp.point.shape[1]
+        assert cp.point.shape[0] == len(cp.body), cp.point.shape[0]
+
+        return cp
+
+
+@jax_dataclasses.pytree_dataclass
+class FrameParameters(JaxsimDataclass):
+    """
+    Class storing the frame parameters of a model.
+
+    Attributes:
+        name: A tuple of strings defining the frame names.
+        body:
+            A vector of integers representing, for each frame, the index of
+            the body (link) to which it is rigidly attached to.
+        transform: The transforms of the frames w.r.t. their parent link.
+
+    Note:
+        Contrarily to LinkParameters and JointParameters, this class is not meant
+        to be created with vmap. This is because the `name` attribute must be `Static`.
+    """
+
+    name: Static[tuple[str, ...]] = dataclasses.field(default_factory=tuple)
+
+    body: Static[tuple[int, ...]] = dataclasses.field(default_factory=tuple)
+
+    transform: jtp.Array = dataclasses.field(default_factory=lambda: jnp.array([]))
+
+    @staticmethod
+    def build_from(model_description: ModelDescription) -> FrameParameters:
+        """
+        Build a FrameParameters object from a model description.
+
+        Args:
+            model_description: The model description to consider.
+
+        Returns:
+            The FrameParameters object.
+        """
+
+        if len(model_description.frames) == 0:
+            return FrameParameters()
+
+        # Extract the frame names.
+        names = tuple(frame.name for frame in model_description.frames)
+
+        # For each frame, extract the index of the link to which it is attached to.
+        parent_link_index_of_frames = tuple(
+            model_description.links_dict[frame.parent.name].index
+            for frame in model_description.frames
+        )
+
+        # For each frame, extract the transform w.r.t. its parent link.
+        transforms = jnp.atleast_3d(
+            jnp.stack([frame.pose for frame in model_description.frames])
+        )
+
+        # Build the FrameParameters object.
+        fp = FrameParameters(
+            name=names,
+            transform=transforms.astype(float),
+            body=parent_link_index_of_frames,
+        )
+
+        assert fp.transform.shape[1:] == (4, 4), fp.transform.shape[1:]
+        assert fp.transform.shape[0] == len(fp.body), fp.transform.shape[0]
+
+        return fp