def main(args):
    cfg_file = os.path.join(args.example_config_path, args.primitive) + ".yaml"
    cfg = get_cfg_defaults()
    cfg.merge_from_file(cfg_file)
    cfg.freeze()

    rospy.init_node('MacroActions')
    signal.signal(signal.SIGINT, signal_handler)

    data = {}
    data['saved_data'] = []
    data['metadata'] = {}

    # parent1, child1 = Pipe()
    # parent2, child2 = Pipe()
    # work_queue = Queue()
    # result_queue = Queue()
    # p1 = Process(target=worker_yumi, args=(child1, work_queue, result_queue, cfg, args))
    # p2 = Process(target=worker_yumi, args=(child2, work_queue, result_queue, cfg, args))
    # p1.start()
    # p2.start()

    # parent1.send("RESET")
    # parent2.send("RESET")

    # print("started workers")
    # time.sleep(15.0)
    # embed()

    # # setup yumi
    # data_seed = 1
    data_seed = args.np_seed
    np.random.seed(data_seed)

    yumi_ar = Robot('yumi_palms',
                    pb=True,
                    pb_cfg={'gui': args.visualize},
                    arm_cfg={'self_collision': False,
                             'seed': data_seed})

    r_gel_id = cfg.RIGHT_GEL_ID
    l_gel_id = cfg.LEFT_GEL_ID

    alpha = cfg.ALPHA
    K = cfg.GEL_CONTACT_STIFFNESS
    restitution = cfg.GEL_RESTITUION

    p.changeDynamics(
        yumi_ar.arm.robot_id,
        r_gel_id,
        restitution=restitution,
        contactStiffness=K,
        contactDamping=alpha*K,
        rollingFriction=args.rolling
    )

    p.changeDynamics(
        yumi_ar.arm.robot_id,
        l_gel_id,
        restitution=restitution,
        contactStiffness=K,
        contactDamping=alpha*K,
        rollingFriction=args.rolling
    )
    dynamics_info = {}
    dynamics_info['contactDamping'] = alpha*K
    dynamics_info['contactStiffness'] = K
    dynamics_info['rollingFriction'] = args.rolling
    dynamics_info['restitution'] = restitution

    yumi_gs = YumiCamsGS(
        yumi_ar,
        cfg,
        exec_thread=False,
        sim_step_repeat=args.step_repeat)

    for _ in range(10):
        yumi_gs.update_joints(cfg.RIGHT_INIT + cfg.LEFT_INIT)

    if args.object:
        obj_id = yumi_ar.pb_client.load_urdf(
            args.config_package_path +
            'descriptions/urdf/'+args.object_name+'.urdf',
            cfg.OBJECT_POSE_3[0:3],
            cfg.OBJECT_POSE_3[3:]
        )

        goal_obj_id = yumi_ar.pb_client.load_urdf(
            args.config_package_path +
            'descriptions/urdf/'+args.object_name+'_trans.urdf',
            cfg.OBJECT_POSE_3[0:3],
            cfg.OBJECT_POSE_3[3:]
        )
        p.setCollisionFilterPair(yumi_ar.arm.robot_id, goal_obj_id, r_gel_id, -1, enableCollision=False)
        p.setCollisionFilterPair(obj_id, goal_obj_id, -1, -1, enableCollision=False)
        p.setCollisionFilterPair(yumi_ar.arm.robot_id, obj_id, r_gel_id, -1, enableCollision=True)
        p.setCollisionFilterPair(yumi_ar.arm.robot_id, obj_id, 27, -1, enableCollision=True)

        yumi_ar.pb_client.reset_body(
            obj_id,
            cfg.OBJECT_POSE_3[:3],
            cfg.OBJECT_POSE_3[3:])

        yumi_ar.pb_client.reset_body(
            goal_obj_id,
            cfg.OBJECT_POSE_3[:3],
            cfg.OBJECT_POSE_3[3:])


    manipulated_object = None
    object_pose1_world = util.list2pose_stamped(cfg.OBJECT_INIT)
    object_pose2_world = util.list2pose_stamped(cfg.OBJECT_FINAL)
    palm_pose_l_object = util.list2pose_stamped(cfg.PALM_LEFT)
    palm_pose_r_object = util.list2pose_stamped(cfg.PALM_RIGHT)

    example_args = {}
    example_args['object_pose1_world'] = object_pose1_world
    example_args['object_pose2_world'] = object_pose2_world
    example_args['palm_pose_l_object'] = palm_pose_l_object
    example_args['palm_pose_r_object'] = palm_pose_r_object
    example_args['object'] = manipulated_object
    example_args['N'] = 60
    example_args['init'] = True
    example_args['table_face'] = 0

    primitive_name = args.primitive

    mesh_file = args.config_package_path + 'descriptions/meshes/objects/' + args.object_name + '_experiments.stl'
    exp_single = SingleArmPrimitives(
        cfg,
        yumi_ar.pb_client.get_client_id(),
        obj_id,
        mesh_file)
    if primitive_name == 'grasp' or primitive_name == 'pivot':
        exp_double = DualArmPrimitives(
            cfg,
            yumi_ar.pb_client.get_client_id(),
            obj_id,
            mesh_file)
        exp_running = exp_double
    else:
        exp_running = exp_single

    # setup macro_planner
    action_planner = ClosedLoopMacroActions(
        cfg,
        yumi_gs,
        obj_id,
        yumi_ar.pb_client.get_client_id(),
        args.config_package_path,
        replan=args.replan,
        object_mesh_file=mesh_file
    )

    data['metadata']['mesh_file'] = mesh_file
    data['metadata']['cfg'] = cfg
    data['metadata']['dynamics'] = dynamics_info
    data['metadata']['cam_cfg'] = yumi_gs.cam_setup_cfg
    data['metadata']['step_repeat'] = args.step_repeat

    delta_z_height = 0.95
    with open(args.config_package_path+'descriptions/urdf/'+args.object_name+'.urdf', 'rb') as f:
        urdf_txt = f.read()

    data['metadata']['object_urdf'] = urdf_txt
    data['metadata']['delta_z_height'] = delta_z_height
    data['metadata']['step_repeat'] = args.step_repeat
    data['metadata']['seed'] = data_seed

    metadata = data['metadata']

    if args.multi:
        cuboid_sampler = CuboidSampler(
            '/root/catkin_ws/src/primitives/objects/cuboids/nominal_cuboid.stl',
            pb_client=yumi_ar.pb_client)
        cuboid_fname_template = '/root/catkin_ws/src/primitives/objects/cuboids/'

        cuboid_manager = MultiBlockManager(
            cuboid_fname_template,
            cuboid_sampler,
            robot_id=yumi_ar.arm.robot_id,
            table_id=27,
            r_gel_id=r_gel_id,
            l_gel_id=l_gel_id)

        yumi_ar.pb_client.remove_body(obj_id)
        yumi_ar.pb_client.remove_body(goal_obj_id)

        cuboid_fname = cuboid_manager.get_cuboid_fname()

        obj_id, sphere_ids, mesh, goal_obj_id = \
            cuboid_sampler.sample_cuboid_pybullet(
                cuboid_fname,
                goal=True,
                keypoints=False)

        cuboid_manager.filter_collisions(obj_id, goal_obj_id)
        action_planner.update_object(obj_id, mesh_file)

    trans_box_lock = threading.RLock()
    goal_viz = GoalVisual(
        trans_box_lock,
        goal_obj_id,
        action_planner.pb_client,
        cfg.OBJECT_POSE_3)

    pickle_path = os.path.join(
        args.data_dir,
        primitive_name,
        args.experiment_name
    )

    if not os.path.exists(pickle_path):
        os.makedirs(pickle_path)

    data_manager = DataManager(pickle_path)

    if args.save_data:
        with open(os.path.join(pickle_path, 'metadata.pkl'), 'wb') as mdata_f:
            pickle.dump(metadata, mdata_f)

    if args.debug:
        if args.multi:
            cuboid_sampler.delete_cuboid(obj_id, goal_obj_id, sphere_ids)

            cuboid_fname = cuboid_manager.get_cuboid_fname()
            obj_id, sphere_ids, mesh, goal_obj_id = cuboid_sampler.sample_cuboid_pybullet(
                cuboid_fname,
                goal=True,
                keypoints=False)

            cuboid_manager.filter_collisions(obj_id, goal_obj_id)
            goal_viz.update_goal_obj(goal_obj_id)
            p.changeDynamics(
                obj_id,
                -1,
                lateralFriction=0.4
            )

            action_planner.update_object(obj_id, mesh_file)
            exp_running.initialize_object(obj_id, cuboid_fname)
            print('Reset multi block!')
        else:
            cuboid_fname = '/root/catkin_ws/src/config/descriptions/meshes/objects/cuboids/realsense_box_experiments.stl'

        for _ in range(args.num_obj_samples):
            if primitive_name == 'pull':
                init_id = exp_running.get_rand_init(ind=2)[-1]
                obj_pose_final = util.list2pose_stamped(exp_running.init_poses[init_id])
                point, normal, face = exp_running.sample_contact(primitive_name)

                world_pose = exp_running.get_palm_poses_world_frame(
                    point,
                    normal,
                    primitive_name=primitive_name)

                obj_pos_world = list(p.getBasePositionAndOrientation(
                    obj_id, yumi_ar.pb_client.get_client_id())[0])
                obj_ori_world = list(p.getBasePositionAndOrientation(
                    obj_id, yumi_ar.pb_client.get_client_id())[1])

                obj_pose_world = util.list2pose_stamped(obj_pos_world + obj_ori_world)
                contact_obj_frame = util.convert_reference_frame(
                    world_pose, obj_pose_world, util.unit_pose())

                example_args['palm_pose_r_object'] = contact_obj_frame
                example_args['object_pose1_world'] = obj_pose_world

                obj_pose_final = util.list2pose_stamped(exp_running.init_poses[init_id])
                obj_pose_final.pose.position.z = obj_pose_world.pose.position.z/1.175
                print("init: ")
                print(util.pose_stamped2list(object_pose1_world))
                print("final: ")
                print(util.pose_stamped2list(obj_pose_final))
                example_args['object_pose2_world'] = obj_pose_final
                example_args['table_face'] = init_id
            elif primitive_name == 'grasp':
                k = 0
                have_contact = False
                contact_face = None
                while True:
                    x, y, dq, q, init_id = exp_running.get_rand_init()
                    obj_pose_world_nom = exp_running.get_obj_pose()[0]

                    palm_poses_world = exp_running.get_palm_poses_world_frame(
                        init_id,
                        obj_pose_world_nom,
                        [x, y, dq])

                    # get_palm_poses_world_frame may adjust the
                    # initial object pose, so need to check it again
                    obj_pose_world = exp_running.get_obj_pose()[0]

                    if palm_poses_world is not None:
                        have_contact = True
                        break
                    k += 1
                    if k >= 10:
                        print("FAILED")
                        break

                if have_contact:
                    obj_pose_final = exp_running.goal_pose_world_frame_mod
                    palm_poses_obj_frame = {}
                    for key in palm_poses_world.keys():
                        palm_poses_obj_frame[key] = util.convert_reference_frame(
                            palm_poses_world[key], obj_pose_world, util.unit_pose())

                    example_args['palm_pose_r_object'] = palm_poses_obj_frame['right']
                    example_args['palm_pose_l_object'] = palm_poses_obj_frame['left']
                    example_args['object_pose1_world'] = obj_pose_world
                    example_args['object_pose2_world'] = obj_pose_final
                    example_args['table_face'] = init_id

            plan = action_planner.get_primitive_plan(primitive_name, example_args, 'right')

            embed()

            import simulation

            for i in range(10):
                simulation.visualize_object(
                    object_pose1_world,
                    filepath="package://config/descriptions/meshes/objects/cuboids/" + cuboid_fname.split('objects/cuboids')[1],
                    name="/object_initial",
                    color=(1., 0., 0., 1.),
                    frame_id="/yumi_body",
                    scale=(1., 1., 1.))
                simulation.visualize_object(
                    object_pose2_world,
                    filepath="package://config/descriptions/meshes/objects/cuboids/" + cuboid_fname.split('objects/cuboids')[1],
                    name="/object_final",
                    color=(0., 0., 1., 1.),
                    frame_id="/yumi_body",
                    scale=(1., 1., 1.))
                rospy.sleep(.1)
            simulation.simulate(plan, cuboid_fname.split('objects/cuboids')[1])
    else:
        global_start = time.time()
        face = 0
        # exp_double.reset_graph(face)
        start_time = time.time()
        success = 0

        for trial in range(args.num_trials):
            k = 0

            if args.multi:
                cuboid_sampler.delete_cuboid(obj_id, goal_obj_id, sphere_ids)

                cuboid_fname = cuboid_manager.get_cuboid_fname()
                obj_id, sphere_ids, mesh, goal_obj_id = cuboid_sampler.sample_cuboid_pybullet(
                    cuboid_fname,
                    goal=True,
                    keypoints=False)

                cuboid_manager.filter_collisions(obj_id, goal_obj_id)
                goal_viz.update_goal_obj(goal_obj_id)
                p.changeDynamics(
                    obj_id,
                    -1,
                    lateralFriction=0.4
                )

                action_planner.update_object(obj_id, mesh_file)
                exp_running.initialize_object(obj_id, cuboid_fname)
                print('Reset multi block!')

            for _ in range(args.num_obj_samples):

                while True:
                    have_contact = False
                    # sample a random stable pose, and get the corresponding
                    # stable orientation index
                    k += 1

                    if primitive_name == 'pull':
                        # init_id = exp_running.get_rand_init()[-1]
                        init_id = exp_running.get_rand_init()[-1]

                        # sample a point on the object that is valid
                        # for the primitive action being executed
                        point, normal, face = exp_running.sample_contact(
                            primitive_name=primitive_name)
                        if point is not None:
                            break
                    elif primitive_name == 'grasp':
                        x, y, dq, q, init_id = exp_double.get_rand_init()
                        obj_pose_world_nom = exp_double.get_obj_pose()[0]

                        palm_poses_world = exp_double.get_palm_poses_world_frame(
                            init_id,
                            obj_pose_world_nom,
                            [x, y, dq])

                        obj_pose_world = exp_double.get_obj_pose()[0]

                        if palm_poses_world is not None:
                            have_contact = True
                            break
                    if k >= 10:
                        print("FAILED")
                        return

                # for _ in range(10):
                #     yumi_gs.update_joints(cfg.RIGHT_INIT + cfg.LEFT_INIT)

                if primitive_name == 'pull':
                    # get the full 6D pose palm in world, at contact location
                    palm_pose_world = exp_running.get_palm_poses_world_frame(
                        point,
                        normal,
                        primitive_name=primitive_name)

                    # get the object pose in the world frame

                    # if trial == 0:
                    #     parent1.send("OBJECT_POSE")
                    # elif trial == 1:
                    #     parent2.send("OBJECT_POSE")

                    obj_pos_world = list(p.getBasePositionAndOrientation(
                        obj_id,
                        yumi_ar.pb_client.get_client_id())[0])
                    obj_ori_world = list(p.getBasePositionAndOrientation(
                        obj_id,
                        yumi_ar.pb_client.get_client_id())[1])

                    obj_pose_world = util.list2pose_stamped(
                        obj_pos_world + obj_ori_world)

                    # obj_pose_world = work_queue.get(block=True)

                    # transform the palm pose from the world frame to the object frame
                    contact_obj_frame = util.convert_reference_frame(
                        palm_pose_world, obj_pose_world, util.unit_pose())

                    # set up inputs to the primitive planner, based on task
                    # including sampled initial object pose and contacts,
                    # and final object pose
                    example_args['palm_pose_r_object'] = contact_obj_frame
                    example_args['object_pose1_world'] = obj_pose_world

                    # obj_pose_final = util.list2pose_stamped(exp_running.init_poses[init_id])

                    x, y, q, _ = exp_running.get_rand_init(execute=False, ind=init_id)
                    final_nominal = exp_running.init_poses[init_id]
                    final_nominal[0] = x
                    final_nominal[1] = y
                    final_nominal[3:] = q
                    obj_pose_final = util.list2pose_stamped(final_nominal)
                    goal_viz.update_goal_state(final_nominal)
                    obj_pose_final.pose.position.z += cfg.TABLE_HEIGHT

                    example_args['object_pose2_world'] = obj_pose_final
                    example_args['table_face'] = init_id
                    example_args['primitive_name'] = primitive_name
                    example_args['N'] = exp_running.calc_n(
                        obj_pose_world, obj_pose_final)
                elif primitive_name == 'grasp':
                    if have_contact:
                        obj_pose_final = exp_double.goal_pose_world_frame_mod
                        palm_poses_obj_frame = {}
                        for key in palm_poses_world.keys():
                            palm_poses_obj_frame[key] = util.convert_reference_frame(
                                palm_poses_world[key], obj_pose_world, util.unit_pose()
                            )
                        example_args['palm_pose_r_object'] = palm_poses_obj_frame['right']
                        example_args['palm_pose_l_object'] = palm_poses_obj_frame['left']
                        example_args['object_pose1_world'] = obj_pose_world
                        example_args['object_pose2_world'] = obj_pose_final
                        example_args['table_face'] = init_id
                    else:
                        continue

                try:
                    obs, pcd = yumi_gs.get_observation(obj_id=obj_id)

                    start = util.pose_stamped2list(obj_pose_world)
                    goal = util.pose_stamped2list(obj_pose_final)

                    keypoints_start = np.array(exp_running.mesh_world.vertices.tolist())
                    keypoints_start_homog = np.hstack(
                        (keypoints_start, np.ones((keypoints_start.shape[0], 1)))
                    )

                    start_mat = util.matrix_from_pose(obj_pose_world)
                    goal_mat = util.matrix_from_pose(obj_pose_final)

                    T_mat = np.matmul(goal_mat, np.linalg.inv(start_mat))
                    keypoints_goal = np.matmul(T_mat, keypoints_start_homog.T).T[:, :3]

                    contact_obj_frame_dict = {}
                    contact_world_frame_dict = {}
                    nearest_pt_world_dict = {}

                    if primitive_name == 'pull':
                        active_arm, inactive_arm = action_planner.get_active_arm(
                            util.pose_stamped2list(obj_pose_world)
                        )

                        # get contact (palm pose object frame)
                        contact_obj_frame_dict[active_arm] = util.pose_stamped2list(contact_obj_frame)
                        contact_world_frame_dict[active_arm] = util.pose_stamped2list(palm_pose_world)
                        contact_pos = open3d.utility.DoubleVector(np.array(contact_world_frame_dict[active_arm][:3]))
                        kdtree = open3d.geometry.KDTreeFlann(pcd)
                        # nearest_pt_ind = kdtree.search_knn_vector_3d(contact_pos, 1)[1][0]
                        # nearest_pt_world_dict[active_arm] = np.asarray(pcd.points)[nearest_pt_ind]

                        contact_obj_frame_dict[inactive_arm] = None
                        contact_world_frame_dict[inactive_arm] = None
                        nearest_pt_world_dict[inactive_arm] = None
                    elif primitive_name == 'grasp':
                        for key in palm_poses_obj_frame.keys():
                            contact_obj_frame_dict[key] = util.pose_stamped2list(palm_poses_obj_frame[key])
                            contact_world_frame_dict[key] = util.pose_stamped2list(palm_poses_world[key])
                            contact_pos = open3d.utility.DoubleVector(np.array(contact_world_frame_dict[key][:3]))
                            kdtree = open3d.geometry.KDTreeFlann(pcd)
                            # nearest_pt_ind = kdtree.search_knn_vector_3d(contact_pos, 1)[1][0]
                            # nearest_pt_world_dict[key] = np.asarray(pcd.points)[nearest_pt_ind]

                    result = action_planner.execute(primitive_name, example_args)
                    if result is not None:
                        print('Trial number: ' + str(trial) + ', reached final: ' + str(result[0]))
                        print('Time so far: ' + str(time.time() - start_time))

                        if result[0]:
                            success += 1
                            sample = {}
                            sample['obs'] = obs
                            sample['start'] = start
                            sample['goal'] = goal
                            sample['keypoints_start'] = keypoints_start
                            sample['keypoints_goal'] = keypoints_goal
                            sample['transformation'] = util.pose_from_matrix(T_mat)
                            sample['contact_obj_frame'] = contact_obj_frame_dict
                            sample['contact_world_frame'] = contact_world_frame_dict
                            # sample['contact_pcd'] = nearest_pt_world_dict
                            sample['result'] = result
                            if primitive_name == 'grasp':
                                sample['goal_face'] = exp_double.goal_face

                            if args.save_data:
                                data_manager.save_observation(sample, str(trial))
                            print("Success: " + str(success))
                    else:
                        continue

                    # data['saved_data'].append(sample)
                #     if trial == 0:
                #         parent1.send("SAMPLE")
                #     elif trial == 1:
                #         parent2.send("SAMPLE")
                #     result = work_queue.get(block=True)

                # if trial == 0:
                #     parent1.send("SAMPLE")
                # elif trial == 1:
                #     parent2.send("SAMPLE")
                # parent1.send("SAMPLE")
                # parent2.send("SAMPLE")

                # start = time.time()
                # done = False
                # result_list = []
                # while (time.time() - start) < cfg.TIMEOUT and not done:
                #     try:
                #         result = result_queue.get(block=True)
                #         result_list.append(result)
                #         if len(result_list) == 2:
                #             done = True
                #     except result_queue.Empty:
                #         continue
                #     time.sleep(0.001)
                except ValueError as e:
                    print("Value error: ")
                    print(e)

                # time.sleep(1.0)

                # pose = util.pose_stamped2list(yumi_gs.compute_fk(yumi_gs.get_jpos(arm='right')))
                # pos, ori = pose[:3], pose[3:]

                # # pose = yumi_gs.get_ee_pose()
                # # pos, ori = pose[0], pose[1]
                # # pos[2] -= 0.0714
                # pos[2] += 0.001
                # r_jnts = yumi_gs.compute_ik(pos, ori, yumi_gs.get_jpos(arm='right'))
                # l_jnts = yumi_gs.get_jpos(arm='left')


                # if r_jnts is not None:
                #     for _ in range(10):
                #         pos[2] += 0.001
                #         r_jnts = yumi_gs.compute_ik(pos, ori, yumi_gs.get_jpos(arm='right'))
                #         l_jnts = yumi_gs.get_jpos(arm='left')

                #         if r_jnts is not None:
                #             yumi_gs.update_joints(list(r_jnts) + l_jnts)
                #         time.sleep(0.1)

                time.sleep(0.1)
                for _ in range(10):
                    yumi_gs.update_joints(cfg.RIGHT_INIT + cfg.LEFT_INIT)

                # for _ in range(10):
                #     j_pos = cfg.RIGHT_INIT + cfg.LEFT_INIT
                #     for ind, jnt_id in enumerate(yumi_ar.arm.arm_jnt_ids):
                #         p.resetJointState(
                #             yumi_ar.arm.robot_id,
                #             jnt_id,
                #             targetValue=j_pos[ind]
                #         )

                # yumi_gs.update_joints(cfg.RIGHT_INIT + cfg.LEFT_INIT)

                # p.resetJointStatesMultiDof(
                #     yumi_ar.arm.robot_id,
                #     yumi_ar.arm.arm_jnt_ids,
                #     targetValues=j_pos)
                # parent1.send("HOME")
                # parent2.send("HOME")

                # time.sleep(1.0)

    print("TOTAL TIME: " + str(time.time() - global_start))
Beispiel #2
0
def main(args):
    cfg_file = os.path.join(args.example_config_path, args.primitive) + ".yaml"
    cfg = get_cfg_defaults()
    cfg.merge_from_file(cfg_file)
    cfg.freeze()

    rospy.init_node('MacroActions')
    signal.signal(signal.SIGINT, signal_handler)

    data_seed = args.np_seed
    primitive_name = args.primitive

    pickle_path = os.path.join(args.data_dir, primitive_name,
                               args.experiment_name)

    if args.save_data:
        suf_i = 0
        original_pickle_path = pickle_path
        while True:
            if os.path.exists(pickle_path):
                suffix = '_%d' % suf_i
                pickle_path = original_pickle_path + suffix
                suf_i += 1
                data_seed += 1
            else:
                os.makedirs(pickle_path)
                break

        if not os.path.exists(pickle_path):
            os.makedirs(pickle_path)

    np.random.seed(data_seed)

    yumi_ar = Robot('yumi_palms',
                    pb=True,
                    pb_cfg={
                        'gui': args.visualize,
                        'opengl_render': False
                    },
                    arm_cfg={
                        'self_collision': False,
                        'seed': data_seed
                    })

    r_gel_id = cfg.RIGHT_GEL_ID
    l_gel_id = cfg.LEFT_GEL_ID
    table_id = cfg.TABLE_ID

    alpha = cfg.ALPHA
    K = cfg.GEL_CONTACT_STIFFNESS
    restitution = cfg.GEL_RESTITUION

    p.changeDynamics(yumi_ar.arm.robot_id,
                     r_gel_id,
                     restitution=restitution,
                     contactStiffness=K,
                     contactDamping=alpha * K,
                     rollingFriction=args.rolling)

    p.changeDynamics(yumi_ar.arm.robot_id,
                     l_gel_id,
                     restitution=restitution,
                     contactStiffness=K,
                     contactDamping=alpha * K,
                     rollingFriction=args.rolling)

    yumi_gs = YumiCamsGS(yumi_ar,
                         cfg,
                         exec_thread=False,
                         sim_step_repeat=args.sim_step_repeat)

    for _ in range(10):
        yumi_gs.update_joints(cfg.RIGHT_INIT + cfg.LEFT_INIT)

    cuboid_sampler = CuboidSampler(os.path.join(
        os.environ['CODE_BASE'],
        'catkin_ws/src/config/descriptions/meshes/objects/cuboids/nominal_cuboid.stl'
    ),
                                   pb_client=yumi_ar.pb_client)
    cuboid_fname_template = os.path.join(
        os.environ['CODE_BASE'],
        'catkin_ws/src/config/descriptions/meshes/objects/cuboids/')

    cuboid_manager = MultiBlockManager(cuboid_fname_template,
                                       cuboid_sampler,
                                       robot_id=yumi_ar.arm.robot_id,
                                       table_id=27,
                                       r_gel_id=r_gel_id,
                                       l_gel_id=l_gel_id)

    if args.multi:
        cuboid_fname = cuboid_manager.get_cuboid_fname()
        # cuboid_fname = '/root/catkin_ws/src/config/descriptions/meshes/objects/cuboids/test_cuboid_smaller_4479.stl'
    else:
        cuboid_fname = args.config_package_path + 'descriptions/meshes/objects/' + \
            args.object_name + '_experiments.stl'
    mesh_file = cuboid_fname
    print("Cuboid file: " + cuboid_fname)

    if args.goal_viz:
        goal_visualization = True
    else:
        goal_visualization = False

    obj_id, sphere_ids, mesh, goal_obj_id = \
        cuboid_sampler.sample_cuboid_pybullet(
            cuboid_fname,
            goal=goal_visualization,
            keypoints=False)

    cuboid_manager.filter_collisions(obj_id, goal_obj_id)

    p.changeDynamics(obj_id, -1, lateralFriction=0.4)

    # goal_face = 0
    goal_faces = [0, 1, 2, 3, 4, 5]
    from random import shuffle
    shuffle(goal_faces)
    goal_face = goal_faces[0]

    exp_single = SingleArmPrimitives(cfg, yumi_ar.pb_client.get_client_id(),
                                     obj_id, cuboid_fname)
    k = 0
    while True:
        k += 1
        if k > 10:
            print('FAILED TO BUILD GRASPING GRAPH')
            return
        try:
            exp_double = DualArmPrimitives(cfg,
                                           yumi_ar.pb_client.get_client_id(),
                                           obj_id,
                                           cuboid_fname,
                                           goal_face=goal_face)
            break
        except ValueError as e:
            print(e)
            yumi_ar.pb_client.remove_body(obj_id)
            if goal_visualization:
                yumi_ar.pb_client.remove_body(goal_obj_id)
            cuboid_fname = cuboid_manager.get_cuboid_fname()
            print("Cuboid file: " + cuboid_fname)

            obj_id, sphere_ids, mesh, goal_obj_id = \
                cuboid_sampler.sample_cuboid_pybullet(
                    cuboid_fname,
                    goal=goal_visualization,
                    keypoints=False)

            cuboid_manager.filter_collisions(obj_id, goal_obj_id)

            p.changeDynamics(obj_id, -1, lateralFriction=0.4)
    if primitive_name == 'grasp':
        exp_running = exp_double
    else:
        exp_running = exp_single

    action_planner = ClosedLoopMacroActions(cfg,
                                            yumi_gs,
                                            obj_id,
                                            yumi_ar.pb_client.get_client_id(),
                                            args.config_package_path,
                                            replan=args.replan,
                                            object_mesh_file=mesh_file)

    if goal_visualization:
        trans_box_lock = threading.RLock()
        goal_viz = GoalVisual(trans_box_lock, goal_obj_id,
                              action_planner.pb_client, cfg.OBJECT_POSE_3)

    action_planner.update_object(obj_id, mesh_file)
    exp_single.initialize_object(obj_id, cuboid_fname)

    dynamics_info = {}
    dynamics_info['contactDamping'] = alpha * K
    dynamics_info['contactStiffness'] = K
    dynamics_info['rollingFriction'] = args.rolling
    dynamics_info['restitution'] = restitution

    data = {}
    data['saved_data'] = []
    data['metadata'] = {}
    data['metadata']['mesh_file'] = mesh_file
    data['metadata']['cfg'] = cfg
    data['metadata']['dynamics'] = dynamics_info
    data['metadata']['cam_cfg'] = yumi_gs.cam_setup_cfg
    data['metadata']['step_repeat'] = args.sim_step_repeat
    data['metadata']['seed'] = data_seed
    data['metadata']['seed_original'] = args.np_seed

    metadata = data['metadata']

    data_manager = DataManager(pickle_path)
    pred_dir = os.path.join(os.getcwd(), 'predictions')
    obs_dir = os.path.join(os.getcwd(), 'observations')

    if args.save_data:
        with open(os.path.join(pickle_path, 'metadata.pkl'), 'wb') as mdata_f:
            pickle.dump(metadata, mdata_f)

    total_trials = 0
    successes = 0
    for _ in range(args.num_blocks):
        # for goal_face in goal_faces:
        for _ in range(1):
            goal_face = np.random.randint(6)
            try:
                print('New object!')
                exp_double.initialize_object(obj_id, cuboid_fname, goal_face)
            except ValueError as e:
                print(e)
                print('Goal face: ' + str(goal_face))
                continue
            for _ in range(args.num_obj_samples):
                total_trials += 1
                if primitive_name == 'grasp':
                    start_face = exp_double.get_valid_ind()
                    if start_face is None:
                        print('Could not find valid start face')
                        continue
                    plan_args = exp_double.get_random_primitive_args(
                        ind=start_face, random_goal=True, execute=True)
                elif primitive_name == 'pull':
                    plan_args = exp_single.get_random_primitive_args(
                        ind=goal_face, random_goal=True, execute=True)

                start_pose = plan_args['object_pose1_world']
                goal_pose = plan_args['object_pose2_world']

                if goal_visualization:
                    goal_viz.update_goal_state(
                        util.pose_stamped2list(goal_pose))
                if args.debug:
                    import simulation

                    plan = action_planner.get_primitive_plan(
                        primitive_name, plan_args, 'right')

                    for i in range(10):
                        simulation.visualize_object(
                            start_pose,
                            filepath=
                            "package://config/descriptions/meshes/objects/cuboids/"
                            + cuboid_fname.split('objects/cuboids')[1],
                            name="/object_initial",
                            color=(1., 0., 0., 1.),
                            frame_id="/yumi_body",
                            scale=(1., 1., 1.))
                        simulation.visualize_object(
                            goal_pose,
                            filepath=
                            "package://config/descriptions/meshes/objects/cuboids/"
                            + cuboid_fname.split('objects/cuboids')[1],
                            name="/object_final",
                            color=(0., 0., 1., 1.),
                            frame_id="/yumi_body",
                            scale=(1., 1., 1.))
                        rospy.sleep(.1)
                    simulation.simulate(
                        plan,
                        cuboid_fname.split('objects/cuboids')[1])
                else:
                    success = False
                    attempts = 0

                    while True:
                        attempts += 1
                        time.sleep(0.1)
                        for _ in range(10):
                            yumi_gs.update_joints(cfg.RIGHT_INIT +
                                                  cfg.LEFT_INIT)

                        p.resetBasePositionAndOrientation(
                            obj_id,
                            util.pose_stamped2list(start_pose)[:3],
                            util.pose_stamped2list(start_pose)[3:])

                        if attempts > 15:
                            break
                        print('attempts: ' + str(attempts))
                        try:
                            obs, pcd = yumi_gs.get_observation(
                                obj_id=obj_id,
                                robot_table_id=(yumi_ar.arm.robot_id,
                                                table_id))

                            obj_pose_world = start_pose
                            obj_pose_final = goal_pose

                            start = util.pose_stamped2list(obj_pose_world)
                            goal = util.pose_stamped2list(obj_pose_final)

                            start_mat = util.matrix_from_pose(obj_pose_world)
                            goal_mat = util.matrix_from_pose(obj_pose_final)

                            T_mat = np.matmul(goal_mat,
                                              np.linalg.inv(start_mat))

                            transformation = np.asarray(util.pose_stamped2list(
                                util.pose_from_matrix(T_mat)),
                                                        dtype=np.float32)
                            # model takes in observation, and predicts:
                            pointcloud_pts = np.asarray(
                                obs['down_pcd_pts'][:100, :], dtype=np.float32)
                            obs_fname = os.path.join(
                                obs_dir,
                                str(total_trials) + '.npz')
                            np.savez(obs_fname,
                                     pointcloud_pts=pointcloud_pts,
                                     transformation=transformation)

                            # embed()

                            got_file = False
                            pred_fname = os.path.join(
                                pred_dir,
                                str(total_trials) + '.npz')
                            start = time.time()
                            while True:
                                try:
                                    prediction = np.load(pred_fname)
                                    got_file = True
                                except:
                                    pass
                                if got_file or (time.time() - start > 300):
                                    break
                                time.sleep(0.01)
                            if not got_file:
                                wait = raw_input(
                                    'waiting for predictions to come back online'
                                )
                                continue
                            os.remove(pred_fname)
                            # embed()

                            ind = np.random.randint(10)
                            # contact_obj_frame_r = prediction['prediction'][ind, :7]
                            # contact_obj_frame_l = prediction['prediction'][ind, 7:]
                            contact_prediction_r = prediction['prediction'][
                                ind, :7]
                            contact_prediction_l = prediction['prediction'][
                                ind, 7:]
                            contact_world_pos_r = contact_prediction_r[:3] + np.mean(
                                pointcloud_pts, axis=0)
                            contact_world_pos_l = contact_prediction_l[:3] + np.mean(
                                pointcloud_pts, axis=0)

                            contact_world_pos_pred = {}
                            contact_world_pos_pred[
                                'right'] = contact_world_pos_r
                            contact_world_pos_pred[
                                'left'] = contact_world_pos_l

                            contact_world_pos_corr = correct_grasp_pos(
                                contact_world_pos_pred, obs['pcd_pts'])

                            contact_world_pos_r = contact_world_pos_corr[
                                'right']
                            contact_world_pos_l = contact_world_pos_corr[
                                'left']

                            contact_world_r = contact_world_pos_r.tolist(
                            ) + contact_prediction_r[3:].tolist()
                            contact_world_l = contact_world_pos_l.tolist(
                            ) + contact_prediction_l[3:].tolist()

                            palm_poses_world = {}
                            # palm_poses_world['right'] = util.convert_reference_frame(
                            #     util.list2pose_stamped(contact_obj_frame_r),
                            #     util.unit_pose(),
                            #     obj_pose_world)
                            # palm_poses_world['left'] = util.convert_reference_frame(
                            #     util.list2pose_stamped(contact_obj_frame_l),
                            #     util.unit_pose(),
                            #     obj_pose_world)
                            palm_poses_world['right'] = util.list2pose_stamped(
                                contact_world_r)
                            palm_poses_world['left'] = util.list2pose_stamped(
                                contact_world_l)

                            # obj_pose_final = self.goal_pose_world_frame_mod``
                            palm_poses_obj_frame = {}
                            # delta = 10e-3
                            penetration_delta = 7.5e-3
                            # delta = np.random.random_sample() * \
                            #     (penetration_delta - 0.5*penetration_delta) + \
                            #     penetration_delta
                            delta = penetration_delta
                            y_normals = action_planner.get_palm_y_normals(
                                palm_poses_world)
                            for key in palm_poses_world.keys():
                                # try to penetrate the object a small amount
                                palm_poses_world[
                                    key].pose.position.x -= delta * y_normals[
                                        key].pose.position.x
                                palm_poses_world[
                                    key].pose.position.y -= delta * y_normals[
                                        key].pose.position.y
                                palm_poses_world[
                                    key].pose.position.z -= delta * y_normals[
                                        key].pose.position.z

                                palm_poses_obj_frame[
                                    key] = util.convert_reference_frame(
                                        palm_poses_world[key], obj_pose_world,
                                        util.unit_pose())

                            # plan_args['palm_pose_r_object'] = util.list2pose_stamped(contact_obj_frame_r)
                            # plan_args['palm_pose_l_object'] = util.list2pose_stamped(contact_obj_frame_l)
                            plan_args[
                                'palm_pose_r_object'] = palm_poses_obj_frame[
                                    'right']
                            plan_args[
                                'palm_pose_l_object'] = palm_poses_obj_frame[
                                    'left']

                            plan = action_planner.get_primitive_plan(
                                primitive_name, plan_args, 'right')

                            # import simulation
                            # for i in range(10):
                            #     simulation.visualize_object(
                            #         start_pose,
                            #         filepath="package://config/descriptions/meshes/objects/cuboids/" +
                            #             cuboid_fname.split('objects/cuboids')[1],
                            #         name="/object_initial",
                            #         color=(1., 0., 0., 1.),
                            #         frame_id="/yumi_body",
                            #         scale=(1., 1., 1.))
                            #     simulation.visualize_object(
                            #         goal_pose,
                            #         filepath="package://config/descriptions/meshes/objects/cuboids/" +
                            #             cuboid_fname.split('objects/cuboids')[1],
                            #         name="/object_final",
                            #         color=(0., 0., 1., 1.),
                            #         frame_id="/yumi_body",
                            #         scale=(1., 1., 1.))
                            #     rospy.sleep(.1)
                            # simulation.simulate(plan, cuboid_fname.split('objects/cuboids')[1])
                            # continue

                            result = action_planner.execute(
                                primitive_name, plan_args)
                            if result is None:
                                continue
                            print('Result: ' + str(result[0]) +
                                  ' Pos Error: ' + str(result[1]) +
                                  ' Ori Error: ' + str(result[2]))
                            if result[0]:
                                successes += 1
                                print('Success rate: ' +
                                      str(successes * 100.0 / total_trials))
                            break

                        except ValueError as e:
                            print("Value error: ")
                            print(e)

                    if args.nice_pull_release:
                        time.sleep(1.0)

                        pose = util.pose_stamped2list(
                            yumi_gs.compute_fk(yumi_gs.get_jpos(arm='right')))
                        pos, ori = pose[:3], pose[3:]

                        pos[2] += 0.001
                        r_jnts = yumi_gs.compute_ik(
                            pos, ori, yumi_gs.get_jpos(arm='right'))
                        l_jnts = yumi_gs.get_jpos(arm='left')

                        if r_jnts is not None:
                            for _ in range(10):
                                pos[2] += 0.001
                                r_jnts = yumi_gs.compute_ik(
                                    pos, ori, yumi_gs.get_jpos(arm='right'))
                                l_jnts = yumi_gs.get_jpos(arm='left')

                                if r_jnts is not None:
                                    yumi_gs.update_joints(
                                        list(r_jnts) + l_jnts)
                                time.sleep(0.1)

                    time.sleep(0.1)
                    for _ in range(10):
                        yumi_gs.update_joints(cfg.RIGHT_INIT + cfg.LEFT_INIT)

        while True:
            try:
                yumi_ar.pb_client.remove_body(obj_id)
                if goal_visualization:
                    yumi_ar.pb_client.remove_body(goal_obj_id)
                cuboid_fname = cuboid_manager.get_cuboid_fname()
                print("Cuboid file: " + cuboid_fname)

                obj_id, sphere_ids, mesh, goal_obj_id = \
                    cuboid_sampler.sample_cuboid_pybullet(
                        cuboid_fname,
                        goal=goal_visualization,
                        keypoints=False)

                cuboid_manager.filter_collisions(obj_id, goal_obj_id)

                p.changeDynamics(obj_id, -1, lateralFriction=0.4)
                action_planner.update_object(obj_id, mesh_file)
                exp_single.initialize_object(obj_id, cuboid_fname)
                # exp_double.initialize_object(obj_id, cuboid_fname, goal_face)
                if goal_visualization:
                    goal_viz.update_goal_obj(goal_obj_id)
                break
            except ValueError as e:
                print(e)
def worker_yumi(child_conn, work_queue, result_queue, cfg, args, seed,
                worker_id, global_info_dict, worker_flag_dict, planning_lock):
    while True:
        try:
            if not child_conn.poll(0.0001):
                continue
            msg = child_conn.recv()
        except (EOFError, KeyboardInterrupt):
            break
        if msg == "RESET":
            data_seed = seed
            np.random.seed(data_seed)

            yumi_ar = Robot('yumi_palms',
                            pb=True,
                            pb_cfg={'gui': args.visualize},
                            arm_cfg={
                                'self_collision': False,
                                'seed': data_seed
                            })

            r_gel_id = cfg.RIGHT_GEL_ID
            l_gel_id = cfg.LEFT_GEL_ID

            alpha = cfg.ALPHA
            K = cfg.GEL_CONTACT_STIFFNESS
            restitution = cfg.GEL_RESTITUION

            p.changeDynamics(yumi_ar.arm.robot_id,
                             r_gel_id,
                             restitution=restitution,
                             contactStiffness=K,
                             contactDamping=alpha * K,
                             rollingFriction=args.rolling)

            p.changeDynamics(yumi_ar.arm.robot_id,
                             l_gel_id,
                             restitution=restitution,
                             contactStiffness=K,
                             contactDamping=alpha * K,
                             rollingFriction=args.rolling)
            dynamics_info = {}
            dynamics_info['contactDamping'] = alpha * K
            dynamics_info['contactStiffness'] = K
            dynamics_info['rollingFriction'] = args.rolling
            dynamics_info['restitution'] = restitution

            yumi_gs = YumiCamsGS(yumi_ar,
                                 cfg,
                                 exec_thread=False,
                                 sim_step_repeat=args.step_repeat)

            for _ in range(10):
                yumi_gs.update_joints(cfg.RIGHT_INIT + cfg.LEFT_INIT)

            obj_id = yumi_ar.pb_client.load_urdf(
                args.config_package_path + 'descriptions/urdf/' +
                args.object_name + '.urdf', cfg.OBJECT_POSE_3[0:3],
                cfg.OBJECT_POSE_3[3:])

            # goal_obj_id = yumi_ar.pb_client.load_urdf(
            #     args.config_package_path +
            #     'descriptions/urdf/'+args.object_name+'_trans.urdf',
            #     cfg.OBJECT_POSE_3[0:3],
            #     cfg.OBJECT_POSE_3[3:]
            # )
            # p.setCollisionFilterPair(yumi_ar.arm.robot_id, goal_obj_id, r_gel_id, -1, enableCollision=False)
            # p.setCollisionFilterPair(obj_id, goal_obj_id, -1, -1, enableCollision=False)
            p.setCollisionFilterPair(yumi_ar.arm.robot_id,
                                     obj_id,
                                     r_gel_id,
                                     -1,
                                     enableCollision=True)
            p.setCollisionFilterPair(yumi_ar.arm.robot_id,
                                     obj_id,
                                     27,
                                     -1,
                                     enableCollision=True)

            yumi_ar.pb_client.reset_body(obj_id, cfg.OBJECT_POSE_3[:3],
                                         cfg.OBJECT_POSE_3[3:])

            # yumi_ar.pb_client.reset_body(
            #     goal_obj_id,
            #     cfg.OBJECT_POSE_3[:3],
            #     cfg.OBJECT_POSE_3[3:])

            primitive_name = args.primitive
            mesh_file = '/root/catkin_ws/src/config/descriptions/meshes/objects/cuboids/realsense_box_experiments.stl'
            cuboid_fname = mesh_file
            face = 0

            # setup macro_planner
            action_planner = ClosedLoopMacroActions(
                cfg,
                yumi_gs,
                obj_id,
                yumi_ar.pb_client.get_client_id(),
                args.config_package_path,
                object_mesh_file=mesh_file,
                replan=args.replan)

            exp_single = SingleArmPrimitives(cfg,
                                             yumi_ar.pb_client.get_client_id(),
                                             obj_id, mesh_file)
            exp_double = DualArmPrimitives(cfg,
                                           yumi_ar.pb_client.get_client_id(),
                                           obj_id, mesh_file)

            if primitive_name == 'grasp' or primitive_name == 'pivot':
                exp_running = exp_double
            else:
                exp_running = exp_single

            action_planner.update_object(obj_id, mesh_file)
            exp_single.initialize_object(obj_id, cuboid_fname)
            exp_double.initialize_object(obj_id, cuboid_fname, face)

            # goal_viz = GoalVisual(
            #     trans_box_lock,
            #     goal_obj_id,
            #     action_planner.pb_client,
            #     cfg.OBJECT_POSE_3)

            pickle_path = os.path.join(args.data_dir, primitive_name,
                                       args.experiment_name)

            if not os.path.exists(pickle_path):
                os.makedirs(pickle_path)

            data_manager = DataManager(pickle_path)
            continue
        if msg == "HOME":
            yumi_gs.update_joints(cfg.RIGHT_INIT + cfg.LEFT_INIT)
            continue
        if msg == "OBJECT_POSE":
            continue
        if msg == "SAMPLE":
            global_info_dict['trial'] += 1
            # print('Success: ' + str(global_info_dict['success']) +
            #         ' Trial number: ' + str(global_info_dict['trial']) +
            #         ' Worker ID: ' + str(worker_id))

            worker_flag_dict[worker_id] = False

            success = False
            start_face = 1
            plan_args = exp_running.get_random_primitive_args(
                ind=start_face, random_goal=False, execute=True)
            try:
                obs, pcd = yumi_gs.get_observation(obj_id=obj_id)

                obj_pose_world = plan_args['object_pose1_world']
                obj_pose_final = plan_args['object_pose2_world']

                contact_obj_frame, contact_world_frame = {}, {}

                contact_obj_frame['right'] = plan_args['palm_pose_r_object']
                contact_world_frame['right'] = plan_args['palm_pose_r_world']
                contact_obj_frame['left'] = plan_args['palm_pose_l_object']
                contact_world_frame['left'] = plan_args['palm_pose_l_world']

                start = util.pose_stamped2list(obj_pose_world)
                goal = util.pose_stamped2list(obj_pose_final)

                keypoints_start = np.array(
                    exp_running.mesh_world.vertices.tolist())
                keypoints_start_homog = np.hstack(
                    (keypoints_start, np.ones((keypoints_start.shape[0], 1))))

                start_mat = util.matrix_from_pose(obj_pose_world)
                goal_mat = util.matrix_from_pose(obj_pose_final)

                T_mat = np.matmul(goal_mat, np.linalg.inv(start_mat))
                keypoints_goal = np.matmul(T_mat,
                                           keypoints_start_homog.T).T[:, :3]

                contact_obj_frame_dict = {}
                contact_world_frame_dict = {}
                nearest_pt_world_dict = {}
                corner_norms = {}
                down_pcd_norms = {}

                if primitive_name == 'pull':
                    # active_arm, inactive_arm = action_planner.get_active_arm(
                    #     util.pose_stamped2list(obj_pose_world)
                    # )
                    active_arm = action_planner.active_arm
                    inactive_arm = action_planner.inactive_arm

                    contact_obj_frame_dict[
                        active_arm] = util.pose_stamped2list(
                            contact_obj_frame[active_arm])
                    contact_world_frame_dict[
                        active_arm] = util.pose_stamped2list(
                            contact_world_frame[active_arm])
                    # contact_pos = open3d.utility.DoubleVector(np.array(contact_world_frame_dict[active_arm][:3]))
                    # kdtree = open3d.geometry.KDTreeFlann(pcd)
                    # nearest_pt_ind = kdtree.search_knn_vector_3d(contact_pos, 1)[1][0]
                    # nearest_pt_world_dict[active_arm] = np.asarray(pcd.points)[nearest_pt_ind]
                    contact_pos = np.array(
                        contact_world_frame_dict[active_arm][:3])

                    corner_dists = (np.asarray(keypoints_start) - contact_pos)
                    corner_norms[active_arm] = np.linalg.norm(corner_dists,
                                                              axis=1)

                    down_pcd_dists = (obs['down_pcd_pts'] - contact_pos)
                    down_pcd_norms[active_arm] = np.linalg.norm(down_pcd_dists,
                                                                axis=1)

                    contact_obj_frame_dict[inactive_arm] = None
                    contact_world_frame_dict[inactive_arm] = None
                    nearest_pt_world_dict[inactive_arm] = None
                    corner_norms[inactive_arm] = None
                    down_pcd_norms[inactive_arm] = None
                elif primitive_name == 'grasp':
                    for key in contact_obj_frame.keys():
                        contact_obj_frame_dict[key] = util.pose_stamped2list(
                            contact_obj_frame[key])
                        contact_world_frame_dict[key] = util.pose_stamped2list(
                            contact_world_frame[key])

                        contact_pos = np.array(
                            contact_world_frame_dict[key][:3])

                        corner_dists = (np.asarray(keypoints_start) -
                                        contact_pos)
                        corner_norms[key] = np.linalg.norm(corner_dists,
                                                           axis=1)

                        down_pcd_dists = (obs['down_pcd_pts'] - contact_pos)
                        down_pcd_norms[key] = np.linalg.norm(down_pcd_dists,
                                                             axis=1)

                        # contact_pos = open3d.utility.DoubleVector(np.array(contact_world_frame_dict[key][:3]))
                        # kdtree = open3d.geometry.KDTreeFlann(pcd)
                        # nearest_pt_ind = kdtree.search_knn_vector_3d(contact_pos, 1)[1][0]
                        # nearest_pt_world_dict[key] = np.asarray(pcd.points)[nearest_pt_ind]

                print('Worker ID: %d, Waiting to plan...' % worker_id)
                while True:
                    time.sleep(0.01)
                    work_queue_empty = work_queue.empty()
                    if work_queue_empty:
                        planner_available = True
                        break
                if planner_available:
                    print('Worker ID: %d, Planning!...' % worker_id)
                    work_queue.put(True)
                    time.sleep(1.0)
                    result = action_planner.execute(primitive_name, plan_args)
                    time.sleep(1.0)
                    while not work_queue.empty():
                        work_queue.get()

                if result is not None:
                    if result[0]:
                        success = True
                        global_info_dict['success'] += 1
                        sample = {}
                        sample['obs'] = obs
                        sample['start'] = start
                        sample['goal'] = goal
                        sample['keypoints_start'] = keypoints_start
                        sample['keypoints_goal'] = keypoints_goal
                        sample['keypoint_dists'] = corner_norms
                        sample['down_pcd_pts'] = obs['down_pcd_pts']
                        sample['down_pcd_dists'] = down_pcd_norms
                        sample['transformation'] = util.pose_from_matrix(T_mat)
                        sample['contact_obj_frame'] = contact_obj_frame_dict
                        sample[
                            'contact_world_frame'] = contact_world_frame_dict
                        # sample['contact_pcd'] = nearest_pt_world_dict
                        sample['result'] = result
                        if primitive_name == 'grasp':
                            sample['goal_face'] = exp_double.goal_face

                        if args.save_data:
                            data_manager.save_observation(
                                sample, str(global_info_dict['trial']))
                        # print("Success: " + str(global_info_dict['success']))
            except ValueError as e:
                print("Value error: ")
                print(e)
                result = None
                time.sleep(1.0)
                while not work_queue.empty():
                    work_queue.get()
            worker_result = {}
            worker_result['result'] = result
            worker_result['id'] = worker_id
            worker_result['success'] = success
            result_queue.put(worker_result)
            child_conn.send('DONE')
            continue
        if msg == "END":
            break
        time.sleep(0.001)
    print('Breaking Worker ID: ' + str(worker_id))
    child_conn.close()
def main(args):
    # get configuration
    cfg_file = osp.join(args.example_config_path, args.primitive) + ".yaml"
    cfg = get_cfg_defaults()
    cfg.merge_from_file(cfg_file)
    cfg.freeze()

    rospy.init_node('EvalSubgoal')
    signal.signal(signal.SIGINT, signal_handler)

    # setup data saving paths
    data_seed = args.np_seed
    primitive_name = args.primitive

    problems_file = '/root/catkin_ws/src/primitives/data/planning/test_problems_0/demo_0.pkl'
    with open(problems_file, 'rb') as f:
        problems_data = pickle.load(f)

    prob_inds = np.arange(len(problems_data), dtype=np.int64).tolist()
    data_inds = np.arange(len(problems_data[0]['problems']),
                          dtype=np.int64).tolist()

    pickle_path = osp.join(args.data_dir, primitive_name, args.experiment_name)

    if args.save_data:
        suf_i = 0
        original_pickle_path = pickle_path
        # while True:
        #     if osp.exists(pickle_path):
        #         suffix = '_%d' % suf_i
        #         pickle_path = original_pickle_path + suffix
        #         suf_i += 1
        #         data_seed += 1
        #     else:
        #         os.makedirs(pickle_path)
        #         break

        if not osp.exists(pickle_path):
            os.makedirs(pickle_path)

    np.random.seed(data_seed)

    # initialize airobot and modify dynamics
    yumi_ar = Robot('yumi_palms',
                    pb=True,
                    pb_cfg={
                        'gui': args.visualize,
                        'opengl_render': False
                    },
                    arm_cfg={
                        'self_collision': False,
                        'seed': data_seed
                    })

    r_gel_id = cfg.RIGHT_GEL_ID
    l_gel_id = cfg.LEFT_GEL_ID
    table_id = cfg.TABLE_ID

    alpha = cfg.ALPHA
    K = cfg.GEL_CONTACT_STIFFNESS
    restitution = cfg.GEL_RESTITUION

    p.changeDynamics(yumi_ar.arm.robot_id,
                     r_gel_id,
                     restitution=restitution,
                     contactStiffness=K,
                     contactDamping=alpha * K,
                     rollingFriction=args.rolling,
                     lateralFriction=0.5)

    p.changeDynamics(yumi_ar.arm.robot_id,
                     l_gel_id,
                     restitution=restitution,
                     contactStiffness=K,
                     contactDamping=alpha * K,
                     rollingFriction=args.rolling,
                     lateralFriction=0.5)

    # initialize PyBullet + MoveIt! + ROS yumi interface
    yumi_gs = YumiCamsGS(yumi_ar,
                         cfg,
                         exec_thread=False,
                         sim_step_repeat=args.sim_step_repeat)

    yumi_ar.arm.go_home(ignore_physics=True)

    # initialize object sampler
    cuboid_sampler = CuboidSampler(osp.join(
        os.environ['CODE_BASE'],
        'catkin_ws/src/config/descriptions/meshes/objects/cuboids/nominal_cuboid.stl'
    ),
                                   pb_client=yumi_ar.pb_client)
    cuboid_fname_template = osp.join(
        os.environ['CODE_BASE'],
        'catkin_ws/src/config/descriptions/meshes/objects/cuboids/')

    cuboid_manager = MultiBlockManager(cuboid_fname_template,
                                       cuboid_sampler,
                                       robot_id=yumi_ar.arm.robot_id,
                                       table_id=table_id,
                                       r_gel_id=r_gel_id,
                                       l_gel_id=l_gel_id)

    if args.multi:
        # cuboid_fname = cuboid_manager.get_cuboid_fname()
        # cuboid_fname = str(osp.join(
        #     '/root/catkin_ws/src/config/descriptions/meshes/objects/cuboids',
        #     problems_data[0]['object_name']))

        # get object name
        k = 0
        prob_inds = copy.deepcopy(
            list(np.arange(len(problems_data), dtype=np.int64)))
        shuffle(prob_inds)
        while True:
            if len(prob_inds) == 0:
                print('Done with test problems!')
                return
            prob_ind = prob_inds.pop()
            obj_name = problems_data[prob_ind]['object_name'].split('.stl')[0]
            if osp.exists(osp.join(pickle_path, obj_name)):
                continue
            os.makedirs(osp.join(pickle_path, obj_name))
            break
        cuboid_fname = str(
            osp.join(
                '/root/catkin_ws/src/config/descriptions/meshes/objects/cuboids',
                obj_name + '.stl'))

    else:
        cuboid_fname = args.config_package_path + 'descriptions/meshes/objects/' + \
            args.object_name + '.stl'
    mesh_file = cuboid_fname
    print("Cuboid file: " + cuboid_fname)

    if args.goal_viz:
        goal_visualization = True
    else:
        goal_visualization = False

    obj_id, sphere_ids, mesh, goal_obj_id = \
        cuboid_sampler.sample_cuboid_pybullet(
            cuboid_fname,
            goal=goal_visualization,
            keypoints=False)

    cuboid_manager.filter_collisions(obj_id, goal_obj_id)

    p.changeDynamics(obj_id, -1, lateralFriction=1.0)

    goal_faces = [0, 1, 2, 3, 4, 5]
    # shuffle(goal_faces)
    goal_face = goal_faces[0]

    # initialize primitive args samplers
    exp_single = SingleArmPrimitives(cfg, yumi_ar.pb_client.get_client_id(),
                                     obj_id, cuboid_fname)
    k = 0
    while True:
        k += 1
        if k > 10:
            print('FAILED TO BUILD GRASPING GRAPH')
            return
        try:
            exp_double = DualArmPrimitives(cfg,
                                           yumi_ar.pb_client.get_client_id(),
                                           obj_id,
                                           cuboid_fname,
                                           goal_face=goal_face)
            break
        except ValueError as e:
            print(e)
            yumi_ar.pb_client.remove_body(obj_id)
            if goal_visualization:
                yumi_ar.pb_client.remove_body(goal_obj_id)
            cuboid_fname = cuboid_manager.get_cuboid_fname()
            print("Cuboid file: " + cuboid_fname)

            obj_id, sphere_ids, mesh, goal_obj_id = \
                cuboid_sampler.sample_cuboid_pybullet(
                    cuboid_fname,
                    goal=goal_visualization,
                    keypoints=False)

            cuboid_manager.filter_collisions(obj_id, goal_obj_id)

            p.changeDynamics(obj_id, -1, lateralFriction=1.0)
    if primitive_name == 'grasp':
        exp_running = exp_double
    else:
        exp_running = exp_single

    # initialize macro action interface
    action_planner = ClosedLoopMacroActions(cfg,
                                            yumi_gs,
                                            obj_id,
                                            yumi_ar.pb_client.get_client_id(),
                                            args.config_package_path,
                                            replan=args.replan,
                                            object_mesh_file=mesh_file)

    if goal_visualization:
        trans_box_lock = threading.RLock()
        goal_viz = GoalVisual(trans_box_lock, goal_obj_id,
                              action_planner.pb_client, cfg.OBJECT_POSE_3)

    action_planner.update_object(obj_id, mesh_file)
    exp_single.initialize_object(obj_id, cuboid_fname)

    # prep save info
    dynamics_info = {}
    dynamics_info['contactDamping'] = alpha * K
    dynamics_info['contactStiffness'] = K
    dynamics_info['rollingFriction'] = args.rolling
    dynamics_info['restitution'] = restitution

    data = {}
    data['saved_data'] = []
    data['metadata'] = {}
    data['metadata']['mesh_file'] = mesh_file
    data['metadata']['cfg'] = cfg
    data['metadata']['dynamics'] = dynamics_info
    data['metadata']['cam_cfg'] = yumi_gs.cam_setup_cfg
    data['metadata']['step_repeat'] = args.sim_step_repeat
    data['metadata']['seed'] = data_seed
    data['metadata']['seed_original'] = args.np_seed

    metadata = data['metadata']

    data_manager = DataManager(pickle_path)
    # pred_dir = osp.join(os.environ['CODE_BASE'], cfg.PREDICTION_DIR)
    # obs_dir = osp.join(os.environ['CODE_BASE'], cfg.OBSERVATION_DIR)
    pred_dir = cfg.PREDICTION_DIR
    obs_dir = cfg.OBSERVATION_DIR
    if not osp.exists(pred_dir):
        os.makedirs(pred_dir)
    if not osp.exists(obs_dir):
        os.makedirs(obs_dir)

    if args.save_data:
        with open(osp.join(pickle_path, 'metadata.pkl'), 'wb') as mdata_f:
            pickle.dump(metadata, mdata_f)

    # prep visualization tools
    palm_mesh_file = osp.join(os.environ['CODE_BASE'], cfg.PALM_MESH_FILE)
    table_mesh_file = osp.join(os.environ['CODE_BASE'], cfg.TABLE_MESH_FILE)
    viz_palms = PalmVis(palm_mesh_file, table_mesh_file, cfg)
    viz_pcd = PCDVis()

    # pull_sampler = PullSamplerBasic()
    pull_sampler = PullSamplerVAEPubSub(obs_dir=obs_dir, pred_dir=pred_dir)
    grasp_sampler = GraspSamplerVAEPubSub(default_target=None,
                                          obs_dir=obs_dir,
                                          pred_dir=pred_dir,
                                          pointnet=args.pointnet)
    # grasp_sampler = GraspSamplerTransVAEPubSub(
    #     None,
    #     obs_dir,
    #     pred_dir,
    #     pointnet=args.pointnet
    # )
    # grasp_sampler = GraspSamplerBasic(
    #     default_target=None
    # )
    parent, child = Pipe(duplex=True)
    work_queue, result_queue = Queue(), Queue()

    experiment_manager = GraspEvalManager(yumi_gs,
                                          yumi_ar.pb_client.get_client_id(),
                                          pickle_path, args.exp_name, parent,
                                          child, work_queue, result_queue, cfg)

    experiment_manager.set_object_id(obj_id, cuboid_fname)

    # begin runs
    total_trials = 0
    total_executions = 0
    total_face_success = 0

    # for _ in range(args.num_blocks):
    for problem_ind in range(1, len(problems_data)):
        for goal_face in goal_faces:
            try:
                exp_double.initialize_object(obj_id, cuboid_fname, goal_face)
            except ValueError as e:
                print('Goal face: ' + str(goal_face), e)
                continue
            for _ in range(args.num_obj_samples):
                yumi_ar.arm.go_home(ignore_physics=True)
                obj_data = experiment_manager.get_object_data()
                if obj_data['trials'] > 0:
                    kvs = {}
                    kvs['trials'] = obj_data['trials']
                    kvs['grasp_success'] = obj_data[
                        'grasp_success'] * 100.0 / obj_data['trials']
                    kvs['mp_success'] = obj_data[
                        'mp_success'] * 100.0 / obj_data['trials']
                    kvs['face_success'] = obj_data[
                        'face_success'] * 100.0 / obj_data['trials']
                    kvs['pos_err'] = np.mean(obj_data['final_pos_error'])
                    kvs['ori_err'] = np.mean(obj_data['final_ori_error'])
                    string = ''

                    for k, v in kvs.items():
                        string += "%s: %.3f, " % (k, v)
                    print(string)

                total_trials += 1
                if primitive_name == 'grasp':
                    start_face = exp_double.get_valid_ind()
                    if start_face is None:
                        print('Could not find valid start face')
                        continue
                    plan_args = exp_double.get_random_primitive_args(
                        ind=start_face, random_goal=True, execute=True)
                elif primitive_name == 'pull':
                    plan_args = exp_single.get_random_primitive_args(
                        ind=goal_face, random_goal=True, execute=True)

                start_pose = plan_args['object_pose1_world']
                goal_pose = plan_args['object_pose2_world']

                if goal_visualization:
                    goal_viz.update_goal_state(
                        util.pose_stamped2list(goal_pose))
                    goal_viz.hide_goal_obj()
                attempts = 0

                # embed()
                # yumi_ar.pb_client.remove_body(obj_id)
                # start_pos = [0.4, 0.0, 0.1]
                # un_norm_ori = np.random.rand(4)
                # start_ori = un_norm_ori/(np.linalg.norm(un_norm_ori))
                # start_pose = util.list2pose_stamped(list(start_pos) + list(start_ori))
                # bandu_names = [
                #     '/root/catkin_ws/src/config/descriptions/bandu/Bandu Block/Bandu Block.urdf',
                #     '/root/catkin_ws/src/config/descriptions/bandu/Big Ring/Big Ring.urdf',
                #     '/root/catkin_ws/src/config/descriptions/bandu/Double Wedge/Double Wedge.urdf',
                #     '/root/catkin_ws/src/config/descriptions/bandu/Egg/Egg.urdf',
                #     '/root/catkin_ws/src/config/descriptions/bandu/Knight Shape/Knight Shape.urdf',
                #     '/root/catkin_ws/src/config/descriptions/bandu/Pencil/Pencil.urdf',
                #     '/root/catkin_ws/src/config/descriptions/bandu/Skewed Rectangular Prism/Skewed Rectangular Prism.urdf',
                #     '/root/catkin_ws/src/config/descriptions/bandu/Skewed Triangular Prism/Skewed Triangular Prism.urdf',
                #     '/root/catkin_ws/src/config/descriptions/bandu/Skewed Wedge/Skewed Wedge.urdf',
                # ]
                # obj_id = yumi_ar.pb_client.load_urdf(
                #     bandu_names[0],
                #     start_pos,
                #     start_ori
                # )

                # pcd1 = trimesh.PointCloud(pointcloud_pts)
                # pcd2 = trimesh.PointCloud(pointcloud_pts[np.where(start_state.pointcloud_mask)[0], :])
                # pcd1.colors = [255, 0, 0, 255]
                # pcd2.colors = [0, 0, 255, 255]
                # scene_full = trimesh.Scene([pcd1, pcd2])
                # scene1 = trimesh.Scene([pcd1])
                # scene2 = trimesh.Scene([pcd2])
                # scene_full.show()

                # embed()
                while True:
                    # if attempts > cfg.ATTEMPT_MAX:
                    if attempts > 4:
                        experiment_manager.set_mp_success(False, attempts)
                        experiment_manager.end_trial(None, False)
                        break
                    # print('attempts: ' + str(attempts))

                    attempts += 1
                    time.sleep(0.1)
                    yumi_ar.arm.go_home(ignore_physics=True)
                    if goal_visualization:
                        goal_viz.update_goal_state(
                            util.pose_stamped2list(goal_pose))
                        goal_viz.hide_goal_obj()
                    time.sleep(1.0)

                    p.resetBasePositionAndOrientation(
                        obj_id,
                        util.pose_stamped2list(start_pose)[:3],
                        util.pose_stamped2list(start_pose)[3:])
                    time.sleep(1.0)

                    obs, pcd = yumi_gs.get_observation(
                        obj_id=obj_id,
                        robot_table_id=(yumi_ar.arm.robot_id, table_id))

                    goal_pose_global = util.pose_stamped2list(goal_pose)
                    goal_mat_global = util.matrix_from_pose(goal_pose)

                    start_mat = util.matrix_from_pose(start_pose)
                    T_mat_global = np.matmul(goal_mat_global,
                                             np.linalg.inv(start_mat))

                    transformation_global = util.pose_stamped2np(
                        util.pose_from_matrix(T_mat_global))
                    # model takes in observation, and predicts:
                    pointcloud_pts = np.asarray(obs['down_pcd_pts'][:100, :],
                                                dtype=np.float32)
                    pointcloud_pts_full = np.asarray(np.concatenate(
                        obs['pcd_pts']),
                                                     dtype=np.float32)
                    table_pts_full = np.concatenate(obs['table_pcd_pts'],
                                                    axis=0)

                    grasp_sampler.update_default_target(
                        table_pts_full[::500, :])

                    # sample from model
                    start_state = PointCloudNode()
                    start_state.set_pointcloud(pcd=pointcloud_pts,
                                               pcd_full=pointcloud_pts_full)
                    if primitive_name == 'grasp':
                        # prediction = grasp_sampler.sample(start_state.pointcloud)
                        prediction = grasp_sampler.sample(
                            state=start_state.pointcloud,
                            state_full=start_state.pointcloud_full)
                    elif primitive_name == 'pull':
                        prediction = pull_sampler.sample(
                            start_state.pointcloud)
                    start_state.pointcloud_mask = prediction['mask']

                    new_state = PointCloudNode()
                    new_state.init_state(start_state,
                                         prediction['transformation'])
                    correction = False
                    if primitive_name == 'grasp':
                        correction = True
                    new_state.init_palms(
                        prediction['palms'],
                        correction=correction,
                        prev_pointcloud=start_state.pointcloud_full)

                    trans_execute = util.pose_from_matrix(
                        new_state.transformation)
                    if args.final_subgoal:
                        trans_execute = util.pose_from_matrix(T_mat_global)

                    if primitive_name == 'grasp':
                        local_plan = grasp_planning_wf(
                            util.list2pose_stamped(new_state.palms[7:]),
                            util.list2pose_stamped(new_state.palms[:7]),
                            trans_execute)
                    elif primitive_name == 'pull':
                        local_plan = pulling_planning_wf(
                            util.list2pose_stamped(new_state.palms[:7]),
                            util.list2pose_stamped(new_state.palms[:7]),
                            trans_execute)

                    if args.rviz_viz:
                        import simulation
                        for i in range(10):
                            simulation.visualize_object(
                                start_pose,
                                filepath=
                                "package://config/descriptions/meshes/objects/cuboids/"
                                + cuboid_fname.split('objects/cuboids')[1],
                                name="/object_initial",
                                color=(1., 0., 0., 1.),
                                frame_id="/yumi_body",
                                scale=(1., 1., 1.))
                            simulation.visualize_object(
                                goal_pose,
                                filepath=
                                "package://config/descriptions/meshes/objects/cuboids/"
                                + cuboid_fname.split('objects/cuboids')[1],
                                name="/object_final",
                                color=(0., 0., 1., 1.),
                                frame_id="/yumi_body",
                                scale=(1., 1., 1.))
                            rospy.sleep(.1)
                        simulation.simulate(
                            local_plan,
                            cuboid_fname.split('objects/cuboids')[1])
                        embed()
                    if args.plotly_viz:
                        plot_data = {}
                        plot_data['start'] = pointcloud_pts
                        plot_data[
                            'object_mask_down'] = start_state.pointcloud_mask

                        fig, _ = viz_pcd.plot_pointcloud(plot_data,
                                                         downsampled=True)
                        fig.show()
                        embed()
                    # embed()
                    # trans_list = []
                    # for i in range(50):
                    #     pred = pull_sampler.sample(start_state.pointcloud)
                    #     trans_list.append(util.pose_stamped2np(util.pose_from_matrix(pred['transformation'])))

                    if args.trimesh_viz:
                        viz_data = {}
                        viz_data[
                            'contact_world_frame_right'] = new_state.palms_raw[:
                                                                               7]
                        viz_data[
                            'contact_world_frame_left'] = new_state.palms_raw[
                                7:]
                        # viz_data['contact_world_frame_left'] = new_state.palms_raw[:7]
                        viz_data['start_vis'] = util.pose_stamped2np(
                            start_pose)
                        viz_data['transformation'] = util.pose_stamped2np(
                            util.pose_from_matrix(
                                prediction['transformation']))
                        # viz_data['transformation'] = np.asarray(trans_list).squeeze()
                        viz_data['mesh_file'] = cuboid_fname
                        viz_data['object_pointcloud'] = pointcloud_pts_full
                        viz_data['start'] = pointcloud_pts
                        # viz_data['start'] = pointcloud_pts_full
                        viz_data['object_mask'] = prediction['mask']
                        embed()

                        scene = viz_palms.vis_palms(viz_data,
                                                    world=True,
                                                    corr=False,
                                                    full_path=True,
                                                    goal_number=1)
                        scene_pcd = viz_palms.vis_palms_pcd(viz_data,
                                                            world=True,
                                                            corr=False,
                                                            full_path=True,
                                                            show_mask=True,
                                                            goal_number=1)
                        scene_pcd.show()
                        # scene.show()
                        embed()

                    real_start_pos = p.getBasePositionAndOrientation(obj_id)[0]
                    real_start_ori = p.getBasePositionAndOrientation(obj_id)[1]
                    real_start_pose = list(real_start_pos) + list(
                        real_start_ori)
                    real_start_mat = util.matrix_from_pose(
                        util.list2pose_stamped(real_start_pose))

                    des_goal_pose = util.transform_pose(
                        util.list2pose_stamped(real_start_pose),
                        util.pose_from_matrix(prediction['transformation']))

                    if goal_visualization:
                        goal_viz.update_goal_state(
                            util.pose_stamped2list(goal_pose))
                        goal_viz.show_goal_obj()

                    # create trial data
                    trial_data = {}
                    trial_data['start_pcd'] = pointcloud_pts_full
                    trial_data['start_pcd_down'] = pointcloud_pts
                    trial_data['start_pcd_mask'] = start_state.pointcloud_mask
                    trial_data['obj_fname'] = cuboid_fname
                    trial_data['start_pose'] = np.asarray(real_start_pose)
                    trial_data['goal_pose'] = util.pose_stamped2np(
                        des_goal_pose)
                    trial_data['goal_pose_global'] = np.asarray(
                        goal_pose_global)
                    trial_data['table_pcd'] = table_pts_full[::500, :]
                    trial_data['trans_des'] = util.pose_stamped2np(
                        util.pose_from_matrix(prediction['transformation']))
                    trial_data['trans_des_global'] = transformation_global

                    # experiment_manager.start_trial()
                    action_planner.active_arm = 'right'
                    action_planner.inactive_arm = 'left'

                    if primitive_name == 'grasp':
                        # try to execute the action
                        yumi_ar.arm.set_jpos([
                            0.9936, -2.1848, -0.9915, 0.8458, 3.7618, 1.5486,
                            0.1127, -1.0777, -2.1187, 0.995, 1.002, -3.6834,
                            1.8132, 2.6405
                        ],
                                             ignore_physics=True)
                        grasp_success = False
                        try:
                            for k, subplan in enumerate(local_plan):
                                time.sleep(1.0)
                                action_planner.playback_dual_arm(
                                    'grasp', subplan, k)
                                if k > 0 and experiment_manager.still_grasping(
                                ):
                                    grasp_success = True

                            real_final_pos = p.getBasePositionAndOrientation(
                                obj_id)[0]
                            real_final_ori = p.getBasePositionAndOrientation(
                                obj_id)[1]
                            real_final_pose = list(real_final_pos) + list(
                                real_final_ori)
                            real_final_mat = util.matrix_from_pose(
                                util.list2pose_stamped(real_final_pose))
                            real_T_mat = np.matmul(
                                real_final_mat, np.linalg.inv(real_start_mat))
                            real_T_pose = util.pose_stamped2np(
                                util.pose_from_matrix(real_T_mat))

                            trial_data['trans_executed'] = real_T_mat
                            trial_data['final_pose'] = real_final_pose
                            experiment_manager.set_mp_success(True, attempts)
                            experiment_manager.end_trial(
                                trial_data, grasp_success)
                            # embed()
                        except ValueError as e:
                            # print('Value Error: ', e)
                            continue
                    elif primitive_name == 'pull':
                        try:
                            yumi_ar.arm.set_jpos(cfg.RIGHT_INIT +
                                                 cfg.LEFT_INIT,
                                                 ignore_physics=True)
                            time.sleep(0.5)
                            action_planner.playback_single_arm(
                                'pull', local_plan[0])
                            time.sleep(0.5)
                            action_planner.single_arm_retract()
                        except ValueError as e:
                            continue

                    time.sleep(3.0)
                    yumi_ar.arm.go_home(ignore_physics=True)
                    break
        embed()

        obj_data = experiment_manager.get_object_data()
        # obj_name = problems_data[problem_ind]['object_name'].split('.stl')[0]
        obj_data_fname = osp.join(pickle_path, obj_name,
                                  obj_name + '_eval_data.pkl')
        # print('Object data: ')
        # for key in obj_data.keys():
        #     print(key, obj_data[key])
        if args.save_data:
            print('Saving to: ' + str(obj_data_fname))
            with open(obj_data_fname, 'wb') as f:
                pickle.dump(obj_data, f)

        yumi_ar.pb_client.remove_body(obj_id)
        if goal_visualization:
            yumi_ar.pb_client.remove_body(goal_obj_id)

        # cuboid_fname = cuboid_manager.get_cuboid_fname()
        # cuboid_fname = str(osp.join(
        #     '/root/catkin_ws/src/config/descriptions/meshes/objects/cuboids',
        #     problems_data[problem_ind]['object_name']))
        while True:
            if len(prob_inds) == 0:
                print('Done with test problems!')
                return
            prob_ind = prob_inds.pop()
            obj_name = problems_data[prob_ind]['object_name'].split('.stl')[0]
            if osp.exists(osp.join(pickle_path, obj_name)):
                continue
            os.makedirs(osp.join(pickle_path, obj_name))
            break
        cuboid_fname = str(
            osp.join(
                '/root/catkin_ws/src/config/descriptions/meshes/objects/cuboids',
                obj_name + '.stl'))

        obj_id, sphere_ids, mesh, goal_obj_id = \
            cuboid_sampler.sample_cuboid_pybullet(
                cuboid_fname,
                goal=goal_visualization,
                keypoints=False)

        cuboid_manager.filter_collisions(obj_id, goal_obj_id)

        p.changeDynamics(obj_id, -1, lateralFriction=1.0)

        action_planner.update_object(obj_id, cuboid_fname)
        exp_single.initialize_object(obj_id, cuboid_fname)
        experiment_manager.set_object_id(obj_id, cuboid_fname)

        if goal_visualization:
            goal_viz.update_goal_obj(goal_obj_id)
Beispiel #5
0
def main(args):
    cfg_file = osp.join(args.example_config_path, args.primitive) + ".yaml"
    cfg = get_cfg_defaults()
    cfg.merge_from_file(cfg_file)
    cfg.freeze()

    rospy.init_node('EvalMultiStep')
    signal.signal(signal.SIGINT, signal_handler)

    data_seed = args.np_seed
    primitive_name = args.primitive

    pickle_path = osp.join(args.data_dir, primitive_name, args.experiment_name)

    if args.save_data:
        suf_i = 0
        original_pickle_path = pickle_path
        # while True:
        #     if osp.exists(pickle_path):
        #         suffix = '_%d' % suf_i
        #         pickle_path = original_pickle_path + suffix
        #         suf_i += 1
        #         data_seed += 1
        #     else:
        #         os.makedirs(pickle_path)
        #         break

        if not osp.exists(pickle_path):
            os.makedirs(pickle_path)

    np.random.seed(data_seed)

    yumi_ar = Robot('yumi_palms',
                    pb=True,
                    pb_cfg={
                        'gui': args.visualize,
                        'opengl_render': False
                    },
                    arm_cfg={
                        'self_collision': False,
                        'seed': data_seed
                    })

    r_gel_id = cfg.RIGHT_GEL_ID
    l_gel_id = cfg.LEFT_GEL_ID
    table_id = cfg.TABLE_ID

    alpha = cfg.ALPHA
    K = cfg.GEL_CONTACT_STIFFNESS
    restitution = cfg.GEL_RESTITUION

    p.changeDynamics(yumi_ar.arm.robot_id,
                     r_gel_id,
                     restitution=restitution,
                     contactStiffness=K,
                     contactDamping=alpha * K,
                     rollingFriction=args.rolling)

    p.changeDynamics(yumi_ar.arm.robot_id,
                     l_gel_id,
                     restitution=restitution,
                     contactStiffness=K,
                     contactDamping=alpha * K,
                     rollingFriction=args.rolling)

    yumi_gs = YumiCamsGS(yumi_ar,
                         cfg,
                         exec_thread=False,
                         sim_step_repeat=args.sim_step_repeat)

    for _ in range(10):
        yumi_gs.update_joints(cfg.RIGHT_INIT + cfg.LEFT_INIT)

    cuboid_sampler = CuboidSampler(osp.join(
        os.environ['CODE_BASE'],
        'catkin_ws/src/config/descriptions/meshes/objects/cuboids/nominal_cuboid.stl'
    ),
                                   pb_client=yumi_ar.pb_client)
    cuboid_fname_template = osp.join(
        os.environ['CODE_BASE'],
        'catkin_ws/src/config/descriptions/meshes/objects/cuboids/')

    cuboid_manager = MultiBlockManager(cuboid_fname_template,
                                       cuboid_sampler,
                                       robot_id=yumi_ar.arm.robot_id,
                                       table_id=table_id,
                                       r_gel_id=r_gel_id,
                                       l_gel_id=l_gel_id)

    if args.multi:
        cuboid_fname = cuboid_manager.get_cuboid_fname()
        # cuboid_fname = '/root/catkin_ws/src/config/descriptions/meshes/objects/cuboids/test_cuboid_smaller_4479.stl'
    else:
        cuboid_fname = args.config_package_path + 'descriptions/meshes/objects/' + \
            args.object_name + '_experiments.stl'
    mesh_file = cuboid_fname

    goal_visualization = False
    if args.goal_viz:
        goal_visualization = True

    obj_id, sphere_ids, mesh, goal_obj_id = \
        cuboid_sampler.sample_cuboid_pybullet(
            cuboid_fname,
            goal=goal_visualization,
            keypoints=False)

    cuboid_manager.filter_collisions(obj_id, goal_obj_id)

    p.changeDynamics(obj_id, -1, lateralFriction=1.0)

    # goal_face = 0
    goal_faces = [0, 1, 2, 3, 4, 5]
    from random import shuffle
    shuffle(goal_faces)
    goal_face = goal_faces[0]

    exp_single = SingleArmPrimitives(cfg, yumi_ar.pb_client.get_client_id(),
                                     obj_id, cuboid_fname)
    k = 0
    while True:
        k += 1
        if k > 10:
            print('FAILED TO BUILD GRASPING GRAPH')
            return
        try:
            exp_double = DualArmPrimitives(cfg,
                                           yumi_ar.pb_client.get_client_id(),
                                           obj_id,
                                           cuboid_fname,
                                           goal_face=goal_face)
            break
        except ValueError as e:
            print(e)
            yumi_ar.pb_client.remove_body(obj_id)
            if goal_visualization:
                yumi_ar.pb_client.remove_body(goal_obj_id)
            cuboid_fname = cuboid_manager.get_cuboid_fname()
            print("Cuboid file: " + cuboid_fname)

            obj_id, sphere_ids, mesh, goal_obj_id = \
                cuboid_sampler.sample_cuboid_pybullet(
                    cuboid_fname,
                    goal=goal_visualization,
                    keypoints=False)

            cuboid_manager.filter_collisions(obj_id, goal_obj_id)

            p.changeDynamics(obj_id, -1, lateralFriction=1.0)

    if primitive_name == 'grasp':
        exp_running = exp_double
    else:
        exp_running = exp_single

    action_planner = ClosedLoopMacroActions(cfg,
                                            yumi_gs,
                                            obj_id,
                                            yumi_ar.pb_client.get_client_id(),
                                            args.config_package_path,
                                            replan=args.replan,
                                            object_mesh_file=mesh_file)

    if goal_visualization:
        trans_box_lock = threading.RLock()
        goal_viz = GoalVisual(trans_box_lock,
                              goal_obj_id,
                              action_planner.pb_client,
                              cfg.OBJECT_POSE_3,
                              show_init=False)

    action_planner.update_object(obj_id, mesh_file)
    exp_single.initialize_object(obj_id, cuboid_fname)

    dynamics_info = {}
    dynamics_info['contactDamping'] = alpha * K
    dynamics_info['contactStiffness'] = K
    dynamics_info['rollingFriction'] = args.rolling
    dynamics_info['restitution'] = restitution

    data = {}
    data['saved_data'] = []
    data['metadata'] = {}
    data['metadata']['mesh_file'] = mesh_file
    data['metadata']['cfg'] = cfg
    data['metadata']['dynamics'] = dynamics_info
    data['metadata']['cam_cfg'] = yumi_gs.cam_setup_cfg
    data['metadata']['step_repeat'] = args.sim_step_repeat
    data['metadata']['seed'] = data_seed
    data['metadata']['seed_original'] = args.np_seed

    metadata = data['metadata']

    data_manager = DataManager(pickle_path)
    pred_dir = osp.join(os.environ['CODE_BASE'], cfg.PREDICTION_DIR)
    obs_dir = osp.join(os.environ['CODE_BASE'], cfg.OBSERVATION_DIR)

    if args.save_data:
        with open(osp.join(pickle_path, 'metadata.pkl'), 'wb') as mdata_f:
            pickle.dump(metadata, mdata_f)

    total_trials = 0
    successes = 0

    # prep visualization tools
    palm_mesh_file = osp.join(os.environ['CODE_BASE'], cfg.PALM_MESH_FILE)
    table_mesh_file = osp.join(os.environ['CODE_BASE'], cfg.TABLE_MESH_FILE)
    viz_palms = PalmVis(palm_mesh_file, table_mesh_file, cfg)
    viz_pcd = PCDVis()

    if args.skeleton == 'pg':
        skeleton = ['pull', 'grasp']
    elif args.skeleton == 'gp':
        skeleton = ['grasp', 'pull']
    elif args.skeleton == 'pgp':
        skeleton = ['pull', 'grasp', 'pull']
    else:
        raise ValueError('Unrecognized plan skeleton!')

    pull_sampler = PullSamplerBasic()
    grasp_sampler = GraspSamplerBasic(None)
    # pull_sampler = PullSamplerVAEPubSub(
    #     obs_dir=obs_dir,
    #     pred_dir=pred_dir
    # )
    # grasp_sampler = GraspSamplerVAEPubSub(
    #     default_target=None,
    #     obs_dir=obs_dir,
    #     pred_dir=pred_dir
    # )
    pull_skill = PullRightSkill(pull_sampler, yumi_gs, pulling_planning_wf)
    pull_skill_no_mp = PullRightSkill(pull_sampler,
                                      yumi_gs,
                                      pulling_planning_wf,
                                      ignore_mp=True)
    grasp_skill = GraspSkill(grasp_sampler, yumi_gs, grasp_planning_wf)
    skills = {}
    # skills['pull'] = pull_skill_no_mp
    skills['pull'] = pull_skill
    skills['grasp'] = grasp_skill

    problems_file = '/root/catkin_ws/src/primitives/data/planning/test_problems_0/demo_0.pkl'
    with open(problems_file, 'rb') as f:
        problems_data = pickle.load(f)

    prob_inds = np.arange(len(problems_data), dtype=np.int64).tolist()
    data_inds = np.arange(len(problems_data[0]['problems']),
                          dtype=np.int64).tolist()

    experiment_manager = GraspEvalManager(yumi_gs,
                                          yumi_ar.pb_client.get_client_id(),
                                          pickle_path, args.exp_name, None,
                                          None, None, None, cfg)

    # experiment_manager.set_object_id(
    #     obj_id,
    #     cuboid_fname
    # )

    total_trial_number = 0
    for _ in range(len(problems_data)):
        # prob_ind = 3

        # obj_fname = str(osp.join(
        #     '/root/catkin_ws/src/config/descriptions/meshes/objects/cuboids',
        #     problems_data[prob_ind]['object_name']))

        # print(obj_fname)
        # for j, problem_data in enumerate(problems_data[prob_ind]['problems']):
        for _ in range(len(problems_data[0]['problems'])):
            total_trial_number += 1
            # prob_ind = 8
            # data_ind = 15
            prob_ind = prob_inds[np.random.randint(len(prob_inds))]
            data_ind = data_inds[np.random.randint(len(data_inds))]
            problem_data = problems_data[prob_ind]['problems'][data_ind]

            obj_fname = str(
                osp.join(
                    '/root/catkin_ws/src/config/descriptions/meshes/objects/cuboids',
                    problems_data[prob_ind]['object_name']))
            obj_name = problems_data[prob_ind]['object_name'].split('.stl')[0]
            print(obj_fname)
            start_pose = problem_data['start_vis'].tolist()

            # put object into work at start_pose, with known obj_fname
            yumi_ar.pb_client.remove_body(obj_id)
            if goal_visualization:
                yumi_ar.pb_client.remove_body(goal_obj_id)

            obj_id, sphere_ids, mesh, goal_obj_id = \
                cuboid_sampler.sample_cuboid_pybullet(
                    obj_fname,
                    goal=goal_visualization,
                    keypoints=False)
            if goal_visualization:
                goal_viz.update_goal_obj(goal_obj_id)
                goal_viz.hide_goal_obj()
                cuboid_manager.filter_collisions(obj_id, goal_obj_id)

            exp_single.initialize_object(obj_id, obj_fname)
            experiment_manager.set_object_id(obj_id, obj_fname)

            p.resetBasePositionAndOrientation(obj_id, start_pose[:3],
                                              start_pose[3:])

            p.changeDynamics(obj_id, -1, lateralFriction=1.0)

            yumi_ar.arm.set_jpos(cfg.RIGHT_OUT_OF_FRAME +
                                 cfg.LEFT_OUT_OF_FRAME,
                                 ignore_physics=True)
            time.sleep(0.2)

            real_start_pos = p.getBasePositionAndOrientation(obj_id)[0]
            real_start_ori = p.getBasePositionAndOrientation(obj_id)[1]
            real_start_pose = list(real_start_pos) + list(real_start_ori)

            transformation_des = util.matrix_from_pose(
                util.list2pose_stamped(
                    problem_data['transformation'].tolist()))

            # #### BEGIN BLOCK FOR GETTING INTRO FIGURE
            # R_3 = common.euler2rot([0.0, 0.0, np.pi/4])
            # t_3 = np.array([0.03, 0.25, 0.0])
            # T_3 = np.eye(4)
            # T_3[:-1, :-1] = R_3
            # # T_3[:-1, -1] = t_3
            # print(T_3)
            # trans_des = np.matmul(T_3, transformation_des)

            # goal_pose = util.pose_stamped2list(util.transform_pose(
            #     util.list2pose_stamped(real_start_pose),
            #     util.pose_from_matrix(trans_des)
            # ))
            # if goal_visualization:
            #     goal_viz.update_goal_state(goal_pose)
            #     goal_viz.show_goal_obj()

            # embed()
            # #### END BLOCK FOR GETTING INTRO FIGURE

            goal_pose = util.pose_stamped2list(
                util.transform_pose(
                    util.list2pose_stamped(real_start_pose),
                    util.list2pose_stamped(problem_data['transformation'])))

            # if skeleton is 'pull' 'grasp' 'pull', add an additional SE(2) transformation to the task
            if args.skeleton == 'pgp':
                while True:
                    x, y, dq = exp_single.get_rand_trans_yaw()

                    goal_pose_2_list = copy.deepcopy(goal_pose)
                    goal_pose_2_list[0] = x
                    goal_pose_2_list[1] = y
                    goal_pose_2_list[3:] = common.quat_multiply(
                        dq, np.asarray(goal_pose[3:]))

                    if goal_pose_2_list[0] > 0.2 and goal_pose_2_list[0] < 0.4 and \
                            goal_pose_2_list[1] > -0.3 and goal_pose_2_list[1] < 0.1:
                        rot = common.quat2rot(dq)
                        T_2 = np.eye(4)
                        T_2[:-1, :-1] = rot
                        T_2[:2, -1] = [x - goal_pose[0], y - goal_pose[1]]
                        break

                goal_pose = goal_pose_2_list
                transformation_des = np.matmul(T_2, transformation_des)

            # # if skeleton is 'grasp' first, invert the desired trans and flip everything
            if args.skeleton == 'gp':
                transformation_des = np.linalg.inv(transformation_des)
                start_tmp = copy.deepcopy(start_pose)
                start_pose = goal_pose
                goal_pose = start_tmp

                p.resetBasePositionAndOrientation(obj_id, start_pose[:3],
                                                  start_pose[3:])

                real_start_pos = p.getBasePositionAndOrientation(obj_id)[0]
                real_start_ori = p.getBasePositionAndOrientation(obj_id)[1]
                real_start_pose = list(real_start_pos) + list(real_start_ori)

                time.sleep(0.5)

            # get observation
            obs, pcd = yumi_gs.get_observation(
                obj_id=obj_id, robot_table_id=(yumi_ar.arm.robot_id, table_id))

            pointcloud_pts = np.asarray(obs['down_pcd_pts'][:100, :],
                                        dtype=np.float32)
            pointcloud_pts_full = np.asarray(np.concatenate(obs['pcd_pts']),
                                             dtype=np.float32)

            grasp_sampler.update_default_target(
                np.concatenate(obs['table_pcd_pts'], axis=0)[::500, :])

            trial_data = {}
            trial_data['start_pcd'] = pointcloud_pts_full
            trial_data['start_pcd_down'] = pointcloud_pts
            trial_data['obj_fname'] = cuboid_fname
            trial_data['start_pose'] = np.asarray(real_start_pose)
            trial_data['goal_pose'] = np.asarray(goal_pose)
            trial_data['goal_pose_global'] = np.asarray(goal_pose)
            trial_data['trans_des_global'] = transformation_des

            trial_data['skeleton'] = args.skeleton

            # plan!
            planner = PointCloudTree(pointcloud_pts,
                                     transformation_des,
                                     skeleton,
                                     skills,
                                     start_pcd_full=pointcloud_pts_full)
            start_plan_time = time.time()
            plan_total = planner.plan()
            trial_data['planning_time'] = time.time() - start_plan_time

            if plan_total is None:
                print('Could not find plan')
                experiment_manager.set_mp_success(False, 1)
                obj_data = experiment_manager.get_object_data()
                # obj_data_fname = osp.join(
                #     pickle_path,
                #     obj_name + '_' + str(total_trial_number),
                #     obj_name + '_' + str(total_trial_number) + '_ms_eval_data.pkl')
                obj_data_fname = osp.join(
                    pickle_path, obj_name + '_' + str(total_trial_number) +
                    '_ms_eval_data.pkl')
                if args.save_data:
                    print('Saving to: ' + str(obj_data_fname))
                    with open(obj_data_fname, 'wb') as f:
                        pickle.dump(obj_data, f)
                continue

            plan = copy.deepcopy(plan_total[1:])

            if args.trimesh_viz:
                # from multistep_planning_eval_cfg import get_cfg_defaults
                # import os.path as osp
                # from eval_utils.visualization_tools import PCDVis, PalmVis
                # cfg = get_cfg_defaults()
                # palm_mesh_file = osp.join(os.environ['CODE_BASE'],
                #                         cfg.PALM_MESH_FILE)
                # table_mesh_file = osp.join(os.environ['CODE_BASE'],
                #                         cfg.TABLE_MESH_FILE)
                # viz_palms = PalmVis(palm_mesh_file, table_mesh_file, cfg)
                # viz_pcd = PCDVis()

                ind = 0

                pcd_data = copy.deepcopy(problem_data)
                pcd_data['start'] = plan_total[ind].pointcloud_full
                pcd_data['object_pointcloud'] = plan_total[ind].pointcloud_full
                pcd_data['transformation'] = np.asarray(
                    util.pose_stamped2list(
                        util.pose_from_matrix(plan_total[ind +
                                                         1].transformation)))
                pcd_data['contact_world_frame_right'] = np.asarray(
                    plan_total[ind + 1].palms[:7])
                if skeleton[ind] == 'grasp':
                    pcd_data['contact_world_frame_left'] = np.asarray(
                        plan_total[ind + 1].palms[:7])
                elif skeleton[ind] == 'pull':
                    pcd_data['contact_world_frame_left'] = np.asarray(
                        plan_total[ind + 1].palms[7:])
                scene = viz_palms.vis_palms_pcd(pcd_data,
                                                world=True,
                                                centered=False,
                                                corr=False)
                scene.show()

                # embed()

            # execute plan if one is found...
            pose_plan = [(real_start_pose,
                          util.list2pose_stamped(real_start_pose))]
            for i in range(1, len(plan) + 1):
                pose = util.transform_pose(
                    pose_plan[i - 1][1],
                    util.pose_from_matrix(plan[i - 1].transformation))
                pose_list = util.pose_stamped2list(pose)
                pose_plan.append((pose_list, pose))

            # get palm poses from plan
            palm_pose_plan = []
            for i, node in enumerate(plan):
                palms = copy.deepcopy(node.palms)
                # palms = copy.deepcopy(node.palms_raw) if node.palms_raw is not None else copy.deepcopy(node.palms)
                if skeleton[i] == 'pull':
                    palms[2] -= 0.002
                palm_pose_plan.append(palms)

            # observe results
            full_plan = []
            for i in range(len(plan)):
                if skeleton[i] == 'pull':
                    local_plan = pulling_planning_wf(
                        util.list2pose_stamped(palm_pose_plan[i]),
                        util.list2pose_stamped(palm_pose_plan[i]),
                        util.pose_from_matrix(plan[i].transformation))
                elif skeleton[i] == 'grasp':
                    local_plan = grasp_planning_wf(
                        util.list2pose_stamped(palm_pose_plan[i][7:]),
                        util.list2pose_stamped(palm_pose_plan[i][:7]),
                        util.pose_from_matrix(plan[i].transformation))
                full_plan.append(local_plan)

            grasp_success = True

            action_planner.active_arm = 'right'
            action_planner.inactive_arm = 'left'

            if goal_visualization:
                goal_viz.update_goal_state(goal_pose)
                goal_viz.show_goal_obj()

            if goal_visualization:
                goal_viz.update_goal_state(goal_pose)
                goal_viz.show_goal_obj()

            real_start_pos = p.getBasePositionAndOrientation(obj_id)[0]
            real_start_ori = p.getBasePositionAndOrientation(obj_id)[1]
            real_start_pose = list(real_start_pos) + list(real_start_ori)
            real_start_mat = util.matrix_from_pose(
                util.list2pose_stamped(real_start_pose))

            # embed()
            try:
                start_playback_time = time.time()
                for playback in range(args.playback_num):
                    if playback > 0 and goal_visualization:
                        goal_viz.hide_goal_obj()

                    yumi_ar.pb_client.reset_body(obj_id, pose_plan[0][0][:3],
                                                 pose_plan[0][0][3:])
                    p.changeDynamics(obj_id, -1, lateralFriction=1.0)
                    for i, skill in enumerate(skeleton):
                        if skill == 'pull':
                            # set arm configuration to good start state
                            yumi_ar.arm.set_jpos(cfg.RIGHT_INIT +
                                                 cfg.LEFT_INIT,
                                                 ignore_physics=True)
                            time.sleep(0.5)

                            # move to making contact, and ensure contact is made
                            # try:
                            #     _, _ = action_planner.single_arm_setup(full_plan[i][0], pre=True)
                            # except ValueError as e:
                            #     print(e)
                            #     break
                            _, _ = action_planner.single_arm_setup(
                                full_plan[i][0], pre=True)
                            start_playback_time = time.time()
                            if not experiment_manager.still_pulling():
                                while True:
                                    if experiment_manager.still_pulling(
                                    ) or time.time(
                                    ) - start_playback_time > 20.0:
                                        break
                                    action_planner.single_arm_approach()
                                    time.sleep(0.075)
                                    new_plan = pulling_planning_wf(
                                        yumi_gs.get_current_tip_poses()
                                        ['left'],
                                        yumi_gs.get_current_tip_poses()
                                        ['right'],
                                        util.pose_from_matrix(
                                            plan[i].transformation))
                                pull_plan = new_plan[0]
                            else:
                                pull_plan = full_plan[i][0]
                            # try:
                            #     action_planner.playback_single_arm('pull', pull_plan, pre=False)
                            # except ValueError as e:
                            #     print(e)
                            #     break
                            action_planner.playback_single_arm('pull',
                                                               pull_plan,
                                                               pre=False)
                            grasp_success = grasp_success and experiment_manager.still_pulling(
                                n=False)
                            print('grasp success: ' + str(grasp_success))
                            time.sleep(0.5)
                            action_planner.single_arm_retract()

                        elif skill == 'grasp':
                            yumi_ar.arm.set_jpos([
                                0.9936, -2.1848, -0.9915, 0.8458, 3.7618,
                                1.5486, 0.1127, -1.0777, -2.1187, 0.995, 1.002,
                                -3.6834, 1.8132, 2.6405
                            ],
                                                 ignore_physics=True)
                            time.sleep(0.5)
                            # try:
                            #     _, _ = action_planner.dual_arm_setup(full_plan[i][0], 0, pre=True)
                            # except ValueError as e:
                            #     print(e)
                            #     break
                            _, _ = action_planner.dual_arm_setup(
                                full_plan[i][0], 0, pre=True)
                            start_playback_time = time.time()
                            if not experiment_manager.still_grasping():
                                while True:
                                    if experiment_manager.still_grasping(
                                    ) or time.time(
                                    ) - start_playback_time > 20.0:
                                        break
                                    action_planner.dual_arm_approach()
                                    time.sleep(0.075)
                                    new_plan = grasp_planning_wf(
                                        yumi_gs.get_current_tip_poses()
                                        ['left'],
                                        yumi_gs.get_current_tip_poses()
                                        ['right'],
                                        util.pose_from_matrix(
                                            plan[i].transformation))
                                grasp_plan = new_plan
                            else:
                                grasp_plan = full_plan[i]
                            for k, subplan in enumerate(grasp_plan):
                                # try:
                                #     action_planner.playback_dual_arm('grasp', subplan, k, pre=False)
                                # except ValueError as e:
                                #     print(e)
                                #     break
                                action_planner.playback_dual_arm('grasp',
                                                                 subplan,
                                                                 k,
                                                                 pre=False)
                                if k == 1:
                                    grasp_success = grasp_success and experiment_manager.still_grasping(
                                        n=False)
                                    print('grasp success: ' +
                                          str(grasp_success))
                                time.sleep(1.0)
            except ValueError as e:
                print(e)
                experiment_manager.set_mp_success(True, 1)
                experiment_manager.set_execute_success(False)
                obj_data = experiment_manager.get_object_data()
                # obj_data_fname = osp.join(
                #     pickle_path,
                #     obj_name + '_' + str(total_trial_number),
                #     obj_name + '_' + str(total_trial_number) + '_ms_eval_data.pkl')
                obj_data_fname = osp.join(
                    pickle_path, obj_name + '_' + str(total_trial_number) +
                    '_ms_eval_data.pkl')
                if args.save_data:
                    print('Saving to: ' + str(obj_data_fname))
                    with open(obj_data_fname, 'wb') as f:
                        pickle.dump(obj_data, f)
                continue

            real_final_pos = p.getBasePositionAndOrientation(obj_id)[0]
            real_final_ori = p.getBasePositionAndOrientation(obj_id)[1]
            real_final_pose = list(real_final_pos) + list(real_final_ori)
            real_final_mat = util.matrix_from_pose(
                util.list2pose_stamped(real_final_pose))
            real_T_mat = np.matmul(real_final_mat,
                                   np.linalg.inv(real_start_mat))
            real_T_pose = util.pose_stamped2np(
                util.pose_from_matrix(real_T_mat))

            trial_data['trans_executed'] = real_T_mat
            trial_data['final_pose'] = real_final_pose

            experiment_manager.set_mp_success(True, 1)
            experiment_manager.set_execute_success(True)
            experiment_manager.end_trial(trial_data, grasp_success)

            time.sleep(3.0)

            obj_data = experiment_manager.get_object_data()

            kvs = {}
            kvs['grasp_success'] = obj_data['grasp_success']
            kvs['pos_err'] = np.mean(obj_data['final_pos_error'])
            kvs['ori_err'] = np.mean(obj_data['final_ori_error'])
            kvs['planning_time'] = obj_data['planning_time']
            string = ''

            for k, v in kvs.items():
                string += "%s: %.3f, " % (k, v)
            print(string)

            # obj_data_fname = osp.join(
            #     pickle_path,
            #     obj_name + '_' + str(total_trial_number),
            #     obj_name + '_' + str(total_trial_number)  + '_ms_eval_data.pkl')
            obj_data_fname = osp.join(
                pickle_path,
                obj_name + '_' + str(total_trial_number) + '_ms_eval_data.pkl')
            if args.save_data:
                print('Saving to: ' + str(obj_data_fname))
                with open(obj_data_fname, 'wb') as f:
                    pickle.dump(obj_data, f)