[Question] Inverse Kinematics on GPU

[hesic73]

In _initialize_episode, I'm trying to initialize the end effector to be positioned directly above the object. My code looks like this:

from mani_skill.envs.sapien_env import BaseEnv
from mani_skill.agents.controllers.pd_ee_pose import PDEEPoseController
from mani_skill.utils.structs import Pose

import torch
from torch import Tensor

import numpy as np

class MyEnv(BaseEnv):

    def _initialize_episode(self, env_idx: Tensor, options: dict):

        arm_ee_delta_pose_controller: PDEEPoseController = self.agent.controller.controllers[
            'arm']
        arm_target_pos = self.obj.pose.p + \
            torch.tensor([0.0, 0.0, 0.2], device=self.device)  # (n, 3)
        arm_target_quat = rpy_to_quaternion(
            torch.tensor([0.0, np.pi/2, 0.0], device=self.device)).repeat(len(env_idx), 1)  # (n, 4)
        arm_target_pose = Pose.create_from_pq(
            p=arm_target_pos, q=arm_target_quat)

        # to root frame
        root_transform = arm_ee_delta_pose_controller.articulation.pose.inv()
        arm_target_pose = root_transform*arm_target_pose

        # gpu ik uses action and q0 instead of target_pose to compute the target_qpos
        # arm_delta_pos = arm_target_pos-self.end_effector_link.pose.p
        # arm_abs_rpy = torch.tensor(
        #     [0.0, np.pi/2, 0.0], device=self.device).repeat(len(env_idx), 1)
        # arm_action = torch.cat(
        #     [arm_delta_pos, arm_abs_rpy], dim=-1)  # (n, 6)

        arm_target_qpos = arm_ee_delta_pose_controller.kinematics.compute_ik(
            target_pose=arm_target_pose,
            q0=arm_ee_delta_pose_controller.articulation.get_qpos(),
            action=arm_action,
        )  # (n, 6)
        if arm_target_qpos is None:
            raise ValueError("IK solution not found")

        palm_target_qpos = torch.zeros(
            (len(env_idx), 16), device=self.device)
        init_qpos = torch.cat(
            [arm_target_qpos, palm_target_qpos], dim=-1)
        self.agent.reset(init_qpos=init_qpos)

def rpy_to_quaternion(euler: torch.Tensor) -> torch.Tensor:
    r, p, y = torch.unbind(euler, dim=-1)
    cy = torch.cos(y * 0.5)
    sy = torch.sin(y * 0.5)
    cp = torch.cos(p * 0.5)
    sp = torch.sin(p * 0.5)
    cr = torch.cos(r * 0.5)
    sr = torch.sin(r * 0.5)

    qw = cr * cp * cy + sr * sp * sy
    qx = sr * cp * cy - cr * sp * sy
    qy = cr * sp * cy + sr * cp * sy
    qz = cr * cp * sy - sr * sp * cy

    quaternion = torch.stack([qw, qx, qy, qz], dim=-1)

    return quaternion

The issue is that this approach only works on the CPU. https://github.com/haosulab/ManiSkill/blob/c1d11f942db4eb12313db8d6ff8e6de171e5b1d2/mani_skill/agents/controllers/utils/kinematics.py#L116

On GPU, Kinematics.compute_ik uses q0 and action to compute the target qpos. What’s the best approach to calculate action in this context?

[hesic73]

I attempts to use pytorch_kinematics.PseudoInverseIK to calculate the target qpos without using the action. And it seems to work. https://github.com/hesic73/ManiSkill/blob/e0ff7a5f092105ca3982086341ee7a5e3669525b/mani_skill/agents/controllers/utils/kinematics.py#L122

[hesic73]

I find that the previous code doesn't handle the indexing.

from mani_skill.envs.sapien_env import BaseEnv
from mani_skill.agents.controllers.pd_ee_pose import PDEEPoseController
from mani_skill.utils.structs import Pose

import torch
from torch import Tensor

import numpy as np

class MyEnv(BaseEnv):

    def _initialize_episode(self, env_idx: Tensor, options: dict):
        arm_ee_delta_pose_controller: PDEEPoseController = self.agent.controller.controllers[
            'arm']

        arm_target_pos = self.obj.pose.p[env_idx] + \
            torch.tensor([0.0, 0.0, 0.2], device=self.device)  # (n, 3)
        arm_target_quat = rpy_to_quaternion(
            torch.tensor([0.0, np.pi/2, 0.0], device=self.device)).repeat(len(env_idx), 1)  # (n, 4)
        arm_target_pose = Pose.create_from_pq(
            p=arm_target_pos, q=arm_target_quat)

        # to root frame
        raw_root_pose = arm_ee_delta_pose_controller.articulation.pose.raw_pose[env_idx]
        root_transform = Pose.create(raw_root_pose).inv()
        arm_target_pose = root_transform*arm_target_pose

        q0 = arm_ee_delta_pose_controller.articulation.get_qpos()[env_idx]

        arm_target_qpos = arm_ee_delta_pose_controller.kinematics.compute_ik(
            target_pose=arm_target_pose,
            q0=q0,
        )  # (n, 6)
        if arm_target_qpos is None:
            raise ValueError("IK solution not found")

        palm_target_qpos = torch.zeros(
            (len(env_idx), 16), device=self.device)
        init_qpos = torch.cat(
            [arm_target_qpos, palm_target_qpos], dim=-1)
        self.agent.reset(init_qpos=init_qpos)

def rpy_to_quaternion(euler: torch.Tensor) -> torch.Tensor:
    r, p, y = torch.unbind(euler, dim=-1)
    cy = torch.cos(y * 0.5)
    sy = torch.sin(y * 0.5)
    cp = torch.cos(p * 0.5)
    sp = torch.sin(p * 0.5)
    cr = torch.cos(r * 0.5)
    sr = torch.sin(r * 0.5)

    qw = cr * cp * cy + sr * sp * sy
    qx = sr * cp * cy - cr * sp * sy
    qy = cr * sp * cy + sr * cp * sy
    qz = cr * cp * sy - sr * sp * cy

    quaternion = torch.stack([qw, qx, qy, qz], dim=-1)

    return quaternion

[hesic73]

After some experiments, I find that pytorch_kinematics.PseudoInverseIK is significantly slower, approximately 4 to 5 times, and the controller is unstable. However, it is capable of solving the target qpos using only the end effector's target pose.

[StoneT2000]

oh this solves #488, that one is probably a batched IK solver (not delta action via inverse jacobians). It should be slower (it's effectively doing motion planning algorithm i think)