Repository: Genesis-Embodied-AI/Genesis
File: genesis/engine/entities/rigid_entity/rigid_entity.py (Python)

Code before:

        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
        -------
        jacobian : torch.Tensor
            The Jacobian matrix of shape (n_envs, 6, entity.n_dofs) or (6, entity.n_dofs) if n_envs == 0.
        """
        if not self._requires_jac_and_IK:
            gs.raise_exception(
                "Inverse kinematics and jacobian are disabled for this entity. Set `morph.requires_jac_and_IK` to True if you need them."
            )

        if self.n_dofs == 0:
            gs.raise_exception("Entity has zero dofs.")

        if local_point is None:
            self._kernel_get_jacobian_zero(link.idx)
        else:
            p_local = torch.as_tensor(local_point, dtype=torch.float64, device="cpu").numpy()
            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, 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._impl_get_jacobian(tgt_link_idx, i_b, p_vec)

Review comment (suggestion):
`p_local = torch.as_tensor(local_point, dtype=gs.tc_float, device=gs.device)`

Code after:

        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
        -------
        jacobian : torch.Tensor
            The Jacobian matrix of shape (n_envs, 6, entity.n_dofs) or (6, entity.n_dofs) if n_envs == 0.
        """
        if not self._requires_jac_and_IK:
            gs.raise_exception(
                "Inverse kinematics and jacobian are disabled for this entity. Set `morph.requires_jac_and_IK` to True if you need them."
            )

        if self.n_dofs == 0:
            gs.raise_exception("Entity has zero dofs.")

        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, 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._impl_get_jacobian(tgt_link_idx, i_b, p_vec)