Utility interface class for inverse kinematics and dynamics computations using PINK library, for all type of robot embodiments.

pink_interface.py

"""Provides utility interface class for inverse kinematics and dynamics computations using PINK library.

If the functionalities of this class are required only intermittently and is not super performance-demanding,
the utility functions in ik_utils.py might be enough.

This is a standalone file.
pypi dependencies: [numpy, pin, pin-pink]
"""
from pink import Configuration, solve_ik, custom_configuration_vector
from pink.tasks import FrameTask, PostureTask, Task
from pinocchio import RobotWrapper, XYZQUATToSE3
import numpy as np
from typing import TypeAlias, Sequence, Mapping

QuatType: TypeAlias = np.ndarray
"""Numpy array representating quaternion in format [x,y,z,w]"""


def get_frame_task_name_for_frame(frame_name: str) -> str:
    return f"frame_task_{frame_name}"


def get_posture_task_name() -> str:
    return "posture_task"


class PinkInterface:
    """Simple interface to solve IK using pinocchio-based PINK library."""

    def __init__(self, robot: RobotWrapper, initial_joint_values: np.ndarray = None):
        """Simple interface to solve IK using pinocchio-based PINK library.

        Args:
            robot (RobotWrapper): The Pinocchio RobotWrapper object to be used.
            initial_joint_values (np.ndarray, optional): The iniiial joint positions to be used as initial
                configuration. Defaults to None (uses whatever is in robot.q0).
        """
        self._robot = robot

        if initial_joint_values is None:
            initial_joint_values = robot.q0

        self.configuration = Configuration(
            model=self._robot.model, data=self._robot.data, q=initial_joint_values
        )
        self.tasks: Mapping[str, Task] = {}
        """Dictionary of all PINK IK tasks, mapping: task name (str) --> pink.tasks.Task object."""

    def add_frame_task(
        self,
        frame_name: str,
        position_cost: float | Sequence[float] = 1.0,
        orientation_cost: float | Sequence[float] = 1.0,
        lm_damping: float = 0.0,
    ):
        """Add a pose task for the specified frame to the IK problem.

        NOTE: only one target can be set for

        Args:
            frame_name (str): The name of the frame (should be present in URDF/RobotWrapper object.)
            position_cost (float | Sequence[float], optional): Strictness associated to the position target
                (range: [0, 1]). Defaults to 1.0.
            orientation_cost (float | Sequence[float], optional): Strictness associated to the orientation target
                (range: [0, 1]). Defaults to 1.0.
            lm_damping (float, optional): The Levenberg-Marquardt damping cost for regularisation. Defaults to 0.0.
        """
        task_name = get_frame_task_name_for_frame(frame_name=frame_name)
        assert (
            task_name not in self.tasks.keys()
        ), f"A frame task for {frame_name} already exists!"

        self.tasks[task_name] = FrameTask(
            body=frame_name,
            position_cost=position_cost,
            orientation_cost=orientation_cost,
            lm_damping=lm_damping,
        )

    def update_frame_task(
        self,
        frame_name: str,
        target_position: np.ndarray,
        target_orientation: QuatType,
        position_cost: float | Sequence[float] = None,
        orientation_cost: float | Sequence[float] = None,
    ):
        """Update the values of a previously created Frame task.

        Args:
            frame_name (str): The name of the frame.
            target_position (np.ndarray): New target position for the frame.
            target_orientation (QuatType): New target orientation quaternion for the frame (x, y, z, w order).
            position_cost (float | Sequence[float], optional): Strictness associated to the position target
                (range: [0, 1]). Defaults to previously set value during task creation.
            orientation_cost (float | Sequence[float], optional): Strictness associated to the orientation target
                (range: [0, 1]). Defaults to previously set value during task creation.
        """
        self.tasks[get_frame_task_name_for_frame(frame_name=frame_name)].set_target(
            transform_target_to_world=XYZQUATToSE3(
                [*target_position, *target_orientation]
            )
        )
        if position_cost is not None:
            self.tasks[
                get_frame_task_name_for_frame(frame_name=frame_name)
            ].set_position_cost(position_cost=position_cost)
        if orientation_cost is not None:
            self.tasks[
                get_frame_task_name_for_frame(frame_name=frame_name)
            ].set_orientation_cost(orientation_cost=orientation_cost)

    def enable_posture_task(
        self, cost: float, lm_damping: float = 0.0, default_q: np.ndarray = None
    ):
        """Enable the posture task (joint targets) for solving IK.

        Args:
            cost (float): Strictness associated to the posture target
                (range: [0, 1]).
            lm_damping (float, optional): The Levenberg-Marquardt damping cost for regularisation. Defaults to 0.0.
            default_q (np.ndarray, optional): Target for posture task. Defaults to the current values in the
                configuration.
        """
        if get_posture_task_name() in self.tasks.keys():
            print("WARNING: posture task already defined. Redefining it!")

        self.tasks[get_posture_task_name()] = PostureTask(
            cost=cost, lm_damping=lm_damping
        )
        if default_q is not None:
            self.update_posture_task(target_q=default_q)
        else:
            self.update_posture_task(target_q=self.configuration.q)

    def update_posture_task(self, target_q: np.ndarray):
        """Update the joint targets for the posture task.

        NOTE: Only possible if posture task is enabled.

        Args:
            target_q (np.ndarray): Target joint positions.

        Raises:
            RuntimeError: Thrown if posture task is not enabled.
        """
        try:
            self.tasks[get_posture_task_name()].set_target(target_q=target_q)
        except KeyError:
            raise RuntimeError(
                "Posture task is not enabled! Enable the task by calling `enable_posture_task()`."
            )

    def update_posture_task_targets_for_joints(self, **kwargs):
        """Update the joint targets for specified joints for the posture task.

        The function takes in keyword arguments, where the key is the name of the
        joint in the pinocchio robot model, and the value is the desired value
        for the joint.

        Raises:
            RuntimeError: Thrown if posture task is not enabled.
        """
        try:
            self.tasks[get_posture_task_name()].set_target(
                custom_configuration_vector(self._robot, **kwargs)
            )
        except KeyError:
            raise RuntimeError(
                "Posture task is not enabled! Enable the task by calling `enable_posture_task()`."
            )

    def update_configuration(self, q: np.ndarray):
        self.configuration.q = q
        self.configuration.update()

    def compute_ik_velocity(self, dt: float = 0.01) -> np.ndarray:
        """Compute the joint velocities that would bring the configuration closer to the IK target solution.

        NOTE: This does not update the self.configuration object.

        Args:
            dt (float, optional): The integration timestep. Defaults to 0.01.

        Returns:
            np.ndarray: Desired joint velocities.
        """
        return solve_ik(self.configuration, self.tasks.values(), dt, solver="quadprog")

    def compute_ik_configuration_vector(
        self, max_iter: int = 200, dt: float = 0.01, ignore_floating_base: bool = True
    ) -> np.ndarray:
        """Compute the IK configuration that would respect all the defined IK tasks.

        Args:
            max_iter (int, optional): Max iteration to be used for integration. Defaults to 200.
            dt (float, optional): The integration timestep. Defaults to 0.01.
            ignore_floating_base (bool, optional): If true, only returns the joint values,
                not the floating base pose. Defaults to True.

        Returns:
            np.ndarray: Target configuration vector.
        """
        for _ in range(int(max_iter)):
            velocity = self.compute_ik_velocity(dt=dt)
            if (
                ignore_floating_base
                and self._robot.model.joints[1].shortname() == "JointModelFreeFlyer"
            ):
                velocity[:6] = 0.0
            if np.all(np.abs(velocity) < 1e-3):
                break
            self.configuration.integrate_inplace(velocity=velocity, dt=dt)
        else:
            print(
                f"Warning: IK gradient has not converged in {int(max_iter)} iterations. Velocity: {velocity}"
            )
        return self.configuration.q.copy()