jaxsim 0.4.3.dev143__py3-none-any.whl → 0.4.3.dev159__py3-none-any.whl

jaxsim/rbda/contacts/common.py

@@ -14,44 +14,15 @@ 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.
+
+    Note:
+        This class is supposed to store only the tunable parameters of the contact
+        model, i.e. all those parameters that can be changed during runtime.
+        If the contact model has also static parameters, they should be stored
+        in the corresponding `ContactModel` class.
     """
 
     @classmethod
@@ -82,12 +53,33 @@ class ContactModel(JaxsimDataclass):
 
     Attributes:
         parameters: The parameters of the contact model.
-        terrain: The terrain model.
+        terrain: The considered terrain.
     """
 
     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,
@@ -99,7 +91,7 @@ class ContactModel(JaxsimDataclass):
         Compute the contact forces.
 
         Args:
-            model: The model to consider.
+            model: The robot model considered by the contact model.
             data: The data of the considered model.
 
         Returns:
@@ -109,6 +101,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
@@ -77,16 +78,6 @@ class RelaxedRigidContactsParams(ContactsParams):
         default_factory=lambda: jnp.array(0.5, dtype=float)
     )
 
-    # Maximum number of iterations
-    max_iterations: jtp.Int = dataclasses.field(
-        default_factory=lambda: jnp.array(50, dtype=int)
-    )
-
-    # Solver tolerance
-    tolerance: jtp.Float = dataclasses.field(
-        default_factory=lambda: jnp.array(1e-6, dtype=float)
-    )
-
     def __hash__(self) -> int:
         from jaxsim.utils.wrappers import HashedNumpyArray
 
@@ -102,8 +93,6 @@ class RelaxedRigidContactsParams(ContactsParams):
                 HashedNumpyArray(self.stiffness),
                 HashedNumpyArray(self.damping),
                 HashedNumpyArray(self.mu),
-                HashedNumpyArray(self.max_iterations),
-                HashedNumpyArray(self.tolerance),
             )
         )
 
@@ -124,8 +113,6 @@ class RelaxedRigidContactsParams(ContactsParams):
         stiffness: jtp.FloatLike | None = None,
         damping: jtp.FloatLike | None = None,
         mu: jtp.FloatLike | None = None,
-        max_iterations: jtp.IntLike | None = None,
-        tolerance: jtp.FloatLike | None = None,
     ) -> Self:
         """Create a `RelaxedRigidContactsParams` instance"""
 
@@ -151,45 +138,88 @@ class RelaxedRigidContactsParams(ContactsParams):
             and jnp.all(self.midpoint >= 0.0)
             and jnp.all(self.power >= 0.0)
             and jnp.all(self.mu >= 0.0)
-            and jnp.all(self.max_iterations > 0)
-            and jnp.all(self.tolerance > 0.0)
         )
 
 
 @jax_dataclasses.pytree_dataclass
-class RelaxedRigidContactsState(ContactsState):
-    """Class storing the state of the relaxed rigid contacts model."""
+class RelaxedRigidContacts(ContactModel):
+    """Relaxed rigid contacts model."""
 
-    def __eq__(self, other: RelaxedRigidContactsState) -> bool:
-        return hash(self) == hash(other)
+    parameters: RelaxedRigidContactsParams = dataclasses.field(
+        default_factory=RelaxedRigidContactsParams.build
+    )
 
-    @classmethod
-    def build(cls: type[Self]) -> Self:
-        """Create a `RelaxedRigidContactsState` instance"""
+    terrain: jax_dataclasses.Static[Terrain] = dataclasses.field(
+        default_factory=FlatTerrain.build
+    )
+
+    _solver_options_keys: jax_dataclasses.Static[tuple[str, ...]] = dataclasses.field(
+        default=("tol", "maxiter", "memory_size"), kw_only=True
+    )
+    _solver_options_values: jax_dataclasses.Static[tuple[Any, ...]] = dataclasses.field(
+        default=(1e-6, 50, 10), kw_only=True
+    )
 
-        return cls()
+    @property
+    def solver_options(self) -> dict[str, Any]:
+
+        return dict(
+            zip(
+                self._solver_options_keys,
+                self._solver_options_values,
+                strict=True,
+            )
+        )
 
     @classmethod
-    def zero(cls: type[Self], **kwargs) -> Self:
-        """Build a zero `RelaxedRigidContactsState` instance from a `JaxSimModel`."""
+    def build(
+        cls: type[Self],
+        parameters: RelaxedRigidContactsParams | None = None,
+        terrain: Terrain | None = None,
+        solver_options: dict[str, Any] | None = None,
+        **kwargs,
+    ) -> Self:
+        """
+        Create a `RelaxedRigidContacts` instance with specified parameters.
 
-        return cls.build()
+        Args:
+            parameters: The parameters of the rigid contacts model.
+            terrain: The considered terrain.
+            solver_options: The options to pass to the L-BFGS solver.
 
-    def valid(self, **kwargs) -> jtp.BoolLike:
-        return True
+        Returns:
+            The `RelaxedRigidContacts` instance.
+        """
 
+        if len(kwargs) != 0:
+            logging.debug(msg=f"Ignoring extra arguments: {kwargs}")
 
-@jax_dataclasses.pytree_dataclass
-class RelaxedRigidContacts(ContactModel):
-    """Relaxed rigid contacts model."""
+        # Get the default solver options.
+        default_solver_options = dict(
+            zip(cls._solver_options_keys, cls._solver_options_values, strict=True)
+        )
 
-    parameters: RelaxedRigidContactsParams = dataclasses.field(
-        default_factory=RelaxedRigidContactsParams
-    )
+        # Create the solver options to set by combining the default solver options
+        # with the user-provided solver options.
+        solver_options = default_solver_options | (solver_options or {})
 
-    terrain: jax_dataclasses.Static[Terrain] = dataclasses.field(
-        default_factory=FlatTerrain
-    )
+        # Make sure that the solver options are hashable.
+        # We need to check this because the solver options are static.
+        try:
+            hash(tuple(solver_options.values()))
+        except TypeError as exc:
+            raise ValueError(
+                "The values of the solver options must be hashable."
+            ) from exc
+
+        return cls(
+            parameters=(
+                parameters or cls.__dataclass_fields__["parameters"].default_factory()
+            ),
+            terrain=terrain or cls.__dataclass_fields__["terrain"].default_factory(),
+            _solver_options_keys=tuple(solver_options.keys()),
+            _solver_options_values=tuple(solver_options.values()),
+        )
 
     @jax.jit
     def compute_contact_forces(
@@ -351,17 +381,23 @@ class RelaxedRigidContacts(ContactModel):
             + D[:, jnp.newaxis] * velocity
         ).flatten()
 
-        # Compute the 3D linear force in C[W] frame
+        # Get the solver options.
+        solver_options = self.solver_options
+
+        # Extract the options corresponding to the convergence criteria.
+        # All the remaining options are passed to the solver.
+        tol = solver_options.pop("tol")
+        maxiter = solver_options.pop("maxiter")
+
+        # Compute the 3D linear force in C[W] frame.
         CW_f_Ci, _ = run_optimization(
             init_params=init_params,
             A=A,
             b=b,
-            maxiter=self.parameters.max_iterations,
-            opt=optax.lbfgs(
-                memory_size=10,
-            ),
+            maxiter=maxiter,
+            opt=optax.lbfgs(**solver_options),
             fun=objective,
-            tol=self.parameters.tolerance,
+            tol=tol,
         )
 
         CW_f_Ci = CW_f_Ci.reshape((-1, 3))

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
@@ -79,39 +80,95 @@ class RigidContactsParams(ContactsParams):
 
 
 @jax_dataclasses.pytree_dataclass
-class RigidContactsState(ContactsState):
-    """Class storing the state of the rigid contacts model."""
+class RigidContacts(ContactModel):
+    """Rigid contacts model."""
 
-    def __eq__(self, other: RigidContactsState) -> bool:
-        return hash(self) == hash(other)
+    parameters: RigidContactsParams = dataclasses.field(
+        default_factory=RigidContactsParams
+    )
 
-    @classmethod
-    def build(cls: type[Self]) -> Self:
-        """Create a `RigidContactsState` instance"""
+    terrain: jax_dataclasses.Static[Terrain] = dataclasses.field(
+        default_factory=FlatTerrain.build
+    )
 
-        return cls()
+    regularization_delassus: jax_dataclasses.Static[float] = dataclasses.field(
+        default=1e-6, kw_only=True
+    )
+
+    _solver_options_keys: jax_dataclasses.Static[tuple[str, ...]] = dataclasses.field(
+        default=("solver_tol",), kw_only=True
+    )
+    _solver_options_values: jax_dataclasses.Static[tuple[Any, ...]] = dataclasses.field(
+        default=(1e-3,), kw_only=True
+    )
+
+    @property
+    def solver_options(self) -> dict[str, Any]:
+
+        return dict(
+            zip(
+                self._solver_options_keys,
+                self._solver_options_values,
+                strict=True,
+            )
+        )
 
     @classmethod
-    def zero(cls: type[Self], **kwargs) -> Self:
-        """Build a zero `RigidContactsState` instance from a `JaxSimModel`."""
+    def build(
+        cls: type[Self],
+        parameters: RigidContactsParams | None = None,
+        terrain: Terrain | None = None,
+        regularization_delassus: jtp.FloatLike | None = None,
+        solver_options: dict[str, Any] | None = None,
+        **kwargs,
+    ) -> Self:
+        """
+        Create a `RigidContacts` instance with specified parameters.
 
-        return cls.build()
+        Args:
+            parameters: The parameters of the rigid contacts model.
+            terrain: The considered terrain.
+            regularization_delassus:
+                The regularization term to add to the diagonal of the Delassus matrix.
+            solver_options: The options to pass to the QP solver.
 
-    def valid(self, **kwargs) -> jtp.BoolLike:
-        return True
+        Returns:
+            The `RigidContacts` instance.
+        """
 
+        if len(kwargs) != 0:
+            logging.debug(msg=f"Ignoring extra arguments: {kwargs}")
 
-@jax_dataclasses.pytree_dataclass
-class RigidContacts(ContactModel):
-    """Rigid contacts model."""
+        # Get the default solver options.
+        default_solver_options = dict(
+            zip(cls._solver_options_keys, cls._solver_options_values, strict=True)
+        )
 
-    parameters: RigidContactsParams = dataclasses.field(
-        default_factory=RigidContactsParams
-    )
+        # Create the solver options to set by combining the default solver options
+        # with the user-provided solver options.
+        solver_options = default_solver_options | (solver_options or {})
 
-    terrain: jax_dataclasses.Static[Terrain] = dataclasses.field(
-        default_factory=FlatTerrain
-    )
+        # Make sure that the solver options are hashable.
+        # We need to check this because the solver options are static.
+        try:
+            hash(tuple(solver_options.values()))
+        except TypeError as exc:
+            raise ValueError(
+                "The values of the solver options must be hashable."
+            ) from exc
+
+        return cls(
+            parameters=(
+                parameters or cls.__dataclass_fields__["parameters"].default_factory()
+            ),
+            terrain=terrain or cls.__dataclass_fields__["terrain"].default_factory(),
+            regularization_delassus=float(
+                regularization_delassus
+                or cls.__dataclass_fields__["regularization_delassus"].default
+            ),
+            _solver_options_keys=tuple(solver_options.keys()),
+            _solver_options_values=tuple(solver_options.values()),
+        )
 
     @staticmethod
     def detect_contacts(
@@ -224,8 +281,6 @@ class RigidContacts(ContactModel):
         *,
         link_forces: jtp.MatrixLike | None = None,
         joint_force_references: jtp.VectorLike | None = None,
-        regularization_term: jtp.FloatLike = 1e-6,
-        solver_tol: jtp.FloatLike = 1e-3,
     ) -> tuple[jtp.Vector, tuple[Any, ...]]:
         """
         Compute the contact forces.
@@ -238,10 +293,6 @@ class RigidContacts(ContactModel):
                 expressed in the same representation of data.
             joint_force_references:
                 Optional `(n_joints,)` vector of joint forces.
-            regularization_term:
-                The regularization term to add to the diagonal of the Delassus
-                matrix for better numerical conditioning.
-            solver_tol: The convergence tolerance to consider in the QP solver.
 
         Returns:
             A tuple containing the contact forces.
@@ -290,10 +341,11 @@ class RigidContacts(ContactModel):
             terrain_height=terrain_height,
         )
 
+        # Compute the Delassus matrix.
         delassus_matrix = RigidContacts._delassus_matrix(M=M, J_WC=J_WC)
 
-        # Add regularization for better numerical conditioning
-        delassus_matrix = delassus_matrix + regularization_term * jnp.eye(
+        # Add regularization for better numerical conditioning.
+        delassus_matrix = delassus_matrix + self.regularization_delassus * jnp.eye(
             delassus_matrix.shape[0]
         )
 
@@ -353,7 +405,7 @@ class RigidContacts(ContactModel):
 
         # Solve the optimization problem
         solution, *_ = qpax.solve_qp(
-            Q=Q, q=q, A=A, b=b, G=G, h=h_bounds, solver_tol=solver_tol
+            Q=Q, q=q, A=A, b=b, G=G, h=h_bounds, **self.solver_options
         )
 
         f_C_lin = solution.reshape(-1, 3)