File size: 2,568 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 60 61 62 |
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils
class BallOnBoxOnContainer(Task):
"""Pick up each ball and place it on the corresponding colored box, which are located in specific positions on a container."""
def __init__(self):
super().__init__()
self.max_steps = 15
self.lang_template = "place the {} box on the container"
self.lang_template_2 = "place the {} ball on the {} box"
self.task_completed_desc = "done placing balls on boxs and box on container."
self.additional_reset()
def reset(self, env):
super().reset(env)
# Add container.
container_size = (0.2, 0.2, 0.06)
container_pose = self.get_random_pose(env, container_size)
container_template = 'container/container-template.urdf'
replace = {'DIM': container_size, 'HALF': (container_size[0] / 2, container_size[1] / 2, container_size[2] / 2)}
container_urdf = self.fill_template(container_template, replace)
env.add_object(container_urdf, container_pose, 'fixed')
# Define colors.
ball_colors = ['red']
box_colors = ['blue']
# Add boxs.
box_size = (0.04, 0.04, 0.06)
box_template = 'box/box-template.urdf'
boxs = []
replace = {'DIM': box_size, 'HALF': (box_size[0] / 2, box_size[1] / 2, box_size[2] / 2), 'COLOR': ball_colors[0]}
box_urdf = self.fill_template(box_template, replace)
box_pose = self.get_random_pose(env, box_size)
box_id = env.add_object(box_urdf, box_pose)
boxs.append(box_id)
# Add balls.
ball_size = (0.04, 0.04, 0.04)
ball_urdf = 'ball/ball.urdf'
balls = []
ball_pose = self.get_random_pose(env, ball_size)
ball_id = env.add_object(ball_urdf, ball_pose, color=box_colors[0])
balls.append(ball_id)
# Goal: place the box on top of the container
self.add_goal(objs=[boxs[0]], matches=np.ones((1, 1)), targ_poses=[container_pose], replace=False,
rotations=True, metric='pose', params=None, step_max_reward=1/2, language_goal=self.lang_template.format(box_colors[0]))
# Goal: place the ball on top of the box
language_goal = self.lang_template_2.format(ball_colors[0], box_colors[0])
self.add_goal(objs=[balls[0]], matches=np.ones((1, 1)), targ_poses=[container_pose], replace=False,
rotations=True, metric='pose', params=None, step_max_reward=1/2, language_goal=language_goal)
|