nbv_sim/run.py
2021-05-26 14:46:12 +02:00

84 lines
2.0 KiB
Python

import argparse
import rospy
from geometry_msgs.msg import Pose
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():
parser = argparse.ArgumentParser()
parser.add_argument(
"--policy",
type=str,
choices=[
"single-view",
"fixed-trajectory",
],
)
parser.add_argument("--rate", type=int, default=10)
return parser
def main():
rospy.init_node("active_grasp")
parser = create_parser()
args = parser.parse_args()
policy = get_policy(args.policy)
gc = GraspController(policy, args.rate)
gc.run()
if __name__ == "__main__":
main()