Merge remote-tracking branch 'origin/main' into pr/25/702_box_raycast

# Conflicts:
#	genesis/ext/pyrender/interaction/viewer_interaction.py

Author: gasnica

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):
@@ -200,7 +217,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
@@ -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()
@@ -395,6 +396,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/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/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)