[FEATURE] Implemented safer GJK and added GJK examples (#1357)

Author: SonSang

examples/collision/pyramid.py

@@ -0,0 +1,60 @@
+import numpy as np
+import genesis as gs
+import argparse
+
+pile_type = "static"
+num_cubes = 5
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--pile_type", type=str, default=pile_type, choices=["static", "falling"])
+parser.add_argument("--num_cubes", type=int, default=num_cubes, choices=range(5, 11))
+parser.add_argument("--cpu", action="store_true", help="Use CPU backend instead of GPU")
+args = parser.parse_args()
+
+pile_type = args.pile_type
+num_cubes = args.num_cubes
+cpu = args.cpu
+backend = gs.cpu if cpu else gs.gpu
+
+gs.init(backend=backend, precision="32")
+
+scene = gs.Scene(
+    sim_options=gs.options.SimOptions(
+        dt=0.01,
+    ),
+    rigid_options=gs.options.RigidOptions(
+        box_box_detection=False,
+        max_collision_pairs=1000,
+        use_gjk_collision=True,
+        enable_mujoco_compatibility=False,
+    ),
+    viewer_options=gs.options.ViewerOptions(
+        camera_pos=(0, -5.5, 2.5),
+        camera_lookat=(0, 0.0, 1.5),
+        camera_fov=30,
+        max_FPS=60,
+    ),
+    show_viewer=True,
+)
+
+plane = scene.add_entity(gs.morphs.Plane(pos=(0, 0, 0)))
+
+# create pyramid of boxes
+box_size = 0.25
+if pile_type == "static":
+    box_spacing = box_size
+else:
+    box_spacing = 1.1 * box_size
+vec_one = np.array([1.0, 1.0, 1.0])
+box_pos_offset = (0 - 0.5, 1, 0.0) + 0.5 * box_size * vec_one
+boxes = {}
+for i in range(num_cubes):
+    for j in range(num_cubes - i):
+        box = scene.add_entity(
+            gs.morphs.Box(size=box_size * vec_one, pos=box_pos_offset + box_spacing * np.array([i + 0.5 * j, 0, j])),
+        )
+
+scene.build()
+
+for i in range(1000):
+    scene.step()

examples/collision/tower.py

@@ -0,0 +1,95 @@
+import numpy as np
+import genesis as gs
+import argparse
+
+object_type = "cylinder"
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--object",
+    type=str,
+    default=object_type,
+    choices=["sphere", "cylinder", "duck"],
+)
+args = parser.parse_args()
+object_type = args.object
+
+gs.init(backend=gs.cpu, precision="32")
+
+scene = gs.Scene(
+    sim_options=gs.options.SimOptions(
+        dt=0.005,
+    ),
+    rigid_options=gs.options.RigidOptions(
+        box_box_detection=False,
+        max_collision_pairs=2000,
+        use_gjk_collision=True,
+        enable_mujoco_compatibility=False,
+    ),
+    viewer_options=gs.options.ViewerOptions(
+        camera_pos=(20, -20, 20),
+        camera_lookat=(0.0, 0.0, 5.0),
+        camera_fov=30,
+        max_FPS=60,
+    ),
+    show_viewer=True,
+)
+
+plane = scene.add_entity(gs.morphs.Plane(pos=(0, 0, 0)))
+
+# create pyramid of boxes
+box_width = 0.25
+box_length = 2.0
+box_height = 0.1
+num_stacks = 50
+height_offset = 0.0
+for i in range(num_stacks):
+    horizontal = i % 2 == 0
+
+    if horizontal:
+        box_size = np.array([box_width, box_length, box_height])
+        box_pos_0 = (-0.4 * box_length, 0, i * (height_offset + box_size[2]) + 0.5 * box_size[2])
+        box_pos_1 = (+0.4 * box_length, 0, i * (height_offset + box_size[2]) + 0.5 * box_size[2])
+    else:
+        box_size = np.array([box_length, box_width, box_height])
+        box_pos_0 = (0, -0.4 * box_length, i * (height_offset + box_size[2]) + 0.5 * box_size[2])
+        box_pos_1 = (0, +0.4 * box_length, i * (height_offset + box_size[2]) + 0.5 * box_size[2])
+
+    scene.add_entity(
+        gs.morphs.Box(size=box_size, pos=box_pos_0),
+    )
+    scene.add_entity(
+        gs.morphs.Box(size=box_size, pos=box_pos_1),
+    )
+
+# Drop a huge mesh
+if object_type == "duck":
+    duck_scale = 1.0
+    duck = scene.add_entity(
+        morph=gs.morphs.Mesh(
+            file="meshes/duck.obj",
+            scale=duck_scale,
+            pos=(0, 0, num_stacks * (height_offset + box_height) + 10 * duck_scale),
+        ),
+    )
+elif object_type == "sphere":
+    sphere_radius = 2.0
+    scene.add_entity(
+        morph=gs.morphs.Sphere(
+            radius=sphere_radius, pos=(0.0, 0.0, num_stacks * (height_offset + box_height) + 5 * sphere_radius)
+        ),
+    )
+elif object_type == "cylinder":
+    cylinder_radius = 2.0
+    cylinder_height = 1.0
+    scene.add_entity(
+        morph=gs.morphs.Cylinder(
+            radius=cylinder_radius,
+            height=cylinder_height,
+            pos=(0.0, 0.0, num_stacks * (height_offset + box_height) + 5 * cylinder_height),
+        ),
+    )
+
+scene.build()
+for i in range(5000):
+    scene.step()

genesis/engine/solvers/rigid/collider_decomp.py

@@ -22,9 +22,14 @@
 
 
 class CCD_ALGORITHM_CODE(IntEnum):
+    # Our MPR (with SDF)
     MPR = 0
-    MPR_SDF = 1
+    # MuJoCo MPR
+    MJ_MPR = 1
+    # Our GJK
     GJK = 2
+    # MuJoCo GJK
+    MJ_GJK = 3
 
 
 @ti.func
@@ -54,15 +59,19 @@ def __init__(self, rigid_solver: "RigidSolver"):
 
         # Identify the convex collision detection (ccd) algorithm
         if self._solver._options.use_gjk_collision:
-            self.ccd_algorithm = CCD_ALGORITHM_CODE.GJK
-        elif self._solver._enable_mujoco_compatibility:
-            self.ccd_algorithm = CCD_ALGORITHM_CODE.MPR
+            if self._solver._enable_mujoco_compatibility:
+                self.ccd_algorithm = CCD_ALGORITHM_CODE.MJ_GJK
+            else:
+                self.ccd_algorithm = CCD_ALGORITHM_CODE.GJK
         else:
-            self.ccd_algorithm = CCD_ALGORITHM_CODE.MPR_SDF
+            if self._solver._enable_mujoco_compatibility:
+                self.ccd_algorithm = CCD_ALGORITHM_CODE.MJ_MPR
+            else:
+                self.ccd_algorithm = CCD_ALGORITHM_CODE.MPR
 
         # FIXME: MPR is necessary because it is used for terrain collision detection
         self._mpr = MPR(rigid_solver)
-        self._gjk = GJK(rigid_solver) if self.ccd_algorithm == CCD_ALGORITHM_CODE.GJK else None
+        self._gjk = GJK(rigid_solver) if self._solver._options.use_gjk_collision else None
 
         # multi contact perturbation and tolerance
         if self._solver._enable_mujoco_compatibility:
@@ -1110,6 +1119,7 @@ def _func_plane_box_contact(self, i_ga, i_gb, i_b):
 
     @ti.func
     def _func_add_contact(self, i_ga, i_gb, normal, contact_pos, penetration, i_b):
+        # print(f"Adding contact {i_ga} {i_gb}, normal:", normal, "contact_pos:", contact_pos, "penetration:", penetration)
         i_col = self.n_contacts[i_b]
 
         if i_col == self._max_contact_pairs:
@@ -1265,8 +1275,8 @@ def _func_convex_convex_contact(self, i_ga, i_gb, i_b):
                         contact_pos = v1 - 0.5 * penetration * normal
                         is_col = penetration > 0
                     else:
-                        ### MPR, MPR + SDF
-                        if ti.static(self.ccd_algorithm != CCD_ALGORITHM_CODE.GJK):
+                        ### MPR, MJ_MPR
+                        if ti.static(self.ccd_algorithm in (CCD_ALGORITHM_CODE.MPR, CCD_ALGORITHM_CODE.MJ_MPR)):
                             # Try using MPR before anything else
                             is_mpr_updated = False
                             is_mpr_guess_direction_available = True
@@ -1301,10 +1311,9 @@ def _func_convex_convex_contact(self, i_ga, i_gb, i_b):
                                 # associated with the point of deepest penetration.
                                 try_sdf = True
 
-                        ### GJK
-                        elif ti.static(self.ccd_algorithm == CCD_ALGORITHM_CODE.GJK):
-                            # If it was not the first detection, only detect single contact point.
-                            self._gjk.func_gjk_contact(i_ga, i_gb, i_b, i_detection == 0)
+                        ### GJK, MJ_GJK
+                        elif ti.static(self.ccd_algorithm in (CCD_ALGORITHM_CODE.GJK, CCD_ALGORITHM_CODE.MJ_GJK)):
+                            self._gjk.func_gjk_contact(i_ga, i_gb, i_b)
 
                             is_col = self._gjk.is_col[i_b] == 1
                             penetration = self._gjk.penetration[i_b]
@@ -1326,7 +1335,7 @@ def _func_convex_convex_contact(self, i_ga, i_gb, i_b):
                                     contact_pos = self._gjk.contact_pos[i_b, 0]
                                     normal = self._gjk.normal[i_b, 0]
 
-                    if ti.static(self.ccd_algorithm == CCD_ALGORITHM_CODE.MPR_SDF):
+                    if ti.static(self.ccd_algorithm == CCD_ALGORITHM_CODE.MPR):
                         if try_sdf:
                             is_col_a = False
                             is_col_b = False
@@ -1392,15 +1401,15 @@ def _func_convex_convex_contact(self, i_ga, i_gb, i_b):
                             axis_0, axis_1 = self._func_contact_orthogonals(i_ga, i_gb, normal, i_b)
                             n_con = 1
 
-                        if ti.static(not self._solver._enable_mujoco_compatibility):
+                        if ti.static(self.ccd_algorithm in (CCD_ALGORITHM_CODE.MPR, CCD_ALGORITHM_CODE.GJK)):
                             self.contact_cache[i_ga, i_gb, i_b].normal = normal
                     else:
                         # Clear collision normal cache if not in contact
                         # self.contact_cache[i_ga, i_gb, i_b].i_va_ws = -1
                         self.contact_cache[i_ga, i_gb, i_b].normal.fill(0.0)
 
                 elif multi_contact and is_col_0 > 0 and is_col > 0:
-                    if ti.static(self.ccd_algorithm == CCD_ALGORITHM_CODE.MPR_SDF):
+                    if ti.static(self.ccd_algorithm in (CCD_ALGORITHM_CODE.MPR, CCD_ALGORITHM_CODE.GJK)):
                         # 1. Project the contact point on both geometries
                         # 2. Revert the effect of small rotation
                         # 3. Update contact point
@@ -1437,9 +1446,9 @@ def _func_convex_convex_contact(self, i_ga, i_gb, i_b):
                         # dynamics since zero-penetration contact points should not induce any force.
                         penetration = normal.dot(contact_point_b - contact_point_a)
 
-                    elif ti.static(self.ccd_algorithm == CCD_ALGORITHM_CODE.GJK):
+                    elif ti.static(self.ccd_algorithm == CCD_ALGORITHM_CODE.MJ_GJK):
                         # Only change penetration to the initial one, because the normal vector could change abruptly
-                        # under GJK-EPA as the nearest simplex is determined by discrete logic, unlike MPR.
+                        # under MuJoCo's GJK-EPA.
                         penetration = penetration_0
 
                     # Discard contact point is repeated
@@ -1457,10 +1466,10 @@ def _func_convex_convex_contact(self, i_ga, i_gb, i_b):
                             self._func_add_contact(i_ga, i_gb, normal, contact_pos, penetration, i_b)
                             n_con = n_con + 1
 
-                    self._solver.geoms_state[i_ga, i_b].pos = ga_pos
-                    self._solver.geoms_state[i_ga, i_b].quat = ga_quat
-                    self._solver.geoms_state[i_gb, i_b].pos = gb_pos
-                    self._solver.geoms_state[i_gb, i_b].quat = gb_quat
+                self._solver.geoms_state[i_ga, i_b].pos = ga_pos
+                self._solver.geoms_state[i_ga, i_b].quat = ga_quat
+                self._solver.geoms_state[i_gb, i_b].pos = gb_pos
+                self._solver.geoms_state[i_gb, i_b].quat = gb_quat
 
     @ti.func
     def _func_rotate_frame(self, i_g, contact_pos, qrot, i_b):

genesis/engine/solvers/rigid/constraint_solver_decomp.py

@@ -35,6 +35,13 @@ def __init__(self, rigid_solver: "RigidSolver"):
         self.ti_n_equalities = ti.field(gs.ti_int, shape=self._solver._batch_shape())
         self.ti_n_equalities.from_numpy(np.full((self._solver._B,), self._solver.n_equalities, dtype=gs.np_int))
 
+        jac_shape = self._solver._batch_shape((self.len_constraints_, self._solver.n_dofs_))
+        if (jac_shape[0] * jac_shape[1] * jac_shape[2]) > np.iinfo(np.int32).max:
+            raise ValueError(
+                f"Jacobian shape {jac_shape} is too large for int32. "
+                "Consider reducing the number of constraints or the number of degrees of freedom."
+            )
+
         self.jac = ti.field(
             dtype=gs.ti_float, shape=self._solver._batch_shape((self.len_constraints_, self._solver.n_dofs_))
         )