Spaces:
Runtime error
Runtime error
| import numpy as np | |
| from cliport.tasks.task import Task | |
| from cliport.utils import utils | |
| import pybullet as p | |
| class PlaceRedInGreen(Task): | |
| """pick up the red blocks and place them into the green bowls amidst other objects.""" | |
| def __init__(self): | |
| super().__init__() | |
| self.max_steps = 10 | |
| self.lang_template = "put the red blocks in a green bowl" | |
| self.task_completed_desc = "done placing blocks in bowls." | |
| self.additional_reset() | |
| def reset(self, env): | |
| super().reset(env) | |
| n_bowls = np.random.randint(1, 4) | |
| n_blocks = np.random.randint(1, n_bowls + 1) | |
| # Add bowls. | |
| # x, y, z dimensions for the asset size | |
| bowl_size = (0.12, 0.12, 0) | |
| bowl_urdf = 'bowl/bowl.urdf' | |
| bowl_poses = [] | |
| for _ in range(n_bowls): | |
| bowl_pose = self.get_random_pose(env, obj_size=bowl_size) | |
| env.add_object(urdf=bowl_urdf, pose=bowl_pose, category='fixed') | |
| bowl_poses.append(bowl_pose) | |
| # Add blocks. | |
| # x, y, z dimensions for the asset size | |
| blocks = [] | |
| block_size = (0.04, 0.04, 0.04) | |
| block_urdf = 'stacking/block.urdf' | |
| for _ in range(n_blocks): | |
| block_pose = self.get_random_pose(env, obj_size=block_size) | |
| block_id = env.add_object(block_urdf, block_pose) | |
| blocks.append(block_id) | |
| # Goal: each red block is in a different green bowl. | |
| self.add_goal(objs=blocks, matches=np.ones((len(blocks), len(bowl_poses))), targ_poses=bowl_poses, replace=False, | |
| rotations=True, metric='pose', params=None, step_max_reward=1) | |
| self.lang_goals.append(self.lang_template) | |
| # Colors of distractor objects. | |
| bowl_colors = [utils.COLORS[c] for c in utils.COLORS if c != 'green'] | |
| block_colors = [utils.COLORS[c] for c in utils.COLORS if c != 'red'] | |
| # Add distractors. | |
| n_distractors = 0 | |
| while n_distractors < 6: | |
| is_block = np.random.rand() > 0.5 | |
| urdf = block_urdf if is_block else bowl_urdf | |
| size = block_size if is_block else bowl_size | |
| colors = block_colors if is_block else bowl_colors | |
| pose = self.get_random_pose(env, obj_size=size) | |
| color = colors[n_distractors % len(colors)] | |
| obj_id = env.add_object(urdf, pose, color=color) | |
| n_distractors += 1 | |