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

jaxsim/rbda/collidable_points.py

@@ -0,0 +1,156 @@
+import jax
+import jax.numpy as jnp
+import jaxlie
+
+import jaxsim.api as js
+import jaxsim.typing as jtp
+from jaxsim.math import Adjoint, Skew
+
+from . import utils
+
+
+def collidable_points_pos_vel(
+    model: js.model.JaxSimModel,
+    *,
+    base_position: jtp.Vector,
+    base_quaternion: jtp.Vector,
+    joint_positions: jtp.Vector,
+    base_linear_velocity: jtp.Vector,
+    base_angular_velocity: jtp.Vector,
+    joint_velocities: jtp.Vector,
+) -> tuple[jtp.Matrix, jtp.Matrix]:
+    """
+
+    Compute the position and linear velocity of the enabled collidable points in the world frame.
+
+    Args:
+        model: The model to consider.
+        base_position: The position of the base link.
+        base_quaternion: The quaternion of the base link.
+        joint_positions: The positions of the joints.
+        base_linear_velocity:
+            The linear velocity of the base link in inertial-fixed representation.
+        base_angular_velocity:
+            The angular velocity of the base link in inertial-fixed representation.
+        joint_velocities: The velocities of the joints.
+
+    Returns:
+        A tuple containing the position and linear velocity of the enabled collidable points.
+    """
+
+    # Get the indices of the enabled collidable points.
+    indices_of_enabled_collidable_points = (
+        model.kin_dyn_parameters.contact_parameters.indices_of_enabled_collidable_points
+    )
+
+    parent_link_idx_of_enabled_collidable_points = jnp.array(
+        model.kin_dyn_parameters.contact_parameters.body, dtype=int
+    )[indices_of_enabled_collidable_points]
+
+    L_p_Ci = model.kin_dyn_parameters.contact_parameters.point[
+        indices_of_enabled_collidable_points
+    ]
+
+    if len(indices_of_enabled_collidable_points) == 0:
+        return jnp.array(0).astype(float), jnp.empty(0).astype(float)
+
+    W_p_B, W_Q_B, s, W_v_WB, ṡ, _, _, _, _, _ = utils.process_inputs(
+        model=model,
+        base_position=base_position,
+        base_quaternion=base_quaternion,
+        joint_positions=joint_positions,
+        base_linear_velocity=base_linear_velocity,
+        base_angular_velocity=base_angular_velocity,
+        joint_velocities=joint_velocities,
+    )
+
+    # Get the parent array λ(i).
+    # Note: λ(0) must not be used, it's initialized to -1.
+    λ = model.kin_dyn_parameters.parent_array
+
+    # Compute the base transform.
+    W_H_B = jaxlie.SE3.from_rotation_and_translation(
+        rotation=jaxlie.SO3(wxyz=W_Q_B),
+        translation=W_p_B,
+    )
+
+    # Compute the parent-to-child adjoints and the motion subspaces of the joints.
+    # These transforms define the relative kinematics of the entire model, including
+    # the base transform for both floating-base and fixed-base models.
+    i_X_λi, S = model.kin_dyn_parameters.joint_transforms_and_motion_subspaces(
+        joint_positions=s, base_transform=W_H_B.as_matrix()
+    )
+
+    # Allocate buffer of transforms world -> link and initialize the base pose.
+    W_X_i = jnp.zeros(shape=(model.number_of_links(), 6, 6))
+    W_X_i = W_X_i.at[0].set(Adjoint.inverse(i_X_λi[0]))
+
+    # Allocate buffer of 6D inertial-fixed velocities and initialize the base velocity.
+    W_v_Wi = jnp.zeros(shape=(model.number_of_links(), 6))
+    W_v_Wi = W_v_Wi.at[0].set(W_v_WB)
+
+    # ====================
+    # Propagate kinematics
+    # ====================
+
+    PropagateTransformsCarry = tuple[jtp.Matrix, jtp.Matrix]
+    propagate_transforms_carry: PropagateTransformsCarry = (W_X_i, W_v_Wi)
+
+    def propagate_kinematics(
+        carry: PropagateTransformsCarry, i: jtp.Int
+    ) -> tuple[PropagateTransformsCarry, None]:
+
+        ii = i - 1
+        W_X_i, W_v_Wi = carry
+
+        # Compute the parent to child 6D transform.
+        λi_X_i = Adjoint.inverse(adjoint=i_X_λi[i])
+
+        # Compute the world to child 6D transform.
+        W_Xi_i = W_X_i[λ[i]] @ λi_X_i
+        W_X_i = W_X_i.at[i].set(W_Xi_i)
+
+        # Propagate the 6D velocity.
+        W_vi_Wi = W_v_Wi[λ[i]] + W_X_i[i] @ (S[i] * ṡ[ii]).squeeze()
+        W_v_Wi = W_v_Wi.at[i].set(W_vi_Wi)
+
+        return (W_X_i, W_v_Wi), None
+
+    (W_X_i, W_v_Wi), _ = (
+        jax.lax.scan(
+            f=propagate_kinematics,
+            init=propagate_transforms_carry,
+            xs=jnp.arange(start=1, stop=model.number_of_links()),
+        )
+        if model.number_of_links() > 1
+        else [(W_X_i, W_v_Wi), None]
+    )
+
+    # ==================================================
+    # Compute position and velocity of collidable points
+    # ==================================================
+
+    def process_point_kinematics(
+        Li_p_C: jtp.Vector, parent_body: jtp.Int
+    ) -> tuple[jtp.Vector, jtp.Vector]:
+
+        # Compute the position of the collidable point.
+        W_p_Ci = (
+            Adjoint.to_transform(adjoint=W_X_i[parent_body]) @ jnp.hstack([Li_p_C, 1])
+        )[0:3]
+
+        # Compute the linear part of the mixed velocity Ci[W]_v_{W,Ci}.
+        CW_vl_WCi = (
+            jnp.block([jnp.eye(3), -Skew.wedge(vector=W_p_Ci).squeeze()])
+            @ W_v_Wi[parent_body].squeeze()
+        )
+
+        return W_p_Ci, CW_vl_WCi
+
+    # Process all the collidable points in parallel.
+    W_p_Ci, CW_vl_WC = jax.vmap(process_point_kinematics)(
+        L_p_Ci,
+        parent_link_idx_of_enabled_collidable_points,
+    )
+
+    return W_p_Ci, CW_vl_WC

jaxsim/rbda/contacts/__init__.py

@@ -0,0 +1,13 @@
+from . import relaxed_rigid, rigid, soft, visco_elastic
+from .common import ContactModel, ContactsParams
+from .relaxed_rigid import RelaxedRigidContacts, RelaxedRigidContactsParams
+from .rigid import RigidContacts, RigidContactsParams
+from .soft import SoftContacts, SoftContactsParams
+from .visco_elastic import ViscoElasticContacts, ViscoElasticContactsParams
+
+ContactParamsTypes = (
+    SoftContactsParams
+    | RigidContactsParams
+    | RelaxedRigidContactsParams
+    | ViscoElasticContactsParams
+)

jaxsim/rbda/contacts/common.py

@@ -0,0 +1,313 @@
+from __future__ import annotations
+
+import abc
+import functools
+
+import jax
+import jax.numpy as jnp
+
+import jaxsim.api as js
+import jaxsim.terrain
+import jaxsim.typing as jtp
+from jaxsim.api.common import ModelDataWithVelocityRepresentation
+from jaxsim.utils import JaxsimDataclass
+
+try:
+    from typing import Self
+except ImportError:
+    from typing_extensions import Self
+
+
+@functools.partial(jax.jit, static_argnames=("terrain",))
+def compute_penetration_data(
+    p: jtp.VectorLike,
+    v: jtp.VectorLike,
+    terrain: jaxsim.terrain.Terrain,
+) -> tuple[jtp.Float, jtp.Float, jtp.Vector]:
+    """
+    Compute the penetration data (depth, rate, and terrain normal) of a collidable point.
+
+    Args:
+        p: The position of the collidable point.
+        v:
+            The linear velocity of the point (linear component of the mixed 6D velocity
+            of the implicit frame `C = (W_p_C, [W])` associated to the point).
+        terrain: The considered terrain.
+
+    Returns:
+        A tuple containing the penetration depth, the penetration velocity,
+        and the considered terrain normal.
+    """
+
+    # Pre-process the position and the linear velocity of the collidable point.
+    W_ṗ_C = jnp.array(v).squeeze()
+    px, py, pz = jnp.array(p).squeeze()
+
+    # Compute the terrain normal and the contact depth.
+    n̂ = terrain.normal(x=px, y=py).squeeze()
+    h = jnp.array([0, 0, terrain.height(x=px, y=py) - pz])
+
+    # Compute the penetration depth normal to the terrain.
+    δ = jnp.maximum(0.0, jnp.dot(h, n̂))
+
+    # Compute the penetration normal velocity.
+    δ_dot = -jnp.dot(W_ṗ_C, n̂)
+
+    # Enforce the penetration rate to be zero when the penetration depth is zero.
+    δ_dot = jnp.where(δ > 0, δ_dot, 0.0)
+
+    return δ, δ_dot, n̂
+
+
+class ContactsParams(JaxsimDataclass):
+    """
+    Abstract class representing the parameters of a contact model.
+
+    Note:
+        This class is supposed to store only the tunable parameters of the contact
+        model, i.e. all those parameters that can be changed during runtime.
+        If the contact model has also static parameters, they should be stored
+        in the corresponding `ContactModel` class.
+    """
+
+    @classmethod
+    @abc.abstractmethod
+    def build(cls: type[Self], **kwargs) -> Self:
+        """
+        Create a `ContactsParams` instance with specified parameters.
+
+        Returns:
+            The `ContactsParams` instance.
+        """
+        pass
+
+    @abc.abstractmethod
+    def valid(self, **kwargs) -> jtp.BoolLike:
+        """
+        Check if the parameters are valid.
+
+        Returns:
+            True if the parameters are valid, False otherwise.
+        """
+        pass
+
+
+class ContactModel(JaxsimDataclass):
+    """
+    Abstract class representing a contact model.
+    """
+
+    @classmethod
+    @abc.abstractmethod
+    def build(
+        cls: type[Self],
+        **kwargs,
+    ) -> Self:
+        """
+        Create a `ContactModel` instance with specified parameters.
+
+        Returns:
+            The `ContactModel` instance.
+        """
+
+        pass
+
+    @abc.abstractmethod
+    def compute_contact_forces(
+        self,
+        model: js.model.JaxSimModel,
+        data: js.data.JaxSimModelData,
+        **kwargs,
+    ) -> tuple[jtp.Matrix, dict[str, jtp.PyTree]]:
+        """
+        Compute the contact forces.
+
+        Args:
+            model: The robot model considered by the contact model.
+            data: The data of the considered model.
+            **kwargs: Optional additional arguments, specific to the contact model.
+
+        Returns:
+            A tuple containing as first element the computed 6D contact force applied to
+            the contact points and expressed in the world frame, and as second element
+            a dictionary of optional additional information.
+        """
+
+        pass
+
+    def compute_link_contact_forces(
+        self,
+        model: js.model.JaxSimModel,
+        data: js.data.JaxSimModelData,
+        **kwargs,
+    ) -> tuple[jtp.Matrix, dict[str, jtp.PyTree]]:
+        """
+        Compute the link contact forces.
+
+        Args:
+            model: The robot model considered by the contact model.
+            data: The data of the considered model.
+            **kwargs: Optional additional arguments, specific to the contact model.
+
+        Returns:
+            A tuple containing as first element the 6D contact force applied to the
+            links and expressed in the frame of the velocity representation of data,
+            and as second element a dictionary of optional additional information.
+        """
+
+        # Compute the contact forces expressed in the inertial frame.
+        # This function, contrarily to `compute_contact_forces`, already handles how
+        # the optional kwargs should be passed to the specific contact models.
+        W_f_C, aux_dict = js.contact.collidable_point_dynamics(
+            model=model, data=data, **kwargs
+        )
+
+        # Compute the 6D forces applied to the links equivalent to the forces applied
+        # to the frames associated to the collidable points.
+        with data.switch_velocity_representation(jaxsim.VelRepr.Inertial):
+
+            W_f_L = self.link_forces_from_contact_forces(
+                model=model, data=data, contact_forces=W_f_C
+            )
+
+        # Store the link forces in the references object for easy conversion.
+        references = js.references.JaxSimModelReferences.build(
+            model=model,
+            data=data,
+            link_forces=W_f_L,
+            velocity_representation=jaxsim.VelRepr.Inertial,
+        )
+
+        # Convert the link forces to the frame corresponding to the velocity
+        # representation of data.
+        with references.switch_velocity_representation(data.velocity_representation):
+            f_L = references.link_forces(model=model, data=data)
+
+        return f_L, aux_dict
+
+    @staticmethod
+    def link_forces_from_contact_forces(
+        model: js.model.JaxSimModel,
+        data: js.data.JaxSimModelData,
+        *,
+        contact_forces: jtp.MatrixLike,
+    ) -> jtp.Matrix:
+        """
+        Compute the link forces from the contact forces.
+
+        Args:
+            model: The robot model considered by the contact model.
+            data: The data of the considered model.
+            contact_forces: The contact forces computed by the contact model.
+
+        Returns:
+            The 6D contact forces applied to the links and expressed in the frame of
+            the velocity representation of data.
+        """
+
+        # Get the object storing the contact parameters of the model.
+        contact_parameters = model.kin_dyn_parameters.contact_parameters
+
+        # Extract the indices corresponding to the enabled collidable points.
+        indices_of_enabled_collidable_points = (
+            contact_parameters.indices_of_enabled_collidable_points
+        )
+
+        # Convert the contact forces to a JAX array.
+        f_C = jnp.atleast_2d(jnp.array(contact_forces, dtype=float).squeeze())
+
+        # Get the pose of the enabled collidable points.
+        W_H_C = js.contact.transforms(model=model, data=data)[
+            indices_of_enabled_collidable_points
+        ]
+
+        # Convert the contact forces to inertial-fixed representation.
+        W_f_C = jax.vmap(
+            lambda f_C, W_H_C: (
+                ModelDataWithVelocityRepresentation.other_representation_to_inertial(
+                    array=f_C,
+                    other_representation=data.velocity_representation,
+                    transform=W_H_C,
+                    is_force=True,
+                )
+            )
+        )(f_C, W_H_C)
+
+        # Construct the vector defining the parent link index of each collidable point.
+        # We use this vector to sum the 6D forces of all collidable points rigidly
+        # attached to the same link.
+        parent_link_index_of_collidable_points = jnp.array(
+            contact_parameters.body, dtype=int
+        )[indices_of_enabled_collidable_points]
+
+        # Create the mask that associate each collidable point to their parent link.
+        # We use this mask to sum the collidable points to the right link.
+        mask = parent_link_index_of_collidable_points[:, jnp.newaxis] == jnp.arange(
+            model.number_of_links()
+        )
+
+        # Sum the forces of all collidable points rigidly attached to a body.
+        # Since the contact forces W_f_C are expressed in the world frame,
+        # we don't need any coordinate transformation.
+        W_f_L = mask.T @ W_f_C
+
+        # Compute the link transforms.
+        W_H_L = (
+            js.model.forward_kinematics(model=model, data=data)
+            if data.velocity_representation is not jaxsim.VelRepr.Inertial
+            else jnp.zeros(shape=(model.number_of_links(), 4, 4))
+        )
+
+        # Convert the inertial-fixed link forces to the velocity representation of data.
+        f_L = jax.vmap(
+            lambda W_f_L, W_H_L: (
+                ModelDataWithVelocityRepresentation.inertial_to_other_representation(
+                    array=W_f_L,
+                    other_representation=data.velocity_representation,
+                    transform=W_H_L,
+                    is_force=True,
+                )
+            )
+        )(W_f_L, W_H_L)
+
+        return f_L
+
+    @classmethod
+    def zero_state_variables(cls, model: js.model.JaxSimModel) -> dict[str, jtp.Array]:
+        """
+        Build zero state variables of the contact model.
+
+        Args:
+            model: The robot model considered by the contact model.
+
+        Note:
+            There are contact models that require to extend the state vector of the
+            integrated ODE system with additional variables. Our integrators are
+            capable of operating on a generic state, as long as it is a PyTree.
+            This method builds the zero state variables of the contact model as a
+            dictionary of JAX arrays.
+
+        Returns:
+            A dictionary storing the zero state variables of the contact model.
+        """
+
+        return {}
+
+    @property
+    def _parameters_class(cls) -> type[ContactsParams]:
+        """
+        Return the class of the contact parameters.
+
+        Returns:
+            The class of the contact parameters.
+        """
+        import importlib
+
+        return getattr(
+            importlib.import_module("jaxsim.rbda.contacts"),
+            (
+                cls.__name__ + "Params"
+                if isinstance(cls, type)
+                else cls.__class__.__name__ + "Params"
+            ),
+        )