Spaces:
Runtime error
Runtime error
| import os | |
| import numpy as np | |
| from cliport.tasks.task import Task | |
| from cliport.utils import utils | |
| import pybullet as p | |
| class PalletizingBoxes(Task): | |
| """Pick up homogeneous fixed-sized boxes and stack them in transposed layers on the pallet.""" | |
| def __init__(self): | |
| super().__init__() | |
| self.max_steps = 30 | |
| self.lang_template = "stack all the boxes on the pallet" | |
| self.task_completed_desc = "done stacking boxes." | |
| self.additional_reset() | |
| def reset(self, env): | |
| super().reset(env) | |
| # Add pallet. | |
| zone_size = (0.3, 0.25, 0.25) | |
| zone_urdf = 'pallet/pallet.urdf' | |
| rotation = utils.eulerXYZ_to_quatXYZW((0, 0, 0)) | |
| zone_pose = ((0.5, 0.25, 0.02), rotation) | |
| env.add_object(zone_urdf, zone_pose, 'fixed') | |
| # Add stack of boxes on pallet. | |
| margin = 0.01 | |
| object_ids = [] | |
| # x, y, z dimensions for the asset size | |
| stack_size = (0.19, 0.19, 0.19) | |
| box_template = 'box/box-template.urdf' | |
| stack_dim = np.int32([2, 3, 3]) | |
| box_size = (stack_size - (stack_dim - 1) * margin) / stack_dim | |
| for z in range(stack_dim[2]): | |
| # Transpose every layer. | |
| stack_dim[0], stack_dim[1] = stack_dim[1], stack_dim[0] | |
| box_size[0], box_size[1] = box_size[1], box_size[0] | |
| # IMPORTANT: Compute object points and store as a dictionary for the `goal` | |
| for y in range(stack_dim[1]): | |
| for x in range(stack_dim[0]): | |
| position = list((x + 0.5, y + 0.5, z + 0.5) * box_size) | |
| position[0] += x * margin - stack_size[0] / 2 | |
| position[1] += y * margin - stack_size[1] / 2 | |
| position[2] += z * margin + 0.03 | |
| pose = (position, (0, 0, 0, 1)) | |
| pose = utils.multiply(zone_pose, pose) | |
| # IMPORTANT: REPLACE THE TEMPLATE URDF | |
| urdf = self.fill_template(box_template, {'DIM': box_size}) | |
| box_id = env.add_object(urdf, pose) | |
| object_ids.append(box_id) | |
| self.color_random_brown(box_id) | |
| # Randomly select top box on pallet and save ground truth pose. | |
| targets = [] | |
| self.steps = [] | |
| boxes = object_ids[:] # make copy | |
| while boxes: | |
| _, height, object_mask = self.get_true_image(env) | |
| top = np.argwhere(height > (np.max(height) - 0.03)) | |
| rpixel = top[int(np.floor(np.random.random() * len(top)))] # y, x | |
| box_id = int(object_mask[rpixel[0], rpixel[1]]) | |
| if box_id in boxes: | |
| position, rotation = p.getBasePositionAndOrientation(box_id) | |
| rposition = np.float32(position) + np.float32([0, -10, 0]) | |
| p.resetBasePositionAndOrientation(box_id, rposition, rotation) | |
| self.steps.append(box_id) | |
| targets.append((position, rotation)) | |
| boxes.remove(box_id) | |
| self.steps.reverse() # Time-reversed depalletizing. | |
| self.add_goal(objs=object_ids, matches=np.eye(len(object_ids)), targ_poses=targets, replace=False, | |
| rotations=True, metric='zone', params=[(zone_pose, zone_size)], step_max_reward=1) | |
| self.lang_goals.append(self.lang_template) | |
| self.spawn_box() | |
| def reward(self): | |
| reward, info = super().reward() | |
| self.spawn_box() | |
| return reward, info |