Add setup.py file

CMakeLists.txt

@@ -6,6 +6,8 @@ find_package(catkin REQUIRED COMPONENTS
   message_generation
 )
 
+catkin_python_setup()
+
 add_message_files(
   FILES
   AABBox.msg

README.md

@@ -1,13 +1 @@
 # active_grasp
-
-First, run the simulation
-
-```
-roslaunch active_grasp simulation.launch
-```
-
-Then you can run a policy.
-
-```
-python3 run.py ...
-```

active_grasp/controller.py

@@ -1,10 +1,10 @@
 import numpy as np
 import rospy
 
-from robot_utils.ros import tf
+from active_grasp.srv import Reset, ResetRequest
+from active_grasp.utils import *
 from robot_utils.ros.panda import PandaGripperClient
 from robot_utils.spatial import Rotation, Transform
-from utils import CartesianPoseControllerClient
 
 
 class GraspController:
@@ -12,6 +12,7 @@ class GraspController:
         self.policy = policy
         self.controller = CartesianPoseControllerClient()
         self.gripper = PandaGripperClient()
+        self.reset_env = rospy.ServiceProxy("reset", Reset)
         self.load_parameters()
 
     def load_parameters(self):
@@ -24,7 +25,10 @@ class GraspController:
             self.execute_grasp(grasp)
 
     def reset(self):
-        raise NotImplementedError
+        req = ResetRequest()
+        res = self.reset_env(req)
+        rospy.sleep(1.0)  # wait for states to be updated
+        return from_bbox_msg(res.bbox)
 
     def explore(self, bbox):
         self.policy.activate(bbox)

active_grasp/policies.py

@@ -18,12 +18,20 @@ from vgn.utils import *
 
 def get_policy(name):
     if name == "single-view":
-        return SingleViewBaseline()
+        return SingleView()
     else:
         raise ValueError("{} policy does not exist.".format(name))
 
 
-class BasePolicy:
+class Policy:
+    def activate(self, bbox):
+        raise NotImplementedError
+
+    def update(self):
+        raise NotImplementedError
+
+
+class BasePolicy(Policy):
     def __init__(self):
         self.cv_bridge = cv_bridge.CvBridge()
         self.vgn = VGN(Path(rospy.get_param("vgn/model")))
@@ -34,6 +42,8 @@ class BasePolicy:
         self.connect_to_camera()
         self.connect_to_rviz()
 
+        self.rate = 2
+
     def load_parameters(self):
         self.base_frame = rospy.get_param("~base_frame_id")
         self.task_frame = rospy.get_param("~frame_id")
@@ -150,15 +160,11 @@ class BasePolicy:
         self.path_pub.publish(MarkerArray([spheres, lines]))
 
 
-class SingleViewBaseline(BasePolicy):
+class SingleView(BasePolicy):
     """
     Process a single image from the initial viewpoint.
     """
 
-    def __init__(self):
-        super().__init__()
-        self.rate = 1
-
     def update(self):
         self.integrate_latest_image()
         self.draw_scene_cloud()

active_grasp/simulation.py

@@ -2,10 +2,10 @@ from pathlib import Path
 import pybullet as p
 import rospkg
 
+from active_grasp.utils import *
 from robot_utils.bullet import *
 from robot_utils.controllers import CartesianPoseController
 from robot_utils.spatial import Rotation, Transform
-from utils import *
 
 
 class Simulation(BtSim):

launch/active_grasp.rviz

@@ -173,6 +173,7 @@ Visualization Manager:
       Show Axes: true
       Show Names: true
       Tree:
+        world:
           panda_link0:
             panda_link1:
               panda_link2:
@@ -189,7 +190,6 @@ Visualization Manager:
                                 {}
                               panda_rightfinger:
                                 {}
-        world:
           task:
             {}
       Update Interval: 0
@@ -320,7 +320,7 @@ Visualization Manager:
   Views:
     Current:
       Class: rviz/Orbit
-      Distance: 1.1227775812149048
+      Distance: 1.5774216651916504
       Enable Stereo Rendering:
         Stereo Eye Separation: 0.05999999865889549
         Stereo Focal Distance: 1
@@ -357,4 +357,4 @@ Window Geometry:
     collapsed: true
   Width: 1275
   X: 61
-  Y: 200
+  Y: 163

scripts/bt_sim_node.py

@@ -8,13 +8,12 @@ from geometry_msgs.msg import PoseStamped
 import numpy as np
 import rospy
 from sensor_msgs.msg import JointState, Image, CameraInfo
-import std_srvs.srv as std_srvs
 import tf2_ros
 
-import active_grasp.srv
+from active_grasp.srv import Reset, ResetResponse
+from active_grasp.simulation import Simulation
+from active_grasp.utils import *
 from robot_utils.ros.conversions import *
-from simulation import Simulation
-from utils import *
 
 
 class BtSimNode:
@@ -31,7 +30,7 @@ class BtSimNode:
         self.broadcast_transforms()
 
     def advertise_services(self):
-        rospy.Service("reset", active_grasp.srv.Reset, self.reset)
+        rospy.Service("reset", Reset, self.reset)
 
     def broadcast_transforms(self):
         self.static_broadcaster = tf2_ros.StaticTransformBroadcaster()
@@ -47,7 +46,7 @@ class BtSimNode:
         self.reset_requested = True
         rospy.sleep(1.0)  # wait for the latest sim step to finish
         bbox = self.sim.reset()
-        res = active_grasp.srv.ResetResponse(to_bbox_msg(bbox))
+        res = ResetResponse(to_bbox_msg(bbox))
         self.step_cnt = 0
         self.reset_requested = False
         return res

scripts/run.py

@@ -1,23 +1,8 @@
 import argparse
 import rospy
 
-
-from active_grasp.srv import Reset, ResetRequest
-from controller import GraspController
-from policies import get_policy
-from utils import *
-
-
-class SimGraspController(GraspController):
-    def __init__(self, policy):
-        super().__init__(policy)
-        self.reset_sim = rospy.ServiceProxy("reset", Reset)
-
-    def reset(self):
-        req = ResetRequest()
-        res = self.reset_sim(req)
-        rospy.sleep(1.0)  # wait for states to be updated
-        return from_bbox_msg(res.bbox)
+from active_grasp.controller import GraspController
+from active_grasp.policies import get_policy
 
 
 def create_parser():
@@ -41,7 +26,7 @@ def main():
     parser = create_parser()
     args = parser.parse_args()
     policy = get_policy(args.policy)
-    controller = SimGraspController(policy)
+    controller = GraspController(policy)
 
     while True:
         controller.run()

setup.py

@@ -0,0 +1,12 @@
+# ! DO NOT MANUALLY INVOKE THIS setup.py, USE CATKIN INSTEAD!
+
+from distutils.core import setup
+
+from catkin_pkg.python_setup import generate_distutils_setup
+
+# Fetch values from package.xml.
+setup_args = generate_distutils_setup(
+    packages=["active_grasp"],
+)
+
+setup(**setup_args)