osc_demo.py

"""Class definition for OpSpace Controller demonstration
"""

from demo_envs import OpSpaceDemoEnv
from demo import Demo
from demo_path import Square, Line, Rotation
import numpy as np
import time
import logging
logger = logging.getLogger(__name__)

import matplotlib.pyplot as plt
import logging
mpl_logger = logging.getLogger('matplotlib')
mpl_logger.setLevel(logging.WARNING)

import perls2.controllers.utils.control_utils as C
import perls2.controllers.utils.transform_utils as T

AXIS_DIM_NUM = {'x': 0, 'y': 1, 'z': 2}

DEFAULT_ROTATION_DELTA_VAL = np.pi / 100.0


def compute_error(goal_state, new_state):
    """ Compute the error between current state and goal state.
    For OpSpace Demos, this is position and orientation error.
    """
    goal_pos = goal_state[:3]
    new_pos = new_state[:3]

    # Check or convert to correct orientation
    goal_ori = T.convert_euler_quat_2mat(goal_state[3:])
    new_ori = T.convert_euler_quat_2mat(new_state[3:])

    pos_error = np.subtract(goal_pos, new_pos)
    ori_error = C.orientation_error(goal_ori, new_ori)

    return np.hstack((pos_error, ori_error))


class OpSpaceDemo(Demo):
    """Operational space demo for testing EEImpedance and EEPosture control.

    Operational space demo uses a cartesian action space defined from the end-effector
    state. The demo generates an environment based on the config file provided.
    It then steps the enviroment with a series of actions based on command-line arguments.

    OpSpace Demos can be used to demonstrate the controllers ability to follow several types
    of paths. These include:
        -Line (in x, y or z direction): straight line
        -Square : square shape defined by delta_val and num_steps in fixed XY plane.
        -Rotation: Holding end-effector position constant while rotating a fixed amount in the
            x, y or z axis of the initial end-effector frame.

    Attributes:
        env (OpSpaceDemoEnv) : environment to step with action generated by demo.
        ctrl_type (str): string identifying type of controller:
            EEImpedance or EEPosture.
        demo_type (str): type of demo to perform:
            Line, Square, Rotation
        test_fn (str): Robot Interface function to perform action specified by demo.
        delta_val (float): delta between consecutive goal states in the demo. [rad] for Rotation
            demos only. For all others [m]
        axis (str): Specifies axis path follows or rotates about for Line or Rotation Demos respectively


    """
    def __init__(self,
                 ctrl_type,
                 demo_type,
                 test_fn,
                 config_file="demo_control_cfg.yaml",
                 **kwargs):
        """Initialize OpSpaceDemo.

        Args:
            ctrl_type (str): string identifying type of controller:
                EEImpedance or EEPosture.
            demo_type (str): type of demo to perform:
                Line, Square, Rotation
            test_fn (str): Robot Interface function to perform action specified by demo.
            config_file (str): filepath for config file to test with demo

        """
        super().__init__(ctrl_type=ctrl_type,
                         demo_type=demo_type,
                         test_fn=test_fn,
                         config_file=config_file,
                         **kwargs)

        self.env = OpSpaceDemoEnv(config=self.config,
                                  use_visualizer=True,
                                  name=None,
                                  test_fn=self.test_fn,
                                  ctrl_type=self.ctrl_type)

        # Connect and reset robot so we can get the correct initial state.
        self.connect_and_reset_robot()

    def make_demo(**kwargs):
        if (kwargs['demo_type'] == "Line"):
            return OpSpaceLineDemo(**kwargs)
        elif (kwargs['demo_type'] == "Square"):
            return OpSpaceSquareDemo(**kwargs)
        elif (kwargs['demo_type'] == "Rotation"):
            return OpSpaceRotationDemo(**kwargs)


    def get_action_kwargs(self, action):
        """Return dict of kwargs specific to function being tested.

        Args:
            action (list): action to be converted in kwargs for robot_interface functions.
        Returns
            action_kwargs (dict): dictionary with key-value pairs corresponding to
                arguments for robot_interface command being tested.
        """
        action_kwargs  = {}
        if self.test_fn =="move_ee_delta":
            action_kwargs['delta'] = action
            action_kwargs['set_pos'] = None
            action_kwargs['set_ori'] = None
            if self.fix_ori:
                action_kwargs['set_ori'] = self.initial_pose[3:]
            if self.fix_pos:
                action_kwargs['set_pos'] = self.initial_pose[:3]
        if self.test_fn=="set_ee_pose":
            action_kwargs['delta'] = None
            action_kwargs['set_pos'] = action[:3]
            action_kwargs['set_ori'] = action[3:]
        return action_kwargs

    def get_state(self):
        """Get robot state for the OpSpace Demo.
        """
        return self.env.robot_interface.ee_pose

    def get_action(self, goal_pose, current_pose):
        """ Return action corresponding to goal and test_fn

        Args:
            goal_pose (list): 7f position and quaternion of goal pose
            current_pose (list): 7f position and quaternion of current
                pose.

        Returns:
            action (list): action corresponding to robot_interface
                function being tested as a list.

        Notes:
            - If test function is "move_ee_delta" returns
                (list 6f) [dx, dy, dz, ax, ay, az] with
                delta position and delta orientation as axis-angle.
            - If test function is "set_ee_pose" returns (list 6f)
                [x,y,z,qx,qy,qz, w] absolute position and orientation
                as quaternion.
        """
        if self.test_fn =="move_ee_delta":
            action = self.get_delta(goal_pose, current_pose)
        elif self.test_fn =="set_ee_pose":
            action = goal_pose
        else:
            raise ValueError("invalid test fn")
        return action

    def run(self):
        """Run the demo. Execute actions in sequence and calculate error.
        """
        self.print_demo_banner()
        self.print_demo_info()
        logger.debug("EE Pose initial:\n{}\n".format(self.env.robot_interface.ee_pose))
        print("--------")
        input("Press Enter to start sending commands.")
        self.step_through_demo()
        self.env.robot_interface.reset()
        self.env.robot_interface.disconnect()

        if self.plot_error:
            self.plot_errors()
        if self.plot_pos:
            self.plot_positions()

        if self.save:
            self.save_data()

    def get_action_list(self):
        if (self.test_fn == "move_ee_delta"):
            return self.path.deltas
        elif (self.test_fn == "set_ee_pose"):
            return self.path.path

    def get_delta(self, goal_pose, current_pose):
        """Get delta between goal pose and current_pose.

        Args: goal_pose (list): 7f pose [x, y, z, qx, qy, qz, w] . Position and quaternion of goal pose.
              current_pose (list): 7f Position and quaternion of current pose.

        Returns: delta (list) 6f [dx, dy, dz, ax, ay, az] delta position and delta orientation
            as axis-angle.
        """

        if len(goal_pose) != 7:
            raise ValueError("Goal pose incorrect dimension should be 7f")
        if len(current_pose) !=7:
            raise ValueError("Current pose incorrect dimension, should be 7f")

        dpos = np.subtract(goal_pose[:3], current_pose[:3])
        goal_mat = T.quat2mat(goal_pose[3:])
        current_mat = T.quat2mat(current_pose[3:])
        delta_mat_T = np.dot(goal_mat, current_mat.T)
        delta_quat = T.mat2quat(np.transpose(delta_mat_T))
        delta_aa = T.quat2axisangle(delta_quat)

        return np.hstack((dpos, delta_aa)).tolist()

    def plot_sequence_arrows(self, axes, x, y, color):
        """Plot a sequence of arrows given a list of x y points.
        """
        for pt_idx in range(len(x)-1):
            dx = np.subtract(x[pt_idx+1], x[pt_idx])
            dy = np.subtract(y[pt_idx+1], y[pt_idx])

            axes.arrow(x[pt_idx], y[pt_idx], dx, dy, color=color)

    def plot_positions(self):
        """ Plot 3 plots showing xy, xz, and yz position.
        Helps for visualizing decoupling.
        """

        goal_x = [goal[0] for goal in self.goal_poses]
        goal_y = [goal[1] for goal in self.goal_poses]
        goal_z = [goal[2] for goal in self.goal_poses]

        state_x = [state[0] for state in self.states]
        state_y = [state[1] for state in self.states]
        state_z = [state[2] for state in self.states]

        fig, (ax_xy, ax_xz, ax_yz) = plt.subplots(1, 3)

        ax_xy.plot(goal_x, goal_y, 'or')
        ax_xy.plot(state_x, state_y, '*b')
        # self.plot_sequence_arrows(ax_xy, state_x, state_y, 'b')
        ax_xy.set_xlabel("x position (m)")
        ax_xy.set_ylabel("y position (m)")
        ax_xy.set_ylim(bottom=-0.1, top=0.4)
        ax_xy.set_xlim(left=0.35, right=0.75)

        ax_xz.plot(goal_x, goal_z, 'or')
        ax_xz.plot(state_x, state_z, '*b')
        ax_xz.set_xlabel("x position (m)")
        ax_xz.set_ylabel("z position (m)")
        ax_xz.set_ylim(bottom=-0.5, top=0.5)
        ax_xz.set_xlim(left=0.35, right=0.75)

        ax_yz.plot(goal_y, goal_z, 'or')
        ax_yz.plot(state_y, state_z, '*b')
        ax_yz.set_xlabel("y position (m)")
        ax_yz.set_ylabel("z position(m)")
        ax_yz.set_ylim(bottom=0, top=2.0)
        ax_yz.set_xlim(left=-0.5, right=0.5)

        fname = self.demo_name + '_pos.png'
        if self.save_fig:
            print("saving figure")
            plt.savefig(fname)
        plt.show()

    def plot_errors(self):
        """ Plot 6 plots showing errors for each dimension.
        x, y, z and qx, qy, qz euler angles (from C.orientation_error)
        """
        errors_x = [error[0] for error in self.errors]
        errors_y = [error[1] for error in self.errors]
        errors_z = [error[2] for error in self.errors]

        errors_qx = [error[3] for error in self.errors]
        errors_qy = [error[4] for error in self.errors]
        errors_qz = [error[5] for error in self.errors]

        fig, ((e_x, e_y, e_z), (e_qx, e_qy, e_qz)) = plt.subplots(2, 3)

        e_x.plot(errors_x)
        e_x.set_title("X error per step.")
        e_x.set_ylabel("error (m)")
        e_x.set_xlabel("step num")

        e_y.plot(errors_y)
        e_y.set_ylabel("error (m)")
        e_y.set_xlabel("step num")
        e_y.set_title("y error per step")

        e_z.plot(errors_z)
        e_z.set_title("z error per step.")
        e_z.set_ylabel("error (m)")
        e_z.set_xlabel("step num")

        e_qx.plot(errors_qx)
        e_qx.set_title("qx error per step.")
        e_qx.set_ylabel("error (rad)")
        e_qx.set_xlabel("step num")

        e_qy.plot(errors_qy)
        e_qy.set_title("qy error per step.")
        e_qy.set_ylabel("error (rad)")
        e_qy.set_xlabel("step num")

        e_qz.plot(errors_qz)
        e_qz.set_title("qz error per step.")
        e_qz.set_ylabel("error (rad)")
        e_qz.set_xlabel("step num")

        fname = self.demo_name + '_error.png'
        if self.save_fig:
            plt.savefig(fname)
        plt.show()
        plt.show()

    def reset_log(self):
        """Reset states, actions and observation logs between demos.
        """
        self.states = []
        self.errors = []
        self.actions = []
        self.observations = []


class OpSpaceLineDemo(OpSpaceDemo):
    """Opspace demo for straight line path.

    End-effector is made to travel in a straight line in either x, y or z direction.

    Attributes:
        env (OpSpaceDemoEnv) : environment to step with action generated by demo.
        ctrl_type (str): string identifying type of controller:
            EEImpedance or EEPosture.
        demo_type (str): type of demo to perform:
            Line, Square, Rotation
        test_fn (str): Robot Interface function to perform action specified by demo.
        delta_val (float): delta [m] between end-effector positions in consecutive goal states.
        axis (str): Specifies axis along which end-effector moves in demo.
    """
    def __init__(self,
                 ctrl_type="EEImpedance",
                 test_fn="move_ee_delta",
                 config_file="demo_control_cfg.yaml",
                 axis='x',
                 path_length=None,
                 delta_val=0.001,
                 num_steps=50, **kwargs):

        super().__init__(ctrl_type=ctrl_type,
                         test_fn=test_fn,
                         demo_type="Line",
                         config_file=config_file,
                         **kwargs)

        # if path_length is defined, use it.
        self.path_length = path_length
        self.axis = axis;
        self.num_steps = num_steps
        self.delta_val = delta_val
        if path_length is not None:
            self.delta_val = self.path_length / self.num_steps

        self.initial_pose = self.env.robot_interface.ee_pose
        self.path = Line(start_pose=self.initial_pose,
                         num_pts=self.num_steps,
                         delta_val=self.delta_val,
                         dim=AXIS_DIM_NUM[self.axis])
        self.fix_ori = True
        self.fix_pos = False

        self.goal_poses = self.path.path
        self.action_list = self.get_action_list()


class OpSpaceLineXYZDemo(OpSpaceLineDemo):
    """Combine OpSpaceLine Demos for x,y and z axes.
    """
    def __init__(self,
                 ctrl_type="EEImpedance",
                 test_fn="move_ee_delta",
                 config_file="demo_control_cfg.yaml",
                 axis='x',
                 path_length=None,
                 delta_val=0.001,
                 num_steps=50, **kwargs):
        super().__init__(ctrl_type=ctrl_type,
                         test_fn=test_fn,
                         config_file=config_file,
                         axis=axis,
                         path_length=path_length,
                         delta_val=delta_val,
                         num_steps=num_steps,
                         **kwargs)

    def print_demo_info(self):
        print("\n\tRunning Line XYZ demo \nwith control type {}. \
            \nTest function {}".format(
            self.ctrl_type, self.demo_type, self.test_fn))
        print("Robot will perform 3 demonstrations: \n moving end-effector in straight line in x, y and then z direction. \nRobot will reset before continuing to next demo.")

    def run(self):
        """Run the demo. Execute actions in sequence and calculate error.
        """
        self.print_demo_banner()
        self.print_demo_info()

        # Run Line demo on x-axis
        self.run_demo_axis('x')
        self.run_demo_axis('y')
        self.run_demo_axis('z')

        print("Demo complete...disconnecting robot.")
        self.env.robot_interface.disconnect()

    def run_demo_axis(self, axis):
        """Run the demo with path along specific axis.

        Execute actions in sequence and log states, observations, errors at each step.
        """
        self.reset_log()
        self.path = Line(start_pose=self.initial_pose,
                         num_pts=self.num_steps,
                         delta_val=self.delta_val,
                         dim=AXIS_DIM_NUM[axis])

        self.goal_poses = self.path.path
        self.action_list = self.get_action_list()

        print("------------------------")
        print("Running Line {} demo \nwith control type {}. \
            \nTest function {}".format(
            axis, self.ctrl_type,  self.test_fn))
        logger.debug("EE Pose initial:\n{}\n".format(self.env.robot_interface.ee_pose))
        print("--------")
        input("Press Enter to start sending commands.")

        self.step_through_demo()
        self.env.robot_interface.reset()
        if self.plot_error:
            self.plot_errors()
        if self.plot_pos:
            self.plot_positions()

        if self.save:
            self.save_data()


class OpSpaceFixedPoseDemo(OpSpaceLineDemo):
    """OpSpace Demo for maintaining initial ee pose.
    """
    def __init__(self,
                 ctrl_type="EEImpedance",
                 test_fn="move_ee_delta",
                 config_file="demo_control_cfg.yaml",
                 num_steps=100, **kwargs):

        super().__init__(
            ctrl_type=ctrl_type,
            test_fn=test_fn,
            config_file=config_file,
            axis='x',
            path_length=None,
            delta_val=0.0,
            num_steps=num_steps, **kwargs)

    def print_demo_info(self):
        print("\n\t Running OSC fixed ee-pose demo \n\t with control type {}. \
            \n\t Test function {}".format(
            self.ctrl_type, self.demo_type, self.test_fn))
        print("Robot will hold initial end-effector pose:\n\t {}".format(self.initial_pose))



class OpSpaceSquareDemo(OpSpaceDemo):
    """OpSpace Demo for square path in xy plane.

    End-effector holds fixed orientation and travels in a square specified by
    delta_val and num_steps.

    """
    def __init__(self,
                 ctrl_type,
                 test_fn,
                 config_file="demo_control_cfg.yaml",
                 delta_val=0.001,
                 num_steps=400,
                 cycles=1,
                 **kwargs):

        kwargs['demo_type'] = "Square"
        super().__init__(ctrl_type=ctrl_type,
                         test_fn=test_fn,
                         config_file=config_file,
                         **kwargs)

        self.delta_val = delta_val
        self.num_steps = num_steps
        self.cycles = cycles
        self.initial_pose = self.env.robot_interface.ee_pose
        self.path = Square(start_pose=self.initial_pose,
                           side_num_pts=int(self.num_steps/4),
                           delta_val=self.delta_val)
        if cycles < 1:
            cycles = 1

        self.goal_poses = self.path.path * cycles
        self.fix_ori = True
        self.fix_pos = False

        self.action_list = self.get_action_list()

class OpSpaceRotationDemo(OpSpaceDemo):
    """OpSpace demo for rotating end-effector while maintaining position.
    """
    def __init__(self,
                 ctrl_type,
                 test_fn,
                 config_file="demo_control_cfg.yaml",
                 axis='x',
                 delta_val=0.00785,
                 num_steps=100,
                 **kwargs):

        kwargs['demo_type'] = "Rotation"
        super().__init__(ctrl_type=ctrl_type,
                         test_fn=test_fn,
                         config_file=config_file,
                         **kwargs)

        self.delta_val = delta_val
        self.num_steps = num_steps
        self.axis = axis
        self.initial_pose = self.env.robot_interface.ee_pose
        self.path = Rotation(
            start_pose=self.initial_pose,
            num_pts=self.num_steps,
            delta_val=self.delta_val,
            dim=AXIS_DIM_NUM[self.axis])

        self.goal_poses = self.path.path
        self.fix_pos = True
        self.fix_ori = False

        self.action_list = self.get_action_list()


class OpSpaceRotationXYZDemo(OpSpaceRotationDemo):
    def __init__(self,
                 ctrl_type,
                 test_fn,
                 config_file="demo_control_cfg.yaml",
                 delta_val=DEFAULT_ROTATION_DELTA_VAL,
                 axis='x',
                 num_steps=100,
                 **kwargs):
        super().__init__(ctrl_type=ctrl_type,
                         test_fn=test_fn,
                         config_file=config_file,
                         axis=axis,
                         delta_val=delta_val,
                         num_steps=num_steps,
                         **kwargs)

    def print_demo_info(self):
        print("\n\tRunning Rotation XYZ demo \nwith control type {}. \
            \nTest function {}".format(
            self.ctrl_type, self.demo_type, self.test_fn))
        print("Robot will perform 3 demos: \n rotating in x, y and then z direction. \nRobot will reset before continuing to next demo.")

    def run(self):
        """Run the demo. Execute actions in sequence and calculate error.
        """
        self.print_demo_banner()
        self.print_demo_info()

        # Run Line demo on x-axis
        self.run_demo_axis('x')
        self.run_demo_axis('y')
        self.run_demo_axis('z')

        print("Demo complete...disconnecting robot.")
        self.env.robot_interface.disconnect()

    def run_demo_axis(self, axis):
        """Run the demo with path along specific axis.

        Execute actions in sequence and log states, observations, errors at each step.
        """
        self.reset_log()
        self.path = Rotation(start_pose=self.initial_pose,
                         num_pts=self.num_steps,
                         delta_val=self.delta_val,
                         dim=AXIS_DIM_NUM[axis])

        self.goal_poses = self.path.path
        self.action_list = self.get_action_list()

        print("------------------------")
        print("Running Rotation {} demo \nwith control type {}. \
            \nTest function {}".format(
            axis, self.ctrl_type,  self.test_fn))
        logger.debug("EE Pose initial:\n{}\n".format(self.env.robot_interface.ee_pose))
        print("--------")
        input("Press Enter to start sending commands.")

        self.step_through_demo()
        self.env.robot_interface.reset()
        if self.plot_error:
            self.plot_errors()
        if self.plot_pos:
            self.plot_positions()

        if self.save:
            self.save_data()


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(
        description="Test controllers and measure errors.")
    parser.add_argument('--ctrl_type',
                        default="EEImpedance",
                        help='Type of controller to test')
    parser.add_argument('--demo_type',
                        default="Line",
                        help='Type of menu to run.')
    parser.add_argument('--test_fn',
                        default='set_ee_pose',
                        help='Function to test',
                        choices=['set_ee_pose', 'move_ee_delta', 'set_joint_delta', 'set_joint_positions', 'set_joint_torques', 'set_joint_velocities'])
    parser.add_argument('--path_length', type=float,
                        default=None, help='length in m of path')
    parser.add_argument('--delta_val',
                        default=0.001, type=float,
                        help="Max step size (m or rad) to take for demo.")
    parser.add_argument('--axis',
                        default='x', type=str,
                        choices=['x', 'y', 'z'],
                        help='axis for demo. Position direction for Line or rotation axis for Rotation')
    parser.add_argument('--num_steps', default=100, type=int,
                        help="max steps for demo.")
    parser.add_argument('--plot_pos', action="store_true",
                        help="whether to plot positions of demo.")
    parser.add_argument('--plot_error', action="store_true",
                        help="whether to plot errors.")
    parser.add_argument('--save', action="store_true",
                        help="whether to store data to file")
    parser.add_argument('--demo_name', default=None,
                        type=str, help="Valid filename for demo.")
    parser.add_argument('--save_fig', action="store_true",
                        help="whether to save pngs of plots")
    parser.add_argument('--fix_ori', action="store_true",
                        help="fix orientation for move_ee_delta")
    parser.add_argument('--fix_pos', action="store_true",
                        help="fix position for move_ee_delta")
    parser.add_argument('--config_file', default='demos/demo_control_cfg.yaml', help='absolute filepath for config file.')
    parser.add_argument('--cycles', type=int, default=1, help="num times to cycle path (only for square)")
    args = parser.parse_args()
    kwargs = vars(args)
    demo = OpSpaceDemo.make_demo(**kwargs)
    demo.run()