Finish plotting for thesis + making video demos

Signed-off-by: An Thai Le <an.thai.le97@gmail.com>
ASUDCT · Sep 20, 2021 · 209adbc · 209adbc
1 parent 4f042b6
commit 209adbc
Show file tree

Hide file tree

Showing 9 changed files with 474 additions and 70 deletions.
diff --git a/scripts/experiments/plot_adaptation_results.py b/scripts/experiments/plot_adaptation_results.py
@@ -0,0 +1,112 @@
+import sys
+from os.path import join, dirname, abspath
+import matplotlib
+import logging
+import numpy as np
+import matplotlib.pyplot as plt
+logging.basicConfig()
+logging.getLogger().setLevel(logging.INFO)
+matplotlib.rcParams['pdf.fonttype'] = 42
+matplotlib.rcParams['ps.fonttype'] = 42
+matplotlib.rcParams['font.size'] = 16
+
+ROOT_DIR = join(dirname(abspath(__file__)), '..', '..')
+sys.path.append(ROOT_DIR)
+from tprmp.utils.loading import load_demos_2d  # noqa
+from tprmp.models.tp_rmp import TPRMP  # noqa
+from tprmp.models.rmp import compute_riemannian_metric  # noqa
+from tprmp.demonstrations.base import Demonstration  # noqa
+from tprmp.visualization.demonstration import plot_frame_2d  # noqa
+from tprmp.demonstrations.manifold import Manifold  # noqa
+from tprmp.demonstrations.frame import Frame  # noqa
+
+mode = 1
+task = 'test'
+demo_path = join(ROOT_DIR, 'data', 'tasks', task, 'demos')
+num_comp = 9
+moving_goal_radius = 0.5
+omega = np.pi
+disturb = True
+disturb_period = [50, 150]
+disturb_magnitude = 10.
+N1, N2 = 3, 5
+limits = [0., 4.5]
+demo_names = ['C', 'C1', 'G', 'hat', 'hat1', 'I', 'I1', 'J', 'L', 'L1', 'P', 'S', 'S1', 'S2', 'U']
+dt = 0.01
+v_eps = 8e-2
+goal_eps = 0.2
+wait = 10
+res = 0.05
+max_steps = 2000
+colormap = 'RdBu' if mode == 1 else 'YlOrBr'
+manifold = Manifold.get_euclidean_manifold(2)
+
+
+def execute(model, start_frame, x0, dx0, origin):
+    x, dx = x0, dx0
+    traj = [x]
+    moving = True
+    t = 0
+    while t < max_steps:
+        if moving:
+            end_pose = origin + moving_goal_radius * np.array([np.cos(omega * t * dt), np.sin(omega * t * dt)])
+            A, b = Demonstration.construct_linear_map(manifold, end_pose)
+            end_frame = Frame(A, b, manifold=manifold)
+            frames = {'start': start_frame, 'end': end_frame}
+        ddx = model.retrieve(x, dx, frames=frames)
+        if disturb and (t >= disturb_period[0] and t <= disturb_period[1]):
+            M = compute_riemannian_metric(x, model._global_mvns, mass_scale=model._mass_scale)
+            v = dx / np.linalg.norm(dx)
+            df = disturb_magnitude * np.array([v[1], v[0]])
+            ddx += np.linalg.inv(M) @ df
+        dx = ddx * dt + dx
+        x = manifold.exp_map(dx * dt, base=x)
+        traj.append(x)
+        goal = model._global_mvns[-1].mean
+        d = np.linalg.norm(manifold.log_map(x, base=goal))
+        if d < goal_eps:
+            moving = False
+            if np.linalg.norm(dx) < v_eps:
+                break
+        t += 1
+    return np.array(traj).T, frames
+
+
+X, Y = np.meshgrid(np.arange(limits[0], limits[1], res), np.arange(limits[0], limits[1], res))
+fig, axs = plt.subplots(N1, N2)
+for n, name in enumerate(demo_names):
+    k, m = int(n / N2), n % N2
+    data_file = join(demo_path, name + '.p')
+    model_file = name + '_' + str(num_comp) + '.p'
+    demos = load_demos_2d(data_file, dt=dt)
+    model = TPRMP.load(task, model_name=model_file)
+    sample = demos[0]
+    frames = sample.get_task_parameters()
+    x0, dx0 = sample.traj[:, 0], np.zeros(2)
+    origin = sample.traj[:, -1]
+    traj, frames = execute(model, frames['start'], x0, dx0, origin)
+    ax = axs[k, m]
+    plt.sca(ax)
+    ax.set_aspect('equal')
+    ax.set_xlim([limits[0], limits[1]])
+    ax.set_ylim([limits[0], limits[1]])
+    ax.plot(traj[0], traj[1], color="b")
+    plot_frame_2d(frames.values())
+    Z = np.zeros_like(X)
+    for i in range(X.shape[0]):
+        for j in range(X.shape[1]):
+            if mode == 1:
+                Z[i, j] = model.compute_potential_field(np.array([X[i, j], Y[i, j]]))
+            else:
+                Z[i, j] = model.compute_dissipation_field(np.array([X[i, j], Y[i, j]]))
+    c = ax.pcolormesh(X, Y, Z, cmap=colormap, shading='auto', vmin=0., vmax=Z.max(), alpha=0.5)
+    fig.colorbar(c, ax=ax)
+    circle = plt.Circle(origin, moving_goal_radius, color='k', fill=False, linestyle='--', alpha=0.7)
+    ax.add_patch(circle)
+
+for ax in fig.axes:
+    try:
+        ax.label_outer()
+    except:  # noqa
+        pass
+plt.show()
diff --git a/scripts/experiments/plot_adaptation_rmpflow_results.py b/scripts/experiments/plot_adaptation_rmpflow_results.py
@@ -0,0 +1,124 @@
+import sys
+from os.path import join, dirname, abspath
+import matplotlib
+import logging
+import numpy as np
+import matplotlib.pyplot as plt
+logging.basicConfig()
+logging.getLogger().setLevel(logging.INFO)
+matplotlib.rcParams['pdf.fonttype'] = 42
+matplotlib.rcParams['ps.fonttype'] = 42
+matplotlib.rcParams['font.size'] = 16
+
+ROOT_DIR = join(dirname(abspath(__file__)), '..', '..')
+sys.path.append(ROOT_DIR)
+from tprmp.utils.loading import load_demos_2d  # noqa
+from tprmp.models.tp_rmp import TPRMP  # noqa
+from tprmp.models.rmp import compute_riemannian_metric  # noqa
+from tprmp.demonstrations.base import Demonstration  # noqa
+from tprmp.models.rmp_tree import RMPLeaf, RMPNode, RMPRoot  # noqa
+from tprmp.models.rmp_models import CollisionAvoidance  # noqa
+from tprmp.visualization.demonstration import plot_frame_2d  # noqa
+from tprmp.demonstrations.manifold import Manifold  # noqa
+from tprmp.demonstrations.frame import Frame  # noqa
+
+mode = 1
+task = 'test'
+demo_path = join(ROOT_DIR, 'data', 'tasks', task, 'demos')
+num_comp = 9
+R = 0.2
+moving_goal_radius = 0.5
+omega = np.pi
+disturb = False
+disturb_period = [50, 150]
+disturb_magnitude = 10.
+N1, N2 = 2, 2
+limits = [0., 4.5]
+demo_names = ['C1', 'L1', 'P', 'U']
+dt = 0.01
+v_eps = 8e-2
+goal_eps = 0.2
+wait = 10
+max_steps = 2100
+res = 0.05
+colormap = 'RdBu' if mode == 1 else 'YlOrBr'
+manifold = Manifold.get_euclidean_manifold(2)
+
+
+def execute(model, root, start_frame, x0, dx0, origin):
+    x, dx = x0, dx0
+    traj = [x]
+    moving = True
+    t = 0
+    while t < max_steps:
+        if moving:
+            end_pose = origin + moving_goal_radius * np.array([np.cos(omega * t * dt), np.sin(omega * t * dt)])
+            A, b = Demonstration.construct_linear_map(manifold, end_pose)
+            end_frame = Frame(A, b, manifold=manifold)
+            frames = {'start': start_frame, 'end': end_frame}
+        model.generate_global_gmm(frames)
+        ddx = root.solve(x, dx)
+        if disturb and (t >= disturb_period[0] and t <= disturb_period[1]):
+            M = compute_riemannian_metric(x, model._global_mvns, mass_scale=model._mass_scale)
+            v = dx / np.linalg.norm(dx)
+            df = disturb_magnitude * np.array([v[1], v[0]])
+            ddx += np.linalg.inv(M) @ df
+        dx = ddx * dt + dx
+        x = manifold.exp_map(dx * dt, base=x)
+        traj.append(x)
+        goal = model._global_mvns[-1].mean
+        d = np.linalg.norm(manifold.log_map(x, base=goal))
+        if d < goal_eps:
+            moving = False
+            if np.linalg.norm(dx) < v_eps:
+                break
+        t += 1
+    return np.array(traj).T, frames
+
+
+X, Y = np.meshgrid(np.arange(limits[0], limits[1], res), np.arange(limits[0], limits[1], res))
+fig, axs = plt.subplots(N1, N2)
+for n, name in enumerate(demo_names):
+    k, m = int(n / N2), n % N2
+    data_file = join(demo_path, name + '.p')
+    model_file = name + '_' + str(num_comp) + '.p'
+    demos = load_demos_2d(data_file, dt=dt)
+    model = TPRMP.load(task, model_name=model_file)
+    sample = demos[0]
+    frames = sample.get_task_parameters()
+    x0, dx0 = sample.traj[:, 0], np.zeros(2)
+    origin = sample.traj[:, -1]
+    # init rmpflow
+    T = sample.traj.shape[1]
+    cpose = sample.traj[:, int(T / 2)]
+    root = RMPRoot('root_space', manifold=manifold)
+    ca_node = CollisionAvoidance('CA_space', parent=root, c=cpose, R=R)
+    tprmp_node = RMPLeaf('TPRMP_space', model.rmp, parent=root, manifold=manifold, psi=lambda x: x, J=lambda x: np.eye(2))
+    traj, frames = execute(model, root, frames['start'], x0, dx0, origin)
+    ax = axs[k, m]
+    plt.sca(ax)
+    ax.set_aspect('equal')
+    ax.set_xlim([limits[0], limits[1]])
+    ax.set_ylim([limits[0], limits[1]])
+    ax.plot(traj[0], traj[1], color="b")
+    plot_frame_2d(frames.values())
+    Z = np.zeros_like(X)
+    for i in range(X.shape[0]):
+        for j in range(X.shape[1]):
+            if mode == 1:
+                Z[i, j] = model.compute_potential_field(np.array([X[i, j], Y[i, j]]))
+            else:
+                Z[i, j] = model.compute_dissipation_field(np.array([X[i, j], Y[i, j]]))
+    c = ax.pcolormesh(X, Y, Z, cmap=colormap, shading='auto', vmin=0., vmax=Z.max(), alpha=0.5)
+    fig.colorbar(c, ax=ax)
+    mcircle = plt.Circle(origin, moving_goal_radius, color='k', fill=False, linestyle='--', alpha=0.7)
+    ax.add_patch(mcircle)
+    ocircle = plt.Circle(cpose, R, color='k', fill=False)
+    ax.add_patch(ocircle)
+
+for ax in fig.axes:
+    try:
+        ax.label_outer()
+    except:  # noqa
+        pass
+plt.show()
diff --git a/scripts/plot_tprmp_2d_moving.py b/scripts/plot_tprmp_2d_moving.py
@@ -2,28 +2,42 @@
 import argparse
 import numpy as np
 import matplotlib.pyplot as plt
-from os.path import join, dirname, abspath
+from os.path import join, dirname, abspath, expanduser
+from matplotlib.animation import FFMpegWriter
+import matplotlib
 import logging
+matplotlib.rcParams['pdf.fonttype'] = 42
+matplotlib.rcParams['ps.fonttype'] = 42
+matplotlib.rcParams['font.size'] = 16
 logging.basicConfig()
 logging.getLogger().setLevel(logging.INFO)
 
 ROOT_DIR = join(dirname(abspath(__file__)), '..')
 sys.path.append(ROOT_DIR)
 from tprmp.utils.loading import load_demos_2d  # noqa
-from tprmp.visualization.demonstration import plot_demo  # noqa
-from tprmp.visualization.dynamics import plot_dissipation_field, plot_potential_field, visualize_rmp # noqa
+from tprmp.visualization.dynamics import plot_dissipation_field, plot_potential_field # noqa
 from tprmp.models.tp_rmp import TPRMP  # noqa
 from tprmp.demonstrations.base import Demonstration  # noqa
 from tprmp.demonstrations.frame import Frame  # noqa
 
 parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                  description='Example run: python test_tprmp.py test.p')
 parser.add_argument('--task', help='The task folder', type=str, default='test')
-parser.add_argument('--data', help='The data file', type=str, default='test3.p')
+parser.add_argument('--demo', help='The data file', type=str, default='U.p')
+parser.add_argument('--data', help='The data file', type=str, default='U_5.p')
 parser.add_argument('--plot', help='The data file', type=str, default='dissipative')
 args = parser.parse_args()
+
+DATA_DIR = join(ROOT_DIR, 'data', 'tasks', args.task, 'demos')
+data_file = join(DATA_DIR, args.demo)
+fps = 30
+video_file = join(expanduser("~"), 'U.mp4')
+metadata = dict(artist='Matplotlib')
+writer = FFMpegWriter(fps=fps, metadata=metadata)
 # parameters
-T = 200
+T = 100
+a = 0.5
+omega = np.pi
 dt = 0.01
 res = 0.05
 max_z = 1000
@@ -33,21 +47,26 @@
 limits = [0., 4.]
 model = TPRMP.load(args.task, model_name=args.data)
 manifold = model.model.manifold
-start_pose = np.array([1.4, 1.1])
-end_pose = np.array([3.1, 2.6])
-plt.figure()
+demos = load_demos_2d(data_file, dt=dt)
+sample = demos[0]
+start_pose = sample.traj[:, 0]
+end_pose = sample.traj[:, -1]
+fig = plt.figure()
+writer.setup(fig, outfile=video_file, dpi=None)
 for t in range(T):
-    start_pose += start_v * dt
-    end_pose += end_v * dt
-    A, b = Demonstration.construct_linear_map(manifold, start_pose)
+    start_p = start_pose + np.array([a * np.cos(omega * t * dt), 0.])
+    end_p = end_pose + np.array([0., a * np.cos(omega * t * dt)])
+    A, b = Demonstration.construct_linear_map(manifold, start_p)
     start_frame = Frame(A, b, manifold=manifold)
-    A, b = Demonstration.construct_linear_map(manifold, end_pose)
+    A, b = Demonstration.construct_linear_map(manifold, end_p)
     end_frame = Frame(A, b, manifold=manifold)
     frames = {'start': start_frame, 'end': end_frame}
     if args.plot == 'potential':
-        plot_potential_field(model, frames, margin=margin, max_z=max_z, three_d=True, res=res, limits=limits, new_fig=False, show=False)
+        plot_potential_field(model, frames, margin=margin, max_z=max_z, three_d=False, res=res, limits=limits, new_fig=False, show=False)
     elif args.plot == 'dissipative':
         plot_dissipation_field(model, frames, margin=margin, res=res, limits=limits, new_fig=False, show=False)
     plt.draw()
+    writer.grab_frame()
     plt.pause(0.00001)
     plt.cla()
+writer.finish()
diff --git a/scripts/record_pick_demo.py b/scripts/record_pick_demo.py
@@ -11,7 +11,7 @@
 sys.path.append(ROOT_DIR)
 
 from tprmp.utils.loading import save_demos # noqa
-from tprmp.envs.tasks import PalletizingBoxes # noqa
+from tprmp.envs.tasks import PickBox # noqa
 from tprmp.envs.gym import Environment  # noqa
 from tprmp.demonstrations.base import Demonstration  # noqa
 from tprmp.demonstrations.frame import Frame  # noqa
@@ -31,14 +31,14 @@
 wait = 0.
 diff = 10
 direct = False
-env = Environment(task=PalletizingBoxes(), disp=True, real_time_step=True, sampling_hz=sampling_hz)
+env = Environment(task=PickBox(), disp=True, real_time_step=True, sampling_hz=sampling_hz)
 
 save_path = join(ROOT_DIR, 'data', 'tasks', 'pick', 'demos')
 os.makedirs(save_path, exist_ok=True)
 save_name = join(save_path, args.n)
 manifold = Manifold.get_manifold_from_name('R^3 x S^3')
-box_id = env.task.goals[0][0][0]
-r = 0.02
+box_id = env.task.box_id
+r = 0.1
 trajs = []
 traj_vels = []
 tags = []
@@ -81,7 +81,7 @@
     ee_frames.append(ee_start)
     num_demo += 1
 # grasp side
-pose_mean = np.array([0.5, -0.25, env.task.box_size[0] / 2])
+pose_mean = np.array([0.5, -0.25, env.task.box_size[1] / 2])
 num_demo = 0
 len_demo = 0
 while True:
@@ -97,7 +97,7 @@
     obj_frame = Frame(A, b, manifold)
     pose = obj_frame.transform(np.append([0., 0., 0.], q_from_euler(np.array([-np.pi/2, 0., 0.]))))
     pose[3:] = q_convert_xyzw(pose[3:])
-    pose[2] += env.task.box_size[0] / 2 + margin
+    pose[2] += env.task.box_size[1] / 2 + margin
     for _ in range(100):  # stabilize arm
         p.stepSimulation()
     input()