Spaces:
Runtime error
Runtime error
File size: 2,327 Bytes
8fc2b4e |
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 51 52 53 54 55 56 57 |
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 AlignCylindersInSquare(Task):
"""Arrange four cylinders of different colors (red, blue, green, yellow)
on the corners of a square facing the center. The square is outlined by
four zones of matching colors at each corner. The red cylinder should be
placed at the red zone and facing the center of the square, and same for
the remaining colors."""
def __init__(self):
super().__init__()
self.max_steps = 20
self.lang_template = "arrange the {color} cylinder at the {color} corner of the square"
self.task_completed_desc = "done arranging cylinders."
self.additional_reset()
def reset(self, env):
super().reset(env)
# Define colors and corresponding zones
colors = ['red', 'blue', 'green', 'yellow']
zones = ['zone-red', 'zone-blue', 'zone-green', 'zone-yellow']
# Add zones at the corners of a square
zone_size = (0.05, 0.05, 0.005)
zone_urdf = 'zone/zone.urdf'
zone_poses = []
for i in range(4):
zone_pose = self.get_random_pose(env, zone_size)
env.add_object(zone_urdf, zone_pose, 'fixed')
zone_poses.append(zone_pose)
# Add cylinders of different colors
cylinder_size = (0.02, 0.02, 0.08)
cylinder_urdf = 'cylinder/cylinder-template.urdf'
cylinders = []
for i in range(4):
cylinder_pose = self.get_random_pose(env, cylinder_size)
color = utils.COLORS[colors[i]]
cylinder_id = env.add_object(cylinder_urdf, cylinder_pose, color=color)
cylinders.append(cylinder_id)
# Goal: each cylinder is in a different zone of the same color
for i in range(4):
self.add_goal(objs=[cylinders[i]], matches=np.ones((1, 1)), targ_poses=[zone_poses[i]], replace=False,
rotations=True, metric='pose', params=None, step_max_reward=1/4)
self.lang_goals.append(self.lang_template.format(color=colors[i])) |