jaxsim 0.4.3.dev143__tar.gz → 0.4.3.dev155__tar.gz

{jaxsim-0.4.3.dev143 → jaxsim-0.4.3.dev155}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev143
+Version: 0.4.3.dev155
 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.dev143 → jaxsim-0.4.3.dev155}/src/jaxsim/_version.py

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

{jaxsim-0.4.3.dev143 → jaxsim-0.4.3.dev155}/src/jaxsim/api/contact.py

@@ -144,7 +144,8 @@ def collidable_point_dynamics(
             The joint force references to apply to the joints.
 
     Returns:
-        The 6D force applied to each collidable point and additional data based on the contact model configured:
+        The 6D force applied to each collidable point and additional data based
+        on the contact model configured:
         - Soft: the material deformation rate.
         - Rigid: no additional data.
         - QuasiRigid: no additional data.
@@ -156,21 +157,13 @@ def collidable_point_dynamics(
     """
 
     # Import privately the contacts classes.
-    from jaxsim.rbda.contacts import (
-        RelaxedRigidContacts,
-        RelaxedRigidContactsState,
-        RigidContacts,
-        RigidContactsState,
-        SoftContacts,
-        SoftContactsState,
-    )
+    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)
-            assert isinstance(data.state.contact, SoftContactsState)
 
             # Compute the 6D force expressed in the inertial frame and applied to each
             # collidable point, and the corresponding material deformation rate.
@@ -187,7 +180,6 @@ def collidable_point_dynamics(
 
         case RigidContacts():
             assert isinstance(model.contact_model, RigidContacts)
-            assert isinstance(data.state.contact, RigidContactsState)
 
             # Compute the 6D force expressed in the inertial frame and applied to each
             # collidable point.
@@ -203,7 +195,6 @@ def collidable_point_dynamics(
 
         case RelaxedRigidContacts():
             assert isinstance(model.contact_model, RelaxedRigidContacts)
-            assert isinstance(data.state.contact, RelaxedRigidContactsState)
 
             # Compute the 6D force expressed in the inertial frame and applied to each
             # collidable point.

{jaxsim-0.4.3.dev143 → jaxsim-0.4.3.dev155}/src/jaxsim/api/data.py

@@ -13,7 +13,6 @@ import jaxsim.api as js
 import jaxsim.math
 import jaxsim.rbda
 import jaxsim.typing as jtp
-from jaxsim.rbda.contacts import SoftContacts
 from jaxsim.utils import Mutability
 from jaxsim.utils.tracing import not_tracing
 
@@ -107,17 +106,17 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
     @staticmethod
     def build(
         model: js.model.JaxSimModel,
-        base_position: jtp.Vector | None = None,
-        base_quaternion: jtp.Vector | None = None,
-        joint_positions: jtp.Vector | None = None,
-        base_linear_velocity: jtp.Vector | None = None,
-        base_angular_velocity: jtp.Vector | None = None,
-        joint_velocities: jtp.Vector | None = None,
+        base_position: jtp.VectorLike | None = None,
+        base_quaternion: jtp.VectorLike | None = None,
+        joint_positions: jtp.VectorLike | None = None,
+        base_linear_velocity: jtp.VectorLike | None = None,
+        base_angular_velocity: jtp.VectorLike | None = None,
+        joint_velocities: jtp.VectorLike | None = None,
         standard_gravity: jtp.FloatLike = jaxsim.math.StandardGravity,
-        contact: jaxsim.rbda.contacts.ContactsState | None = None,
         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:
         """
         Create a `JaxSimModelData` object with the given state.
@@ -133,56 +132,73 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
                 The base angular velocity in the selected representation.
             joint_velocities: The joint velocities.
             standard_gravity: The standard gravity constant.
-            contact: The state of the soft contacts.
             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
+                considered by default in JaxSim.
 
         Returns:
-            A `JaxSimModelData` object with the given state.
+            A `JaxSimModelData` initialized with the given state.
         """
 
         base_position = jnp.array(
-            base_position if base_position is not None else jnp.zeros(3)
+            base_position if base_position is not None else jnp.zeros(3),
+            dtype=float,
         ).squeeze()
 
         base_quaternion = jnp.array(
-            base_quaternion
-            if base_quaternion is not None
-            else jnp.array([1.0, 0, 0, 0])
+            (
+                base_quaternion
+                if base_quaternion is not None
+                else jnp.array([1.0, 0, 0, 0])
+            ),
+            dtype=float,
         ).squeeze()
 
         base_linear_velocity = jnp.array(
-            base_linear_velocity if base_linear_velocity is not None else jnp.zeros(3)
+            base_linear_velocity if base_linear_velocity is not None else jnp.zeros(3),
+            dtype=float,
         ).squeeze()
 
         base_angular_velocity = jnp.array(
-            base_angular_velocity if base_angular_velocity is not None else jnp.zeros(3)
+            (
+                base_angular_velocity
+                if base_angular_velocity is not None
+                else jnp.zeros(3)
+            ),
+            dtype=float,
         ).squeeze()
 
         gravity = jnp.zeros(3).at[2].set(-standard_gravity)
 
         joint_positions = jnp.atleast_1d(
-            joint_positions.squeeze()
-            if joint_positions is not None
-            else jnp.zeros(model.dofs())
+            jnp.array(
+                (
+                    joint_positions
+                    if joint_positions is not None
+                    else jnp.zeros(model.dofs())
+                ),
+                dtype=float,
+            ).squeeze()
         )
 
         joint_velocities = jnp.atleast_1d(
-            joint_velocities.squeeze()
-            if joint_velocities is not None
-            else jnp.zeros(model.dofs())
+            jnp.array(
+                (
+                    joint_velocities
+                    if joint_velocities is not None
+                    else jnp.zeros(model.dofs())
+                ),
+                dtype=float,
+            ).squeeze()
         )
 
-        time_ns = (
-            jnp.array(
-                time * 1e9,
-                dtype=jnp.uint64 if jax.config.read("jax_enable_x64") else jnp.uint32,
-            )
-            if time is not None
-            else jnp.array(
-                0, dtype=jnp.uint64 if jax.config.read("jax_enable_x64") else jnp.uint32
-            )
+        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(
@@ -194,21 +210,22 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             other_representation=velocity_representation,
             transform=W_H_B,
             is_force=False,
-        )
+        ).astype(float)
 
         ode_state = ODEState.build_from_jaxsim_model(
             model=model,
-            base_position=base_position.astype(float),
-            base_quaternion=base_quaternion.astype(float),
-            joint_positions=joint_positions.astype(float),
-            base_linear_velocity=v_WB[0:3].astype(float),
-            base_angular_velocity=v_WB[3:6].astype(float),
-            joint_velocities=joint_velocities.astype(float),
-            tangential_deformation=(
-                contact.tangential_deformation
-                if contact is not None and isinstance(model.contact_model, SoftContacts)
-                else None
-            ),
+            base_position=base_position,
+            base_quaternion=base_quaternion,
+            joint_positions=joint_positions,
+            base_linear_velocity=v_WB[0:3],
+            base_angular_velocity=v_WB[3:6],
+            joint_velocities=joint_velocities,
+            # Unpack all the additional ODE states. If the contact model requires an
+            # additional state that is not explicitly passed to this builder, ODEState
+            # automatically populates that state with zeroed variables.
+            # This is not true for any other custom state that the user might want to
+            # pass to the integrator.
+            **(extended_ode_state if extended_ode_state else {}),
         )
 
         if not ode_state.valid(model=model):
@@ -220,13 +237,14 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
                 contacts_params = js.contact.estimate_good_soft_contacts_parameters(
                     model=model, standard_gravity=standard_gravity
                 )
+
             else:
                 contacts_params = model.contact_model.parameters
 
         return JaxSimModelData(
             time_ns=time_ns,
             state=ode_state,
-            gravity=gravity.astype(float),
+            gravity=gravity,
             contacts_params=contacts_params,
             velocity_representation=velocity_representation,
         )

{jaxsim-0.4.3.dev143 → jaxsim-0.4.3.dev155}/src/jaxsim/api/model.py

@@ -33,7 +33,7 @@ class JaxSimModel(JaxsimDataclass):
     model_name: Static[str]
 
     terrain: Static[jaxsim.terrain.Terrain] = dataclasses.field(
-        default=jaxsim.terrain.FlatTerrain(), repr=False
+        default_factory=jaxsim.terrain.FlatTerrain.build, repr=False
     )
 
     contact_model: jaxsim.rbda.contacts.ContactModel | None = dataclasses.field(
@@ -101,13 +101,14 @@ class JaxSimModel(JaxsimDataclass):
                 A path to an SDF/URDF file, a string containing
                 its content, or a pre-parsed/pre-built rod model.
             model_name:
-                The optional name of the model that overrides the one in
-                the description.
-            terrain:
-                The optional terrain to consider.
+                The name of the model. If not specified, it is read from the description.
+            terrain: The terrain to consider (the default is a flat infinite plane).
+            contact_model:
+                The contact model to consider.
+                If not specified, a soft contacts model is used.
             is_urdf:
-                The optional flag to force the model description to be parsed as a
-                URDF or a SDF. This is otherwise automatically inferred.
+                The optional flag to force the model description to be parsed as a URDF.
+                This is usually automatically inferred.
             considered_joints:
                 The list of joints to consider. If None, all joints are considered.
 
@@ -120,7 +121,7 @@ class JaxSimModel(JaxsimDataclass):
         # Parse the input resource (either a path to file or a string with the URDF/SDF)
         # and build the -intermediate- model description.
         intermediate_description = jaxsim.parsers.rod.build_model_description(
-            model_description=model_description
+            model_description=model_description, is_urdf=is_urdf
         )
 
         # Lump links together if not all joints are considered.
@@ -160,11 +161,11 @@ class JaxSimModel(JaxsimDataclass):
                 The intermediate model description defining the kinematics and dynamics
                 of the model.
             model_name:
-                The optional name of the model overriding the physics model name.
-            terrain:
-                The optional terrain to consider.
+                The name of the model. If not specified, it is read from the description.
+            terrain: The terrain to consider (the default is a flat infinite plane).
             contact_model:
-                The optional contact model to consider. If None, the soft contact model is used.
+                The contact model to consider.
+                If not specified, a soft contacts model is used.
 
         Returns:
             The built Model object.
@@ -173,21 +174,31 @@ class JaxSimModel(JaxsimDataclass):
         # Set the model name (if not provided, use the one from the model description).
         model_name = model_name if model_name is not None else model_description.name
 
-        # Set the terrain (if not provided, use the default flat terrain).
-        terrain = terrain or JaxSimModel.__dataclass_fields__["terrain"].default
-        contact_model = contact_model or jaxsim.rbda.contacts.SoftContacts(
-            terrain=terrain
+        # Consider the default terrain (a flat infinite plane) if not specified.
+        terrain = (
+            terrain or JaxSimModel.__dataclass_fields__["terrain"].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.
+        contact_model = contact_model or jaxsim.rbda.contacts.SoftContacts.build(
+            terrain=terrain, parameters=None
         )
 
         # Build the model.
         model = JaxSimModel(
             model_name=model_name,
-            _description=wrappers.HashlessObject(obj=model_description),
             kin_dyn_parameters=js.kin_dyn_parameters.KynDynParameters.build(
                 model_description=model_description
             ),
             terrain=terrain,
             contact_model=contact_model,
+            # The following is wrapped as hashless since it's a static argument, and we
+            # don't want to trigger recompilation if it changes. All relevant parameters
+            # needed to compute kinematics and dynamics quantities are stored in the
+            # kin_dyn_parameters attribute.
+            _description=wrappers.HashlessObject(obj=model_description),
         )
 
         return model

{jaxsim-0.4.3.dev143 → jaxsim-0.4.3.dev155}/src/jaxsim/api/ode.py

@@ -370,9 +370,8 @@ def system_dynamics(
         corresponding derivative, and the dictionary of auxiliary data returned
         by the system dynamics evaluation.
     """
-    from jaxsim.rbda.contacts.relaxed_rigid import RelaxedRigidContacts
-    from jaxsim.rbda.contacts.rigid import RigidContacts
-    from jaxsim.rbda.contacts.soft import SoftContacts
+
+    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(
@@ -382,17 +381,20 @@ def system_dynamics(
         link_forces=link_forces,
     )
 
-    ode_state_kwargs = {}
+    # Initialize the dictionary storing the derivative of the additional state variables
+    # that extend the state vector of the integrated ODE system.
+    extended_ode_state = {}
 
     match model.contact_model:
+
         case SoftContacts():
-            ode_state_kwargs["tangential_deformation"] = aux_dict["m_dot"]
+            extended_ode_state["tangential_deformation"] = aux_dict["m_dot"]
 
         case RigidContacts() | RelaxedRigidContacts():
             pass
 
         case _:
-            raise ValueError("Unable to determine contact state class prefix.")
+            raise ValueError(f"Invalid contact model {model.contact_model}")
 
     # Extract the velocities.
     W_ṗ_B, W_Q̇_B, ṡ = system_position_dynamics(
@@ -412,7 +414,7 @@ def system_dynamics(
         base_linear_velocity=W_v̇_WB[0:3],
         base_angular_velocity=W_v̇_WB[3:6],
         joint_velocities=s̈,
-        **ode_state_kwargs,
+        **extended_ode_state,
     )
 
     return ode_state_derivative, aux_dict

{jaxsim-0.4.3.dev143 → jaxsim-0.4.3.dev155}/src/jaxsim/api/ode_data.py

@@ -1,19 +1,13 @@
 from __future__ import annotations
 
+import dataclasses
+
+import jax
 import jax.numpy as jnp
 import jax_dataclasses
 
 import jaxsim.api as js
 import jaxsim.typing as jtp
-from jaxsim.rbda.contacts import (
-    ContactsState,
-    RelaxedRigidContacts,
-    RelaxedRigidContactsState,
-    RigidContacts,
-    RigidContactsState,
-    SoftContacts,
-    SoftContactsState,
-)
 from jaxsim.utils import JaxsimDataclass
 
 # =============================================================================
@@ -125,15 +119,18 @@ class ODEState(JaxsimDataclass):
 
     Attributes:
         physics_model: The state of the physics model.
-        contact: The state of the contacts model.
+        extended:
+            Additional state variables extending the state vector corresponding to
+            equations of motion. These extended variables are passed to the integrator.
     """
 
     physics_model: PhysicsModelState
-    contact: ContactsState
+
+    extended: dict[str, jtp.PyTree] = dataclasses.field(default_factory=dict)
 
     @staticmethod
     def build_from_jaxsim_model(
-        model: js.model.JaxSimModel | None = None,
+        model: js.model.JaxSimModel,
         joint_positions: jtp.Vector | None = None,
         joint_velocities: jtp.Vector | None = None,
         base_position: jtp.Vector | None = None,
@@ -155,7 +152,15 @@ class ODEState(JaxsimDataclass):
                 The linear velocity of the base link in inertial-fixed representation.
             base_angular_velocity:
                 The angular velocity of the base link in inertial-fixed representation.
-            kwargs: Additional arguments needed to build the contact state.
+            kwargs:
+                Additional arguments corresponding variables extending the default
+                state vector of the physics model.
+
+        Note:
+            Kwargs can be used to supply any additional state variables that are passed
+            to the integrator. This is useful to extend the default system dynamics,
+            for example if the contact model requires additional state variables or to
+            simulate additional dynamics like actuators or muscoloskeletal models.
 
         Returns:
             The `ODEState` built from the `JaxSimModel`.
@@ -165,29 +170,11 @@ class ODEState(JaxsimDataclass):
             `JaxSimModel` and initialized to zero.
         """
 
-        # Get the contact model from the `JaxSimModel`.
-        match model.contact_model:
-
-            case SoftContacts():
-
-                tangential_deformation = kwargs.get("tangential_deformation", None)
-
-                contact = SoftContactsState.build_from_jaxsim_model(
-                    model=model,
-                    **(
-                        dict(tangential_deformation=tangential_deformation)
-                        if tangential_deformation is not None
-                        else dict()
-                    ),
-                )
-            case RigidContacts():
-                contact = RigidContactsState.build()
+        # Initialize the extended state with the optional contact state.
+        extended_state = model.contact_model.zero_state_variables(model=model)
 
-            case RelaxedRigidContacts():
-                contact = RelaxedRigidContactsState.build()
-
-            case _:
-                raise ValueError("Unsupported contact model.")
+        # Override the default extended state with optional kwargs.
+        extended_state |= kwargs
 
         return ODEState.build(
             model=model,
@@ -200,13 +187,13 @@ class ODEState(JaxsimDataclass):
                 base_linear_velocity=base_linear_velocity,
                 base_angular_velocity=base_angular_velocity,
             ),
-            contact=contact,
+            extended_state=extended_state,
         )
 
     @staticmethod
     def build(
         physics_model_state: PhysicsModelState | None = None,
-        contact: ContactsState | None = None,
+        extended_state: dict[str, jtp.PyTree] | None = None,
         model: js.model.JaxSimModel | None = None,
     ) -> ODEState:
         """
@@ -214,62 +201,60 @@ class ODEState(JaxsimDataclass):
 
         Args:
             physics_model_state: The state of the physics model.
-            contact: The state of the contacts model.
+            extended_state: Additional state variables extending the state vector.
             model: The `JaxSimModel` associated with the ODE state.
 
         Returns:
             A `ODEState` instance.
         """
 
+        # Build a zero state for the physics model if not provided.
         physics_model_state = (
             physics_model_state
             if physics_model_state is not None
             else PhysicsModelState.zero(model=model)
         )
 
-        # Get the contact model from the `JaxSimModel`.
-        match contact:
-            case (
-                SoftContactsState() | RigidContactsState() | RelaxedRigidContactsState()
-            ):
-                pass
-            case None:
-                contact = SoftContactsState.zero(model=model)
-            case _:
-                raise ValueError("Unable to determine contact state class prefix.")
-
-        return ODEState(physics_model=physics_model_state, contact=contact)
+        return ODEState(
+            physics_model=physics_model_state,
+            extended=extended_state,
+        )
 
     @staticmethod
     def zero(model: js.model.JaxSimModel, data: js.data.JaxSimModelData) -> ODEState:
         """
-        Build a zero `ODEState` from a `JaxSimModel`.
+        Build a zero `ODEState` corresponding to a `JaxSimModel`.
 
         Args:
-            model: The `JaxSimModel` associated with the ODE state.
+            model: The model to consider.
+            data: The data of the considered model.
 
         Returns:
             A zero `ODEState` instance.
         """
 
-        model_state = ODEState.build(
-            model=model, contact=data.state.contact.zero(model=model)
+        ode_state = ODEState.build(
+            model=model,
+            extended_state=jax.tree.map(
+                lambda x: jnp.zeros_like(x), data.state.extended
+            ),
         )
 
-        return model_state
+        return ode_state
 
     def valid(self, model: js.model.JaxSimModel) -> bool:
         """
         Check if the `ODEState` is valid for a given `JaxSimModel`.
 
         Args:
-            model: The `JaxSimModel` to validate the `ODEState` against.
+            model: The model to validate this `ODEState` against.
 
         Returns:
             `True` if the ODE state is valid for the given model, `False` otherwise.
         """
 
-        return self.physics_model.valid(model=model) and self.contact.valid(model=model)
+        # TODO: should we validate the extended state?
+        return self.physics_model.valid(model=model)
 
 
 # ==================================================

jaxsim-0.4.3.dev155/src/jaxsim/rbda/contacts/__init__.py

@@ -0,0 +1,5 @@
+from . import relaxed_rigid, rigid, soft
+from .common import ContactModel, ContactsParams
+from .relaxed_rigid import RelaxedRigidContacts, RelaxedRigidContactsParams
+from .rigid import RigidContacts, RigidContactsParams
+from .soft import SoftContacts, SoftContactsParams