jaxsim 0.4.3.dev181__py3-none-any.whl → 0.4.3.dev200__py3-none-any.whl

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.3.dev181'
-__version_tuple__ = version_tuple = (0, 4, 3, 'dev181')
+__version__ = version = '0.4.3.dev200'
+__version_tuple__ = version_tuple = (0, 4, 3, 'dev200')

jaxsim/api/contact.py

@@ -36,11 +36,10 @@ def collidable_point_kinematics(
         the linear component of the mixed 6D frame velocity.
     """
 
-    from jaxsim.rbda import collidable_points
-
     # Switch to inertial-fixed since the RBDAs expect velocities in this representation.
     with data.switch_velocity_representation(VelRepr.Inertial):
-        W_p_Ci, W_ṗ_Ci = collidable_points.collidable_points_pos_vel(
+
+        W_p_Ci, W_ṗ_Ci = jaxsim.rbda.collidable_points.collidable_points_pos_vel(
             model=model,
             base_position=data.base_position(),
             base_quaternion=data.base_orientation(dcm=False),
@@ -304,6 +303,15 @@ def in_contact(
 
 
 def estimate_good_soft_contacts_parameters(
+    *args, **kwargs
+) -> jaxsim.rbda.contacts.ContactParamsTypes:
+
+    msg = "This method is deprecated, please use `{}`."
+    logging.warning(msg.format(estimate_good_contact_parameters.__name__))
+    return estimate_good_contact_parameters(*args, **kwargs)
+
+
+def estimate_good_contact_parameters(
     model: js.model.JaxSimModel,
     *,
     standard_gravity: jtp.FloatLike = jaxsim.math.StandardGravity,
@@ -312,14 +320,9 @@ def estimate_good_soft_contacts_parameters(
     damping_ratio: jtp.FloatLike = 1.0,
     max_penetration: jtp.FloatLike | None = None,
     **kwargs,
-) -> (
-    jaxsim.rbda.contacts.RelaxedRigidContactsParams
-    | jaxsim.rbda.contacts.RigidContactsParams
-    | jaxsim.rbda.contacts.SoftContactsParams
-    | jaxsim.rbda.contacts.ViscoElasticContactsParams
-):
+) -> jaxsim.rbda.contacts.ContactParamsTypes:
     """
-    Estimate good parameters for soft-like contact models.
+    Estimate good contact parameters.
 
     Args:
         model: The model to consider.
@@ -332,12 +335,19 @@ def estimate_good_soft_contacts_parameters(
         max_penetration:
             The maximum penetration allowed in steady state when the robot is
             supported by the configured number of active collidable points.
+        kwargs:
+            Additional model-specific parameters passed to the builder method of
+            the parameters class.
 
     Returns:
-        The estimated good soft contacts parameters.
+        The estimated good contacts parameters.
+
+    Note:
+        This is primarily a convenience function for soft-like contact models.
+        However, it provides with some good default parameters also for the other ones.
 
     Note:
-        This method provides a good starting point for the soft contacts parameters.
+        This method provides a good set of contacts parameters.
         The user is encouraged to fine-tune the parameters based on the
         specific application.
     """
@@ -364,6 +374,7 @@ def estimate_good_soft_contacts_parameters(
     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)
     )
 
@@ -381,8 +392,11 @@ def estimate_good_soft_contacts_parameters(
                 max_penetration=max_δ,
                 number_of_active_collidable_points_steady_state=nc,
                 damping_ratio=damping_ratio,
-                p=model.contact_model.parameters.p,
-                q=model.contact_model.parameters.q,
+                **dict(
+                    p=model.contact_model.parameters.p,
+                    q=model.contact_model.parameters.q,
+                )
+                | kwargs,
             )
 
         case contacts.ViscoElasticContacts():
@@ -396,15 +410,40 @@ def estimate_good_soft_contacts_parameters(
                     max_penetration=max_δ,
                     number_of_active_collidable_points_steady_state=nc,
                     damping_ratio=damping_ratio,
-                    p=model.contact_model.parameters.p,
-                    q=model.contact_model.parameters.q,
-                    **kwargs,
+                    **dict(
+                        p=model.contact_model.parameters.p,
+                        q=model.contact_model.parameters.q,
+                    )
+                    | kwargs,
                 )
             )
 
+        case contacts.RigidContacts():
+            assert isinstance(model.contact_model, contacts.RigidContacts)
+
+            # Disable Baumgarte stabilization by default since it does not play
+            # well with the forward Euler integrator.
+            K = kwargs.get("K", 0.0)
+
+            parameters = contacts.RigidContactsParams.build(
+                mu=static_friction_coefficient,
+                **dict(
+                    K=K,
+                    D=2 * jnp.sqrt(K),
+                )
+                | kwargs,
+            )
+
+        case contacts.RelaxedRigidContacts():
+            assert isinstance(model.contact_model, contacts.RelaxedRigidContacts)
+
+            parameters = contacts.RelaxedRigidContactsParams.build(
+                mu=static_friction_coefficient,
+                **kwargs,
+            )
+
         case _:
-            logging.warning("The active contact model is not soft-like, no-op.")
-            parameters = model.contact_model.parameters
+            raise ValueError(f"Invalid contact model: {model.contact_model}")
 
     return parameters
 

jaxsim/api/data.py

@@ -34,7 +34,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
     state: ODEState
 
-    gravity: jtp.Array
+    gravity: jtp.Vector
 
     contacts_params: jaxsim.rbda.contacts.ContactsParams = dataclasses.field(repr=False)
 
@@ -224,7 +224,8 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
                 jaxsim.rbda.contacts.SoftContacts
                 | jaxsim.rbda.contacts.ViscoElasticContacts,
             ):
-                contacts_params = js.contact.estimate_good_soft_contacts_parameters(
+
+                contacts_params = js.contact.estimate_good_contact_parameters(
                     model=model, standard_gravity=standard_gravity
                 )
 

jaxsim/api/model.py

@@ -40,6 +40,8 @@ class JaxSimModel(JaxsimDataclass):
         default_factory=jaxsim.terrain.FlatTerrain.build, repr=False
     )
 
+    # Note that this is the default contact model.
+    # Its parameters, if any, are then overridden from those stored in JaxSimModelData.
     contact_model: jaxsim.rbda.contacts.ContactModel | None = dataclasses.field(
         default=None, repr=False
     )
@@ -2044,24 +2046,18 @@ def step(
                 M = js.model.free_floating_mass_matrix(model, data_tf)
                 W_p_C = js.contact.collidable_point_positions(model, data_tf)
 
-                # Compute the height of the terrain below each collidable point.
-                px, py, _ = W_p_C.T
-                terrain_height = jax.vmap(model.terrain.height)(px, py)
-
-                # Compute the contact state.
-                inactive_collidable_points, _ = (
-                    jaxsim.rbda.contacts.RigidContacts.detect_contacts(
-                        W_p_C=W_p_C,
-                        terrain_height=terrain_height,
-                    )
-                )
+                # Compute the penetration depth of the collidable points.
+                δ, *_ = jax.vmap(
+                    jaxsim.rbda.contacts.common.compute_penetration_data,
+                    in_axes=(0, 0, None),
+                )(W_p_C, jnp.zeros_like(W_p_C), model.terrain)
 
                 # Compute the impact velocity.
                 # It may be discontinuous in case new contacts are made.
                 BW_nu_post_impact = (
                     jaxsim.rbda.contacts.RigidContacts.compute_impact_velocity(
                         data=data_tf,
-                        inactive_collidable_points=inactive_collidable_points,
+                        inactive_collidable_points=(δ <= 0),
                         M=M,
                         J_WC=J_WC,
                     )

jaxsim/mujoco/__init__.py

@@ -1,3 +1,4 @@
 from .loaders import RodModelToMjcf, SdfToMjcf, UrdfToMjcf
 from .model import MujocoModelHelper
+from .utils import mujoco_data_from_jaxsim
 from .visualizer import MujocoVideoRecorder, MujocoVisualizer

jaxsim/mujoco/loaders.py

@@ -646,7 +646,7 @@ class MujocoCamera:
     def build(cls, **kwargs) -> MujocoCamera:
 
         if not all(isinstance(value, str) for value in kwargs.values()):
-            raise ValueError("Values must be strings")
+            raise ValueError(f"Values must be strings: {kwargs}")
 
         return cls(**kwargs)
 

jaxsim/mujoco/model.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import functools
 import pathlib
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from typing import Any
 
 import mujoco as mj
@@ -107,7 +107,8 @@ class MujocoModelHelper:
             size = [float(el) for el in hfield_element["@size"].split(" ")]
             size[0], size[1] = heightmap_radius_xy
             size[2] = 1.0
-            size[3] = max(0, -min(hfield))
+            # The following could be zero but Mujoco complains if it's exactly zero.
+            size[3] = max(0.000_001, -min(hfield))
 
             # Replace the 'size' attribute.
             hfields_dict[heightmap_name]["@size"] = " ".join(str(el) for el in size)
@@ -315,7 +316,7 @@ class MujocoModelHelper:
         self.data.qpos[sl] = position
 
     def set_joint_positions(
-        self, joint_names: list[str], positions: npt.NDArray | list[npt.NDArray]
+        self, joint_names: Sequence[str], positions: npt.NDArray | list[npt.NDArray]
     ) -> None:
         """Set the positions of multiple joints."""
 

jaxsim/mujoco/utils.py

@@ -0,0 +1,101 @@
+import mujoco as mj
+import numpy as np
+
+from . import MujocoModelHelper
+
+
+def mujoco_data_from_jaxsim(
+    mujoco_model: mj.MjModel,
+    jaxsim_model,
+    jaxsim_data,
+    mujoco_data: mj.MjData | None = None,
+    update_removed_joints: bool = True,
+) -> mj.MjData:
+    """
+    Create a Mujoco data object from a JaxSim model and data objects.
+
+    Args:
+        mujoco_model: The Mujoco model object corresponding to the JaxSim model.
+        jaxsim_model: The JaxSim model object from which the Mujoco model was created.
+        jaxsim_data: The JaxSim data object containing the state of the model.
+        mujoco_data: An optional Mujoco data object. If None, a new one will be created.
+        update_removed_joints:
+            If True, the positions of the joints that have been removed during the
+            model reduction process will be set to their initial values.
+
+    Returns:
+        The Mujoco data object containing the state of the JaxSim model.
+
+    Note:
+        This method is useful to initialize a Mujoco data object used for visualization
+        with the state of a JaxSim model. In particular, this function takes care of
+        initializing the positions of the joints that have been removed during the
+        model reduction process. After the initial creation of the Mujoco data object,
+        it's faster to update the state using an external MujocoModelHelper object.
+    """
+
+    # The package `jaxsim.mujoco` is supposed to be jax-independent.
+    # We import all the JaxSim resources privately.
+    import jaxsim.api as js
+
+    if not isinstance(jaxsim_model, js.model.JaxSimModel):
+        raise ValueError("The `jaxsim_model` argument must be a JaxSimModel object.")
+
+    if not isinstance(jaxsim_data, js.data.JaxSimModelData):
+        raise ValueError("The `jaxsim_data` argument must be a JaxSimModelData object.")
+
+    # Create the helper to operate on the Mujoco model and data.
+    model_helper = MujocoModelHelper(model=mujoco_model, data=mujoco_data)
+
+    # If the model is fixed-base, the Mujoco model won't have the joint corresponding
+    # to the floating base, and the helper would raise an exception.
+    if jaxsim_model.floating_base():
+
+        # Set the model position.
+        model_helper.set_base_position(position=np.array(jaxsim_data.base_position()))
+
+        # Set the model orientation.
+        model_helper.set_base_orientation(
+            orientation=np.array(jaxsim_data.base_orientation())
+        )
+
+    # Set the joint positions.
+    if jaxsim_model.dofs() > 0:
+
+        model_helper.set_joint_positions(
+            joint_names=list(jaxsim_model.joint_names()),
+            positions=np.array(
+                jaxsim_data.joint_positions(
+                    model=jaxsim_model, joint_names=jaxsim_model.joint_names()
+                )
+            ),
+        )
+
+    # Updating these joints is not necessary after the first time.
+    # Users can disable this update after initialization.
+    if update_removed_joints:
+
+        # Create a dictionary with the joints that have been removed for various reasons
+        # (like link lumping due to model reduction).
+        joints_removed_dict = {
+            j.name: j
+            for j in jaxsim_model.description._joints_removed
+            if j.name not in set(jaxsim_model.joint_names())
+        }
+
+        # Set the positions of the removed joints.
+        _ = [
+            model_helper.set_joint_position(
+                position=joints_removed_dict[joint_name].initial_position,
+                joint_name=joint_name,
+            )
+            # Select all original joint that have been removed from the JaxSim model
+            # that are still present in the Mujoco model.
+            for joint_name in joints_removed_dict
+            if joint_name in model_helper.joint_names()
+        ]
+
+    # Return the mujoco data with updated kinematics.
+    mj.mj_forward(mujoco_model, model_helper.data)
+
+    return model_helper.data

jaxsim/mujoco/visualizer.py

@@ -89,7 +89,7 @@ class MujocoVideoRecorder:
         if not exist_ok and path.is_file():
             raise FileExistsError(f"The file '{path}' already exists.")
 
-        media.write_video(path=path, images=self.frames, fps=self.fps)
+        media.write_video(path=path, images=np.array(self.frames), fps=self.fps)
 
     @staticmethod
     def compute_down_sampling(original_fps: int, target_min_fps: int) -> int:

jaxsim/rbda/contacts/__init__.py

@@ -4,3 +4,10 @@ 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

@@ -1,8 +1,12 @@
 from __future__ import annotations
 
 import abc
+import functools
 from typing import Any
 
+import jax
+import jax.numpy as jnp
+
 import jaxsim.api as js
 import jaxsim.terrain
 import jaxsim.typing as jtp
@@ -14,6 +18,47 @@ 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.
@@ -86,7 +131,7 @@ class ContactModel(JaxsimDataclass):
         model: js.model.JaxSimModel,
         data: js.data.JaxSimModelData,
         **kwargs,
-    ) -> tuple[jtp.Vector, tuple[Any, ...]]:
+    ) -> tuple[jtp.Matrix, tuple[Any, ...]]:
         """
         Compute the contact forces.
 
@@ -95,8 +140,9 @@ class ContactModel(JaxsimDataclass):
             data: The data of the considered model.
 
         Returns:
-            A tuple containing as first element the computed 6D contact force applied to the contact point and expressed in the world frame,
-            and as second element a tuple of optional additional information.
+            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 tuple of optional additional information.
         """
 
         pass