Env wise gravity

Author: LeonLiu4

genesis/engine/simulator.py

@@ -105,7 +105,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)
@@ -184,9 +185,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)
@@ -411,7 +418,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

@@ -17,14 +17,11 @@ def __init__(self, scene: "Scene", sim: "Simulator", options):
         self._uid = gs.UID()
         self._sim = sim
         self._scene = scene
-        self._dt: float = options.dt
-        self._substep_dt: float = options.dt / sim.substeps
+        self._dt = options.dt
+        self._substep_dt = 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/rigid/rigid_solver_decomp.py

@@ -2844,9 +2844,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)
@@ -2883,7 +2881,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)
@@ -4532,7 +4530,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)