Merge branch 'main' into easydiffrigid

Author: SonSang

genesis/engine/scene.py

@@ -273,6 +273,22 @@ def _validate_options(
         if not isinstance(renderer_options, RendererOptions):
             gs.raise_exception("`renderer` should be an instance of `gs.renderers.Renderer`.")
 
+        # Validate rigid_options against sim_options
+        if rigid_options.box_box_detection is None:
+            rigid_options.box_box_detection = not sim_options.requires_grad
+        elif rigid_options.box_box_detection and sim_options.requires_grad:
+            gs.raise_exception(
+                "`rigid_options.box_box_detection` cannot be True when `sim_options.requires_grad` is True."
+            )
+        if not rigid_options.use_gjk_collision and sim_options.requires_grad:
+            gs.raise_exception(
+                "`rigid_options.use_gjk_collision` cannot be False when `sim_options.requires_grad` is True."
+            )
+        if rigid_options.enable_mujoco_compatibility and sim_options.requires_grad:
+            gs.raise_exception(
+                "`rigid_options.enable_mujoco_compatibility` cannot be True when `sim_options.requires_grad` is True."
+            )
+
     def destroy(self):
         if getattr(self, "_recorder_manager", None) is not None:
             if self._recorder_manager.is_recording:

genesis/engine/solvers/pbd_solver.py

@@ -923,6 +923,7 @@ def set_animate_particles_by_link(
         envs_idx = self._scene._sanitize_envs_idx(envs_idx)
         self._sim._coupler.kernel_attach_pbd_to_rigid_link(particles_idx, envs_idx, link_idx, links_state)
 
+    @ti.kernel
     def _kernel_get_particles_vel(
         self,
         particle_start: ti.i32,

genesis/engine/solvers/rigid/collider_decomp.py

@@ -196,18 +196,18 @@ def support_driver(
     v = ti.Vector.zero(gs.ti_float, 3)
     geom_type = geoms_info.type[i_g]
     if geom_type == gs.GEOM_TYPE.SPHERE:
-        v = support_field._func_support_sphere(geoms_state, geoms_info, direction, i_g, i_b, False)
+        v, v_, vid = support_field._func_support_sphere(geoms_state, geoms_info, direction, i_g, i_b, False)
     elif geom_type == gs.GEOM_TYPE.ELLIPSOID:
         v = support_field._func_support_ellipsoid(geoms_state, geoms_info, direction, i_g, i_b)
     elif geom_type == gs.GEOM_TYPE.CAPSULE:
         v = support_field._func_support_capsule(geoms_state, geoms_info, direction, i_g, i_b, False)
     elif geom_type == gs.GEOM_TYPE.BOX:
-        v, _ = support_field._func_support_box(geoms_state, geoms_info, direction, i_g, i_b)
+        v, v_, vid = support_field._func_support_box(geoms_state, geoms_info, direction, i_g, i_b)
     elif geom_type == gs.GEOM_TYPE.TERRAIN:
         if ti.static(collider_static_config.has_terrain):
             v, _ = support_field._func_support_prism(collider_state, direction, i_g, i_b)
     else:
-        v, _ = support_field._func_support_world(
+        v, v_, vid = support_field._func_support_world(
             geoms_state, geoms_info, support_field_info, support_field_static_config, direction, i_g, i_b
         )
     return v

genesis/engine/solvers/rigid/diff_gjk_decomp.py

@@ -247,6 +247,7 @@ def build(self):
         self._static_rigid_sim_cache_key = array_class.get_static_rigid_sim_cache_key(self)
         self._static_rigid_sim_config = self.StaticRigidSimConfig(
             para_level=self.sim._para_level,
+            requires_grad=getattr(self.sim.options, "requires_grad", False),
             use_hibernation=getattr(self, "_use_hibernation", False),
             use_contact_island=getattr(self, "_use_contact_island", False),
             batch_links_info=getattr(self._options, "batch_links_info", False),

genesis/engine/solvers/rigid/gjk_decomp.py

@@ -139,9 +139,9 @@ def _func_support_world(
     g_pos = geoms_state.pos[i_g, i_b]
     g_quat = geoms_state.quat[i_g, i_b]
     d_mesh = gu.ti_transform_by_quat(d, gu.ti_inv_quat(g_quat))
-    v, vid = _func_support_mesh(support_field_info, support_field_static_config, d_mesh, i_g)
-    v_ = gu.ti_transform_by_trans_quat(v, g_pos, g_quat)
-    return v_, vid
+    v_, vid = _func_support_mesh(support_field_info, support_field_static_config, d_mesh, i_g)
+    v = gu.ti_transform_by_trans_quat(v_, g_pos, g_quat)
+    return v, v_, vid
 
 
 @ti.func
@@ -207,10 +207,18 @@ def _func_support_sphere(
     sphere_radius = geoms_info.data[i_g][0]
 
     # Shrink the sphere to a point
-    res = sphere_center
+    v = sphere_center
+    v_ = ti.Vector.zero(gs.ti_float, 3)
+    vid = -1
     if not shrink:
-        res += d * sphere_radius
-    return res
+        v += d * sphere_radius
+
+        # Local position of the support point
+        g_quat = geoms_state.quat[i_g, i_b]
+        local_d = gu.ti_inv_transform_by_quat(d, g_quat)
+        v_ = local_d * sphere_radius
+
+    return v, v_, vid
 
 
 @ti.func
@@ -308,7 +316,7 @@ def _func_support_box(
     vid = (v_[0] > 0.0) * 1 + (v_[1] > 0.0) * 2 + (v_[2] > 0.0) * 4
     vid += geoms_info.vert_start[i_g]
     v = gu.ti_transform_by_trans_quat(v_, g_pos, g_quat)
-    return v, vid
+    return v, v_, vid
 
 
 @ti.func

genesis/engine/solvers/rigid/mpr_decomp.py

@@ -315,7 +315,7 @@ class RigidOptions(Options):
     contact_resolve_time: Optional[float] = None
     constraint_timeconst: float = 0.01
     use_contact_island: bool = False
-    box_box_detection: bool = True
+    box_box_detection: Optional[bool] = None
 
     # hibernation threshold
     use_hibernation: bool = False

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -360,15 +360,15 @@ class StructContactData:
     link_b: V_ANNOTATION
 
 
-def get_contact_data(solver, max_contact_pairs):
+def get_contact_data(solver, max_contact_pairs, requires_grad):
     f_batch = solver._batch_shape
     max_contact_pairs_ = max(1, max_contact_pairs)
     kwargs = {
         "geom_a": V(dtype=gs.ti_int, shape=f_batch(max_contact_pairs_)),
         "geom_b": V(dtype=gs.ti_int, shape=f_batch(max_contact_pairs_)),
-        "penetration": V(dtype=gs.ti_float, shape=f_batch(max_contact_pairs_)),
-        "normal": V_VEC(3, dtype=gs.ti_float, shape=f_batch(max_contact_pairs_)),
-        "pos": V_VEC(3, dtype=gs.ti_float, shape=f_batch(max_contact_pairs_)),
+        "normal": V(dtype=gs.ti_vec3, shape=f_batch(max_contact_pairs_), needs_grad=requires_grad),
+        "pos": V(dtype=gs.ti_vec3, shape=f_batch(max_contact_pairs_), needs_grad=requires_grad),
+        "penetration": V(dtype=gs.ti_float, shape=f_batch(max_contact_pairs_), needs_grad=requires_grad),
         "friction": V(dtype=gs.ti_float, shape=f_batch(max_contact_pairs_)),
         "sol_params": V_VEC(7, dtype=gs.ti_float, shape=f_batch(max_contact_pairs_)),
         "force": V(dtype=gs.ti_vec3, shape=f_batch(max_contact_pairs_)),
@@ -389,6 +389,62 @@ def __init__(self):
         return ClassContactData()
 
 
+@dataclasses.dataclass
+class StructDiffContactInput:
+    ### Non-differentiable input data
+    # Geom id of the two geometries
+    geom_a: V_ANNOTATION
+    geom_b: V_ANNOTATION
+    # Local positions of the 3 vertices from the two geometries that define the face on the Minkowski difference
+    local_pos1_a: V_ANNOTATION
+    local_pos1_b: V_ANNOTATION
+    local_pos1_c: V_ANNOTATION
+    local_pos2_a: V_ANNOTATION
+    local_pos2_b: V_ANNOTATION
+    local_pos2_c: V_ANNOTATION
+    # Local positions of the 1 vertex from the two geometries that define the support point for the face above
+    w_local_pos1: V_ANNOTATION
+    w_local_pos2: V_ANNOTATION
+    # Reference id of the contact point, which is needed for the backward pass
+    ref_id: V_ANNOTATION
+    # Flag whether the contact data can be computed in numerically stable way in both the forward and backward passes
+    valid: V_ANNOTATION
+    ### Differentiable input data
+    # Reference penetration depth, which is needed for computing the weight of the contact point
+    ref_penetration: V_ANNOTATION
+
+
+def get_diff_contact_input(solver, max_contacts_per_pair):
+    _B = solver._B
+    kwargs = {
+        "geom_a": V(dtype=gs.ti_int, shape=(_B, max_contacts_per_pair)),
+        "geom_b": V(dtype=gs.ti_int, shape=(_B, max_contacts_per_pair)),
+        "local_pos1_a": V_VEC(3, dtype=gs.ti_float, shape=(_B, max_contacts_per_pair)),
+        "local_pos1_b": V_VEC(3, dtype=gs.ti_float, shape=(_B, max_contacts_per_pair)),
+        "local_pos1_c": V_VEC(3, dtype=gs.ti_float, shape=(_B, max_contacts_per_pair)),
+        "local_pos2_a": V_VEC(3, dtype=gs.ti_float, shape=(_B, max_contacts_per_pair)),
+        "local_pos2_b": V_VEC(3, dtype=gs.ti_float, shape=(_B, max_contacts_per_pair)),
+        "local_pos2_c": V_VEC(3, dtype=gs.ti_float, shape=(_B, max_contacts_per_pair)),
+        "w_local_pos1": V_VEC(3, dtype=gs.ti_float, shape=(_B, max_contacts_per_pair)),
+        "w_local_pos2": V_VEC(3, dtype=gs.ti_float, shape=(_B, max_contacts_per_pair)),
+        "ref_id": V(dtype=gs.ti_int, shape=(_B, max_contacts_per_pair)),
+        "valid": V(dtype=gs.ti_int, shape=(_B, max_contacts_per_pair)),
+        "ref_penetration": V(dtype=gs.ti_float, shape=(_B, max_contacts_per_pair), needs_grad=True),
+    }
+
+    if use_ndarray:
+        return StructDiffContactInput(**kwargs)
+    else:
+
+        @ti.data_oriented
+        class ClassDiffContactInput:
+            def __init__(self):
+                for k, v in kwargs.items():
+                    setattr(self, k, v)
+
+        return ClassDiffContactInput()
+
+
 @dataclasses.dataclass
 class StructSortBuffer:
     value: V_ANNOTATION
@@ -547,19 +603,22 @@ class StructColliderState:
     n_contacts_hibernated: V_ANNOTATION
     first_time: V_ANNOTATION
     contact_cache: StructContactCache
+    # Input data for differentiable contact detection used in the backward pass
+    diff_contact_input: StructDiffContactInput
 
 
-def get_collider_state(solver, n_possible_pairs, collider_static_config):
+def get_collider_state(solver, static_rigid_sim_config, n_possible_pairs, collider_static_config):
     _B = solver._B
     f_batch = solver._batch_shape
     n_geoms = solver.n_geoms_
     max_collision_pairs = min(solver._max_collision_pairs, n_possible_pairs)
     max_collision_pairs_broad = max_collision_pairs * collider_static_config.max_collision_pairs_broad_k
     max_contact_pairs = max_collision_pairs * collider_static_config.n_contacts_per_pair
+    requires_grad = static_rigid_sim_config.requires_grad
 
     ############## broad phase SAP ##############
 
-    contact_data = get_contact_data(solver, max_contact_pairs)
+    contact_data = get_contact_data(solver, max_contact_pairs, requires_grad)
     sort_buffer = get_sort_buffer(solver)
     contact_cache = get_contact_cache(solver)
     kwargs = {
@@ -584,6 +643,7 @@ def get_collider_state(solver, n_possible_pairs, collider_static_config):
         "n_contacts_hibernated": V(dtype=gs.ti_int, shape=_B),
         "first_time": V(dtype=gs.ti_int, shape=_B),
         "contact_cache": contact_cache,
+        "diff_contact_input": get_diff_contact_input(solver, max_contact_pairs if requires_grad else 1),
     }
 
     if use_ndarray:
@@ -722,6 +782,8 @@ class StructMDVertex:
     # Vertex of the Minkowski difference
     obj1: V_ANNOTATION
     obj2: V_ANNOTATION
+    local_obj1: V_ANNOTATION
+    local_obj2: V_ANNOTATION
     id1: V_ANNOTATION
     id2: V_ANNOTATION
     mink: V_ANNOTATION
@@ -732,6 +794,8 @@ def get_gjk_simplex_vertex(solver):
     kwargs = {
         "obj1": V_VEC(3, dtype=gs.ti_float, shape=(_B, 4)),
         "obj2": V_VEC(3, dtype=gs.ti_float, shape=(_B, 4)),
+        "local_obj1": V_VEC(3, dtype=gs.ti_float, shape=(_B, 4)),
+        "local_obj2": V_VEC(3, dtype=gs.ti_float, shape=(_B, 4)),
         "id1": V(dtype=gs.ti_int, shape=(_B, 4)),
         "id2": V(dtype=gs.ti_int, shape=(_B, 4)),
         "mink": V_VEC(3, dtype=gs.ti_float, shape=(_B, 4)),
@@ -756,6 +820,8 @@ def get_epa_polytope_vertex(solver, gjk_static_config):
     kwargs = {
         "obj1": V_VEC(3, dtype=gs.ti_float, shape=(_B, max_num_polytope_verts)),
         "obj2": V_VEC(3, dtype=gs.ti_float, shape=(_B, max_num_polytope_verts)),
+        "local_obj1": V_VEC(3, dtype=gs.ti_float, shape=(_B, max_num_polytope_verts)),
+        "local_obj2": V_VEC(3, dtype=gs.ti_float, shape=(_B, max_num_polytope_verts)),
         "id1": V(dtype=gs.ti_int, shape=(_B, max_num_polytope_verts)),
         "id2": V(dtype=gs.ti_int, shape=(_B, max_num_polytope_verts)),
         "mink": V_VEC(3, dtype=gs.ti_float, shape=(_B, max_num_polytope_verts)),
@@ -865,6 +931,7 @@ class StructEPAPolytopeFace:
     normal: V_ANNOTATION
     dist2: V_ANNOTATION
     map_idx: V_ANNOTATION
+    visited: V_ANNOTATION
 
 
 def get_epa_polytope_face(solver, polytope_max_faces):
@@ -876,6 +943,7 @@ def get_epa_polytope_face(solver, polytope_max_faces):
         "normal": V_VEC(3, dtype=gs.ti_float, shape=(_B, polytope_max_faces)),
         "dist2": V(dtype=gs.ti_float, shape=(_B, polytope_max_faces)),
         "map_idx": V(dtype=gs.ti_int, shape=(_B, polytope_max_faces)),
+        "visited": V(dtype=gs.ti_int, shape=(_B, polytope_max_faces)),
     }
 
     if use_ndarray:
@@ -1060,6 +1128,10 @@ class StructGJKState:
     is_col: V_ANNOTATION
     penetration: V_ANNOTATION
     distance: V_ANNOTATION
+    # Differentiable contact detection
+    diff_contact_input: StructDiffContactInput
+    n_diff_contact_input: V_ANNOTATION
+    diff_penetration: V_ANNOTATION
 
 
 def get_gjk_state(solver, static_rigid_sim_config, gjk_static_config):
@@ -1068,6 +1140,7 @@ def get_gjk_state(solver, static_rigid_sim_config, gjk_static_config):
     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
+    requires_grad = solver._static_rigid_sim_config.requires_grad
 
     ### GJK simplex
     simplex_vertex = get_gjk_simplex_vertex(solver)
@@ -1141,6 +1214,9 @@ def get_gjk_state(solver, static_rigid_sim_config, gjk_static_config):
             "is_col": V(dtype=gs.ti_bool, shape=(_B,)),
             "penetration": V(dtype=gs.ti_float, shape=(_B,)),
             "distance": V(dtype=gs.ti_float, shape=(_B,)),
+            "diff_contact_input": get_diff_contact_input(solver, max(1, max_contacts_per_pair if requires_grad else 1)),
+            "n_diff_contact_input": V(dtype=gs.ti_int, shape=(_B,)),
+            "diff_penetration": V(dtype=gs.ti_float, shape=(_B, max_contacts_per_pair)),
         }
     )
 
@@ -1799,9 +1875,10 @@ class StructGeomsState:
 
 def get_geoms_state(solver):
     shape = solver._batch_shape(solver.n_geoms_)
+    requires_grad = solver._static_rigid_sim_config.requires_grad
     kwargs = {
-        "pos": V(dtype=gs.ti_vec3, shape=shape),
-        "quat": V(dtype=gs.ti_vec4, shape=shape),
+        "pos": V(dtype=gs.ti_vec3, shape=shape, needs_grad=requires_grad),
+        "quat": V(dtype=gs.ti_vec4, shape=shape, needs_grad=requires_grad),
         "aabb_min": V(dtype=gs.ti_vec3, shape=shape),
         "aabb_max": V(dtype=gs.ti_vec3, shape=shape),
         "verts_updated": V(dtype=gs.ti_bool, shape=shape),
@@ -2381,3 +2458,4 @@ def __init__(self, solver):
 LinksStateAdjointCache = ti.template() if not use_ndarray else StructLinksState
 JointsStateAdjointCache = ti.template() if not use_ndarray else StructJointsState
 GeomsStateAdjointCache = ti.template() if not use_ndarray else StructGeomsState
+DiffContactInput = ti.template() if not use_ndarray else StructDiffContactInput