| 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 import primitives | |
| from cliport.tasks.grippers import Spatula | |
| from cliport.tasks.task import Task | |
| from cliport.utils import utils | |
| class SweepRedBlocksIntoZone(Task): | |
| """Sweep a pile of red blocks into a designated zone.""" | |
| def __init__(self): | |
| super().__init__() | |
| self.max_steps = 20 | |
| self.lang_template = "sweep the pile of red blocks into the green square" | |
| self.task_completed_desc = "done sweeping." | |
| self.primitive = primitives.push | |
| self.ee = Spatula | |
| self.additional_reset() | |
| def reset(self, env): | |
| super().__init__(env) | |
| # Add goal zone. | |
| zone_size = (0.12, 0.12, 0) | |
| zone_pose = self.get_random_pose(env, zone_size) | |
| env.add_object('zone/zone.urdf', zone_pose, 'fixed') | |
| # Add pile of red blocks with `make_piles` function | |
| block_urdf = 'stacking/block.urdf' | |
| block_size = (0.04, 0.04, 0.04) | |
| block_color = utils.COLORS['red'] | |
| pile_pose = self.get_random_pose(env, block_size) | |
| pile_ids = self.make_piles(env, block_urdf, block_size, block_color, pile_pose, num_piles=1, num_objs=5) | |
| # Add goal | |
| self.add_goal(objs=pile_ids, matches=np.ones((5, 1)), targ_poses=[zone_pose], replace=True, | |
| rotations=False, metric='zone', params=[(zone_pose, zone_size)], step_max_reward=1, language_goal=self.lang_template) |