-
Notifications
You must be signed in to change notification settings - Fork 109
/
create_ir_database.py
executable file
·152 lines (133 loc) · 6.52 KB
/
create_ir_database.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
#!/usr/bin/env python
import argparse
import random
import time
from pybullet_tools.pr2_utils import set_arm_conf, get_other_arm, arm_conf, REST_LEFT_ARM, \
get_carry_conf, get_gripper_link, GET_GRASPS, IR_FILENAME, get_database_file, DRAKE_PR2_URDF, \
set_group_conf, get_group_conf, get_base_pose
from pybullet_tools.utils import create_box, disconnect, add_data_path, connect, get_movable_joints, get_joint_positions, \
sample_placement, set_pose, multiply, invert, set_joint_positions, pairwise_collision, inverse_kinematics, \
get_link_pose, get_body_name, write_pickle, uniform_pose_generator, set_base_values, \
load_pybullet, HideOutput, wait_if_gui, draw_point, point_from_pose, has_gui, elapsed_time, \
sub_inverse_kinematics, BodySaver
from pybullet_tools.pr2_problems import create_table
from pybullet_tools.ikfast.pr2.ik import pr2_inverse_kinematics, is_ik_compiled
from pybullet_tools.ikfast.utils import USE_CURRENT
from pybullet_tools.pr2_primitives import get_stable_gen, get_grasp_gen, get_ik_ir_gen
def save_inverse_reachability(robot, arm, grasp_type, tool_link, gripper_from_base_list):
# TODO: store value of torso and roll joint for the IK database. Sample the roll joint.
# TODO: hash the pr2 urdf as well
filename = IR_FILENAME.format(grasp_type, arm)
path = get_database_file(filename)
data = {
'filename': filename,
'robot': get_body_name(robot),
'grasp_type': grasp_type,
'arm': arm,
'torso': get_group_conf(robot, 'torso'),
'carry_conf': get_carry_conf(arm, grasp_type),
'tool_link': tool_link,
'ikfast': is_ik_compiled(),
'gripper_from_base': gripper_from_base_list,
}
write_pickle(path, data)
if has_gui():
handles = []
for gripper_from_base in gripper_from_base_list:
handles.extend(draw_point(point_from_pose(gripper_from_base), color=(1, 0, 0)))
wait_if_gui()
return path
#######################################################
def create_inverse_reachability(robot, body, table, arm, grasp_type, max_attempts=500, num_samples=500):
tool_link = get_gripper_link(robot, arm)
robot_saver = BodySaver(robot)
gripper_from_base_list = []
grasps = GET_GRASPS[grasp_type](body)
start_time = time.time()
while len(gripper_from_base_list) < num_samples:
box_pose = sample_placement(body, table)
set_pose(body, box_pose)
grasp_pose = random.choice(grasps)
gripper_pose = multiply(box_pose, invert(grasp_pose))
for attempt in range(max_attempts):
robot_saver.restore()
base_conf = next(uniform_pose_generator(robot, gripper_pose)) #, reachable_range=(0., 1.)))
#set_base_values(robot, base_conf)
set_group_conf(robot, 'base', base_conf)
if pairwise_collision(robot, table):
continue
grasp_conf = pr2_inverse_kinematics(robot, arm, gripper_pose) #, nearby_conf=USE_CURRENT)
#conf = inverse_kinematics(robot, link, gripper_pose)
if (grasp_conf is None) or pairwise_collision(robot, table):
continue
gripper_from_base = multiply(invert(get_link_pose(robot, tool_link)), get_base_pose(robot))
#wait_if_gui()
gripper_from_base_list.append(gripper_from_base)
print('{} / {} | {} attempts | [{:.3f}]'.format(
len(gripper_from_base_list), num_samples, attempt, elapsed_time(start_time)))
wait_if_gui()
break
else:
print('Failed to find a kinematic solution after {} attempts'.format(max_attempts))
return save_inverse_reachability(robot, arm, grasp_type, tool_link, gripper_from_base_list)
#######################################################
class MockProblem(object):
def __init__(self, robot, fixed=[], grasp_types=[]):
self.robot = robot
self.fixed = fixed
self.grasp_types = grasp_types
def create_inverse_reachability2(robot, body, table, arm, grasp_type, max_attempts=500, num_samples=500):
tool_link = get_gripper_link(robot, arm)
problem = MockProblem(robot, fixed=[table], grasp_types=[grasp_type])
placement_gen_fn = get_stable_gen(problem)
grasp_gen_fn = get_grasp_gen(problem, collisions=True)
ik_ir_fn = get_ik_ir_gen(problem, max_attempts=max_attempts, learned=False, teleport=True)
placement_gen = placement_gen_fn(body, table)
grasps = list(grasp_gen_fn(body))
print('Grasps:', len(grasps))
# TODO: sample the torso height
# TODO: consider IK with respect to the torso frame
start_time = time.time()
gripper_from_base_list = []
while len(gripper_from_base_list) < num_samples:
[(p,)] = next(placement_gen)
(g,) = random.choice(grasps)
output = next(ik_ir_fn(arm, body, p, g), None)
if output is None:
print('Failed to find a solution after {} attempts'.format(max_attempts))
else:
(_, ac) = output
[at,] = ac.commands
at.path[-1].assign()
gripper_from_base = multiply(invert(get_link_pose(robot, tool_link)), get_base_pose(robot))
gripper_from_base_list.append(gripper_from_base)
print('{} / {} [{:.3f}]'.format(
len(gripper_from_base_list), num_samples, elapsed_time(start_time)))
wait_if_gui()
return save_inverse_reachability(robot, arm, grasp_type, tool_link, gripper_from_base_list)
#######################################################
def main():
parser = argparse.ArgumentParser() # Automatically includes help
parser.add_argument('-arm', required=True)
parser.add_argument('-grasp', required=True)
parser.add_argument('-viewer', action='store_true', help='enable viewer.')
args = parser.parse_args()
arm = args.arm
other_arm = get_other_arm(arm)
grasp_type = args.grasp
connect(use_gui=args.viewer)
add_data_path()
with HideOutput():
robot = load_pybullet(DRAKE_PR2_URDF)
set_group_conf(robot, 'torso', [0.2])
set_arm_conf(robot, arm, get_carry_conf(arm, grasp_type))
set_arm_conf(robot, other_arm, arm_conf(other_arm, REST_LEFT_ARM))
#plane = p.loadURDF("plane.urdf")
#table = p.loadURDF("table/table.urdf", 0, 0, 0, 0, 0, 0.707107, 0.707107)
table = create_table()
box = create_box(.07, .07, .14)
#create_inverse_reachability(robot, box, table, arm=arm, grasp_type=grasp_type)
create_inverse_reachability2(robot, box, table, arm=arm, grasp_type=grasp_type)
disconnect()
if __name__ == '__main__':
main()