Merge branch 'main' into yiling/250713_aux_kernels

Author: YilingQiao

genesis/engine/solvers/rigid/array_class.py

@@ -224,6 +224,165 @@ def __init__(self, f_batch):
         self.simplex_size = ti.field(gs.ti_int, shape=f_batch())
 
 
+# =========================================== GJK ===========================================
+@ti.data_oriented
+class GJKState:
+    def __init__(self, solver, static_rigid_sim_config, gjk_static_config):
+        _B = solver._B
+        polytope_max_faces = gjk_static_config.polytope_max_faces
+        max_contacts_per_pair = gjk_static_config.max_contacts_per_pair
+        max_contact_polygon_verts = gjk_static_config.max_contact_polygon_verts
+
+        # Cache to store the previous support points for support mesh function.
+        self.support_mesh_prev_vertex_id = ti.field(dtype=gs.ti_int, shape=(_B, 2))
+
+        ### GJK simplex
+        struct_simplex_vertex = ti.types.struct(
+            # Support points on the two objects
+            obj1=gs.ti_vec3,
+            obj2=gs.ti_vec3,
+            # Support point IDs on the two objects
+            id1=gs.ti_int,
+            id2=gs.ti_int,
+            # Vertex on Minkowski difference
+            mink=gs.ti_vec3,
+        )
+        struct_simplex = ti.types.struct(
+            # Number of vertices in the simplex
+            nverts=gs.ti_int,
+            # Distance from the origin to the simplex
+            dist=gs.ti_float,
+        )
+        struct_simplex_buffer = ti.types.struct(
+            # Normals of the simplex faces
+            normal=gs.ti_vec3,
+            # Signed distances of the simplex faces from the origin
+            sdist=gs.ti_float,
+        )
+        self.simplex_vertex = struct_simplex_vertex.field(shape=(_B, 4))
+        self.simplex_buffer = struct_simplex_buffer.field(shape=(_B, 4))
+        self.simplex = struct_simplex.field(shape=(_B,))
+
+        # Only when we enable MuJoCo compatibility, we use the simplex vertex and buffer for intersection checks.
+        if static_rigid_sim_config.enable_mujoco_compatibility:
+            self.simplex_vertex_intersect = struct_simplex_vertex.field(shape=(_B, 4))
+            self.simplex_buffer_intersect = struct_simplex_buffer.field(shape=(_B, 4))
+            self.nsimplex = ti.field(dtype=gs.ti_int, shape=(_B,))
+
+        # In safe GJK, if the initial simplex is degenerate and the geometries are discrete, we go through vertices
+        # on the Minkowski difference to find a vertex that would make a valid simplex. To prevent iterating through
+        # the same vertices again during initial simplex construction, we keep the vertex ID of the last vertex that
+        # we searched, so that we can start searching from the next vertex.
+        self.last_searched_simplex_vertex_id = ti.field(dtype=gs.ti_int, shape=(_B,))
+
+        ### EPA polytope
+        struct_polytope_vertex = struct_simplex_vertex
+        struct_polytope_face = ti.types.struct(
+            # Indices of the vertices forming the face on the polytope
+            verts_idx=gs.ti_ivec3,
+            # Indices of adjacent faces, one for each edge: [v1,v2], [v2,v3], [v3,v1]
+            adj_idx=gs.ti_ivec3,
+            # Projection of the origin onto the face, can be used as face normal
+            normal=gs.ti_vec3,
+            # Square of 2-norm of the normal vector, negative means deleted face
+            dist2=gs.ti_float,
+            # Index of the face in the polytope map, -1 for not in the map, -2 for deleted
+            map_idx=gs.ti_int,
+        )
+        # Horizon is used for representing the faces to delete when the polytope is expanded by inserting a new vertex.
+        struct_polytope_horizon_data = ti.types.struct(
+            # Indices of faces on horizon
+            face_idx=gs.ti_int,
+            # Corresponding edge of each face on the horizon
+            edge_idx=gs.ti_int,
+        )
+        struct_polytope = ti.types.struct(
+            # Number of vertices in the polytope
+            nverts=gs.ti_int,
+            # Number of faces in the polytope (it could include deleted faces)
+            nfaces=gs.ti_int,
+            # Number of faces in the polytope map (only valid faces on polytope)
+            nfaces_map=gs.ti_int,
+            # Number of edges in the horizon
+            horizon_nedges=gs.ti_int,
+            # Support point on the Minkowski difference where the horizon is created
+            horizon_w=gs.ti_vec3,
+        )
+
+        self.polytope = struct_polytope.field(shape=(_B,))
+        self.polytope_verts = struct_polytope_vertex.field(shape=(_B, 5 + gjk_static_config.epa_max_iterations))
+        self.polytope_faces = struct_polytope_face.field(shape=(_B, polytope_max_faces))
+        self.polytope_horizon_data = struct_polytope_horizon_data.field(
+            shape=(_B, 6 + gjk_static_config.epa_max_iterations)
+        )
+
+        # Face indices that form the polytope. The first [nfaces_map] indices are the faces that form the polytope.
+        self.polytope_faces_map = ti.Vector.field(n=polytope_max_faces, dtype=gs.ti_int, shape=(_B,))
+
+        # Stack to use for visiting faces during the horizon construction. The size is (# max faces * 3),
+        # because a face has 3 edges.
+        self.polytope_horizon_stack = struct_polytope_horizon_data.field(shape=(_B, polytope_max_faces * 3))
+
+        # Data structures for multi-contact detection based on MuJoCo's implementation.
+        if gjk_static_config.enable_mujoco_multi_contact:
+            struct_contact_face = ti.types.struct(
+                # Vertices from the two colliding faces
+                vert1=gs.ti_vec3,
+                vert2=gs.ti_vec3,
+                endverts=gs.ti_vec3,
+                # Normals of the two colliding faces
+                normal1=gs.ti_vec3,
+                normal2=gs.ti_vec3,
+                # Face ID of the two colliding faces
+                id1=gs.ti_int,
+                id2=gs.ti_int,
+            )
+            # Struct for storing temp. contact normals
+            struct_contact_normal = ti.types.struct(
+                endverts=gs.ti_vec3,
+                # Normal vector of the contact point
+                normal=gs.ti_vec3,
+                # Face ID
+                id=gs.ti_int,
+            )
+            struct_contact_halfspace = ti.types.struct(
+                # Halfspace normal
+                normal=gs.ti_vec3,
+                # Halfspace distance from the origin
+                dist=gs.ti_float,
+            )
+            self.contact_faces = struct_contact_face.field(shape=(_B, max_contact_polygon_verts))
+            self.contact_normals = struct_contact_normal.field(shape=(_B, max_contact_polygon_verts))
+            self.contact_halfspaces = struct_contact_halfspace.field(shape=(_B, max_contact_polygon_verts))
+            self.contact_clipped_polygons = gs.ti_vec3.field(shape=(_B, 2, max_contact_polygon_verts))
+
+        # Whether or not the MuJoCo's contact manifold detection algorithm was used for the current pair.
+        self.multi_contact_flag = ti.field(dtype=gs.ti_int, shape=(_B,))
+
+        ### Final results
+        # Witness information
+        struct_witness = ti.types.struct(
+            # Witness points on the two objects
+            point_obj1=gs.ti_vec3,
+            point_obj2=gs.ti_vec3,
+        )
+        self.witness = struct_witness.field(shape=(_B, max_contacts_per_pair))
+        self.n_witness = ti.field(dtype=gs.ti_int, shape=(_B,))
+
+        # Contact information, the namings are the same as those from the calling function. Even if they could be
+        # redundant, we keep them for easier use from the calling function.
+        self.n_contacts = ti.field(dtype=gs.ti_int, shape=(_B,))
+        self.contact_pos = gs.ti_vec3.field(shape=(_B, max_contacts_per_pair))
+        self.normal = gs.ti_vec3.field(shape=(_B, max_contacts_per_pair))
+        self.is_col = ti.field(dtype=gs.ti_int, shape=(_B,))
+        self.penetration = ti.field(dtype=gs.ti_float, shape=(_B,))
+
+        # Distance between the two objects.
+        # If the objects are separated, the distance is positive.
+        # If the objects are intersecting, the distance is negative (depth).
+        self.distance = ti.field(dtype=gs.ti_float, shape=(_B,))
+
+
 # =========================================== SupportField ===========================================
 @ti.data_oriented
 class SupportFieldInfo:

genesis/engine/solvers/rigid/collider_decomp.py

@@ -367,13 +367,16 @@ def detection(self) -> None:
             self._solver.geoms_info,
             self._solver.geoms_init_AABB,
             self._solver.verts_info,
+            self._solver.faces_info,
             self._solver._rigid_global_info,
             self._solver._static_rigid_sim_config,
             self._collider_state,
             self._collider_info,
             self._collider_static_config,
             self._mpr._mpr_state,
             self._mpr._mpr_static_config,
+            self._gjk._gjk_state if self._gjk is not None else None,
+            self._gjk._gjk_static_config if self._gjk is not None else None,
             self._support_field._support_field_info,
             self._support_field._support_field_static_config,
             self._mpr,
@@ -1171,13 +1174,16 @@ def _func_narrow_phase_convex_vs_convex(
         geoms_info: array_class.GeomsInfo,
         geoms_init_AABB: array_class.GeomsInitAABB,
         verts_info: array_class.VertsInfo,
+        faces_info: array_class.FacesInfo,
         rigid_global_info: ti.template(),
         static_rigid_sim_config: ti.template(),
         collider_state: ti.template(),
         collider_info: ti.template(),
         collider_static_config: ti.template(),
         mpr_state: ti.template(),
         mpr_static_config: ti.template(),
+        gjk_state: ti.template(),
+        gjk_static_config: ti.template(),
         support_field_info: ti.template(),
         support_field_static_config: ti.template(),
         # FIXME: We need mpr, gjk, sdf, and support_field for now to call their class functions. After migration is
@@ -1224,13 +1230,16 @@ def _func_narrow_phase_convex_vs_convex(
                             geoms_info,
                             geoms_init_AABB,
                             verts_info,
+                            faces_info,
                             rigid_global_info,
                             static_rigid_sim_config,
                             collider_state,
                             collider_info,
                             collider_static_config,
                             mpr_state,
                             mpr_static_config,
+                            gjk_state,
+                            gjk_static_config,
                             support_field_info,
                             support_field_static_config,
                             mpr,
@@ -1252,13 +1261,16 @@ def _func_narrow_phase_convex_vs_convex(
                                 geoms_info,
                                 geoms_init_AABB,
                                 verts_info,
+                                faces_info,
                                 rigid_global_info,
                                 static_rigid_sim_config,
                                 collider_state,
                                 collider_info,
                                 collider_static_config,
                                 mpr_state,
                                 mpr_static_config,
+                                gjk_state,
+                                gjk_static_config,
                                 support_field_info,
                                 support_field_static_config,
                                 mpr,
@@ -1768,13 +1780,16 @@ def _func_convex_convex_contact(
         geoms_info: array_class.GeomsInfo,
         geoms_init_AABB: array_class.GeomsInitAABB,
         verts_info: array_class.VertsInfo,
+        faces_info: array_class.FacesInfo,
         rigid_global_info: ti.template(),
         static_rigid_sim_config: ti.template(),
         collider_state: ti.template(),
         collider_info: ti.template(),
         collider_static_config: ti.template(),
         mpr_state: ti.template(),
         mpr_static_config: ti.template(),
+        gjk_state: ti.template(),
+        gjk_static_config: ti.template(),
         support_field_info: ti.template(),
         support_field_static_config: ti.template(),
         mpr: ti.template(),
@@ -1933,21 +1948,37 @@ def _func_convex_convex_contact(
                         elif ti.static(
                             collider_static_config.ccd_algorithm in (CCD_ALGORITHM_CODE.GJK, CCD_ALGORITHM_CODE.MJ_GJK)
                         ):
-                            gjk.func_gjk_contact(i_ga, i_gb, i_b)
+                            gjk.func_gjk_contact(
+                                geoms_state,
+                                geoms_info,
+                                verts_info,
+                                faces_info,
+                                static_rigid_sim_config,
+                                collider_state,
+                                collider_static_config,
+                                gjk_state,
+                                gjk_static_config,
+                                support_field_info,
+                                support_field_static_config,
+                                support_field,
+                                i_ga,
+                                i_gb,
+                                i_b,
+                            )
 
-                            is_col = gjk.is_col[i_b] == 1
-                            penetration = gjk.penetration[i_b]
-                            n_contacts = gjk.n_contacts[i_b]
+                            is_col = gjk_state.is_col[i_b] == 1
+                            penetration = gjk_state.penetration[i_b]
+                            n_contacts = gjk_state.n_contacts[i_b]
 
                             if is_col:
-                                if gjk.multi_contact_flag[i_b]:
+                                if gjk_state.multi_contact_flag[i_b]:
                                     # Used MuJoCo's multi-contact algorithm to find multiple contact points. Therefore,
                                     # add the discovered contact points and stop multi-contact search.
                                     for i_c in range(n_contacts):
                                         # Ignore contact points if the number of contacts exceeds the limit.
                                         if i_c < ti.static(collider_static_config.n_contacts_per_pair):
-                                            contact_pos = gjk.contact_pos[i_b, i_c]
-                                            normal = gjk.normal[i_b, i_c]
+                                            contact_pos = gjk_state.contact_pos[i_b, i_c]
+                                            normal = gjk_state.normal[i_b, i_c]
                                             self_unused._func_add_contact(
                                                 geoms_state,
                                                 geoms_info,
@@ -1963,8 +1994,8 @@ def _func_convex_convex_contact(
 
                                     break
                                 else:
-                                    contact_pos = gjk.contact_pos[i_b, 0]
-                                    normal = gjk.normal[i_b, 0]
+                                    contact_pos = gjk_state.contact_pos[i_b, 0]
+                                    normal = gjk_state.normal[i_b, 0]
 
                     if ti.static(collider_static_config.ccd_algorithm == CCD_ALGORITHM_CODE.MPR):
                         if try_sdf: