File size: 1,637 Bytes
ff66cf3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
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) |