Spaces:
Runtime error
Runtime error
| import os | |
| import numpy as np | |
| from cliport.tasks.task import Task | |
| from cliport.utils import utils | |
| class AlignBoxCorner(Task): | |
| """Pick up the randomly sized box and align one of its corners to the L-shaped marker on the tabletop.""" | |
| def __init__(self): | |
| super().__init__() | |
| self.max_steps = 3 | |
| self.lang_template = "align the brown box with the green corner" | |
| self.task_completed_desc = "done with alignment" | |
| self.additional_reset() | |
| def reset(self, env): | |
| super().reset(env) | |
| # Generate randomly shaped box. | |
| box_size = self.get_random_size(0.05, 0.15, 0.05, 0.15, 0.01, 0.06) | |
| # Add corner. | |
| dimx = (box_size[0] / 2 - 0.025 + 0.0025, box_size[0] / 2 + 0.0025) | |
| dimy = (box_size[1] / 2 + 0.0025, box_size[1] / 2 - 0.025 + 0.0025) | |
| corner_template = 'corner/corner-template.urdf' | |
| replace = {'DIMX': dimx, 'DIMY': dimy} | |
| # IMPORTANT: REPLACE THE TEMPLATE URDF | |
| corner_urdf = self.fill_template(corner_template, replace) | |
| corner_size = (box_size[0], box_size[1], 0) | |
| corner_pose = self.get_random_pose(env, corner_size) | |
| env.add_object(corner_urdf, corner_pose, 'fixed') | |
| # Add possible placing poses. | |
| theta = utils.quatXYZW_to_eulerXYZ(corner_pose[1])[2] | |
| fip_rot = utils.eulerXYZ_to_quatXYZW((0, 0, theta + np.pi)) | |
| pose1 = (corner_pose[0], fip_rot) | |
| alt_x = (box_size[0] / 2) - (box_size[1] / 2) | |
| alt_y = (box_size[1] / 2) - (box_size[0] / 2) | |
| alt_pos = (alt_x, alt_y, 0) | |
| alt_rot0 = utils.eulerXYZ_to_quatXYZW((0, 0, np.pi / 2)) | |
| alt_rot1 = utils.eulerXYZ_to_quatXYZW((0, 0, 3 * np.pi / 2)) | |
| pose2 = utils.multiply(corner_pose, (alt_pos, alt_rot0)) | |
| pose3 = utils.multiply(corner_pose, (alt_pos, alt_rot1)) | |
| # Add box. | |
| box_template = 'box/box-template.urdf' | |
| # IMPORTANT: REPLACE THE TEMPLATE URDF | |
| box_urdf = self.fill_template(box_template, {'DIM': np.float32(box_size)}) | |
| box_pose = self.get_random_pose(env, box_size) | |
| box_id = env.add_object(box_urdf, box_pose) | |
| self.color_random_brown(box_id) | |
| # Goal: box is aligned with corner (1 of 4 possible poses). | |
| self.add_goal(objs=[box_id], matches=np.int32([[1, 1, 1, 1]]), targ_poses=[corner_pose, pose1, pose2, pose3], replace=False, | |
| rotations=True, metric='pose', params=None, step_max_reward=1, symmetries=[2 * np.pi]) | |
| self.lang_goals.append(self.lang_template) | |