[FEATURE] Environment-wise gravity (Genesis-Embodied-AI#1324)

Author: LeonLiu4

genesis/engine/coupler.py

@@ -200,7 +200,7 @@ def mpm_grid_op(self, f: ti.i32, t: ti.f32):
                 vel_mpm = (1 / self.mpm_solver.grid[f, I, i_b].mass) * self.mpm_solver.grid[f, I, i_b].vel_in
 
                 # gravity
-                vel_mpm += self.mpm_solver.substep_dt * self.mpm_solver._gravity[None]
+                vel_mpm += self.mpm_solver.substep_dt * self.mpm_solver._gravity[i_b]
 
                 pos = (I + self.mpm_solver.grid_offset) * self.mpm_solver.dx
                 mass_mpm = self.mpm_solver.grid[f, I, i_b].mass / self.mpm_solver._p_vol_scale

genesis/engine/simulator.py

@@ -395,6 +395,10 @@ def get_state(self):
 
         return state
 
+    def set_gravity(self, gravity, envs_idx=None):
+        for solver in self._solvers:
+            solver.set_gravity(gravity, envs_idx)
+
     # ------------------------------------------------------------------------------------
     # ----------------------------------- properties -------------------------------------
     # ------------------------------------------------------------------------------------

genesis/engine/solvers/base_solver.py

@@ -21,13 +21,8 @@ def __init__(self, scene: "Scene", sim: "Simulator", options):
         self._scene = scene
         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._init_gravity = getattr(options, "gravity", None)
+        self._gravity = None
         self._entities: list[Entity] = gs.List()
 
         # force fields
@@ -36,6 +31,26 @@ def __init__(self, scene: "Scene", sim: "Simulator", options):
     def _add_force_field(self, force_field):
         self._ffs.append(force_field)
 
+    def build(self):
+        self._B = self._sim._B
+        if self._init_gravity is not None:
+            g_np = np.asarray(self._init_gravity, dtype=gs.np_float)
+            g_np = np.repeat(g_np[None], self._B, axis=0)
+            self._gravity = ti.Vector.field(3, dtype=gs.ti_float, shape=self._B)
+            self._gravity.from_numpy(g_np)
+
+    @gs.assert_built
+    def set_gravity(self, gravity, envs_idx=None):
+        if self._gravity is None:
+            return
+        g = np.asarray(gravity, dtype=gs.np_float)
+        if envs_idx is None:
+            if g.ndim == 1:
+                g = np.repeat(g[None], self._B, axis=0)
+            self._gravity.from_numpy(g)
+        else:
+            self._gravity[envs_idx] = g
+
     def dump_ckpt_to_numpy(self) -> dict[str, np.ndarray]:
         arrays: dict[str, np.ndarray] = {}
 

genesis/engine/solvers/fem_solver.py

@@ -276,6 +276,7 @@ def reset_grad(self):
             entity.reset_grad()
 
     def build(self):
+        super().build()
         self.n_envs = self.sim.n_envs
         self._B = self.sim._B
 
@@ -389,7 +390,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 +417,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/mpm_solver.py

@@ -196,6 +196,7 @@ def reset_grad(self):
             entity.reset_grad()
 
     def build(self):
+        super().build()
         # particles and entities
         self._B = self._sim._B
 

genesis/engine/solvers/pbd_solver.py

@@ -208,6 +208,7 @@ def init_ckpt(self):
         self._ckpt = dict()
 
     def build(self):
+        super().build()
         self._B = self._sim._B
         self._n_particles = self.n_particles
         self._n_fluid_particles = self.n_fluid_particles
@@ -371,7 +372,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

@@ -156,6 +156,7 @@ def add_entity(self, idx, material, morph, surface, visualize_contact) -> Entity
         return entity
 
     def build(self):
+        super().build()
         self.n_envs = self.sim.n_envs
         self._B = self.sim._B
         self._para_level = self.sim._para_level
@@ -2929,9 +2930,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)
@@ -2968,7 +2967,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)
@@ -4617,7 +4616,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/sf_solver.py

@@ -41,6 +41,7 @@ def set_jets(self, jets):
         self.jets = jets
 
     def build(self):
+        super().build()
         if self.is_active():
             self.t = 0.0
             self.setup_fields()

genesis/engine/solvers/sph_solver.py

@@ -128,6 +128,7 @@ def reset_grad(self):
         pass
 
     def build(self):
+        super().build()
         self._B = self._sim._B
 
         # particles and entities
@@ -275,7 +276,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,

genesis/engine/solvers/tool_solver.py

@@ -32,6 +32,7 @@ def __init__(self, scene, sim, options):
         self.setup_boundary()
 
     def build(self):
+        super().build()
         for entity in self._entities:
             entity.build()