forked from caelan/pybullet-planning
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_kuka_pick.py
executable file
·78 lines (67 loc) · 2.62 KB
/
test_kuka_pick.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
#!/usr/bin/env python
from __future__ import print_function
from pybullet_tools.kuka_primitives import BodyPose, BodyConf, Command, get_grasp_gen, \
get_ik_fn, get_free_motion_gen, get_holding_motion_gen
from pybullet_tools.utils import WorldSaver, enable_gravity, connect, dump_world, set_pose, \
draw_global_system, draw_pose, set_camera_pose, Pose, Point, set_default_camera, stable_z, \
BLOCK_URDF, load_model, wait_if_gui, disconnect, DRAKE_IIWA_URDF, wait_if_gui, update_state, disable_real_time, HideOutput
def plan(robot, block, fixed, teleport):
grasp_gen = get_grasp_gen(robot, 'top')
ik_fn = get_ik_fn(robot, fixed=fixed, teleport=teleport)
free_motion_fn = get_free_motion_gen(robot, fixed=([block] + fixed), teleport=teleport)
holding_motion_fn = get_holding_motion_gen(robot, fixed=fixed, teleport=teleport)
pose0 = BodyPose(block)
conf0 = BodyConf(robot)
saved_world = WorldSaver()
for grasp, in grasp_gen(block):
saved_world.restore()
result1 = ik_fn(block, pose0, grasp)
if result1 is None:
continue
conf1, path2 = result1
pose0.assign()
result2 = free_motion_fn(conf0, conf1)
if result2 is None:
continue
path1, = result2
result3 = holding_motion_fn(conf1, conf0, block, grasp)
if result3 is None:
continue
path3, = result3
return Command(path1.body_paths +
path2.body_paths +
path3.body_paths)
return None
def main(display='execute'): # control | execute | step
connect(use_gui=True)
disable_real_time()
draw_global_system()
with HideOutput():
robot = load_model(DRAKE_IIWA_URDF) # KUKA_IIWA_URDF | DRAKE_IIWA_URDF
floor = load_model('models/short_floor.urdf')
block = load_model(BLOCK_URDF, fixed_base=False)
set_pose(block, Pose(Point(y=0.5, z=stable_z(block, floor))))
set_default_camera(distance=2)
dump_world()
saved_world = WorldSaver()
command = plan(robot, block, fixed=[floor], teleport=False)
if (command is None) or (display is None):
print('Unable to find a plan!')
return
saved_world.restore()
update_state()
wait_if_gui('{}?'.format(display))
if display == 'control':
enable_gravity()
command.control(real_time=False, dt=0)
elif display == 'execute':
command.refine(num_steps=10).execute(time_step=0.005)
elif display == 'step':
command.step()
else:
raise ValueError(display)
print('Quit?')
wait_if_gui()
disconnect()
if __name__ == '__main__':
main()