jaxsim 0.4.3.dev161__tar.gz → 0.4.3.dev177__tar.gz

{jaxsim-0.4.3.dev161 → jaxsim-0.4.3.dev177}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev161
+Version: 0.4.3.dev177
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>
 Maintainer-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippo.ferretti@iit.it>

{jaxsim-0.4.3.dev161 → jaxsim-0.4.3.dev177}/src/jaxsim/_version.py

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

{jaxsim-0.4.3.dev161 → jaxsim-0.4.3.dev177}/src/jaxsim/api/contact.py

@@ -6,9 +6,12 @@ import jax
 import jax.numpy as jnp
 
 import jaxsim.api as js
+import jaxsim.exceptions
 import jaxsim.terrain
 import jaxsim.typing as jtp
+from jaxsim import logging
 from jaxsim.math import Adjoint, Cross, Transform
+from jaxsim.rbda import contacts
 
 from .common import VelRepr
 
@@ -156,14 +159,11 @@ def collidable_point_dynamics(
         Instead, the 6D forces are returned in the active representation.
     """
 
-    # Import privately the contacts classes.
-    from jaxsim.rbda.contacts import RelaxedRigidContacts, RigidContacts, SoftContacts
-
     # Build the soft contact model.
     match model.contact_model:
 
-        case SoftContacts():
-            assert isinstance(model.contact_model, SoftContacts)
+        case contacts.SoftContacts():
+            assert isinstance(model.contact_model, contacts.SoftContacts)
 
             # Compute the 6D force expressed in the inertial frame and applied to each
             # collidable point, and the corresponding material deformation rate.
@@ -178,8 +178,8 @@ def collidable_point_dynamics(
             # of the ODE system. We need to pass its dynamics to the integrator.
             aux_data = dict(m_dot=CW_ṁ)
 
-        case RigidContacts():
-            assert isinstance(model.contact_model, RigidContacts)
+        case contacts.RigidContacts():
+            assert isinstance(model.contact_model, contacts.RigidContacts)
 
             # Compute the 6D force expressed in the inertial frame and applied to each
             # collidable point.
@@ -192,8 +192,8 @@ def collidable_point_dynamics(
 
             aux_data = dict()
 
-        case RelaxedRigidContacts():
-            assert isinstance(model.contact_model, RelaxedRigidContacts)
+        case contacts.RelaxedRigidContacts():
+            assert isinstance(model.contact_model, contacts.RelaxedRigidContacts)
 
             # Compute the 6D force expressed in the inertial frame and applied to each
             # collidable point.
@@ -206,6 +206,31 @@ def collidable_point_dynamics(
 
             aux_data = dict()
 
+        case contacts.ViscoElasticContacts():
+            assert isinstance(model.contact_model, contacts.ViscoElasticContacts)
+
+            # It is not yet clear how to pass the time step to this stage.
+            # A possibility is to restrict the integrator to only forward Euler
+            # and store the Δt inside the model.
+            module = jaxsim.rbda.contacts.visco_elastic.step.__module__
+            name = jaxsim.rbda.contacts.visco_elastic.step.__name__
+            msg = "You need to use the custom '{}.{}' function with this contact model."
+            jaxsim.exceptions.raise_runtime_error_if(
+                condition=True, msg=msg.format(module, name)
+            )
+
+            # Compute the 6D force expressed in the inertial frame and applied to each
+            # collidable point.
+            W_f_Ci, (W_f̿_Ci, m_tf) = model.contact_model.compute_contact_forces(
+                model=model,
+                data=data,
+                dt=None,  # TODO
+                link_forces=link_forces,
+                joint_force_references=joint_force_references,
+            )
+
+            aux_data = dict(W_f_avg2_C=W_f̿_Ci, m_tf=m_tf)
+
         case _:
             raise ValueError(f"Invalid contact model {model.contact_model}")
 
@@ -278,7 +303,6 @@ def in_contact(
     return links_in_contact
 
 
-@jax.jit
 def estimate_good_soft_contacts_parameters(
     model: js.model.JaxSimModel,
     *,
@@ -287,9 +311,15 @@ def estimate_good_soft_contacts_parameters(
     number_of_active_collidable_points_steady_state: jtp.IntLike = 1,
     damping_ratio: jtp.FloatLike = 1.0,
     max_penetration: jtp.FloatLike | None = None,
-) -> jaxsim.rbda.contacts.SoftContactsParams:
+    **kwargs,
+) -> (
+    jaxsim.rbda.contacts.RelaxedRigidContactsParams
+    | jaxsim.rbda.contacts.RigidContactsParams
+    | jaxsim.rbda.contacts.SoftContactsParams
+    | jaxsim.rbda.contacts.ViscoElasticContactsParams
+):
     """
-    Estimate good soft contacts parameters for the given model.
+    Estimate good parameters for soft-like contact models.
 
     Args:
         model: The model to consider.
@@ -313,7 +343,10 @@ def estimate_good_soft_contacts_parameters(
     """
 
     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,
@@ -338,21 +371,39 @@ def estimate_good_soft_contacts_parameters(
 
     match model.contact_model:
 
-        case jaxsim.rbda.contacts.SoftContacts():
-            assert isinstance(model.contact_model, jaxsim.rbda.contacts.SoftContacts)
+        case contacts.SoftContacts():
+            assert isinstance(model.contact_model, contacts.SoftContacts)
+
+            parameters = contacts.SoftContactsParams.build_default_from_jaxsim_model(
+                model=model,
+                standard_gravity=standard_gravity,
+                static_friction_coefficient=static_friction_coefficient,
+                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,
+            )
+
+        case contacts.ViscoElasticContacts():
+            assert isinstance(model.contact_model, contacts.ViscoElasticContacts)
 
             parameters = (
-                jaxsim.rbda.contacts.SoftContactsParams.build_default_from_jaxsim_model(
+                contacts.ViscoElasticContactsParams.build_default_from_jaxsim_model(
                     model=model,
                     standard_gravity=standard_gravity,
                     static_friction_coefficient=static_friction_coefficient,
                     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,
                 )
             )
 
         case _:
+            logging.warning("The active contact model is not soft-like, no-op.")
             parameters = model.contact_model.parameters
 
     return parameters

{jaxsim-0.4.3.dev161 → jaxsim-0.4.3.dev177}/src/jaxsim/api/data.py

@@ -38,12 +38,6 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
     contacts_params: jaxsim.rbda.contacts.ContactsParams = dataclasses.field(repr=False)
 
-    time_ns: jtp.Int = dataclasses.field(
-        default_factory=lambda: jnp.array(
-            0, dtype=jnp.uint64 if jax.config.read("jax_enable_x64") else jnp.uint32
-        ),
-    )
-
     def __hash__(self) -> int:
 
         from jaxsim.utils.wrappers import HashedNumpyArray
@@ -52,7 +46,6 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             (
                 hash(self.state),
                 HashedNumpyArray.hash_of_array(self.gravity),
-                HashedNumpyArray.hash_of_array(self.time_ns),
                 hash(self.contacts_params),
             )
         )
@@ -115,7 +108,6 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
         standard_gravity: jtp.FloatLike = jaxsim.math.StandardGravity,
         contacts_params: jaxsim.rbda.contacts.ContactsParams | None = None,
         velocity_representation: VelRepr = VelRepr.Inertial,
-        time: jtp.FloatLike | None = None,
         extended_ode_state: dict[str, jtp.PyTree] | None = None,
     ) -> JaxSimModelData:
         """
@@ -134,7 +126,6 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             standard_gravity: The standard gravity constant.
             contacts_params: The parameters of the soft contacts.
             velocity_representation: The velocity representation to use.
-            time: The time at which the state is created.
             extended_ode_state:
                 Additional user-defined state variables that are not part of the
                 standard `ODEState` object. Useful to extend the system dynamics
@@ -196,11 +187,6 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             ).squeeze()
         )
 
-        time_ns = jnp.array(
-            time * 1e9 if time is not None else 0.0,
-            dtype=jnp.uint64 if jax.config.read("jax_enable_x64") else jnp.uint32,
-        )
-
         W_H_B = jaxsim.math.Transform.from_quaternion_and_translation(
             translation=base_position, quaternion=base_quaternion
         )
@@ -233,7 +219,11 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         if contacts_params is None:
 
-            if isinstance(model.contact_model, jaxsim.rbda.contacts.SoftContacts):
+            if isinstance(
+                model.contact_model,
+                jaxsim.rbda.contacts.SoftContacts
+                | jaxsim.rbda.contacts.ViscoElasticContacts,
+            ):
                 contacts_params = js.contact.estimate_good_soft_contacts_parameters(
                     model=model, standard_gravity=standard_gravity
                 )
@@ -242,7 +232,6 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
                 contacts_params = model.contact_model.parameters
 
         return JaxSimModelData(
-            time_ns=time_ns,
             state=ode_state,
             gravity=gravity,
             contacts_params=contacts_params,
@@ -253,16 +242,6 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
     # Extract quantities
     # ==================
 
-    def time(self) -> jtp.Float:
-        """
-        Get the simulated time.
-
-        Returns:
-            The simulated time in seconds.
-        """
-
-        return self.time_ns.astype(float) / 1e9
-
     def standard_gravity(self) -> jtp.Float:
         """
         Get the standard gravity constant.

{jaxsim-0.4.3.dev161 → jaxsim-0.4.3.dev177}/src/jaxsim/api/kin_dyn_parameters.py

@@ -5,6 +5,8 @@ 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
@@ -753,6 +755,13 @@ class ContactParameters(JaxsimDataclass):
 
     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 np.where(np.array(self.enabled))[0]
+
     @staticmethod
     def build_from(model_description: ModelDescription) -> ContactParameters:
         """
@@ -785,7 +794,11 @@ class ContactParameters(JaxsimDataclass):
         )
 
         # Build the ContactParameters object.
-        cp = ContactParameters(point=points, body=link_index_of_points)
+        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]

{jaxsim-0.4.3.dev161 → jaxsim-0.4.3.dev177}/src/jaxsim/api/model.py

@@ -32,6 +32,10 @@ class JaxSimModel(JaxsimDataclass):
 
     model_name: Static[str]
 
+    time_step: jaxsim.integrators.TimeStep = dataclasses.field(
+        default_factory=lambda: jnp.array(0.001, dtype=float),
+    )
+
     terrain: Static[jaxsim.terrain.Terrain] = dataclasses.field(
         default_factory=jaxsim.terrain.FlatTerrain.build, repr=False
     )
@@ -64,6 +68,9 @@ class JaxSimModel(JaxsimDataclass):
         if self.model_name != other.model_name:
             return False
 
+        if self.time_step != other.time_step:
+            return False
+
         if self.kin_dyn_parameters != other.kin_dyn_parameters:
             return False
 
@@ -74,6 +81,7 @@ class JaxSimModel(JaxsimDataclass):
         return hash(
             (
                 hash(self.model_name),
+                hash(float(self.time_step)),
                 hash(self.kin_dyn_parameters),
                 hash(self.contact_model),
             )
@@ -88,6 +96,7 @@ class JaxSimModel(JaxsimDataclass):
         model_description: str | pathlib.Path | rod.Model,
         model_name: str | None = None,
         *,
+        time_step: jtp.FloatLike | None = None,
         terrain: jaxsim.terrain.Terrain | None = None,
         contact_model: jaxsim.rbda.contacts.ContactModel | None = None,
         is_urdf: bool | None = None,
@@ -102,6 +111,9 @@ class JaxSimModel(JaxsimDataclass):
                 its content, or a pre-parsed/pre-built rod model.
             model_name:
                 The name of the model. If not specified, it is read from the description.
+            time_step:
+                The default time step to consider for the simulation. It can be
+                manually overridden in the function that steps the simulation.
             terrain: The terrain to consider (the default is a flat infinite plane).
             contact_model:
                 The contact model to consider.
@@ -135,6 +147,7 @@ class JaxSimModel(JaxsimDataclass):
         model = JaxSimModel.build(
             model_description=intermediate_description,
             model_name=model_name,
+            time_step=time_step,
             terrain=terrain,
             contact_model=contact_model,
         )
@@ -150,6 +163,7 @@ class JaxSimModel(JaxsimDataclass):
         model_description: ModelDescription,
         model_name: str | None = None,
         *,
+        time_step: jtp.FloatLike | None = None,
         terrain: jaxsim.terrain.Terrain | None = None,
         contact_model: jaxsim.rbda.contacts.ContactModel | None = None,
     ) -> JaxSimModel:
@@ -162,6 +176,9 @@ class JaxSimModel(JaxsimDataclass):
                 of the model.
             model_name:
                 The name of the model. If not specified, it is read from the description.
+            time_step:
+                The default time step to consider for the simulation. It can be
+                manually overridden in the function that steps the simulation.
             terrain: The terrain to consider (the default is a flat infinite plane).
             contact_model:
                 The contact model to consider.
@@ -179,6 +196,11 @@ class JaxSimModel(JaxsimDataclass):
             terrain or JaxSimModel.__dataclass_fields__["terrain"].default_factory()
         )
 
+        # Consider the default time step if not specified.
+        time_step = (
+            time_step or JaxSimModel.__dataclass_fields__["time_step"].default_factory()
+        )
+
         # Create the default contact model.
         # It will be populated with an initial estimation of good parameters.
         # While these might not be the best, they are a good starting point.
@@ -192,6 +214,7 @@ class JaxSimModel(JaxsimDataclass):
             kin_dyn_parameters=js.kin_dyn_parameters.KynDynParameters.build(
                 model_description=model_description
             ),
+            time_step=time_step,
             terrain=terrain,
             contact_model=contact_model,
             # The following is wrapped as hashless since it's a static argument, and we
@@ -1915,8 +1938,9 @@ def step(
     model: JaxSimModel,
     data: js.data.JaxSimModelData,
     *,
-    dt: jtp.FloatLike,
     integrator: jaxsim.integrators.Integrator,
+    t0: jtp.FloatLike = 0.0,
+    dt: jtp.FloatLike | None = None,
     integrator_state: dict[str, Any] | None = None,
     link_forces: jtp.MatrixLike | None = None,
     joint_force_references: jtp.VectorLike | None = None,
@@ -1928,9 +1952,10 @@ def step(
     Args:
         model: The model to consider.
         data: The data of the considered model.
-        dt: The time step to consider.
         integrator: The integrator to use.
         integrator_state: The state of the integrator.
+        t0: The initial time to consider. Only relevant for time-dependent dynamics.
+        dt: The time step to consider. If not specified, it is read from the model.
         link_forces:
             The 6D forces to apply to the links expressed in the frame corresponding to
             the velocity representation of `data`.
@@ -1951,17 +1976,20 @@ def step(
 
     integrator_state = integrator_state if integrator_state is not None else dict()
 
-    # Extract the initial resources.
-    t0_ns = data.time_ns
+    # Initialize the time-related variables.
     state_t0 = data.state
-    integrator_state_x0 = integrator_state
+    t0 = jnp.array(t0, dtype=float)
+    dt = jnp.array(dt if dt is not None else model.time_step).astype(float)
+
+    # Rename the integrator state.
+    integrator_state_t0 = integrator_state
 
     # Step the dynamics forward.
     state_tf, integrator_state_tf = integrator.step(
         x0=state_t0,
-        t0=jnp.array(t0_ns / 1e9).astype(float),
+        t0=t0,
         dt=dt,
-        params=integrator_state_x0,
+        params=integrator_state_t0,
         # Always inject the current (model, data) pair into the system dynamics
         # considered by the integrator, and include the input variables represented
         # by the pair (joint_force_references, link_forces).
@@ -1980,24 +2008,8 @@ def step(
         ),
     )
 
-    tf_ns = t0_ns + jnp.array(dt * 1e9, dtype=t0_ns.dtype)
-    tf_ns = jnp.where(tf_ns >= t0_ns, tf_ns, jnp.array(0, dtype=t0_ns.dtype))
-
-    jax.lax.cond(
-        pred=tf_ns < t0_ns,
-        true_fun=lambda: jax.debug.print(
-            "The simulation time overflowed, resetting simulation time to 0."
-        ),
-        false_fun=lambda: None,
-    )
-
-    data_tf = (
-        # Store the new state of the model and the new time.
-        data.replace(
-            state=state_tf,
-            time_ns=tf_ns,
-        )
-    )
+    # Store the new state of the model.
+    data_tf = data.replace(state=state_tf)
 
     # Post process the simulation state, if needed.
     match model.contact_model:
@@ -2064,7 +2076,4 @@ def step(
                 velocity_representation=data.velocity_representation, validate=False
             )
 
-    return (
-        data_tf,
-        integrator_state_tf,
-    )
+    return data_tf, integrator_state_tf

{jaxsim-0.4.3.dev161 → jaxsim-0.4.3.dev177}/src/jaxsim/api/ode.py

@@ -8,6 +8,7 @@ import jaxsim.rbda
 import jaxsim.typing as jtp
 from jaxsim.integrators import Time
 from jaxsim.math import Quaternion
+from jaxsim.rbda import contacts
 
 from .common import VelRepr
 from .ode_data import ODEState
@@ -62,7 +63,6 @@ def wrap_system_dynamics_for_integration(
         # Update the state and time stored inside data.
         with data_f.editable(validate=True) as data_rw:
             data_rw.state = x
-            data_rw.time_ns = jnp.array(t * 1e9).astype(data_rw.time_ns.dtype)
 
         # Evaluate the system dynamics, allowing to override the kwargs originally
         # passed when the closure was created.
@@ -371,8 +371,6 @@ def system_dynamics(
         by the system dynamics evaluation.
     """
 
-    from jaxsim.rbda.contacts import RelaxedRigidContacts, RigidContacts, SoftContacts
-
     # Compute the accelerations and the material deformation rate.
     W_v̇_WB, s̈, aux_dict = system_velocity_dynamics(
         model=model,
@@ -387,10 +385,18 @@ def system_dynamics(
 
     match model.contact_model:
 
-        case SoftContacts():
+        case contacts.SoftContacts():
             extended_ode_state["tangential_deformation"] = aux_dict["m_dot"]
 
-        case RigidContacts() | RelaxedRigidContacts():
+        case contacts.ViscoElasticContacts():
+
+            extended_ode_state["contacts_state"] = {
+                "tangential_deformation": jnp.zeros_like(
+                    data.state.extended["tangential_deformation"]
+                )
+            }
+
+        case contacts.RigidContacts() | contacts.RelaxedRigidContacts():
             pass
 
         case _:

{jaxsim-0.4.3.dev161 → jaxsim-0.4.3.dev177}/src/jaxsim/rbda/contacts/__init__.py

@@ -1,5 +1,6 @@
-from . import relaxed_rigid, rigid, soft
+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