diff --git a/active_grasp/policy.py b/active_grasp/policy.py index 9fc1e7f..90810c6 100644 --- a/active_grasp/policy.py +++ b/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): diff --git a/scripts/run.py b/scripts/run.py index 9a2ccbd..0af15b3 100644 --- a/scripts/run.py +++ b/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])