| import numpy as np |
| import os |
| import pybullet as p |
| import random |
| from cliport.tasks import primitives |
| from cliport.tasks.grippers import Spatula |
| from cliport.tasks.task import Task |
| from cliport.utils import utils |
| import numpy as np |
| from cliport.tasks.task import Task |
| from cliport.utils import utils |
|
|
| class BuildCircle(Task): |
| """Pick up six blocks of different colors (red, blue, green, yellow, orange, and purple) |
| and place them on a tabletop in a circle arrangement. The arrangement should start with |
| red at the top and continue clockwise in this order: blue, green, yellow, orange, and finally purple.""" |
|
|
| def __init__(self): |
| super().__init__() |
| self.max_steps = 20 |
| self.lang_template = "place the {color} block at the {position} of the circle" |
| self.task_completed_desc = "done building circle." |
| self.colors = ['red', 'blue', 'green', 'yellow', 'orange', 'purple'] |
| self.positions = ['top', 'top right', 'bottom right', 'bottom', 'bottom left', 'top left'] |
| self.additional_reset() |
|
|
| def reset(self, env): |
| super().reset(env) |
|
|
| # Add blocks. |
| block_size = (0.04, 0.04, 0.04) |
| block_urdf = 'block/block.urdf' |
| blocks = [] |
| for i in range(6): |
| block_pose = self.get_random_pose(env, block_size) |
| block_id = env.add_object(block_urdf, block_pose, color=utils.COLORS[self.colors[i]]) |
| blocks.append(block_id) |
|
|
| # Define target poses for the blocks in a circle arrangement. |
| radius = 0.1 |
| center = (0.5, 0.5, 0) |
| angles = np.linspace(0, 2*np.pi, 7)[:-1] |
| targ_poses = [(center[0] + radius*np.cos(angle), center[1] + radius*np.sin(angle), block_size[2]/2) for angle in angles] |
| targ_poses = [(pose, (0, 0, 0, 1)) for pose in targ_poses] # add default quaternion for orientation |
|
|
| # Add goals. |
| for i in range(6): |
| language_goal = self.lang_template.format(color=self.colors[i], position=self.positions[i]) |
| self.add_goal(objs=[blocks[i]], matches=np.ones((1, 1)), targ_poses=[targ_poses[i]], replace=False, |
| rotations=False, metric='pose', params=None, step_max_reward=1/6, language_goal=language_goal) |
