Merge remote-tracking branch 'origin/main' into HEAD

Author: Kashu7100

.github/workflows/generic-cpu.yml .github/workflows/generic.yml.github/workflows/generic-cpu.yml renamed to .github/workflows/generic.yml

@@ -1,4 +1,4 @@
-name: Generic - CPU only
+name: Generic
 
 on:
   pull_request:

.github/workflows/linux-gpu.yml .github/workflows/production.yml.github/workflows/linux-gpu.yml renamed to .github/workflows/production.yml

@@ -1,4 +1,4 @@
-name: Linux x86 - Nvidia GPU
+name: Production
 
 on:
   # Trigger the workflow on push on the master branch, or for any pull request

.pre-commit-config.yaml

@@ -1,8 +1,5 @@
 repos:
   - repo: https://github.com/psf/black
-    rev: 25.1.0  # Use the latest version or specify the version you prefer
+    rev: 25.1.0
     hooks:
       - id: black
-        args:
-          - --line-length=120
-          - .

examples/rigid/suction_cup.py

@@ -96,8 +96,8 @@ def main():
 
     # add suction / weld constraint
     rigid = scene.sim.rigid_solver
-    link_cube = np.array([cube.get_link("box_baselink").idx], dtype=gs.np_int)
-    link_franka = np.array([franka.get_link("hand").idx], dtype=gs.np_int)
+    link_cube = cube.get_link("box_baselink").idx
+    link_franka = franka.get_link("hand").idx
     rigid.add_weld_constraint(link_cube, link_franka)
 
     # lift

genesis/engine/entities/emitter.py

@@ -150,7 +150,7 @@ def emit(
             positions = gu.transform_by_trans_R(
                 positions,
                 pos,
-                gu.z_up_to_R(direction) @ gu.axis_angle_to_R(np.array([0, 0, 1]), theta),
+                gu.z_up_to_R(direction) @ gu.axis_angle_to_R(np.array([0.0, 0.0, 1.0], dtype=gs.np_float), theta),
             ).astype(gs.np_float)
 
             positions = np.tile(positions[np.newaxis], (self._sim._B, 1, 1))
@@ -206,17 +206,21 @@ def emit_omni(self, source_radius=0.1, pos=(0.5, 0.5, 1.0), speed=1.0, particle_
         Parameters:
         ----------
         source_radius: float, optional
-            The radius of the sphere source. Particles will be emitted from a shell with inner radius using 0.8 * source_radius and outer radius using source_radius.
+            The radius of the sphere source. Particles will be emitted from a shell with inner radius using
+            '0.8 * source_radius' and outer radius using source_radius.
         pos: array_like, shape=(3,)
             The center of the sphere source.
         speed: float
             The speed of the emitted particles.
         particle_size: float | None
-            The size (diameter) of the emitted particles. The actual number of particles emitted is determined by the volume of the sphere source and the size of the particles. If None, the solver's particle size is used. Note that this particle size only affects computation for number of particles emitted, not the actual size of the particles in simulation and rendering.
+            The size (diameter) of the emitted particles. The actual number of particles emitted is determined by the
+            volume of the sphere source and the size of the particles. If None, the solver's particle size is used.
+            Note that this particle size only affects computation for number of particles emitted, not the actual size
+            of the particles in simulation and rendering.
         """
         assert self._entity is not None
 
-        pos = np.array(pos)
+        pos = np.asarray(pos, dtype=gs.np_float)
 
         if particle_size is None:
             particle_size = self._solver.particle_size
@@ -227,20 +231,20 @@ def emit_omni(self, source_radius=0.1, pos=(0.5, 0.5, 1.0), speed=1.0, particle_
             inner_radius=source_radius * 0.4,
             sampler=self._entity.sampler,
         )
-
-        positions = gu.transform_by_T(positions_, gu.trans_to_T(pos)).astype(gs.np_float)
+        positions = pos + positions_
 
         if not self._solver.boundary.is_inside(positions):
             gs.raise_exception("Emitted particles are outside the boundary.")
 
-        n_particles = len(positions)
-        dists = np.linalg.norm(positions_, axis=1, keepdims=True)
-        positions[np.where(dists < gs.EPS)[0]] = np.array([gs.EPS, gs.EPS, gs.EPS])
-        vels = (positions_ / dists * speed).astype(gs.np_float)
+        dists = np.linalg.norm(positions_, axis=1)
+        positions[dists < gs.EPS] = gs.EPS
+        vels = (speed / (dists + gs.EPS)) * positions_
 
+        n_particles = len(positions)
         if n_particles > self._entity.n_particles:
             gs.logger.warning(
-                f"Number of particles to emit ({n_particles}) at the current step is larger than the maximum number of particles ({self._entity.n_particles})."
+                f"Number of particles to emit ({n_particles}) at the current step is larger than the maximum number "
+                f"of particles ({self._entity.n_particles})."
             )
 
         self._solver._kernel_set_particles_pos(

genesis/engine/entities/mpm_entity.py

@@ -89,14 +89,9 @@ def set_pos(self, f, pos):
         f : int
             The current substep index.
         pos : gs.Tensor
-            A tensor of shape (n_particles, 3) representing particle positions.
+            A tensor of shape (n_envs, n_particles, 3) representing particle positions.
         """
-        self.solver._kernel_set_particles_pos(
-            f,
-            self._particle_start,
-            self._n_particles,
-            pos,
-        )
+        self.solver._kernel_set_particles_pos(f, self._particle_start, self._n_particles, pos)
 
     def set_pos_grad(self, f, pos_grad):
         """

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -796,14 +796,17 @@ def _add_equality(self, name, type, objs_name, data, sol_params):
     # ------------------------------------------------------------------------------------
 
     @gs.assert_built
-    def get_jacobian(self, link):
+    def get_jacobian(self, link, local_point=None):
         """
-        Get the Jacobian matrix for a target link.
+        Get the spatial Jacobian for a point on a target link.
 
         Parameters
         ----------
         link : RigidLink
             The target link.
+        local_point : torch.Tensor or None, shape (3,)
+            Coordinates of the point in the link’s *local* frame.
+            If None, the link origin is used (back-compat).
 
         Returns
         -------
@@ -818,31 +821,48 @@ def get_jacobian(self, link):
         if self.n_dofs == 0:
             gs.raise_exception("Entity has zero dofs.")
 
-        self._kernel_get_jacobian(link.idx)
+        if local_point is None:
+            self._kernel_get_jacobian_zero(link.idx)
+        else:
+            p_local = torch.as_tensor(local_point, dtype=gs.tc_float, device=gs.device)
+            if p_local.shape != (3,):
+                gs.raise_exception("Must be a vector of length 3")
+            self._kernel_get_jacobian(link.idx, p_local)
 
         jacobian = self._jacobian.to_torch(gs.device).permute(2, 0, 1)
         if self._solver.n_envs == 0:
             jacobian = jacobian.squeeze(0)
 
         return jacobian
 
+    @ti.func
+    def _impl_get_jacobian(self, tgt_link_idx, i_b, p_vec):
+        self._func_get_jacobian(
+            tgt_link_idx,
+            i_b,
+            p_vec,
+            ti.Vector.one(gs.ti_int, 3),
+            ti.Vector.one(gs.ti_int, 3),
+        )
+
     @ti.kernel
-    def _kernel_get_jacobian(self, tgt_link_idx: ti.i32):
-        ti.loop_config(serialize=self._solver._para_level < gs.PARA_LEVEL.ALL)
+    def _kernel_get_jacobian(self, tgt_link_idx: ti.i32, p_local: ti.types.ndarray()):
+        p_vec = ti.Vector([p_local[0], p_local[1], p_local[2]], dt=gs.ti_float)
         for i_b in range(self._solver._B):
-            self._func_get_jacobian(
-                tgt_link_idx,
-                i_b,
-                ti.Vector.one(gs.ti_int, 3),
-                ti.Vector.one(gs.ti_int, 3),
-            )
+            self._impl_get_jacobian(tgt_link_idx, i_b, p_vec)
+
+    @ti.kernel
+    def _kernel_get_jacobian_zero(self, tgt_link_idx: ti.i32):
+        for i_b in range(self._solver._B):
+            self._impl_get_jacobian(tgt_link_idx, i_b, ti.Vector.zero(gs.ti_float, 3))
 
     @ti.func
-    def _func_get_jacobian(self, tgt_link_idx, i_b, pos_mask, rot_mask):
+    def _func_get_jacobian(self, tgt_link_idx, i_b, p_local, pos_mask, rot_mask):
         for i_row, i_d in ti.ndrange(6, self.n_dofs):
             self._jacobian[i_row, i_d, i_b] = 0.0
 
-        tgt_link_pos = self._solver.links_state[tgt_link_idx, i_b].pos
+        tgt_link_state = self._solver.links_state[tgt_link_idx, i_b]
+        tgt_link_pos = tgt_link_state.pos + gu.ti_transform_by_quat(p_local, tgt_link_state.quat)
         i_l = tgt_link_idx
         while i_l > -1:
             I_l = [i_l, i_b] if ti.static(self.solver._options.batch_links_info) else i_l
@@ -1318,7 +1338,7 @@ def _kernel_inverse_kinematics(
                         # update jacobian for ee link
                         i_l_ee = links_idx[i_ee]
                         self._func_get_jacobian(
-                            i_l_ee, i_b, pos_mask, rot_mask
+                            i_l_ee, i_b, ti.Vector.zero(gs.ti_float, 3), pos_mask, rot_mask
                         )  # NOTE: we still compute jacobian for all dofs as we haven't found a clean way to implement this
 
                         # copy to multi-link jacobian (only for the effective n_dofs instead of self.n_dofs)

genesis/engine/entities/rigid_entity/rigid_link.py

@@ -55,24 +55,24 @@ def __init__(
         self._entity_idx_in_solver = entity.idx
 
         self._uid = gs.UID()
-        self._idx = idx
-        self._parent_idx = parent_idx
-        self._root_idx = root_idx
-        self._child_idxs = list()
-        self._invweight = invweight
-
-        self._joint_start = joint_start
-        self._n_joints = n_joints
-
-        self._geom_start = geom_start
-        self._cell_start = cell_start
-        self._vert_start = vert_start
-        self._face_start = face_start
-        self._edge_start = edge_start
-        self._verts_state_start = verts_state_start
-        self._vgeom_start = vgeom_start
-        self._vvert_start = vvert_start
-        self._vface_start = vface_start
+        self._idx: int = idx
+        self._parent_idx: int = parent_idx  # -1 if no parent
+        self._root_idx: int | None = root_idx  # None if no root
+        self._child_idxs: list[int] = list()
+        self._invweight: float | None = invweight
+
+        self._joint_start: int = joint_start
+        self._n_joints: int = n_joints
+
+        self._geom_start: int = geom_start
+        self._cell_start: int = cell_start
+        self._vert_start: int = vert_start
+        self._face_start: int = face_start
+        self._edge_start: int = edge_start
+        self._verts_state_start: int = verts_state_start
+        self._vgeom_start: int = vgeom_start
+        self._vvert_start: int = vvert_start
+        self._vface_start: int = vface_start
 
         # Link position & rotation at creation time:
         self._pos: ArrayLike = pos

genesis/engine/entities/sph_entity.py

@@ -92,12 +92,7 @@ def set_pos(self, f, pos):
         pos : ndarray
             Array of particle positions of shape (n_envs, n_particles, 3).
         """
-        self.solver._kernel_set_particles_pos(
-            f,
-            self._particle_start,
-            self._n_particles,
-            pos,
-        )
+        self.solver._kernel_set_particles_pos(f, self._particle_start, self._n_particles, pos)
 
     def set_pos_grad(self, f: ti.i32, pos_grad: ti.types.ndarray()):
         """

genesis/engine/solvers/mpm_solver.py

@@ -720,9 +720,9 @@ def _kernel_set_particles_pos(
                 self.particles[f, i_global, i_b].pos[k] = pos[i_b, i_p, k]
 
             # we restore these whenever directly setting positions
-            self.particles[f, i_global, i_b].vel = ti.Vector.zero(gs.ti_float, 3)
+            self.particles[f, i_global, i_b].vel.fill(0.0)
             self.particles[f, i_global, i_b].F = ti.Matrix.identity(gs.ti_float, 3)
-            self.particles[f, i_global, i_b].C = ti.Matrix.zero(gs.ti_float, 3, 3)
+            self.particles[f, i_global, i_b].C.fill(0.0)
             self.particles[f, i_global, i_b].Jp = self.particles_info[i_global].default_Jp
 
     @ti.kernel