File size: 2,251 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 45 46 47 48 49 50 |
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
class BuildCircle(Task):
"""Pick up six blocks of different colors (red, blue, green, yellow, orange, and purple)
and place them on a tabletop in a circle arrangement. The arrangement should start with
red at the top and continue clockwise in this order: blue, green, yellow, orange, and finally purple."""
def __init__(self):
super().__init__()
self.max_steps = 20
self.lang_template = "place the {color} block at the {position} of the circle"
self.task_completed_desc = "done building circle."
self.colors = ['red', 'blue', 'green', 'yellow', 'orange', 'purple']
self.positions = ['top', 'top right', 'bottom right', 'bottom', 'bottom left', 'top left']
self.additional_reset()
def reset(self, env):
super().reset(env)
# Add blocks.
block_size = (0.04, 0.04, 0.04)
block_urdf = 'block/block.urdf'
blocks = []
for i in range(6):
block_pose = self.get_random_pose(env, block_size)
block_id = env.add_object(block_urdf, block_pose, color=utils.COLORS[self.colors[i]])
blocks.append(block_id)
# Define target poses for the blocks in a circle arrangement.
radius = 0.1
center = (0.5, 0.5, 0)
angles = np.linspace(0, 2*np.pi, 7)[:-1]
targ_poses = [(center[0] + radius*np.cos(angle), center[1] + radius*np.sin(angle), block_size[2]/2) for angle in angles]
targ_poses = [(pose, (0, 0, 0, 1)) for pose in targ_poses] # add default quaternion for orientation
# Add goals.
for i in range(6):
language_goal = self.lang_template.format(color=self.colors[i], position=self.positions[i])
self.add_goal(objs=[blocks[i]], matches=np.ones((1, 1)), targ_poses=[targ_poses[i]], replace=False,
rotations=False, metric='pose', params=None, step_max_reward=1/6, language_goal=language_goal) |