phantom/submodules/phantom-robosuite/robosuite/robots/bimanual.py

import copy
import os
from collections import OrderedDict

import numpy as np

import robosuite.utils.transform_utils as T
from robosuite.controllers import controller_factory, load_controller_config
from robosuite.models.grippers import gripper_factory
from robosuite.robots.manipulator import Manipulator
from robosuite.utils.buffers import DeltaBuffer, RingBuffer
from robosuite.utils.observables import Observable, sensor


class Bimanual(Manipulator):
    """
    Initializes a bimanual robot simulation object.

    Args:
        robot_type (str): Specification for specific robot arm to be instantiated within this env (e.g: "Panda")

        idn (int or str): Unique ID of this robot. Should be different from others

        controller_config (dict or list of dict --> dict of dict): If set, contains relevant controller parameters
            for creating custom controllers. Else, uses the default controller for this specific task. Should either
            be single dict if same controller is to be used for both robot arms or else it should be a list of length 2.

            :NOTE: In the latter case, assumes convention of [right, left]

        initial_qpos (sequence of float): If set, determines the initial joint positions of the robot to be
            instantiated for the task

        initialization_noise (dict): Dict containing the initialization noise parameters. The expected keys and
            corresponding value types are specified below:

            :`'magnitude'`: The scale factor of uni-variate random noise applied to each of a robot's given initial
                joint positions. Setting this value to "None" or 0.0 results in no noise being applied.
                If "gaussian" type of noise is applied then this magnitude scales the standard deviation applied,
                If "uniform" type of noise is applied then this magnitude sets the bounds of the sampling range
            :`'type'`: Type of noise to apply. Can either specify "gaussian" or "uniform"

            :Note: Specifying None will automatically create the required dict with "magnitude" set to 0.0

        mount_type (str): type of mount, used to instantiate mount models from mount factory.
            Default is "default", which is the default mount associated with this robot's corresponding model.
            None results in no mount, and any other (valid) model overrides the default mount.

        gripper_type (str or list of str --> dict): type of gripper, used to instantiate
            gripper models from gripper factory. Default is "default", which is the default gripper associated
            within the 'robot' specification. None removes the gripper, and any other (valid) model overrides the
            default gripper. Should either be single str if same gripper type is to be used for both arms or else
            it should be a list of length 2

            :NOTE: In the latter case, assumes convention of [right, left]

        control_freq (float): how many control signals to receive
            in every second. This sets the amount of simulation time
            that passes between every action input.
    """

    def __init__(
        self,
        robot_type: str,
        idn=0,
        controller_config=None,
        initial_qpos=None,
        initialization_noise=None,
        mount_type="default",
        gripper_type="default",
        control_freq=20,
    ):

        self.controller = self._input2dict(None)
        self.controller_config = self._input2dict(copy.deepcopy(controller_config))
        self.gripper = self._input2dict(None)
        self.gripper_type = self._input2dict(gripper_type)
        self.has_gripper = self._input2dict([gripper_type is not None for _, gripper_type in self.gripper_type.items()])

        self.gripper_joints = self._input2dict(None)  # xml joint names for gripper
        self._ref_gripper_joint_pos_indexes = self._input2dict(None)  # xml gripper joint position indexes in mjsim
        self._ref_gripper_joint_vel_indexes = self._input2dict(None)  # xml gripper joint velocity indexes in mjsim
        self._ref_joint_gripper_actuator_indexes = self._input2dict(
            None
        )  # xml gripper (pos) actuator indexes for robot in mjsim
        self.eef_rot_offset = self._input2dict(None)  # rotation offsets from final arm link to gripper (quat)
        self.eef_site_id = self._input2dict(None)  # xml element id for eef in mjsim
        self.eef_cylinder_id = self._input2dict(None)  # xml element id for eef cylinder in mjsim
        self.torques = None  # Current torques being applied

        self.recent_ee_forcetorques = self._input2dict(None)  # Current and last forces / torques sensed at eef
        self.recent_ee_pose = self._input2dict(None)  # Current and last eef pose (pos + ori (quat))
        self.recent_ee_vel = self._input2dict(None)  # Current and last eef velocity
        self.recent_ee_vel_buffer = self._input2dict(None)  # RingBuffer holding prior 10 values of velocity values
        self.recent_ee_acc = self._input2dict(None)  # Current and last eef acceleration

        super().__init__(
            robot_type=robot_type,
            idn=idn,
            initial_qpos=initial_qpos,
            initialization_noise=initialization_noise,
            mount_type=mount_type,
            control_freq=control_freq,
        )

    def _load_controller(self):
        """
        Loads controller to be used for dynamic trajectories
        """
        # Flag for loading urdf once (only applicable for IK controllers)
        urdf_loaded = False

        # Load controller configs for both left and right arm
        for arm in self.arms:
            # First, load the default controller if none is specified
            if not self.controller_config[arm]:
                # Need to update default for a single agent
                controller_path = os.path.join(
                    os.path.dirname(__file__),
                    "..",
                    "controllers/config/{}.json".format(self.robot_model.default_controller_config[arm]),
                )
                self.controller_config[arm] = load_controller_config(custom_fpath=controller_path)

            # Assert that the controller config is a dict file:
            #             NOTE: "type" must be one of: {JOINT_POSITION, JOINT_TORQUE, JOINT_VELOCITY,
            #                                           OSC_POSITION, OSC_POSE, IK_POSE}
            assert (
                type(self.controller_config[arm]) == dict
            ), "Inputted controller config must be a dict! Instead, got type: {}".format(
                type(self.controller_config[arm])
            )

            # Add to the controller dict additional relevant params:
            #   the robot name, mujoco sim, eef_name, actuator_range, joint_indexes, timestep (model) freq,
            #   policy (control) freq, and ndim (# joints)
            self.controller_config[arm]["robot_name"] = self.name
            self.controller_config[arm]["sim"] = self.sim
            self.controller_config[arm]["eef_name"] = self.gripper[arm].important_sites["grip_site"]
            self.controller_config[arm]["eef_rot_offset"] = self.eef_rot_offset[arm]
            self.controller_config[arm]["ndim"] = self._joint_split_idx
            self.controller_config[arm]["policy_freq"] = self.control_freq
            (start, end) = (None, self._joint_split_idx) if arm == "right" else (self._joint_split_idx, None)
            self.controller_config[arm]["joint_indexes"] = {
                "joints": self.joint_indexes[start:end],
                "qpos": self._ref_joint_pos_indexes[start:end],
                "qvel": self._ref_joint_vel_indexes[start:end],
            }
            self.controller_config[arm]["actuator_range"] = (
                self.torque_limits[0][start:end],
                self.torque_limits[1][start:end],
            )

            # Only load urdf the first time this controller gets called
            self.controller_config[arm]["load_urdf"] = True if not urdf_loaded else False
            urdf_loaded = True

            # Instantiate the relevant controller
            self.controller[arm] = controller_factory(self.controller_config[arm]["type"], self.controller_config[arm])

    def load_model(self):
        """
        Loads robot and optionally add grippers.
        """
        # First, run the superclass method to load the relevant model
        super().load_model()

        # Verify that the loaded model is of the correct type for this robot
        if self.robot_model.arm_type != "bimanual":
            raise TypeError(
                "Error loading robot model: Incompatible arm type specified for this robot. "
                "Requested model arm type: {}, robot arm type: {}".format(self.robot_model.arm_type, type(self))
            )

        # Now, load the gripper if necessary
        for arm in self.arms:
            if self.has_gripper[arm]:
                if self.gripper_type[arm] == "default":
                    # Load the default gripper from the robot file
                    self.gripper[arm] = gripper_factory(
                        self.robot_model.default_gripper[arm], idn="_".join((str(self.idn), arm))
                    )
                else:
                    # Load user-specified gripper
                    self.gripper[arm] = gripper_factory(self.gripper_type[arm], idn="_".join((str(self.idn), arm)))
            else:
                # Load null gripper
                self.gripper[arm] = gripper_factory(None, idn="_".join((str(self.idn), arm)))
            # Grab eef rotation offset
            self.eef_rot_offset[arm] = T.quat_multiply(
                self.robot_model.hand_rotation_offset[arm], self.gripper[arm].rotation_offset
            )
            # Add this gripper to the robot model
            self.robot_model.add_gripper(self.gripper[arm], self.robot_model.eef_name[arm])

    def reset(self, deterministic=False):
        """
        Sets initial pose of arm and grippers. Overrides gripper joint configuration if we're using a
        deterministic reset (e.g.: hard reset from xml file)

        Args:
            deterministic (bool): If true, will not randomize initializations within the sim
        """
        # First, run the superclass method to reset the position and controller
        super().reset(deterministic)

        # Setup arm-specific values
        for arm in self.arms:
            # Now, reset the grippers if necessary
            if self.has_gripper[arm]:
                if not deterministic:
                    self.sim.data.qpos[self._ref_gripper_joint_pos_indexes[arm]] = self.gripper[arm].init_qpos

                self.gripper[arm].current_action = np.zeros(self.gripper[arm].dof)

            # Update base pos / ori references in controller (technically only needs to be called once)
            self.controller[arm].update_base_pose(self.base_pos, self.base_ori)
            # Setup buffers for eef values
            self.recent_ee_forcetorques[arm] = DeltaBuffer(dim=6)
            self.recent_ee_pose[arm] = DeltaBuffer(dim=7)
            self.recent_ee_vel[arm] = DeltaBuffer(dim=6)
            self.recent_ee_vel_buffer[arm] = RingBuffer(dim=6, length=10)
            self.recent_ee_acc[arm] = DeltaBuffer(dim=6)

    def setup_references(self):
        """
        Sets up necessary reference for robots, grippers, and objects.

        Note that this should get called during every reset from the environment
        """
        # First, run the superclass method to setup references for joint-related values / indexes
        super().setup_references()

        # Now, add references to gripper if necessary
        # indices for grippers in qpos, qvel
        for arm in self.arms:
            if self.has_gripper[arm]:
                self.gripper_joints[arm] = list(self.gripper[arm].joints)
                self._ref_gripper_joint_pos_indexes[arm] = [
                    self.sim.model.get_joint_qpos_addr(x) for x in self.gripper_joints[arm]
                ]
                self._ref_gripper_joint_vel_indexes[arm] = [
                    self.sim.model.get_joint_qvel_addr(x) for x in self.gripper_joints[arm]
                ]
                self._ref_joint_gripper_actuator_indexes[arm] = [
                    self.sim.model.actuator_name2id(actuator) for actuator in self.gripper[arm].actuators
                ]

            # IDs of sites for eef visualization
            self.eef_site_id[arm] = self.sim.model.site_name2id(self.gripper[arm].important_sites["grip_site"])
            self.eef_cylinder_id[arm] = self.sim.model.site_name2id(self.gripper[arm].important_sites["grip_cylinder"])

    def control(self, action, policy_step=False):
        """
        Actuate the robot with the
        passed joint velocities and gripper control.

        Args:
            action (np.array): The control to apply to the robot. The first @self.robot_model.dof dimensions should
                be the desired normalized joint velocities and if the robot has a gripper, the next @self.gripper.dof
                dimensions should be actuation controls for the gripper.

                :NOTE: Assumes inputted actions are of form:
                    [right_arm_control, right_gripper_control, left_arm_control, left_gripper_control]

            policy_step (bool): Whether a new policy step (action) is being taken

        Raises:
            AssertionError: [Invalid action dimension]
        """
        # clip actions into valid range
        assert len(action) == self.action_dim, "environment got invalid action dimension -- expected {}, got {}".format(
            self.action_dim, len(action)
        )

        self.torques = np.array([])
        # Now execute actions for each arm
        for arm in self.arms:
            # Make sure to split action space correctly
            (start, end) = (None, self._action_split_idx) if arm == "right" else (self._action_split_idx, None)
            sub_action = action[start:end]

            gripper_action = None
            if self.has_gripper[arm]:
                # get all indexes past controller dimension indexes
                gripper_action = sub_action[self.controller[arm].control_dim :]
                sub_action = sub_action[: self.controller[arm].control_dim]

            # Update the controller goal if this is a new policy step
            if policy_step:
                self.controller[arm].set_goal(sub_action)

            # Now run the controller for a step and add it to the torques
            self.torques = np.concatenate((self.torques, self.controller[arm].run_controller()))

            # Get gripper action, if applicable
            if self.has_gripper[arm]:
                self.grip_action(gripper=self.gripper[arm], gripper_action=gripper_action)

        # Clip the torques
        low, high = self.torque_limits
        self.torques = np.clip(self.torques, low, high)

        # Apply joint torque control
        self.sim.data.ctrl[self._ref_joint_actuator_indexes] = self.torques

        # If this is a policy step, also update buffers holding recent values of interest
        if policy_step:
            # Update proprioceptive values
            self.recent_qpos.push(self._joint_positions)
            self.recent_actions.push(action)
            self.recent_torques.push(self.torques)

            for arm in self.arms:
                # Update arm-specific proprioceptive values
                self.recent_ee_forcetorques[arm].push(np.concatenate((self.ee_force[arm], self.ee_torque[arm])))
                self.recent_ee_pose[arm].push(
                    np.concatenate((self.controller[arm].ee_pos, T.mat2quat(self.controller[arm].ee_ori_mat)))
                )
                self.recent_ee_vel[arm].push(
                    np.concatenate((self.controller[arm].ee_pos_vel, self.controller[arm].ee_ori_vel))
                )

                # Estimation of eef acceleration (averaged derivative of recent velocities)
                self.recent_ee_vel_buffer[arm].push(
                    np.concatenate((self.controller[arm].ee_pos_vel, self.controller[arm].ee_ori_vel))
                )
                diffs = np.vstack(
                    [
                        self.recent_ee_acc[arm].current,
                        self.control_freq * np.diff(self.recent_ee_vel_buffer[arm].buf, axis=0),
                    ]
                )
                ee_acc = np.array([np.convolve(col, np.ones(10) / 10.0, mode="valid")[0] for col in diffs.transpose()])
                self.recent_ee_acc[arm].push(ee_acc)

    def _visualize_grippers(self, visible):
        """
        Visualizes the gripper site(s) if applicable.

        Args:
            visible (bool): True if visualizing the gripper for this arm.
        """
        for arm in self.arms:
            self.gripper[arm].set_sites_visibility(sim=self.sim, visible=visible)

    def setup_observables(self):
        """
        Sets up observables to be used for this robot

        Returns:
            OrderedDict: Dictionary mapping observable names to its corresponding Observable object
        """
        # Get general robot observables first
        observables = super().setup_observables()

        # Get prefix from robot model to avoid naming clashes for multiple robots and define observables modality
        pf = self.robot_model.naming_prefix
        modality = f"{pf}proprio"
        sensors = []
        names = []

        for arm in self.arms:
            # Add in eef info
            arm_sensors, arm_sensor_names = self._create_arm_sensors(arm=arm, modality=modality)
            sensors += arm_sensors
            names += arm_sensor_names

        # Create observables for this robot
        for name, s in zip(names, sensors):
            observables[name] = Observable(
                name=name,
                sensor=s,
                sampling_rate=self.control_freq,
            )

        return observables

    def _create_arm_sensors(self, arm, modality):
        """
        Helper function to create sensors for a given arm. This is abstracted in a separate function call so that we
        don't have local function naming collisions during the _setup_observables() call.

        Args:
            arm (str): Arm to create sensors for
            modality (str): Modality to assign to all sensors

        Returns:
            2-tuple:
                sensors (list): Array of sensors for the given arm
                names (list): array of corresponding observable names
        """
        pf = self.robot_model.naming_prefix

        # eef features
        @sensor(modality=modality)
        def eef_pos(obs_cache):
            return np.array(self.sim.data.site_xpos[self.eef_site_id[arm]])

        @sensor(modality=modality)
        def eef_quat(obs_cache):
            return T.convert_quat(self.sim.data.get_body_xquat(self.robot_model.eef_name[arm]), to="xyzw")

        sensors = [eef_pos, eef_quat]
        names = [f"{pf}{arm}_eef_pos", f"{pf}{arm}_eef_quat"]

        # add in gripper sensors if this robot has a gripper
        if self.has_gripper[arm]:

            @sensor(modality=modality)
            def gripper_qpos(obs_cache):
                return np.array([self.sim.data.qpos[x] for x in self._ref_gripper_joint_pos_indexes[arm]])

            @sensor(modality=modality)
            def gripper_qvel(obs_cache):
                return np.array([self.sim.data.qvel[x] for x in self._ref_gripper_joint_vel_indexes[arm]])

            sensors += [gripper_qpos, gripper_qvel]
            names += [f"{pf}{arm}_gripper_qpos", f"{pf}{arm}_gripper_qvel"]

        return sensors, names

    def _input2dict(self, inp):
        """
        Helper function that converts an input that is either a single value or a list into a dict with keys for
        each arm: "right", "left"

        Args:
            inp (str or list or None): Input value to be converted to dict

            :Note: If inp is a list, then assumes format is [right, left]

        Returns:
            dict: Inputs mapped for each robot arm
        """
        # First, convert to list if necessary
        if type(inp) is not list:
            inp = [inp for _ in range(2)]
        # Now, convert list to dict and return
        return {key: value for key, value in zip(self.arms, inp)}

    @property
    def arms(self):
        """
        Returns name of arms used as naming convention throughout this module

        Returns:
            2-tuple: ('right', 'left')
        """
        return "right", "left"

    @property
    def action_limits(self):
        """
        Action lower/upper limits per dimension.

        Returns:
            2-tuple:

                - (np.array) minimum (low) action values
                - (np.array) maximum (high) action values
        """
        # Action limits based on controller limits
        low, high = [], []
        for arm in self.arms:
            low_g, high_g = (
                ([-1] * self.gripper[arm].dof, [1] * self.gripper[arm].dof) if self.has_gripper[arm] else ([], [])
            )
            low_c, high_c = self.controller[arm].control_limits
            low, high = np.concatenate([low, low_c, low_g]), np.concatenate([high, high_c, high_g])
        return low, high

    @property
    def ee_ft_integral(self):
        """
        Returns:
            dict: each arm-specific entry specifies the integral over time of the applied ee force-torque for that arm
        """
        vals = {}
        for arm in self.arms:
            vals[arm] = np.abs((1.0 / self.control_freq) * self.recent_ee_forcetorques[arm].average)
        return vals

    @property
    def ee_force(self):
        """
        Returns:
            dict: each arm-specific entry specifies the force applied at the force sensor at the robot arm's eef
        """
        vals = {}
        for arm in self.arms:
            vals[arm] = self.get_sensor_measurement(self.gripper[arm].important_sensors["force_ee"])
        return vals

    @property
    def ee_torque(self):
        """
        Returns:
            dict: each arm-specific entry specifies the torque applied at the torque sensor at the robot arm's eef
        """
        vals = {}
        for arm in self.arms:
            vals[arm] = self.get_sensor_measurement(self.gripper[arm].important_sensors["torque_ee"])
        return vals

    @property
    def _hand_pose(self):
        """
        Returns:
            dict: each arm-specific entry specifies the eef pose in base frame of robot.
        """
        vals = {}
        for arm in self.arms:
            vals[arm] = self.pose_in_base_from_name(self.robot_model.eef_name[arm])
        return vals

    @property
    def _hand_quat(self):
        """
        Returns:
            dict: each arm-specific entry specifies the eef quaternion in base frame of robot.
        """
        vals = {}
        orns = self._hand_orn
        for arm in self.arms:
            vals[arm] = T.mat2quat(orns[arm])
        return vals

    @property
    def _hand_total_velocity(self):
        """
        Returns:
            dict: each arm-specific entry specifies the total eef velocity (linear + angular) in the base frame
            as a numpy array of shape (6,)
        """
        vals = {}
        for arm in self.arms:
            # Determine correct start, end points based on arm
            (start, end) = (None, self._joint_split_idx) if arm == "right" else (self._joint_split_idx, None)

            # Use jacobian to translate joint velocities to end effector velocities.
            Jp = self.sim.data.get_body_jacp(self.robot_model.eef_name[arm]).reshape((3, -1))
            Jp_joint = Jp[:, self._ref_joint_vel_indexes[start:end]]

            Jr = self.sim.data.get_body_jacr(self.robot_model.eef_name[arm]).reshape((3, -1))
            Jr_joint = Jr[:, self._ref_joint_vel_indexes[start:end]]

            eef_lin_vel = Jp_joint.dot(self._joint_velocities)
            eef_rot_vel = Jr_joint.dot(self._joint_velocities)
            vals[arm] = np.concatenate([eef_lin_vel, eef_rot_vel])
        return vals

    @property
    def _hand_pos(self):
        """
        Returns:
            dict: each arm-specific entry specifies the position of eef in base frame of robot.
        """
        vals = {}
        poses = self._hand_pose
        for arm in self.arms:
            eef_pose_in_base = poses[arm]
            vals[arm] = eef_pose_in_base[:3, 3]
        return vals

    @property
    def _hand_orn(self):
        """
        Returns:
            dict: each arm-specific entry specifies the orientation of eef in base frame of robot as a rotation matrix.
        """
        vals = {}
        poses = self._hand_pose
        for arm in self.arms:
            eef_pose_in_base = poses[arm]
            vals[arm] = eef_pose_in_base[:3, :3]
        return vals

    @property
    def _hand_vel(self):
        """
        Returns:
            dict: each arm-specific entry specifies the velocity of eef in base frame of robot.
        """
        vels = self._hand_total_velocity
        for arm in self.arms:
            vels[arm] = vels[arm][:3]
        return vels

    @property
    def _hand_ang_vel(self):
        """
        Returns:
            dict: each arm-specific entry specifies the angular velocity of eef in base frame of robot.
        """
        vels = self._hand_total_velocity
        for arm in self.arms:
            vels[arm] = vels[arm][3:]
        return vels

    @property
    def _action_split_idx(self):
        """
        Grabs the index that correctly splits the right arm from the left arm actions

        :NOTE: Assumes inputted actions are of form:
            [right_arm_control, right_gripper_control, left_arm_control, left_gripper_control]

        Returns:
            int: Index splitting right from left arm actions
        """
        return (
            self.controller["right"].control_dim + self.gripper["right"].dof
            if self.has_gripper["right"]
            else self.controller["right"].control_dim
        )

    @property
    def _joint_split_idx(self):
        """
        Returns:
            int: the index that correctly splits the right arm from the left arm joints
        """
        return int(len(self.robot_joints) / 2)