Add C-space score function

active_grasp/nbv.py

@@ -56,7 +56,7 @@ class NextBestView(MultiViewPolicy):
     def __init__(self):
         super().__init__()
         self.min_z_dist = rospy.get_param("~camera/min_z_dist")
-        self.max_views = 40
+        self.max_views = 100
 
     def activate(self, bbox, view_sphere):
         super().activate(bbox, view_sphere)
@@ -78,7 +78,7 @@ class NextBestView(MultiViewPolicy):
             self.done = True
         else:
             with Timer("state_update"):
-                self.integrate(img, x)
+                self.integrate(img, x, q)
             views = self.view_candidates
             with Timer("ig_computation"):
                 gains = [self.ig_fn(v) for v in views]

active_grasp/policy.py

@@ -20,6 +20,7 @@ class Policy:
 
     def load_parameters(self):
         self.base_frame = rospy.get_param("~base_frame_id")
+        self.T_grasp_ee = Transform.from_list(rospy.get_param("~ee_grasp_offset")).inv()
         self.cam_frame = rospy.get_param("~camera/frame_id")
         self.task_frame = "task"
         info_topic = rospy.get_param("~camera/info_topic")
@@ -29,6 +30,7 @@ class Policy:
 
     def init_robot_model(self):
         self.model = KDLModel.from_parameter_server(self.base_frame, self.cam_frame)
+        self.ee_model = KDLModel.from_parameter_server(self.base_frame, "panda_link8")
 
     def init_visualizer(self):
         self.vis = Visualizer()
@@ -65,7 +67,7 @@ class Policy:
     def update(self, img, x, q):
         raise NotImplementedError
 
-    def sort_grasps(self, in_grasps):
+    def sort_grasps(self, in_grasps, q):
         # Transforms grasps into base frame, checks whether they lie on the target, and sorts by their score
         grasps, scores = [], []
 
@@ -80,15 +82,17 @@ class Policy:
             tip = pose.rotation.apply([0, 0, 0.05]) + pose.translation
             if self.bbox.is_inside(tip):
                 grasp.pose = pose
-                grasps.append(grasp)
+                q_grasp = self.ee_model.ik(q, pose * self.T_grasp_ee)
-                scores.append(self.score_fn(grasp))
+                if q_grasp is not None:
+                    grasps.append(grasp)
+                    scores.append(self.score_fn(grasp, q, q_grasp))
 
         grasps, scores = np.asarray(grasps), np.asarray(scores)
         indices = np.argsort(-scores)
         return grasps[indices], scores[indices]
 
-    def score_fn(self, grasp):
+    def score_fn(self, grasp, q, q_grasp):
-        return grasp.quality
+        return -np.linalg.norm(q - q_grasp)
 
 
 class SingleViewPolicy(Policy):
@@ -105,10 +109,14 @@ class SingleViewPolicy(Policy):
             self.vis.quality(self.task_frame, voxel_size, out.qual, 0.5)
 
             grasps = select_grid(voxel_size, out, threshold=self.qual_threshold)
-            grasps, scores = self.sort_grasps(grasps)
+            grasps, scores = self.sort_grasps(grasps, q)
-            self.vis.grasps(self.base_frame, grasps, scores)
+
+            if len(grasps) > 0:
+                smin, smax = np.min(scores), np.max(scores)
+                self.best_grasp = grasps[0]
+                self.vis.grasps(self.base_frame, grasps, scores, smin, smax)
+                self.vis.best_grasp(self.base_frame, grasps[0], scores[0], smin, smax)
 
-            self.best_grasp = grasps[0] if len(grasps) > 0 else None
             self.done = True
 
 
@@ -118,7 +126,7 @@ class MultiViewPolicy(Policy):
         self.T = 10  # Window size of grasp prediction history
         self.qual_hist = np.zeros((self.T,) + (40,) * 3, np.float32)
 
-    def integrate(self, img, x):
+    def integrate(self, img, x, q):
         self.views.append(x)
         self.vis.path(self.base_frame, self.views)
 
@@ -137,12 +145,13 @@ class MultiViewPolicy(Policy):
 
         with Timer("grasp_selection"):
             grasps = select_grid(voxel_size, out, threshold=self.qual_threshold)
-            grasps, scores = self.sort_grasps(grasps)
+            grasps, scores = self.sort_grasps(grasps, q)
-        self.vis.grasps(self.base_frame, grasps, scores)
 
         if len(grasps) > 0:
+            smin, smax = np.min(scores), np.max(scores)
             self.best_grasp = grasps[0]
-            self.vis.best_grasp(self.base_frame, grasps[0], scores[0])
+            self.vis.grasps(self.base_frame, grasps, scores, smin, smax)
+            self.vis.best_grasp(self.base_frame, grasps[0], scores[0], smin, smax)
 
 
 def compute_error(x_d, x):

active_grasp/visualization.py

@@ -1,4 +1,4 @@
-from geometry_msgs.msg import PoseArray
+from geometry_msgs.msg import PoseStamped
 import matplotlib.colors
 import numpy as np
 import rospy
@@ -27,6 +27,7 @@ class Visualizer:
             latch=True,
             queue_size=1,
         )
+        self.pose_pub = rospy.Publisher("pose", PoseStamped, queue_size=1)
         self.quality_pub = rospy.Publisher("quality", PointCloud2, queue_size=1)
 
     def clear(self):
@@ -76,9 +77,6 @@ class Visualizer:
         self.draw(create_grasp_markers(frame, grasp, color, "best_grasp", radius=0.006))
 
     def grasps(self, frame, grasps, scores, smin=0.9, smax=1.0, alpha=0.8):
-        if len(grasps) == 0:
-            return
-
         markers = []
         for i, (grasp, score) in enumerate(zip(grasps, scores)):
             color = cmap((score - smin) / (smax - smin))
@@ -137,6 +135,10 @@ class Visualizer:
         )
         self.draw([marker])
 
+    def pose(self, frame, pose):
+        msg = to_pose_stamped_msg(pose, frame)
+        self.pose_pub.publish(msg)
+
     def quality(self, frame, voxel_size, quality, threshold=0.9):
         points, values = grid_to_map_cloud(voxel_size, quality, threshold)
         values = (values - 0.9) / (1.0 - 0.9)  # to increase contrast