[Kamino] fix kamino contact capacity (#3732)

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: rubengrandia

newton/_src/solvers/kamino/_src/core/builder.py

@@ -1757,7 +1757,7 @@ def compute_required_contact_capacity(
                 else:
                     world_max_contacts[geom1.wid] += num_contacts
 
-        # Override the per-world maximum contacts if specified in the settings
+        # Cap per-world totals when a per-world maximum is specified
         if max_contacts_per_world is not None:
             for w in range(self.num_worlds):
                 world_max_contacts[w] = min(world_max_contacts[w], max_contacts_per_world)

newton/_src/solvers/kamino/_src/core/conversions.py

@@ -826,8 +826,8 @@ def compute_required_contact_capacity(
         max_contacts_per_pair: Optional maximum number of contacts to allocate per shape pair.
             If `None`, no per-pair limit is applied.
         max_contacts_per_world: Optional maximum number of contacts to allocate per world.
-            If `None`, no per-world limit is applied, otherwise it will
-            override the computed per-world requirements if it is larger.
+            If `None`, no per-world limit is applied, otherwise caps the computed
+            per-world requirements at this value.
 
     Returns:
         (model_required_contacts, world_required_contacts):
@@ -860,7 +860,7 @@ def compute_required_contact_capacity(
     )
     world_max_contacts = world_max_contacts_wp.numpy()
 
-    # Override the per-world maximum contacts if specified in the settings
+    # Cap per-world totals when a per-world maximum is specified
     if max_contacts_per_world is not None:
         world_max_contacts = np.minimum(world_max_contacts, max_contacts_per_world)
 

newton/_src/solvers/kamino/_src/core/shapes.py

@@ -620,96 +620,125 @@ def max_contacts_for_shape_pair(type_a: int, type_b: int) -> tuple[int, int]:
     Count the number of potential contact points for a collision pair in both
     directions of the collision pair (collisions from A to B and from B to A).
 
-    Inputs must be canonicalized such that the type of shape A is less than or equal to the type of shape B.
+    Shape types are canonicalized such that the type of shape A is less than or equal to the type of shape B.
 
     Args:
         type_a: First shape type as :class:`GeoType` integer value.
         type_b: Second shape type as :class:`GeoType` integer value.
 
     Returns:
         Number of contact points for collisions between A->B and B->A.
+        Each component independently bounds the contacts generated in its
+        direction. The reverse capacity is zero only when the pair's
+        narrow-phase implementation emits contacts in canonical order alone;
+        otherwise it reserves capacity for the reverse pass.
     """
     # Ensure the shape types are ordered canonically
     if type_a > type_b:
         type_a, type_b = type_b, type_a
 
-    if type_a == GeoType.SPHERE:
-        return 1, 0
+    return _max_contacts_for_shape_pair_impl(type_a, type_b)
+
+
+@wp.func
+def _max_contacts_for_shape_pair_impl(type_a: int, type_b: int) -> tuple[int, int]:
+    """
+    Return the contact capacity for a canonical shape pair without reordering.
+
+    This is used for testing purposes, asserting that type_a > type_b always returns (0, 0).
+    This enforces that the implementation doesn't accidentally specify an unreachable pair.
+    """
+    if type_a == GeoType.PLANE:
+        if type_b == GeoType.HFIELD:
+            return _MESH_CONVEX_MAX, 0
+        elif type_b == GeoType.SPHERE:
+            return 1, 0
+        elif type_b == GeoType.CAPSULE:
+            return 2, 0
+        elif type_b == GeoType.ELLIPSOID:
+            return 1, 0
+        elif type_b == GeoType.CYLINDER:
+            return 4, 0
+        elif type_b == GeoType.BOX:
+            return 4, 0
+        elif type_b == GeoType.MESH or type_b == GeoType.CONVEX_MESH:
+            return _MESH_CONVEX_MAX, 0
+        elif type_b == GeoType.CONE:
+            return 5, 0
+
+    elif type_a == GeoType.HFIELD:
+        if type_b == GeoType.MESH:
+            return _MESH_MESH_MAX, 0
+        elif type_b >= GeoType.HFIELD:
+            return _MESH_CONVEX_MAX, 0
+
+    elif type_a == GeoType.SPHERE:
+        if type_b == GeoType.MESH or type_b == GeoType.CONVEX_MESH:
+            return _MESH_CONVEX_MAX, 0
+        elif type_b >= GeoType.SPHERE:
+            return 1, 0
 
     elif type_a == GeoType.CAPSULE:
         if type_b == GeoType.CAPSULE:
-            return 2, 2
+            return 2, 0
         elif type_b == GeoType.ELLIPSOID:
-            return 8, 8
+            return 1, 0
         elif type_b == GeoType.CYLINDER:
-            return 4, 4
+            return 5, 0
         elif type_b == GeoType.BOX:
-            return 8, 8
+            return 5, 0
         elif type_b == GeoType.MESH or type_b == GeoType.CONVEX_MESH:
             return _MESH_CONVEX_MAX, 0
         elif type_b == GeoType.CONE:
-            return 4, 4
-        elif type_b == GeoType.PLANE:
-            return 8, 8
+            return 5, 0
 
     elif type_a == GeoType.ELLIPSOID:
         if type_b == GeoType.ELLIPSOID:
-            return 4, 4
+            return 1, 0
         elif type_b == GeoType.CYLINDER:
-            return 4, 4
+            return 1, 0
         elif type_b == GeoType.BOX:
-            return 8, 8
+            return 1, 0
         elif type_b == GeoType.MESH or type_b == GeoType.CONVEX_MESH:
             return _MESH_CONVEX_MAX, 0
         elif type_b == GeoType.CONE:
-            return 8, 8
-        elif type_b == GeoType.PLANE:
-            return 4, 4
+            return 1, 0
 
     elif type_a == GeoType.CYLINDER:
         if type_b == GeoType.CYLINDER:
-            return 4, 4
+            return 5, 0
         elif type_b == GeoType.BOX:
-            return 8, 8
+            return 5, 0
         elif type_b == GeoType.MESH or type_b == GeoType.CONVEX_MESH:
             return _MESH_CONVEX_MAX, 0
         elif type_b == GeoType.CONE:
-            return 4, 4
-        elif type_b == GeoType.PLANE:
-            return 6, 6
+            return 5, 0
 
     elif type_a == GeoType.BOX:
         if type_b == GeoType.BOX:
-            return 12, 12
+            return 8, 0
         elif type_b == GeoType.MESH or type_b == GeoType.CONVEX_MESH:
             return _MESH_CONVEX_MAX, 0
         elif type_b == GeoType.CONE:
-            return 8, 8
-        elif type_b == GeoType.PLANE:
-            return 12, 12
+            return 5, 0
 
-    elif type_a == GeoType.MESH or type_a == GeoType.CONVEX_MESH:
-        if type_b == GeoType.HFIELD:
+    elif type_a == GeoType.MESH:
+        if type_b == GeoType.MESH:
             return _MESH_MESH_MAX, 0
         elif type_b == GeoType.CONE:
             return _MESH_CONVEX_MAX, 0
-        elif type_b == GeoType.PLANE:
-            return _MESH_CONVEX_MAX, 0
-        else:
+        elif type_b == GeoType.CONVEX_MESH:
             return _MESH_MESH_MAX, 0
 
-    elif type_a == GeoType.HFIELD:
-        # Heightfield vs convex primitives
-        return _MESH_CONVEX_MAX, 0
-
     elif type_a == GeoType.CONE:
         if type_b == GeoType.CONE:
-            return 4, 4
-        elif type_b == GeoType.PLANE:
-            return 8, 8
+            return 5, 0
+        elif type_b == GeoType.CONVEX_MESH:
+            return _MESH_CONVEX_MAX, 0
 
-    elif type_a == GeoType.PLANE:
-        pass
+    elif type_a == GeoType.CONVEX_MESH:
+        if type_b == GeoType.CONVEX_MESH:
+            return _MESH_MESH_MAX, 0
 
     # unsupported type combination
     return 0, 0

newton/_src/solvers/kamino/_src/geometry/contacts.py

@@ -1352,17 +1352,18 @@ def convert_contacts_newton_to_kamino(
     # convert exceeds the capacity of the output contacts.
     if contacts_in.rigid_contact_max > contacts_out.model_max_contacts_host:
         msg.warning(
-            "Newton `rigid_contact_max` (%d) exceeds Kamino `model_max_contacts_host` (%d); contacts will be truncated.",
+            "Newton `rigid_contact_max` (%d) exceeds Kamino `model_max_contacts_host` (%d); active contacts may be truncated.",
             contacts_in.rigid_contact_max,
             contacts_out.model_max_contacts_host,
         )
 
     # Skip conversion of contact forces if not requested
     contacts_in_force = contacts_in.force if convert_forces else None
 
-    # Set the maximum number of contacts to convert to the smallest of the
-    # number of contacts detected and the maximum capacity of the output contacts.
-    max_converted_contacts = min(contacts_in.rigid_contact_max, contacts_out.model_max_contacts_host)
+    # Scan every Newton contact slot so saturated worlds cannot prevent later
+    # worlds from filling their own capacity. The kernel skips inactive slots
+    # and enforces per-world and model limits.
+    max_converted_contacts = contacts_in.rigid_contact_max
 
     # Clear the output contacts to reset the active contact
     # counts and reset contact data to sentinel values.
@@ -1374,9 +1375,7 @@ def convert_contacts_newton_to_kamino(
         restitution_mix_mode=MaterialMixMode.from_string(restitution_mix_mode),
     )
 
-    # Launch the conversion kernel to convert Newton contacts to Kamino's format
-    # NOTE: To reduce overhead, the total thread count is set to the smallest of
-    # the number of contacts detected and the maximum capacity of the output contacts.
+    # Launch the conversion kernel to convert Newton contacts to Kamino's format.
     wp.launch(
         kernel=_convert_contacts_newton_to_kamino,
         dim=max_converted_contacts,
@@ -1494,11 +1493,10 @@ def convert_contacts_kamino_to_newton(
             f"contacts_out.device={contacts_out.device}"
         )
 
-    # Issue warning to the user if the number of contacts to
-    # convert exceeds the capacity of the output contacts.
-    if contacts_in.model_max_contacts_host > contacts_out.rigid_contact_max:
+    # Only Kamino-generated contacts can be truncated during export.
+    if clear_output and contacts_in.model_max_contacts_host > contacts_out.rigid_contact_max:
         msg.warning(
-            "Kamino `model_max_contacts_host` (%d) exceeds Newton `rigid_contact_max` (%d); contacts will be truncated.",
+            "Kamino `model_max_contacts_host` (%d) exceeds Newton `rigid_contact_max` (%d); active contacts may be truncated.",
             contacts_in.model_max_contacts_host,
             contacts_out.rigid_contact_max,
         )

newton/_src/solvers/kamino/_src/geometry/detector.py

@@ -147,6 +147,86 @@ def __repr__(self):
         return self.__str__()
 
 
+###
+# Contact capacity helpers
+###
+
+# Conservative heuristics for the fallback allocation path when pair-based
+# capacity metadata is unavailable (``model_minimum_contacts == 0``).
+_EXPLICIT_CONTACTS_PER_PAIR = 10
+_DYNAMIC_CONTACTS_PER_COLLIDABLE = 20
+
+
+def _cap_world_contacts_at_total(world_max_contacts: list[int], max_total: int) -> list[int]:
+    """Scale per-world contact budgets down so their sum does not exceed ``max_total``."""
+    total = sum(world_max_contacts)
+    if total <= max_total:
+        return list(world_max_contacts)
+    if max_total <= 0:
+        return [0] * len(world_max_contacts)
+
+    capped = [0] * len(world_max_contacts)
+    remainders: list[tuple[float, int]] = []
+    assigned = 0
+    for i, count in enumerate(world_max_contacts):
+        scaled = count * max_total / total
+        floor = int(scaled)
+        capped[i] = floor
+        assigned += floor
+        remainders.append((scaled - floor, i))
+    for _, i in sorted(remainders, key=lambda item: item[0], reverse=True):
+        if assigned >= max_total:
+            break
+        capped[i] += 1
+        assigned += 1
+    return capped
+
+
+def _estimate_fallback_world_max_contacts(
+    model: ModelKamino,
+    config: CollisionDetectorConfig,
+) -> list[int]:
+    """Estimate per-world contact capacity from geometry when pair metadata is unavailable."""
+    num_worlds = model.size.num_worlds
+    world_max_contacts = [0] * num_worlds
+
+    if config.broadphase == "explicit" and model.geoms.collidable_pairs is not None:
+        pairs = model.geoms.collidable_pairs.numpy()
+        wid = model.geoms.wid.numpy()
+        for pair in pairs:
+            g0, g1 = int(pair[0]), int(pair[1])
+            world_id = int(wid[g0]) if wid[g0] >= 0 else int(wid[g1])
+            if 0 <= world_id < num_worlds:
+                world_max_contacts[world_id] += _EXPLICIT_CONTACTS_PER_PAIR
+    else:
+        wid = model.geoms.wid.numpy()
+        group = model.geoms.group.numpy()
+        for geom_id in range(len(wid)):
+            world_id = int(wid[geom_id])
+            if 0 <= world_id < num_worlds and group[geom_id] > 0:
+                world_max_contacts[world_id] += _DYNAMIC_CONTACTS_PER_COLLIDABLE
+
+    return world_max_contacts
+
+
+def _resolve_contact_capacity(
+    model: ModelKamino,
+    config: CollisionDetectorConfig,
+) -> tuple[int, list[int]]:
+    """Resolve model- and per-world contact budgets from geometry and config caps."""
+    if model.geoms.model_minimum_contacts > 0:
+        world_max_contacts = list(model.geoms.world_minimum_contacts)
+    else:
+        world_max_contacts = _estimate_fallback_world_max_contacts(model, config)
+
+    model_max_contacts = sum(world_max_contacts)
+    if model_max_contacts > config.max_contacts:
+        world_max_contacts = _cap_world_contacts_at_total(world_max_contacts, config.max_contacts)
+        model_max_contacts = sum(world_max_contacts)
+
+    return model_max_contacts, world_max_contacts
+
+
 ###
 # Interfaces
 ###
@@ -303,41 +383,22 @@ def finalize(
         # Configure the collision detection pipeline type based on the config
         self._pipeline_type = CollisionPipelineType.from_string(self._config.pipeline)
 
-        # TODO: FIX THIS SO THAT PER-WORLD MAX IS ACTUALLY BASED ON THE NUM OF COLLIDABLE
-        # GOEMS IN EACH WORLD, INSTEAD OF JUST DIVIDING THE MODEL MAX BY THE NUM WORLDS
-        # For collision pipeline, we don't multiply by per-pair factors since broad phase
-        # discovers pairs dynamically. Users can provide rigid_contact_max explicitly,
-        # otherwise it is estimated from shape count and broad phase mode.
-        if self._model.geoms.model_minimum_contacts > 0:
-            self._model_max_contacts = self._model.geoms.model_minimum_contacts
-            self._world_max_contacts = self._model.geoms.world_minimum_contacts
-        else:
-            # Estimate based on broad phase mode and available information
-            if self._config.broadphase == "explicit" and self._model.geoms.collidable_pairs is not None:
-                # For EXPLICIT mode, we know the maximum possible pairs
-                # Estimate ~10 contacts per shape pair (conservative for mesh-mesh contacts)
-                self._model_max_contacts = max(self._config.max_contacts, self._model.geoms.num_collidable_pairs * 10)
-            else:
-                # For NXN/SAP dynamic broad phase, estimate based on shape count
-                # Assume each shape contacts ~20 others on average (conservative estimate)
-                # This scales much better than O(N²) while still being safe
-                self._model_max_contacts = max(self._config.max_contacts, self._model.geoms.num_collidable * 20)
-
-            # Set the world max contacts to be the same for all worlds in the model
-            num_worlds = self._model.size.num_worlds
-            self._world_max_contacts = [self._model_max_contacts // num_worlds] * num_worlds
-
-        # Override per-world max contacts if config specifies it.
+        # Resolve contact capacity.
         if self._config.max_contacts_per_world is not None:
+            # Use the explicit per-world override when available.
             num_worlds = self._model.size.num_worlds
             per_world = self._config.max_contacts_per_world
             self._world_max_contacts = [per_world] * num_worlds
             self._model_max_contacts = per_world * num_worlds
+        else:
+            # Otherwise estimate per world from geometry.
+            # ``max_contacts`` caps the model total.
+            self._model_max_contacts, self._world_max_contacts = _resolve_contact_capacity(self._model, self._config)
 
         # Create the contacts interface which will allocate all contacts data arrays
         # NOTE: If internal allocations happen, then they will contain
         # the contacts generated by the collision detection pipelines
-        self._contacts = ContactsKamino(capacity=self._world_max_contacts, device=self._device)
+        self._contacts = ContactsKamino(capacity=list(self._world_max_contacts), device=self._device)
 
         # Proceed with allocations only if the model admits contacts, which
         # occurs when collision geometries defined in the builder and model