[BUG FIX] Various minor cleanup and bug fixes. (Genesis-Embodied-AI#1390)

* Cleanup pre-commit.
* Cleanup rigid solver.
* Cleanup 'Emitter.emit_omit' implementation. 
* Fix some wrong dtypes.
* Fix PBD solver batching.
* Fix 'add_weld_constraint'.

Author: duburcqa

.pre-commit-config.yaml

@@ -1,8 +1,5 @@
 repos:
   - repo: https://github.com/psf/black
-    rev: 25.1.0  # Use the latest version or specify the version you prefer
+    rev: 25.1.0
     hooks:
       - id: black
-        args:
-          - --line-length=120
-          - .

examples/rigid/suction_cup.py

@@ -96,8 +96,8 @@ def main():
 
     # add suction / weld constraint
     rigid = scene.sim.rigid_solver
-    link_cube = np.array([cube.get_link("box_baselink").idx], dtype=gs.np_int)
-    link_franka = np.array([franka.get_link("hand").idx], dtype=gs.np_int)
+    link_cube = cube.get_link("box_baselink").idx
+    link_franka = franka.get_link("hand").idx
     rigid.add_weld_constraint(link_cube, link_franka)
 
     # lift

genesis/engine/entities/emitter.py

@@ -150,7 +150,7 @@ def emit(
             positions = gu.transform_by_trans_R(
                 positions,
                 pos,
-                gu.z_up_to_R(direction) @ gu.axis_angle_to_R(np.array([0, 0, 1]), theta),
+                gu.z_up_to_R(direction) @ gu.axis_angle_to_R(np.array([0.0, 0.0, 1.0], dtype=gs.np_float), theta),
             ).astype(gs.np_float)
 
             positions = np.tile(positions[np.newaxis], (self._sim._B, 1, 1))
@@ -206,17 +206,21 @@ def emit_omni(self, source_radius=0.1, pos=(0.5, 0.5, 1.0), speed=1.0, particle_
         Parameters:
         ----------
         source_radius: float, optional
-            The radius of the sphere source. Particles will be emitted from a shell with inner radius using 0.8 * source_radius and outer radius using source_radius.
+            The radius of the sphere source. Particles will be emitted from a shell with inner radius using
+            '0.8 * source_radius' and outer radius using source_radius.
         pos: array_like, shape=(3,)
             The center of the sphere source.
         speed: float
             The speed of the emitted particles.
         particle_size: float | None
-            The size (diameter) of the emitted particles. The actual number of particles emitted is determined by the volume of the sphere source and the size of the particles. If None, the solver's particle size is used. Note that this particle size only affects computation for number of particles emitted, not the actual size of the particles in simulation and rendering.
+            The size (diameter) of the emitted particles. The actual number of particles emitted is determined by the
+            volume of the sphere source and the size of the particles. If None, the solver's particle size is used.
+            Note that this particle size only affects computation for number of particles emitted, not the actual size
+            of the particles in simulation and rendering.
         """
         assert self._entity is not None
 
-        pos = np.array(pos)
+        pos = np.asarray(pos, dtype=gs.np_float)
 
         if particle_size is None:
             particle_size = self._solver.particle_size
@@ -227,20 +231,20 @@ def emit_omni(self, source_radius=0.1, pos=(0.5, 0.5, 1.0), speed=1.0, particle_
             inner_radius=source_radius * 0.4,
             sampler=self._entity.sampler,
         )
-
-        positions = gu.transform_by_T(positions_, gu.trans_to_T(pos)).astype(gs.np_float)
+        positions = pos + positions_
 
         if not self._solver.boundary.is_inside(positions):
             gs.raise_exception("Emitted particles are outside the boundary.")
 
-        n_particles = len(positions)
-        dists = np.linalg.norm(positions_, axis=1, keepdims=True)
-        positions[np.where(dists < gs.EPS)[0]] = np.array([gs.EPS, gs.EPS, gs.EPS])
-        vels = (positions_ / dists * speed).astype(gs.np_float)
+        dists = np.linalg.norm(positions_, axis=1)
+        positions[dists < gs.EPS] = gs.EPS
+        vels = (speed / (dists + gs.EPS)) * positions_
 
+        n_particles = len(positions)
         if n_particles > self._entity.n_particles:
             gs.logger.warning(
-                f"Number of particles to emit ({n_particles}) at the current step is larger than the maximum number of particles ({self._entity.n_particles})."
+                f"Number of particles to emit ({n_particles}) at the current step is larger than the maximum number "
+                f"of particles ({self._entity.n_particles})."
             )
 
         self._solver._kernel_set_particles_pos(

genesis/engine/entities/mpm_entity.py

@@ -89,14 +89,9 @@ def set_pos(self, f, pos):
         f : int
             The current substep index.
         pos : gs.Tensor
-            A tensor of shape (n_particles, 3) representing particle positions.
+            A tensor of shape (n_envs, n_particles, 3) representing particle positions.
         """
-        self.solver._kernel_set_particles_pos(
-            f,
-            self._particle_start,
-            self._n_particles,
-            pos,
-        )
+        self.solver._kernel_set_particles_pos(f, self._particle_start, self._n_particles, pos)
 
     def set_pos_grad(self, f, pos_grad):
         """

genesis/engine/entities/sph_entity.py

@@ -92,12 +92,7 @@ def set_pos(self, f, pos):
         pos : ndarray
             Array of particle positions of shape (n_envs, n_particles, 3).
         """
-        self.solver._kernel_set_particles_pos(
-            f,
-            self._particle_start,
-            self._n_particles,
-            pos,
-        )
+        self.solver._kernel_set_particles_pos(f, self._particle_start, self._n_particles, pos)
 
     def set_pos_grad(self, f: ti.i32, pos_grad: ti.types.ndarray()):
         """

genesis/engine/solvers/mpm_solver.py

@@ -720,9 +720,9 @@ def _kernel_set_particles_pos(
                 self.particles[f, i_global, i_b].pos[k] = pos[i_b, i_p, k]
 
             # we restore these whenever directly setting positions
-            self.particles[f, i_global, i_b].vel = ti.Vector.zero(gs.ti_float, 3)
+            self.particles[f, i_global, i_b].vel.fill(0.0)
             self.particles[f, i_global, i_b].F = ti.Matrix.identity(gs.ti_float, 3)
-            self.particles[f, i_global, i_b].C = ti.Matrix.zero(gs.ti_float, 3, 3)
+            self.particles[f, i_global, i_b].C.fill(0.0)
             self.particles[f, i_global, i_b].Jp = self.particles_info[i_global].default_Jp
 
     @ti.kernel

genesis/engine/solvers/pbd_solver.py

@@ -840,8 +840,8 @@ def _kernel_set_particles_pos(
         for i_p, i_b in ti.ndrange(n_particles, self._B):
             i_global = i_p + particle_start
             for k in ti.static(range(3)):
-                self.particles[i_global, i_b].pos[k] = pos[i_p, k]
-            self.particles[i_global, i_b].vel = ti.Vector.zero(gs.ti_float, 3)
+                self.particles[i_global, i_b].pos[k] = pos[i_b, i_p, k]
+            self.particles[i_global, i_b].vel.fill(0.0)
 
     @ti.kernel
     def _kernel_set_particles_vel(

genesis/engine/solvers/rigid/collider_decomp.py

@@ -613,12 +613,12 @@ def _func_contact_mpr_terrain(self, i_ga, i_gb, i_b):
                 r_max = ti.min(self._solver.terrain_rc[0] - 1, r_max)
                 c_max = ti.min(self._solver.terrain_rc[1] - 1, c_max)
 
-                cnt = 0
+                n_con = 0
                 for r in range(r_min, r_max):
                     nvert = 0
                     for c in range(c_min, c_max + 1):
                         for i in range(2):
-                            if cnt < self._n_contacts_per_pair:
+                            if n_con < ti.static(self._n_contacts_per_pair):
                                 nvert = nvert + 1
                                 self.add_prism_vert(
                                     sh * (r + i) + self._solver.terrain_xyz_maxmin[3],
@@ -651,17 +651,17 @@ def _func_contact_mpr_terrain(self, i_ga, i_gb, i_b):
                                         contact_pos = contact_pos + gb_pos
 
                                         valid = True
-                                        i_col = self.n_contacts[i_b]
-                                        for j in range(cnt):
+                                        i_c = self.n_contacts[i_b]
+                                        for j in range(n_con):
                                             if (
-                                                contact_pos - self.contact_data[i_col - j - 1, i_b].pos
+                                                contact_pos - self.contact_data[i_c - j - 1, i_b].pos
                                             ).norm() < tolerance:
                                                 valid = False
                                                 break
 
                                         if valid:
                                             self._func_add_contact(i_ga, i_gb, normal, contact_pos, penetration, i_b)
-                                            cnt = cnt + 1
+                                            n_con = n_con + 1
 
         self._solver.geoms_state[i_ga, i_b].pos, self._solver.geoms_state[i_ga, i_b].quat = ga_pos, ga_quat
         self._solver.geoms_state[i_gb, i_b].pos, self._solver.geoms_state[i_gb, i_b].quat = gb_pos, gb_quat
@@ -787,12 +787,13 @@ def _func_broad_phase(self):
                                     self.contact_cache[i_ga, i_gb, i_b].normal.fill(0.0)
                                 continue
 
-                            if self.n_broad_pairs[i_b] == self._max_collision_pairs:
+                            i_p = self.n_broad_pairs[i_b]
+                            if i_p == self._max_collision_pairs:
                                 # print(self._warn_msg_max_collision_pairs)
                                 break
-                            self.broad_collision_pairs[self.n_broad_pairs[i_b], i_b][0] = i_ga
-                            self.broad_collision_pairs[self.n_broad_pairs[i_b], i_b][1] = i_gb
-                            self.n_broad_pairs[i_b] = self.n_broad_pairs[i_b] + 1
+                            self.broad_collision_pairs[i_p, i_b][0] = i_ga
+                            self.broad_collision_pairs[i_p, i_b][1] = i_gb
+                            self.n_broad_pairs[i_b] += 1
 
                         self.active_buffer[n_active, i_b] = self.sort_buffer[i, i_b].i_g
                         n_active = n_active + 1
@@ -1058,9 +1059,9 @@ def _func_narrow_phase_nonconvex_vs_nonterrain(self):
 
                                             # Discard contact point is repeated
                                             repeated = False
-                                            for i_con in range(n_con):
+                                            for i_c in range(n_con):
                                                 if not repeated:
-                                                    idx_prev = self.n_contacts[i_b] - 1 - i_con
+                                                    idx_prev = self.n_contacts[i_b] - 1 - i_c
                                                     prev_contact = self.contact_data[idx_prev, i_b].pos
                                                     if (contact_pos - prev_contact).norm() < tolerance:
                                                         repeated = True
@@ -1106,7 +1107,7 @@ def _func_plane_box_contact(self, i_ga, i_gb, i_b):
                 contact_pos_0 = contact_pos
                 tolerance = self._func_compute_tolerance(i_ga, i_gb, i_b)
                 for i_v in range(gb_info.vert_start, gb_info.vert_end):
-                    if n_con < self._n_contacts_per_pair:
+                    if n_con < ti.static(self._n_contacts_per_pair):
                         pos_corner = gu.ti_transform_by_trans_quat(
                             self._solver.verts_info[i_v].init_pos, gb_state.pos, gb_state.quat
                         )
@@ -1119,10 +1120,8 @@ def _func_plane_box_contact(self, i_ga, i_gb, i_b):
 
     @ti.func
     def _func_add_contact(self, i_ga, i_gb, normal, contact_pos, penetration, i_b):
-        # print(f"Adding contact {i_ga} {i_gb}, normal:", normal, "contact_pos:", contact_pos, "penetration:", penetration)
-        i_col = self.n_contacts[i_b]
-
-        if i_col == self._max_contact_pairs:
+        i_c = self.n_contacts[i_b]
+        if i_c == self._max_contact_pairs:
             # FIXME: 'ti.static_print' cannot be used as it will be printed systematically, completely ignoring guard
             # condition, while 'print' is slowing down the kernel even if every called in practice...
             # print(self._warn_msg_max_collision_pairs)
@@ -1135,17 +1134,17 @@ def _func_add_contact(self, i_ga, i_gb, normal, contact_pos, penetration, i_b):
             friction_b = gb_info.friction * self._solver.geoms_state[i_gb, i_b].friction_ratio
 
             # b to a
-            self.contact_data[i_col, i_b].geom_a = i_ga
-            self.contact_data[i_col, i_b].geom_b = i_gb
-            self.contact_data[i_col, i_b].normal = normal
-            self.contact_data[i_col, i_b].pos = contact_pos
-            self.contact_data[i_col, i_b].penetration = penetration
-            self.contact_data[i_col, i_b].friction = ti.max(ti.max(friction_a, friction_b), 1e-2)
-            self.contact_data[i_col, i_b].sol_params = 0.5 * (ga_info.sol_params + gb_info.sol_params)
-            self.contact_data[i_col, i_b].link_a = ga_info.link_idx
-            self.contact_data[i_col, i_b].link_b = gb_info.link_idx
-
-            self.n_contacts[i_b] = i_col + 1
+            self.contact_data[i_c, i_b].geom_a = i_ga
+            self.contact_data[i_c, i_b].geom_b = i_gb
+            self.contact_data[i_c, i_b].normal = normal
+            self.contact_data[i_c, i_b].pos = contact_pos
+            self.contact_data[i_c, i_b].penetration = penetration
+            self.contact_data[i_c, i_b].friction = ti.max(ti.max(friction_a, friction_b), 1e-2)
+            self.contact_data[i_c, i_b].sol_params = 0.5 * (ga_info.sol_params + gb_info.sol_params)
+            self.contact_data[i_c, i_b].link_a = ga_info.link_idx
+            self.contact_data[i_c, i_b].link_b = gb_info.link_idx
+
+            self.n_contacts[i_b] = i_c + 1
 
     @ti.func
     def _func_compute_tolerance(self, i_ga, i_gb, i_b):
@@ -1325,7 +1324,7 @@ def _func_convex_convex_contact(self, i_ga, i_gb, i_b):
                                     # add the discovered contact points and stop multi-contact search.
                                     for i_c in range(n_contacts):
                                         # Ignore contact points if the number of contacts exceeds the limit.
-                                        if i_c < self._n_contacts_per_pair:
+                                        if i_c < ti.static(self._n_contacts_per_pair):
                                             contact_pos = self._gjk.contact_pos[i_b, i_c]
                                             normal = self._gjk.normal[i_b, i_c]
                                             self._func_add_contact(i_ga, i_gb, normal, contact_pos, penetration, i_b)
@@ -1453,9 +1452,9 @@ def _func_convex_convex_contact(self, i_ga, i_gb, i_b):
 
                     # Discard contact point is repeated
                     repeated = False
-                    for i_con in range(n_con):
+                    for i_c in range(n_con):
                         if not repeated:
-                            idx_prev = self.n_contacts[i_b] - 1 - i_con
+                            idx_prev = self.n_contacts[i_b] - 1 - i_c
                             prev_contact = self.contact_data[idx_prev, i_b].pos
                             if (contact_pos - prev_contact).norm() < tolerance:
                                 repeated = True