-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
57 lines (48 loc) · 1.96 KB
/
utils.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
import json
from pybullet_tools.parse_json import parse_robot, parse_body
from pybullet_tools.utils import set_joint_positions, \
wait_if_gui, wait_for_duration, get_collision_fn
from pybullet_tools.pr2_utils import get_disabled_collisions
import pybullet as p
import numpy as np
def load_env(env_file):
# load robot and obstacles defined in a json file
with open(env_file, 'r') as f:
env_json = json.loads(f.read())
robots = {robot['name']: parse_robot(robot) for robot in env_json['robots']}
bodies = {body['name']: parse_body(body) for body in env_json['bodies']}
return robots, bodies
def get_collision_fn_PR2(robot, joints, obstacles):
# check robot collision with environment
disabled_collisions = get_disabled_collisions(robot)
return get_collision_fn(robot, joints, obstacles=obstacles, attachments=[], \
self_collisions=True, disabled_collisions=disabled_collisions)
def execute_trajectory(robot, joints, path, sleep=None):
# Move the robot according to a given path
if path is None:
print('Path is empty')
return
print('Moving forward')
for bq in path:
set_joint_positions(robot, joints, bq)
if sleep is None:
wait_if_gui('Continue?')
else:
wait_for_duration(sleep)
#print('Finished')
def draw_sphere_marker(position, radius, color):
vs_id = p.createVisualShape(p.GEOM_SPHERE, radius=radius, rgbaColor=color)
marker_id = p.createMultiBody(basePosition=position, baseCollisionShapeIndex=-1, baseVisualShapeIndex=vs_id)
return marker_id
def draw_line(start, end, width, color):
line_id = p.addUserDebugLine(start, end, color, width)
return line_id
def dist(vect1, vect2):
distance = 0
for num1, num2 in zip(vect1, vect2):
distance += (num1 - num2)**2
return np.sqrt(distance)
def distNode(node1, node2):
if node1.isObs or node2.isObs:
return np.inf
return dist(node1.get_state(), node2.get_state())