Spaces:
Runtime error
Runtime error
| 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 | |
| import pybullet as p | |
| class AlignSpheresInBoxes(Task): | |
| """Pick up each sphere and place it into the box of the same color. Navigate around the blocks.""" | |
| def __init__(self): | |
| super().__init__() | |
| self.max_steps = 20 | |
| self.lang_template = "put the {color} sphere in the {color} box" | |
| self.task_completed_desc = "done aligning spheres in boxes." | |
| self.additional_reset() | |
| def reset(self, env): | |
| super().reset(env) | |
| # Add boxes. | |
| # x, y, z dimensions for the asset size | |
| box_size = (0.12, 0.12, 0.12) | |
| box_urdf = 'box/box-template.urdf' | |
| box_poses = [] | |
| box_colors = ['red', 'blue', 'green', 'yellow'] | |
| for color in box_colors: | |
| box_pose = self.get_random_pose(env, box_size) | |
| env.add_object(box_urdf, box_pose, category='fixed', color=utils.COLORS[color]) | |
| box_poses.append(box_pose) | |
| # Add spheres. | |
| # x, y, z dimensions for the asset size | |
| spheres = [] | |
| sphere_size = (0.04, 0.04, 0.04) | |
| sphere_urdf = 'sphere/sphere-template.urdf' | |
| for color in box_colors: | |
| sphere_pose = self.get_random_pose(env, sphere_size) | |
| sphere_id = env.add_object(sphere_urdf, sphere_pose, color=utils.COLORS[color]) | |
| spheres.append(sphere_id) | |
| # Add blocks. | |
| # x, y, z dimensions for the asset size | |
| blocks = [] | |
| block_size = (0.04, 0.04, 0.04) | |
| block_urdf = 'block/small.urdf' | |
| for _ in range(10): | |
| block_pose = self.get_random_pose(env, block_size) | |
| block_id = env.add_object(block_urdf, block_pose) | |
| blocks.append(block_id) | |
| # Goal: each sphere is in a box of the same color. | |
| for i in range(len(spheres)): | |
| self.add_goal(objs=[spheres[i]], matches=np.ones((1, 1)), targ_poses=[box_poses[i]], replace=False, | |
| rotations=True, metric='pose', params=None, step_max_reward=1/len(spheres)) | |
| self.lang_goals.append(self.lang_template.format(color=box_colors[i])) |