def get_target_zone_corners(self):
"""Determine corners of target zone.
Assumes we follow this structure in some order:
c2 --- c3
| |
| |
c1 --- c4
The actual xyz positions depend on the sampled `self.zone_pose`.
We use this for (among other things) the target points for the
ClothPickPlace demonstrator. Make sure the clockwise ordering is
consistent with labels, for the reward and demonstrator. NOTE: the
zone stays clockwise in terms of our c1 -> c2 -> c3 -> c4 ordering
(from a top down view) but the CLOTH may have been flipped. For now I
assume that can be handled in the action by considering the possible
counter-clockwise map.
"""
el2 = self._zone_length / 2
self.c1_position = U.apply(self.zone_pose, (-el2, -el2, 0))
self.c2_position = U.apply(self.zone_pose, (-el2, el2, 0))
self.c3_position = U.apply(self.zone_pose, (el2, el2, 0))
self.c4_position = U.apply(self.zone_pose, (el2, -el2, 0))
self._corner_targets_xy = np.array([
self.c1_position[:2],
self.c2_position[:2],
self.c3_position[:2],
self.c4_position[:2],
])
def reward(self):
"""Get delta rewards for current timestep.
Returns:
A tuple consisting of the scalar (delta) reward, plus `extras`
dict which has extra task-dependent info from the process of
computing rewards that gives us finer-grained details. Use
`extras` for further data analysis.
"""
reward, info = 0, {}
# Unpack next goal step.
objs, matches, targs, _, _, metric, params, max_reward = self.goals[0]
# Evaluate by matching object poses.
if metric == 'pose':
step_reward = 0
for i in range(len(objs)):
object_id, (symmetry, _) = objs[i]
pose = p.getBasePositionAndOrientation(object_id)
targets_i = np.argwhere(matches[i, :]).reshape(-1)
for j in targets_i:
target_pose = targs[j]
if self.is_match(pose, target_pose, symmetry):
step_reward += max_reward / len(objs)
break
# Evaluate by measuring object intersection with zone.
elif metric == 'zone':
zone_pts, total_pts = 0, 0
obj_pts, zones = params
for zone_pose, zone_size in zones:
# Count valid points in zone.
for obj_id in obj_pts:
pts = obj_pts[obj_id]
obj_pose = p.getBasePositionAndOrientation(obj_id)
world_to_zone = utils.invert(zone_pose)
obj_to_zone = utils.multiply(world_to_zone, obj_pose)
pts = np.float32(utils.apply(obj_to_zone, pts))
if len(zone_size) > 1:
valid_pts = np.logical_and.reduce([
pts[0, :] > -zone_size[0] / 2, pts[0, :] < zone_size[0] / 2,
pts[1, :] > -zone_size[1] / 2, pts[1, :] < zone_size[1] / 2,
pts[2, :] < self.bounds[2, 1]])
zone_pts += np.sum(np.float32(valid_pts))
total_pts += pts.shape[1]
step_reward = max_reward * (zone_pts / total_pts)
# Get cumulative rewards and return delta.
reward = self.progress + step_reward - self._rewards
self._rewards = self.progress + step_reward
# Move to next goal step if current goal step is complete.
if np.abs(max_reward - step_reward) < 0.01:
self.progress += max_reward # Update task progress.
self.goals.pop(0)
return reward, info
def reset(self, env):
self.total_rewards = 0
# Add goal post.
line_urdf = 'assets/line/line.urdf'
self.zone_size = (0.5, 1, 0.2)
self.zone_pose = ((0.5, -0.86, 0),
utils.get_pybullet_quaternion_from_rot((0, 0, 0)))
env.add_object(line_urdf, self.zone_pose, fixed=True)
# Add box.
box_size = self.random_size(0.05, 0.25, 0.05, 0.15, 0.05, 0.05)
box_template = 'assets/box/box-template.urdf'
box_urdf = self.fill_template(box_template, {'DIM': box_size})
px = self.bounds[0, 0] + 0.1 + np.random.rand() * 0.3
position = (px, 0.3, box_size[2] / 2)
theta = np.random.rand() * np.pi / 4 - np.pi / 8
rotation = utils.get_pybullet_quaternion_from_rot((0, 0, theta))
box_pose = (position, rotation)
box_id = env.add_object(box_urdf, box_pose)
os.remove(box_urdf)
self.color_random_brown(box_id)
self.object_points = {box_id: self.get_object_points(box_id)}
# Move end effector to start position next to box.
box_dim = p.getVisualShapeData(box_id)[0][3]
start_position = np.array([0, box_dim[1] / 2 + 0.01, 0])
start_position = utils.apply(box_pose, start_position)
rotation = tuple(env.home_pose[3:])
joints = env.solve_ik((tuple(start_position) + rotation))
for i in range(len(env.joints)):
p.resetJointState(env.ur5, env.joints[i], joints[i])
def reset(self, env):
# Add stand.
base_size = (0.12, 0.36, 0.01)
base_urdf = 'assets/hanoi/stand.urdf'
base_pose = self.random_pose(env, base_size)
env.add_object(base_urdf, base_pose, fixed=True)
# Rod positions in base coordinates.
rod_positions = ((0, -0.12, 0.03), (0, 0, 0.03), (0, 0.12, 0.03))
# Add disks.
num_disks = 3
for i in range(num_disks):
disk_urdf = 'assets/hanoi/disk%d.urdf' % i
position = utils.apply(base_pose, rod_positions[0])
z_offset = 0.015 * (num_disks - i - 2)
position = (position[0], position[1], position[2] + z_offset)
env.add_object(disk_urdf, (position, base_pose[1]))
# Solve Hanoi sequence with dynamic programming.
hanoi_steps = [] # [[object index, from rod, to rod], ...]
def solve_hanoi(n, t0, t1, t2):
if n == 0:
hanoi_steps.append([n, t0, t1])
return
solve_hanoi(n - 1, t0, t2, t1)
hanoi_steps.append([n, t0, t1])
solve_hanoi(n - 1, t2, t1, t0)
solve_hanoi(num_disks - 1, 0, 2, 1)
self.num_steps = len(hanoi_steps)
# Construct goal sequence [{object id : (symmetry, pose)}, ...]
self.goal = {'places': {}, 'steps': []}
for step in hanoi_steps:
object_id = env.objects[step[0]]
rod_position = rod_positions[step[2]]
place_position = utils.apply(base_pose, rod_position)
place_pose = (place_position,
utils.get_pybullet_quaternion_from_rot((0, 0, 0)))
place_id = len(self.goal['places'])
self.goal['places'][place_id] = place_pose
self.goal['steps'].append({object_id: (0, [place_id])})
def reset(self, env):
super().reset(env)
# Add stand.
base_size = (0.12, 0.36, 0.01)
base_urdf = 'assets/hanoi/stand.urdf'
base_pose = self.get_random_pose(env, base_size)
env.add_object(base_urdf, base_pose, 'fixed')
# Rod positions in base coordinates.
rod_pos = ((0, -0.12, 0.03), (0, 0, 0.03), (0, 0.12, 0.03))
# Add disks.
disks = []
n_disks = 3
for i in range(n_disks):
disk_urdf = 'assets/hanoi/disk%d.urdf' % i
pos = utils.apply(base_pose, rod_pos[0])
z = 0.015 * (n_disks - i - 2)
pos = (pos[0], pos[1], pos[2] + z)
disks.append(env.add_object(disk_urdf, (pos, base_pose[1])))
# Solve Hanoi sequence with dynamic programming.
hanoi_steps = [] # [[object index, from rod, to rod], ...]
def solve_hanoi(n, t0, t1, t2):
if n == 0:
hanoi_steps.append([n, t0, t1])
return
solve_hanoi(n - 1, t0, t2, t1)
hanoi_steps.append([n, t0, t1])
solve_hanoi(n - 1, t2, t1, t0)
solve_hanoi(n_disks - 1, 0, 2, 1)
# Goal: pick and place disks using Hanoi sequence.
for step in hanoi_steps:
disk_id = disks[step[0]]
targ_pos = rod_pos[step[2]]
targ_pos = utils.apply(base_pose, targ_pos)
targ_pose = (targ_pos, (0, 0, 0, 1))
self.goals.append(
([(disk_id, (0, None))], np.int32([[1]]), [targ_pose], False,
True, 'pose', None, 1 / len(hanoi_steps)))
def reset(self, env):
super().reset(env)
# Add base.
base_size = (0.05, 0.15, 0.005)
base_urdf = 'assets/stacking/stand.urdf'
base_pose = self.get_random_pose(env, base_size)
env.add_object(base_urdf, base_pose, 'fixed')
# Block colors.
colors = [
utils.COLORS['purple'], utils.COLORS['blue'],
utils.COLORS['green'], utils.COLORS['yellow'],
utils.COLORS['orange'], utils.COLORS['red']
]
# Add blocks.
objs = []
# sym = np.pi / 2
block_size = (0.04, 0.04, 0.04)
block_urdf = 'assets/stacking/block.urdf'
for i in range(6):
block_pose = self.get_random_pose(env, block_size)
block_id = env.add_object(block_urdf, block_pose)
p.changeVisualShape(block_id, -1, rgbaColor=colors[i] + [1])
objs.append((block_id, (np.pi / 2, None)))
# Associate placement locations for goals.
place_pos = [(0, -0.05, 0.03), (0, 0, 0.03), (0, 0.05, 0.03),
(0, -0.025, 0.08), (0, 0.025, 0.08), (0, 0, 0.13)]
targs = [(utils.apply(base_pose, i), base_pose[1]) for i in place_pos]
# Goal: blocks are stacked in a pyramid (bottom row: green, blue, purple).
self.goals.append((objs[:3], np.ones(
(3, 3)), targs[:3], False, True, 'pose', None, 1 / 2))
# Goal: blocks are stacked in a pyramid (middle row: yellow, orange).
self.goals.append((objs[3:5], np.ones(
(2, 2)), targs[3:5], False, True, 'pose', None, 1 / 3))
# Goal: blocks are stacked in a pyramid (top row: red).
self.goals.append((objs[5:], np.ones(
(1, 1)), targs[5:], False, True, 'pose', None, 1 / 6))
def reset(self, env):
# Add base.
base_size = (0.05, 0.15, 0.005)
base_urdf = 'assets/stacking/stand.urdf'
base_pose = self.random_pose(env, base_size)
env.add_object(base_urdf, base_pose, fixed=True)
# Block colors.
colors = [
utils.COLORS['purple'], utils.COLORS['blue'],
utils.COLORS['green'], utils.COLORS['yellow'],
utils.COLORS['orange'], utils.COLORS['red']
]
# Add blocks.
block_ids = []
block_size = (0.04, 0.04, 0.04)
block_urdf = 'assets/stacking/block.urdf'
for i in range(6):
block_pose = self.random_pose(env, block_size)
block_id = env.add_object(block_urdf, block_pose)
p.changeVisualShape(block_id, -1, rgbaColor=colors[i] + [1])
block_ids.append(block_id)
# Associate placement locations.
self.num_steps = 6
self.goal = {'places': {}, 'steps': []}
self.goal['places'] = {
0: (utils.apply(base_pose, (0, -0.05, 0.03)), base_pose[1]),
1: (utils.apply(base_pose, (0, 0, 0.03)), base_pose[1]),
2: (utils.apply(base_pose, (0, 0.05, 0.03)), base_pose[1]),
3: (utils.apply(base_pose, (0, -0.025, 0.08)), base_pose[1]),
4: (utils.apply(base_pose, (0, 0.025, 0.08)), base_pose[1]),
5: (utils.apply(base_pose, (0, 0, 0.13)), base_pose[1])
}
block_symmetry = np.pi / 2
self.goal['steps'] = [{
block_ids[0]: (block_symmetry, [0, 1, 2]),
block_ids[1]: (block_symmetry, [0, 1, 2]),
block_ids[2]: (block_symmetry, [0, 1, 2])
}, {
block_ids[3]: (block_symmetry, [3, 4]),
block_ids[4]: (block_symmetry, [3, 4])
}, {
block_ids[5]: (block_symmetry, [5])
}]
def add_cable(self, env, size_range, info):
"""Add a cable to the env, consisting of rigids beads.
Add each bead ID to (a) env.objects, (b) object_points, and (c)
cable_bead_ids. Use (b) because the demonstrator checks it to pick
the bead farthest from a goal, and it is also used to tally up beads
within the zone (to compute reward). Use (c) to distinguish between
bead vs non-bead objects in case we add other items.
Args:
env: A ravens environment.
size_range: Used to indicate the area of the target, so the beads
avoid spawning there.
info: Stores relevant stuff, such as for ground-truth targets.
"""
num_parts = self.num_parts
radius = self.radius
length = self.length
# Add beaded cable.
distance = length / num_parts
position, _ = self.random_pose(env, size_range)
position = np.float32(position)
part_shape = p.createCollisionShape(p.GEOM_BOX,
halfExtents=[radius] * 3)
part_visual = p.createVisualShape(p.GEOM_SPHERE, radius=radius * 1.5)
# Iterate through parts and create constraints as needed.
for i in range(num_parts):
position[2] += distance
parent_frame = (0, 0, distance)
part_id = p.createMultiBody(0.1,
part_shape,
part_visual,
basePosition=position)
if i > 0:
constraint_id = p.createConstraint(
parentBodyUniqueId=env.objects[-1],
parentLinkIndex=-1,
childBodyUniqueId=part_id,
childLinkIndex=-1,
jointType=p.JOINT_POINT2POINT,
jointAxis=(0, 0, 0),
parentFramePosition=parent_frame,
childFramePosition=(0, 0, 0))
p.changeConstraint(constraint_id, maxForce=100)
# Colors
if (i > 0) and (i < num_parts - 1):
p.changeVisualShape(part_id, -1, rgbaColor=self.color_bead)
elif i == num_parts - 1:
p.changeVisualShape(part_id, -1, rgbaColor=self.color_end)
# Add objects in a consistent manner.
self.cable_bead_ids.append(part_id)
env.objects.append(part_id)
self.object_points[part_id] = np.float32(([0], [0], [0]))
# Get target placing positions for each cable bead, if applicable.
if (self._name == 'cable-shape'
or self._name == 'cable-shape-notarget'
or self._name == 'cable-line-notarget'):
# ----------------------------------------------------------- #
# Here, zone_pose = square_pose, unlike Ravens cable, where the
# zone_pose is shifted so that its center matches the straight
# line segment center. For `true_position`, we use `zone_pose`
# but apply the correct offset to deal with the sides. Note
# that `length` is the size of a fully smoothed cable, BUT we
# made a rectangle with each side <= length.
# ----------------------------------------------------------- #
lx = info['lengthx']
ly = info['lengthy']
r = radius
if info['nb_sides'] == 1:
# Here it's just a straight line on the 'lx' side.
x_coord = lx / 2 - (distance * i)
y_coord = 0
true_position = (x_coord - r, y_coord, 0)
elif info['nb_sides'] == 2:
# Start from lx side, go 'left' to the pivot point, then on
# the ly side, go 'upwards' but offset by `i`. For radius
# offset, I just got this by tuning. XD
if i < info['cutoff']:
x_coord = lx / 2 - (distance * i)
y_coord = -ly / 2
true_position = (x_coord - r, y_coord, 0)
else:
x_coord = -lx / 2
y_coord = -ly / 2 + (distance * (i - info['cutoff']))
true_position = (x_coord, y_coord + r, 0)
elif info['nb_sides'] == 3:
# Start from positive lx, positive ly, go down to first
# pivot. Then go left to the second pivot, then up again.
# For v1, division by two is because we assume BOTH of the
# 'ly edges' were divided by two.
v1 = (self.num_parts - info['cutoff']) / 2
v2 = self.num_parts - v1
if i < v1:
x_coord = lx / 2
y_coord = ly / 2 - (distance * i)
true_position = (x_coord, y_coord - r, 0)
elif i < v2:
x_coord = lx / 2 - (distance * (i - v1))
y_coord = -ly / 2
true_position = (x_coord - r, y_coord, 0)
else:
x_coord = -lx / 2
y_coord = -ly / 2 + (distance * (i - v2))
true_position = (x_coord, y_coord + r, 0)
elif info['nb_sides'] == 4:
# I think this is similar to the 2-side case: we start in
# the same direction and go counter-clockwise.
v1 = info['cutoff'] / 2
v2 = num_parts / 2
v3 = (num_parts + info['cutoff']) / 2
if i < v1:
x_coord = lx / 2 - (distance * i)
y_coord = -ly / 2
true_position = (x_coord, y_coord, 0)
elif i < v2:
x_coord = -lx / 2
y_coord = -ly / 2 + (distance * (i - v1))
true_position = (x_coord, y_coord, 0)
elif i < v3:
x_coord = -lx / 2 + (distance * (i - v2))
y_coord = ly / 2
true_position = (x_coord, y_coord, 0)
else:
x_coord = lx / 2
y_coord = ly / 2 - (distance * (i - v3))
true_position = (x_coord, y_coord, 0)
# Map true_position onto the workspace from zone_pose.
true_position = U.apply(self.zone_pose, true_position)
# See `cable.py`: just get the places and steps set.
self.goal['places'][part_id] = (true_position, (0, 0, 0, 1.))
symmetry = 0
self.goal['steps'][0][part_id] = (symmetry, [part_id])
# Debugging target zones.
if self.target_debug_markers:
sq_pose = ((true_position[0], true_position[1], 0.002),
(0, 0, 0, 1))
sq_template = 'assets/square/square-template-allsides-blue.urdf'
replace = {'DIM': (0.003, ), 'HALF': (0.003 / 2, )}
urdf = self.fill_template(sq_template, replace)
env.add_object(urdf, sq_pose, fixed=True)
os.remove(urdf)
else:
print(f'Warning, env {self._name} will not have goals.')
def reset(self, env):
# Add kit.
kit_size = (0.28, 0.2, 0.005)
kit_urdf = 'assets/kitting/kit.urdf'
kit_pose = self.random_pose(env, kit_size)
env.add_object(kit_urdf, kit_pose, fixed=True)
# num_shapes = 20
# train_split = 14
num_objects = 5
if self.mode == 'train':
object_shapes = np.random.choice(self.train_set, num_objects)
else:
if self.homogeneous:
object_shapes = [np.random.choice(self.test_set)] * num_objects
else:
object_shapes = np.random.choice(self.test_set, num_objects)
self.num_steps = num_objects
self.goal = {'places': {}, 'steps': [{}]}
colors = [
utils.COLORS['purple'], utils.COLORS['blue'],
utils.COLORS['green'], utils.COLORS['yellow'], utils.COLORS['red']
]
symmetries = [
2 * np.pi, 2 * np.pi, 2 * np.pi / 3, np.pi / 2, np.pi / 2,
2 * np.pi, np.pi, 2 * np.pi / 5, np.pi, np.pi / 2, 2 * np.pi / 5,
0, 2 * np.pi, 2 * np.pi, 2 * np.pi, 2 * np.pi, 0, 2 * np.pi / 6,
2 * np.pi, 2 * np.pi
]
# Build kit.
place_positions = [[-0.09, 0.045, 0.0014], [0, 0.045, 0.0014],
[0.09, 0.045, 0.0014], [-0.045, -0.045, 0.0014],
[0.045, -0.045, 0.0014]]
object_template = 'assets/kitting/object-template.urdf'
for i in range(num_objects):
shape = f'{object_shapes[i]:02d}.obj'
scale = [0.003, 0.003, 0.0001] # .0005
position = utils.apply(kit_pose, place_positions[i])
theta = np.random.rand() * 2 * np.pi
rotation = utils.get_pybullet_quaternion_from_rot((0, 0, theta))
replace = {
'FNAME': (shape, ),
'SCALE': scale,
'COLOR': (0.2, 0.2, 0.2)
}
urdf = self.fill_template(object_template, replace)
_ = env.add_object(urdf, (position, rotation), fixed=True)
os.remove(urdf)
self.goal['places'][i] = (position, rotation)
# Add objects.
for i in range(num_objects):
ishape = object_shapes[i]
size = (0.08, 0.08, 0.02)
pose = self.random_pose(env, size)
shape = f'{ishape:02d}.obj'
scale = [0.003, 0.003, 0.001] # .0005
replace = {'FNAME': (shape, ), 'SCALE': scale, 'COLOR': colors[i]}
urdf = self.fill_template(object_template, replace)
block_id = env.add_object(urdf, pose)
os.remove(urdf)
places = np.argwhere(ishape == object_shapes).reshape(-1)
self.goal['steps'][0][block_id] = (symmetries[ishape], places)
def reset(self, env):
super().reset(env)
# Add kit.
kit_size = (0.28, 0.2, 0.005)
kit_urdf = 'assets/kitting/kit.urdf'
kit_pose = self.get_random_pose(env, kit_size)
env.add_object(kit_urdf, kit_pose, 'fixed')
n_objects = 5
if self.mode == 'train':
obj_shapes = np.random.choice(self.train_set, n_objects)
else:
if self.homogeneous:
obj_shapes = [np.random.choice(self.test_set)] * n_objects
else:
obj_shapes = np.random.choice(self.test_set, n_objects)
colors = [
utils.COLORS['purple'], utils.COLORS['blue'], utils.COLORS['green'],
utils.COLORS['yellow'], utils.COLORS['red']
]
symmetry = [
2 * np.pi, 2 * np.pi, 2 * np.pi / 3, np.pi / 2, np.pi / 2, 2 * np.pi,
np.pi, 2 * np.pi / 5, np.pi, np.pi / 2, 2 * np.pi / 5, 0, 2 * np.pi,
2 * np.pi, 2 * np.pi, 2 * np.pi, 0, 2 * np.pi / 6, 2 * np.pi, 2 * np.pi
]
# Build kit.
targets = []
targ_pos = [[-0.09, 0.045, 0.0014], [0, 0.045, 0.0014],
[0.09, 0.045, 0.0014], [-0.045, -0.045, 0.0014],
[0.045, -0.045, 0.0014]]
template = 'assets/kitting/object-template.urdf'
for i in range(n_objects):
shape = f'{obj_shapes[i]:02d}.obj'
scale = [0.003, 0.003, 0.0001] # .0005
pos = utils.apply(kit_pose, targ_pos[i])
theta = np.random.rand() * 2 * np.pi
rot = utils.eulerXYZ_to_quatXYZW((0, 0, theta))
replace = {'FNAME': (shape,), 'SCALE': scale, 'COLOR': (0.2, 0.2, 0.2)}
urdf = self.fill_template(template, replace)
env.add_object(urdf, (pos, rot), 'fixed')
os.remove(urdf)
targets.append((pos, rot))
# Add objects.
objects = []
matches = []
# objects, syms, matcheses = [], [], []
for i in range(n_objects):
shape = obj_shapes[i]
size = (0.08, 0.08, 0.02)
pose = self.get_random_pose(env, size)
fname = f'{shape:02d}.obj'
scale = [0.003, 0.003, 0.001] # .0005
replace = {'FNAME': (fname,), 'SCALE': scale, 'COLOR': colors[i]}
urdf = self.fill_template(template, replace)
block_id = env.add_object(urdf, pose)
os.remove(urdf)
objects.append((block_id, (symmetry[shape], None)))
# objects[block_id] = symmetry[shape]
match = np.zeros(len(targets))
match[np.argwhere(obj_shapes == shape).reshape(-1)] = 1
matches.append(match)
# print(targets)
# exit()
# matches.append(list(np.argwhere(obj_shapes == shape).reshape(-1)))
matches = np.int32(matches)
# print(matcheses)
# exit()
# Add goal.
# self.goals.append((objects, syms, targets, 'matches', 'pose', 1.))
# Goal: objects are placed in their respective kit locations.
# print(objects)
# print(matches)
# print(targets)
# exit()
self.goals.append((objects, matches, targets, False, True, 'pose', None, 1))
def reset(self, env):
self.total_rewards = 0
self.goal = {'places': {}, 'steps': [{}]}
num_parts = 20
radius = 0.005
length = 2 * radius * num_parts * np.sqrt(2)
square_size = (length, length, 0)
square_pose = self.random_pose(env, square_size)
square_template = 'assets/square/square-template.urdf'
replace = {'DIM': (length,), 'HALF': (length / 2 - 0.005,)}
urdf = self.fill_template(square_template, replace)
env.add_object(urdf, square_pose, fixed=True)
os.remove(urdf)
# Add goal line.
# line_template = 'assets/line/line-template.urdf'
self.zone_size = (length, 0.03, 0.2)
zone_range = (self.zone_size[0], self.zone_size[1], 0.001)
zone_position = (0, length / 2, 0.001)
zone_position = utils.apply(square_pose, zone_position)
self.zone_pose = (zone_position, square_pose[1])
# urdf = self.fill_template(line_template, {'DIM': (length,)})
# env.add_object(urdf, self.zone_pose, fixed=True)
# os.remove(urdf)
# Add beaded cable.
distance = length / num_parts
position, _ = self.random_pose(env, zone_range)
position = np.float32(position)
part_shape = p.createCollisionShape(p.GEOM_BOX, halfExtents=[radius] * 3)
part_visual = p.createVisualShape(p.GEOM_SPHERE, radius=radius * 1.5)
# part_visual = p.createVisualShape(p.GEOM_BOX, halfExtents=[radius] * 3)
self.object_points = {}
for i in range(num_parts):
position[2] += distance
part_id = p.createMultiBody(
0.1, part_shape, part_visual, basePosition=position)
if env.objects:
constraint_id = p.createConstraint(
parentBodyUniqueId=env.objects[-1],
parentLinkIndex=-1,
childBodyUniqueId=part_id,
childLinkIndex=-1,
jointType=p.JOINT_POINT2POINT,
jointAxis=(0, 0, 0),
parentFramePosition=(0, 0, distance),
childFramePosition=(0, 0, 0))
p.changeConstraint(constraint_id, maxForce=100)
if (i > 0) and (i < num_parts - 1):
color = utils.COLORS['red'] + [1]
p.changeVisualShape(part_id, -1, rgbaColor=color)
env.objects.append(part_id)
self.object_points[part_id] = np.float32((0, 0, 0)).reshape(3, 1)
true_position = (radius + distance * i - length / 2, 0, 0)
true_position = utils.apply(self.zone_pose, true_position)
self.goal['places'][part_id] = (true_position, (0, 0, 0, 1.))
symmetry = 0 # zone-evaluation: symmetry does not matter
self.goal['steps'][0][part_id] = (symmetry, [part_id])
# Wait for beaded cable to settle.
env.start()
time.sleep(1)
env.pause()
def reset(self, env):
super().reset(env)
n_parts = 20
radius = 0.005
length = 2 * radius * n_parts * np.sqrt(2)
# Add 3-sided square.
square_size = (length, length, 0)
square_pose = self.get_random_pose(env, square_size)
square_template = 'assets/square/square-template.urdf'
replace = {'DIM': (length, ), 'HALF': (length / 2 - 0.005, )}
urdf = self.fill_template(square_template, replace)
env.add_object(urdf, square_pose, 'fixed')
os.remove(urdf)
# Get corner points of square.
corner0 = (length / 2, length / 2, 0.001)
corner1 = (-length / 2, length / 2, 0.001)
corner0 = utils.apply(square_pose, corner0)
corner1 = utils.apply(square_pose, corner1)
# Add cable (series of articulated small blocks).
increment = (np.float32(corner1) - np.float32(corner0)) / n_parts
position, _ = self.get_random_pose(env, (0.1, 0.1, 0.1))
position = np.float32(position)
part_shape = p.createCollisionShape(p.GEOM_BOX,
halfExtents=[radius] * 3)
part_visual = p.createVisualShape(p.GEOM_SPHERE, radius=radius * 1.5)
parent_id = -1
targets = []
objects = []
for i in range(n_parts):
position[2] += np.linalg.norm(increment)
part_id = p.createMultiBody(0.1,
part_shape,
part_visual,
basePosition=position)
if parent_id > -1:
constraint_id = p.createConstraint(
parentBodyUniqueId=parent_id,
parentLinkIndex=-1,
childBodyUniqueId=part_id,
childLinkIndex=-1,
jointType=p.JOINT_POINT2POINT,
jointAxis=(0, 0, 0),
parentFramePosition=(0, 0, np.linalg.norm(increment)),
childFramePosition=(0, 0, 0))
p.changeConstraint(constraint_id, maxForce=100)
if (i > 0) and (i < n_parts - 1):
color = utils.COLORS['red'] + [1]
p.changeVisualShape(part_id, -1, rgbaColor=color)
env.obj_ids['rigid'].append(part_id)
parent_id = part_id
target_xyz = np.float32(corner0) + i * increment + increment / 2
objects.append((part_id, (0, None)))
targets.append((target_xyz, (0, 0, 0, 1)))
matches = np.clip(np.eye(n_parts) + np.eye(n_parts)[::-1], 0, 1)
self.goals.append(
(objects, matches, targets, False, False, 'pose', None, 1))
for i in range(480):
p.stepSimulation()
