Spaces:
Runtime error
Runtime error
File size: 2,481 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 58 59 60 61 62 63 64 65 66 67 |
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
import pybullet as p
class AlignCylindersInZones(Task):
"""Place four differently colored cylinders each into a matching colored zone.
The zones are surrounded by small blocks, which the robot needs to move out of the way
without knocking them out of their respective zones."""
def __init__(self):
super().__init__()
self.max_steps = 20
self.lang_template = "place the {color} cylinder in the {color} zone"
self.task_completed_desc = "done aligning cylinders."
self.additional_reset()
def reset(self, env):
super().reset(env)
# Cylinder colors.
colors = ['red', 'blue', 'green', 'yellow']
# Add cylinders.
# x, y, z dimensions for the asset size
cylinder_size = (0.04, 0.04, 0.04)
cylinder_urdf = 'cylinder/cylinder-template.urdf'
cylinders = []
for i in range(4):
cylinder_pose = self.get_random_pose(env, cylinder_size)
cylinder_id = env.add_object(cylinder_urdf, cylinder_pose, color=utils.COLORS[colors[i]])
cylinders.append(cylinder_id)
# Add zones.
# x, y, z dimensions for the asset size
zone_size = (0.12, 0.12, 0)
zone_urdf = 'zone/zone.urdf'
zones = []
for i in range(4):
zone_pose = self.get_random_pose(env, zone_size)
env.add_object(zone_urdf, zone_pose, color=utils.COLORS[colors[i]], category='fixed')
zones.append(zone_pose)
# Add small blocks around the zones.
# x, y, z dimensions for the asset size
block_size = (0.02, 0.02, 0.02)
block_urdf = 'block/small.urdf'
for _ in range(16):
block_pose = self.get_random_pose(env, block_size)
env.add_object(block_urdf, block_pose)
# Goal: each cylinder is in a matching colored zone.
for i in range(4):
self.add_goal(objs=[cylinders[i]], matches=np.ones((1, 1)), targ_poses=[zones[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])) |