Spaces:
Runtime error
Runtime error
File size: 3,849 Bytes
8fc2b4e |
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 |
import os
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils
class AssemblingKitsSeq(Task):
""" Precisely place each specified shape in the specified hole following the order prescribed in the
language instruction at each timestep."""
def __init__(self):
super().__init__()
self.max_steps = 7
self.homogeneous = False
self.lang_template = "put the {color} {obj} in the {loc}{obj} hole"
self.task_completed_desc = "done assembling kit."
self.additional_reset()
def reset(self, env):
super().reset(env)
# Add kit.
kit_size = (0.28, 0.2, 0.005)
kit_urdf = 'kitting/kit.urdf'
kit_pose = self.get_random_pose(env, kit_size)
env.add_object(kit_urdf, kit_pose, 'fixed')
# Shape Names:
shapes = utils.assembling_kit_shapes
n_objects = 5
obj_shapes = self.get_kitting_shapes(n_objects)
colors, color_names = utils.get_colors(mode=self.mode)
# Build kit.
targets = []
targets_spatial_desc = []
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 = 'kitting/object-template.urdf'
for i in range(n_objects):
shape = os.path.join(self.assets_root, 'kitting',
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]}
# IMPORTANT: REPLACE THE TEMPLATE URDF
urdf = self.fill_template(template, replace)
env.add_object(urdf, (pos, rot), 'fixed')
targets.append((pos, rot))
# Decide spatial description based on the location of the hole (top-down view).
shape_type = obj_shapes[i]
if list(obj_shapes).count(obj_shapes[i]) > 1:
duplicate_shapes = [j for j, o in enumerate(obj_shapes) if i != j and o == shape_type]
other_poses = [utils.apply(kit_pose, targ_pos[d]) for d in duplicate_shapes]
if all(pos[0] < op[0] and abs(pos[0]-op[0]) > abs(pos[1]-op[1]) for op in other_poses):
spatial_desc = "top "
elif all(pos[0] > op[0] and abs(pos[0]-op[0]) > abs(pos[1]-op[1]) for op in other_poses):
spatial_desc = "bottom "
elif all(pos[1] < op[1] for op in other_poses):
spatial_desc = "left "
elif all(pos[1] > op[1] for op in other_poses):
spatial_desc = "right "
else:
spatial_desc = "middle "
targets_spatial_desc.append(spatial_desc)
else:
targets_spatial_desc.append("")
# Add objects.
objects, matches = self.make_kitting_objects(env, targets=targets, obj_shapes=obj_shapes, n_objects=n_objects, colors=colors)
target_idxs = list(range(n_objects))
np.random.shuffle(target_idxs)
for i in target_idxs:
self.add_goal(objs=[objects[i]], matches=np.int32([[1]]), targ_poses=[targets[i]], replace=False,
rotations=True, metric='pose', params=None, step_max_reward=1 / n_objects)
self.lang_goals.append(self.lang_template.format(color=color_names[i],
obj=shapes[obj_shapes[i]],
loc=targets_spatial_desc[i]))
self.max_steps = n_objects
|