diff --git a/bt_sim_node.py b/bt_sim_node.py index 49bbe15..04c6991 100755 --- a/bt_sim_node.py +++ b/bt_sim_node.py @@ -3,17 +3,14 @@ import actionlib import argparse import cv_bridge -import numpy as np -import rospy -import tf2_ros - -import control_msgs.msg import franka_gripper.msg from geometry_msgs.msg import PoseStamped +import numpy as np +import rospy from sensor_msgs.msg import JointState, Image, CameraInfo -import std_srvs.srv +import std_srvs.srv as std_srvs +import tf2_ros -from robot_utils.controllers import CartesianPoseController from robot_utils.ros.conversions import * from simulation import Simulation @@ -21,76 +18,36 @@ from simulation import Simulation class BtSimNode: def __init__(self, gui): self.sim = Simulation(gui=gui) - - self.controller_update_rate = 60 - self.joint_state_publish_rate = 60 - self.camera_publish_rate = 5 - - self.cv_bridge = cv_bridge.CvBridge() - - self.setup_robot_topics() - self.setup_camera_topics() - self.setup_controllers() - self.setup_gripper_interface() - self.broadcast_transforms() - - rospy.Service("reset", std_srvs.srv.Trigger, self.reset) + self.robot_state_interface = RobotStateInterface(self.sim.arm, self.sim.gripper) + self.arm_interface = ArmInterface(self.sim.arm, self.sim.controller) + self.gripper_interface = GripperInterface(self.sim.gripper) + self.camera_interface = CameraInterface(self.sim.camera) self.step_cnt = 0 self.reset_requested = False - def setup_robot_topics(self): - self.joint_pub = rospy.Publisher("joint_states", JointState, queue_size=10) + self.advertise_services() + self.broadcast_transforms() - def setup_controllers(self): - self.cartesian_pose_controller = CartesianPoseController(self.sim.arm) - rospy.Subscriber("command", PoseStamped, self.target_pose_cb) - - def target_pose_cb(self, msg): - assert msg.header.frame_id == self.sim.arm.base_frame - self.cartesian_pose_controller.x_d = from_pose_msg(msg.pose) - - def setup_gripper_interface(self): - self.gripper_interface = GripperInterface(self.sim.gripper) - - def setup_camera_topics(self): - self.cam_info_msg = to_camera_info_msg(self.sim.camera.intrinsic) - self.cam_info_msg.header.frame_id = "cam_optical_frame" - self.cam_info_pub = rospy.Publisher( - "/cam/depth/camera_info", - CameraInfo, - queue_size=10, - ) - self.depth_pub = rospy.Publisher("/cam/depth/image_raw", Image, queue_size=10) + def advertise_services(self): + rospy.Service("reset", std_srvs.Trigger, self.reset) def broadcast_transforms(self): self.static_broadcaster = tf2_ros.StaticTransformBroadcaster() - - msgs = [] - msg = geometry_msgs.msg.TransformStamped() - msg.header.stamp = rospy.Time.now() - msg.header.frame_id = "world" - msg.child_frame_id = "panda_link0" - msg.transform = to_transform_msg(self.sim.T_W_B) - msgs.append(msg) - - msg = geometry_msgs.msg.TransformStamped() - msg.header.stamp = rospy.Time.now() - msg.header.frame_id = "world" - msg.child_frame_id = "task" - msg.transform = to_transform_msg(Transform.translation(self.sim.origin)) - - msgs.append(msg) - + msgs = [ + to_transform_stamped_msg(self.sim.T_W_B, "world", "panda_link0"), + to_transform_stamped_msg( + Transform.translation(self.sim.origin), "world", "task" + ), + ] self.static_broadcaster.sendTransform(msgs) def reset(self, req): self.reset_requested = True rospy.sleep(1.0) # wait for the latest sim step to finish self.sim.reset() - self.cartesian_pose_controller.x_d = self.sim.arm.pose() self.step_cnt = 0 self.reset_requested = False - return std_srvs.srv.TriggerResponse(success=True) + return std_srvs.TriggerResponse(success=True) def run(self): rate = rospy.Rate(self.sim.rate) @@ -102,18 +59,22 @@ class BtSimNode: rate.sleep() def handle_updates(self): - if self.step_cnt % int(self.sim.rate / self.controller_update_rate) == 0: - self.cartesian_pose_controller.update() - self.gripper_interface.update(1.0 / self.controller_update_rate) - if self.step_cnt % int(self.sim.rate / self.joint_state_publish_rate) == 0: - self.publish_joint_state() - if self.step_cnt % int(self.sim.rate / self.camera_publish_rate) == 0: - self.publish_cam_info() - self.publish_cam_imgs() + self.robot_state_interface.update() + self.arm_interface.update() + self.gripper_interface.update(self.sim.dt) + if self.step_cnt % int(self.sim.rate / 5) == 0: + self.camera_interface.update() - def publish_joint_state(self): - q, dq = self.sim.arm.get_state() - width = self.sim.gripper.read() + +class RobotStateInterface: + def __init__(self, arm, gripper): + self.arm = arm + self.gripper = gripper + self.joint_pub = rospy.Publisher("joint_states", JointState, queue_size=10) + + def update(self): + q, _ = self.arm.get_state() + width = self.gripper.read() msg = JointState() msg.header.stamp = rospy.Time.now() msg.name = ["panda_joint{}".format(i) for i in range(1, 8)] + [ @@ -121,18 +82,23 @@ class BtSimNode: "panda_finger_joint2", ] msg.position = np.r_[q, 0.5 * width, 0.5 * width] - msg.velocity = dq self.joint_pub.publish(msg) - def publish_cam_info(self): - self.cam_info_msg.header.stamp = rospy.Time.now() - self.cam_info_pub.publish(self.cam_info_msg) - def publish_cam_imgs(self): - _, depth = self.sim.camera.get_image() - depth_msg = self.cv_bridge.cv2_to_imgmsg(depth) - depth_msg.header.stamp = rospy.Time.now() - self.depth_pub.publish(depth_msg) +class ArmInterface: + def __init__(self, arm, controller): + self.arm = arm + self.controller = controller + rospy.Subscriber("command", PoseStamped, self.target_cb) + + def update(self): + q, _ = self.arm.get_state() + cmd = self.controller.update(q) + self.arm.set_desired_joint_velocities(cmd) + + def target_cb(self, msg): + assert msg.header.frame_id == self.arm.base_frame + self.controller.x_d = from_pose_msg(msg.pose) class GripperInterface: @@ -143,61 +109,59 @@ class GripperInterface: franka_gripper.msg.MoveAction, auto_start=False, ) - self.move_server.register_goal_callback(self.goal_cb) + self.move_server.register_goal_callback(self.move_action_goal_cb) self.move_server.start() - def goal_cb(self): + self.grasp_server = actionlib.SimpleActionServer( + "/franka_gripper/grasp", + franka_gripper.msg.GraspAction, + auto_start=False, + ) + self.grasp_server.register_goal_callback(self.grasp_action_goal_cb) + self.grasp_server.start() + + def move_action_goal_cb(self): + self.elapsed_time_since_move_action_goal = 0.0 goal = self.move_server.accept_new_goal() - self.gripper.move(goal.width) - self.elapsed_time = 0.0 + self.gripper.set_desired_width(goal.width) + + def grasp_action_goal_cb(self): + self.elapsed_time_since_grasp_action_goal = 0.0 + goal = self.grasp_server.accept_new_goal() + self.gripper.set_desired_width(goal.width) def update(self, dt): if self.move_server.is_active(): - self.elapsed_time += dt - if self.elapsed_time > 1.0: + self.elapsed_time_since_move_action_goal += dt + if self.elapsed_time_since_move_action_goal > 1.0: self.move_server.set_succeeded() + if self.grasp_server.is_active(): + self.elapsed_time_since_grasp_action_goal += dt + if self.elapsed_time_since_grasp_action_goal > 1.0: + self.grasp_server.set_succeeded() -class JointTrajectoryController: - def __init__(self, arm, stopped=False): - self.arm = arm - self.action_server = actionlib.SimpleActionServer( - "position_joint_trajectory_controller/follow_joint_trajectory", - control_msgs.msg.FollowJointTrajectoryAction, - auto_start=False, +class CameraInterface: + def __init__(self, camera): + self.camera = camera + self.cv_bridge = cv_bridge.CvBridge() + self.cam_info_msg = to_camera_info_msg(self.camera.intrinsic) + self.cam_info_msg.header.frame_id = "cam_optical_frame" + self.cam_info_pub = rospy.Publisher( + "/cam/depth/camera_info", + CameraInfo, + queue_size=10, ) - self.action_server.register_goal_callback(self.goal_cb) - self.action_server.start() + self.depth_pub = rospy.Publisher("/cam/depth/image_raw", Image, queue_size=10) - self.is_running = False - if not stopped: - self.start() - - def goal_cb(self): - goal = self.action_server.accept_new_goal() - self.elapsed_time = 0.0 - self.points = iter(goal.trajectory.points) - self.next_point = next(self.points) - - def start(self): - self.is_running = True - - def stop(self): - self.is_running = False - - def update(self, dt): - if not (self.is_running and self.action_server.is_active()): - return - - self.elapsed_time += dt - if self.elapsed_time > self.next_point.time_from_start.to_sec(): - try: - self.next_point = next(self.points) - except StopIteration: - self.action_server.set_succeeded() - return - - self.arm.set_desired_joint_positions(self.next_point.positions) + def update(self): + stamp = rospy.Time.now() + self.cam_info_msg.header.stamp = stamp + self.cam_info_pub.publish(self.cam_info_msg) + _, depth = self.camera.get_image() + depth_msg = self.cv_bridge.cv2_to_imgmsg(depth) + depth_msg.header.stamp = stamp + self.depth_pub.publish(depth_msg) def create_parser(): @@ -208,10 +172,8 @@ def create_parser(): def main(): rospy.init_node("bt_sim") - parser = create_parser() args, _ = parser.parse_known_args() - server = BtSimNode(args.gui) server.run() diff --git a/config/active_grasp.yaml b/config/active_grasp.yaml index 7fd374e..6ff8ce8 100644 --- a/config/active_grasp.yaml +++ b/config/active_grasp.yaml @@ -9,5 +9,5 @@ active_grasp: info_topic: /cam/depth/camera_info depth_topic: /cam/depth/image_raw vgn: - model: $(find vgn)/data/models/vgn_conv.pth + model: $(find vgn)/assets/models/vgn_conv.pth finger_depth: 0.05 diff --git a/controller.py b/controller.py new file mode 100644 index 0000000..deeed8f --- /dev/null +++ b/controller.py @@ -0,0 +1,71 @@ +import numpy as np +import rospy + +from robot_utils.ros import tf +from robot_utils.ros.panda import PandaGripperClient +from robot_utils.spatial import Rotation, Transform +from utils import CartesianPoseControllerClient + + +class GraspController: + def __init__(self, policy): + self.policy = policy + self.controller = CartesianPoseControllerClient() + self.gripper = PandaGripperClient() + self.load_parameters() + + def load_parameters(self): + self.T_G_EE = Transform.from_list(rospy.get_param("~ee_grasp_offset")).inv() + + def run(self): + self.reset() + grasp = self.explore() + if grasp: + self.execute_grasp(grasp) + + def reset(self): + raise NotImplementedError + + def explore(self): + self.policy.activate() + r = rospy.Rate(self.policy.rate) + while True: + cmd = self.policy.update() + if self.policy.done: + break + self.controller.send_target(cmd) + r.sleep() + return self.policy.best_grasp + + def execute_grasp(self, grasp): + T_B_G = self.postprocess(grasp) + + self.gripper.move(0.08) + + # Move to an initial pose offset. + self.controller.send_target( + T_B_G * Transform.translation([0, 0, -0.05]) * self.T_G_EE + ) + rospy.sleep(3.0) + + # Approach grasp pose. + self.controller.send_target(T_B_G * self.T_G_EE) + rospy.sleep(1.0) + + # Close the fingers. + self.gripper.grasp() + + # Lift the object. + target = Transform.translation([0, 0, 0.2]) * T_B_G * self.T_G_EE + self.controller.send_target(target) + rospy.sleep(2.0) + + # Check whether the object remains in the hand + return self.gripper.read() > 0.005 + + def postprocess(self, T_B_G): + # Ensure that the camera is pointing forward. + rot = T_B_G.rotation + if rot.as_matrix()[:, 0][0] < 0: + T_B_G.rotation = rot * Rotation.from_euler("z", np.pi) + return T_B_G diff --git a/launch/simulation.launch b/launch/active_grasp.launch similarity index 66% rename from launch/simulation.launch rename to launch/active_grasp.launch index b61c993..6d055b9 100644 --- a/launch/simulation.launch +++ b/launch/active_grasp.launch @@ -1,19 +1,12 @@ - - - - + - - - - - + diff --git a/launch/panda_visualization.rviz b/launch/active_grasp.rviz similarity index 81% rename from launch/panda_visualization.rviz rename to launch/active_grasp.rviz index 74e45a6..4776f8c 100644 --- a/launch/panda_visualization.rviz +++ b/launch/active_grasp.rviz @@ -4,9 +4,9 @@ Panels: Name: Displays Property Tree Widget: Expanded: - - /Global Options1 + - /TF1/Tree1 Splitter Ratio: 0.5 - Tree Height: 549 + Tree Height: 547 - Class: rviz/Selection Name: Selection - Class: rviz/Tool Properties @@ -25,7 +25,7 @@ Panels: Experimental: false Name: Time SyncMode: 0 - SyncSource: Scene Cloud + SyncSource: "" Preferences: PromptSaveOnExit: true Toolbars: @@ -137,13 +137,13 @@ Visualization Manager: Frames: All Enabled: false cam_optical_frame: - Value: true + Value: false panda_hand: - Value: true + Value: false panda_leftfinger: Value: false panda_link0: - Value: true + Value: false panda_link1: Value: false panda_link2: @@ -164,6 +164,8 @@ Visualization Manager: Value: false task: Value: true + world: + Value: true Marker Alpha: 1 Marker Scale: 0.5 Name: TF @@ -171,34 +173,27 @@ Visualization Manager: Show Axes: true Show Names: true Tree: - panda_link0: - panda_link1: - panda_link2: - panda_link3: - panda_link4: - panda_link5: - panda_link6: - panda_link7: - panda_link8: - panda_hand: - cam_optical_frame: - {} - panda_leftfinger: - {} - panda_rightfinger: - {} + world: + panda_link0: + panda_link1: + panda_link2: + panda_link3: + panda_link4: + panda_link5: + panda_link6: + panda_link7: + panda_link8: + panda_hand: + cam_optical_frame: + {} + panda_leftfinger: + {} + panda_rightfinger: + {} task: {} Update Interval: 0 Value: true - - Class: rviz/Marker - Enabled: true - Marker Topic: /workspace - Name: Workspace - Namespaces: - "": true - Queue Size: 100 - Value: true - Alpha: 0.20000000298023224 Autocompute Intensity Bounds: false Autocompute Value Bounds: @@ -214,7 +209,9 @@ Visualization Manager: Enabled: false Invert Rainbow: false Max Color: 255; 255; 255 + Max Intensity: 4096 Min Color: 0; 0; 0 + Min Intensity: 0 Name: TSDF Position Transformer: XYZ Queue Size: 10 @@ -243,7 +240,7 @@ Visualization Manager: Invert Rainbow: false Max Color: 255; 255; 255 Min Color: 0; 0; 0 - Name: Scene Cloud + Name: SceneCloud Position Transformer: XYZ Queue Size: 10 Selectable: true @@ -255,12 +252,36 @@ Visualization Manager: Use Fixed Frame: true Use rainbow: true Value: true + - Alpha: 1 + Axes Length: 0.05000000074505806 + Axes Radius: 0.004999999888241291 + Class: rviz/Pose + Color: 255; 25; 0 + Enabled: true + Head Length: 0.30000001192092896 + Head Radius: 0.10000000149011612 + Name: EETarget + Queue Size: 10 + Shaft Length: 1 + Shaft Radius: 0.05000000074505806 + Shape: Axes + Topic: /command + Unreliable: false + Value: true + - Class: rviz/MarkerArray + Enabled: true + Marker Topic: /path + Name: CameraPath + Namespaces: + {} + Queue Size: 100 + Value: true - Class: rviz/MarkerArray Enabled: true Marker Topic: /grasps - Name: Predicted Grasps + Name: PredictedGrasps Namespaces: - "": true + {} Queue Size: 100 Value: true Enabled: true @@ -291,7 +312,7 @@ Visualization Manager: Views: Current: Class: rviz/Orbit - Distance: 1.5115783214569092 + Distance: 1.5376757383346558 Enable Stereo Rendering: Stereo Eye Separation: 0.05999999865889549 Stereo Focal Distance: 1 @@ -307,17 +328,17 @@ Visualization Manager: Invert Z Axis: false Name: Current View Near Clip Distance: 0.009999999776482582 - Pitch: 0.20979644358158112 + Pitch: 0.21479728817939758 Target Frame: - Yaw: 5.238524913787842 + Yaw: 5.193514823913574 Saved: ~ Window Geometry: Displays: collapsed: false - Height: 846 + Height: 844 Hide Left Dock: false Hide Right Dock: true - QMainWindow State: 000000ff00000000fd000000040000000000000156000002b0fc0200000008fb0000001200530065006c0065006300740069006f006e00000001e10000009b0000005c00fffffffb0000001e0054006f006f006c002000500072006f007000650072007400690065007302000001ed000001df00000185000000a3fb000000120056006900650077007300200054006f006f02000001df000002110000018500000122fb000000200054006f006f006c002000500072006f0070006500720074006900650073003203000002880000011d000002210000017afb000000100044006900730070006c006100790073010000003d000002b0000000c900fffffffb0000002000730065006c0065006300740069006f006e00200062007500660066006500720200000138000000aa0000023a00000294fb00000014005700690064006500530074006500720065006f02000000e6000000d2000003ee0000030bfb0000000c004b0069006e0065006300740200000186000001060000030c00000261000000010000010f000002b0fc0200000003fb0000001e0054006f006f006c002000500072006f00700065007200740069006500730100000041000000780000000000000000fb0000000a00560069006500770073000000003d000002b0000000a400fffffffb0000001200530065006c0065006300740069006f006e010000025a000000b200000000000000000000000200000490000000a9fc0100000001fb0000000a00560069006500770073030000004e00000080000002e10000019700000003000004b00000003efc0100000002fb0000000800540069006d00650100000000000004b0000002eb00fffffffb0000000800540069006d0065010000000000000450000000000000000000000354000002b000000004000000040000000800000008fc00000002000000020000000000000002000000010000000a0054006f006f006c00730100000000ffffffff0000000000000000 + QMainWindow State: 000000ff00000000fd000000040000000000000156000002aefc0200000008fb0000001200530065006c0065006300740069006f006e00000001e10000009b0000005c00fffffffb0000001e0054006f006f006c002000500072006f007000650072007400690065007302000001ed000001df00000185000000a3fb000000120056006900650077007300200054006f006f02000001df000002110000018500000122fb000000200054006f006f006c002000500072006f0070006500720074006900650073003203000002880000011d000002210000017afb000000100044006900730070006c006100790073010000003d000002ae000000c900fffffffb0000002000730065006c0065006300740069006f006e00200062007500660066006500720200000138000000aa0000023a00000294fb00000014005700690064006500530074006500720065006f02000000e6000000d2000003ee0000030bfb0000000c004b0069006e0065006300740200000186000001060000030c00000261000000010000010f000002b0fc0200000003fb0000001e0054006f006f006c002000500072006f00700065007200740069006500730100000041000000780000000000000000fb0000000a00560069006500770073000000003d000002b0000000a400fffffffb0000001200530065006c0065006300740069006f006e010000025a000000b200000000000000000000000200000490000000a9fc0100000001fb0000000a00560069006500770073030000004e00000080000002e10000019700000003000004f30000003efc0100000002fb0000000800540069006d00650100000000000004f3000002eb00fffffffb0000000800540069006d0065010000000000000450000000000000000000000397000002ae00000004000000040000000800000008fc00000002000000020000000000000002000000010000000a0054006f006f006c00730100000000ffffffff0000000000000000 Selection: collapsed: false Time: @@ -326,6 +347,6 @@ Window Geometry: collapsed: false Views: collapsed: true - Width: 1200 - X: 720 - Y: 27 + Width: 1267 + X: 100 + Y: 457 diff --git a/policies.py b/policies.py index d275890..d1d350e 100644 --- a/policies.py +++ b/policies.py @@ -2,259 +2,128 @@ import cv_bridge import numpy as np from pathlib import Path import rospy -import scipy.interpolate - -from geometry_msgs.msg import PoseStamped from sensor_msgs.msg import Image, CameraInfo -import std_srvs.srv from visualization_msgs.msg import Marker, MarkerArray - -from robot_utils.perception import * -from robot_utils.spatial import Rotation, Transform -from robot_utils.ros.conversions import * -from robot_utils.ros.panda import PandaGripperClient from robot_utils.ros import tf +from robot_utils.ros.conversions import * from robot_utils.ros.rviz import * +from robot_utils.spatial import Transform +from vgn.detection import VGN, compute_grasps +from vgn.perception import UniformTSDFVolume import vgn.vis -from vgn.detection import VGN, compute_grasps - -def get_controller(name): +def get_policy(name): if name == "single-view": return SingleViewBaseline() - elif name == "fixed-trajectory": - return FixedTrajectoryBaseline() - elif name == "mvp": - return MultiViewPicking() else: raise ValueError("{} policy does not exist.".format(name)) -class BaseController: +class BasePolicy: def __init__(self): - self.frame = rospy.get_param("~frame_id") - self.length = rospy.get_param("~length") - self.cv_bridge = cv_bridge.CvBridge() - self.reset_client = rospy.ServiceProxy("/reset", std_srvs.srv.Trigger) - self.tsdf = UniformTSDFVolume(0.3, 40) self.vgn = VGN(Path(rospy.get_param("vgn/model"))) - self._setup_robot_connection() - self._setup_camera_connection() - self._setup_rviz_connection() - self._lookup_transforms() + self.load_parameters() + self.lookup_transforms() + self.connect_to_camera() + self.connect_to_rviz() - def run(self): - self.reset() - self.explore() - self.execute_grasp() - # self.release_object() - - def reset(self): - self._reset_env() - self._clear_rviz() - rospy.sleep(1.0) # wait for states to be updated - self._init_policy() - - self.viewpoints = [] - self.done = False - self.best_grasp = None - - def explore(self): - r = rospy.Rate(self.rate) - while not self.done: - self._update() - r.sleep() - - def execute_grasp(self): - if not self.best_grasp: - return - - self.gripper.move(0.08) - - # Ensure that the camera is pointing forward. - T_O_G = self.best_grasp.pose - rot = T_O_G.rotation - if rot.as_matrix()[:, 0][0] < 0: - T_O_G.rotation = rot * Rotation.from_euler("z", np.pi) - - # Move to an initial pose offset. - target = T_O_G * Transform.translation([0, 0, -0.05]) * self.T_G_EE - self._send_target_pose(target) - rospy.sleep(3.0) - - # Approach grasp pose. - self._send_target_pose(T_O_G * self.T_G_EE) - rospy.sleep(1.0) - - # Close the fingers. - self.gripper.move(0.0) - - # Lift the object. - target = Transform.translation([0, 0, 0.2]) * T_O_G * self.T_G_EE - self._send_target_pose(target) - rospy.sleep(2.0) - - def _setup_robot_connection(self): + def load_parameters(self): self.base_frame = rospy.get_param("~base_frame_id") - self.ee_frame = rospy.get_param("~ee_frame_id") - self.T_G_EE = Transform.from_list(rospy.get_param("~ee_grasp_offset")).inv() - self.target_pose_pub = rospy.Publisher("/command", PoseStamped, queue_size=10) - self.gripper = PandaGripperClient() + self.task_frame = rospy.get_param("~frame_id") + self.cam_frame = rospy.get_param("~camera/frame_id") + self.info_topic = rospy.get_param("~camera/info_topic") + self.depth_topic = rospy.get_param("~camera/depth_topic") - 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)) + def lookup_transforms(self): + tf._init_listener() + rospy.sleep(1.0) # wait to receive transforms + self.T_B_task = tf.lookup(self.base_frame, self.task_frame) + + def connect_to_camera(self): + msg = rospy.wait_for_message(self.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) + rospy.Subscriber(self.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 = tf.lookup( - self._cam_frame_id, - self.frame, + def sensor_cb(self, msg): + self.depth_img = self.cv_bridge.imgmsg_to_cv2(msg).astype(np.float32) + self.extrinsic = tf.lookup( + self.cam_frame, + self.task_frame, msg.header.stamp, - rospy.Duration(0.1), + rospy.Duration(0.2), ) - def _setup_rviz_connection(self): + def connect_to_rviz(self): self.path_pub = rospy.Publisher("path", MarkerArray, queue_size=1, latch=True) - def _lookup_transforms(self): - self.T_B_O = tf.lookup(self.base_frame, self.frame) + def activate(self): + self.viewpoints = [] + self.done = False + self.best_grasp = None # grasp pose defined w.r.t. the robot's base frame - def _reset_env(self): - req = std_srvs.srv.TriggerRequest() - self.reset_client(req) - - def _clear_rviz(self): - vgn.vis.clear() - self.path_pub.publish(DELETE_MARKER_ARRAY_MSG) - - def _init_policy(self): + def update(self): raise NotImplementedError - def _update(self): - raise NotImplementedError - - def _draw_camera_path(self, frustum=False): - identity = Transform.identity() - color = np.r_[31, 119, 180] / 255.0 - - # Spheres for each viewpoint - scale = 0.01 * np.ones(3) - spheres = create_marker(Marker.SPHERE_LIST, self.frame, identity, scale, color) - spheres.id = 0 - spheres.points = [to_point_msg(p.translation) for p in self.viewpoints] - - # Line strip connecting viewpoints - scale = [0.005, 0.0, 0.0] - lines = create_marker(Marker.LINE_STRIP, self.frame, identity, scale, color) - lines.id = 1 - lines.points = [to_point_msg(p.translation) for p in self.viewpoints] - - markers = [spheres, lines] - - # Frustums - if frustum: - for i, pose in enumerate(self.viewpoints): - msg = create_cam_marker(self.intrinsic, pose, self.frame) - msg.id = i + 2 - markers.append(msg) - - self.path_pub.publish(MarkerArray(markers)) - - def _draw_scene_cloud(self): - cloud = self.tsdf.get_scene_cloud() - vgn.vis.draw_points(np.asarray(cloud.points)) - - def _integrate_latest_image(self): - self.viewpoints.append(self.last_extrinsic.inv()) + def integrate_latest_image(self): + self.viewpoints.append(self.extrinsic.inv()) self.tsdf.integrate( - self.last_depth_img, + self.depth_img, self.intrinsic, - self.last_extrinsic, + self.extrinsic, ) - def _predict_best_grasp(self): + 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) vgn.vis.draw_grasps(grasps, 0.05) - return grasps[0] if len(grasps) > 0 else None + return self.T_B_task * grasps[0].pose if len(grasps) > 0 else None - def _send_target_pose(self, target): - """Target is expected to be given w.r.t. the task frame.""" - msg = PoseStamped() - msg.header.frame_id = self.base_frame - msg.pose = to_pose_msg(self.T_B_O * target) - self.target_pose_pub.publish(msg) + def draw_scene_cloud(self): + cloud = self.tsdf.get_scene_cloud() + vgn.vis.draw_points(np.asarray(cloud.points)) + + def draw_camera_path(self): + identity = Transform.identity() + color = np.r_[31, 119, 180] / 255.0 + + # Spheres for each viewpoint + scale = 0.01 * np.ones(3) + spheres = create_marker( + Marker.SPHERE_LIST, self.task_frame, identity, scale, color + ) + spheres.id = 0 + spheres.points = [to_point_msg(p.translation) for p in self.viewpoints] + + # Line strip connecting viewpoints + scale = [0.005, 0.0, 0.0] + lines = create_marker( + Marker.LINE_STRIP, self.task_frame, identity, scale, color + ) + lines.id = 1 + lines.points = [to_point_msg(p.translation) for p in self.viewpoints] + + self.path_pub.publish(MarkerArray([spheres, lines])) -class SingleViewBaseline(BaseController): +class SingleViewBaseline(BasePolicy): """ - Integrate a single image from the initial viewpoint. + Process a single image from the initial viewpoint. """ def __init__(self): super().__init__() self.rate = 1 - def _init_policy(self): - pass - - def _update(self): - self._integrate_latest_image() - self._draw_scene_cloud() - self._draw_camera_path(frustum=True) - self.best_grasp = self._predict_best_grasp() + def update(self): + self.integrate_latest_image() + self.draw_scene_cloud() + self.draw_camera_path() + self.best_grasp = self.predict_best_grasp() self.done = True - - -class FixedTrajectoryBaseline(BaseController): - """Follow a pre-defined circular trajectory.""" - - 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 _init_policy(self): - self.tic = rospy.Time.now() - x0 = tf.lookup(self.frame, self.ee_frame) - self.center = np.r_[x0.translation[0] + self.radius, x0.translation[1:]] - self.target = x0 - - def _update(self): - elapsed_time = (rospy.Time.now() - self.tic).to_sec() - if elapsed_time > self.duration: - self.best_grasp = self._predict_best_grasp() - self.done = True - else: - # Update state - self._integrate_latest_image() - - # Compute next viewpoint - 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._send_target_pose(self.target) - - # Draw - self._draw_scene_cloud() - self._draw_camera_path() - - -class MultiViewPicking(BaseController): - pass diff --git a/run.py b/run.py deleted file mode 100644 index c3787df..0000000 --- a/run.py +++ /dev/null @@ -1,32 +0,0 @@ -import argparse -import rospy - -from policies import get_controller - - -def create_parser(): - parser = argparse.ArgumentParser() - parser.add_argument( - "--policy", - type=str, - choices=[ - "single-view", - "fixed-trajectory", - "mvp", - ], - ) - return parser - - -def main(): - rospy.init_node("active_grasp") - - parser = create_parser() - args = parser.parse_args() - - controller = get_controller(args.policy) - controller.run() - - -if __name__ == "__main__": - main() diff --git a/sim_grasp.py b/sim_grasp.py new file mode 100644 index 0000000..89b07ef --- /dev/null +++ b/sim_grasp.py @@ -0,0 +1,49 @@ +import argparse +import rospy +import std_srvs.srv as std_srvs + + +from controller import GraspController +from policies import get_policy + + +class SimGraspController(GraspController): + def __init__(self, policy): + super().__init__(policy) + self.reset_sim = rospy.ServiceProxy("/reset", std_srvs.Trigger) + + def reset(self): + req = std_srvs.TriggerRequest() + self.reset_sim(req) + rospy.sleep(1.0) # wait for states to be updated + + +def create_parser(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--policy", + type=str, + choices=[ + "single-view", + "top", + "alignment", + "random", + "fixed-trajectory", + ], + ) + return parser + + +def main(): + rospy.init_node("active_grasp") + parser = create_parser() + args = parser.parse_args() + policy = get_policy(args.policy) + controller = SimGraspController(policy) + + while True: + controller.run() + + +if __name__ == "__main__": + main() diff --git a/simulation.py b/simulation.py index 3ecf89e..6e3c831 100644 --- a/simulation.py +++ b/simulation.py @@ -3,43 +3,53 @@ import pybullet as p import rospkg from robot_utils.bullet import * +from robot_utils.controllers import CartesianPoseController from robot_utils.spatial import Rotation, Transform -class Simulation(BtManipulationSim): +class Simulation(BtSim): def __init__(self, gui=True): super().__init__(gui=gui, sleep=False) p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0) p.resetDebugVisualizerCamera(1.4, 50, -35, [0.0, 0.0, 0.6]) + self.object_uids = [] + self.find_object_urdfs() - self.add_table() - self.add_robot() + self.load_table() + self.load_robot() + self.load_controller() + + self.reset() def find_object_urdfs(self): rospack = rospkg.RosPack() - root = Path(rospack.get_path("vgn")) / "data/urdfs/packed/test" + root = Path(rospack.get_path("vgn")) / "assets/urdfs/packed/test" self.urdfs = [str(f) for f in root.iterdir() if f.suffix == ".urdf"] - def add_table(self): + def load_table(self): p.loadURDF("plane.urdf") ori = Rotation.from_rotvec(np.array([0, 0, np.pi / 2])).as_quat() p.loadURDF("table/table.urdf", baseOrientation=ori, useFixedBase=True) self.length = 0.3 self.origin = [-0.3, -0.5 * self.length, 0.5] - def add_robot(self): - self.T_W_B = Transform(Rotation.identity(), np.r_[-0.6, 0.0, 0.5]) + def load_robot(self): + self.T_W_B = Transform(Rotation.identity(), np.r_[-0.6, 0.0, 0.4]) self.arm = BtPandaArm(self.T_W_B) self.gripper = BtPandaGripper(self.arm) + self.model = Model(self.arm.urdf_path, self.arm.base_frame, self.arm.ee_frame) self.camera = BtCamera(320, 240, 1.047, 0.1, 1.0, self.arm.uid, 11) + def load_controller(self): + self.controller = CartesianPoseController(self.model, self.arm.ee_frame, None) + def reset(self): self.remove_all_objects() self.set_initial_arm_configuration() urdfs = np.random.choice(self.urdfs, 4) origin = np.r_[self.origin[:2], 0.625] - self.add_random_arrangement(urdfs, origin, self.length, 0.8) + self.random_object_arrangement(urdfs, origin, self.length, 0.8) def set_initial_arm_configuration(self): q = self.arm.configurations["ready"] @@ -47,3 +57,69 @@ class Simulation(BtManipulationSim): q[5] = np.deg2rad(np.random.uniform(90, 105)) for i, q_i in enumerate(q): p.resetJointState(self.arm.uid, i, q_i, 0) + p.resetJointState(self.arm.uid, 9, 0.04, 0) + p.resetJointState(self.arm.uid, 10, 0.04, 0) + x0 = self.model.pose(self.arm.ee_frame, q) + self.controller.x_d = x0 + + def load_object(self, urdf, ori, pos, scale=1.0): + uid = p.loadURDF(str(urdf), pos, ori.as_quat(), globalScaling=scale) + self.object_uids.append(uid) + return uid + + def remove_object(self, uid): + p.removeBody(uid) + self.object_uids.remove(uid) + + def remove_all_objects(self): + for uid in list(self.object_uids): + self.remove_object(uid) + + def random_object_arrangement(self, urdfs, origin, length, scale=1.0, attempts=10): + for urdf in urdfs: + # Load the object + uid = self.load_object(urdf, Rotation.identity(), [0.0, 0.0, 0.0], scale) + lower, upper = p.getAABB(uid) + z_offset = 0.5 * (upper[2] - lower[2]) + 0.002 + state_id = p.saveState() + for _ in range(attempts): + # Try to place the object without collision + ori = Rotation.from_rotvec([0.0, 0.0, np.random.uniform(0, 2 * np.pi)]) + offset = np.r_[np.random.uniform(0.2, 0.8, 2) * length, z_offset] + p.resetBasePositionAndOrientation(uid, origin + offset, ori.as_quat()) + self.step() + if not p.getContactPoints(uid): + break + else: + p.restoreState(stateId=state_id) + else: + # No placement found, remove the object + self.remove_object(uid) + + +class CartesianPoseController: + def __init__(self, model, frame, x0): + self._model = model + self._frame = frame + + self.kp = np.ones(6) * 4.0 + self.max_linear_vel = 0.2 + self.max_angular_vel = 1.57 + + self.x_d = x0 + + def update(self, q): + x = self._model.pose(self._frame, q) + error = np.zeros(6) + error[:3] = self.x_d.translation - x.translation + error[3:] = (self.x_d.rotation * x.rotation.inv()).as_rotvec() + dx = self._limit_rate(self.kp * error) + J_pinv = np.linalg.pinv(self._model.jacobian(self._frame, q)) + cmd = np.dot(J_pinv, dx) + return cmd + + def _limit_rate(self, dx): + linear, angular = dx[:3], dx[3:] + linear = np.clip(linear, -self.max_linear_vel, self.max_linear_vel) + angular = np.clip(angular, -self.max_angular_vel, self.max_angular_vel) + return np.r_[linear, angular] diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..7ef2723 --- /dev/null +++ b/utils.py @@ -0,0 +1,13 @@ +from geometry_msgs.msg import PoseStamped +import rospy + +from robot_utils.ros.conversions import * + + +class CartesianPoseControllerClient: + def __init__(self, topic="/command"): + self.target_pub = rospy.Publisher(topic, PoseStamped, queue_size=10) + + def send_target(self, pose): + msg = to_pose_stamped_msg(pose, "panda_link0") + self.target_pub.publish(msg)