Simplify grasp controller

config/active_grasp.yaml

@@ -1,5 +1,6 @@
 active_grasp:
   frame_id: task
+  length: 0.3
   base_frame_id: panda_link0
   ee_frame_id: panda_hand
   ee_grasp_offset: [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.065]

policies.py

@@ -6,16 +6,18 @@ import scipy.interpolate
 
 from geometry_msgs.msg import Pose
 from sensor_msgs.msg import Image, CameraInfo
+import std_srvs.srv
 
 from robot_tools.spatial import Rotation, Transform
 from robot_tools.ros.conversions import *
+from robot_tools.ros.panda import PandaGripperClient
 from robot_tools.ros.tf import TransformTree
 from robot_tools.perception import *
 from vgn import vis
 from vgn.detection import VGN, compute_grasps
 
 
-def get_policy(name):
+def get_controller(name):
     if name == "single-view":
         return SingleViewBaseline()
     elif name == "fixed-trajectory":
@@ -26,131 +28,156 @@ def get_policy(name):
         raise ValueError("{} policy does not exist.".format(name))
 
 
-class Policy:
+class BaseController:
     def __init__(self):
-        params = rospy.get_param("active_grasp")
+        self.frame_id = rospy.get_param("~frame_id")
+        self.length = rospy.get_param("~length")
 
-        self.frame_id = params["frame_id"]
-
-        # Robot
-        self.base_frame_id = params["base_frame_id"]
-        self.ee_frame_id = params["ee_frame_id"]
-        self.tf = TransformTree()
-        self.H_EE_G = Transform.from_list(params["ee_grasp_offset"])
-        self.target_pose_pub = rospy.Publisher("/cmd", Pose, queue_size=10)
-
-        # Camera
-        self.cam_frame_id = rospy.get_param("~camera/frame_id")
-        info_topic = rospy.get_param("~camera/info_topic")
-        depth_topic = rospy.get_param("~camera/depth_topic")
-        msg = rospy.wait_for_message(info_topic, CameraInfo, rospy.Duration(2.0))
-        self.intrinsic = from_camera_info_msg(msg)
         self.cv_bridge = cv_bridge.CvBridge()
+        self.tf = TransformTree()
+
+        self.reset_client = rospy.ServiceProxy("/reset", std_srvs.srv.Trigger)
 
-        # TSDF
         self.tsdf = UniformTSDFVolume(0.3, 40)
-
-        # VGN
         self.vgn = VGN(Path(rospy.get_param("vgn/model")))
 
-        rospy.sleep(1.0)
+        self.setup_robot_connection()
-        self.H_B_T = self.tf.lookup(self.base_frame_id, self.frame_id, rospy.Time.now())
+        self.setup_camera_connection()
+        self.lookup_transforms()
+
+    def setup_robot_connection(self):
+        self.base_frame_id = rospy.get_param("~base_frame_id")
+        self.ee_frame_id = rospy.get_param("~ee_frame_id")
+        self.ee_grasp_offset = Transform.from_list(rospy.get_param("~ee_grasp_offset"))
+        self.target_pose_pub = rospy.Publisher("/cmd", Pose, queue_size=10)
+        self.gripper = PandaGripperClient()
+
+    def setup_camera_connection(self):
+        self.cam_frame_id = rospy.get_param("~camera/frame_id")
+        info_topic = rospy.get_param("~camera/info_topic")
+        msg = rospy.wait_for_message(info_topic, CameraInfo, rospy.Duration(2.0))
+        self.intrinsic = from_camera_info_msg(msg)
+        depth_topic = rospy.get_param("~camera/depth_topic")
         rospy.Subscriber(depth_topic, Image, self.sensor_cb, queue_size=1)
 
     def sensor_cb(self, msg):
         self.last_depth_img = self.cv_bridge.imgmsg_to_cv2(msg).astype(np.float32)
         self.last_extrinsic = self.tf.lookup(
-            self.cam_frame_id, self.frame_id, msg.header.stamp, rospy.Duration(0.1)
+            self.cam_frame_id,
+            self.frame_id,
+            msg.header.stamp,
+            rospy.Duration(0.1),
         )
 
-    def plan_best_grasp(self):
+    def lookup_transforms(self):
-        tsdf_grid = self.tsdf.get_grid()
+        self.T_B_O = self.tf.lookup(self.base_frame_id, self.frame_id, rospy.Time.now())
-        out = self.vgn.predict(tsdf_grid)
-        grasps = compute_grasps(self.tsdf.voxel_size, out)
 
-        vis.draw_tsdf(tsdf_grid, self.tsdf.voxel_size)
+    def run(self):
-        vis.draw_grasps(grasps, 0.05)
+        self.reset()
+        self.explore()
+        self.execute_grasp()
 
-        # Ensure that the camera is pointing forward.
+    def reset(self):
-        if len(grasps) == 0:
+        req = std_srvs.srv.TriggerRequest()
-            return
+        self.reset_client(req)
-        grasp = grasps[0]
+        rospy.sleep(1.0)  # wait for states to be updated
-        rot = grasp.pose.rotation
-        axis = rot.as_matrix()[:, 0]
-        if axis[0] < 0:
-            grasp.pose.rotation = rot * Rotation.from_euler("z", np.pi)
-
-        # Compute target pose of the EE
-        H_T_G = grasp.pose
-        H_B_EE = self.H_B_T * H_T_G * self.H_EE_G.inv()
-        return H_B_EE
-
-
-class SingleViewBaseline(Policy):
-    def __init__(sel):
-        super().__init__()
-
-    def start(self):
         self.done = False
 
+    def explore(self):
+        r = rospy.Rate(self.rate)
+        while not self.done:
+            self.update()
+            r.sleep()
+
+    def update(self):
+        raise NotImplementedError
+
+    def execute_grasp(self):
+        if not self.best_grasp:
+            return
+        grasp = self.best_grasp
+
+        # Ensure that the camera is pointing forward.
+        rot = grasp.pose.rotation
+        if rot.as_matrix()[:, 0][0] < 0:
+            grasp.pose.rotation = rot * Rotation.from_euler("z", np.pi)
+        target = self.T_B_O * grasp.pose * self.ee_grasp_offset.inv()
+
+        self.gripper.move(0.08)
+        rospy.sleep(1.0)  # TODO properly implement the simulated move server
+        self.target_pose_pub.publish(to_pose_msg(target))
+        rospy.sleep(2.0)
+        self.gripper.move(0.0)
+        rospy.sleep(1.0)  # TODO
+        target.translation[2] += 0.3
+        self.target_pose_pub.publish(to_pose_msg(target))
+        rospy.sleep(2.0)
+
+    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)
+        vis.draw_grasps(grasps, 0.05)
+        return grasps[0] if len(grasps) > 0 else None
+
+
+class SingleViewBaseline(BaseController):
+    def __init__(self):
+        super().__init__()
+        self.rate = 1
+
+    def reset(self):
+        super().reset()
+
     def update(self):
-        # Integrate image
         self.tsdf.integrate(
             self.last_depth_img,
             self.intrinsic,
             self.last_extrinsic,
         )
-
-        # Visualize reconstruction
         cloud = self.tsdf.get_scene_cloud()
         vis.draw_points(np.asarray(cloud.points))
-
+        self.best_grasp = self.predict_best_grasp()
-        # Plan grasp
-        self.best_grasp = self.plan_best_grasp()
         self.done = True
-        return
 
 
-class FixedTrajectoryBaseline(Policy):
+class FixedTrajectoryBaseline(BaseController):
     def __init__(self):
         super().__init__()
+        self.rate = 10
         self.duration = 4.0
         self.radius = 0.1
         self.m = scipy.interpolate.interp1d([0, self.duration], [np.pi, 3.0 * np.pi])
 
-    def start(self):
+    def reset(self):
+        super().reset()
         self.tic = rospy.Time.now()
         timeout = rospy.Duration(0.1)
         x0 = self.tf.lookup(self.base_frame_id, self.ee_frame_id, self.tic, timeout)
-        self.origin = np.r_[x0.translation[0] + self.radius, x0.translation[1:]]
+        self.center = np.r_[x0.translation[0] + self.radius, x0.translation[1:]]
         self.target = x0
-        self.done = False
 
     def update(self):
         elapsed_time = (rospy.Time.now() - self.tic).to_sec()
-
-        # Integrate image
-        self.tsdf.integrate(
-            self.last_depth_img,
-            self.intrinsic,
-            self.last_extrinsic,
-        )
-
-        # Visualize current integration
-        cloud = self.tsdf.get_scene_cloud()
-        vis.draw_points(np.asarray(cloud.points))
-
         if elapsed_time > self.duration:
-            self.best_grasp = self.plan_best_grasp()
+            self.best_grasp = self.predict_best_grasp()
             self.done = True
-            return
+        else:
-
+            self.tsdf.integrate(
-        t = self.m(elapsed_time)
+                self.last_depth_img,
-        self.target.translation = (
+                self.intrinsic,
-            self.origin + np.r_[self.radius * np.cos(t), self.radius * np.sin(t), 0.0]
+                self.last_extrinsic,
-        )
+            )
-        self.target_pose_pub.publish(to_pose_msg(self.target))
+            cloud = self.tsdf.get_scene_cloud()
+            vis.draw_points(np.asarray(cloud.points))
+            t = self.m(elapsed_time)
+            self.target.translation = (
+                self.center
+                + np.r_[self.radius * np.cos(t), self.radius * np.sin(t), 0.0]
+            )
+            self.target_pose_pub.publish(to_pose_msg(self.target))
 
 
-class MultiViewPicking(Policy):
+class MultiViewPicking(BaseController):
     pass

run.py

@@ -1,57 +1,7 @@
 import argparse
 import rospy
 
-from geometry_msgs.msg import Pose
+from policies import get_controller
-import std_srvs.srv
-
-from policies import get_policy
-from robot_tools.ros.conversions import *
-from robot_tools.ros.panda import PandaGripperClient
-
-
-class GraspController:
-    def __init__(self, policy, rate):
-        self.policy = policy
-        self.rate = rate
-        self.reset_client = rospy.ServiceProxy("/reset", std_srvs.srv.Trigger)
-        self.target_pose_pub = rospy.Publisher("/cmd", Pose, queue_size=10)
-        self.gripper = PandaGripperClient()
-        rospy.sleep(1.0)
-
-    def run(self):
-        self.reset()
-        self.explore()
-        self.execute_grasp()
-
-    def reset(self):
-        req = std_srvs.srv.TriggerRequest()
-        self.reset_client(req)
-        rospy.sleep(1.0)  # wait for states to be updated
-
-    def explore(self):
-        r = rospy.Rate(self.rate)
-        self.policy.start()
-        while not self.policy.done:
-            self.policy.update()
-            r.sleep()
-
-    def execute_grasp(self):
-        self.gripper.move(0.08)
-        rospy.sleep(1.0)
-        target = self.policy.best_grasp
-
-        if not target:
-            return
-
-        self.target_pose_pub.publish(to_pose_msg(target))
-        rospy.sleep(2.0)
-        self.gripper.move(0.0)
-        rospy.sleep(1.0)
-        target.translation[2] += 0.1
-        self.target_pose_pub.publish(to_pose_msg(target))
-        rospy.sleep(2.0)
-        self.gripper.move(0.08)
-        rospy.sleep(1.0)
 
 
 def create_parser():
@@ -64,7 +14,6 @@ def create_parser():
             "fixed-trajectory",
         ],
     )
-    parser.add_argument("--rate", type=int, default=10)
     return parser
 
 
@@ -74,9 +23,8 @@ def main():
     parser = create_parser()
     args = parser.parse_args()
 
-    policy = get_policy(args.policy)
+    controller = get_controller(args.policy)
-    gc = GraspController(policy, args.rate)
+    controller.run()
-    gc.run()
 
 
 if __name__ == "__main__":

simulation.py

@@ -37,9 +37,9 @@ class Simulation(BtManipulationSim):
 
     def reset(self):
         self.remove_all_objects()
+        self.set_initial_arm_configuration()
         urdfs = np.random.choice(self.urdfs, 4)
         self.add_random_arrangement(urdfs, np.r_[self.origin[:2], 0.625], self.length)
-        self.set_initial_arm_configuration()
 
     def set_initial_arm_configuration(self):
         q = self.arm.configurations["ready"]