[MISC] Improve robust many box unit test. (#2001)

* Improve robust many box unit test.
* Improve unit test coverage.
* More efficient 'kernel_update_all_verts'.
* Disable fastcache on the CI because it causes crashes.

Author: duburcqa

.github/workflows/generic.yml

@@ -68,6 +68,8 @@ jobs:
       PY_COLORS: "1"
       GS_CACHE_FILE_PATH: ".cache/genesis"
       GS_ENABLE_NDARRAY: ${{ matrix.GS_ENABLE_NDARRAY }}
+      # FIXME: Disabling fastcache because it causes crashes for some reason...
+      GS_ENABLE_FASTCACHE: "0"
       GS_TORCH_FORCE_CPU_DEVICE: ${{ startsWith(matrix.OS, 'macos-') && '1' || '0' }}
       TI_OFFLINE_CACHE: "1"
       TI_OFFLINE_CACHE_CLEANING_POLICY: "never"

genesis/engine/couplers/sap_coupler.py

@@ -601,6 +601,7 @@ def precompute(self, i_step):
 
         if self.rigid_solver.is_active:
             kernel_update_all_verts(
+                geoms_info=self.rigid_solver.geoms_info,
                 geoms_state=self.rigid_solver.geoms_state,
                 verts_info=self.rigid_solver.verts_info,
                 free_verts_state=self.rigid_solver.free_verts_state,

genesis/engine/sensors/raycaster.py

@@ -334,6 +334,7 @@ def _update_bvh(cls, shared_metadata: RaycasterSharedMetadata):
         from genesis.engine.solvers.rigid.rigid_solver_decomp import kernel_update_all_verts
 
         kernel_update_all_verts(
+            geoms_info=shared_metadata.solver.geoms_info,
             geoms_state=shared_metadata.solver.geoms_state,
             verts_info=shared_metadata.solver.verts_info,
             free_verts_state=shared_metadata.solver.free_verts_state,

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -5005,9 +5005,9 @@ def func_update_verts_for_geom(
     fixed_verts_state: array_class.VertsState,
 ):
     if not geoms_state.verts_updated[i_g, i_b]:
-        i_g_start = geoms_info.vert_start[i_g]
-        if verts_info.is_fixed[i_g_start]:
-            for i_v in range(i_g_start, geoms_info.vert_end[i_g]):
+        i_v_start = geoms_info.vert_start[i_g]
+        if verts_info.is_fixed[i_v_start]:
+            for i_v in range(i_v_start, geoms_info.vert_end[i_g]):
                 verts_state_idx = verts_info.verts_state_idx[i_v]
                 fixed_verts_state.pos[verts_state_idx] = gu.ti_transform_by_trans_quat(
                     verts_info.init_pos[i_v], geoms_state.pos[i_g, i_b], geoms_state.quat[i_g, i_b]
@@ -5016,7 +5016,7 @@ def func_update_verts_for_geom(
             for j_b in range(_B):
                 geoms_state.verts_updated[i_g, j_b] = True
         else:
-            for i_v in range(i_g_start, geoms_info.vert_end[i_g]):
+            for i_v in range(i_v_start, geoms_info.vert_end[i_g]):
                 verts_state_idx = verts_info.verts_state_idx[i_v]
                 free_verts_state.pos[verts_state_idx, i_b] = gu.ti_transform_by_trans_quat(
                     verts_info.init_pos[i_v], geoms_state.pos[i_g, i_b], geoms_state.quat[i_g, i_b]
@@ -5026,34 +5026,26 @@ def func_update_verts_for_geom(
 
 @ti.func
 def func_update_all_verts(
+    geoms_info: array_class.GeomsInfo,
     geoms_state: array_class.GeomsState,
     verts_info: array_class.VertsInfo,
     free_verts_state: array_class.VertsState,
     fixed_verts_state: array_class.VertsState,
 ):
-    n_verts = verts_info.geom_idx.shape[0]
-    _B = geoms_state.pos.shape[1]
-    for i_v, i_b in ti.ndrange(n_verts, _B):
-        i_g = verts_info.geom_idx[i_v]
-        verts_state_idx = verts_info.verts_state_idx[i_v]
-        if verts_info.is_fixed[i_v]:
-            fixed_verts_state.pos[verts_state_idx] = gu.ti_transform_by_trans_quat(
-                verts_info.init_pos[i_v], geoms_state.pos[i_g, i_b], geoms_state.quat[i_g, i_b]
-            )
-        else:
-            free_verts_state.pos[verts_state_idx, i_b] = gu.ti_transform_by_trans_quat(
-                verts_info.init_pos[i_v], geoms_state.pos[i_g, i_b], geoms_state.quat[i_g, i_b]
-            )
+    n_geoms, _B = geoms_state.pos.shape
+    for i_g, i_b in ti.ndrange(n_geoms, _B):
+        func_update_verts_for_geom(i_g, i_b, geoms_state, geoms_info, verts_info, free_verts_state, fixed_verts_state)
 
 
 @ti.kernel(fastcache=gs.use_fastcache)
 def kernel_update_all_verts(
+    geoms_info: array_class.GeomsInfo,
     geoms_state: array_class.GeomsState,
     verts_info: array_class.VertsInfo,
     free_verts_state: array_class.VertsState,
     fixed_verts_state: array_class.VertsState,
 ):
-    func_update_all_verts(geoms_state, verts_info, free_verts_state, fixed_verts_state)
+    func_update_all_verts(geoms_info, geoms_state, verts_info, free_verts_state, fixed_verts_state)
 
 
 @ti.kernel
@@ -5115,7 +5107,6 @@ def func_hibernate__for_all_awake_islands_either_hiberanate_or_update_aabb_sort_
     static_rigid_sim_config: ti.template(),
     contact_island_state: array_class.ContactIslandState,
 ) -> None:
-
     n_entities = entities_state.hibernated.shape[0]
     _B = entities_state.hibernated.shape[1]
 

tests/test_rigid_physics.py

@@ -877,6 +877,7 @@ def test_many_boxes_dynamics(box_box_detection, gjk_collision, dynamics, show_vi
             dt=0.01,
         ),
         rigid_options=gs.options.RigidOptions(
+            max_collision_pairs=1000,
             box_box_detection=box_box_detection,
             use_gjk_collision=gjk_collision,
         ),
@@ -893,7 +894,7 @@ def test_many_boxes_dynamics(box_box_detection, gjk_collision, dynamics, show_vi
         i, j, k = int(n / 25), int(n / 5) % 5, n % 5
         scene.add_entity(
             gs.morphs.Box(
-                pos=(i * 1.01, j * 1.01, k * 1.01 + 0.5),
+                pos=(i * (1.0 - 1e-3), j * (1.0 - 1e-3), k * (1.0 - 1e-3) + 0.5),
                 size=(1.0, 1.0, 1.0),
             ),
             surface=gs.surfaces.Default(
@@ -920,7 +921,7 @@ def test_many_boxes_dynamics(box_box_detection, gjk_collision, dynamics, show_vi
             assert qpos[:2].norm() < 20.0
             assert qpos[2] < 5.0
         else:
-            qpos0 = np.array((i * 1.01, j * 1.01, k * 1.01 + 0.5))
+            qpos0 = np.array((i * (1.0 - 1e-3), j * (1.0 - 1e-3), k * (1.0 - 1e-3) + 0.5))
             assert_allclose(qpos[:3], qpos0, atol=0.05)
             assert_allclose(qpos[3:], 0, atol=0.03)
 
@@ -1220,6 +1221,13 @@ def test_set_root_pose(batch_fixed_verts, relative, show_viewer, tol):
             batch_fixed_verts=batch_fixed_verts,
         ),
     )
+    sphere = scene.add_entity(
+        gs.morphs.Sphere(
+            radius=0.04,
+            batch_fixed_verts=False,
+            fixed=True,
+        ),
+    )
     cube = scene.add_entity(
         gs.morphs.Box(
             size=(0.04, 0.04, 0.04),
@@ -1244,6 +1252,7 @@ def test_set_root_pose(batch_fixed_verts, relative, show_viewer, tol):
             scene.visualizer.update()
     cube.set_pos(pos_delta[[0]] + (0.0, 0.0, 0.16), envs_idx=[0])
     cube.set_pos(pos_delta[[1]] + (0.0, 0.0, 0.11), envs_idx=[1])
+    sphere.set_pos(np.tile(pos_delta[[0]], (2, 1)) + 1.0)
     quat_delta = np.random.rand(2, 4)
     with nullcontext() if batch_fixed_verts else pytest.raises(gs.GenesisException):
         robot.set_quat(quat_delta)
@@ -1257,6 +1266,10 @@ def test_set_root_pose(batch_fixed_verts, relative, show_viewer, tol):
     if show_viewer:
         scene.visualizer.update()
 
+    sphere_aabb, sphere_base_aabb = sphere.get_AABB(), sphere.geoms[0].get_AABB()
+    assert_allclose(sphere_aabb.mean(dim=-2), pos_delta[0] + 1.0, tol=tol)
+    assert_allclose(sphere.get_AABB(), sphere.geoms[0].get_AABB(), tol=tol)
+
     # Simulate for a while to check if the dynamic object is colliding with the static one
     if batch_fixed_verts:
         has_collided = torch.tensor([False, False], dtype=torch.bool, device=gs.device)