goal plotting

Author: FedeClaudi

proj/environment/manager.py

@@ -111,10 +111,15 @@ def _save_results(self):
         )
 
         # save the last frame as a results image
-        last_frame = [
-            f for f in self.frames_folder.glob("*.png") if f.is_file()
-        ][-1]
-        shutil.copy(str(last_frame), str(self.datafolder / "final_frame.png"))
+        try:
+            last_frame = [
+                f for f in self.frames_folder.glob("*.png") if f.is_file()
+            ][-1]
+            shutil.copy(
+                str(last_frame), str(self.datafolder / "final_frame.png")
+            )
+        except IndexError:
+            pass  # no frames were saved
 
         # save cost history
         pd.DataFrame(self.cost_history).to_hdf(
@@ -149,7 +154,7 @@ def conclude(self):
         self._log_conf()
         self._save_results()
 
-        if self.model.PLOT_LIVE:
+        if self.model.LIVE_PLOT:
             self._save_video()
 
         # save summary plot

proj/environment/plotter.py

@@ -238,11 +238,26 @@ def _plot_cost(self, keep_s=1.2):
                 )
         ax.legend()
 
-    # def _plot_goal(self, goal):
-    #     goal = self.model._goal(*goal)
-    #     x = self.model.curr_x
+    def _plot_goal(self, goal):
+        self.goal_ax.clear()
+        goal = self.model._state(*goal)
+        x = self.model.curr_x
+
+        for n, k in enumerate(x._fields):
+            self.goal_ax.bar(n, goal._asdict()[k], color=colors[k], alpha=0.7)
+            self.goal_ax.scatter(
+                n,
+                x._asdict()[k],
+                color=desaturate_color(colors[k]),
+                s=200,
+                zorder=99,
+                lw=1,
+                edgecolors="white",
+                label=k,
+            )
 
-    #     for k in self.model.curr_x._fields:
+        self.goal_ax.legend()
+        self.goal_ax.set(xticks=[])
 
     def visualize_world_live(self, curr_goals, elapsed=None):
         ax = self.xy_ax
@@ -289,6 +304,9 @@ def visualize_world_live(self, curr_goals, elapsed=None):
         # plot cost
         self._plot_cost()
 
+        # plot goal
+        self._plot_goal(curr_goals[0, :])
+
         # display plot
         self.f.canvas.draw()
         plt.pause(0.00001)

proj/model/config.py

@@ -6,7 +6,7 @@ class Config:
 
     USE_FAST = True  # if true use cumba's methods
     SPAWN_TYPE = "trajectory"
-    LIVE_PLOT = False
+    LIVE_PLOT = True
 
     # ----------------------------- Simulation params ---------------------------- #
     dt = 0.005