Exemplo n.º 1
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)
Exemplo n.º 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
    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)

    yumi_gs = YumiGelslimPybulet(yumi_ar, cfg, exec_thread=False)

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

    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/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)

    cuboid_fname = cuboid_manager.get_cuboid_fname()
    print("Cuboid file: " + cuboid_fname)
    # cuboid_fname = os.path.join(cuboid_fname_template, 'test_cuboid_smaller_4206.stl')

    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)

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

    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
    goal_face = 0

    # mesh_file = args.config_package_path + 'descriptions/meshes/objects/' + \
    #     args.object_name + '_experiments.stl'
    mesh_file = cuboid_fname
    exp_single = SingleArmPrimitives(cfg, yumi_ar.pb_client.get_client_id(),
                                     obj_id, mesh_file)
    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,
                                           mesh_file,
                                           goal_face=goal_face)
            break
        except ValueError as e:
            print(e)
            yumi_ar.pb_client.remove_body(obj_id)
            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=True,
                    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)

    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)
    exp_double.initialize_object(obj_id, cuboid_fname, goal_face)
    print('Reset multi block!')

    for _ in range(args.num_blocks):
        for _ in range(args.num_obj_samples):
            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=None, random_goal=True, execute=True)

            start_pose = plan_args['object_pose1_world']
            goal_pose = plan_args['object_pose2_world']
            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:
                try:
                    result = action_planner.execute(primitive_name, plan_args)
                    if result is not None:
                        print(str(result[0]))

                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)
                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=True,
                        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)
                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)
def main(args):
    pull_cfg_file = os.path.join(args.example_config_path, 'pull') + ".yaml"
    pull_cfg = get_cfg_defaults()
    pull_cfg.merge_from_file(pull_cfg_file)
    pull_cfg.freeze()

    grasp_cfg_file = os.path.join(args.example_config_path, 'grasp') + ".yaml"
    grasp_cfg = get_cfg_defaults()
    grasp_cfg.merge_from_file(grasp_cfg_file)
    grasp_cfg.freeze()

    # cfg = grasp_cfg
    cfg = pull_cfg

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

    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)

    yumi_gs = YumiGelslimPybulet(yumi_ar, cfg, exec_thread=False)

    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:])

    # 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
    face = np.random.randint(6)
    # face = 3

    mesh_file = args.config_package_path + 'descriptions/meshes/objects/cuboids/' + \
        args.object_name + '_experiments.stl'
    cuboid_fname = mesh_file
    exp_single = SingleArmPrimitives(pull_cfg,
                                     yumi_ar.pb_client.get_client_id(), obj_id,
                                     mesh_file)

    exp_double = DualArmPrimitives(grasp_cfg,
                                   yumi_ar.pb_client.get_client_id(),
                                   obj_id,
                                   mesh_file,
                                   goal_face=face)

    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)

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

    multistep_planner = MultiStepManager(cfg, grasp_cfg, pull_cfg)

    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)
    print('Reset multi block!')

    for _ in range(args.num_obj_samples):
        # get grasp sample
        # start_face_index = np.random.randint(len(grasp_cfg.VALID_GRASP_PAIRS[face]))
        # start_face = grasp_cfg.VALID_GRASP_PAIRS[face][start_face_index]

        # grasp_args = exp_double.get_random_primitive_args(ind=start_face,
        #                                                   random_goal=True,
        #                                                   execute=False)
        # # pull_args_start = exp_single.get_random_primitive_args(ind=cfg.GRASP_TO_PULL[start_face],
        #                                                        random_goal=True)
        # pull_args_goal = exp_single.get_random_primitive_args(ind=cfg.GRASP_TO_PULL[face],
        #                                                       random_goal=True)

        # pull_args_start['object_pose2_world'] = grasp_args['object_pose1_world']
        # pull_args_goal['object_pose1_world'] = grasp_args['object_pose2_world']

        # full_args = [pull_args_start, grasp_args, pull_args_goal]
        # full_args = [grasp_args]

        full_args = multistep_planner.get_args(exp_single,
                                               exp_double,
                                               face,
                                               execute=False)

        # obj_pose_final = exp_running.goal_pose_world_frame_mod
        # palm_poses_obj_frame = {}
        # y_normals = action_planner.get_palm_y_normals(palm_poses_world)
        # delta = 2e-3
        # for key in palm_poses_world.keys():
        #     palm_pose_world = palm_poses_world[key]

        #     # try to penetrate the object a small amount
        #     palm_pose_world.pose.position.x -= delta*y_normals[key].pose.position.x
        #     palm_pose_world.pose.position.y -= delta*y_normals[key].pose.position.y
        #     palm_pose_world.pose.position.z -= delta*y_normals[key].pose.position.z

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

        valid_subplans = 0
        valid_plans = []
        for plan_args in full_args:
            plan = action_planner.get_primitive_plan(plan_args['name'],
                                                     plan_args, 'right')
            goal_viz.update_goal_state(
                util.pose_stamped2list(plan_args['object_pose2_world']))

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

            if args.debug:
                for _ 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:
                primitive_name = plan_args['name']
                start_joints = {}
                if primitive_name == 'pull':
                    start_joints['right'] = pull_cfg.RIGHT_INIT
                    start_joints['left'] = pull_cfg.LEFT_INIT
                elif primitive_name == 'grasp':
                    start_joints['right'] = grasp_cfg.RIGHT_INIT
                    start_joints['left'] = grasp_cfg.LEFT_INIT
                subplan_valid = action_planner.full_mp_check(
                    plan, primitive_name, start_joints=start_joints)
                if subplan_valid:
                    print("subplan valid!")
                    valid_subplans += 1
                    valid_plans.append(plan)
                else:
                    print("not valid, primitive: " + primitive_name)

        if valid_subplans == len(full_args) and not args.debug:
            yumi_ar.pb_client.reset_body(
                obj_id,
                util.pose_stamped2list(full_args[0]['object_pose1_world'])[:3],
                util.pose_stamped2list(full_args[0]['object_pose1_world'])[3:])
            for plan_args in full_args:
                # teleport to start state, to avoid colliding with object during transit
                primitive_name = plan_args['name']
                time.sleep(0.1)
                for _ in range(10):
                    if primitive_name == 'pull':
                        # yumi_gs.update_joints(pull_cfg.RIGHT_INIT + pull_cfg.LEFT_INIT)
                        yumi_gs.yumi_pb.arm.set_jpos(pull_cfg.RIGHT_INIT +
                                                     pull_cfg.LEFT_INIT,
                                                     ignore_physics=True)
                    elif primitive_name == 'grasp':
                        # yumi_gs.update_joints(grasp_cfg.RIGHT_INIT + grasp_cfg.LEFT_INIT)
                        yumi_gs.yumi_pb.arm.set_jpos(grasp_cfg.RIGHT_INIT +
                                                     grasp_cfg.LEFT_INIT,
                                                     ignore_physics=True)
                try:
                    # should update the plan args start state to be the current simulator state
                    result = action_planner.execute(primitive_name, plan_args)
                    if result is not None:
                        print(str(result[0]))

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

                time.sleep(1.0)

                if primitive_name == 'pull':
                    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)