haptic-skin/plot_traj.py at main · VT-Collab/haptic-skin · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
import time
import numpy as np
import pickle
import argparse
import matplotlib.pyplot as plt
from pyquaternion import Quaternion
from scipy import interpolate
from cycler import cycler
from utils import *


Panda = TrajectoryClient()

def region(ax, start_marg, end_marg):
    ax.axvspan(start_marg, end_marg, color='#a1d99b')

X = {}
Y = {}
Z = {}
Quat = {}
n = len(USERS)
colors = np.linspace(0.3, 1, n).tolist()

for method in ["none", "GUI", "local", "global"]:
    fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(12,4))
    fig.suptitle(method.upper())

    x = {}
    y = {}
    z = {}
    quat = {}

    for user_n in USERS:
        filename = "data/demos/user_" + str(user_n) + "_" + method +  ".pkl"
        file = open(filename, "rb")
        data = pickle.load(file)

        ee = data["ee positions"]
        points = np.array(ee).T

        R = data["rotation matrix"]
        quat_diff = []
        for r in R:
            quat_diff.append(Quaternion.absolute_distance(Panda.rot2quat(R_desire),
                                                          Panda.rot2quat(r)))

        x["user_" + str(user_n)] = points[0,:] - points[0,0]
        y["user_" + str(user_n)] = points[1,:]
        z["user_" + str(user_n)] = points[2,:]
        quat["user_" + str(user_n)] = quat_diff

    X[method] = x
    Y[method] = y
    Z[method] = z
    Quat[method] = quat

    ### plot user demonstrations ###
    for idx, user_n in enumerate(USERS):

        # distance
        ax1.plot(X[method]["user_" + str(user_n)], Y[method]["user_" + str(user_n)],
                            label="user_" + str(user_n), c=str(colors[idx]))
        ax1.set_xlabel('X [m]')
        ax1.set_ylabel('Y [m]')
        ax1.set_ylim(-0.5, 0.5)
        ax1.title.set_text('Distance From User')
        if method == "GUI":
            region(ax1, 0.8, 1.1)
        elif method == "global":
            region(ax1, 0.5, 0.8)
        elif method == "local":
            region(ax1, 0.8, 1.1)

        # height
        ax2.plot(X[method]["user_" + str(user_n)], Z[method]["user_" + str(user_n)],
                            c=str(colors[idx]))
        ax2.set_xlabel('X [m]')
        ax2.set_ylabel('Z [m]')
        ax2.set_ylim(0, 0.8)
        ax2.title.set_text('Height From Table')
        if method == "GUI":
            region(ax2, 0.2, 0.5)
        elif method == "global":
            region(ax2, 0.2, 0.5)
        elif method == "local":
            region(ax2, 0.5, 0.8)

        # orientation
        ax3.plot(X[method]["user_" + str(user_n)], Quat[method]["user_" + str(user_n)],
                            c=str(colors[idx]))
        ax3.set_xlabel('X [m]')
        ax3.set_ylabel('Quaternion')
        ax3.set_ylim(0, 1.5)
        ax3.title.set_text('End-Effector Orientation')
        if method == "GUI":
            region(ax3, 0.5, 0.8)
        elif method == "global":
            region(ax3, 0.8, 1.1)
        elif method == "local":
            region(ax3, 0.2, 0.5)


    plt.tight_layout()
    plt.savefig("results_plot/" + method + ".png")