File size: 2,496 Bytes
ff66cf3 |
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 |
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils
class BlockOnCylinderOnPallet(Task):
"""Pick up each block and place it on the corresponding colored cylinder, which are located in specific positions on a pallet."""
def __init__(self):
super().__init__()
self.max_steps = 15
self.lang_template = "place the {} cylinder on the pallet"
self.lang_template_2 = "place the {} block on the {} cylinder"
self.task_completed_desc = "done placing blocks on cylinders and cylinder on pallet."
self.additional_reset()
def reset(self, env):
super().reset(env)
# Add pallet.
pallet_size = (0.35, 0.35, 0.01)
pallet_pose = self.get_random_pose(env, pallet_size)
pallet_urdf = 'pallet/pallet.urdf'
env.add_object(pallet_urdf, pallet_pose, 'fixed')
# Define colors.
block_colors = ['red']
cylinder_colors = ['blue']
# Add cylinders.
cylinder_size = (0.04, 0.04, 0.06)
cylinder_template = 'cylinder/cylinder-template.urdf'
cylinders = []
replace = {'DIM': cylinder_size, 'HALF': (cylinder_size[0] / 2, cylinder_size[1] / 2, cylinder_size[2] / 2), 'COLOR': block_colors[0]}
cylinder_urdf = self.fill_template(cylinder_template, replace)
cylinder_pose = self.get_random_pose(env, cylinder_size)
cylinder_id = env.add_object(cylinder_urdf, cylinder_pose)
cylinders.append(cylinder_id)
# Add blocks.
block_size = (0.04, 0.04, 0.04)
block_urdf = 'block/block.urdf'
blocks = []
block_pose = self.get_random_pose(env, block_size)
block_id = env.add_object(block_urdf, block_pose, color=cylinder_colors[0])
blocks.append(block_id)
# Goal: place the cylinder on top of the pallet
self.add_goal(objs=[cylinders[0]], matches=np.ones((1, 1)), targ_poses=[pallet_pose], replace=False,
rotations=True, metric='pose', params=None, step_max_reward=1/2, language_goal=self.lang_template.format(cylinder_colors[0]))
# Goal: place the block on top of the cylinder
language_goal = self.lang_template_2.format(block_colors[0], cylinder_colors[0])
self.add_goal(objs=[blocks[0]], matches=np.ones((1, 1)), targ_poses=[pallet_pose], replace=False,
rotations=True, metric='pose', params=None, step_max_reward=1/2, language_goal=language_goal)
|