Compute mean over recent grasp predictions

active_grasp/policy.py

@@ -2,11 +2,12 @@ import numpy as np
 from sensor_msgs.msg import CameraInfo
 from pathlib import Path
 import rospy
+import warnings
 
 from .visualization import Visualizer
 from robot_helpers.ros import tf
 from robot_helpers.ros.conversions import *
-from vgn.detection import VGN, compute_grasps
+from vgn.detection import *
 from vgn.perception import UniformTSDFVolume
 from vgn.utils import *
 
@@ -20,8 +21,9 @@ class Policy:
 
 
 class BasePolicy(Policy):
-    def __init__(self, rate=5):
+    def __init__(self, rate=5, filter_grasps=False):
         self.rate = rate
+        self.filter_grasps = filter_grasps
         self.load_parameters()
         self.init_visualizer()
 
@@ -32,6 +34,7 @@ class BasePolicy(Policy):
         msg = rospy.wait_for_message(info_topic, CameraInfo, rospy.Duration(2.0))
         self.intrinsic = from_camera_info_msg(msg)
         self.vgn = VGN(Path(rospy.get_param("vgn/model")))
+        self.score_fn = lambda g: g.pose.translation[2]
 
     def init_visualizer(self):
         self.visualizer = Visualizer(self.base_frame)
@@ -45,7 +48,15 @@ class BasePolicy(Policy):
         tf.broadcast(self.T_base_task, self.base_frame, self.task_frame)
         rospy.sleep(0.1)  # wait for the transform to be published
 
-        self.tsdf = UniformTSDFVolume(0.3, 40)
+        N, self.T = 40, 10  # spatial and temporal resolution
+        grid_shape = (N,) * 3
+
+        self.tsdf = UniformTSDFVolume(0.3, N)
+
+        self.qual_hist = np.zeros((self.T,) + grid_shape, np.float32)
+        self.rot_hist = np.zeros((self.T, 4) + grid_shape, np.float32)
+        self.width_hist = np.zeros((self.T,) + grid_shape, np.float32)
+
         self.viewpoints = []
         self.done = False
         self.best_grasp = None
@@ -59,16 +70,48 @@ class BasePolicy(Policy):
         self.visualizer.scene_cloud(self.task_frame, self.tsdf.get_scene_cloud())
         self.visualizer.path(self.viewpoints)
 
-    def compute_best_grasp(self):
-        return self.predict_best_grasp()
+        if self.filter_grasps:
+            tsdf_grid = self.tsdf.get_grid()
+            out = self.vgn.predict(tsdf_grid)
+            t = (len(self.viewpoints) - 1) % self.T
+            self.qual_hist[t, ...] = out.qual
+            self.rot_hist[t, ...] = out.rot
+            self.width_hist[t, ...] = out.width
 
-    def predict_best_grasp(self):
-        tsdf_grid = self.tsdf.get_grid()
-        out = self.vgn.predict(tsdf_grid)
-        score_fn = lambda g: g.pose.translation[2]
-        grasps = compute_grasps(self.tsdf.voxel_size, out, score_fn, max_filter_size=3)
+    def compute_best_grasp(self):
+        if self.filter_grasps:
+            T = len(self.viewpoints) if len(self.viewpoints) // self.T == 0 else self.T
+            mask = self.qual_hist[:T, ...] > 0.0
+            # The next line prints a warning since some voxels have no grasp
+            # predictions resulting in empty slices.
+            qual = np.mean(self.qual_hist[:T, ...], axis=0, where=mask)
+            qual = threshold_quality(qual, 0.9)
+            index_list = select_local_maxima(qual, 3)
+
+            grasps = []
+            for (i, j, k) in index_list:
+                ts = np.flatnonzero(self.qual_hist[:T, i, j, k])
+                if len(ts) < 3:
+                    continue
+                oris = Rotation.from_quat([self.rot_hist[t, :, i, j, k] for t in ts])
+                ori = oris.mean()
+                # TODO check variance as well
+                pos = np.array([i, j, k], dtype=np.float64)
+                width = self.width_hist[ts, i, j, k].mean()
+                quality = self.qual_hist[ts, i, j, k].mean()
+                grasps.append(Grasp(Transform(ori, pos), width, quality))
+        else:
+            tsdf_grid = self.tsdf.get_grid()
+            out = self.vgn.predict(tsdf_grid)
+            qual = threshold_quality(out.qual, 0.9)
+            index_list = select_local_maxima(qual, 3)
+            grasps = [select_at(out, i) for i in index_list]
+
+        grasps = [from_voxel_coordinates(g, self.tsdf.voxel_size) for g in grasps]
         grasps = self.transform_grasps_to_base_frame(grasps)
         grasps = self.select_grasps_on_target_object(grasps)
+        grasps = sort_grasps(grasps, self.score_fn)
+
         return grasps[0] if len(grasps) > 0 else None
 
     def transform_grasps_to_base_frame(self, grasps):

scripts/run.py

@@ -16,7 +16,7 @@ def main():
     parser = create_parser()
     args = parser.parse_args()
 
-    policy = make(args.policy, args.rate)
+    policy = make(args.policy, args.rate, args.filter_grasps)
     controller = GraspController(policy)
     logger = Logger(args)
 
@@ -33,6 +33,7 @@ def create_parser():
     parser.add_argument("--runs", type=int, default=10)
     parser.add_argument("--logdir", type=Path, default="logs")
     parser.add_argument("--rate", type=int, default=5)
+    parser.add_argument("--filter-grasps", action="store_true")
     parser.add_argument("--seed", type=int, default=12)
     return parser
 
@@ -40,8 +41,14 @@ def create_parser():
 class Logger:
     def __init__(self, args):
         stamp = datetime.now().strftime("%y%m%d-%H%M%S")
-        descr = "policy={},rate={}".format(args.policy, args.rate)
-        self.path = args.logdir / (stamp + "_" + descr + ".csv")
+        name = "{}_policy={},rate={},filter-grasps={},seed={}.csv".format(
+            stamp,
+            args.policy,
+            args.rate,
+            args.filter_grasps,
+            args.seed,
+        )
+        self.path = args.logdir / name
 
     def log_run(self, info):
         df = pd.DataFrame.from_records([info])