Share view computation among all policies

active_grasp/baselines.py

@@ -1,7 +1,6 @@
 import numpy as np
 
-from .policy import SingleViewPolicy, MultiViewPolicy
+from .policy import SingleViewPolicy, MultiViewPolicy, compute_error
-from vgn.utils import look_at
 
 
 class InitialView(SingleViewPolicy):
@@ -14,23 +13,19 @@ class InitialView(SingleViewPolicy):
 class TopView(SingleViewPolicy):
     def activate(self, bbox):
         super().activate(bbox)
-        eye = np.r_[self.center[:2], self.bbox.max[2] + self.min_z_dist]
+        self.x_d = self.view_sphere.get_view(0.0, 0.0)
-        up = np.r_[1.0, 0.0, 0.0]
-        self.x_d = look_at(eye, self.center, up)
         self.done = False if self.is_view_feasible(self.x_d) else True
 
 
 class TopTrajectory(MultiViewPolicy):
     def activate(self, bbox):
         super().activate(bbox)
-        eye = np.r_[self.center[:2], self.bbox.max[2] + self.min_z_dist]
+        self.x_d = self.view_sphere.get_view(0.0, 0.0)
-        up = np.r_[1.0, 0.0, 0.0]
-        self.x_d = look_at(eye, self.center, up)
         self.done = False if self.is_view_feasible(self.x_d) else True
 
     def update(self, img, x):
         self.integrate(img, x)
-        linear, angular = self.compute_error(self.x_d, x)
+        linear, angular = compute_error(self.x_d, x)
         if np.linalg.norm(linear) < 0.02:
             self.done = True
             return np.zeros(6)

active_grasp/bbox.py

@@ -9,6 +9,7 @@ class AABBox:
     def __init__(self, bbox_min, bbox_max):
         self.min = np.asarray(bbox_min)
         self.max = np.asarray(bbox_max)
+        self.center = 0.5 * (self.min + self.max)
 
     @property
     def corners(self):

active_grasp/nbv.py

@@ -2,8 +2,8 @@ import itertools
 import numpy as np
 import rospy
 
-from .policy import MultiViewPolicy
+from .policy import MultiViewPolicy, compute_error
-from vgn.utils import look_at, spherical_to_cartesian
+from vgn.utils import look_at
 
 
 class NextBestView(MultiViewPolicy):
@@ -16,14 +16,6 @@ class NextBestView(MultiViewPolicy):
     def activate(self, bbox):
         super().activate(bbox)
         self.generate_view_candidates()
-        # self.vis.views(
-        #     self.base_frame,
-        #     self.intrinsic,
-        #     self.view_candidates,
-        #     np.ones(len(self.view_candidates)),
-        # )
-        # rospy.sleep(1.0)
-        # return
 
     def update(self, img, x):
         if len(self.views) > self.max_views:
@@ -38,7 +30,7 @@ class NextBestView(MultiViewPolicy):
             self.vis.views(self.base_frame, self.intrinsic, views, utilities)
             i = np.argmax(utilities)
             nbv, ig = views[i], gains[i]
-            cmd = self.compute_velocity_cmd(*self.compute_error(nbv, x))
+            cmd = self.compute_velocity_cmd(*compute_error(nbv, x))
             if self.best_grasp:
                 R, t = self.best_grasp.pose.rotation, self.best_grasp.pose.translation
                 if np.linalg.norm(t - x.translation) < self.min_z_dist:
@@ -49,20 +41,15 @@ class NextBestView(MultiViewPolicy):
                 eye = R.apply([0.0, 0.0, -0.2]) + t
                 up = np.r_[1.0, 0.0, 0.0]
                 x_d = look_at(eye, center, up)
-                cmd = self.compute_velocity_cmd(*self.compute_error(x_d, x))
+                cmd = self.compute_velocity_cmd(*compute_error(x_d, x))
             return cmd
 
     def generate_view_candidates(self):
-        center = np.r_[self.center[:2], self.bbox.max[2]]
-        r = self.min_z_dist
         thetas = np.arange(1, 4) * np.deg2rad(30)
         phis = np.arange(1, 6) * np.deg2rad(60)
         self.view_candidates = []
         for theta, phi in itertools.product(thetas, phis):
-            eye = center + spherical_to_cartesian(r, theta, phi)
+            view = self.view_sphere.get_view(theta, phi)
-            target = self.center
-            up = np.r_[1.0, 0.0, 0.0]
-            view = look_at(eye, target, up)
             if self.is_view_feasible(view):
                 self.view_candidates.append(view)
 

active_grasp/policy.py

@@ -8,7 +8,7 @@ from robot_helpers.ros import tf
 from robot_helpers.ros.conversions import *
 from vgn.detection import *
 from vgn.perception import UniformTSDFVolume
-from vgn.utils import *
+from vgn.utils import look_at, spherical_to_cartesian
 
 
 class Policy:
@@ -37,26 +37,31 @@ class Policy:
         self.T_cam_ee = tf.lookup(self.cam_frame, "panda_link8")
 
     def activate(self, bbox):
-        self.bbox = bbox
         self.vis.clear()
+
+        self.bbox = bbox
+        self.vis.bbox(self.base_frame, self.bbox)
+
+        self.view_sphere = ViewSphere(bbox)
+
         self.calibrate_task_frame()
-        self.vis.bbox(self.base_frame, bbox)
+
-        self.vis.workspace(self.task_frame, 0.3)
         self.tsdf = UniformTSDFVolume(0.3, 40)
         self.vgn = VGN(Path(rospy.get_param("vgn/model")))
+
         self.views = []
         self.best_grasp = None
         self.done = False
 
     def calibrate_task_frame(self):
-        self.center = 0.5 * (self.bbox.min + self.bbox.max)
+        self.T_base_task = Transform.translation(self.bbox.center - np.full(3, 0.15))
-        self.T_base_task = Transform.translation(self.center - np.full(3, 0.15))
         self.T_task_base = self.T_base_task.inv()
         tf.broadcast(self.T_base_task, self.base_frame, self.task_frame)
-        rospy.sleep(0.5)
+        rospy.sleep(0.5)  # Wait for tf tree to be updated.
+        self.vis.workspace(self.task_frame, 0.3)
 
-    def score_fn(self, grasp):
+    def update(self, img, pose):
-        return grasp.quality
+        raise NotImplementedError
 
     def sort_grasps(self, in_grasps):
         # Transforms grasps into base frame, checks whether they lie on the target, and sorts by their score
@@ -81,19 +86,14 @@ class Policy:
         indices = np.argsort(-scores)
         return grasps[indices], scores[indices]
 
-    def update(self, img, pose):
+    def score_fn(self, grasp):
-        raise NotImplementedError
+        return grasp.quality
 
     def is_view_feasible(self, view):
         # Check whether MoveIt can find a trajectory to the given view
         success, _ = self.moveit.plan(view * self.T_cam_ee)
         return success
 
-    def compute_error(self, x_d, x):
-        linear = x_d.translation - x.translation
-        angular = (x_d.rotation * x.rotation.inv()).as_rotvec()
-        return linear, angular
-
     def compute_velocity_cmd(self, linear, angular):
         kp = 4.0
         linear = kp * linear
@@ -104,7 +104,7 @@ class Policy:
 
 class SingleViewPolicy(Policy):
     def update(self, img, x):
-        linear, angular = self.compute_error(self.x_d, x)
+        linear, angular = compute_error(self.x_d, x)
 
         if np.linalg.norm(linear) < 0.02:
             self.views.append(x)
@@ -152,6 +152,27 @@ class MultiViewPolicy(Policy):
         self.vis.grasps(self.base_frame, grasps, scores)
 
 
+class ViewSphere:
+    # Define sphere for view generation on top of the bbox.
+    # TODO an ellipse around the bbox's center would be even nicer ;)
+
+    def __init__(self, bbox):
+        self.center = np.r_[bbox.center[:2], bbox.max[2]]
+        self.r = rospy.get_param("~camera/min_z_dist")
+        self.target = bbox.center
+
+    def get_view(self, theta, phi):
+        eye = self.center + spherical_to_cartesian(self.r, theta, phi)
+        up = np.r_[1.0, 0.0, 0.0]
+        return look_at(eye, self.target, up)
+
+
+def compute_error(x_d, x):
+    linear = x_d.translation - x.translation
+    angular = (x_d.rotation * x.rotation.inv()).as_rotvec()
+    return linear, angular
+
+
 registry = {}
 
 

cfg/scenes/test1.yaml

@@ -1,4 +1,4 @@
-center: [0.5, 0.0, 0.20]
+center: [0.5, 0.1, 0.20]
 q:  [0.0, -1.39, 0.0, -2.36, 0.0, 1.57, 0.79]
 objects:
 - object_id: 006_mustard_bottle