[MISC] Add error code to rigid solver. (#1979)

* Add outlier rejection mechanism to FPS meter.
* Add option to override number of broad phase collision pairs.
* Add error code to rigid solver.
* Fix number of contact per pair if box-box collision algo is enabled.
* Fix performance drop.
* Remove redundant contact points in box-box.

Author: duburcqa

genesis/engine/simulator.py

@@ -287,6 +287,8 @@ def step(self, in_backward=False):
 
         if self.rigid_solver.is_active:
             self.rigid_solver.clear_external_force()
+            if self._cur_substep_global % 10 == 0:
+                self.rigid_solver.check_errno()
 
         self._sensor_manager.step()
 

genesis/engine/solvers/rigid/collider_decomp.py

@@ -263,6 +263,7 @@ def build(self):
         # FIXME: AvatarSolver should not inherite from RigidSolver.
         if type(self) is RigidSolver or self.is_active:
             self.data_manager = array_class.DataManager(self)
+            self._errno = self.data_manager.errno
 
             self._rigid_global_info = self.data_manager.rigid_global_info
             if self._use_hibernation:
@@ -871,6 +872,21 @@ def substep(self):
                 contact_island_state=self.constraint_solver.contact_island.contact_island_state,
             )
 
+    def check_errno(self):
+        match kernel_get_errno(self._errno):
+            case 1:
+                max_collision_pairs_broad = self.collider._collider_info.max_collision_pairs_broad[None]
+                gs.raise_exception(
+                    f"Exceeding max number of broad phase candidate contact pairs ({max_contact_pairs}). Please increase "
+                    f"the value of RigidSolver's option 'multiplier_collision_broad_phase'."
+                )
+            case 2:
+                max_contact_pairs = self.collider._collider_info.max_contact_pairs[None]
+                gs.raise_exception(
+                    f"Exceeding max number of contact pairs ({max_contact_pairs}). Please increase the value of "
+                    "RigidSolver's option 'max_collision_pairs'."
+                )
+
     def _kernel_detect_collision(self):
         self.collider.clear()
         self.collider.detection()
@@ -1251,6 +1267,7 @@ def set_state(self, f, state, envs_idx=None):
                 static_rigid_sim_config=self._static_rigid_sim_config,
             )
 
+            self._errno[None] = 0
             self.collider.reset(envs_idx)
             self.collider.clear(envs_idx)
             if self.constraint_solver is not None:
@@ -1609,6 +1626,7 @@ def set_qpos(self, qpos, qs_idx=None, envs_idx=None, *, skip_forward=False, unsa
             qpos = qpos.unsqueeze(0)
         kernel_set_qpos(qpos, qs_idx, envs_idx, self._rigid_global_info, self._static_rigid_sim_config)
 
+        self._errno[None] = 0
         self.collider.reset(envs_idx)
         self.collider.clear(envs_idx)
         if self.constraint_solver is not None:
@@ -1837,6 +1855,7 @@ def set_dofs_position(self, position, dofs_idx=None, envs_idx=None, *, skip_forw
             self._static_rigid_sim_config,
         )
 
+        self._errno[None] = 0
         self.collider.reset(envs_idx)
         self.collider.clear(envs_idx)
         if self.constraint_solver is not None:
@@ -6923,3 +6942,8 @@ def kernel_set_geoms_friction(
     ti.loop_config(serialize=ti.static(static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL))
     for i_g_ in ti.ndrange(geoms_idx.shape[0]):
         geoms_info.friction[geoms_idx[i_g_]] = friction[i_g_]
+
+
+@ti.kernel(fastcache=gs.use_fastcache)
+def kernel_get_errno(errno: array_class.V_ANNOTATION) -> ti.i32:
+    return errno[None]

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -364,6 +364,7 @@ class RigidOptions(Options):
     enable_adjacent_collision: bool = False
     disable_constraint: bool = False
     max_collision_pairs: int = 300
+    multiplier_collision_broad_phase: int = 8
     integrator: gs.integrator = gs.integrator.approximate_implicitfast
     IK_max_targets: int = 6
 

genesis/options/solvers.py

@@ -232,15 +232,14 @@ def get_constraint_state(constraint_solver, solver):
 
     efc_AR_shape = maybe_shape((len_constraints_, len_constraints_, _B), solver._options.noslip_iterations > 0)
     efc_b_shape = maybe_shape((len_constraints_, _B), solver._options.noslip_iterations > 0)
+    jac_relevant_dofs_shape = (len_constraints_, solver.n_dofs_, _B)
+    jac_n_relevant_dofs_shape = (len_constraints_, _B)
 
     return StructConstraintState(
         n_constraints=V(dtype=gs.ti_int, shape=(_B,)),
         ti_n_equalities=V(dtype=gs.ti_int, shape=(_B,)),
-        jac=V(dtype=gs.ti_float, shape=(len_constraints_, solver.n_dofs_, _B)),
         diag=V(dtype=gs.ti_float, shape=(len_constraints_, _B)),
         aref=V(dtype=gs.ti_float, shape=(len_constraints_, _B)),
-        jac_relevant_dofs=V(dtype=gs.ti_int, shape=(len_constraints_, solver.n_dofs_, _B)),
-        jac_n_relevant_dofs=V(dtype=gs.ti_int, shape=(len_constraints_, _B)),
         n_constraints_equality=V(dtype=gs.ti_int, shape=(_B,)),
         n_constraints_frictionloss=V(dtype=gs.ti_int, shape=(_B,)),
         improved=V(dtype=gs.ti_int, shape=(_B,)),
@@ -293,6 +292,10 @@ def get_constraint_state(constraint_solver, solver):
         bw_Ju=V(dtype=gs.ti_float, shape=maybe_shape((len_constraints_, _B), solver._requires_grad)),
         bw_y=V(dtype=gs.ti_float, shape=maybe_shape((len_constraints_, _B), solver._requires_grad)),
         bw_w=V(dtype=gs.ti_float, shape=maybe_shape((len_constraints_, _B), solver._requires_grad)),
+        # /!\ Moving allocation of these tensors at the end improves runtime speed by ~5-10%  /!\
+        jac=V(dtype=gs.ti_float, shape=(len_constraints_, solver.n_dofs_, _B)),
+        jac_relevant_dofs=V(dtype=gs.ti_int, shape=jac_relevant_dofs_shape),
+        jac_n_relevant_dofs=V(dtype=gs.ti_int, shape=jac_n_relevant_dofs_shape),
     )
 
 
@@ -518,8 +521,8 @@ def get_collider_state(
         broad_collision_pairs=V_VEC(2, dtype=gs.ti_int, shape=(max(max_collision_pairs_broad, 1), _B)),
         active_buffer_awake=V(dtype=gs.ti_int, shape=(n_geoms, _B)),
         active_buffer_hib=V(dtype=gs.ti_int, shape=(n_geoms, _B)),
-        box_depth=V(dtype=gs.ti_float, shape=(collider_info.box_MAXCONPAIR[None], _B)),
-        box_points=V_VEC(3, dtype=gs.ti_float, shape=(collider_info.box_MAXCONPAIR[None], _B)),
+        box_depth=V(dtype=gs.ti_float, shape=(collider_static_config.n_contacts_per_pair, _B)),
+        box_points=V_VEC(3, dtype=gs.ti_float, shape=(collider_static_config.n_contacts_per_pair, _B)),
         box_pts=V_VEC(3, dtype=gs.ti_float, shape=(6, _B)),
         box_lines=V_VEC(6, dtype=gs.ti_float, shape=(4, _B)),
         box_linesu=V_VEC(6, dtype=gs.ti_float, shape=(4, _B)),
@@ -555,8 +558,6 @@ class StructColliderInfo(metaclass=BASE_METACLASS):
     mc_perturbation: V_ANNOTATION
     mc_tolerance: V_ANNOTATION
     mpr_to_sdf_overlap_ratio: V_ANNOTATION
-    # maximum number of contact points for box-box collision detection
-    box_MAXCONPAIR: V_ANNOTATION
     # differentiable contact tolerance
     diff_pos_tolerance: V_ANNOTATION
     diff_normal_tolerance: V_ANNOTATION
@@ -586,7 +587,6 @@ def get_collider_info(solver, n_vert_neighbors, collider_static_config, **kwargs
         mc_perturbation=V_SCALAR_FROM(dtype=gs.ti_float, value=kwargs["mc_perturbation"]),
         mc_tolerance=V_SCALAR_FROM(dtype=gs.ti_float, value=kwargs["mc_tolerance"]),
         mpr_to_sdf_overlap_ratio=V_SCALAR_FROM(dtype=gs.ti_float, value=kwargs["mpr_to_sdf_overlap_ratio"]),
-        box_MAXCONPAIR=V_SCALAR_FROM(dtype=gs.ti_int, value=kwargs["box_MAXCONPAIR"]),
         diff_pos_tolerance=V_SCALAR_FROM(dtype=gs.ti_float, value=kwargs["diff_pos_tolerance"]),
         diff_normal_tolerance=V_SCALAR_FROM(dtype=gs.ti_float, value=kwargs["diff_normal_tolerance"]),
     )
@@ -1731,6 +1731,7 @@ class StructRigidSimStaticConfig(metaclass=AutoInitMeta):
 class DataManager:
     def __init__(self, solver):
         self.rigid_global_info = get_rigid_global_info(solver)
+
         self.dofs_info = get_dofs_info(solver)
         self.dofs_state = get_dofs_state(solver)
         self.links_info = get_links_info(solver)
@@ -1758,6 +1759,8 @@ def __init__(self, solver):
         self.entities_info = get_entities_info(solver)
         self.entities_state = get_entities_state(solver)
 
+        self.errno = V_SCALAR_FROM(dtype=gs.ti_int, value=0)
+
 
 DofsState = StructDofsState if gs.use_ndarray else ti.template()
 DofsInfo = StructDofsInfo if gs.use_ndarray else ti.template()

genesis/utils/array_class.py

@@ -177,13 +177,17 @@ def sleep(self):
 
 
 class FPSTracker:
-    def __init__(self, n_envs, alpha=0.95, minimum_interval_seconds: float | None = 0.1):
+    def __init__(
+        self, n_envs, alpha=0.95, minimum_interval_seconds: float | None = 0.05, outlier_threshold: float = 1.5
+    ):
         self.last_time = None
         self.n_envs = n_envs
         self.dt_ema = None
         self.alpha = alpha
         self.minimum_interval_seconds = minimum_interval_seconds
+        self.outlier_threshold = outlier_threshold
         self.steps_since_last_print: int = 0
+        self.total_fps = 0.0
 
     def step(self, current_time: float | None = None) -> float | None:
         if not current_time:
@@ -197,13 +201,21 @@ def step(self, current_time: float | None = None) -> float | None:
 
         self.steps_since_last_print += 1
 
+        # Skip if update is too soon
         if self.minimum_interval_seconds and current_time - self.last_time < self.minimum_interval_seconds:
             return None
 
+        # Outlier rejection
+        if self.dt_ema is not None:
+            if dt > self.dt_ema * self.outlier_threshold or dt * self.outlier_threshold < self.dt_ema:
+                self.dt_ema = dt
+
+        # EMA update
         if self.dt_ema:
             self.dt_ema = self.alpha * self.dt_ema + (1 - self.alpha) * dt
         else:
             self.dt_ema = dt
+
         fps = 1 / self.dt_ema * self.steps_since_last_print
         if self.n_envs > 0:
             self.total_fps = fps * self.n_envs

genesis/utils/tools.py

@@ -6,6 +6,7 @@
 import torch
 
 import genesis as gs
+from genesis.utils.misc import tensor_to_array
 
 from .utils import assert_allclose
 
@@ -75,7 +76,7 @@ def test_muscle(n_envs, muscle_material, show_viewer):
         pos = worm.get_state().pos[0, worm.get_el2v()].mean(1)
         n_units = worm.n_elements
 
-    pos = pos.cpu().numpy()
+    pos = tensor_to_array(pos)
     pos_max, pos_min = pos.max(0), pos.min(0)
     pos_range = pos_max - pos_min
     lu_thresh, fh_thresh = 0.3, 0.6

tests/test_deformable_physics.py

@@ -470,7 +470,9 @@ def test_dynamic_weld(show_viewer, tol):
             size=(0.04, 0.04, 0.04),
             pos=(0.65, 0.0, 0.02),
         ),
-        surface=gs.surfaces.Plastic(color=(1, 0, 0)),
+        surface=gs.surfaces.Plastic(
+            color=(1, 0, 0),
+        ),
     )
     robot = scene.add_entity(
         gs.morphs.MJCF(
@@ -2114,7 +2116,6 @@ def test_terrain_generation(request, show_viewer):
         name="my_terrain",
     )
     # FIXME: Collision detection is very unstable for 'stepping_stones' pattern.
-    # np.random.seed(4)
     terrain = scene.add_entity(gs.morphs.Terrain(**terrain_kwargs))
     obj = scene.add_entity(
         morph=gs.morphs.Box(
@@ -2133,7 +2134,7 @@ def test_terrain_generation(request, show_viewer):
     obj.set_pos(obj_pos_init_rel + torch.tensor(TERRAIN_OFFSET))
 
     # Drop the objects and simulate for a while.
-    for _ in range(500):
+    for _ in range(600):
         scene.step()
 
     # Check that objects are not moving anymore
@@ -2161,7 +2162,7 @@ def test_terrain_generation(request, show_viewer):
 
 
 @pytest.mark.required
-def test_discrete_obstacles_terrain():
+def test_terrain_discrete_obstacles():
     scene = gs.Scene()
     terrain = scene.add_entity(
         gs.morphs.Terrain(
@@ -2182,9 +2183,9 @@ def test_discrete_obstacles_terrain():
     height_field = terrain.geoms[0].metadata["height_field"]
     platform = height_field[5:7, 5:7]
 
-    assert height_field.max().item() == 2.0
-    assert height_field.min().item() == -2.0
-    assert (platform == 0.0).all()
+    assert height_field.max() == 2.0
+    assert height_field.min() == -2.0
+    assert (platform < gs.EPS).all()
 
 
 def test_mesh_to_heightfield(tmp_path, show_viewer):
@@ -3175,7 +3176,13 @@ def test_noslip_iterations(scale, show_viewer, tol):
         rigid_options=gs.options.RigidOptions(
             noslip_iterations=5,
         ),
-        profiling_options=gs.options.ProfilingOptions(show_FPS=False),
+        viewer_options=gs.options.ViewerOptions(
+            camera_pos=(3 * scale, 3 * scale, 3 * scale),
+            camera_lookat=(scale, 0.0, 0.0),
+        ),
+        profiling_options=gs.options.ProfilingOptions(
+            show_FPS=False,
+        ),
         show_viewer=show_viewer,
     )
 
@@ -3188,6 +3195,10 @@ def test_noslip_iterations(scale, show_viewer, tol):
                     pos=(i * scale, 0, 0),
                     fixed=(i == 0),
                 ),
+                surface=gs.surfaces.Default(
+                    color=(*np.random.rand(3), 0.7),
+                ),
+                visualize_contact=True,
             )
         )
     scene.build()
@@ -3204,7 +3215,7 @@ def test_noslip_iterations(scale, show_viewer, tol):
         boxes[2].control_dofs_force(np.array([-safety / coeff_f * n_box * rho * scale**3 * g]), np.array([0]))
         scene.step()
 
-    box_1_z = boxes[1].get_qpos().cpu().numpy()[2]
+    box_1_z = tensor_to_array(boxes[1].get_qpos())[2]
     # allow some small sliding due to first few frames
     # scale = 0.1 is less stable than bigger scale
     assert_allclose(box_1_z, 0.0, atol=4e-2 * scale)
@@ -3216,7 +3227,7 @@ def test_noslip_iterations(scale, show_viewer, tol):
         boxes[2].control_dofs_force(np.array([-safety / coeff_f * n_box * rho * scale**3 * g]), np.array([0]))
         scene.step()
 
-    box_1_z = boxes[1].get_qpos().cpu().numpy()[2]
+    box_1_z = tensor_to_array(boxes[1].get_qpos())[2]
     # it will slip away
     assert box_1_z < -scale
 

tests/test_rigid_physics.py

@@ -362,7 +362,7 @@ def test_pyrender_vec3():
     # repr and tensor conversion
     t = a.as_tensor()
     assert isinstance(t, torch.Tensor)
-    assert_allclose(t.cpu().numpy(), a.v, tol=gs.EPS)
+    assert_allclose(t, a.v, tol=gs.EPS)
 
     # --- Quat tests ---
     q = Quat.from_wxyz(1.0, 0.0, 0.0, 0.0)  # identity
@@ -407,12 +407,12 @@ def test_pyrender_vec3():
     # tensor conversion
     tq = qz.as_tensor()
     assert isinstance(tq, torch.Tensor)
-    assert_allclose(tq.cpu().numpy(), qz.v, tol=gs.EPS)
+    assert_allclose(tq, qz.v, tol=gs.EPS)
 
 
 def test_fps_tracker():
     n_envs = 23
-    tracker = FPSTracker(alpha=0, n_envs=n_envs)
+    tracker = FPSTracker(alpha=0.0, minimum_interval_seconds=0.1, n_envs=n_envs)
     tracker.step(current_time=10.0)
     assert not tracker.step(current_time=10.0)
     assert not tracker.step(current_time=10.0)