[BUG FIX] Fix gravity setters. (#1498)

Author: Milotrince

genesis/engine/simulator.py

@@ -403,7 +403,8 @@ def get_state(self):
 
     def set_gravity(self, gravity, envs_idx=None):
         for solver in self._solvers:
-            solver.set_gravity(gravity, envs_idx)
+            if solver.is_active():
+                solver.set_gravity(gravity, envs_idx)
 
     # ------------------------------------------------------------------------------------
     # ----------------------------------- properties -------------------------------------

genesis/engine/solvers/base_solver.py

@@ -42,14 +42,26 @@ def build(self):
     @gs.assert_built
     def set_gravity(self, gravity, envs_idx=None):
         if self._gravity is None:
+            gs.logger.debug("Gravity is not defined, skipping `set_gravity`.")
             return
         g = np.asarray(gravity, dtype=gs.np_float)
         if envs_idx is None:
             if g.ndim == 1:
                 g = np.tile(g, (self._B, 1))
+            assert g.shape == (self._B, 3), "Input gravity array should match (n_envs, 3)"
             self._gravity.from_numpy(g)
         else:
-            self._gravity[envs_idx] = g
+            envs_idx = np.atleast_1d(np.array(envs_idx, dtype=gs.np_int))
+            if g.ndim == 1:
+                g = np.tile(g, (len(envs_idx), 1))
+            assert g.shape == (len(envs_idx), 3), "Input gravity array should match (len(envs_idx), 3)"
+            self._kernel_set_gravity(g, envs_idx)
+
+    @ti.kernel
+    def _kernel_set_gravity(self, gravity: ti.types.ndarray(), envs_idx: ti.types.ndarray()):
+        for i_b_ in range(envs_idx.shape[0]):
+            for j in ti.static(range(3)):
+                self._gravity[envs_idx[i_b_]][j] = gravity[i_b_, j]
 
     def dump_ckpt_to_numpy(self) -> dict[str, np.ndarray]:
         arrays: dict[str, np.ndarray] = {}

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -2194,16 +2194,9 @@ def update_vgeoms(self):
 
     @gs.assert_built
     def set_gravity(self, gravity, envs_idx=None):
-        if not hasattr(self, "_rigid_global_info"):
-            super().set_gravity(gravity, envs_idx)
-            return
-        g = np.asarray(gravity, dtype=gs.np_float)
-        if envs_idx is None:
-            if g.ndim == 1:
-                g = np.tile(g, (self._B, 1))
-            self._rigid_global_info.gravity.from_numpy(g)
-        else:
-            self._rigid_global_info.gravity[envs_idx] = g
+        super().set_gravity(gravity, envs_idx)
+        if hasattr(self, "_rigid_global_info"):
+            self._rigid_global_info.gravity.copy_from(self._gravity)
 
     def rigid_entity_inverse_kinematics(
         self,

tests/test_rigid_physics.py

@@ -2177,18 +2177,29 @@ def test_gravity(show_viewer, tol):
     )
 
     sphere = scene.add_entity(gs.morphs.Sphere())
-    scene.build(n_envs=2)
+    scene.build(n_envs=3)
 
-    scene.sim.set_gravity(torch.tensor([0.0, 0.0, -9.8]), envs_idx=0)
-    scene.sim.set_gravity(torch.tensor([0.0, 0.0, 9.8]), envs_idx=1)
+    scene.sim.set_gravity(torch.tensor([0.0, 0.0, 0.0]))
+    scene.sim.set_gravity(torch.tensor([[1.0, 0.0, 0.0], [0.0, 2.0, 0.0]]), envs_idx=[0, 1])
+    scene.sim.set_gravity(torch.tensor([0.0, 0.0, 3.0]), envs_idx=2)
 
-    for _ in range(200):
-        scene.step()
+    with np.testing.assert_raises(AssertionError):
+        scene.sim.set_gravity(torch.tensor([0.0, -10.0]))
+
+    with np.testing.assert_raises(AssertionError):
+        scene.sim.set_gravity(torch.tensor([[0.0, 0.0, -10.0], [0.0, 0.0, -10.0]]), envs_idx=1)
 
-    first_pos = sphere.get_dofs_position()[0, 2]
-    second_pos = sphere.get_dofs_position()[1, 2]
+    scene.step()
 
-    assert_allclose(first_pos * -1, second_pos, tol=tol)
+    assert_allclose(
+        [
+            [1.0, 0.0, 0.0],
+            [0.0, 2.0, 0.0],
+            [0.0, 0.0, 3.0],
+        ],
+        sphere.get_links_acc().squeeze(),
+        tol=tol,
+    )
 
 
 @pytest.mark.required