Add test scenes

active_grasp/__init__.py

@@ -5,5 +5,4 @@ from .nbv import NextBestView
 register("initial-view", InitialView)
 register("top-view", TopView)
 register("top-trajectory", TopTrajectory)
-register("circular-trajectory", CircularTrajectory)
 register("nbv", NextBestView)

active_grasp/baselines.py

@@ -1,6 +1,4 @@
 import numpy as np
-import rospy
-import scipy.interpolate
 
 from .policy import SingleViewPolicy, MultiViewPolicy
 from vgn.utils import look_at
@@ -16,7 +14,7 @@ class InitialView(SingleViewPolicy):
 class TopView(SingleViewPolicy):
     def activate(self, bbox):
         super().activate(bbox)
-        eye = np.r_[self.center[:2], self.center[2] + 0.3]
+        eye = np.r_[self.center[:2], self.bbox.max[2] + self.min_z_dist]
         up = np.r_[1.0, 0.0, 0.0]
         self.x_d = look_at(eye, self.center, up)
 
@@ -24,7 +22,7 @@ class TopView(SingleViewPolicy):
 class TopTrajectory(MultiViewPolicy):
     def activate(self, bbox):
         super().activate(bbox)
-        eye = np.r_[self.center[:2], self.center[2] + 0.3]
+        eye = np.r_[self.center[:2], self.bbox.max[2] + self.min_z_dist]
         up = np.r_[1.0, 0.0, 0.0]
         self.x_d = look_at(eye, self.center, up)
 
@@ -36,30 +34,3 @@ class TopTrajectory(MultiViewPolicy):
             return np.zeros(6)
         else:
             return self.compute_velocity_cmd(linear, angular)
-
-
-class CircularTrajectory(MultiViewPolicy):
-    def __init__(self, rate):
-        super().__init__(rate)
-        self.r = 0.08
-        self.h = 0.3
-        self.duration = 2.0 * np.pi * self.r / self.linear_vel
-        self.m = scipy.interpolate.interp1d([0.0, self.duration], [np.pi, 3.0 * np.pi])
-
-    def activate(self, bbox):
-        super().activate(bbox)
-        self.tic = rospy.Time.now()
-
-    def update(self, img, x):
-        self.integrate(img, x)
-        elapsed_time = (rospy.Time.now() - self.tic).to_sec()
-        if elapsed_time > self.duration:
-            self.done = True
-            return np.zeros(6)
-        else:
-            t = self.m(elapsed_time)
-            eye = self.center + np.r_[self.r * np.cos(t), self.r * np.sin(t), self.h]
-            up = np.r_[1.0, 0.0, 0.0]
-            x_d = look_at(eye, self.center, up)
-            linear, angular = self.compute_error(x_d, x)
-            return self.compute_velocity_cmd(linear, angular)

active_grasp/policy.py

@@ -33,9 +33,10 @@ class Policy:
 
     def activate(self, bbox):
         self.bbox = bbox
-        self.calibrate_task_frame()
         self.vis.clear()
+        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 = []
@@ -47,7 +48,7 @@ class Policy:
         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.1)
+        rospy.sleep(0.5)
 
     def compute_error(self, x_d, x):
         linear = x_d.translation - x.translation
@@ -67,6 +68,13 @@ class Policy:
 
         for grasp in in_grasps:
             pose = self.T_base_task * grasp.pose
+
+            R, t = pose.rotation, pose.translation
+
+            # Filter out artifacts close to the support
+            if t[2] < self.bbox.min[2] + 0.04:
+                continue
+
             tip = pose.rotation.apply([0, 0, 0.05]) + pose.translation
             if self.bbox.is_inside(tip):
                 grasp.pose = pose
@@ -97,7 +105,7 @@ class SingleViewPolicy(Policy):
             self.vis.map_cloud(self.task_frame, self.tsdf.get_map_cloud())
 
             out = self.vgn.predict(tsdf_grid)
-            self.vis.quality(self.task_frame, voxel_size, out.qual)
+            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)
@@ -122,6 +130,7 @@ class MultiViewPolicy(Policy):
 
         tsdf_grid, voxel_size = self.tsdf.get_grid(), self.tsdf.voxel_size
         out = self.vgn.predict(tsdf_grid)
+        self.vis.quality(self.task_frame, self.tsdf.voxel_size, out.qual, 0.5)
 
         grasps = select_grid(voxel_size, out, threshold=self.qual_threshold)
         grasps, scores = self.sort_grasps(grasps)

active_grasp/simulation.py

@@ -40,15 +40,15 @@ class Simulation:
         self.arm = BtPandaArm(panda_urdf)
         self.gripper = BtPandaGripper(self.arm)
         self.model = Model(panda_urdf, self.arm.base_frame, self.arm.ee_frame)
-        self.camera = BtCamera(320, 240, 0.96, 0.2, 1.0, self.arm.uid, 11)
+        self.camera = BtCamera(320, 240, 0.96, 0.01, 1.0, self.arm.uid, 11)
 
     def seed(self, seed):
         self.rng = np.random.default_rng(seed)
 
     def reset(self):
-        self.scene.reset(rng=self.rng)
         q = self.scene.sample_initial_configuration(self.rng)
         self.set_arm_configuration(q)
+        self.scene.reset(rng=self.rng)
         uid = self.select_target()
         bbox = self.get_target_bbox(uid)
         return bbox
@@ -179,10 +179,13 @@ class CustomScene(Scene):
                 self.get_ycb_urdf_path(object["object_id"]),
                 Rotation.from_euler("xyz", object["rpy"], degrees=True),
                 self.center + np.asarray(object["xyz"]),
+                object.get("scale", 1),
             )
+        for _ in range(60):
+            p.stepSimulation()
 
     def sample_initial_configuration(self, rng):
-        return self.scene["configuration"]
+        return self.scene["q"]
 
 
 def get_scene(scene_id):

active_grasp/visualization.py

@@ -137,8 +137,8 @@ class Visualizer:
         )
         self.draw([marker])
 
-    def quality(self, frame, voxel_size, quality):
-        points, values = grid_to_map_cloud(voxel_size, quality, threshold=0.9)
+    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
         msg = to_cloud_msg(frame, points, intensities=values)
         self.quality_pub.publish(msg)
@@ -169,6 +169,28 @@ class Visualizer:
             markers.append(marker)
         self.draw(markers)
 
+    def workspace(self, frame, size):
+        scale = size * 0.005
+        pose = Transform.identity()
+        scale = [scale, 0.0, 0.0]
+        color = [0.5, 0.5, 0.5]
+        lines = [
+            ([0.0, 0.0, 0.0], [size, 0.0, 0.0]),
+            ([size, 0.0, 0.0], [size, size, 0.0]),
+            ([size, size, 0.0], [0.0, size, 0.0]),
+            ([0.0, size, 0.0], [0.0, 0.0, 0.0]),
+            ([0.0, 0.0, size], [size, 0.0, size]),
+            ([size, 0.0, size], [size, size, size]),
+            ([size, size, size], [0.0, size, size]),
+            ([0.0, size, size], [0.0, 0.0, size]),
+            ([0.0, 0.0, 0.0], [0.0, 0.0, size]),
+            ([size, 0.0, 0.0], [size, 0.0, size]),
+            ([size, size, 0.0], [size, size, size]),
+            ([0.0, size, 0.0], [0.0, size, size]),
+        ]
+        msg = create_line_list_marker(frame, pose, scale, color, lines, ns="workspace")
+        self.draw([msg])
+
 
 def create_cam_view_marker(
     frame, pose, scale, color, intrinsic, near, far, ns="", id=0

cfg/active_grasp.yaml

@@ -13,7 +13,7 @@ grasp_controller:
     frame_id: camera_optical_frame
     info_topic: /camera/depth/camera_info
     depth_topic: /camera/depth/image_raw
-    min_z_dist: 0.22
+    min_z_dist: 0.26
     
 vgn:
   model: $(find vgn)/assets/models/vgn_conv.pth

cfg/scenes/test.yaml

@@ -1,9 +0,0 @@
-center: [0.55, 0.0, 0.3]
-objects:
-- object_id: 003_cracker_box
-  rpy: [0.0, 0.0, 0.0]
-  xyz: [-0.08, 0.08, 0.002]
-- object_id: 010_potted_meat_can
-  rpy: [0.0, 0.0, 0.0]
-  xyz: [0.0, 0.0, 0.002]
-configuration: [0.0, -1.04, 0.17, -1.91, 0.0, 1.57, 0.79]

cfg/scenes/test1.yaml

@@ -0,0 +1,6 @@
+center: [0.5, 0.0, 0.20]
+q:  [0.0, -1.39, 0.0, -2.36, 0.0, 1.57, 0.79]
+objects:
+- object_id: 006_mustard_bottle
+  xyz: [0.05, 0.1, 0.0]
+  rpy: [0, 0, -50]

cfg/scenes/test2.yaml

@@ -0,0 +1,6 @@
+center: [0.5, 0.0, 0.40]
+q:  [0.0, -1.39, 0.0, -2.36, 0.0, 1.57, 0.79]
+objects:
+- object_id: 006_mustard_bottle
+  xyz: [0.05, 0.1, 0.0]
+  rpy: [0, 0, -50]

cfg/scenes/test3.yaml

@@ -0,0 +1,7 @@
+center: [0.45, 0.0, 0.20]
+q:  [0.0, -1.39, 0.0, -2.36, 0.0, 1.57, 0.79]
+objects:
+- object_id: 019_pitcher_base
+  xyz: [0, 0.1, 0]
+  rpy: [0, 0, 60]
+  scale: 1.0

cfg/scenes/test4.yaml

@@ -0,0 +1,12 @@
+center: [0.5, 0.05, 0.2]
+q:  [0.0, -1.39, 0.0, -2.36, 0.0, 1.57, 0.79]
+objects:
+- object_id: 004_sugar_box
+  xyz: [-0.04, -0.03, 0.0]
+  rpy: [0, 0, 30]
+- object_id: 008_pudding_box
+  xyz: [-0.02, -0.06, 0.19]
+  rpy: [0, 0, 45]
+- object_id: 010_potted_meat_can
+  xyz: [0.04, 0.03, 0.0]
+  rpy: [0, 0, 60]

scripts/reset.py

@@ -0,0 +1,11 @@
+import rospy
+
+from active_grasp.srv import *
+
+rospy.init_node("test")
+
+seed = rospy.ServiceProxy("seed", Seed)
+reset = rospy.ServiceProxy("reset", Reset)
+
+seed(SeedRequest(1))
+reset(ResetRequest())