| 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 ContainerPyramidConstruction(Task): | |
| """Construct a pyramid of containers with specific color and shape arrangement.""" | |
| def __init__(self): | |
| super().__init__() | |
| self.max_steps = 10 | |
| self.lang_template = "stack the {color} {shape} container on the {level} level" | |
| self.task_completed_desc = "done constructing container pyramid." | |
| 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 container shapes and colors. | |
| shapes = ['square', 'circle', 'triangle'] | |
| colors = ['red', 'blue', 'green'] | |
| levels = ['bottom', 'middle', 'top'] | |
| # Add containers. | |
| container_urdf = 'container/container-template.urdf' | |
| containers = [] | |
| for i in range(3): | |
| for j in range(3 - i): | |
| # Define container size and color. | |
| container_size = (0.1, 0.1, 0.1) | |
| container_color = colors[i] | |
| # Define container pose. | |
| x_offset = 0.05 * (j - (3 - i - 1) / 2) | |
| z_offset = 0.1 * i | |
| container_pose = (pallet_pose[0] + x_offset, pallet_pose[1], pallet_pose[2] + z_offset) | |
| # Add container to the environment. | |
| replace = {'DIM': container_size, 'HALF': (container_size[0] / 2, container_size[1] / 2, container_size[2] / 2)} | |
| container_urdf_filled = self.fill_template(container_urdf, replace) | |
| container_id = env.add_object(container_urdf_filled, container_pose, color=container_color) | |
| containers.append(container_id) | |
| # Add goals. | |
| for i in range(3): | |
| for j in range(3 - i): | |
| idx = int(i * (i + 1) / 2 + j) | |
| language_goal = self.lang_template.format(color=colors[i], shape=shapes[j], level=levels[i]) | |
| self.add_goal(objs=[containers[idx]], matches=np.ones((1, 1)), targ_poses=[pallet_pose], replace=False, | |
| rotations=True, metric='pose', params=None, step_max_reward=1 / 6, language_goal=language_goal) |