HiroIshida/tinyfk

An example of scikit robot wrapper

HiroIshida opened this issue · 0 comments

import tinyfk
import skrobot
from skrobot.coordinates.math import rpy_angle
@dataclass(frozen=True)
class KinematicsSolver:  # a wrapper of tinyfk
    # assume using python 3.6 or higher. and use cpython.
    kinsol: tinyfk.RobotModel
    control_joint_names: List[str]
    endeffector_link_names: List[str]
    joint_name_id_table: Dict[str, int]
    link_name_id_table: Dict[str, int]

    @classmethod
    def construct(cls, robot_model: skrobot.model.RobotModel, control_joint_names: List[str], endeffector_names: List[str]) -> 'KinematicsSolver':
        kinsol = tinyfk.RobotModel(robot_model.urdf_path)

        joint_names = [j.name for j in robot_model.joint_list]
        joint_ids = kinsol.get_joint_ids(joint_names)
        joint_name_id_table = {name: idd for name, idd in zip(joint_names, joint_ids)}

        link_names = [j.name for j in robot_model.link_list]
        link_ids = kinsol.get_link_ids(link_names)
        link_name_id_table = {name: idd for name, idd in zip(link_names, link_ids)}
        return cls(kinsol, control_joint_names, endeffector_names, joint_name_id_table, link_name_id_table)

    def get_control_joint_ids(self) -> List[int]:
        return self.kinsol.get_joint_ids(self.control_joint_names)

    def get_endeffector_link_ids(self) -> List[int]:
        return self.kinsol.get_link_ids(self.endeffector_link_names)

    def solve_ik(self, coords: Coordinates, robot_model: skrobot.model.RobotModel) -> Optional[np.ndarray]:

        if robot_model is not None:
            av = robot_model.angle_vector()
            assert len(av) == len(self.joint_name_id_table)
            self.kinsol.set_joint_angles(list(self.joint_name_id_table.values()), av)

        init_angle = np.array([robot_model.__dict__[jname].joint_angle() for jname in self.control_joint_names])

        xyz = coords.worldpos()
        ypr = np.flip(rpy_angle(coords.worldrot()))[1]
        #from IPython import embed; embed()
        xyzrpy = np.hstack([xyz, ypr])
        elink_id = self.get_endeffector_link_ids()[0]
        joint_ids = self.get_control_joint_ids()
        angles = self.kinsol.solve_inverse_kinematics(xyzrpy, init_angle, elink_id, joint_ids)

        for joint_name, angle in zip(self.control_joint_names, angles):
            robot_model.__dict__[joint_name].joint_angle(angle)
        return robot_model.angle_vector()