Add robot configuration to the state

active_grasp/baselines.py

@@ -4,9 +4,9 @@ from .policy import SingleViewPolicy, MultiViewPolicy, compute_error
 
 
 class InitialView(SingleViewPolicy):
-    def update(self, img, pose):
+    def update(self, img, x, q):
-        self.x_d = pose
+        self.x_d = x
-        super().update(img, pose)
+        super().update(img, x, q)
 
 
 class TopView(SingleViewPolicy):
@@ -22,8 +22,8 @@ class TopTrajectory(MultiViewPolicy):
         self.x_d = self.view_sphere.get_view(0.0, 0.0)
         self.done = False if self.is_feasible(self.x_d) else True
 
-    def update(self, img, x):
+    def update(self, img, x, q):
-        self.integrate(img, x)
+        self.integrate(img, x, q)
         linear, _ = compute_error(self.x_d, x)
         if np.linalg.norm(linear) < 0.02:
             self.done = True

active_grasp/controller.py

@@ -11,7 +11,7 @@ from .timer import Timer
 from active_grasp.srv import Reset, ResetRequest
 from robot_helpers.ros import tf
 from robot_helpers.ros.conversions import *
-from robot_helpers.ros.panda import PandaGripperClient
+from robot_helpers.ros.panda import PandaArmClient, PandaGripperClient
 from robot_helpers.ros.moveit import MoveItClient
 from robot_helpers.spatial import Rotation, Transform
 from vgn.utils import look_at, cartesian_to_spherical, spherical_to_cartesian
@@ -43,8 +43,9 @@ class GraspController:
         )
 
     def init_robot_connection(self):
-        self.cartesian_vel_pub = rospy.Publisher("command", Twist, queue_size=10)
+        self.arm = PandaArmClient()
         self.gripper = PandaGripperClient()
+        self.cartesian_vel_pub = rospy.Publisher("command", Twist, queue_size=10)
 
     def init_moveit(self):
         self.moveit = MoveItClient("panda_arm")
@@ -95,18 +96,19 @@ class GraspController:
         timer = rospy.Timer(rospy.Duration(1.0 / self.control_rate), self.send_vel_cmd)
         r = rospy.Rate(self.policy_rate)
         while not self.policy.done:
-            img, pose = self.get_state()
+            img, pose, q = self.get_state()
-            self.policy.update(img, pose)
+            self.policy.update(img, pose, q)
             r.sleep()
         rospy.sleep(0.1)  # Wait for a zero command to be sent to the robot.
         timer.shutdown()
         return self.policy.best_grasp
 
     def get_state(self):
+        q, _ = self.arm.get_state()
         msg = copy.deepcopy(self.latest_depth_msg)
         img = self.cv_bridge.imgmsg_to_cv2(msg).astype(np.float32)
         pose = tf.lookup(self.base_frame, self.cam_frame, msg.header.stamp)
-        return img, pose
+        return img, pose, q
 
     def send_vel_cmd(self, event):
         if self.policy.x_d is None or self.policy.done:

active_grasp/nbv.py

@@ -73,7 +73,7 @@ class NextBestView(MultiViewPolicy):
             if self.is_feasible(view):
                 self.view_candidates.append(view)
 
-    def update(self, img, x):
+    def update(self, img, x, q):
         if len(self.views) > self.max_views or self.best_grasp_prediction_is_stable():
             self.done = True
         else:

active_grasp/policy.py

@@ -62,7 +62,7 @@ class Policy:
         rospy.sleep(0.5)  # Wait for tf tree to be updated
         self.vis.workspace(self.task_frame, 0.3)
 
-    def update(self, img, pose):
+    def update(self, img, x, q):
         raise NotImplementedError
 
     def sort_grasps(self, in_grasps):
@@ -92,7 +92,7 @@ class Policy:
 
 
 class SingleViewPolicy(Policy):
-    def update(self, img, x):
+    def update(self, img, x, q):
         linear, _ = compute_error(self.x_d, x)
         if np.linalg.norm(linear) < 0.02:
             self.views.append(x)