[FEATURE] Environment-wise gravity

genesis/engine/simulator.py

@@ -106,7 +106,8 @@ def __init__(
         self._steps_local = options._steps_local
 
         self._cur_substep_global = 0
-        self._gravity = np.array(options.gravity)
+        self._g_np = np.asarray(options.gravity, dtype=np.float32)
+        self._g_ti = None
 
         # solvers
         self.tool_solver = ToolSolver(self.scene, self, self.tool_options)
@@ -188,9 +189,15 @@ def build(self):
         self._B = self.scene._B
         self._para_level = self.scene._para_level
 
+        if self._g_np.ndim == 1:
+            self._g_np = np.repeat(self._g_np[None], self._B, axis=0)
+        self._g_ti = ti.Vector.field(3, gs.ti_float, shape=self._B)
+        self._g_ti.from_numpy(self._g_np)
+
         # solvers
         self._rigid_only = self.rigid_solver.is_active()
         for solver in self._solvers:
+            solver._finalize_batch(self._B)
             solver.build()
             if solver.is_active():
                 self._active_solvers.append(solver)
@@ -418,7 +425,7 @@ def scene(self):
     @property
     def gravity(self):
         """The gravity vector."""
-        return self._gravity
+        return self._g_ti
 
     @property
     def requires_grad(self):

genesis/engine/solvers/base_solver.py

@@ -20,11 +20,8 @@ def __init__(self, scene: "Scene", sim: "Simulator", options):
         self._dt: float = options.dt
         self._substep_dt: float = options.dt / sim.substeps
 
-        if hasattr(options, "gravity"):
-            self._gravity = ti.field(dtype=gs.ti_vec3, shape=())
-            self._gravity.from_numpy(np.array(options.gravity, dtype=gs.np_float))
-        else:
-            self._gravity = None
+        self._gravity_cfg = np.asarray(options.gravity, gs.np_float) if hasattr(options, "gravity") else None
+        self._gravity = None
 
         self._entities: list[Entity] = gs.List()
 
@@ -34,6 +31,15 @@ def __init__(self, scene: "Scene", sim: "Simulator", options):
     def _add_force_field(self, force_field):
         self._ffs.append(force_field)
 
+    def _finalize_batch(self, B: int):
+        if self._gravity_cfg is None or self._gravity is not None:
+            return
+        self._gravity = ti.field(dtype=gs.ti_vec3, shape=(B,))
+        g = self._gravity_cfg
+        if g.ndim == 1:
+            g = np.tile(g[None, :], (B, 1))
+        self._gravity.from_numpy(g[..., None])
+
     # ------------------------------------------------------------------------------------
     # ----------------------------------- properties -------------------------------------
     # ------------------------------------------------------------------------------------

genesis/engine/solvers/fem_solver.py

@@ -389,7 +389,7 @@ def apply_uniform_force(self, f: ti.i32):
             #       however, this inevitably damp the gravity.
             self.elements_v[f + 1, i_v, i_b].vel *= ti.exp(-dt * self.damping)
             # Add gravity (avoiding damping on gravity)
-            self.elements_v[f + 1, i_v, i_b].vel += dt * self._gravity[None]
+            self.elements_v[f + 1, i_v, i_b].vel += dt * self._gravity[i_b]
 
     @ti.kernel
     def compute_pos(self, f: ti.i32):
@@ -416,7 +416,7 @@ def init_pos_and_inertia(self, f: ti.i32):
         dt = self.substep_dt
         for i_v, i_b in ti.ndrange(self.n_vertices, self._B):
             self.elements_v_energy[i_b, i_v].inertia = (
-                self.elements_v[f, i_v, i_b].pos + self.elements_v[f, i_v, i_b].vel * dt + self._gravity[None] * dt**2
+                self.elements_v[f, i_v, i_b].pos + self.elements_v[f, i_v, i_b].vel * dt + self._gravity[i_b] * dt**2
             )
             self.elements_v[f + 1, i_v, i_b].pos = self.elements_v[f, i_v, i_b].pos
 

genesis/engine/solvers/pbd_solver.py

@@ -371,7 +371,7 @@ def _kernel_apply_external_force(self, f: ti.i32, t: ti.f32):
         for i_p, i_b in ti.ndrange(self._n_particles, self._B):
             if self.particles[i_p, i_b].free:
                 # gravity
-                self.particles[i_p, i_b].vel = self.particles[i_p, i_b].vel + self._gravity[None] * self._substep_dt
+                self.particles[i_p, i_b].vel = self.particles[i_p, i_b].vel + self._gravity[i_b] * self._substep_dt
 
                 # external force fields
                 acc = ti.Vector.zero(gs.ti_float, 3)

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -2932,9 +2932,7 @@ def _func_update_acc(self, update_cacc: ti.template()):
                         i_p = self.links_info[I_l].parent_idx
 
                         if i_p == -1:
-                            self.links_state[i_l, i_b].cdd_vel = -self._gravity[None] * (
-                                1 - e_info.gravity_compensation
-                            )
+                            self.links_state[i_l, i_b].cdd_vel = -self._gravity[i_b] * (1 - e_info.gravity_compensation)
                             self.links_state[i_l, i_b].cdd_ang = ti.Vector.zero(gs.ti_float, 3)
                             if ti.static(update_cacc):
                                 self.links_state[i_l, i_b].cacc_lin = ti.Vector.zero(gs.ti_float, 3)
@@ -2971,7 +2969,7 @@ def _func_update_acc(self, update_cacc: ti.template()):
                     i_p = self.links_info[I_l].parent_idx
 
                     if i_p == -1:
-                        self.links_state[i_l, i_b].cdd_vel = -self._gravity[None] * (1 - e_info.gravity_compensation)
+                        self.links_state[i_l, i_b].cdd_vel = -self._gravity[i_b] * (1 - e_info.gravity_compensation)
                         self.links_state[i_l, i_b].cdd_ang = ti.Vector.zero(gs.ti_float, 3)
                         if ti.static(update_cacc):
                             self.links_state[i_l, i_b].cacc_lin = ti.Vector.zero(gs.ti_float, 3)
@@ -4620,7 +4618,7 @@ def _kernel_get_links_acc(
             # Mimick IMU accelerometer signal if requested
             if mimick_imu:
                 # Subtract gravity
-                acc_classic_lin -= self._gravity[None]
+                acc_classic_lin -= self._gravity[i_b]
 
                 # Move the resulting linear acceleration in local links frame
                 acc_classic_lin = gu.ti_inv_transform_by_quat(acc_classic_lin, self.links_state[i_l, i_b].quat)

genesis/engine/solvers/sph_solver.py

@@ -275,7 +275,7 @@ def _task_compute_non_pressure_forces(self, i, j, ret: ti.template(), i_b: ti.i3
     def _kernel_compute_non_pressure_forces(self, f: ti.i32, t: ti.f32):
         for i_p, i_b in ti.ndrange(self._n_particles, self._B):
             if self.particles_ng_reordered[i_p, i_b].active:
-                acc = self._gravity[None]
+                acc = self._gravity[i_b]
                 self.sh.for_all_neighbors(
                     i_p,
                     self.particles_reordered.pos,

tests/test_rigid_physics.py

@@ -2200,6 +2200,36 @@ def test_urdf_mimic(show_viewer, tol):
     assert_allclose(gs_qpos[-1], gs_qpos[-2], tol=tol)
 
 
+@pytest.mark.required
+@pytest.mark.parametrize("backend", [gs.cpu])
+def test_gravity(show_viewer):
+    base_rpm = 14468.429183500699
+    scene = gs.Scene(
+        show_viewer=show_viewer,
+        sim_options=gs.options.SimOptions(
+            dt=0.01,
+            substeps=1,
+            gravity=[(0.0, 0.0, -9.8), (0.0, 0.0, -10.00)],
+        ),
+    )
+
+    drone = scene.add_entity(
+        gs.morphs.Drone(file="urdf/drones/cf2x.urdf", pos=(0, 0, 1.0)),
+    )
+
+    scene.build(n_envs=2)
+
+    for _ in range(500):
+        drone.set_propellels_rpm([[base_rpm, base_rpm, base_rpm, base_rpm], [base_rpm, base_rpm, base_rpm, base_rpm]])
+        scene.step()
+
+    first_pos = drone.get_dofs_position()[0, 2]
+    second_pos = drone.get_dofs_position()[1, 2]
+    assert_allclose(
+        second_pos, first_pos - 2.5, tol=scene.sim_options.dt
+    )  # Relax the tolerance due to time integration's error
+
+
 @pytest.mark.required
 @pytest.mark.parametrize("backend", [gs.cpu])
 def test_drone_hover_same_with_and_without_substeps(show_viewer, tol):