merge

Author: YilingQiao

genesis/engine/entities/drone_entity.py

@@ -75,15 +75,14 @@ def set_propellels_rpm(self, propellels_rpm):
             gs.raise_exception("`propellels_rpm` cannot be negative.")
         self._propellers_revs = (self._propellers_revs + propellels_rpm) % (60 / self.solver.dt)
 
-        self.solver._kernel_set_drone_rpm(
+        self.solver.set_drone_rpm(
             self._n_propellers,
             self._propellers_link_idxs,
             propellels_rpm,
             self._propellers_spin,
             self.KF,
             self.KM,
             self._model == "RACE",
-            self.solver.links_state,
         )
 
     def update_propeller_vgeoms(self):
@@ -93,7 +92,7 @@ def update_propeller_vgeoms(self):
         This method is a no-op if animation is disabled due to missing visual geometry.
         """
         if self._animate_propellers:
-            self.solver._update_drone_propeller_vgeoms(
+            self.solver.update_drone_propeller_vgeoms(
                 self._n_propellers, self._propellers_vgeom_idxs, self._propellers_revs, self._propellers_spin
             )
 

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -1222,7 +1222,7 @@ def inverse_kinematics_multilink(
         if envs_idx is None:
             envs_idx = torch.zeros(1, dtype=gs.tc_int, device=gs.device)
 
-        self._kernel_inverse_kinematics(
+        self._solver.rigid_entity_inverse_kinematics(
             links_idx,
             poss,
             quats,
@@ -1245,7 +1245,9 @@ def inverse_kinematics_multilink(
             max_step_size,
             respect_joint_limit,
             envs_idx,
+            self,
         )
+
         qpos = self._IK_qpos_best.to_torch(gs.device).transpose(1, 0)
         if self._solver.n_envs > 0:
             qpos = qpos[envs_idx]
@@ -1261,273 +1263,6 @@ def inverse_kinematics_multilink(
             return qpos, error_pose
         return qpos
 
-    @ti.kernel
-    def _kernel_inverse_kinematics(
-        self,
-        links_idx: ti.types.ndarray(),
-        poss: ti.types.ndarray(),
-        quats: ti.types.ndarray(),
-        n_links: ti.i32,
-        dofs_idx: ti.types.ndarray(),
-        n_dofs: ti.i32,
-        links_idx_by_dofs: ti.types.ndarray(),
-        n_links_by_dofs: ti.i32,
-        custom_init_qpos: ti.i32,
-        init_qpos: ti.types.ndarray(),
-        max_samples: ti.i32,
-        max_solver_iters: ti.i32,
-        damping: ti.f32,
-        pos_tol: ti.f32,
-        rot_tol: ti.f32,
-        pos_mask_: ti.types.ndarray(),
-        rot_mask_: ti.types.ndarray(),
-        link_pos_mask: ti.types.ndarray(),
-        link_rot_mask: ti.types.ndarray(),
-        max_step_size: ti.f32,
-        respect_joint_limit: ti.i32,
-        envs_idx: ti.types.ndarray(),
-    ):
-        # convert to ti Vector
-        pos_mask = ti.Vector([pos_mask_[0], pos_mask_[1], pos_mask_[2]], dt=gs.ti_float)
-        rot_mask = ti.Vector([rot_mask_[0], rot_mask_[1], rot_mask_[2]], dt=gs.ti_float)
-        n_error_dims = 6 * n_links
-
-        for i_b in envs_idx:
-            # save original qpos
-            for i_q in range(self.n_qs):
-                self._IK_qpos_orig[i_q, i_b] = self._solver.qpos[i_q + self._q_start, i_b]
-
-            if custom_init_qpos:
-                for i_q in range(self.n_qs):
-                    self._solver.qpos[i_q + self._q_start, i_b] = init_qpos[i_b, i_q]
-
-            for i_error in range(n_error_dims):
-                self._IK_err_pose_best[i_error, i_b] = 1e4
-
-            solved = False
-            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.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):
-                        i_l_ee = links_idx[i_ee]
-
-                        tgt_pos_i = ti.Vector([poss[i_ee, i_b, 0], poss[i_ee, i_b, 1], poss[i_ee, i_b, 2]])
-                        err_pos_i = tgt_pos_i - self._solver.links_state.pos[i_l_ee, i_b]
-                        for k in range(3):
-                            err_pos_i[k] *= pos_mask[k] * link_pos_mask[i_ee]
-                        if err_pos_i.norm() > pos_tol:
-                            solved = False
-
-                        tgt_quat_i = ti.Vector(
-                            [quats[i_ee, i_b, 0], quats[i_ee, i_b, 1], quats[i_ee, i_b, 2], quats[i_ee, i_b, 3]]
-                        )
-                        err_rot_i = gu.ti_quat_to_rotvec(
-                            gu.ti_transform_quat_by_quat(
-                                gu.ti_inv_quat(self._solver.links_state.quat[i_l_ee, i_b]), tgt_quat_i
-                            )
-                        )
-                        for k in range(3):
-                            err_rot_i[k] *= rot_mask[k] * link_rot_mask[i_ee]
-                        if err_rot_i.norm() > rot_tol:
-                            solved = False
-
-                        # put into multi-link error array
-                        for k in range(3):
-                            self._IK_err_pose[i_ee * 6 + k, i_b] = err_pos_i[k]
-                            self._IK_err_pose[i_ee * 6 + k + 3, i_b] = err_rot_i[k]
-
-                    if solved:
-                        break
-
-                    # compute multi-link jacobian
-                    for i_ee in range(n_links):
-                        # update jacobian for ee link
-                        i_l_ee = links_idx[i_ee]
-                        self._func_get_jacobian(
-                            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)
-                        for i_dof in range(n_dofs):
-                            for i_error in ti.static(range(6)):
-                                i_row = i_ee * 6 + i_error
-                                i_dof_ = dofs_idx[i_dof]
-                                self._IK_jacobian[i_row, i_dof, i_b] = self._jacobian[i_error, i_dof_, i_b]
-
-                    # compute dq = jac.T @ inverse(jac @ jac.T + diag) @ error (only for the effective n_dofs instead of self.n_dofs)
-                    lu.mat_transpose(self._IK_jacobian, self._IK_jacobian_T, n_error_dims, n_dofs, i_b)
-                    lu.mat_mul(
-                        self._IK_jacobian,
-                        self._IK_jacobian_T,
-                        self._IK_mat,
-                        n_error_dims,
-                        n_dofs,
-                        n_error_dims,
-                        i_b,
-                    )
-                    lu.mat_add_eye(self._IK_mat, damping**2, n_error_dims, i_b)
-                    lu.mat_inverse(self._IK_mat, self._IK_L, self._IK_U, self._IK_y, self._IK_inv, n_error_dims, i_b)
-                    lu.mat_mul_vec(self._IK_inv, self._IK_err_pose, self._IK_vec, n_error_dims, n_error_dims, i_b)
-
-                    for i in range(self.n_dofs):  # IK_delta_qpos = IK_jacobian_T @ IK_vec
-                        self._IK_delta_qpos[i, i_b] = 0
-                    for i in range(n_dofs):
-                        for j in range(n_error_dims):
-                            i_ = dofs_idx[
-                                i
-                            ]  # NOTE: IK_delta_qpos uses the original indexing instead of the effective n_dofs
-                            self._IK_delta_qpos[i_, i_b] += self._IK_jacobian_T[i, j, i_b] * self._IK_vec[j, i_b]
-
-                    for i in range(self.n_dofs):
-                        self._IK_delta_qpos[i, i_b] = ti.math.clamp(
-                            self._IK_delta_qpos[i, i_b], -max_step_size, max_step_size
-                        )
-
-                    # update q
-                    self._solver._func_integrate_dq_entity(
-                        self._IK_delta_qpos, self._idx_in_solver, i_b, respect_joint_limit
-                    )
-
-                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.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]
-
-                        tgt_pos_i = ti.Vector([poss[i_ee, i_b, 0], poss[i_ee, i_b, 1], poss[i_ee, i_b, 2]])
-                        err_pos_i = tgt_pos_i - self._solver.links_state.pos[i_l_ee, i_b]
-                        for k in range(3):
-                            err_pos_i[k] *= pos_mask[k] * link_pos_mask[i_ee]
-                        if err_pos_i.norm() > pos_tol:
-                            solved = False
-
-                        tgt_quat_i = ti.Vector(
-                            [quats[i_ee, i_b, 0], quats[i_ee, i_b, 1], quats[i_ee, i_b, 2], quats[i_ee, i_b, 3]]
-                        )
-                        err_rot_i = gu.ti_quat_to_rotvec(
-                            gu.ti_transform_quat_by_quat(
-                                gu.ti_inv_quat(self._solver.links_state.quat[i_l_ee, i_b]), tgt_quat_i
-                            )
-                        )
-                        for k in range(3):
-                            err_rot_i[k] *= rot_mask[k] * link_rot_mask[i_ee]
-                        if err_rot_i.norm() > rot_tol:
-                            solved = False
-
-                        # put into multi-link error array
-                        for k in range(3):
-                            self._IK_err_pose[i_ee * 6 + k, i_b] = err_pos_i[k]
-                            self._IK_err_pose[i_ee * 6 + k + 3, i_b] = err_rot_i[k]
-
-                if solved:
-                    for i_q in range(self.n_qs):
-                        self._IK_qpos_best[i_q, i_b] = self._solver.qpos[i_q + self._q_start, i_b]
-                    for i_error in range(n_error_dims):
-                        self._IK_err_pose_best[i_error, i_b] = self._IK_err_pose[i_error, i_b]
-                    break
-
-                else:
-                    # copy to _IK_qpos if this sample is better
-                    improved = True
-                    for i_ee in range(n_links):
-                        error_pos_i = ti.Vector([self._IK_err_pose[i_ee * 6 + i_error, i_b] for i_error in range(3)])
-                        error_rot_i = ti.Vector([self._IK_err_pose[i_ee * 6 + i_error, i_b] for i_error in range(3, 6)])
-                        error_pos_best = ti.Vector(
-                            [self._IK_err_pose_best[i_ee * 6 + i_error, i_b] for i_error in range(3)]
-                        )
-                        error_rot_best = ti.Vector(
-                            [self._IK_err_pose_best[i_ee * 6 + i_error, i_b] for i_error in range(3, 6)]
-                        )
-                        if error_pos_i.norm() > error_pos_best.norm() or error_rot_i.norm() > error_rot_best.norm():
-                            improved = False
-                            break
-
-                    if improved:
-                        for i_q in range(self.n_qs):
-                            self._IK_qpos_best[i_q, i_b] = self._solver.qpos[i_q + self._q_start, i_b]
-                        for i_error in range(n_error_dims):
-                            self._IK_err_pose_best[i_error, i_b] = self._IK_err_pose[i_error, i_b]
-
-                    # Resample init q
-                    if respect_joint_limit and i_sample < max_samples - 1:
-                        for _i_l in range(n_links_by_dofs):
-                            i_l = links_idx_by_dofs[_i_l]
-                            I_l = [i_l, i_b] if ti.static(self.solver._options.batch_links_info) else i_l
-
-                            for i_j in range(
-                                self._solver.links_info.joint_start[I_l], self._solver.links_info.joint_end[I_l]
-                            ):
-                                I_j = [i_j, i_b] if ti.static(self.solver._options.batch_joints_info) else i_j
-
-                                I_dof_start = (
-                                    [self._solver.joints_info.dof_start[I_j], i_b]
-                                    if ti.static(self.solver._options.batch_dofs_info)
-                                    else self._solver.joints_info.dof_start[I_j]
-                                )
-                                q_start = self._solver.joints_info.q_start[I_j]
-                                dof_limit = self._solver.dofs_info.limit[I_dof_start]
-
-                                if self._solver.joints_info.type[I_j] == gs.JOINT_TYPE.FREE:
-                                    pass
-
-                                elif (
-                                    self._solver.joints_info.type[I_j] == gs.JOINT_TYPE.REVOLUTE
-                                    or self._solver.joints_info.type[I_j] == gs.JOINT_TYPE.PRISMATIC
-                                ):
-                                    if ti.math.isinf(dof_limit[0]) or ti.math.isinf(dof_limit[1]):
-                                        pass
-                                    else:
-                                        self._solver.qpos[q_start, i_b] = dof_limit[0] + ti.random() * (
-                                            dof_limit[1] - dof_limit[0]
-                                        )
-                    else:
-                        pass  # When respect_joint_limit=False, we can simply continue from the last solution
-
-            # 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.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):
         """

genesis/engine/entities/rigid_entity/rigid_geom.py

@@ -388,6 +388,7 @@ def get_verts(self):
         """
         Get the vertices of the geom in world frame.
         """
+        self._solver.update_verts_for_geom(self._idx)
         if self.is_free:
             tensor = torch.empty(
                 self._solver._batch_shape((self.n_verts, 3), True), dtype=gs.tc_float, device=gs.device
@@ -402,17 +403,12 @@ def get_verts(self):
 
     @ti.kernel
     def _kernel_get_free_verts(self, tensor: ti.types.ndarray()):
-        for i_b in range(self._solver._B):
-            self._solver._func_update_verts_for_geom(self._idx, i_b)
-
         for i_v, j, i_b in ti.ndrange(self.n_verts, 3, self._solver._B):
             idx_vert = i_v + self._verts_state_start
             tensor[i_b, i_v, j] = self._solver.free_verts_state.pos[idx_vert, i_b][j]
 
     @ti.kernel
     def _kernel_get_fixed_verts(self, tensor: ti.types.ndarray()):
-        self._solver._func_update_verts_for_geom(self._idx, 0)
-
         for i_v, j in ti.ndrange(self.n_verts, 3):
             idx_vert = i_v + self._verts_state_start
             tensor[i_v, j] = self._solver.fixed_verts_state.pos[idx_vert][j]

genesis/engine/simulator.py

@@ -273,11 +273,7 @@ def step(self, in_backward=False):
                     self.save_ckpt()
 
         if self.rigid_solver.is_active():
-            self.rigid_solver._kernel_clear_external_force(
-                links_state=self.rigid_solver.links_state,
-                rigid_global_info=self.rigid_solver._rigid_global_info,
-                static_rigid_sim_config=self.rigid_solver._static_rigid_sim_config,
-            )
+            self.rigid_solver.clear_external_force()
 
     def _step_grad(self):
         for _ in range(self._substeps - 1, -1, -1):