[MISC] Fix differentiable unit tests not running on the CI. (#1963)

* Fix differentiable unit tests not running on the CI.
* Fix AD support.

Author: duburcqa

.github/workflows/production.yml

@@ -39,6 +39,9 @@ jobs:
       matrix:
         GS_ENABLE_NDARRAY: ["0", "1"]
 
+    env:
+      GS_ENABLE_NDARRAY: ${{ matrix.GS_ENABLE_NDARRAY }}
+
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -75,7 +78,7 @@ jobs:
               pip install --no-input --extra-index-url https://pypi.nvidia.com/ omniverse-kit
               pip install --no-input ".[dev,render,usd]"
 
-              pytest -v --backend gpu --dev --forked ./tests
+              pytest -v -rs --backend gpu --dev --forked ./tests
           EOF
       - name: Kill srun job systematically
         if: always()

genesis/engine/entities/emitter.py

@@ -230,7 +230,7 @@ def emit_omni(self, source_radius=0.1, pos=(0.5, 0.5, 1.0), speed=1.0, particle_
 
         n_particles = len(positions)
         if n_particles > self._entity.n_particles:
-            gs.logger.warning(
+            gs.raise_exception(
                 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})."
             )

genesis/engine/solvers/mpm_solver.py

@@ -398,10 +398,9 @@ def g2p(
         geoms_info: array_class.GeomsInfo,
         links_state: array_class.LinksState,
         rigid_global_info: array_class.RigidGlobalInfo,
-    ) -> ti.i32:
-        is_success = True
+    ):
         for i_p, i_b in ti.ndrange(self._n_particles, self._B):
-            if is_success and self.particles_ng[f, i_p, i_b].active:
+            if self.particles_ng[f, i_p, i_b].active:
                 base = ti.floor(self.particles[f, i_p, i_b].pos * self._inv_dx - 0.5).cast(gs.ti_int)
                 fx = self.particles[f, i_p, i_b].pos * self._inv_dx - base.cast(gs.ti_float)
                 w = [0.5 * (1.5 - fx) ** 2, 0.75 - (fx - 1.0) ** 2, 0.5 * (fx - 0.5) ** 2]
@@ -433,19 +432,16 @@ def g2p(
                     new_vel += weight * grid_vel
                     new_C += 4 * self._inv_dx * weight * grid_vel.outer_product(dpos)
 
-                if not ti.math.isnan(new_vel).any():
-                    # compute actual new_pos with new_vel
-                    new_pos = self.particles[f, i_p, i_b].pos + self.substep_dt * new_vel
+                # compute actual new_pos with new_vel
+                new_pos = self.particles[f, i_p, i_b].pos + self.substep_dt * new_vel
 
-                    # impose boundary for safety, in case simulation explodes and tries to access illegal cell address
-                    new_pos, new_vel = self.boundary.impose_pos_vel(new_pos, new_vel)
+                # impose boundary for safety, in case simulation explodes and tries to access illegal cell address
+                new_pos, new_vel = self.boundary.impose_pos_vel(new_pos, new_vel)
 
-                    # advect to next frame
-                    self.particles[f + 1, i_p, i_b].vel = new_vel
-                    self.particles[f + 1, i_p, i_b].C = new_C
-                    self.particles[f + 1, i_p, i_b].pos = new_pos
-                else:
-                    is_success = False
+                # advect to next frame
+                self.particles[f + 1, i_p, i_b].vel = new_vel
+                self.particles[f + 1, i_p, i_b].C = new_C
+                self.particles[f + 1, i_p, i_b].pos = new_pos
             else:
                 self.particles[f + 1, i_p, i_b].vel = self.particles[f, i_p, i_b].vel
                 self.particles[f + 1, i_p, i_b].pos = self.particles[f, i_p, i_b].pos
@@ -455,6 +451,12 @@ def g2p(
 
             self.particles_ng[f + 1, i_p, i_b].active = self.particles_ng[f, i_p, i_b].active
 
+    @ti.kernel
+    def _is_state_valid(self, f: ti.i32) -> ti.i32:
+        is_success = True
+        for i_p, i_b in ti.ndrange(self._n_particles, self._B):
+            if ti.math.isnan(self.particles[f, i_p, i_b].pos).any():
+                is_success = False
         return is_success
 
     # ------------------------------------------------------------------------------------
@@ -497,13 +499,13 @@ def substep_pre_coupling_grad(self, f):
         self.compute_F_tmp.grad(f)
 
     def substep_post_coupling(self, f):
-        is_success = self.g2p(
+        self.g2p(
             f,
             self.sim.coupler.rigid_solver.geoms_info,
             self.sim.coupler.rigid_solver.links_state,
             self.sim.coupler.rigid_solver._rigid_global_info,
         )
-        if not is_success:
+        if not self._is_state_valid(f):
             gs.raise_exception(
                 "NaN detected in MPM states. Try reducing the time step size or adjusting simulation parameters."
             )

tests/test_grad.py

@@ -326,11 +326,9 @@ def compute_loss(input_mass, input_jac, input_aref, input_efc_D, input_force):
         assert_allclose(dL_error, 0.0, atol=RTOL)
 
 
+@pytest.mark.required
 @pytest.mark.parametrize("backend", [gs.cpu, gs.gpu])
 def test_differentiable_push(precision, show_viewer):
-    # FIXME: Wait for fix to be merged in GsTaichi: https://github.com/Genesis-Embodied-AI/gstaichi/pull/225
-    if sys.platform == "darwin" and gs.backend != gs.cpu:
-        pytest.skip(reason="GsTaichi does not support AutoDiff on non-CPU backend on Mac OS for now.")
     if sys.platform == "linux" and gs.backend == gs.cpu and precision == "64":
         pytest.skip(reason="GsTaichi segfault when using AutoDiff on CPU backend on Linux for now.")