[MISC] Add MPR vs GJK benchmark. (#1988)

duburcqa · web-flow · commit e8801d4648fc · 2025-11-16T17:21:54.000+01:00
* Minor cleanup.
* Add MPR vs GJK benchmark.
diff --git a/examples/collision/pyramid.py b/examples/collision/pyramid.py
@@ -31,7 +31,7 @@ def main():
     # create pyramid of boxes
     box_size = 0.25
     box_spacing = (1.0 - 1e-3 + 0.1 * (args.pile_type == "static")) * box_size
-    box_pos_offset = (0 - 0.5, 1, 0.0) + 0.5 * np.array([box_size, box_size, box_size])
+    box_pos_offset = (-0.5, 1, 0.0) + 0.5 * np.array([box_size, box_size, box_size])
     boxes = {}
     for i in range(args.num_cubes):
         for j in range(args.num_cubes - i):
diff --git a/examples/collision/tower.py b/examples/collision/tower.py
@@ -1,7 +1,7 @@
-import numpy as np
-import genesis as gs
 import argparse
 
+import genesis as gs
+
 
 def main():
     parser = argparse.ArgumentParser()
@@ -34,11 +34,11 @@ def main():
     num_stacks = 50
     for i in range(num_stacks):
         if i % 2 == 0:  # horizontal stack
-            box_size = np.array([box_width, box_length, box_height])
+            box_size = (box_width, box_length, box_height)
             box_pos_0 = (-0.4 * box_length, 0, i * (box_height - 1e-3) + 0.5 * box_height)
             box_pos_1 = (+0.4 * box_length, 0, i * (box_height - 1e-3) + 0.5 * box_height)
         else:  # vertical stack
-            box_size = np.array([box_length, box_width, box_height])
+            box_size = (box_length, box_width, box_height)
             box_pos_0 = (0, -0.4 * box_length, i * (box_height - 1e-3) + 0.5 * box_height)
             box_pos_1 = (0, +0.4 * box_length, i * (box_height - 1e-3) + 0.5 * box_height)
         for box_pos in (box_pos_0, box_pos_1):
@@ -68,8 +68,7 @@ def main():
             ),
         )
     elif object_type == "cylinder":
-        cylinder_radius = 2.0
-        cylinder_height = 1.0
+        cylinder_radius, cylinder_height = 2.0, 1.0
         scene.add_entity(
             morph=gs.morphs.Cylinder(
                 radius=cylinder_radius,
diff --git a/genesis/engine/solvers/rigid/collider_decomp.py b/genesis/engine/solvers/rigid/collider_decomp.py
@@ -2370,9 +2370,10 @@ def func_convex_convex_contact(
                         for i_mpr in range(2):
                             if i_mpr == 1:
                                 # Try without warm-start if no contact was detected using it.
-                                # When penetration depth is very shallow, MPR may wrongly classify two geometries as not in
-                                # contact while they actually are. This helps to improve contact persistence without increasing
-                                # much the overall computational cost since the fallback should not be triggered very often.
+                                # When penetration depth is very shallow, MPR may wrongly classify two geometries as not
+                                # in contact while they actually are. This helps to improve contact persistence without
+                                # increasing much the overall computational cost since the fallback should not be
+                                # triggered very often.
                                 is_mpr_guess_direction_available = (ti.abs(normal_ws) > EPS).any()
                                 if (i_detection == 0) and not is_col and is_mpr_guess_direction_available:
                                     normal_ws = ti.Vector.zero(gs.ti_float, 3)
@@ -2401,16 +2402,17 @@ def func_convex_convex_contact(
                         # Fallback on SDF if collision is detected by MPR but no collision direction was cached and the
                         # initial penetration is already quite large, because the contact information provided by MPR
                         # may be unreliable in such a case.
-                        if is_col and penetration > tolerance:
-                            # MPR cannot handle collision detection for fully enclosed geometries. Falling back to SDF.
-                            # Note that SDF does not take into account to direction of interest. As such, it cannot be
-                            # used reliably for anything else than the point of deepest penetration.
-                            prefer_sdf = (
-                                collider_info.mc_tolerance[None] * penetration
-                                >= collider_info.mpr_to_sdf_overlap_ratio[None] * tolerance
-                            )
-                            if prefer_sdf or not is_mpr_guess_direction_available:
-                                try_sdf = True
+                        if ti.static(collider_static_config.ccd_algorithm == CCD_ALGORITHM_CODE.MPR):
+                            if is_col and penetration > tolerance:
+                                # MPR cannot handle collision detection for fully enclosed geometries. Falling back to
+                                # SDF. Note that SDF does not take into account to direction of interest. As such, it
+                                # cannot be used reliably for anything else than the point of deepest penetration.
+                                prefer_sdf = (
+                                    collider_info.mc_tolerance[None] * penetration
+                                    >= collider_info.mpr_to_sdf_overlap_ratio[None] * tolerance
+                                )
+                                if prefer_sdf or not is_mpr_guess_direction_available:
+                                    try_sdf = True
 
                     ### GJK, MJ_GJK
                     elif ti.static(
@@ -2608,7 +2610,7 @@ def func_convex_convex_contact(
                         errno,
                     )
                     if multi_contact:
-                        # perturb geom_a around two orthogonal axes to find multiple contacts
+                        # Perturb geom_a around two orthogonal axes to find multiple contacts
                         axis_0, axis_1 = func_contact_orthogonals(
                             i_ga,
                             i_gb,
diff --git a/genesis/engine/solvers/rigid/constraint_solver_decomp.py b/genesis/engine/solvers/rigid/constraint_solver_decomp.py
@@ -560,9 +560,10 @@ def func_equality_connect(
     for i_3 in range(3):
         n_con = ti.atomic_add(constraint_state.n_constraints[i_b], 1)
         ti.atomic_add(constraint_state.n_constraints_equality[i_b], 1)
+        con_n_relevant_dofs = 0
 
         if ti.static(static_rigid_sim_config.sparse_solve):
-            for i_d_ in range(collider_state.jac_n_relevant_dofs[n_con, i_b]):
+            for i_d_ in range(constraint_state.jac_n_relevant_dofs[n_con, i_b]):
                 i_d = constraint_state.jac_relevant_dofs[n_con, i_d_, i_b]
                 constraint_state.jac[n_con, i_d, i_b] = gs.ti_float(0.0)
         else:
@@ -1231,9 +1232,8 @@ def func_nt_hessian_direct(
     n_entities = entities_info.n_links.shape[0]
 
     # H = M + J'*D*J
-    for i_d1 in range(n_dofs):
-        for i_d2 in range(n_dofs):
-            constraint_state.nt_H[i_d1, i_d2, i_b] = gs.ti_float(0.0)
+    for i_d1, i_d2 in ti.ndrange(n_dofs, n_dofs):
+        constraint_state.nt_H[i_d1, i_d2, i_b] = gs.ti_float(0.0)
 
     if ti.static(static_rigid_sim_config.sparse_solve):
         for i_c in range(constraint_state.n_constraints[i_b]):
@@ -1251,17 +1251,16 @@ def func_nt_hessian_direct(
                             * constraint_state.active[i_c, i_b]
                         )
     else:
-        for i_c in range(constraint_state.n_constraints[i_b]):
-            for i_d1 in range(n_dofs):
-                if ti.abs(constraint_state.jac[i_c, i_d1, i_b]) > EPS:
-                    for i_d2 in range(i_d1 + 1):
-                        constraint_state.nt_H[i_d1, i_d2, i_b] = (
-                            constraint_state.nt_H[i_d1, i_d2, i_b]
-                            + constraint_state.jac[i_c, i_d2, i_b]
-                            * constraint_state.jac[i_c, i_d1, i_b]
-                            * constraint_state.efc_D[i_c, i_b]
-                            * constraint_state.active[i_c, i_b]
-                        )
+        for i_d1, i_c in ti.ndrange(n_dofs, constraint_state.n_constraints[i_b]):
+            if ti.abs(constraint_state.jac[i_c, i_d1, i_b]) > EPS:
+                for i_d2 in range(i_d1 + 1):
+                    constraint_state.nt_H[i_d1, i_d2, i_b] = (
+                        constraint_state.nt_H[i_d1, i_d2, i_b]
+                        + constraint_state.jac[i_c, i_d2, i_b]
+                        * constraint_state.jac[i_c, i_d1, i_b]
+                        * constraint_state.efc_D[i_c, i_b]
+                        * constraint_state.active[i_c, i_b]
+                    )
 
     for i_d1 in range(n_dofs):
         for i_d2 in range(i_d1 + 1, n_dofs):
@@ -1404,6 +1403,7 @@ def func_solve(
 ):
     _B = constraint_state.grad.shape[1]
     n_dofs = constraint_state.grad.shape[0]
+
     ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
     for i_b in range(_B):
         # this safeguard seems not necessary in normal execution
@@ -1419,8 +1419,10 @@ def func_solve(
                     constraint_state=constraint_state,
                     static_rigid_sim_config=static_rigid_sim_config,
                 )
-                if constraint_state.improved[i_b] < 1:
+                if not constraint_state.improved[i_b]:
                     break
+        else:
+            constraint_state.improved[i_b] = False
 
 
 @ti.func
@@ -1540,7 +1542,7 @@ def func_no_linesearch(i_b, constraint_state: array_class.ConstraintState):
     func_ls_init(i_b)
     n_dofs = constraint_state.search.shape[0]
 
-    constraint_state.improved[i_b] = 1
+    constraint_state.improved[i_b] = True
     for i_d in range(n_dofs):
         constraint_state.qacc[i_d, i_b] = constraint_state.qacc[i_d, i_b] + constraint_state.search[i_d, i_b]
         constraint_state.Ma[i_d, i_b] = constraint_state.Ma[i_d, i_b] + constraint_state.mv[i_d, i_b]
@@ -1810,7 +1812,7 @@ def func_solve_body(
     )
 
     if ti.abs(alpha) < rigid_global_info.EPS[None]:
-        constraint_state.improved[i_b] = 0
+        constraint_state.improved[i_b] = False
     else:
         for i_d in range(n_dofs):
             constraint_state.qacc[i_d, i_b] = (
@@ -1861,9 +1863,9 @@ def func_solve_body(
             gradient += constraint_state.grad[i_d, i_b] * constraint_state.grad[i_d, i_b]
         gradient = ti.sqrt(gradient)
         if gradient < tol_scaled or improvement < tol_scaled:
-            constraint_state.improved[i_b] = 0
+            constraint_state.improved[i_b] = False
         else:
-            constraint_state.improved[i_b] = 1
+            constraint_state.improved[i_b] = True
 
             if ti.static(static_rigid_sim_config.solver_type == gs.constraint_solver.Newton):
                 for i_d in range(n_dofs):
diff --git a/genesis/engine/solvers/rigid/constraint_solver_decomp_island.py b/genesis/engine/solvers/rigid/constraint_solver_decomp_island.py
@@ -49,7 +49,7 @@ def __init__(self, rigid_solver: "RigidSolver"):
             self.jac_n_relevant_dofs = ti.field(gs.ti_int, shape=(self.len_constraints_, self._B))
 
         self.n_constraints = ti.field(gs.ti_int, shape=(self._B,))
-        self.improved = ti.field(gs.ti_int, shape=(self._B,))
+        self.improved = ti.field(gs.ti_bool, shape=(self._B,))
 
         self.Jaref = ti.field(dtype=gs.ti_float, shape=(self.len_constraints_, self._B))
         self.Ma = ti.field(dtype=gs.ti_float, shape=(self._solver.n_dofs_, self._B))
@@ -566,7 +566,7 @@ def _func_solve(self, i_island: int, i_b: int):
             tol_scaled = (self._solver.meaninertia[i_b] * ti.max(1, self._solver.n_dofs)) * self.tolerance
             for it in range(self.iterations):
                 self._func_solve_body(i_island, i_b)
-                if self.improved[i_b] < 1:
+                if not self.improved[i_b]:
                     break
 
                 gradient = gs.ti_float(0.0)
@@ -868,9 +868,9 @@ def _func_solve_body(self, island, i_b):
         alpha = self._func_linesearch(island, i_b)
 
         if ti.abs(alpha) < gs.EPS:
-            self.improved[i_b] = 0
+            self.improved[i_b] = False
         else:
-            self.improved[i_b] = 1
+            self.improved[i_b] = True
             for i_island_entity in range(self.contact_island.island_entity[island, i_b].n):
                 i_e_ = self.contact_island.island_entity[island, i_b].start + i_island_entity
                 i_e = self.contact_island.entity_id[i_e_, i_b]
diff --git a/genesis/utils/array_class.py b/genesis/utils/array_class.py
@@ -109,25 +109,31 @@ class StructRigidGlobalInfo(metaclass=BASE_METACLASS):
 def get_rigid_global_info(solver):
     _B = solver._B
 
+    mass_mat_shape = (solver.n_dofs_, solver.n_dofs_, _B)
+    if math.prod(mass_mat_shape) > np.iinfo(np.int32).max:
+        gs.raise_exception(
+            f"Mass matrix shape (n_dofs={solver.n_dofs_}, n_dofs={solver.n_dofs_}, n_envs={_B}) is too large."
+        )
+
     return StructRigidGlobalInfo(
+        envs_offset=V_VEC(3, dtype=gs.ti_float, shape=(_B,)),
+        gravity=V_VEC(3, dtype=gs.ti_float, shape=(_B,)),
+        meaninertia=V(dtype=gs.ti_float, shape=(_B,)),
         n_awake_dofs=V(dtype=gs.ti_int, shape=(_B,)),
-        awake_dofs=V(dtype=gs.ti_int, shape=(solver.n_dofs_, _B)),
         n_awake_entities=V(dtype=gs.ti_int, shape=(_B,)),
-        awake_entities=V(dtype=gs.ti_int, shape=(solver.n_entities_, _B)),
         n_awake_links=V(dtype=gs.ti_int, shape=(_B,)),
+        awake_dofs=V(dtype=gs.ti_int, shape=(solver.n_dofs_, _B)),
+        awake_entities=V(dtype=gs.ti_int, shape=(solver.n_entities_, _B)),
         awake_links=V(dtype=gs.ti_int, shape=(solver.n_links_, _B)),
         qpos0=V(dtype=gs.ti_float, shape=(solver.n_qs_, _B)),
         qpos=V(dtype=gs.ti_float, shape=(solver.n_qs_, _B)),
         links_T=V_MAT(n=4, m=4, dtype=gs.ti_float, shape=(solver.n_links_,)),
-        envs_offset=V_VEC(3, dtype=gs.ti_float, shape=(_B,)),
         geoms_init_AABB=V_VEC(3, dtype=gs.ti_float, shape=(solver.n_geoms_, 8)),
-        mass_mat=V(dtype=gs.ti_float, shape=(solver.n_dofs_, solver.n_dofs_, _B)),
-        mass_mat_L=V(dtype=gs.ti_float, shape=(solver.n_dofs_, solver.n_dofs_, _B)),
         mass_mat_D_inv=V(dtype=gs.ti_float, shape=(solver.n_dofs_, _B)),
         mass_mat_mask=V(dtype=gs.ti_bool, shape=(solver.n_entities_, _B)),
-        meaninertia=V(dtype=gs.ti_float, shape=(_B,)),
         mass_parent_mask=V(dtype=gs.ti_float, shape=(solver.n_dofs_, solver.n_dofs_)),
-        gravity=V_VEC(3, dtype=gs.ti_float, shape=(_B,)),
+        mass_mat=V(dtype=gs.ti_float, shape=mass_mat_shape),
+        mass_mat_L=V(dtype=gs.ti_float, shape=mass_mat_shape),
         substep_dt=V_SCALAR_FROM(dtype=gs.ti_float, value=solver._substep_dt),
         iterations=V_SCALAR_FROM(dtype=gs.ti_int, value=solver._options.iterations),
         tolerance=V_SCALAR_FROM(dtype=gs.ti_float, value=solver._options.tolerance),
@@ -239,7 +245,7 @@ def get_constraint_state(constraint_solver, solver):
         ti_n_equalities=V(dtype=gs.ti_int, shape=(_B,)),
         n_constraints_equality=V(dtype=gs.ti_int, shape=(_B,)),
         n_constraints_frictionloss=V(dtype=gs.ti_int, shape=(_B,)),
-        improved=V(dtype=gs.ti_int, shape=(_B,)),
+        improved=V(dtype=gs.ti_bool, shape=(_B,)),
         cost_ws=V(dtype=gs.ti_float, shape=(_B,)),
         gauss=V(dtype=gs.ti_float, shape=(_B,)),
         cost=V(dtype=gs.ti_float, shape=(_B,)),
diff --git a/tests/test_rigid_benchmarks.py b/tests/test_rigid_benchmarks.py
