Merge pull request #34 from ami-iit/refactor/soft_contacts

Rewrite soft-contacts algorithm

src/jaxsim/high_level/link.py

@@ -241,4 +241,4 @@ def add_com_external_force(
         )
 
     def in_contact(self) -> jtp.Bool:
-        return not jnp.allclose(self.external_force(), 0)
+        return self.parent_model.in_contact()[self.index()]

src/jaxsim/high_level/model.py

@@ -402,6 +402,41 @@ def joints(self, joint_names: List[str] = None) -> List[high_level.joint.Joint]:
 
         return [self._joints[name] for name in joint_names]
 
+    def in_contact(
+        self,
+        link_names: Optional[List[str]] = None,
+        terrain: Terrain = FlatTerrain(),
+    ) -> jtp.Vector:
+        """"""
+
+        link_names = link_names if link_names is not None else self.link_names()
+
+        if set(link_names) - set(self._links.keys()) != set():
+            raise ValueError("One or more link names are not part of the model")
+
+        from jaxsim.physics.algos.soft_contacts import collidable_points_pos_vel
+
+        W_p_Ci, _ = collidable_points_pos_vel(
+            model=self.physics_model,
+            q=self.data.model_state.joint_positions,
+            qd=self.data.model_state.joint_velocities,
+            xfb=self.data.model_state.xfb(),
+        )
+
+        terrain_height = jax.vmap(terrain.height)(W_p_Ci[0, :], W_p_Ci[1, :])
+
+        below_terrain = W_p_Ci[2, :] <= terrain_height
+
+        links_in_contact = jax.vmap(
+            lambda link_index: jnp.where(
+                self.physics_model.gc.body == link_index,
+                below_terrain,
+                jnp.zeros_like(below_terrain, dtype=bool),
+            ).any()
+        )(jnp.array([link.index() for link in self.links(link_names=link_names)]))
+
+        return links_in_contact
+
     # ==================
     # Vectorized methods
     # ==================
@@ -441,6 +476,14 @@ def joint_velocities(self, joint_names: List[str] = None) -> jtp.Vector:
             self._joint_indices(joint_names=joint_names)
         ]
 
+    def joint_generalized_forces_targets(
+        self, joint_names: List[str] = None
+    ) -> jtp.Vector:
+        if self.dofs() == 0 and (joint_names is None or len(joint_names) == 0):
+            return jnp.array([])
+
+        return self.data.model_input.tau[self._joint_indices(joint_names=joint_names)]
+
     def joint_limits(
         self, joint_names: List[str] = None
     ) -> Tuple[jtp.Vector, jtp.Vector]: