jaxsim 0.7.1.dev49__py3-none-any.whl → 0.7.1.dev53__py3-none-any.whl

jaxsim/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.7.1.dev49'
-__version_tuple__ = version_tuple = (0, 7, 1, 'dev49')
+__version__ = version = '0.7.1.dev53'
+__version_tuple__ = version_tuple = (0, 7, 1, 'dev53')

jaxsim/api/contact.py

@@ -191,39 +191,22 @@ def estimate_good_contact_parameters(
         The user is encouraged to fine-tune the parameters based on the
         specific application.
     """
-
-    def estimate_model_height(model: js.model.JaxSimModel) -> jtp.Float:
-        """
-        Displacement between the CoM and the lowest collidable point using zero
-        joint positions.
-        """
-
-        zero_data = js.data.JaxSimModelData.build(
-            model=model,
-        )
-
+    if max_penetration is None:
+        zero_data = js.data.JaxSimModelData.build(model=model)
         W_pz_CoM = js.com.com_position(model=model, data=zero_data)[2]
-
         if model.floating_base():
             W_pz_C = collidable_point_positions(model=model, data=zero_data)[:, -1]
-            return 2 * (W_pz_CoM - W_pz_C.min())
-
-        return 2 * W_pz_CoM
+            W_pz_CoM = W_pz_CoM - W_pz_C.min()
 
-    max_δ = (
-        max_penetration
-        if max_penetration is not None
-        # Consider as default a 0.5% of the model height.
-        else 0.005 * estimate_model_height(model=model)
-    )
+        # Consider as default a 1% of the model center of mass height.
+        max_penetration = 0.01 * W_pz_CoM
 
     nc = number_of_active_collidable_points_steady_state
-
     return model.contact_model._parameters_class().build_default_from_jaxsim_model(
         model=model,
         standard_gravity=standard_gravity,
         static_friction_coefficient=static_friction_coefficient,
-        max_penetration=max_δ,
+        max_penetration=max_penetration,
         number_of_active_collidable_points_steady_state=nc,
         damping_ratio=damping_ratio,
     )
@@ -569,6 +552,39 @@ def link_contact_forces(
     # to the frames associated to the collidable points.
     W_f_L = link_forces_from_contact_forces(model=model, contact_forces=W_f_C)
 
+    # Process constraint wrenches if present.
+    if "constr_wrenches_inertial" in aux_dict:
+        wrench_pair_constr_inertial = aux_dict["constr_wrenches_inertial"]
+
+        # Retrieve the constraint map from the model's kinematic parameters.
+        constraint_map = model.kin_dyn_parameters.constraints
+
+        # Extract the frame indices of the constraints.
+        frame_idxs_1 = constraint_map.frame_idxs_1
+        frame_idxs_2 = constraint_map.frame_idxs_2
+
+        n_kin_constraints = frame_idxs_1.shape[0]
+
+        if n_kin_constraints > 0:
+            parent_link_indices = jax.vmap(
+                lambda frame_idx_1, frame_idx_2: jnp.array(
+                    (
+                        js.frame.idx_of_parent_link(model, frame_index=frame_idx_1),
+                        js.frame.idx_of_parent_link(model, frame_index=frame_idx_2),
+                    )
+                )
+            )(frame_idxs_1, frame_idxs_2)
+
+            # Apply each constraint wrench to its corresponding parent link in W_f_L.
+            def apply_wrench(W_f_L, parent_indices, wrench_pair):
+                W_f_L = W_f_L.at[parent_indices[0]].add(wrench_pair[0])
+                W_f_L = W_f_L.at[parent_indices[1]].add(wrench_pair[1])
+                return W_f_L
+
+            W_f_L = jax.vmap(apply_wrench, in_axes=(None, 0, 0))(
+                W_f_L, parent_link_indices, wrench_pair_constr_inertial
+            ).sum(axis=0)
+
     return W_f_L, aux_dict
 
 

jaxsim/api/kin_dyn_parameters.py

@@ -34,6 +34,7 @@ class KinDynParameters(JaxsimDataclass):
         joint_model: The joint model of the model.
         joint_parameters: The parameters of the joints.
         hw_link_metadata: The hardware parameters of the model links.
+        constraints: The kinematic constraints of the model. They can be used only with Relaxed-Rigid contact model.
     """
 
     # Static
@@ -58,6 +59,9 @@ class KinDynParameters(JaxsimDataclass):
     # Model hardware parameters
     hw_link_metadata: HwLinkMetadata | None = dataclasses.field(default=None)
 
+    # Kinematic constraints
+    constraints: ConstraintMap | None = dataclasses.field(default=None)
+
     @property
     def motion_subspaces(self) -> jtp.Matrix:
         r"""
@@ -80,12 +84,15 @@ class KinDynParameters(JaxsimDataclass):
         return self._support_body_array_bool.get()
 
     @staticmethod
-    def build(model_description: ModelDescription) -> KinDynParameters:
+    def build(
+        model_description: ModelDescription, constraints: ConstraintMap | None
+    ) -> KinDynParameters:
         """
         Construct the kinematic and dynamic parameters of the model.
 
         Args:
             model_description: The parsed model description to consider.
+            constraints: An object of type ConstraintMap specifying the kinematic constraint of the model.
 
         Returns:
             The kinematic and dynamic parameters of the model.
@@ -253,6 +260,12 @@ class KinDynParameters(JaxsimDataclass):
 
         motion_subspaces = jnp.vstack([jnp.zeros((6, 1))[jnp.newaxis, ...], S_J])
 
+        # ===========
+        # Constraints
+        # ===========
+
+        constraints = ConstraintMap() if constraints is None else constraints
+
         # =================================
         # Build and return KinDynParameters
         # =================================
@@ -267,6 +280,7 @@ class KinDynParameters(JaxsimDataclass):
             joint_parameters=joint_parameters,
             contact_parameters=contact_parameters,
             frame_parameters=frame_parameters,
+            constraints=constraints,
         )
 
     def __eq__(self, other: KinDynParameters) -> bool:
@@ -1168,3 +1182,85 @@ class ScalingFactors(JaxsimDataclass):
 
     dims: jtp.Vector
     density: jtp.Float
+
+
+@dataclasses.dataclass(frozen=True)
+class ConstraintType:
+    """
+    Enumeration of all supported constraint types.
+    """
+
+    Weld: ClassVar[int] = 0
+    # TODO: handle Connect constraint
+    # Connect: ClassVar[int] = 1
+
+
+@jax_dataclasses.pytree_dataclass
+class ConstraintMap(JaxsimDataclass):
+    """
+    Class storing the kinematic constraints of a model.
+    """
+
+    frame_idxs_1: jtp.Int = dataclasses.field(
+        default_factory=lambda: jnp.array([], dtype=int)
+    )
+    frame_idxs_2: jtp.Int = dataclasses.field(
+        default_factory=lambda: jnp.array([], dtype=int)
+    )
+    constraint_types: jtp.Int = dataclasses.field(
+        default_factory=lambda: jnp.array([], dtype=int)
+    )
+    K_P: jtp.Float = dataclasses.field(
+        default_factory=lambda: jnp.array([], dtype=float)
+    )
+    K_D: jtp.Float = dataclasses.field(
+        default_factory=lambda: jnp.array([], dtype=float)
+    )
+
+    def add_constraint(
+        self,
+        frame_idx_1: int,
+        frame_idx_2: int,
+        constraint_type: int,
+        K_P: float | None = None,
+        K_D: float | None = None,
+    ) -> ConstraintMap:
+        """
+        Add a constraint to the constraint map.
+
+        Args:
+            frame_idx_1: The index of the first frame.
+            frame_idx_2: The index of the second frame.
+            constraint_type: The type of constraint.
+            K_P: The proportional gain for Baumgarte stabilization (default: 1000).
+            K_D: The derivative gain for Baumgarte stabilization (default: 2 * sqrt(K_P)).
+
+        Returns:
+            A new ConstraintMap instance with the added constraint.
+
+        Note:
+            Since this method returns a new instance of ConstraintMap with the new constraint,
+            it will trigger recompilations in JIT-compiled functions.
+        """
+
+        # Set default values for Baumgarte coefficients if not provided
+        if K_P is None:
+            K_P = 1000
+        if K_D is None:
+            K_D = 2 * np.sqrt(K_P)
+
+        # Create new arrays with the input elements appended
+        new_frame_idxs_1 = jnp.append(self.frame_idxs_1, frame_idx_1)
+        new_frame_idxs2 = jnp.append(self.frame_idxs_2, frame_idx_2)
+        new_constraint_types = jnp.append(self.constraint_types, constraint_type)
+        new_K_P = jnp.append(self.K_P, K_P)
+        new_K_D = jnp.append(self.K_D, K_D)
+
+        # Return a new ConstraintMap object with updated attributes
+        return ConstraintMap(
+            frame_idxs_1=new_frame_idxs_1,
+            frame_idxs_2=new_frame_idxs2,
+            constraint_types=new_constraint_types,
+            K_P=new_K_P,
+            K_D=new_K_D,
+        )

jaxsim/api/model.py

@@ -141,6 +141,7 @@ class JaxSimModel(JaxsimDataclass):
         is_urdf: bool | None = None,
         considered_joints: Sequence[str] | None = None,
         gravity: jtp.FloatLike = jaxsim.math.STANDARD_GRAVITY,
+        constraints: jaxsim.rbda.kinematic_constraints.ConstraintMap | None = None,
     ) -> JaxSimModel:
         """
         Build a Model object from a model description.
@@ -167,6 +168,9 @@ class JaxSimModel(JaxsimDataclass):
             considered_joints:
                 The list of joints to consider. If None, all joints are considered.
             gravity: The gravity constant. Normally passed as a positive value.
+            constraints:
+                An object of type ConstraintMap containing the kinematic constraints to consider. If None, no constraints are considered.
+                Note that constraints can be used only with RelaxedRigidContacts.
 
         Returns:
             The built Model object.
@@ -198,6 +202,7 @@ class JaxSimModel(JaxsimDataclass):
             contact_params=contact_params,
             integrator=integrator,
             gravity=-gravity,
+            constraints=constraints,
         )
 
         # Store the origin of the model, in case downstream logic needs it.
@@ -225,6 +230,7 @@ class JaxSimModel(JaxsimDataclass):
         actuation_params: jaxsim.rbda.actuation.ActuationParams | None = None,
         integrator: IntegratorType | None = None,
         gravity: jtp.FloatLike = jaxsim.math.STANDARD_GRAVITY,
+        constraints: jaxsim.rbda.kinematic_constraints.ConstraintMap | None = None,
     ) -> JaxSimModel:
         """
         Build a Model object from an intermediate model description.
@@ -247,6 +253,8 @@ class JaxSimModel(JaxsimDataclass):
             actuation_params: The parameters of the actuation model.
             integrator: The integrator to use for the simulation.
             gravity: The gravity constant.
+            constraints:
+                An object of type ConstraintMap containing the kinematic constraints to consider. If None, no constraints are considered.
 
         Returns:
             The built Model object.
@@ -278,6 +286,14 @@ class JaxSimModel(JaxsimDataclass):
             else jaxsim.rbda.contacts.RelaxedRigidContacts.build()
         )
 
+        if constraints is not None and not isinstance(
+            contact_model, jaxsim.rbda.contacts.RelaxedRigidContacts
+        ):
+            constraints = None
+            logging.warning(
+                f"Contact model {contact_model.__class__.__name__} does not support kinematic constraints. Use RelaxedRigidContacts instead."
+            )
+
         if contact_params is None:
             contact_params = contact_model._parameters_class()
 
@@ -295,7 +311,7 @@ class JaxSimModel(JaxsimDataclass):
         model = cls(
             model_name=model_name,
             kin_dyn_parameters=js.kin_dyn_parameters.KinDynParameters.build(
-                model_description=model_description
+                model_description=model_description, constraints=constraints
             ),
             time_step=time_step,
             terrain=terrain,
@@ -777,6 +793,7 @@ def reduce(
         actuation_params=model.actuation_params,
         gravity=model.gravity,
         integrator=model.integrator,
+        constraints=model.kin_dyn_parameters.constraints,
     )
 
     with reduced_model.mutable_context(mutability=Mutability.MUTABLE_NO_VALIDATION):
@@ -2330,7 +2347,11 @@ def update_hw_parameters(
 
     Returns:
         The updated JaxSimModel object with modified hardware parameters.
+
+    Note:
+        This function can be used only with models using Relax-Rigid contact model.
     """
+
     kin_dyn_params: KinDynParameters = model.kin_dyn_parameters
     link_parameters: LinkParameters = kin_dyn_params.link_parameters
     hw_link_metadata: HwLinkMetadata = kin_dyn_params.hw_link_metadata

jaxsim/math/rotation.py

@@ -82,3 +82,17 @@ class Rotation:
         R = c * jnp.eye(3) - s * Skew.wedge(u) + c1 * u @ u.T
 
         return R.transpose()
+
+    @staticmethod
+    def log_vee(R: jnp.ndarray) -> jtp.Vector:
+        """
+        Compute the logarithm map of an SO(3) rotation matrix.
+
+        Args:
+            R: The SO(3) rotation matrix.
+
+        Returns:
+            The corresponding so(3) tangent vector.
+        """
+
+        return jaxlie.SO3.from_matrix(R).log()

jaxsim/mujoco/loaders.py

@@ -5,6 +5,7 @@ import warnings
 from collections.abc import Sequence
 from typing import Any
 
+import jaxlie
 import mujoco as mj
 import numpy as np
 import rod.urdf.exporter
@@ -279,11 +280,13 @@ class RodModelToMjcf:
         # Add a floating joint if floating-base
         # -------------------------------------
 
+        base_link_name = rod_model.get_canonical_link()
+
         if not rod_model.is_fixed_base():
             considered_joints |= {"world_to_base"}
             urdf_string = RodModelToMjcf.add_floating_joint(
                 urdf_string=urdf_string,
-                base_link_name=rod_model.get_canonical_link(),
+                base_link_name=base_link_name,
                 floating_joint_name="world_to_base",
             )
 
@@ -379,6 +382,30 @@ class RodModelToMjcf:
         # Find the <mujoco> element (might be the root itself).
         mujoco_element: ET._Element = next(iter(root.iter("mujoco")))
 
+        # --------------
+        # Add the frames
+        # --------------
+
+        for frame in rod_model.frames():
+            frame: rod.Frame
+            parent_name = frame.attached_to
+            parent_element = mujoco_element.find(f".//body[@name='{parent_name}']")
+
+            if parent_element is None and parent_name == base_link_name:
+                parent_element = mujoco_element.find(".//worldbody")
+
+            if parent_element is not None:
+                quat = jaxlie.SO3.from_rpy_radians(*frame.pose.rpy).wxyz
+                _ = ET.SubElement(
+                    parent_element,
+                    "site",
+                    name=frame.name,
+                    pos=" ".join(map(str, frame.pose.xyz)),
+                    quat=" ".join(map(str, quat)),
+                )
+            else:
+                warnings.warn(f"Parent link '{parent_name}' not found", stacklevel=2)
+
         # --------------
         # Add the motors
         # --------------

jaxsim/rbda/__init__.py

@@ -8,4 +8,10 @@ from .jacobian import (
     jacobian_derivative_full_doubly_left,
     jacobian_full_doubly_left,
 )
+from .kinematic_constraints import (
+    compute_constraint_baumgarte_term,
+    compute_constraint_jacobians,
+    compute_constraint_jacobians_derivative,
+    compute_constraint_transforms,
+)
 from .rnea import rnea

jaxsim/rbda/contacts/common.py

@@ -18,6 +18,10 @@ except ImportError:
     from typing_extensions import Self
 
 
+MAX_STIFFNESS = 1e6
+MAX_DAMPING = 1e4
+
+
 @functools.partial(jax.jit, static_argnames=("terrain",))
 def compute_penetration_data(
     p: jtp.VectorLike,
@@ -133,28 +137,30 @@ class ContactsParams(JaxsimDataclass):
         ξ = damping_ratio
         δ_max = max_penetration
         μc = static_friction_coefficient
+        nc = number_of_active_collidable_points_steady_state
 
         # Compute the total mass of the model.
         m = jnp.array(model.kin_dyn_parameters.link_parameters.mass).sum()
 
-        # Rename the standard gravity.
-        g = standard_gravity
-
-        # Compute the average support force on each collidable point.
-        f_average = m * g / number_of_active_collidable_points_steady_state
-
         # Compute the stiffness to get the desired steady-state penetration.
         # Note that this is dependent on the non-linear exponent used in
         # the damping term of the Hunt/Crossley model.
-        K = f_average / jnp.power(δ_max, 1 + p) if stiffness is None else stiffness
+        if stiffness is None:
+            # Compute the average support force on each collidable point.
+            f_average = m * standard_gravity / nc
+
+            stiffness = f_average / jnp.power(δ_max, 1 + p)
+            stiffness = jnp.clip(stiffness, 0, MAX_STIFFNESS)
 
         # Compute the damping using the damping ratio.
-        critical_damping = 2 * jnp.sqrt(K * m)
-        D = ξ * critical_damping if damping is None else damping
+        critical_damping = 2 * jnp.sqrt(stiffness * m)
+        if damping is None:
+            damping = ξ * critical_damping
+            damping = jnp.clip(damping, 0, MAX_DAMPING)
 
         return self.build(
-            K=K,
-            D=D,
+            K=stiffness,
+            D=damping,
             mu=μc,
             p=p,
             q=q,

jaxsim/rbda/contacts/relaxed_rigid.py

@@ -12,6 +12,13 @@ import optax
 import jaxsim.api as js
 import jaxsim.typing as jtp
 from jaxsim.api.common import ModelDataWithVelocityRepresentation, VelRepr
+from jaxsim.rbda.kinematic_constraints import (
+    ConstraintMap,
+    compute_constraint_baumgarte_term,
+    compute_constraint_jacobians,
+    compute_constraint_jacobians_derivative,
+    compute_constraint_transforms,
+)
 
 from . import common, soft
 
@@ -333,10 +340,22 @@ class RelaxedRigidContacts(common.ContactModel):
         # Compute the position in the constraint frame.
         position_constraint = jax.vmap(lambda δ, n̂: -δ * n̂)(δ, n̂)
 
+        # Compute the regularization terms.
+        a_ref, r, *_ = self._regularizers(
+            model=model,
+            position_constraint=position_constraint,
+            velocity_constraint=velocity,
+            parameters=model.contact_params,
+        )
+
         # Compute the transforms of the implicit frames corresponding to the
         # collidable points.
         W_H_C = js.contact.transforms(model=model, data=data)
 
+        # Retrieve the kinematic constraints, if any.
+        kin_constraints: ConstraintMap = model.kin_dyn_parameters.constraints
+        n_kin_constraints: int = 6 * kin_constraints.frame_idxs_1.shape[0]
+
         with (
             data.switch_velocity_representation(VelRepr.Mixed),
             references.switch_velocity_representation(VelRepr.Mixed),
@@ -354,30 +373,103 @@ class RelaxedRigidContacts(common.ContactModel):
 
             M = js.model.free_floating_mass_matrix(model=model, data=data)
 
+            # Compute the linear part of the Jacobian of the collidable points
             Jl_WC = jnp.vstack(
                 jax.vmap(lambda J, δ: J * (δ > 0))(
                     js.contact.jacobian(model=model, data=data)[:, :3, :], δ
                 )
             )
 
-            J̇_WC = jnp.vstack(
+            # Compute the linear part of the Jacobian derivative of the collidable points
+            J̇l_WC = jnp.vstack(
                 jax.vmap(lambda J̇, δ: J̇ * (δ > 0))(
                     js.contact.jacobian_derivative(model=model, data=data)[:, :3], δ
                 ),
             )
 
-        # Compute the regularization terms.
-        a_ref, R, *_ = self._regularizers(
-            model=model,
-            position_constraint=position_constraint,
-            velocity_constraint=velocity,
-            parameters=model.contact_params,
-        )
+        # Check if there are any kinematic constraints
+        if n_kin_constraints > 0:
+            with (
+                data.switch_velocity_representation(VelRepr.Mixed),
+                references.switch_velocity_representation(VelRepr.Mixed),
+            ):
+                J_constr = jax.vmap(
+                    compute_constraint_jacobians, in_axes=(None, None, 0)
+                )(model, data, kin_constraints)
+
+                J̇_constr = jax.vmap(
+                    compute_constraint_jacobians_derivative, in_axes=(None, None, 0)
+                )(model, data, kin_constraints)
+
+                W_H_constr_pairs = jax.vmap(
+                    compute_constraint_transforms, in_axes=(None, None, 0)
+                )(model, data, kin_constraints)
+
+                constr_baumgarte_term = jnp.ravel(
+                    jax.vmap(
+                        compute_constraint_baumgarte_term,
+                        in_axes=(0, None, 0, 0),
+                    )(
+                        J_constr,
+                        BW_ν,
+                        W_H_constr_pairs,
+                        kin_constraints,
+                    ),
+                )
+
+                J_constr = jnp.vstack(J_constr)
+                J̇_constr = jnp.vstack(J̇_constr)
+
+                R = jnp.diag(jnp.hstack([r, jnp.zeros(n_kin_constraints)]))
+                a_ref = jnp.hstack([a_ref, -constr_baumgarte_term])
+
+                J = jnp.vstack([Jl_WC, J_constr])
+                J̇ = jnp.vstack([J̇l_WC, J̇_constr])
+
+        else:
+            R = jnp.diag(r)
+
+            J = Jl_WC
+            J̇ = J̇l_WC
+
+        # Compute the Delassus matrix for contacts (mixed representation).
+        G_contacts = Jl_WC @ jnp.linalg.pinv(M) @ Jl_WC.T
 
-        # Compute the Delassus matrix and the free mixed linear acceleration of
-        # the collidable points.
-        G = Jl_WC @ jnp.linalg.pinv(M) @ Jl_WC.T
-        CW_al_free_WC = Jl_WC @ BW_ν̇_free + J̇_WC @ BW_ν
+        # Compute the Delassus matrix for constraints (inertial representation) if constraints exist.
+        with data.switch_velocity_representation(VelRepr.Inertial):
+            if n_kin_constraints > 0:
+                G_constraints = (
+                    J_constr
+                    @ jnp.linalg.pinv(
+                        js.model.free_floating_mass_matrix(
+                            model=model,
+                            data=data,
+                        )
+                    )
+                    @ J_constr.T
+                )
+            else:
+                G_constraints = jnp.zeros((0, 0))
+
+        # Combine the Delassus matrices for contacts and constraints if constraints exist.
+        if G_constraints.shape[0] > 0:
+            G = jnp.block(
+                [
+                    [
+                        G_contacts,
+                        jnp.zeros((G_contacts.shape[0], G_constraints.shape[1])),
+                    ],
+                    [
+                        jnp.zeros((G_constraints.shape[0], G_contacts.shape[1])),
+                        G_constraints,
+                    ],
+                ]
+            )
+        else:
+            G = G_contacts
+
+        # Compute the free mixed linear acceleration of the collidable points.
+        CW_al_free_WC = J @ BW_ν̇_free + J̇ @ BW_ν
 
         # Calculate quantities for the linear optimization problem.
         A = G + R
@@ -469,6 +561,9 @@ class RelaxedRigidContacts(common.ContactModel):
             )[0]
         )(position, velocity).flatten()
 
+        if n_kin_constraints > 0:
+            init_params = jnp.hstack([init_params, jnp.zeros(n_kin_constraints)])
+
         # Get the solver options.
         solver_options = self.solver_options
 
@@ -494,8 +589,20 @@ class RelaxedRigidContacts(common.ContactModel):
             has_aux=True,
         )
 
-        # Reshape the optimized solution to be a matrix of 3D contact forces.
-        CW_fl_C = solution.reshape(-1, 3)
+        if n_kin_constraints > 0:
+            # Extract the last n_kin_constr values from the solution and split them into 6D wrenches
+            kin_constr_wrench_inertial = solution[-n_kin_constraints:].reshape(-1, 6)
+
+            # Form an array of tuples with each wrench and its opposite using jax constructs
+            kin_constr_wrench_pairs_inertial = jnp.stack(
+                (kin_constr_wrench_inertial, -kin_constr_wrench_inertial), axis=1
+            )
+
+            # Reshape the optimized solution to be a matrix of 3D contact forces.
+            CW_fl_C = solution[0:-n_kin_constraints].reshape(-1, 3)
+        else:
+            kin_constr_wrench_pairs_inertial = jnp.zeros((0, 2, 6))
+            CW_fl_C = solution.reshape(-1, 3)
 
         # Convert the contact forces from mixed to inertial-fixed representation.
         W_f_C = ModelDataWithVelocityRepresentation.other_representation_to_inertial(
@@ -505,7 +612,9 @@ class RelaxedRigidContacts(common.ContactModel):
             is_force=True,
         )
 
-        return W_f_C, {}
+        return W_f_C, {
+            "constr_wrenches_inertial": kin_constr_wrench_pairs_inertial,
+        }
 
     @staticmethod
     def _regularizers(
@@ -635,4 +744,4 @@ class RelaxedRigidContacts(common.ContactModel):
             ),
         )
 
-        return a_ref, jnp.diag(R), K, D
+        return a_ref, R, K, D

jaxsim/rbda/kinematic_constraints.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+import jax.numpy as jnp
+
+import jaxsim.api as js
+import jaxsim.typing as jtp
+from jaxsim.api.common import VelRepr
+from jaxsim.api.kin_dyn_parameters import ConstraintMap
+from jaxsim.math.rotation import Rotation
+
+
+def compute_constraint_jacobians(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    constraint: ConstraintMap,
+) -> jtp.Matrix:
+    """
+    Compute the constraint Jacobian matrix representing the kinematic constraints between two frames.
+
+    Args:
+        model: The JaxSim model.
+        data: The data of the considered model.
+        constraint: The considered constraint.
+
+    Returns:
+        The resulting constraint Jacobian matrix representing the kinematic constraint
+        between the two specified frames, in inertial representation.
+    """
+
+    J_WF1 = js.frame.jacobian(
+        model=model,
+        data=data,
+        frame_index=constraint.frame_idxs_1,
+        output_vel_repr=VelRepr.Inertial,
+    )
+    J_WF2 = js.frame.jacobian(
+        model=model,
+        data=data,
+        frame_index=constraint.frame_idxs_2,
+        output_vel_repr=VelRepr.Inertial,
+    )
+
+    return J_WF1 - J_WF2
+
+
+def compute_constraint_jacobians_derivative(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    constraint: ConstraintMap,
+) -> jtp.Matrix:
+    """
+    Compute the derivative of the constraint Jacobian matrix representing the kinematic constraints between two frames.
+
+    Args:
+        model: The JaxSim model.
+        data: The data of the considered model.
+        constraint: The considered constraint.
+
+    Returns:
+        The resulting constraint Jacobian derivative matrix representing the kinematic constraint
+        between the two specified frames, in inertial representation.
+    """
+
+    J̇_WF1 = js.frame.jacobian_derivative(
+        model=model,
+        data=data,
+        frame_index=constraint.frame_idxs_1,
+        output_vel_repr=VelRepr.Inertial,
+    )
+    J̇_WF2 = js.frame.jacobian_derivative(
+        model=model,
+        data=data,
+        frame_index=constraint.frame_idxs_2,
+        output_vel_repr=VelRepr.Inertial,
+    )
+
+    return J̇_WF1 - J̇_WF2
+
+
+def compute_constraint_baumgarte_term(
+    J_constr: jtp.Matrix,
+    nu: jtp.Vector,
+    W_H_F_constr: tuple[jtp.Matrix, jtp.Matrix],
+    constraint: ConstraintMap,
+) -> jtp.Vector:
+    """
+    Compute the Baumgarte stabilization term for kinematic constraints.
+
+    Args:
+        J_constr: The constraint Jacobian matrix.
+        nu: The generalized velocity vector.
+        W_H_F_constr: A tuple containing the homogeneous transformation matrices
+            of two frames (W_H_F1 and W_H_F2) with respect to the world frame.
+        K_P: The proportional gain for position and orientation error correction.
+        K_D: The derivative gain for velocity error correction.
+        constraint: The considered constraint.
+
+    Returns:
+        The computed Baumgarte stabilization term.
+    """
+    W_H_F1, W_H_F2 = W_H_F_constr
+
+    W_p_F1 = W_H_F1[0:3, 3]
+    W_p_F2 = W_H_F2[0:3, 3]
+
+    W_R_F1 = W_H_F1[0:3, 0:3]
+    W_R_F2 = W_H_F2[0:3, 0:3]
+
+    K_P = constraint.K_P
+    K_D = constraint.K_D
+
+    vel_error = J_constr @ nu
+    position_error = W_p_F1 - W_p_F2
+    R_error = W_R_F2.T @ W_R_F1
+    orientation_error = Rotation.log_vee(R_error)
+
+    baumgarte_term = (
+        K_P * jnp.concatenate([position_error, orientation_error]) + K_D * vel_error
+    )
+
+    return baumgarte_term
+
+
+def compute_constraint_transforms(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    constraint: ConstraintMap,
+) -> jtp.Matrix:
+    """
+    Compute the transformation matrices for a given kinematic constraint between two frames.
+
+    Args:
+        model: The JaxSim model.
+        data: The data of the considered model.
+        constraint: The considered constraint.
+
+    Returns:
+        A matrix containing the tuple of transformation matrices of the two frames.
+    """
+
+    W_H_F1 = js.frame.transform(
+        model=model, data=data, frame_index=constraint.frame_idxs_1
+    )
+    W_H_F2 = js.frame.transform(
+        model=model, data=data, frame_index=constraint.frame_idxs_2
+    )
+
+    return jnp.array((W_H_F1, W_H_F2))

{jaxsim-0.7.1.dev49.dist-info → jaxsim-0.7.1.dev53.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxsim
-Version: 0.7.1.dev49
+Version: 0.7.1.dev53
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippoluca.ferretti@outlook.com>
 Maintainer-email: Filippo Luca Ferretti <filippo.ferretti@iit.it>, Alessandro Croci <alessandro.croci@iit.it>

{jaxsim-0.7.1.dev49.dist-info → jaxsim-0.7.1.dev53.dist-info}/RECORD

@@ -1,5 +1,5 @@
 jaxsim/__init__.py,sha256=EKeysKN-7UswwJLCl7n6qIBKQIVUtYsCMYu_tCoFn7g,3628
-jaxsim/_version.py,sha256=YV4uYC9h68LubyOgYdRGhoUvNTSB-Q4AvQqllU5Yg2o,526
+jaxsim/_version.py,sha256=LFxKSa9ekR80WFGFfy_25o7U4fTH7Li5aXFdTjAxVI0,526
 jaxsim/exceptions.py,sha256=MQ3LRMfVMX2-g3qYj7mUVNV9OLlIA48TANJegbcQyXI,2641
 jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
 jaxsim/typing.py,sha256=7msl8t5Jt09RNYfKdPJtpjLfWurldcycDappb045Eso,761
@@ -7,14 +7,14 @@ jaxsim/api/__init__.py,sha256=4skzcTTejLFfZ_JE6yEEyNxObpXnR5u-bYsn2lBEx-4,234
 jaxsim/api/actuation_model.py,sha256=GgLi-yhEpsczwhMNYIlMvet8tirmmED6S7AumbSbk4U,3705
 jaxsim/api/com.py,sha256=47a9SSaXY540RCkVnHodwLNUrodIfJIkguIYdSEQVwQ,13697
 jaxsim/api/common.py,sha256=yTaRXDYkXmISBOhZ93f9TssR0p4wq7qj7B6OsvYzRME,6942
-jaxsim/api/contact.py,sha256=dlKKDQUG-KQ5qQaYBv2NmZLDb1OnJdltZv8MWXkD_W0,20969
+jaxsim/api/contact.py,sha256=H4ltV0RYv3AbaKPtu0pfyWhcsm4FiEBUEfBtpusEpHw,22087
 jaxsim/api/data.py,sha256=9pxug2gFIZPwqi9kNYXhEziA5IQBB9KNNwIfPfc_kAU,23249
 jaxsim/api/frame.py,sha256=4wg6GsyBQgYhSvc-ry_31JsaL66sZt3TtgwjB7NrHmk,14583
 jaxsim/api/integrators.py,sha256=sHdTWw2Z-Q7jggA8zRkA1KYYd4HNIozXPwNvFwt0YBs,9011
 jaxsim/api/joint.py,sha256=AnqlNWmBOay-gsoo0y4AbfFQ2OCJm-8T1E0IMhZeLoY,7457
-jaxsim/api/kin_dyn_parameters.py,sha256=VrQvH5MXxyiZRkUtXrm99bAZ8bs9Ry02uWMYc-5ZL9Q,39263
+jaxsim/api/kin_dyn_parameters.py,sha256=73Bixrqm481d058xcEBMj8C5xA9ndLDQdn1YE55JR8M,42488
 jaxsim/api/link.py,sha256=bSZOYJDY9HJMgY25VzevTTsdFZTJc6yRRpslc5FhGHE,12782
-jaxsim/api/model.py,sha256=QmLwVogrYQXP97F_6Sz04Wu_gnavuXDA4CNIgDfve8o,86158
+jaxsim/api/model.py,sha256=HaP9Yql_wz2iUGFdMVTtUuI2-CNr5wyFJLbL6rh1Jyo,87302
 jaxsim/api/ode.py,sha256=fp20_LK9lXw2DfNkQgrfJmtd_iBXDNzZkOn0u5Pm8Qw,6190
 jaxsim/api/references.py,sha256=-vd50y3v-jkXAsILS432etIKV6e2EUE2oOeLHuUrfuQ,20399
 jaxsim/math/__init__.py,sha256=dNozvtm8WsB7nxw4uK29yQQKPcDUEczr2zcHoZfMItc,383
@@ -23,13 +23,13 @@ jaxsim/math/cross.py,sha256=AM4HauuuT09q2TN42qvdXhJ9LvtCh0e7ZyLjP-7sANs,1498
 jaxsim/math/inertia.py,sha256=T-iAjPYSD_72R0ZG8GDJhe5i3Jc3ojhlbBRSscTdCKg,1577
 jaxsim/math/joint_model.py,sha256=vBnwXSsw2LCb2Tr5wl2iCo0KvLqcibBbeKcsoH5r9tk,6990
 jaxsim/math/quaternion.py,sha256=MSaZywzJDxs2te1ZELeIcupKSFIA9q_pdXy7fDAEqM4,4539
-jaxsim/math/rotation.py,sha256=TEUtT3X2tFieNxdlccup1pfaTgCTtfX-hTNotd8-nNk,1892
+jaxsim/math/rotation.py,sha256=XMBEnyyWF7jAjRvcFPT1SNj4l-ls0nhxZImp-VGMnrM,2220
 jaxsim/math/skew.py,sha256=z_9YN-NDHL3n4KXWNbzTSMkFDZ0SDpz4RUcwwYFOaao,1402
 jaxsim/math/transform.py,sha256=d0_m_obmUOmnI8Bte0ktvibR9Hv9M9qpg8tVuLON2g0,3192
 jaxsim/math/utils.py,sha256=JgJrBPeuCvi0969VqoNsyk3CflQiLzopngKDjl6RfiE,1898
 jaxsim/mujoco/__init__.py,sha256=1kAWzYOS7nP29S5FGyWPqiAnPf4yPSoaPW-WBGBjVV0,214
 jaxsim/mujoco/__main__.py,sha256=GBmB7J-zj75ZnFyuAAmpSOpbxi_HhHhWJeot3ljGDJY,5291
-jaxsim/mujoco/loaders.py,sha256=OCk1T11iIm3qZUibNpo_bxxLgaGSkCpLt7ae_ND0ExA,23272
+jaxsim/mujoco/loaders.py,sha256=CY-AANwrKRKOhM5E9EDi_ofQYb5jiN4YOleEu8Z0S9A,24231
 jaxsim/mujoco/model.py,sha256=bRXo1uhWDN37sP9qdejr_2vq_PXHQ7p6eyBlFff_JcE,16492
 jaxsim/mujoco/utils.py,sha256=q75OSjxLU2BromVUejt0DVnSbrV5D177YW6LkOdu78g,8823
 jaxsim/mujoco/visualizer.py,sha256=cmI6DhFb1XC7oEtg_wl-s-U56dWHA-F7GlBD6EDYTyA,7744
@@ -44,19 +44,20 @@ jaxsim/parsers/rod/__init__.py,sha256=G2vqlLajBLUc4gyzXwsEI2Wsi4TMOIF9bLDFeT6KrG
 jaxsim/parsers/rod/meshes.py,sha256=yAXefG73_zqbVKRUdlcz9yFmypjDIpiP9cO96PeAozE,2842
 jaxsim/parsers/rod/parser.py,sha256=rmj-W5ekdcIe_Lu66nxTKkgwxF7vGDZdNkfcYYlU-Yc,14398
 jaxsim/parsers/rod/utils.py,sha256=wmD-wCF1lLO8pknX7A3a8CGt9wDlTS_xCqQulcZ_XlM,8242
-jaxsim/rbda/__init__.py,sha256=ksfupKZzeJyysxrbyMyEfszUdBH6LfCfkSz3KLfODhY,328
+jaxsim/rbda/__init__.py,sha256=1gjy9Q7uwZSK4Cqe4gqRHlJ_tYLtF1GvfQEoYx27rsk,520
 jaxsim/rbda/aba.py,sha256=jxvss3XL8pBaT40bWG5pHcH9f1DDl3LoRqdnpFSlCWo,9030
 jaxsim/rbda/collidable_points.py,sha256=XyeV1I43GL22j03rkNVocaIPOGYirt3PiDHrFMndziQ,2070
 jaxsim/rbda/crba.py,sha256=DC9kBXMG1qXaoAdo8K7OCnVHT_YUaL_t6Li56sRf8ro,5093
 jaxsim/rbda/forward_kinematics.py,sha256=qem7Yp-B2oNVOsU3Q2CWV2tbfZKJOCAdDozFgaPB8tg,3838
 jaxsim/rbda/jacobian.py,sha256=EaMvf073UnLWJGXm4UZIlYd4erulFAGgj_pp89k6xic,11113
+jaxsim/rbda/kinematic_constraints.py,sha256=qXxjenDBV67oMaK4dl5MOfDOXKLMZ6K5zbm6qC1Lvz8,4234
 jaxsim/rbda/rnea.py,sha256=lMU7xxdPqGGzk0QwteB-IYjL4auHOpd78C1YqAXlp9s,7588
 jaxsim/rbda/utils.py,sha256=6JwEDQqLMsBX7CUmPYEhdPEscXmGbWVYg6xEriPOgvE,5587
 jaxsim/rbda/actuation/__init__.py,sha256=zWqB8VBHadbyf8FAuhQtcfWdetGjfVxuNDwIeUqNOS4,36
 jaxsim/rbda/actuation/common.py,sha256=aGFqO4VTgQLsTJyOtVuoa_otT_RbkckmG3rq7wjOyB4,462
 jaxsim/rbda/contacts/__init__.py,sha256=resrBkTdOA-1YMdcdUH2RATEhAf_Ye6MQNtjG3ClMYQ,371
-jaxsim/rbda/contacts/common.py,sha256=qVm3Ghoytg1HAeykNrYw5-4rQJ4Mv7h0Pk75ETzGXyc,9045
-jaxsim/rbda/contacts/relaxed_rigid.py,sha256=RjeLF06Pp19qio447U9z5EdhdM6nyMh-ISQX_2-vdaE,21349
+jaxsim/rbda/contacts/common.py,sha256=e-NDoC1N9PK5B0e1fQkYtxS495GSpRCGZRzB6-53WS4,9207
+jaxsim/rbda/contacts/relaxed_rigid.py,sha256=oBQV-oAbx6_D8xhddCRLbV5VV7Lz75levRS5mxBnYVg,25621
 jaxsim/rbda/contacts/rigid.py,sha256=I_TjsT-84ywou2idYioqekeQbIOHTdHI54vot8ijMPk,17605
 jaxsim/rbda/contacts/soft.py,sha256=a7NYMknPfWKfCdbVu83ttDu1u_gssIRvxe9L1622tM0,15284
 jaxsim/terrain/__init__.py,sha256=f7lVX-iNpH_wkkjef9Qpjh19TTAUOQw76EiLYJDVizc,78
@@ -65,8 +66,8 @@ jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=XzmZeIibcaOzaxpprsGSxH3UrM66PAO456rFV91sNXg,11453
 jaxsim/utils/tracing.py,sha256=Btwxdfhb7fJLk3r5PlQkGYj60Y2KbFT1gANGIA697FU,530
 jaxsim/utils/wrappers.py,sha256=3IMwydqFgmSPqeuUQ3PRmdhDc1IoT6XC23jPC_LjWXs,4175
-jaxsim-0.7.1.dev49.dist-info/licenses/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.7.1.dev49.dist-info/METADATA,sha256=c1M6f3ufhWSqOwFuveQj65wnlXdNdHZ24weT_KaKJFY,17851
-jaxsim-0.7.1.dev49.dist-info/WHEEL,sha256=zaaOINJESkSfm_4HQVc5ssNzHCPXhJm0kEUakpsEHaU,91
-jaxsim-0.7.1.dev49.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.7.1.dev49.dist-info/RECORD,,
+jaxsim-0.7.1.dev53.dist-info/licenses/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.7.1.dev53.dist-info/METADATA,sha256=qaL63IhImfWKV-9NalfteDwSmo6w7f-LOH73phZr_do,17851
+jaxsim-0.7.1.dev53.dist-info/WHEEL,sha256=zaaOINJESkSfm_4HQVc5ssNzHCPXhJm0kEUakpsEHaU,91
+jaxsim-0.7.1.dev53.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.7.1.dev53.dist-info/RECORD,,