Fix plotting

Signed-off-by: An Thai Le <an.thai.le97@gmail.com>
ASUDCT · Jul 16, 2021 · 37b523b · 37b523b
1 parent d4b6f14
commit 37b523b
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 3 deletions.
diff --git a/scripts/plot_demos.py b/scripts/plot_demos.py
@@ -16,4 +16,4 @@
 DATA_DIR = join(ROOT_DIR, 'data', 'tasks', args.task, 'demos')
 data_file = join(DATA_DIR, args.data)
 demos = load_demos(data_file)
-plot_demo(demos, only_global=False, plot_quat=False, new_fig=True, show=True)
+plot_demo(demos, only_global=False, plot_quat=False, new_fig=True, new_ax=True, show=True)
diff --git a/tprmp/models/tp_hsmm.py b/tprmp/models/tp_hsmm.py
@@ -111,7 +111,7 @@ def plot_model(self, demos, tagging=True, plot_quat=False, show=True):
         tag_to_demos = TPHSMM.compute_tag_to_demo(demos)
         for tag in tag_to_demos:  # plot one demo for each tag
             demo = tag_to_demos[tag][0]
-            plot_demo(demo, plot_quat=plot_quat, new_fig=True, show=False)
+            plot_demo(demo, plot_quat=plot_quat, new_fig=True, new_ax=True, show=False)
             plot_gmm(self, demo.get_task_parameters(), plot_quat=plot_quat, tag=tag if tagging else None, new_fig=False, show=show)
 
     def set_params(self, model_params):

diff --git a/tprmp/visualization/demonstration.py b/tprmp/visualization/demonstration.py
@@ -29,7 +29,11 @@ def _plot_traj_global(demos, **kwargs):
     legend = kwargs.get('legend', True)
     plot_frames = kwargs.get('plot_frames', True)
     title = kwargs.get('title', 'Global frame')
-    ax = plt.gca()
+    new_ax = kwargs.get('new_ax', False)
+    if new_ax:
+        ax = plt.subplot(111, projection="3d")
+    else:
+        ax = plt.gca()
     plt.title(title)
     demo_tags = list(set([demo.tag for demo in demos]))
     tag_map = {v: i for i, v in enumerate(demo_tags)}