[MISC] Cleanup PARA_LEVEL. Add 'GS_PARA_LEVEL' env var. (#1968)

* Cleanup PARA_LEVEL. 
* Add 'GS_PARA_LEVEL' env var.
* Add global EPS to Data manager to please GsTaichi fastcache.

Author: duburcqa

genesis/engine/couplers/legacy_coupler.py

@@ -190,6 +190,7 @@ def _func_collide_with_rigid_geom(
             normal_rigid = sdf_decomp.sdf_func_normal_world(
                 geoms_state=geoms_state,
                 geoms_info=geoms_info,
+                rigid_global_info=rigid_global_info,
                 collider_static_config=collider_static_config,
                 sdf_info=sdf_info,
                 pos_world=pos_world,
@@ -241,6 +242,7 @@ def _func_collide_with_rigid_geom_robust(
         normal_rigid = sdf_decomp.sdf_func_normal_world(
             geoms_state=geoms_state,
             geoms_info=geoms_info,
+            rigid_global_info=rigid_global_info,
             collider_static_config=collider_static_config,
             sdf_info=sdf_info,
             pos_world=pos_world,
@@ -500,6 +502,7 @@ def mpm_surface_to_particle(
         geoms_state: array_class.GeomsState,
         geoms_info: array_class.GeomsInfo,
         sdf_info: array_class.SDFInfo,
+        rigid_global_info: array_class.RigidGlobalInfo,
         collider_static_config: ti.template(),
     ):
         for i_p, i_b in ti.ndrange(self.mpm_solver.n_particles, self.mpm_solver._B):
@@ -509,6 +512,7 @@ def mpm_surface_to_particle(
                         sdf_normal = sdf_decomp.sdf_func_normal_world(
                             geoms_state=geoms_state,
                             geoms_info=geoms_info,
+                            rigid_global_info=rigid_global_info,
                             collider_static_config=collider_static_config,
                             sdf_info=sdf_info,
                             pos_world=self.mpm_solver.particles[f, i_p, i_b].pos,
@@ -730,6 +734,7 @@ def kernel_pbd_rigid_collide(
         geoms_info: array_class.GeomsInfo,
         links_state: array_class.LinksState,
         sdf_info: array_class.SDFInfo,
+        rigid_global_info: array_class.RigidGlobalInfo,
         collider_static_config: ti.template(),
     ):
         for i_p, i_b in ti.ndrange(self.pbd_solver._n_particles, self.sph_solver._B):
@@ -753,6 +758,7 @@ def kernel_pbd_rigid_collide(
                             geoms_info,
                             links_state,
                             sdf_info,
+                            rigid_global_info,
                             collider_static_config,
                         )
 
@@ -856,6 +862,7 @@ def _func_pbd_collide_with_rigid_geom(
         geoms_info: array_class.GeomsInfo,
         links_state: array_class.LinksState,
         sdf_info: array_class.SDFInfo,
+        rigid_global_info: array_class.RigidGlobalInfo,
         collider_static_config: ti.template(),
     ):
         """
@@ -879,6 +886,7 @@ def _func_pbd_collide_with_rigid_geom(
         contact_normal = sdf_decomp.sdf_func_normal_world(
             geoms_state=geoms_state,
             geoms_info=geoms_info,
+            rigid_global_info=rigid_global_info,
             collider_static_config=collider_static_config,
             sdf_info=sdf_info,
             pos_world=pos_world,
@@ -926,6 +934,7 @@ def preprocess(self, f):
                 self.rigid_solver.geoms_state,
                 self.rigid_solver.geoms_info,
                 self.rigid_solver.sdf._sdf_info,
+                self.rigid_solver._rigid_global_info,
                 self.rigid_solver.collider._collider_static_config,
             )
 
@@ -963,6 +972,7 @@ def couple(self, f):
                 geoms_info=self.rigid_solver.geoms_info,
                 links_state=self.rigid_solver.links_state,
                 sdf_info=self.rigid_solver.sdf._sdf_info,
+                rigid_global_info=self.rigid_solver._rigid_global_info,
                 collider_static_config=self.rigid_solver.collider._collider_static_config,
             )
 

genesis/engine/couplers/sap_coupler.py

@@ -377,14 +377,14 @@ def rigid_update_volume_verts_pressure_gradient(
             i_g = self.rigid_volume_verts_geom_idx[i_v]
             pos = geoms_state.pos[i_g, i_b]
             quat = geoms_state.quat[i_g, i_b]
-            R = gu.ti_quat_to_R(quat)
+            R = gu.ti_quat_to_R(quat, gs.EPS)
             self.rigid_volume_verts[i_b, i_v] = R @ self.rigid_volume_verts_rest[i_v] + pos
 
         for i_b, i_e in ti.ndrange(self._B, self.n_rigid_volume_elems):
             i_g = self.rigid_volume_elems_geom_idx[i_e]
             pos = geoms_state.pos[i_g, i_b]
             quat = geoms_state.quat[i_g, i_b]
-            R = gu.ti_quat_to_R(quat)
+            R = gu.ti_quat_to_R(quat, gs.EPS)
             self.rigid_pressure_gradient[i_b, i_e] = R @ self.rigid_pressure_gradient_rest[i_e]
 
     @ti.kernel

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -1431,7 +1431,10 @@ def _kernel_forward_kinematics(
             self._IK_qpos_orig[qs_idx[i_q_], envs_idx[i_b_]] = rigid_global_info.qpos[qs_idx[i_q_], envs_idx[i_b_]]
             # set new qpos
             rigid_global_info.qpos[qs_idx[i_q_], envs_idx[i_b_]] = qpos[i_b_, i_q_]
-            # run FK
+
+        # run FK
+        ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
+        for i_b_ in range(envs_idx.shape[0]):
             gs.engine.solvers.rigid.rigid_solver_decomp.func_forward_kinematics_entity(
                 self._idx_in_solver,
                 envs_idx[i_b_],
@@ -1446,18 +1449,21 @@ def _kernel_forward_kinematics(
                 static_rigid_sim_config,
             )
 
-        ti.loop_config(serialize=ti.static(static_rigid_sim_config.para_level < gs.PARA_LEVEL.PARTIAL))
+        ti.loop_config(serialize=ti.static(static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL))
         for i_l_, i_b_ in ti.ndrange(links_idx.shape[0], envs_idx.shape[0]):
             for i in ti.static(range(3)):
                 links_pos[i_b_, i_l_, i] = links_state.pos[links_idx[i_l_], envs_idx[i_b_]][i]
             for i in ti.static(range(4)):
                 links_quat[i_b_, i_l_, i] = links_state.quat[links_idx[i_l_], envs_idx[i_b_]][i]
 
+        # restore original qpos
         ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
         for i_q_, i_b_ in ti.ndrange(qs_idx.shape[0], envs_idx.shape[0]):
-            # restore original qpos
             rigid_global_info.qpos[qs_idx[i_q_], envs_idx[i_b_]] = self._IK_qpos_orig[qs_idx[i_q_], envs_idx[i_b_]]
-            # run FK
+
+        # run FK
+        ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
+        for i_b_ in range(envs_idx.shape[0]):
             gs.engine.solvers.rigid.rigid_solver_decomp.func_forward_kinematics_entity(
                 self._idx_in_solver,
                 envs_idx[i_b_],

genesis/engine/entities/tool_entity/mesh.py

@@ -114,7 +114,7 @@ def normal(self, f, pos_world, i_b):
         normal_vec_mesh = R_voxels_to_mesh @ normal_vec_voxels
 
         normal_vec_world = gu.ti_transform_by_quat(normal_vec_mesh, self.entity.quat[f, i_b])
-        normal_vec_world = gu.ti_normalize(normal_vec_world)
+        normal_vec_world = gu.ti_normalize(normal_vec_world, gs.EPS)
 
         return normal_vec_world
 
@@ -131,7 +131,7 @@ def normal_(self, pos_voxels):
             dec[i] -= delta
             normal_vec[i] = (self.sdf_(inc) - self.sdf_(dec)) / (2 * delta)
 
-        normal_vec = gu.ti_normalize(normal_vec)
+        normal_vec = gu.ti_normalize(normal_vec, gs.EPS)
 
         return normal_vec
 
@@ -195,7 +195,7 @@ def pbd_collide(self, f, pos_world, thickness, dt):
 
             sdf_voxels = self.sdf_(pos_voxels)
             normal_vec_voxels = self.normal_(pos_voxels)
-            normal_vec_voxels_ = gu.ti_normalize(normal_vec_voxels)
+            normal_vec_voxels_ = gu.ti_normalize(normal_vec_voxels, gs.EPS)
 
             vec_voxels = -sdf_voxels * normal_vec_voxels_
             R_voxels_to_mesh = self.T_mesh_to_sdf[None][:3, :3].inverse()

genesis/engine/entities/tool_entity/tool_entity.py

@@ -182,7 +182,7 @@ def advect(self, f: ti.i32):
             )
             # rotate in world coordinates about itself.
             self.quat[f + 1, i_b] = ti_transform_quat_by_quat(
-                self.quat[f, i_b], ti_rotvec_to_quat(self.ang[f, i_b] * self._solver.substep_dt)
+                self.quat[f, i_b], ti_rotvec_to_quat(self.ang[f, i_b] * self._solver.substep_dt, gs.EPS)
             )
 
             self.vel[f + 1, i_b] = self.vel[f, i_b]

genesis/engine/scene.py

@@ -871,15 +871,18 @@ def _parallelize(
             - for non-batched env, we only parallelize certain loops that have big loop size
             - for batched env, we parallelize all loops
         - When using cpu, we serialize everything.
-            - This is emprically as fast as parallel loops even with big batchsize (tested up to B=10000), because invoking multiple cpu processes cannot utilize all cpu usage.
-            - In order to exploit full cpu power, users are encouraged to launch multiple processes manually, and each will use a single cpu thred.
+            - Parallelization only provides a boost for n_envs >= num_threads and ti_num_threads > 1.
+              It is always disabled by default but can be enforced by setting the env var `GS_PARA_LEVEL=2`.
+            - In order to exploit full cpu power, users are encouraged to launch multiple processes manually and set
+              env var `TI_NUM_THREADS=1`, so that each process uses a single cpu thread.
         """
         if gs.backend == gs.cpu:
-            self._para_level = gs.PARA_LEVEL.NEVER
-        elif self.n_envs == 0:
-            self._para_level = gs.PARA_LEVEL.PARTIAL
+            para_level = gs.PARA_LEVEL.NEVER
+        elif self.n_envs <= 1:
+            para_level = gs.PARA_LEVEL.PARTIAL
         else:
-            self._para_level = gs.PARA_LEVEL.ALL
+            para_level = gs.PARA_LEVEL.ALL
+        self._para_level = int(os.environ.get("GS_PARA_LEVEL", para_level))
 
     @gs.assert_built
     def reset(self, state: SimState | None = None, envs_idx=None):

genesis/engine/sensors/raycaster.py

@@ -193,7 +193,7 @@ def kernel_cast_rays(
         ray_start_world = ti_transform_by_trans_quat(ray_start_local, link_pos, link_quat)
 
         ray_dir_local = ti.math.vec3(ray_directions[i_p, 0], ray_directions[i_p, 1], ray_directions[i_p, 2])
-        ray_direction_world = ti_normalize(ti_transform_by_quat(ray_dir_local, link_quat))
+        ray_direction_world = ti_normalize(ti_transform_by_quat(ray_dir_local, link_quat), gs.EPS)
 
         # --- 2. BVH Traversal ---
         # FIXME: this duplicates the logic in LBVH.query() which also does traversal
@@ -269,7 +269,7 @@ def kernel_cast_rays(
             else:
                 # Local frame output along provided local ray direction
                 hit_point = dist * ti_normalize(
-                    ti.math.vec3(ray_directions[i_p, 0], ray_directions[i_p, 1], ray_directions[i_p, 2])
+                    ti.math.vec3(ray_directions[i_p, 0], ray_directions[i_p, 1], ray_directions[i_p, 2]), gs.EPS
                 )
                 output_hits[i_b, i_p_offset + i_p_sensor * 3 + 0] = hit_point.x
                 output_hits[i_b, i_p_offset + i_p_sensor * 3 + 1] = hit_point.y

genesis/engine/solvers/mpm_solver.py

@@ -369,6 +369,7 @@ def p2g(
                                 sdf_normal_cell = sdf_decomp.sdf_func_normal_world(
                                     geoms_state=geoms_state,
                                     geoms_info=geoms_info,
+                                    rigid_global_info=rigid_global_info,
                                     collider_static_config=collider_static_config,
                                     sdf_info=sdf_info,
                                     pos_world=cell_pos,