Spaces:
Build error
Build error
| 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) | |