[FEATURE] SAP coupler and Hydroelastic contact model. (#1295)

---------

Co-authored-by: Alexis DUBURCQ <alexis.duburcq@gmail.com>
Co-authored-by: YilingQiao <49262224+YilingQiao@users.noreply.github.com>

Author: 3 people

genesis/engine/coupler.py

@@ -2,6 +2,8 @@
 import taichi as ti
 import torch
 
+import igl
+
 import genesis as gs
 import genesis.utils.element as eu
 import genesis.utils.geom as gu
@@ -11,6 +13,7 @@
 from genesis.utils.misc import to_gs_tensor
 
 from .base_entity import Entity
+from genesis.engine.coupler import SAPCoupler
 
 
 @ti.data_oriented
@@ -82,6 +85,9 @@ def __init__(self, scene, solver, material, morph, surface, idx, v_start=0, el_s
         unique_el = all_el[unique_idcs]
         self._surface_el_np = unique_el[cnt == 1].astype(gs.np_int)
 
+        if isinstance(self.sim.coupler, SAPCoupler):
+            self.compute_pressure_field()
+
         self.init_tgt_vars()
         self.init_ckpt()
 
@@ -371,6 +377,7 @@ def _add_to_solver(self, in_backward=False):
             mat_mu=self._material.mu,
             mat_lam=self._material.lam,
             mat_rho=self._material.rho,
+            mat_friction_mu=self._material.friction_mu,
             n_surfaces=self._n_surfaces,
             v_start=self._v_start,
             el_start=self._el_start,
@@ -382,6 +389,29 @@ def _add_to_solver(self, in_backward=False):
         )
         self.active = True
 
+    def compute_pressure_field(self):
+        """
+        Compute the pressure field for the FEM entity based on its tetrahedral elements.
+
+        For hydroelastic contact: https://drake.mit.edu/doxygen_cxx/group__hydroelastic__user__guide.html
+
+        Notes
+        -----
+        https://github.com/RobotLocomotion/drake/blob/master/geometry/proximity/make_mesh_field.cc
+        TODO: Add margin support
+        Drake's implementation of margin seems buggy.
+        """
+        init_positions = self.init_positions.cpu().numpy()
+        signed_distance, *_ = igl.signed_distance(init_positions, init_positions, self._surface_tri_np)
+        unsigned_distance = np.abs(signed_distance)
+        max_distance = np.max(unsigned_distance)
+        if max_distance < gs.EPS:
+            gs.raise_exception(
+                f"Pressure field max distance is too small: {max_distance}. "
+                "This might be due to a mesh having no internal vertices."
+            )
+        self.pressure_field_np = unsigned_distance / max_distance * self.material._hydroelastic_modulus  # normalize
+
     # ------------------------------------------------------------------------------------
     # ---------------------------- checkpoint and buffer ---------------------------------
     # ------------------------------------------------------------------------------------

genesis/engine/entities/fem_entity.py

@@ -23,19 +23,32 @@ class Base(Material):
         Poisson ratio, describing the material's volume change under stress. Default is 0.2.
     rho: float, optional
         Material density (kg/m^3). Default is 1000.
+    hydroelastic_modulus: float, optional
+        Hydroelastic modulus for hydroelastic contact. Default is 1e7.
+    friction_mu: float, optional
+        Friction coefficient. Default is 0.1.
+    hessian_invariant: bool, optional
+        If True, Hessian is computed only once. Default is False.
     """
 
     def __init__(
         self,
-        E=1e6,  # Young's modulus
-        nu=0.2,  # Poisson's ratio
-        rho=1000.0,  # density (kg/m^3)
+        E=1e6,
+        nu=0.2,
+        rho=1000.0,
+        hydroelastic_modulus=1e7,
+        friction_mu=0.1,
+        hessian_invariant=False,
     ):
         super().__init__()
 
         self._E = E
         self._nu = nu
         self._rho = rho
+        self._hydroelastic_modulus = hydroelastic_modulus
+        self._friction_mu = friction_mu
+        self.hessian_invariant = hessian_invariant
+        self.hessian_ready = False
 
         # lame parameters: https://github.com/taichi-dev/taichi_elements/blob/d19678869a28b09a32ef415b162e35dc929b792d/engine/mpm_solver.py#L203
         self._mu = E / (2.0 * (1.0 + nu))
@@ -84,3 +97,13 @@ def lam(self):
     def rho(self):
         """The rest density."""
         return self._rho
+
+    @property
+    def contact_stiffness(self):
+        """The contact stiffness."""
+        return self._hydroelastic_modulus
+
+    @property
+    def friction_mu(self):
+        """The friction coefficient."""
+        return self._friction_mu

genesis/engine/materials/FEM/base.py

@@ -27,6 +27,10 @@ class Elastic(Base):
         Poisson ratio, describing the material's volume change under stress. Default is 0.2.
     rho: float, optional
         Material density (kg/m^3). Default is 1000.
+    hydroelastic_modulus: float, optional
+        Hydroelastic modulus for hydroelastic contact. Default is 1e7.
+    friction_mu: float, optional
+        Friction coefficient. Default is 0.1.
     model: str, optional
         Constitutive model to use for stress computation. Options are:
         - 'linear': Linear elasticity model
@@ -39,23 +43,26 @@ def __init__(
         E=1e6,  # Young's modulus
         nu=0.2,  # Poisson's ratio
         rho=1000.0,  # density (kg/m^3)
+        hydroelastic_modulus=1e7,  # hydroelastic_modulus for hydroelastic contact
+        friction_mu=0.1,
         model="linear",
     ):
-        super().__init__(E, nu, rho)
+        super().__init__(E, nu, rho, hydroelastic_modulus, friction_mu)
 
         if model == "linear":
             self.update_stress = self.update_stress_linear
             self.compute_energy_gradient_hessian = self.compute_energy_gradient_hessian_linear
+            self.compute_energy_gradient = self.compute_energy_gradient_linear
             self.compute_energy = self.compute_energy_linear
-        elif model == "stable_neohookean":
-            self.update_stress = self.update_stress_stable_neohookean
-            self.compute_energy_gradient_hessian = self.compute_energy_gradient_hessian_stable_neohookean
-            self.compute_energy = self.compute_energy_stable_neohookean
-        elif model == "stable_neohooken":
-            gs.logger.warning("The 'stable_neohooken' model is deprecated. Use 'stable_neohookean' instead.")
+            self.hessian_invariant = True
+        elif model in ("stable_neohookean", "stable_neohooken"):
             self.update_stress = self.update_stress_stable_neohookean
             self.compute_energy_gradient_hessian = self.compute_energy_gradient_hessian_stable_neohookean
+            self.compute_energy_gradient = self.compute_energy_gradient_stable_neohookean
             self.compute_energy = self.compute_energy_stable_neohookean
+            self.hessian_invariant = False
+            if model == "stable_neohooken":
+                gs.logger.warning("The 'stable_neohooken' model is deprecated. Use 'stable_neohookean' instead.")
         else:
             gs.raise_exception(f"Unrecognized constitutive model: {model}")
 
@@ -147,6 +154,48 @@ def compute_energy_gradient_hessian_linear(self, mu, lam, J, F, actu, m_dir, i_e
         hessian_field[i_b, 1, 2, i_e][1, 2] = hessian_field[i_b, 2, 1, i_e][2, 1] = lam
         return energy, gradient
 
+    @ti.func
+    def compute_energy_gradient_linear(self, mu, lam, J, F, actu, m_dir, i_e, i_b):
+        """
+        Compute the energy, gradient for linear elasticity.
+
+        Parameters
+        ----------
+        mu: float
+            The first Lame parameter (shear modulus).
+        lam: float
+            The second Lame parameter (related to volume change).
+        J: float
+            The determinant of the deformation gradient F.
+        F: ti.Matrix
+            The deformation gradient matrix.
+        actu: ti.Matrix
+            The activation matrix (not used in linear elasticity).
+        m_dir: ti.Matrix
+            The material direction (not used in linear elasticity).
+
+        Returns
+        -------
+        energy: float
+            The computed energy.
+        gradient: ti.Matrix
+            The gradient of the energy with respect to the deformation gradient F.
+
+        Notes
+        -------
+        This implementation assumes small deformations and linear stress-strain relationship.
+        It is adapted from the HOBAKv1 implementation for linear elasticity:
+        https://github.com/theodorekim/HOBAKv1/blob/8420c51b795735d8fb912e0f8810f935d96fb636/src/Hyperelastic/Volume/LINEAR.cpp
+        """
+        I = ti.Matrix.identity(dt=gs.ti_float, n=3)
+        eps = 0.5 * (F + F.transpose()) - I
+        trEps = eps.trace()
+        energy = mu * eps.norm_sqr() + 0.5 * lam * trEps**2
+
+        gradient = 2.0 * mu * eps + lam * trEps * I
+
+        return energy, gradient
+
     @ti.func
     def compute_energy_linear(self, mu, lam, J, F, actu, m_dir):
         """
@@ -176,8 +225,7 @@ def compute_energy_linear(self, mu, lam, J, F, actu, m_dir):
         -------
         This implementation assumes small deformations and linear stress-strain relationship.
         It is adapted from the HOBAKv1 implementation for linear elasticity:
-        https://github.com/theodorekim/HOBAKv1/blob/main/src/Hyperelastic/Volume/LINEAR.cpp
-
+        https://github.com/theodorekim/HOBAKv1/blob/8420c51b795735d8fb912e0f8810f935d96fb636/src/Hyperelastic/Volume/LINEAR.cpp
         """
         I = ti.Matrix.identity(dt=gs.ti_float, n=3)
         eps = 0.5 * (F + F.transpose()) - I
@@ -228,6 +276,45 @@ def compute_energy_gradient_hessian_stable_neohookean(self, mu, lam, J, F, actu,
         """
         raise NotImplementedError("Hessian computation is not implemented for stable_neohookean model.")
 
+    @ti.func
+    def compute_energy_gradient_stable_neohookean(self, mu, lam, J, F, actu, m_dir, i_e, i_b):
+        """
+        Compute the energy, gradient for the stable Neo-Hookean model.
+
+        Parameters
+        ----------
+        mu: float
+            The first Lame parameter (shear modulus).
+        lam: float
+            The second Lame parameter (related to volume change).
+        J: float
+            The determinant of the deformation gradient F.
+        F: ti.Matrix
+            The deformation gradient matrix.
+        actu: ti.Matrix
+            The activation matrix (not used in stable Neo-Hookean).
+        m_dir: ti.Matrix
+            The material direction (not used in stable Neo-Hookean).
+
+        Returns
+        -------
+        energy: float
+            The computed energy.
+        gradient: ti.Matrix
+            The gradient of the energy with respect to the deformation gradient F.
+
+        Raises
+        -------
+        NotImplementedError
+            This implementation does not compute the Gradient for the stable Neo-Hookean model.
+
+        Notes
+        -------
+        This implementation is adapted from the HOBAKv1 stable Neo-Hookean model:
+        https://github.com/theodorekim/HOBAKv1/blob/8420c51b795735d8fb912e0f8810f935d96fb636/src/Hyperelastic/Volume/SNH.cpp
+        """
+        raise NotImplementedError("Gradient computation is not implemented for stable_neohookean model.")
+
     @ti.func
     def compute_energy_stable_neohookean(self, mu, lam, J, F, actu, m_dir):
         """
@@ -256,7 +343,7 @@ def compute_energy_stable_neohookean(self, mu, lam, J, F, actu, m_dir):
         Notes
         -------
         This implementation is adapted from the HOBAKv1 stable Neo-Hookean model:
-        https://github.com/theodorekim/HOBAKv1/blob/main/src/Hyperelastic/Volume/SNH.cpp
+        https://github.com/theodorekim/HOBAKv1/blob/8420c51b795735d8fb912e0f8810f935d96fb636/src/Hyperelastic/Volume/SNH.cpp
         """
         _lambda = lam + mu
         _alpha = 1.0 + mu / _lambda

genesis/engine/materials/FEM/elastic.py

@@ -8,6 +8,7 @@
 from genesis.options.solvers import (
     AvatarOptions,
     CouplerOptions,
+    SAPCouplerOptions,
     FEMOptions,
     MPMOptions,
     PBDOptions,
@@ -19,7 +20,7 @@
 )
 from genesis.repr_base import RBC
 
-from .coupler import Coupler
+from .coupler import Coupler, SAPCoupler
 from .entities import HybridEntity
 from .solvers import (
     AvatarSolver,
@@ -133,7 +134,10 @@ def __init__(
         self._active_solvers: list[Solver] = gs.List()
 
         # coupler
-        self._coupler = Coupler(self, self.coupler_options)
+        if isinstance(self.coupler_options, SAPCouplerOptions):
+            self._coupler = SAPCoupler(self, self.coupler_options)
+        else:
+            self._coupler = Coupler(self, self.coupler_options)
 
         # states
         self._queried_states = QueriedStates()