jaxsim 0.4.3.dev139__py3-none-any.whl → 0.4.3.dev155__py3-none-any.whl

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
 
 # =============================================================================
@@ -38,16 +32,16 @@ class ODEInput(JaxsimDataclass):
     @staticmethod
     def build_from_jaxsim_model(
         model: js.model.JaxSimModel | None = None,
-        joint_forces: jtp.VectorLike | None = None,
         link_forces: jtp.MatrixLike | None = None,
+        joint_force_references: jtp.VectorLike | None = None,
     ) -> ODEInput:
         """
         Build an `ODEInput` from a `JaxSimModel`.
 
         Args:
             model: The `JaxSimModel` associated with the ODE input.
-            joint_forces: The vector of joint forces.
             link_forces: The matrix of external forces applied to the links.
+            joint_force_references: The vector of joint force references.
 
         Returns:
             The `ODEInput` built from the `JaxSimModel`.
@@ -60,8 +54,8 @@ class ODEInput(JaxsimDataclass):
         return ODEInput.build(
             physics_model_input=PhysicsModelInput.build_from_jaxsim_model(
                 model=model,
-                joint_forces=joint_forces,
                 link_forces=link_forces,
+                joint_force_references=joint_force_references,
             ),
             model=model,
         )
@@ -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()
-
-            case RelaxedRigidContacts():
-                contact = RelaxedRigidContactsState.build()
+        # Initialize the extended state with the optional contact state.
+        extended_state = model.contact_model.zero_state_variables(model=model)
 
-            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)
 
 
 # ==================================================
@@ -526,16 +511,16 @@ class PhysicsModelInput(JaxsimDataclass):
     @staticmethod
     def build_from_jaxsim_model(
         model: js.model.JaxSimModel | None = None,
-        joint_forces: jtp.VectorLike | None = None,
         link_forces: jtp.MatrixLike | None = None,
+        joint_force_references: jtp.VectorLike | None = None,
     ) -> PhysicsModelInput:
         """
         Build a `PhysicsModelInput` from a `JaxSimModel`.
 
         Args:
             model: The `JaxSimModel` associated with the input.
-            joint_forces: The vector of joint forces.
             link_forces: The matrix of external forces applied to the links.
+            joint_force_references: The vector of joint force references.
 
         Returns:
             A `PhysicsModelInput` instance.
@@ -546,7 +531,7 @@ class PhysicsModelInput(JaxsimDataclass):
         """
 
         return PhysicsModelInput.build(
-            joint_forces=joint_forces,
+            joint_force_references=joint_force_references,
             link_forces=link_forces,
             number_of_dofs=model.dofs(),
             number_of_links=model.number_of_links(),
@@ -554,8 +539,8 @@ class PhysicsModelInput(JaxsimDataclass):
 
     @staticmethod
     def build(
-        joint_forces: jtp.VectorLike | None = None,
         link_forces: jtp.MatrixLike | None = None,
+        joint_force_references: jtp.VectorLike | None = None,
         number_of_dofs: jtp.Int | None = None,
         number_of_links: jtp.Int | None = None,
     ) -> PhysicsModelInput:
@@ -563,8 +548,8 @@ class PhysicsModelInput(JaxsimDataclass):
         Build a `PhysicsModelInput`.
 
         Args:
-            joint_forces: The vector of joint forces.
             link_forces: The matrix of external forces applied to the links.
+            joint_force_references: The vector of joint force references.
             number_of_dofs: The number of degrees of freedom of the model.
             number_of_links: The number of links of the model.
 
@@ -572,19 +557,21 @@ class PhysicsModelInput(JaxsimDataclass):
             A `PhysicsModelInput` instance.
         """
 
-        joint_forces = (
-            joint_forces if joint_forces is not None else jnp.zeros(number_of_dofs)
-        )
+        joint_force_references = jnp.atleast_1d(
+            jnp.array(joint_force_references, dtype=float).squeeze()
+            if joint_force_references is not None
+            else jnp.zeros(number_of_dofs)
+        ).astype(float)
 
-        link_forces = (
-            link_forces
+        link_forces = jnp.atleast_2d(
+            jnp.array(link_forces, dtype=float).squeeze()
             if link_forces is not None
             else jnp.zeros(shape=(number_of_links, 6))
-        )
+        ).astype(float)
 
         return PhysicsModelInput(
-            tau=jnp.array(joint_forces, dtype=float),
-            f_ext=jnp.array(link_forces, dtype=float),
+            tau=joint_force_references,
+            f_ext=link_forces,
         )
 
     @staticmethod

jaxsim/api/references.py

@@ -55,8 +55,8 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
     @staticmethod
     def build(
         model: js.model.JaxSimModel,
-        joint_force_references: jtp.Vector | None = None,
-        link_forces: jtp.Matrix | None = None,
+        joint_force_references: jtp.VectorLike | None = None,
+        link_forces: jtp.MatrixLike | None = None,
         data: js.data.JaxSimModelData | None = None,
         velocity_representation: VelRepr | None = None,
     ) -> JaxSimModelReferences:
@@ -78,14 +78,14 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
 
         # Create or adjust joint force references.
         joint_force_references = jnp.atleast_1d(
-            joint_force_references.squeeze()
+            jnp.array(joint_force_references, dtype=float).squeeze()
             if joint_force_references is not None
             else jnp.zeros(model.dofs())
         ).astype(float)
 
         # Create or adjust link forces.
         f_L = jnp.atleast_2d(
-            link_forces.squeeze()
+            jnp.array(link_forces, dtype=float).squeeze()
             if link_forces is not None
             else jnp.zeros((model.number_of_links(), 6))
         ).astype(float)
@@ -299,9 +299,9 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
             A new `JaxSimModelReferences` object with the given joint force references.
         """
 
-        forces = jnp.array(forces)
+        forces = jnp.atleast_1d(jnp.array(forces, dtype=float).squeeze())
 
-        def replace(forces: jtp.VectorLike) -> JaxSimModelReferences:
+        def replace(forces: jtp.Vector) -> JaxSimModelReferences:
             return self.replace(
                 validate=True,
                 input=self.input.replace(

jaxsim/rbda/contacts/__init__.py

@@ -1,9 +1,5 @@
 from . import relaxed_rigid, rigid, soft
-from .common import ContactModel, ContactsParams, ContactsState
-from .relaxed_rigid import (
-    RelaxedRigidContacts,
-    RelaxedRigidContactsParams,
-    RelaxedRigidContactsState,
-)
-from .rigid import RigidContacts, RigidContactsParams, RigidContactsState
-from .soft import SoftContacts, SoftContactsParams, SoftContactsState
+from .common import ContactModel, ContactsParams
+from .relaxed_rigid import RelaxedRigidContacts, RelaxedRigidContactsParams
+from .rigid import RigidContacts, RigidContactsParams
+from .soft import SoftContacts, SoftContactsParams

jaxsim/rbda/contacts/common.py

@@ -14,41 +14,6 @@ except ImportError:
     from typing_extensions import Self
 
 
-class ContactsState(JaxsimDataclass):
-    """
-    Abstract class storing the state of the contacts model.
-    """
-
-    @classmethod
-    @abc.abstractmethod
-    def build(cls: type[Self], **kwargs) -> Self:
-        """
-        Build the contact state object.
-
-        Returns:
-            The contact state object.
-        """
-        pass
-
-    @classmethod
-    @abc.abstractmethod
-    def zero(cls: type[Self], **kwargs) -> Self:
-        """
-        Build a zero contact state.
-
-        Returns:
-            The zero contact state.
-        """
-        pass
-
-    @abc.abstractmethod
-    def valid(self, **kwargs) -> jtp.BoolLike:
-        """
-        Check if the contacts state is valid.
-        """
-        pass
-
-
 class ContactsParams(JaxsimDataclass):
     """
     Abstract class representing the parameters of a contact model.
@@ -88,6 +53,27 @@ class ContactModel(JaxsimDataclass):
     parameters: ContactsParams
     terrain: jaxsim.terrain.Terrain
 
+    @classmethod
+    @abc.abstractmethod
+    def build(
+        cls: type[Self],
+        parameters: ContactsParams,
+        terrain: jaxsim.terrain.Terrain,
+        **kwargs,
+    ) -> Self:
+        """
+        Create a `ContactModel` instance with specified parameters.
+
+        Args:
+            parameters: The parameters of the contact model.
+            terrain: The considered terrain.
+
+        Returns:
+            The `ContactModel` instance.
+        """
+
+        pass
+
     @abc.abstractmethod
     def compute_contact_forces(
         self,
@@ -109,6 +95,27 @@ class ContactModel(JaxsimDataclass):
 
         pass
 
+    @classmethod
+    def zero_state_variables(cls, model: js.model.JaxSimModel) -> dict[str, jtp.Array]:
+        """
+        Build zero state variables of the contact model.
+
+        Args:
+            model: The robot model considered by the contact model.
+
+        Note:
+            There are contact models that require to extend the state vector of the
+            integrated ODE system with additional variables. Our integrators are
+            capable of operating on a generic state, as long as it is a PyTree.
+            This method builds the zero state variables of the contact model as a
+            dictionary of JAX arrays.
+
+        Returns:
+            A dictionary storing the zero state variables of the contact model.
+        """
+
+        return {}
+
     def initialize_model_and_data(
         self,
         model: js.model.JaxSimModel,

jaxsim/rbda/contacts/relaxed_rigid.py

@@ -11,11 +11,12 @@ import optax
 
 import jaxsim.api as js
 import jaxsim.typing as jtp
+from jaxsim import logging
 from jaxsim.api.common import VelRepr
 from jaxsim.math import Adjoint
 from jaxsim.terrain.terrain import FlatTerrain, Terrain
 
-from .common import ContactModel, ContactsParams, ContactsState
+from .common import ContactModel, ContactsParams
 
 try:
     from typing import Self
@@ -156,41 +157,46 @@ class RelaxedRigidContactsParams(ContactsParams):
         )
 
 
-@jax_dataclasses.pytree_dataclass
-class RelaxedRigidContactsState(ContactsState):
-    """Class storing the state of the relaxed rigid contacts model."""
-
-    def __eq__(self, other: RelaxedRigidContactsState) -> bool:
-        return hash(self) == hash(other)
-
-    @classmethod
-    def build(cls: type[Self]) -> Self:
-        """Create a `RelaxedRigidContactsState` instance"""
-
-        return cls()
-
-    @classmethod
-    def zero(cls: type[Self], **kwargs) -> Self:
-        """Build a zero `RelaxedRigidContactsState` instance from a `JaxSimModel`."""
-
-        return cls.build()
-
-    def valid(self, **kwargs) -> jtp.BoolLike:
-        return True
-
-
 @jax_dataclasses.pytree_dataclass
 class RelaxedRigidContacts(ContactModel):
     """Relaxed rigid contacts model."""
 
     parameters: RelaxedRigidContactsParams = dataclasses.field(
-        default_factory=RelaxedRigidContactsParams
+        default_factory=RelaxedRigidContactsParams.build
     )
 
     terrain: jax_dataclasses.Static[Terrain] = dataclasses.field(
-        default_factory=FlatTerrain
+        default_factory=FlatTerrain.build
     )
 
+    @classmethod
+    def build(
+        cls: type[Self],
+        parameters: RelaxedRigidContactsParams | None = None,
+        terrain: Terrain | None = None,
+        **kwargs,
+    ) -> Self:
+        """
+        Create a `RelaxedRigidContacts` instance with specified parameters.
+
+        Args:
+            parameters: The parameters of the rigid contacts model.
+            terrain: The considered terrain.
+
+        Returns:
+            The `RelaxedRigidContacts` instance.
+        """
+
+        if len(kwargs) != 0:
+            logging.debug(msg=f"Ignoring extra arguments: {kwargs}")
+
+        return cls(
+            parameters=(
+                parameters or cls.__dataclass_fields__["parameters"].default_factory()
+            ),
+            terrain=terrain or cls.__dataclass_fields__["terrain"].default_factory(),
+        )
+
     @jax.jit
     def compute_contact_forces(
         self,
@@ -274,7 +280,9 @@ class RelaxedRigidContacts(ContactModel):
                     model=model,
                     data=data,
                     link_forces=references.link_forces(model=model, data=data),
-                    joint_forces=references.joint_force_references(model=model),
+                    joint_force_references=references.joint_force_references(
+                        model=model
+                    ),
                 )
             )
             BW_ν = data.generalized_velocity()

jaxsim/rbda/contacts/rigid.py

@@ -9,10 +9,11 @@ import jax_dataclasses
 
 import jaxsim.api as js
 import jaxsim.typing as jtp
+from jaxsim import logging
 from jaxsim.api.common import ModelDataWithVelocityRepresentation, VelRepr
 from jaxsim.terrain import FlatTerrain, Terrain
 
-from .common import ContactModel, ContactsParams, ContactsState
+from .common import ContactModel, ContactsParams
 
 try:
     from typing import Self
@@ -78,29 +79,6 @@ class RigidContactsParams(ContactsParams):
         )
 
 
-@jax_dataclasses.pytree_dataclass
-class RigidContactsState(ContactsState):
-    """Class storing the state of the rigid contacts model."""
-
-    def __eq__(self, other: RigidContactsState) -> bool:
-        return hash(self) == hash(other)
-
-    @classmethod
-    def build(cls: type[Self]) -> Self:
-        """Create a `RigidContactsState` instance"""
-
-        return cls()
-
-    @classmethod
-    def zero(cls: type[Self], **kwargs) -> Self:
-        """Build a zero `RigidContactsState` instance from a `JaxSimModel`."""
-
-        return cls.build()
-
-    def valid(self, **kwargs) -> jtp.BoolLike:
-        return True
-
-
 @jax_dataclasses.pytree_dataclass
 class RigidContacts(ContactModel):
     """Rigid contacts model."""
@@ -110,9 +88,37 @@ class RigidContacts(ContactModel):
     )
 
     terrain: jax_dataclasses.Static[Terrain] = dataclasses.field(
-        default_factory=FlatTerrain
+        default_factory=FlatTerrain.build
     )
 
+    @classmethod
+    def build(
+        cls: type[Self],
+        parameters: RigidContactsParams | None = None,
+        terrain: Terrain | None = None,
+        **kwargs,
+    ) -> Self:
+        """
+        Create a `RigidContacts` instance with specified parameters.
+
+        Args:
+            parameters: The parameters of the rigid contacts model.
+            terrain: The considered terrain.
+
+        Returns:
+            The `RigidContacts` instance.
+        """
+
+        if len(kwargs) != 0:
+            logging.debug(msg=f"Ignoring extra arguments: {kwargs}")
+
+        return cls(
+            parameters=(
+                parameters or cls.__dataclass_fields__["parameters"].default_factory()
+            ),
+            terrain=terrain or cls.__dataclass_fields__["terrain"].default_factory(),
+        )
+
     @staticmethod
     def detect_contacts(
         W_p_C: jtp.ArrayLike,
@@ -313,8 +319,10 @@ class RigidContacts(ContactModel):
                 js.ode.system_acceleration(
                     model=model,
                     data=data,
-                    joint_forces=references.joint_force_references(model=model),
                     link_forces=references.link_forces(model=model, data=data),
+                    joint_force_references=references.joint_force_references(
+                        model=model
+                    ),
                 )
             )