Save collision shape to link transform

Author: flferretti

src/jaxsim/api/kin_dyn_parameters.py

@@ -783,24 +783,41 @@ class ContactParameters(JaxsimDataclass):
 
     Attributes:
         body:
-            A tuple of integers representing, for each collidable point, the index of
-            the body (link) to which it is rigidly attached to.
-        point:
-            The translations between the link frame and the collidable point, expressed
-            in the coordinates of the parent link frame.
-        enabled:
-            A tuple of booleans representing, for each collidable point, whether it is
-            enabled or not in contact models.
+            A tuple of integers representing, for each collision shape, the index of
+            the link to which it is rigidly attached to.
+        transform:
+            The 4x4 homogeneous transformation matrices representing the pose of each
+            collision shape with respect to the parent link frame.
+        shape_size:
+            The size parameters of each collidable shape.
+        shape_type:
+            The type of each collidable shape (sphere, box, cylinder, etc.).
 
     Note:
         Contrarily to LinkParameters and JointParameters, this class is not meant
         to be created with vmap. This is because the `body` attribute must be `Static`.
     """
 
-    center: jtp.Vector = dataclasses.field(default_factory=lambda: jnp.array([]))
+    body: Static[tuple[int, ...]] = dataclasses.field(default_factory=tuple)
+
+    transform: jtp.Matrix = dataclasses.field(default_factory=lambda: jnp.array([]))
     shape_size: jtp.Vector = dataclasses.field(default_factory=lambda: jnp.array([]))
     shape_type: jtp.Vector = dataclasses.field(default_factory=lambda: jnp.array([]))
 
+    @property
+    def center(self) -> jtp.Array:
+        """Extract translation vectors from transformation matrices."""
+        if self.transform.size == 0:
+            return jnp.array([])
+        return self.transform[:, :3, 3]
+
+    @property
+    def orientation(self) -> jtp.Array:
+        """Extract rotation matrices from transformation matrices."""
+        if self.transform.size == 0:
+            return jnp.array([])
+        return self.transform[:, :3, :3]
+
     @staticmethod
     def build_from(model_description: ModelDescription) -> ContactParameters:
         """
@@ -816,7 +833,12 @@ def build_from(model_description: ModelDescription) -> ContactParameters:
         if len(model_description.collision_shapes) == 0:
             return ContactParameters()
 
-        shape_types, shape_sizes, centers = [], [], []
+        shape_types, shape_sizes, transforms, parent_link_indices = (
+            [],
+            [],
+            [],
+            [],
+        )
 
         # Assume the link_parameters and the collision_shapes are in the same order.
         for collision in model_description.collision_shapes:
@@ -828,11 +850,17 @@ def build_from(model_description: ModelDescription) -> ContactParameters:
 
             shape_sizes.append(collision.size.squeeze())
 
-            centers.append(collision.center)
+            transforms.append(collision.transform)
+
+            # Get the parent link index for this collision shape.
+            parent_link_indices.append(
+                model_description.links_dict[collision.parent_link].index
+            )
 
         # Build the ContactParameters object.
         return ContactParameters(
-            center=jnp.array(centers, dtype=float),
+            body=tuple(parent_link_indices),
+            transform=jnp.array(transforms, dtype=float),
             shape_type=jnp.array(shape_types, dtype=int),
             shape_size=jnp.array(shape_sizes, dtype=float),
         )

src/jaxsim/parsers/descriptions/collision.py

@@ -3,6 +3,8 @@
 import dataclasses
 from abc import ABC
 
+import numpy as np
+
 import jaxsim.typing as jtp
 
 
@@ -15,16 +17,16 @@ class CollisionShape(ABC):
     It is not intended to be instantiated directly.
     """
 
-    center: jtp.VectorLike
     size: jtp.VectorLike
     parent_link: str
+    transform: jtp.MatrixLike = dataclasses.field(default_factory=lambda: np.eye(4))
 
     def __hash__(self) -> int:
         return hash(
             (
-                hash(tuple(self.center.tolist())),
                 hash(tuple(self.size.tolist())),
                 hash(self.parent_link),
+                hash(tuple(self.transform.flatten().tolist())),
             )
         )
 
@@ -35,59 +37,33 @@ def __eq__(self, other: CollisionShape) -> bool:
 
         return hash(self) == hash(other)
 
+    @property
+    def center(self) -> jtp.Vector:
+        """Extract the translation from the transformation matrix."""
+        return self.transform[:3, 3]
+
+    @property
+    def orientation(self) -> jtp.Matrix:
+        """Extract the rotation matrix from the transformation matrix."""
+        return self.transform[:3, :3]
+
 
 @dataclasses.dataclass
 class BoxCollision(CollisionShape):
     """
     Represents a box-shaped collision shape.
     """
 
-    @property
-    def x(self) -> float:
-        return self.size[0]
-
-    @property
-    def y(self) -> float:
-        return self.size[1]
-
-    @property
-    def z(self) -> float:
-        return self.size[2]
-
-    @x.setter
-    def x(self, value: float) -> None:
-        self.size[0] = value
-
-    @y.setter
-    def y(self, value: float) -> None:
-        self.size[1] = value
-
-    @z.setter
-    def z(self, value: float) -> None:
-        self.size[2] = value
-
 
 @dataclasses.dataclass
 class SphereCollision(CollisionShape):
     """
     Represents a spherical collision shape.
     """
 
-    @property
-    def radius(self) -> float:
-        return self.size[0]
-
 
 @dataclasses.dataclass
 class CylinderCollision(CollisionShape):
     """
     Represents a cylindrical collision shape.
     """
-
-    @property
-    def radius(self) -> float:
-        return self.size[0]
-
-    @property
-    def height(self) -> float:
-        return self.size[1]

src/jaxsim/parsers/rod/parser.py

@@ -346,7 +346,7 @@ def extract_model_data(
             # Fill with unsupported collision shape
             collisions.append(
                 descriptions.collision.CollisionShape(
-                    center=jnp.array([0.0, 0.0, 0.0]),
+                    transform=jnp.eye(4),
                     size=jnp.array([0.0, 0.0, 0.0]),
                     parent_link=link.name,
                 )

src/jaxsim/parsers/rod/utils.py

@@ -105,11 +105,9 @@ def create_box_collision(
 
     H = collision.pose.transform() if collision.pose is not None else np.eye(4)
 
-    center = H[:3, 3]
-
     return descriptions.BoxCollision(
         size=np.array([x, y, z]),
-        center=center,
+        transform=H,
         parent_link=link_description.name,
     )
 
@@ -132,11 +130,9 @@ def create_sphere_collision(
 
     H = collision.pose.transform() if collision.pose is not None else np.eye(4)
 
-    center_wrt_link = (H @ np.hstack([0, 0, 0, 1.0]))[0:-1]
-
     return descriptions.SphereCollision(
         size=np.array([r] * 3),
-        center=center_wrt_link,
+        transform=H,
         parent_link=link_description.name,
     )
 
@@ -160,10 +156,8 @@ def create_cylinder_collision(
 
     H = collision.pose.transform() if collision.pose is not None else np.eye(4)
 
-    center_wrt_link = (H @ np.hstack([0, 0, 0, 1.0]))[0:-1]
-
     return descriptions.CylinderCollision(
         size=np.array([r, l, 0]),
-        center=center_wrt_link,
+        transform=H,
         parent_link=link_description.name,
     )

src/jaxsim/rbda/contacts/common.py

@@ -32,7 +32,7 @@
 def compute_penetration_data(
     model: js.model.JaxSimModel,
     *,
-    shape_offset: jtp.Vector,
+    shape_transform: jtp.Matrix,
     shape_type: CollidableShapeType,
     shape_size: jtp.Vector,
     link_transforms: jtp.Matrix,
@@ -43,7 +43,7 @@ def compute_penetration_data(
 
     Args:
         model: The model to consider.
-        shape_offset: The offset of the collidable shape with respect to the link frame.
+        shape_transform: The 4x4 transform of the collidable shape with respect to the link frame.
         shape_type: The type of the collidable shape.
         shape_size: The size parameters of the collidable shape.
         link_transforms: The transforms from the world frame to each link.
@@ -55,10 +55,11 @@ def compute_penetration_data(
         expressed in mixed representation.
     """
 
-    W_H_L, W_ṗ_L = link_transforms, link_velocities
+    W_H_L, W_ṗ_L = link_transforms, link_velocities
 
-    # Offset the collision shape origin.
-    W_H_L = W_H_L.at[:3, 3].set(W_H_L[:3, 3] + shape_offset @ W_H_L[:3, :3].T)
+    # Apply the collision shape transform.
+    # This computes W_H_S where S is the collision shape frame.
+    W_H_S = W_H_L @ shape_transform
 
     # Pre-process the position and the linear velocity of the collidable point.
     # Note that we consider 3 candidate contact points also for spherical shapes,
@@ -69,7 +70,7 @@ def compute_penetration_data(
         (box_plane, cylinder_plane, sphere_plane),
         model.terrain,
         shape_size,
-        W_H_L,
+        W_H_S,
     )
 
     W_p_C = W_H_C[:, :3, 3]

tests/test_api_contact.py

@@ -97,16 +97,24 @@ def test_contact_jacobian_derivative(
         velocity_representation=velocity_representation,
     )
 
-    W_H_L = data._link_transforms
+    body_indices = np.array(model.kin_dyn_parameters.contact_parameters.body)
+
+    # Get link transforms for each collision shape
+    W_H_L = data._link_transforms[body_indices]
+
+    # Get contact point positions (shape: num_collision_shapes, 3, 3)
     W_p_C = js.contact.contact_point_positions(model=model, data=data)
 
-    # Vectorize over the 3 points for one link
-    transform_points = jax.vmap(
-        lambda H, p: H @ jnp.hstack([p, 1.0]), in_axes=(None, 0)
-    )
+    # Transform contact points from world to link frame
+    # For each collision shape, transform its 3 contact points
+    def transform_to_link_frame(W_H_L_i, W_p_Ci):
+        """Transform 3 contact points from world to link frame."""
 
-    # Vectorize over the links
-    L_p_Ci = jax.vmap(transform_points, in_axes=(0, 0))(W_H_L, W_p_C)[..., :3]
+        L_H_W = jnp.linalg.inv(W_H_L_i)
+        return jax.vmap(lambda p: (L_H_W @ jnp.hstack([p, 1.0]))[:3])(W_p_Ci)
+
+    # Apply to all collision shapes: shape (num_collision_shapes, 3, 3)
+    L_p_Ci = jax.vmap(transform_to_link_frame)(W_H_L, W_p_C)
 
     # =====
     # Tests
@@ -115,16 +123,15 @@ def test_contact_jacobian_derivative(
     # Load the model in ROD.
     rod_model = rod.Sdf.load(sdf=model.built_from).model
 
-    # Add dummy frames on the contact shapes.
-
-    for idx, link_name, points in zip(
-        np.arange(model.number_of_links()), model.link_names(), L_p_Ci, strict=True
+    for shape_idx, (link_idx, points) in enumerate(
+        zip(body_indices, L_p_Ci, strict=True)
     ):
-        # points: shape (3, 3) for this link
+        link_name = model.link_names()[link_idx]
+
         for j, p in enumerate(points):
             rod_model.add_frame(
                 frame=rod.Frame(
-                    name=f"contact_shape_{idx}_{j}",
+                    name=f"contact_shape_{shape_idx}_{j}",
                     attached_to=link_name,
                     pose=rod.Pose(
                         relative_to=link_name,
@@ -154,11 +161,12 @@ def test_contact_jacobian_derivative(
     )
 
     # Extract the indexes of the frames attached to the contact shapes.
+    num_collision_shapes = len(model.kin_dyn_parameters.contact_parameters.body)
     frame_idxs = js.frame.names_to_idxs(
         model=model_with_frames,
         frame_names=(
-            f"contact_shape_{idx}_{j}"
-            for idx in np.arange(model.number_of_links())
+            f"contact_shape_{shape_idx}_{j}"
+            for shape_idx in range(num_collision_shapes)
             for j in range(3)
         ),
     )