jaxsim 0.7.1.dev49__tar.gz → 0.7.1.dev53__tar.gz

{jaxsim-0.7.1.dev49 → jaxsim-0.7.1.dev53}/PKG-INFO

@@ -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 → jaxsim-0.7.1.dev53}/examples/assets/cartpole.urdf

@@ -58,6 +58,18 @@
             </geometry>
         </collision>
     </link>
+    <link name="cart_frame"/>
+    <link name="rail_frame"/>
+    <joint name="cart_frame_joint" type="fixed">
+        <parent link="cart" />
+        <child link="cart_frame" />
+        <origin xyz="0.0 0.0 0.0" rpy="0.0 0.0 0.0" />
+    </joint>
+    <joint name="rail_frame_joint" type="fixed">
+        <parent link="rail" />
+        <child link="rail_frame" />
+        <origin xyz="0.0 0.0 1.2" rpy="0.0 0.0 0.0" />
+    </joint>
     <joint name="world_to_rail" type="fixed">
         <origin xyz="0.0 0.0 0.0" rpy="0.0 0.0 0.0"/>
         <parent link="world"/>

{jaxsim-0.7.1.dev49 → jaxsim-0.7.1.dev53}/src/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-0.7.1.dev49 → jaxsim-0.7.1.dev53}/src/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-0.7.1.dev49 → jaxsim-0.7.1.dev53}/src/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-0.7.1.dev49 → jaxsim-0.7.1.dev53}/src/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-0.7.1.dev49 → jaxsim-0.7.1.dev53}/src/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-0.7.1.dev49 → jaxsim-0.7.1.dev53}/src/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-0.7.1.dev49 → jaxsim-0.7.1.dev53}/src/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-0.7.1.dev49 → jaxsim-0.7.1.dev53}/src/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,