Merge branch 'main' into feature/jacobian

Author: YilingQiao

README.md

@@ -18,6 +18,7 @@
 # Genesis
 
 ## 🔥 News
+- [2025-07-02] The development of Genesis is now officially supported by [Genesis AI](https://genesis-ai.company/).
 - [2025-01-09] We released a [detailed performance benchmarking and comparison report](https://github.com/zhouxian/genesis-speed-benchmark) on Genesis, together with all the test scripts.
 - [2025-01-08] Released v0.2.1 🎊 🎉
 - [2025-01-08] Created [Discord](https://discord.gg/nukCuhB47p) and [Wechat](https://drive.google.com/uc?export=view&id=1ZS9nnbQ-t1IwkzJlENBYqYIIOOZhXuBZ) group.

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()

genesis/engine/solvers/rigid/collider_decomp.py

@@ -2072,7 +2072,7 @@ def get_contacts(self, as_tensor: bool = True, to_torch: bool = True):
         # Early return if already pre-computed
         contacts_info = self._contacts_info_cache.get((as_tensor, to_torch))
         if contacts_info is not None:
-            return contacts_info
+            return contacts_info.copy()
 
         # Find out how much dynamic memory must be allocated
         n_contacts = tuple(self.n_contacts.to_numpy())
@@ -2140,7 +2140,7 @@ def get_contacts(self, as_tensor: bool = True, to_torch: bool = True):
         # Cache contact information before returning
         self._contacts_info_cache[(as_tensor, to_torch)] = contacts_info
 
-        return contacts_info
+        return contacts_info.copy()
 
     @ti.kernel
     def _kernel_get_contacts(self, is_padded: ti.template(), iout: ti.types.ndarray(), fout: ti.types.ndarray()):

genesis/ext/pyrender/interaction/viewer_interaction.py

@@ -3,11 +3,9 @@
 
 import numpy as np
 
-from pyglet.event import EVENT_HANDLE_STATE
-
 from .ray import Plane, Ray, RayHit
 from .vec3 import Vec3
-from .viewer_interaction_base import ViewerInteractionBase
+from .viewer_interaction_base import ViewerInteractionBase, EVENT_HANDLE_STATE
 
 if TYPE_CHECKING:
     from genesis.engine.scene import Scene

genesis/ext/pyrender/interaction/viewer_interaction_base.py

@@ -1,7 +1,10 @@
-from pyglet.event import EVENT_HANDLE_STATE
+from typing import Union, Literal
 
 import genesis as gs
 
+
+EVENT_HANDLE_STATE = Union[Literal[True], None]
+
 # Note: Viewer window is based on pyglet.window.Window, mouse events are defined in pyglet.window.BaseWindow
 
 class ViewerInteractionBase():

genesis/ext/pyrender/viewer.py

@@ -41,7 +41,6 @@
 import pyglet
 from moviepy.video.io.ffmpeg_writer import FFMPEG_VideoWriter
 from pyglet import clock
-from pyglet.event import EVENT_HANDLE_STATE
 
 from .camera import IntrinsicsCamera, OrthographicCamera, PerspectiveCamera
 from .constants import (
@@ -57,7 +56,7 @@
     TextAlign,
 )
 from .interaction.viewer_interaction import ViewerInteraction
-from .interaction.viewer_interaction_base import ViewerInteractionBase
+from .interaction.viewer_interaction_base import ViewerInteractionBase, EVENT_HANDLE_STATE
 from .light import DirectionalLight
 from .node import Node
 from .renderer import Renderer