Spaces:
Runtime error
Runtime error
| 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])) |