[MISC] Only enable GJK by default if gradient computation is required. (#1984)

* Increase default number of compile threads.
* Do not disable self-collision in benchmarks anymore.
* Only enable GJK by default if gradient computation is required.

Author: duburcqa

genesis/__init__.py

@@ -214,6 +214,7 @@ def init(
     if ti_num_cpu_threads is not None:
         taichi_kwargs.update(
             cpu_max_num_threads=int(ti_num_cpu_threads),
+            num_compile_threads=int(ti_num_cpu_threads),
         )
 
     if seed is not None:

genesis/engine/scene.py

@@ -283,7 +283,9 @@ def _validate_options(
             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:
+        if rigid_options.use_gjk_collision is None:
+            rigid_options.use_gjk_collision = sim_options.requires_grad
+        elif 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."
             )

genesis/options/solvers.py

@@ -349,7 +349,8 @@ class RigidOptions(Options):
     max_dynamic_constraints : int, optional
         Maximum number of dynamic constraints (like suction cup). Defaults to 8.
     use_gjk_collision: bool, optional
-        Whether to use GJK for collision detection instead of MPR. Defaults to True.
+        Whether to use GJK for collision detection instead of MPR. More stable but much slower. Defaults to
+        `sim_options.requires_grad`.
 
     Warning
     -------
@@ -400,7 +401,7 @@ class RigidOptions(Options):
     enable_mujoco_compatibility: bool = False
 
     # GJK collision detection
-    use_gjk_collision: bool = True
+    use_gjk_collision: Optional[bool] = None
 
     def __init__(self, **data):
         super().__init__(**data)

tests/test_grad.py

@@ -140,7 +140,6 @@ def compute_dL_error(dL_dx, x_type):
 # stable way.
 @pytest.mark.required
 @pytest.mark.precision("64")
-@pytest.mark.parametrize("backend", [gs.cpu])
 def test_diff_solver(backend, monkeypatch):
     from genesis.engine.solvers.rigid.constraint_solver_decomp import func_init_solver, func_solve
     from genesis.engine.solvers.rigid.rigid_solver_decomp import kernel_step_1

tests/test_rigid_benchmarks.py

@@ -330,7 +330,6 @@ def anymal_c(solver, n_envs, gjk):
         rigid_options=gs.options.RigidOptions(
             **get_rigid_solver_options(
                 dt=STEP_DT,
-                enable_self_collision=False,
                 **(dict(constraint_solver=solver) if solver is not None else {}),
                 **(dict(use_gjk_collision=gjk) if gjk is not None else {}),
             )
@@ -403,7 +402,6 @@ def batched_franka(solver, n_envs, gjk):
         rigid_options=gs.options.RigidOptions(
             **get_rigid_solver_options(
                 dt=STEP_DT,
-                enable_self_collision=False,
                 **(dict(constraint_solver=solver) if solver is not None else {}),
                 **(dict(use_gjk_collision=gjk) if gjk is not None else {}),
             )
@@ -455,7 +453,6 @@ def random(solver, n_envs, gjk):
         rigid_options=gs.options.RigidOptions(
             **get_rigid_solver_options(
                 dt=STEP_DT,
-                enable_self_collision=False,
                 **(dict(constraint_solver=solver) if solver is not None else {}),
                 **(dict(use_gjk_collision=gjk) if gjk is not None else {}),
             )
@@ -626,10 +623,12 @@ def box_pyramid(solver, n_envs, n_cubes, enable_island, gjk, enable_mujoco_compa
 @pytest.mark.parametrize(
     "runnable, solver, gjk, n_envs, backend",
     [
+        ("anymal_c", None, True, 30000, gs.gpu),
         ("anymal_c", gs.constraint_solver.CG, None, 30000, gs.gpu),
         ("anymal_c", gs.constraint_solver.Newton, None, 30000, gs.gpu),
         ("anymal_c", None, None, 0, gs.gpu),
         ("anymal_c", None, None, 0, gs.cpu),
+        ("batched_franka", None, True, 30000, gs.gpu),
         ("batched_franka", gs.constraint_solver.CG, None, 30000, gs.gpu),
         ("batched_franka", gs.constraint_solver.Newton, None, 30000, gs.gpu),
         ("batched_franka", None, None, 0, gs.gpu),

tests/test_rigid_physics.py

@@ -1756,8 +1756,8 @@ def test_set_dofs_frictionloss_physics(gs_sim, tol):
 def test_frictionloss_advanced(show_viewer, tol):
     scene = gs.Scene(
         viewer_options=gs.options.ViewerOptions(
-            camera_pos=(1.0, 0.25, 0.75),
-            camera_lookat=(0.0, 0.0, 0.0),
+            camera_pos=(1.4, 0.7, 1.4),
+            camera_lookat=(0.6, 0.0, 0.0),
         ),
         show_viewer=show_viewer,
         show_FPS=False,
@@ -1768,21 +1768,25 @@ def test_frictionloss_advanced(show_viewer, tol):
         morph=gs.morphs.MJCF(
             file=f"{asset_path}/SO101/so101_new_calib.xml",
         ),
+        # vis_mode="collision",
     )
     box = scene.add_entity(
         gs.morphs.Box(
+            pos=(0.1, 0.0, 0.6),
             size=(0.025, 0.025, 0.025),
         ),
     )
     scene.build()
 
     scene.reset()
-    box.set_pos(torch.tensor((0.1, 0.0, 1.0), dtype=gs.tc_float, device=gs.device))
-    for _ in range(200):
+    for _ in range(230):
         scene.step()
 
     assert_allclose(robot.get_contacts()["position"][:, 2].min(), 0.0, tol=1e-4)
-    # assert_allclose(robot.get_AABB()[2], 0.0, tol=1e-3)
+    assert_allclose(robot.get_AABB()[0, 2], 0.0, tol=2e-4)
+    box_pos = box.get_pos()
+    assert box_pos[0] > 0.6
+    assert_allclose(box_pos[1:], 0.0, tol=0.02)
     assert_allclose(box.get_dofs_velocity(), 0.0, tol=tol)
 
 
@@ -3112,8 +3116,8 @@ def test_contype_conaffinity(show_viewer, tol):
         scene.step()
 
     assert_allclose(box1.get_pos(), (0.0, 0.0, 0.25), atol=5e-4)
-    assert_allclose(box2.get_pos(), (0.0, 0.0, 0.75), atol=1e-3)
-    assert_allclose(box2.get_pos(), box3.get_pos(), atol=1e-4)
+    assert_allclose(box2.get_pos(), (0.0, 0.0, 0.75), atol=2e-3)
+    assert_allclose(box2.get_pos(), box3.get_pos(), atol=2e-3)
     assert_allclose(scene.rigid_solver.get_links_acc(slice(box4.link_start, box4.link_end)), GRAVITY, atol=tol)