Source code for concepts.dm.crowhat.robots.panda_robot_interfaces

#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File   : panda_robot_interfaces.py
# Author : Jiayuan Mao
# Email  : maojiayuan@gmail.com
# Date   : 03/30/2024
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.

from typing import Optional, Tuple, Callable

import numpy as np

from concepts.hw_interface.franka.remote_client import FrankaRemoteClient
from concepts.simulator.pybullet.components.panda.panda_robot import PandaRobot

from concepts.dm.crowhat.robots.robot_interfaces import RobotControllerExecutionContext, RobotControllerExecutionFailed, PrimitiveRobotControllerInterface, ContactPrimitiveRobotControlInterface, RobotArmJointTrajectory
from concepts.dm.crowhat.pybullet_interfaces.pybullet_sim_interfaces import PyBulletSimulationControllerInterface, PyBulletPhysicalControllerInterface

__all__ = [
    'PandaRobotPrimitiveControllerInterface', 'PandaRobotContactPrimitiveControllerInterface',
    'PyBulletPandaRobotPrimitiveControllerInterface', 'PyBulletPandaRobotContactPrimitiveControllerInterface',
    'PhysicalPandaRobotPrimitiveControllerInterface', 'PhysicalPandaRobotContactPrimitiveControllerInterface'
]




class PandaRobotPrimitiveControllerInterface(PrimitiveRobotControllerInterface):


    def get_qpos(self) -> np.ndarray:
        raise NotImplementedError




    def get_qvel(self) -> np.ndarray:
        raise NotImplementedError




    def get_ee_pose(self) -> Tuple[np.ndarray, np.ndarray]:
        raise NotImplementedError




    def move_home(self) -> bool:
        raise NotImplementedError




    def move_qpos(self, qpos: np.ndarray) -> None:
        raise NotImplementedError




    def move_qpos_trajectory(self, trajectory: RobotArmJointTrajectory) -> None:
        raise NotImplementedError




    def open_gripper(self) -> None:
        raise NotImplementedError




    def close_gripper(self) -> None:
        raise NotImplementedError




    def grasp(self, width: float = 0.05, force: float = 40) -> None:
        raise NotImplementedError






class PandaRobotContactPrimitiveControllerInterface(ContactPrimitiveRobotControlInterface):


    def __init__(self, primitive_interface: PandaRobotPrimitiveControllerInterface):
        super().__init__()
        self._primitive_interface = primitive_interface


    @property
    def primitive_interface(self) -> PandaRobotPrimitiveControllerInterface:
        return self._primitive_interface



    def grasp(self, approaching_trajectory: RobotArmJointTrajectory, width: float = 0.05, force: float = 40) -> None:
        self.primitive_interface.move_qpos_trajectory(approaching_trajectory)
        self.primitive_interface.grasp(width, force)




    def place(self, placing_trajectory: RobotArmJointTrajectory) -> None:
        self.primitive_interface.move_qpos_trajectory(placing_trajectory)
        self.primitive_interface.open_gripper()






class PyBulletPandaRobotPrimitiveControllerInterface(PandaRobotPrimitiveControllerInterface):


    def __init__(self, panda_robot):
        super().__init__()
        self._panda_robot = panda_robot


    @property
    def panda_robot(self) -> PandaRobot:
        return self._panda_robot



    def get_qpos(self) -> np.ndarray:
        return self.panda_robot.get_qpos()




    def get_qvel(self) -> np.ndarray:
        raise NotImplementedError('qvel for panda robot is not implemented.')




    def get_ee_pose(self) -> Tuple[np.ndarray, np.ndarray]:
        return self.panda_robot.get_ee_pose()




    def move_home(self, **kwargs) -> None:
        self.get_update_default_parameters('move_home', kwargs)
        with RobotControllerExecutionContext('move_home', **kwargs) as ctx:
            ctx.monitor(self.panda_robot.move_home(**kwargs))




    def move_qpos(self, qpos: np.ndarray, **kwargs) -> None:
        self.get_update_default_parameters('move_qpos', kwargs)
        with RobotControllerExecutionContext('move_qpos', qpos, **kwargs) as ctx:
            ctx.monitor(self.panda_robot.move_qpos(qpos, **kwargs))




    def move_pose(self, pos: np.ndarray, quat: np.ndarray, **kwargs) -> None:
        self.get_update_default_parameters('move_pose', kwargs)
        with RobotControllerExecutionContext('move_pose', pos, quat, **kwargs) as ctx:
            ctx.monitor(self.panda_robot.move_pose(pos, quat, **kwargs))




    def move_qpos_trajectory(self, trajectory: RobotArmJointTrajectory, **kwargs) -> None:
        self.get_update_default_parameters('move_qpos_trajectory', kwargs)
        with RobotControllerExecutionContext('move_qpos_trajectory', trajectory, **kwargs) as ctx:
            ctx.monitor(self.panda_robot.move_qpos_path_v2(trajectory, **kwargs))




    def open_gripper(self, **kwargs) -> None:
        self.get_update_default_parameters('open_gripper', kwargs)
        with RobotControllerExecutionContext('open_gripper', **kwargs) as ctx:
            ctx.monitor(self.panda_robot.open_gripper_free(**kwargs))




    def close_gripper(self, **kwargs) -> None:
        self.get_update_default_parameters('close_gripper', kwargs)
        with RobotControllerExecutionContext('close_gripper', **kwargs) as ctx:
            ctx.monitor(self.panda_robot.close_gripper_free(**kwargs))




    def grasp(self, width: float = 0.05, force: float = 40, **kwargs) -> None:
        # NB(Jiayuan Mao @ 2024/03/30): width and force parameters are ignored in the PyBullet-based implementation because it uses magic attachments.
        self.get_update_default_parameters('grasp', kwargs)
        with RobotControllerExecutionContext('grasp', **kwargs) as ctx:
            ctx.monitor(self.panda_robot.grasp(**kwargs))




    def attach_physical_interface(self, physical_interface: PyBulletPhysicalControllerInterface):
        physical_interface.register_controller('move_home', self.move_home)
        physical_interface.register_controller('move_qpos', self.move_qpos)
        physical_interface.register_controller('move_pose', self.move_pose)
        physical_interface.register_controller('move_qpos_trajectory', self.move_qpos_trajectory)
        physical_interface.register_controller('open_gripper', self.open_gripper)
        physical_interface.register_controller('close_gripper', self.close_gripper)
        physical_interface.register_controller('grasp', self.grasp)




    def attach_simulation_interface(self, simulation_interface: PyBulletSimulationControllerInterface):
        simulation_interface.register_controller('move_home', self.move_home)
        simulation_interface.register_controller('move_qpos', self.move_qpos)
        simulation_interface.register_controller('move_pose', self.move_pose)
        simulation_interface.register_controller('move_qpos_trajectory', self.move_qpos_trajectory)
        simulation_interface.register_controller('open_gripper', self.open_gripper)
        simulation_interface.register_controller('close_gripper', self.close_gripper)
        simulation_interface.register_controller('grasp', self.grasp)






class PyBulletPandaRobotContactPrimitiveControllerInterface(PandaRobotContactPrimitiveControllerInterface):


    def __init__(self, primitive_interface: Optional[PyBulletPandaRobotPrimitiveControllerInterface] = None, timeout_multiplier: float = 1.0):
        super().__init__(primitive_interface)
        self.timeout_multiplier = timeout_multiplier
        self.set_default_parameters('grasp', move_home_timeout=30, trajectory_timeout=50)
        self.set_default_parameters('place', trajectory_timeout=50)


    @property
    def primitive_interface(self) -> PyBulletPandaRobotPrimitiveControllerInterface:
        return self._primitive_interface

    _primitive_interface: PyBulletPandaRobotPrimitiveControllerInterface



    def grasp(self, approaching_trajectory: RobotArmJointTrajectory, width: float = 0.05, force: float = 40, **kwargs):
        self.get_update_default_parameters('grasp', kwargs)
        move_home_timeout = kwargs['move_home_timeout']
        trajectory_timeout = kwargs['trajectory_timeout']

        try:
            self.primitive_interface.move_home(timeout=move_home_timeout)
            self._smooth_move_qpos_trajectory(approaching_trajectory, timeout=trajectory_timeout)
            self.primitive_interface.grasp(width, force)
            self._smooth_move_qpos_trajectory(approaching_trajectory[:-5][::-1], timeout=trajectory_timeout)
        except RobotControllerExecutionFailed:
            raise




    def place(self, placing_trajectory: RobotArmJointTrajectory, **kwargs):
        self.get_update_default_parameters('place', kwargs)
        trajectory_timeout = kwargs['trajectory_timeout']

        try:
            self._smooth_move_qpos_trajectory(placing_trajectory, timeout=trajectory_timeout)
            self.primitive_interface.open_gripper()
            self.primitive_interface.move_home()
        except RobotControllerExecutionFailed:
            raise


    def _smooth_move_qpos_trajectory(self, qt: RobotArmJointTrajectory, timeout: float = 10):
        robot = self.primitive_interface.panda_robot
        qt = qt.copy()
        qt.insert(0, robot.get_qpos())

        smooth_qt = [qt[0]]
        cspace = robot.get_configuration_space()
        for qpos1, qpos2 in zip(qt[:-1], qt[1:]):
            smooth_qt.extend(cspace.gen_path(qpos1, qpos2)[1][1:])

        with RobotControllerExecutionContext('move_qpos_trajectory', smooth_qt) as ctx:
            ctx.monitor(robot.move_qpos_path_v2(smooth_qt, timeout=timeout * self.timeout_multiplier))



    def attach_physical_interface(self, physical_interface: PyBulletPhysicalControllerInterface):
        physical_interface.register_controller('ctl_grasp', self.grasp)
        physical_interface.register_controller('ctl_place', self.place)




    def attach_simulation_interface(self, simulation_interface: PyBulletSimulationControllerInterface):
        simulation_interface.register_controller('ctl_grasp', self.grasp)
        simulation_interface.register_controller('ctl_place', self.place)






class PhysicalPandaRobotPrimitiveControllerInterface(PandaRobotPrimitiveControllerInterface):


    def __init__(self, panda_robot: FrankaRemoteClient):
        super().__init__()
        self._panda_robot = panda_robot


    @property
    def panda_robot(self) -> FrankaRemoteClient:
        return self._panda_robot



    def get_qpos(self) -> np.ndarray:
        return self.panda_robot.get_qpos()




    def get_qvel(self) -> np.ndarray:
        return self.panda_robot.get_qvel()




    def get_ee_pose(self) -> Tuple[np.ndarray, np.ndarray]:
        return self.panda_robot.get_ee_pose()




    def move_home(self) -> bool:
        return self.panda_robot.move_home()




    def move_qpos(self, qpos: np.ndarray, **kwargs) -> None:
        self.get_update_default_parameters('move_qpos', kwargs)
        self.panda_robot.move_qpos(qpos, **kwargs)




    def move_qpos_trajectory(self, trajectory: RobotArmJointTrajectory, **kwargs) -> None:
        self.get_update_default_parameters('move_qpos_trajectory', kwargs)
        self.panda_robot.move_qpos_trajectory(trajectory, **kwargs)




    def open_gripper(self, **kwargs) -> None:
        self.get_update_default_parameters('open_gripper', kwargs)
        self.panda_robot.open_gripper(**kwargs)




    def close_gripper(self, **kwargs) -> None:
        self.get_update_default_parameters('close_gripper', kwargs)
        self.panda_robot.close_gripper(**kwargs)




    def grasp(self, width: float = 0.05, force: float = 40, **kwargs) -> None:
        self.get_update_default_parameters('grasp', kwargs)
        self.panda_robot.grasp(width, force, **kwargs)






class PhysicalPandaRobotContactPrimitiveControllerInterface(PandaRobotContactPrimitiveControllerInterface):


    def __init__(self, primitive_interface: PhysicalPandaRobotPrimitiveControllerInterface):
        super().__init__(primitive_interface)




    def attach_simulation_interface(self, simulation_interface: PyBulletPhysicalControllerInterface):
        simulation_interface.register_controller('ctl_grasp', self.grasp)
        simulation_interface.register_controller('ctl_place', self.place)




    def grasp(self, approaching_trajectory: RobotArmJointTrajectory, width: float = 0.05, force: float = 40, **kwargs) -> None:
        self.get_update_default_parameters('grasp', kwargs)
        self.primitive_interface.move_qpos_trajectory(approaching_trajectory)
        self.primitive_interface.grasp(width, force)




    def place(self, placing_trajectory: RobotArmJointTrajectory, **kwargs) -> None:
        self.get_update_default_parameters('place', kwargs)
        self.primitive_interface.move_qpos_trajectory(placing_trajectory)
        self.primitive_interface.open_gripper()