Fix support of Apple Metal.

Author: duburcqa

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -560,7 +560,9 @@ def _init_jac_and_IK(self):
             else:
                 q_limit_lower.append(joint.dofs_limit[:, 0])
                 q_limit_upper.append(joint.dofs_limit[:, 1])
-        self.q_limit = np.array([np.concatenate(q_limit_lower), np.concatenate(q_limit_upper)])
+        self.q_limit = np.stack(
+            (np.concatenate(q_limit_lower), np.concatenate(q_limit_upper)), axis=0, dtype=gs.np_float
+        )
 
         # for storing intermediate results
         self._IK_n_tgts = self._solver._options.IK_max_targets
@@ -1287,7 +1289,6 @@ def _kernel_inverse_kinematics(
         rot_mask = ti.Vector([rot_mask_[0], rot_mask_[1], rot_mask_[2]], dt=gs.ti_float)
         n_error_dims = 6 * n_links
 
-        ti.loop_config(serialize=self._solver._para_level < gs.PARA_LEVEL.ALL)
         for i_b in envs_idx:
             # save original qpos
             for i_q in range(self.n_qs):
@@ -1579,7 +1580,6 @@ def _kernel_forward_kinematics(
         links_idx: ti.types.ndarray(),
         envs_idx: ti.types.ndarray(),
     ):
-
         ti.loop_config(serialize=self._solver._para_level < gs.PARA_LEVEL.ALL)
         for i_q_, i_b_ in ti.ndrange(qs_idx.shape[0], envs_idx.shape[0]):
             # save original qpos
@@ -2195,7 +2195,7 @@ def _get_idx(self, idx_local, idx_local_max, idx_global_start=0, *, unsafe=False
                 (idx_local.stop if idx_local.stop is not None else idx_local_max) + idx_global_start,
                 idx_local.step or 1,
             )
-        elif isinstance(idx_local, int):
+        elif isinstance(idx_local, (int, np.integer)):
             idx_global = idx_local + idx_global_start
         elif isinstance(idx_local, (list, tuple)):
             try:

genesis/engine/solvers/fem_solver.py

@@ -295,7 +295,7 @@ def _init_surface_info(self):
         surface_vertices_mapping = np.full(self.n_vertices, -1, dtype=np.int32)
         surface_vertices_mapping[surface_vertices_np] = np.arange(len(surface_vertices_np))
         mass = igl.massmatrix(pos_np, surface_vertices_mapping[surface_triangles_np])
-        surface_vert_mass_np = mass.diagonal()
+        surface_vert_mass_np = mass.diagonal().astype(gs.np_float, copy=False)
         self.surface_vert_mass = ti.field(
             dtype=gs.ti_float,
             shape=(len(surface_vertices_np),),

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -4563,7 +4563,7 @@ def _sanitize_1D_io_variables(
                 inputs_idx.stop if inputs_idx.stop is not None else input_size,
                 inputs_idx.step or 1,
             )
-        elif isinstance(inputs_idx, int):
+        elif isinstance(inputs_idx, (int, np.integer)):
             inputs_idx = [inputs_idx]
 
         is_preallocated = tensor is not None
@@ -4648,7 +4648,7 @@ def _sanitize_2D_io_variables(
                 inputs_idx.stop if inputs_idx.stop is not None else input_size,
                 inputs_idx.step or 1,
             )
-        elif isinstance(inputs_idx, int):
+        elif isinstance(inputs_idx, (int, np.integer)):
             inputs_idx = [inputs_idx]
 
         is_preallocated = tensor is not None

genesis/engine/states/entities.py

@@ -12,7 +12,7 @@ def __init__(self, entity, s_global):
         self.s_global = s_global
 
         args = {
-            "dtype": float,
+            "dtype": gs.tc_float,
             "requires_grad": self.entity.scene.requires_grad,
             "scene": self.entity.scene,
         }
@@ -48,7 +48,7 @@ def __init__(self, entity, s_global):
         self._s_global = s_global
         base_shape = (self.entity.sim._B, self._entity.n_particles)
         args = {
-            "dtype": float,
+            "dtype": gs.tc_float,
             "requires_grad": self._entity.scene.requires_grad,
             "scene": self._entity.scene,
         }
@@ -115,7 +115,7 @@ def __init__(self, entity, s_global):
         self._s_global = s_global
         base_shape = (self.entity.sim._B, self._entity.n_particles)
         args = {
-            "dtype": float,
+            "dtype": gs.tc_float,
             "requires_grad": False,
             "scene": self._entity.scene,
         }
@@ -151,7 +151,7 @@ def __init__(self, entity, s_global):
         base_shape = (self.entity.sim._B, self._entity.n_vertices, 3)
 
         args = {
-            "dtype": float,
+            "dtype": gs.tc_float,
             "requires_grad": False,
             "scene": self.entity.scene,
         }

genesis/engine/states/solvers.py

@@ -51,7 +51,7 @@ class RigidSolverState:
     def __init__(self, scene):
         self.scene = scene
         args = {
-            "dtype": float,
+            "dtype": gs.tc_float,
             "requires_grad": scene.requires_grad,
             "scene": self.scene,
         }
@@ -84,7 +84,7 @@ class AvatarSolverState:
     def __init__(self, scene):
         self.scene = scene
         args = {
-            "dtype": float,
+            "dtype": gs.tc_float,
             "requires_grad": scene.requires_grad,
             "scene": self.scene,
         }
@@ -137,7 +137,7 @@ class MPMSolverState(RBC):
     def __init__(self, scene):
         self._scene = scene
         args = {
-            "dtype": float,
+            "dtype": gs.tc_float,
             "requires_grad": scene.requires_grad,
             "scene": self._scene,
         }
@@ -146,7 +146,7 @@ def __init__(self, scene):
         self._C = gs.zeros((scene.sim._B, scene.sim.mpm_solver.n_particles, 3, 3), **args)
         self._F = gs.zeros((scene.sim._B, scene.sim.mpm_solver.n_particles, 3, 3), **args)
         self._Jp = gs.zeros((scene.sim._B, scene.sim.mpm_solver.n_particles), **args)
-        args["dtype"] = bool
+        args["dtype"] = gs.tc_bool
         args["requires_grad"] = False
         self._active = gs.zeros((scene.sim._B, scene.sim.mpm_solver.n_particles), **args)
 
@@ -197,13 +197,13 @@ class SPHSolverState:
     def __init__(self, scene):
         self._scene = scene
         args = {
-            "dtype": float,
+            "dtype": gs.tc_float,
             "requires_grad": scene.requires_grad,
             "scene": self._scene,
         }
         self._pos = gs.zeros((scene.sim._B, scene.sim.sph_solver.n_particles, 3), **args)
         self._vel = gs.zeros((self._scene.sim._B, scene.sim.sph_solver.n_particles, 3), **args)
-        args["dtype"] = bool
+        args["dtype"] = gs.tc_bool
         args["requires_grad"] = False
         self._active = gs.zeros((self._scene.sim._B, scene.sim.sph_solver.n_particles), **args)
 
@@ -232,13 +232,13 @@ class PBDSolverState:
     def __init__(self, scene):
         self._scene = scene
         args = {
-            "dtype": float,
+            "dtype": flgs.tc_floatoat,
             "requires_grad": scene.requires_grad,
             "scene": self._scene,
         }
         self._pos = gs.zeros((scene.sim._B, scene.sim.pbd_solver.n_particles, 3), **args)
         self._vel = gs.zeros((self._scene.sim._B, scene.sim.pbd_solver.n_particles, 3), **args)
-        args["dtype"] = bool
+        args["dtype"] = gs.tc_bool
         args["requires_grad"] = False
         self._free = gs.zeros((self._scene.sim._B, scene.sim.pbd_solver.n_particles), **args)
 
@@ -269,7 +269,7 @@ def __init__(self, scene):
         }
         self._pos = gs.zeros((scene.sim._B, scene.sim.fem_solver.n_vertices, 3), **args)
         self._vel = gs.zeros((scene.sim._B, scene.sim.fem_solver.n_vertices, 3), **args)
-        args["dtype"] = bool
+        args["dtype"] = gs.tc_bool
         args["requires_grad"] = False
         self._active = gs.zeros((scene.sim._B, scene.sim.fem_solver.n_elements), **args)
 

genesis/utils/misc.py

@@ -467,7 +467,7 @@ def ti_field_to_torch(
             if mask is None or isinstance(mask, slice):
                 # Slices are always valid by default. Nothing to check.
                 is_out_of_bounds = False
-            elif isinstance(mask, int):
+            elif isinstance(mask, (int, np.integer)):
                 # Do not allow negative indexing for consistency with Taichi
                 is_out_of_bounds = not (0 <= mask < _field_shape[i])
             elif isinstance(mask, torch.Tensor):
@@ -487,12 +487,12 @@ def ti_field_to_torch(
     # Must convert masks to torch if not slice or int since torch will do it anyway.
     # Note that being contiguous is not required and does not affect performance.
     must_allocate = False
-    is_row_mask_tensor = not (row_mask is None or isinstance(row_mask, (slice, int)))
+    is_row_mask_tensor = not (row_mask is None or isinstance(row_mask, (slice, int, np.integer)))
     if is_row_mask_tensor:
         _row_mask = torch.as_tensor(row_mask, dtype=gs.tc_int, device=gs.device)
         must_allocate = _row_mask is not row_mask
         row_mask = _row_mask
-    is_col_mask_tensor = not (col_mask is None or isinstance(col_mask, (slice, int)))
+    is_col_mask_tensor = not (col_mask is None or isinstance(col_mask, (slice, int, np.integer)))
     if is_col_mask_tensor:
         _col_mask = torch.as_tensor(col_mask, dtype=gs.tc_int, device=gs.device)
         must_allocate = _col_mask is not col_mask
@@ -532,8 +532,8 @@ def ti_field_to_torch(
     # Extract slice if necessary.
     # Note that unsqueeze is MUCH faster than indexing with `[row_mask]` to keep batch dimensions,
     # because this required allocating GPU data.
-    is_single_col = (is_col_mask_tensor and col_mask.ndim == 0) or isinstance(col_mask, int)
-    is_single_row = (is_row_mask_tensor and row_mask.ndim == 0) or isinstance(row_mask, int)
+    is_single_col = (is_col_mask_tensor and col_mask.ndim == 0) or isinstance(col_mask, (int, np.integer))
+    is_single_row = (is_row_mask_tensor and row_mask.ndim == 0) or isinstance(row_mask, (int, np.integer))
     try:
         if is_col_mask_tensor and is_row_mask_tensor:
             if not is_single_col and not is_single_row: