solver-specific gravity and fixed unit test

Author: LeonLiu4

genesis/engine/solvers/base_solver.py

@@ -19,6 +19,7 @@ def __init__(self, scene: "Scene", sim: "Simulator", options):
         self._scene = scene
         self._dt: float = options.dt
         self._substep_dt: float = options.dt / sim.substeps
+        self._init_gravity = getattr(options, "gravity", None)
         self._gravity = None
         self._entities: list[Entity] = gs.List()
 
@@ -28,11 +29,15 @@ def __init__(self, scene: "Scene", sim: "Simulator", options):
     def _add_force_field(self, force_field):
         self._ffs.append(force_field)
 
-    def build(self, B: int):
-        self._B = B
-        g_np = np.asarray(self._sim._gravity)
-        g_np = np.repeat(g_np[None], B, axis=0)
-        self._gravity = ti.Vector.field(3, dtype=gs.ti_float, shape=B)
+    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)
+        else:
+            g_np = np.asarray(self._sim._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)
 
     def set_gravity(self, gravity, envs_idx=None):

genesis/engine/solvers/fem_solver.py

@@ -276,9 +276,9 @@ def reset_grad(self):
             entity.reset_grad()
 
     def build(self):
+        super().build()
         self.n_envs = self.sim.n_envs
         self._B = self.sim._B
-        super().build(self._B)
 
         # batch fields
         self.init_batch_fields()

genesis/engine/solvers/mpm_solver.py

@@ -196,9 +196,10 @@ def reset_grad(self):
             entity.reset_grad()
 
     def build(self):
+        super().build()
+
         # particles and entities
         self._B = self._sim._B
-        super().build(self._B)
 
         self._n_particles = self.n_particles
         self._n_vverts = self.n_vverts

genesis/engine/solvers/pbd_solver.py

@@ -208,8 +208,8 @@ def init_ckpt(self):
         self._ckpt = dict()
 
     def build(self):
+        super().build()
         self._B = self._sim._B
-        super().build(self._B)
         self._n_particles = self.n_particles
         self._n_fluid_particles = self.n_fluid_particles
         self._n_edges = self.n_edges

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -156,9 +156,10 @@ 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
-        super().build(self._B)
         self._para_level = self.sim._para_level
 
         for entity in self._entities:

genesis/engine/solvers/sf_solver.py

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

genesis/engine/solvers/sph_solver.py

@@ -128,8 +128,9 @@ def reset_grad(self):
         pass
 
     def build(self):
+        super().build()
+
         self._B = self._sim._B
-        super().build(self._B)
 
         # particles and entities
         self._n_particles = self.n_particles

genesis/engine/solvers/tool_solver.py

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

tests/test_rigid_physics.py

@@ -2205,11 +2205,6 @@ def test_urdf_mimic(show_viewer, tol):
 def test_gravity(show_viewer, tol):
     scene = gs.Scene(
         show_viewer=show_viewer,
-        sim_options=gs.options.SimOptions(
-            dt=0.01,
-            substeps=1,
-            gravity=(0.0, 0.0, -9.8),
-        ),
     )
 
     sphere = scene.add_entity(gs.morphs.Sphere())