fixed minor error

Author: SonSang

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -1879,6 +1879,7 @@ def get_verts(self):
             The vertices of the entity (using collision geoms).
         """
 
+        self._update_verts_for_geom()
         if self.is_free:
             tensor = torch.empty(
                 self._solver._batch_shape((self.n_verts, 3), True), dtype=gs.tc_float, device=gs.device
@@ -1891,22 +1892,20 @@ def get_verts(self):
             self._kernel_get_fixed_verts(tensor)
         return tensor
 
-    @ti.kernel
-    def _kernel_get_free_verts(self, tensor: ti.types.ndarray()):
-        for i_g_, i_b in ti.ndrange(self.n_geoms, self._solver._B):
+    @gs.assert_built
+    def _update_verts_for_geom(self):
+        for i_g_ in range(self.n_geoms):
             i_g = i_g_ + self._geom_start
-            self._solver._func_update_verts_for_geom(i_g, i_b)
+            self._solver.update_verts_for_geom(i_g)
 
+    @ti.kernel
+    def _kernel_get_free_verts(self, tensor: ti.types.ndarray()):
         for i, j, b in ti.ndrange(self.n_verts, 3, self._solver._B):
             idx_vert = i + self._verts_state_start
             tensor[b, i, j] = self._solver.free_verts_state.pos[idx_vert, b][j]
 
     @ti.kernel
     def _kernel_get_fixed_verts(self, tensor: ti.types.ndarray()):
-        for i_g_ in range(self.n_geoms):
-            i_g = i_g_ + self._geom_start
-            self._solver._func_update_verts_for_geom(i_g, 0)
-
         for i, j in ti.ndrange(self.n_verts, 3):
             idx_vert = i + self._verts_state_start
             tensor[i, j] = self._solver.fixed_verts_state.pos[idx_vert][j]

genesis/engine/entities/rigid_entity/rigid_link.py

@@ -294,6 +294,7 @@ def get_verts(self):
         """
         Get the vertices of the link's collision body (concatenation of all `link.geoms`) in the world frame.
         """
+        self._update_verts_for_geom()
         if self.is_free:
             tensor = torch.empty(
                 self._solver._batch_shape((self.n_verts, 3), True), dtype=gs.tc_float, device=gs.device
@@ -306,22 +307,20 @@ def get_verts(self):
             self._kernel_get_fixed_verts(tensor)
         return tensor
 
-    @ti.kernel
-    def _kernel_get_free_verts(self, tensor: ti.types.ndarray()):
-        for i_g_, i_b in ti.ndrange(self.n_geoms, self._solver._B):
+    @gs.assert_built
+    def _update_verts_for_geom(self):
+        for i_g_ in range(self.n_geoms):
             i_g = i_g_ + self._geom_start
-            self._solver._func_update_verts_for_geom(i_g, i_b)
+            self._solver.update_verts_for_geom(i_g)
 
+    @ti.kernel
+    def _kernel_get_free_verts(self, tensor: ti.types.ndarray()):
         for i, j, b in ti.ndrange(self.n_verts, 3, self._solver._B):
             idx_vert = i + self._verts_state_start
             tensor[b, i, j] = self._solver.free_verts_state.pos[idx_vert, b][j]
 
     @ti.kernel
     def _kernel_get_fixed_verts(self, tensor: ti.types.ndarray()):
-        for i_g_ in range(self.n_geoms):
-            i_g = i_g_ + self._geom_start
-            self._solver._func_update_verts_for_geom(i_g, 0)
-
         for i, j in ti.ndrange(self.n_verts, 3):
             idx_vert = i + self._verts_state_start
             tensor[i, j] = self._solver.fixed_verts_state.pos[idx_vert][j]

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -4778,21 +4778,34 @@ def kernel_update_verts_for_geom(
 ):
     _B = geoms_state.verts_updated.shape[1]
     for i_b in range(_B):
-        if not geoms_state.verts_updated[i_g, i_b]:
-            if geoms_info.is_free[i_g]:
-                for i_v in range(geoms_info.vert_start[i_g], geoms_info.vert_end[i_g]):
-                    verts_state_idx = verts_info.verts_state_idx[i_v]
-                    free_verts_state.pos[verts_state_idx, i_b] = gu.ti_transform_by_trans_quat(
-                        verts_info.init_pos[i_v], geoms_state.pos[i_g, i_b], geoms_state.quat[i_g, i_b]
-                    )
-                geoms_state.verts_updated[i_g, i_b] = 1
-            elif i_b == 0:
-                for i_v in range(geoms_info.vert_start[i_g], geoms_info.vert_end[i_g]):
-                    verts_state_idx = verts_info.verts_state_idx[i_v]
-                    fixed_verts_state.pos[verts_state_idx] = gu.ti_transform_by_trans_quat(
-                        verts_info.init_pos[i_v], geoms_state.pos[i_g, i_b], geoms_state.quat[i_g, i_b]
-                    )
-                geoms_state.verts_updated[i_g, 0] = 1
+        func_update_verts_for_geom(i_g, i_b, geoms_state, geoms_info, verts_info, free_verts_state, fixed_verts_state)
+
+
+@ti.func
+def func_update_verts_for_geom(
+    i_g: ti.i32,
+    i_b: ti.i32,
+    geoms_state: array_class.GeomsState,
+    geoms_info: array_class.GeomsInfo,
+    verts_info: array_class.VertsInfo,
+    free_verts_state: array_class.FreeVertsState,
+    fixed_verts_state: array_class.FixedVertsState,
+):
+    if not geoms_state.verts_updated[i_g, i_b]:
+        if geoms_info.is_free[i_g]:
+            for i_v in range(geoms_info.vert_start[i_g], geoms_info.vert_end[i_g]):
+                verts_state_idx = verts_info.verts_state_idx[i_v]
+                free_verts_state.pos[verts_state_idx, i_b] = gu.ti_transform_by_trans_quat(
+                    verts_info.init_pos[i_v], geoms_state.pos[i_g, i_b], geoms_state.quat[i_g, i_b]
+                )
+            geoms_state.verts_updated[i_g, i_b] = 1
+        elif i_b == 0:
+            for i_v in range(geoms_info.vert_start[i_g], geoms_info.vert_end[i_g]):
+                verts_state_idx = verts_info.verts_state_idx[i_v]
+                fixed_verts_state.pos[verts_state_idx] = gu.ti_transform_by_trans_quat(
+                    verts_info.init_pos[i_v], geoms_state.pos[i_g, i_b], geoms_state.quat[i_g, i_b]
+                )
+            geoms_state.verts_updated[i_g, 0] = 1
 
 
 @ti.func