Merge branch 'main' into feature/get-weld

Author: LeonLiu4

genesis/__init__.py

@@ -118,6 +118,15 @@ def init(
     np_int = np.int32
     tc_int = torch.int32
 
+    # Bool
+    # Note that `ti.u1` is broken on Apple Metal and output garbage.
+    global ti_bool
+    global np_bool
+    global tc_bool
+    ti_bool = ti.u1 if backend != gs_backend.metal else ti.i32
+    np_bool = np.bool_
+    tc_bool = torch.bool
+
     # let's use GLSL convention: https://learnwebgl.brown37.net/12_shader_language/glsl_data_types.html
     global ti_vec2
     ti_vec2 = ti.types.vector(2, ti_float)

genesis/engine/bvh.py

@@ -125,7 +125,7 @@ def __init__(self, aabb: AABB, max_n_query_result_per_aabb: int = 8):
         self.aabb_min = ti.field(gs.ti_vec3, shape=(self.n_batches))
         self.aabb_max = ti.field(gs.ti_vec3, shape=(self.n_batches))
         self.scale = ti.field(gs.ti_vec3, shape=(self.n_batches))
-        self.morton_codes = ti.field(ti.u64, shape=(self.n_batches, self.n_aabbs))
+        self.morton_codes = ti.field(ti.types.vector(2, ti.u32), shape=(self.n_batches, self.n_aabbs))
 
         # Histogram for radix sort
         self.hist = ti.field(ti.u32, shape=(self.n_batches, 256))
@@ -134,7 +134,7 @@ def __init__(self, aabb: AABB, max_n_query_result_per_aabb: int = 8):
         # Offset for radix sort
         self.offset = ti.field(ti.u32, shape=(self.n_batches, self.n_aabbs))
         # Temporary storage for radix sort
-        self.tmp_morton_codes = ti.field(ti.u64, shape=(self.n_batches, self.n_aabbs))
+        self.tmp_morton_codes = ti.field(ti.types.vector(2, ti.u32), shape=(self.n_batches, self.n_aabbs))
 
         @ti.dataclass
         class Node:
@@ -158,8 +158,8 @@ class Node:
         # Nodes of the BVH, first n_aabbs - 1 are internal nodes, last n_aabbs are leaf nodes
         self.nodes = self.Node.field(shape=(self.n_batches, self.n_aabbs * 2 - 1))
         # Whether an internal node has been visited during traversal
-        self.internal_node_active = ti.field(ti.u1, shape=(self.n_batches, self.n_aabbs - 1))
-        self.internal_node_ready = ti.field(ti.u1, shape=(self.n_batches, self.n_aabbs - 1))
+        self.internal_node_active = ti.field(gs.ti_bool, shape=(self.n_batches, self.n_aabbs - 1))
+        self.internal_node_ready = ti.field(gs.ti_bool, shape=(self.n_batches, self.n_aabbs - 1))
 
         # Query results, vec3 of batch id, self id, query id
         self.query_result = ti.field(gs.ti_ivec3, shape=(self.max_n_query_results))
@@ -213,7 +213,7 @@ def compute_morton_codes(self):
             morton_code_y = self.expand_bits(morton_code_y)
             morton_code_z = self.expand_bits(morton_code_z)
             morton_code = (morton_code_x << 2) | (morton_code_y << 1) | (morton_code_z)
-            self.morton_codes[i_b, i_a] = (ti.u64(morton_code) << 32) | ti.u64(i_a)
+            self.morton_codes[i_b, i_a] = ti.Vector([morton_code, i_a], dt=ti.u32)
 
     @ti.func
     def expand_bits(self, v):
@@ -243,7 +243,7 @@ def _kernel_radix_sort_morton_codes_one_round(self, i: int):
             # This is now sequential
             # TODO Parallelize, need to use groups to handle data to remain stable, could be not worth it
             for i_a in range(self.n_aabbs):
-                code = (self.morton_codes[i_b, i_a] >> (i * 8)) & 0xFF
+                code = (self.morton_codes[i_b, i_a][1 - (i // 4)] >> ((i % 4) * 8)) & 0xFF
                 self.offset[i_b, i_a] = ti.atomic_add(self.hist[i_b, ti.i32(code)], 1)
 
         # Compute prefix sum
@@ -254,7 +254,7 @@ def _kernel_radix_sort_morton_codes_one_round(self, i: int):
 
         # Reorder morton codes
         for i_b, i_a in ti.ndrange(self.n_batches, self.n_aabbs):
-            code = (self.morton_codes[i_b, i_a] >> (i * 8)) & 0xFF
+            code = (self.morton_codes[i_b, i_a][1 - (i // 4)] >> ((i % 4) * 8)) & 0xFF
             idx = ti.i32(self.offset[i_b, i_a] + self.prefix_sum[i_b, ti.i32(code)])
             self.tmp_morton_codes[i_b, idx] = self.morton_codes[i_b, i_a]
 
@@ -318,10 +318,13 @@ def delta(self, i, j, i_b):
         result = -1
         if j >= 0 and j < self.n_aabbs:
             result = 64
-            x = self.morton_codes[i_b, ti.i32(i)] ^ self.morton_codes[i_b, ti.i32(j)]
-            for b in range(64):
-                if x & (ti.u64(1) << (63 - b)):
-                    result = b
+            for i_bit in range(2):
+                x = self.morton_codes[i_b, ti.i32(i)][i_bit] ^ self.morton_codes[i_b, ti.i32(j)][i_bit]
+                for b in range(32):
+                    if x & (1 << (31 - b)):
+                        result = b + 32 * i_bit
+                        break
+                if result != 64:
                     break
         return result
 
@@ -342,15 +345,15 @@ def _kernel_compute_bounds_init(self):
         self.internal_node_ready.fill(False)
 
         for i_b, i in ti.ndrange(self.n_batches, self.n_aabbs):
-            idx = ti.i32(self.morton_codes[i_b, i])
+            idx = ti.i32(self.morton_codes[i_b, i][1])
             self.nodes[i_b, i + self.n_aabbs - 1].bound.min = self.aabbs[i_b, idx].min
             self.nodes[i_b, i + self.n_aabbs - 1].bound.max = self.aabbs[i_b, idx].max
             parent_idx = self.nodes[i_b, i + self.n_aabbs - 1].parent
             if parent_idx != -1:
                 self.internal_node_active[i_b, parent_idx] = True
 
     @ti.kernel
-    def _kernel_compute_bounds_one_layer(self) -> ti.u1:
+    def _kernel_compute_bounds_one_layer(self) -> ti.i32:
         for i_b, i in ti.ndrange(self.n_batches, self.n_aabbs - 1):
             if self.internal_node_active[i_b, i]:
                 left_bound = self.nodes[i_b, self.nodes[i_b, i].left].bound
@@ -395,10 +398,8 @@ def query(self, aabbs: ti.template()):
                     if node.left == -1 and node.right == -1:
                         idx = ti.atomic_add(self.query_result_count[None], 1)
                         if idx < self.max_n_query_results:
-                            code = self.morton_codes[i_b, node_idx - (self.n_aabbs - 1)]
-                            self.query_result[idx] = gs.ti_ivec3(
-                                i_b, ti.i32(code & ti.u64(0xFFFFFFFF)), i_q
-                            )  # Store the AABB index
+                            code = self.morton_codes[i_b, node_idx - (self.n_aabbs - 1)][1]
+                            self.query_result[idx] = gs.ti_ivec3(i_b, ti.i32(code), i_q)  # Store the AABB index
                     else:
                         # Push children onto the stack
                         if node.right != -1:

genesis/engine/coupler.py

@@ -694,7 +694,7 @@ def init_fem_fields(self):
         self.fem_pressure.from_numpy(fem_pressure_np)
         self.fem_pressure_gradient = ti.field(gs.ti_vec3, shape=(fem_solver._B, fem_solver.n_elements))
         self.fem_floor_contact_pair_type = ti.types.struct(
-            active=ti.u1,  # whether the contact pair is active
+            active=gs.ti_bool,  # whether the contact pair is active
             batch_idx=gs.ti_int,  # batch index
             geom_idx=gs.ti_int,  # index of the FEM element
             intersection_code=gs.ti_int,  # intersection code for the element
@@ -747,7 +747,7 @@ def init_fem_fields(self):
 
     def init_sap_fields(self):
         self.batch_active = ti.field(
-            dtype=ti.u1,
+            dtype=gs.ti_bool,
             shape=self.sim._B,
             needs_grad=False,
         )
@@ -757,7 +757,7 @@ def init_sap_fields(self):
 
     def init_pcg_fields(self):
         self.batch_pcg_active = ti.field(
-            dtype=ti.u1,
+            dtype=gs.ti_bool,
             shape=self.sim._B,
             needs_grad=False,
         )
@@ -796,7 +796,7 @@ def init_pcg_fields(self):
 
     def init_linesearch_fields(self):
         self.batch_linesearch_active = ti.field(
-            dtype=ti.u1,
+            dtype=gs.ti_bool,
             shape=self.sim._B,
             needs_grad=False,
         )

genesis/engine/entities/pbd_entity.py

@@ -95,7 +95,7 @@ def _kernel_add_particles_edges_to_solver(
         edges: ti.types.ndarray(),
         edges_len_rest: ti.types.ndarray(),
         mat_type: ti.i32,
-        active: ti.u1,
+        active: ti.i32,
     ):
         for i_p_ in range(self.n_particles):
             i_p = i_p_ + self._particle_start
@@ -587,7 +587,7 @@ def _kernel_add_particles_to_solver(
         particles: ti.types.ndarray(),
         rho: ti.float32,
         mat_type: ti.i32,
-        active: ti.u1,
+        active: ti.i32,
     ):
         for i_p_ in range(self._n_particles):
             i_p = i_p_ + self._particle_start

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -407,6 +407,9 @@ def _load_scene(self, morph, surface):
             j_info["dofs_stiffness"] = np.zeros(6)
             j_info["dofs_invweight"] = np.zeros(6)
             j_info["dofs_damping"] = np.zeros(6)
+            if isinstance(morph, gs.morphs.Drone):
+                mass_tot = sum(l_info["inertial_mass"] for l_info in l_infos)
+                j_info["dofs_damping"][3:] = mass_tot * morph.default_base_ang_damping_scale
             j_info["dofs_armature"] = np.zeros(6)
             j_info["dofs_kp"] = np.zeros((6,), dtype=gs.np_float)
             j_info["dofs_kv"] = np.zeros((6,), dtype=gs.np_float)
@@ -1113,13 +1116,13 @@ def inverse_kinematics_multilink(
             gs.raise_exception("Target link not provided.")
 
         if len(poss) == 0:
-            poss = [None] * n_links
+            poss = [None for _ in range(n_links)]
             pos_mask = [False, False, False]
         elif len(poss) != n_links:
             gs.raise_exception("Accepting only `poss` with length equal to `links` or empty list.")
 
         if len(quats) == 0:
-            quats = [None] * n_links
+            quats = [None for _ in range(n_links)]
             rot_mask = [False, False, False]
         elif len(quats) != n_links:
             gs.raise_exception("Accepting only `quats` with length equal to `links` or empty list.")
@@ -1300,8 +1303,19 @@ def _kernel_inverse_kinematics(
             for i_sample in range(max_samples):
                 for _ in range(max_solver_iters):
                     # run FK to update link states using current q
-                    self._solver._func_forward_kinematics_entity(self._idx_in_solver, i_b)
-
+                    self._solver._func_forward_kinematics_entity(
+                        self._idx_in_solver,
+                        i_b,
+                        self._solver.links_state,
+                        self._solver.links_info,
+                        self._solver.joints_state,
+                        self._solver.joints_info,
+                        self._solver.dofs_state,
+                        self._solver.dofs_info,
+                        self._solver.entities_info,
+                        self._solver._rigid_global_info,
+                        self._solver._static_rigid_sim_config,
+                    )
                     # compute error
                     solved = True
                     for i_ee in range(n_links):
@@ -1386,7 +1400,19 @@ def _kernel_inverse_kinematics(
 
                 if not solved:
                     # re-compute final error if exited not due to solved
-                    self._solver._func_forward_kinematics_entity(self._idx_in_solver, i_b)
+                    self._solver._func_forward_kinematics_entity(
+                        self._idx_in_solver,
+                        i_b,
+                        self._solver.links_state,
+                        self._solver.links_info,
+                        self._solver.joints_state,
+                        self._solver.joints_info,
+                        self._solver.dofs_state,
+                        self._solver.dofs_info,
+                        self._solver.entities_info,
+                        self._solver._rigid_global_info,
+                        self._solver._static_rigid_sim_config,
+                    )
                     solved = True
                     for i_ee in range(n_links):
                         i_l_ee = links_idx[i_ee]
@@ -1480,7 +1506,19 @@ def _kernel_inverse_kinematics(
             # restore original qpos and link state
             for i_q in range(self.n_qs):
                 self._solver.qpos[i_q + self._q_start, i_b] = self._IK_qpos_orig[i_q, i_b]
-            self._solver._func_forward_kinematics_entity(self._idx_in_solver, i_b)
+            self._solver._func_forward_kinematics_entity(
+                self._idx_in_solver,
+                i_b,
+                self._solver.links_state,
+                self._solver.links_info,
+                self._solver.joints_state,
+                self._solver.joints_info,
+                self._solver.dofs_state,
+                self._solver.dofs_info,
+                self._solver.entities_info,
+                self._solver._rigid_global_info,
+                self._solver._static_rigid_sim_config,
+            )
 
     @gs.assert_built
     def forward_kinematics(self, qpos, qs_idx_local=None, links_idx_local=None, envs_idx=None):
@@ -1549,7 +1587,19 @@ def _kernel_forward_kinematics(
             # set new qpos
             self._solver.qpos[qs_idx[i_q_], envs_idx[i_b_]] = qpos[i_b_, i_q_]
             # run FK
-            self._solver._func_forward_kinematics_entity(self._idx_in_solver, envs_idx[i_b_])
+            self._solver._func_forward_kinematics_entity(
+                self._idx_in_solver,
+                envs_idx[i_b_],
+                self._solver.links_state,
+                self._solver.links_info,
+                self._solver.joints_state,
+                self._solver.joints_info,
+                self._solver.dofs_state,
+                self._solver.dofs_info,
+                self._solver.entities_info,
+                self._solver._rigid_global_info,
+                self._solver._static_rigid_sim_config,
+            )
 
         ti.loop_config(serialize=self._solver._para_level < gs.PARA_LEVEL.PARTIAL)
         for i_l_, i_b_ in ti.ndrange(links_idx.shape[0], envs_idx.shape[0]):
@@ -1563,7 +1613,19 @@ def _kernel_forward_kinematics(
             # restore original qpos
             self._solver.qpos[qs_idx[i_q_], envs_idx[i_b_]] = self._IK_qpos_orig[qs_idx[i_q_], envs_idx[i_b_]]
             # run FK
-            self._solver._func_forward_kinematics_entity(self._idx_in_solver, envs_idx[i_b_])
+            self._solver._func_forward_kinematics_entity(
+                self._idx_in_solver,
+                envs_idx[i_b_],
+                self._solver.links_state,
+                self._solver.links_info,
+                self._solver.joints_state,
+                self._solver.joints_info,
+                self._solver.dofs_state,
+                self._solver.dofs_info,
+                self._solver.entities_info,
+                self._solver._rigid_global_info,
+                self._solver._static_rigid_sim_config,
+            )
 
     # ------------------------------------------------------------------------------------
     # --------------------------------- motion planing -----------------------------------

genesis/engine/entities/rigid_entity/rigid_joint.py

@@ -67,13 +67,6 @@ def __init__(
         self._dofs_kv = dofs_kv
         self._dofs_force_range = dofs_force_range
 
-        # NOTE: temp hack to use 0 damping/armature for drone
-        if isinstance(self._entity, gs.engine.entities.DroneEntity) and self._type == gs.JOINT_TYPE.FREE:
-            import numpy as np
-
-            self._dofs_damping = np.zeros_like(self._dofs_damping)
-            self._dofs_armature = np.zeros_like(self._dofs_armature)
-
     # ------------------------------------------------------------------------------------
     # -------------------------------- real-time state -----------------------------------
     # ------------------------------------------------------------------------------------

genesis/engine/force_fields.py

@@ -16,7 +16,7 @@ class ForceField(RBC):
     """
 
     def __init__(self):
-        self._active = ti.field(ti.u1, shape=())
+        self._active = ti.field(gs.ti_bool, shape=())
         self._active[None] = False
 
     def activate(self):

genesis/engine/solvers/avatar_solver.py

@@ -58,8 +58,30 @@ def _kernel_step(self):
     @ti.kernel
     def _kernel_forward_kinematics_links_geoms(self, envs_idx: ti.types.ndarray()):
         for i_b in envs_idx:
-            self._func_forward_kinematics(i_b)
-            self._func_update_geoms(i_b)
+            self._func_forward_kinematics(
+                i_b,
+                self.links_state,
+                self.links_info,
+                self.joints_state,
+                self.joints_info,
+                self.dofs_state,
+                self.dofs_info,
+                self.entities_info,
+                self._rigid_global_info,
+                self._static_rigid_sim_config,
+            )
+            self._func_update_geoms(
+                i_b,
+                self.links_state,
+                self.links_info,
+                self.joints_state,
+                self.joints_info,
+                self.dofs_state,
+                self.dofs_info,
+                self.entities_info,
+                self._rigid_global_info,
+                self._static_rigid_sim_config,
+            )
 
     @ti.func
     def _func_detect_collision(self):