Merge pull request #105 from traversaro/fix103

Avoid to use as Static attributes classes that do not have a __eq__ method that returns a scalar bool

src/jaxsim/api/ode.py

@@ -123,7 +123,7 @@ def system_velocity_dynamics(
     # Initialize the derivative of the tangential deformation ṁ ∈ ℝ^{n_c × 3}.
     ṁ = jnp.zeros_like(data.state.soft_contacts.tangential_deformation).astype(float)
 
-    if model.physics_model.gc.body.size > 0:
+    if len(model.physics_model.gc.body) > 0:
         # Compute the position and linear velocities (mixed representation) of
         # all collidable points belonging to the robot.
         W_p_Ci, W_ṗ_Ci = Contact.collidable_point_kinematics(model=model, data=data)
@@ -140,7 +140,11 @@ def system_velocity_dynamics(
         # we don't need any coordinate transformation.
         W_f_Li_terrain = jax.vmap(
             lambda nc: (
-                jnp.vstack(jnp.equal(model.physics_model.gc.body, nc).astype(int))
+                jnp.vstack(
+                    jnp.equal(
+                        np.array(model.physics_model.gc.body, dtype=int), nc
+                    ).astype(int)
+                )
                 * W_f_Ci
             ).sum(axis=0)
         )(jnp.arange(model.number_of_links()))

src/jaxsim/parsers/descriptions/collision.py

@@ -50,6 +50,14 @@ def change_link(
             enabled=self.enabled,
         )
 
+    def __eq__(self, other):
+        retval = (
+            self.parent_link == other.parent_link
+            and (self.position == other.position).all()
+            and self.enabled == other.enabled
+        )
+        return retval
+
     def __str__(self):
         return (
             f"{self.__class__.__name__}("
@@ -93,6 +101,9 @@ class BoxCollision(CollisionShape):
 
     center: npt.NDArray
 
+    def __eq__(self, other):
+        return (self.center == other.center).all() and super().__eq__(other)
+
 
 @dataclasses.dataclass
 class SphereCollision(CollisionShape):
@@ -105,3 +116,6 @@ class SphereCollision(CollisionShape):
     """
 
     center: npt.NDArray
+
+    def __eq__(self, other):
+        return (self.center == other.center).all() and super().__eq__(other)

src/jaxsim/parsers/descriptions/link.py

@@ -38,6 +38,17 @@ class LinkDescription(JaxsimDataclass):
     def __hash__(self) -> int:
         return hash(self.__repr__())
 
+    def __eq__(self, other) -> bool:
+        return (
+            self.name == other.name
+            and self.mass == other.mass
+            and (self.inertia == other.inertia).all()
+            and self.index == other.index
+            and self.parent == other.parent
+            and (self.pose == other.pose).all()
+            and self.children == other.children
+        )
+
     @property
     def name_and_index(self) -> str:
         """

src/jaxsim/parsers/kinematic_graph.py

@@ -34,6 +34,11 @@ class RootPose(NamedTuple):
     root_position: npt.NDArray = np.zeros(3)
     root_quaternion: npt.NDArray = np.array([1.0, 0, 0, 0])
 
+    def __eq__(self, other):
+        return (self.root_position == other.root_position).all() and (
+            self.root_quaternion == other.root_quaternion
+        ).all()
+
 
 @dataclasses.dataclass(frozen=True)
 class KinematicGraph:

src/jaxsim/physics/algos/soft_contacts.py

@@ -58,7 +58,7 @@ def build_from_physics_model(
 
         return SoftContactsState.build(
             tangential_deformation=tangential_deformation,
-            number_of_collidable_points=physics_model.gc.body.size,
+            number_of_collidable_points=len(physics_model.gc.body),
         )
 
     @staticmethod
@@ -95,7 +95,7 @@ def valid(
         return check_valid_shape(
             what="tangential_deformation",
             shape=self.tangential_deformation.shape,
-            expected_shape=(3, physics_model.gc.body.size),
+            expected_shape=(3, len(physics_model.gc.body)),
             valid=True,
         )
 
@@ -237,7 +237,7 @@ def process_point_kinematics(
 
     # Process all the collidable points in parallel
     W_p_Ci, CW_v_WC = jax.vmap(process_point_kinematics)(
-        model.gc.point.T, model.gc.body
+        model.gc.point.T, np.array(model.gc.body, dtype=int)
     )
 
     return W_p_Ci.transpose(), CW_v_WC.transpose()

src/jaxsim/physics/algos/terrain.py

@@ -46,23 +46,21 @@ def height(self, x: float, y: float) -> float:
 
 @jax_dataclasses.pytree_dataclass
 class PlaneTerrain(Terrain):
-    plane_normal: jtp.Vector = jax_dataclasses.field(
-        default_factory=lambda: jnp.array([0, 0, 1.0])
-    )
+    plane_normal: list = jax_dataclasses.field(default_factory=lambda: [0, 0, 1.0])
 
     @staticmethod
-    def build(plane_normal: jtp.Vector) -> "PlaneTerrain":
+    def build(plane_normal: list) -> "PlaneTerrain":
         """
         Create a PlaneTerrain instance with a specified plane normal vector.
 
         Args:
-            plane_normal (jtp.Vector): The normal vector of the terrain plane.
+            plane_normal (list): The normal vector of the terrain plane.
 
         Returns:
             PlaneTerrain: A PlaneTerrain instance.
         """
 
-        return PlaneTerrain(plane_normal=jnp.array(plane_normal, dtype=float))
+        return PlaneTerrain(plane_normal=plane_normal)
 
     def height(self, x: float, y: float) -> float:
         """

src/jaxsim/physics/model/ground_contact.py

@@ -23,9 +23,7 @@ class GroundContact:
     """
 
     point: npt.NDArray = dataclasses.field(default_factory=lambda: jnp.array([]))
-    body: Static[npt.NDArray] = dataclasses.field(
-        default_factory=lambda: np.array([], dtype=int)
-    )
+    body: Static[list] = dataclasses.field(default_factory=lambda: [])
 
     @staticmethod
     def build_from(
@@ -42,9 +40,9 @@ def build_from(
 
         # Build the GroundContact attributes
         points = jnp.vstack([cp.position for cp in collidable_points]).T
-        link_index_of_points = np.array(
-            [links_dict[cp.parent_link.name].index for cp in collidable_points]
-        )
+        link_index_of_points = [
+            links_dict[cp.parent_link.name].index for cp in collidable_points
+        ]
 
         # Build the object
         gc = GroundContact(point=points, body=link_index_of_points)

src/jaxsim/simulation/ode.py

@@ -123,7 +123,7 @@ def dx_dt(
         ode_state.soft_contacts.tangential_deformation
     )
 
-    if physics_model.gc.body.size > 0:
+    if len(physics_model.gc.body) > 0:
         (
             contact_forces_links,
             tangential_deformation_dot,