[MISC] Improve reset (#1350)

Author: Kashu7100

genesis/engine/coupler.py

@@ -1,4 +1,5 @@
 from typing import TYPE_CHECKING
+
 import numpy as np
 import taichi as ti
 
@@ -83,14 +84,30 @@ def build(self) -> None:
         self._dx = 1 / 1024
         self._stencil_size = int(np.floor(self._dx / self.sph_solver.hash_grid_cell_size) + 2)
 
-        self.reset()
+        self.reset(envs_idx=self.sim.scene._envs_idx)
 
-    def reset(self) -> None:
+    def reset(self, envs_idx=None) -> None:
         if self._rigid_mpm and self.mpm_solver.enable_CPIC:
-            self.mpm_rigid_normal.fill(0)
+            if envs_idx is None:
+                self.mpm_rigid_normal.fill(0)
+            else:
+                self._kernel_reset_mpm(envs_idx)
 
         if self._rigid_sph:
-            self.sph_rigid_normal.fill(0)
+            if envs_idx is None:
+                self.sph_rigid_normal.fill(0)
+            else:
+                self._kernel_reset_sph(envs_idx)
+
+    @ti.kernel
+    def _kernel_reset_mpm(self, envs_idx: ti.types.ndarray()):
+        for i_p, i_g, i_b_ in ti.ndrange(self.mpm_solver.n_particles, self.rigid_solver.n_geoms, envs_idx.shape[0]):
+            self.mpm_rigid_normal[i_p, i_g, envs_idx[i_b_]] = 0.0
+
+    @ti.kernel
+    def _kernel_reset_sph(self, envs_idx: ti.types.ndarray()):
+        for i_p, i_g, i_b_ in ti.ndrange(self.sph_solver.n_particles, self.rigid_solver.n_geoms, envs_idx.shape[0]):
+            self.sph_rigid_normal[i_p, i_g, envs_idx[i_b_]] = 0.0
 
     @ti.func
     def _func_collide_with_rigid(self, f, pos_world, vel, mass, i_b):
@@ -667,7 +684,7 @@ def build(self) -> None:
         self.init_pcg_fields()
         self.init_linesearch_fields()
 
-    def reset(self):
+    def reset(self, envs_idx=None) -> None:
         pass
 
     def init_fem_fields(self):

genesis/engine/scene.py

@@ -12,6 +12,7 @@
 from genesis.engine.force_fields import ForceField
 from genesis.engine.materials.base import Material
 from genesis.engine.entities import Emitter
+from genesis.engine.states.solvers import SimState
 from genesis.engine.simulator import Simulator
 from genesis.options import (
     AvatarOptions,
@@ -695,23 +696,27 @@ def _parallelize(
             self._para_level = gs.PARA_LEVEL.ALL
 
     @gs.assert_built
-    def reset(self, state: dict | None = None, envs_idx=None):
+    def reset(self, state: SimState | None = None, envs_idx=None):
         """
         Resets the scene to its initial state.
 
         Parameters
         ----------
-        state : dict | None
-            The state to reset the scene to. If None, the scene will be reset to its initial state. If this is given, the scene's registerered initial state will be updated to this state.
+        state : SimState | None
+            The state to reset the scene to. If None, the scene will be reset to its initial state.
+            If this is given, the scene's registerered initial state will be updated to this state.
+        envs_idx : None | array_like, optional
+            The indices of the environments. If None, all environments will be considered. Defaults to None.
         """
-        gs.logger.info(f"Resetting Scene ~~~<{self._uid}>~~~.")
-        self._reset(state, envs_idx)
+        gs.logger.debug(f"Resetting Scene ~~~<{self._uid}>~~~.")
+        self._reset(state, envs_idx=envs_idx)
 
-    def _reset(self, state=None, envs_idx=None):
+    def _reset(self, state: SimState | None = None, *, envs_idx=None):
         if self._is_built:
             if state is None:
                 state = self._init_state
             else:
+                assert isinstance(state, SimState), "state must be a SimState object"
                 self._init_state = state
             self._sim.reset(state, envs_idx)
         else:

genesis/engine/simulator.py

@@ -22,6 +22,7 @@
 
 from .coupler import Coupler, SAPCoupler
 from .entities import HybridEntity
+from .solvers.base_solver import Solver
 from .solvers import (
     AvatarSolver,
     FEMSolver,
@@ -205,12 +206,12 @@ def build(self):
             if isinstance(entity, HybridEntity):
                 entity.build()
 
-    def reset(self, state, envs_idx=None):
+    def reset(self, state: SimState, envs_idx=None):
         for solver, solver_state in zip(self._solvers, state):
-            solver.set_state(0, solver_state, envs_idx)
+            if solver.n_entities > 0:
+                solver.set_state(0, solver_state, envs_idx)
 
-        # TODO: keeping as is for now, since coupler is currently for non-batched scenes
-        self.coupler.reset()
+        self.coupler.reset(envs_idx=envs_idx)
 
         # TODO: keeping as is for now
         self.reset_grad()