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 BallTowerInZones(Task): | |
| """Pick up each sphere and stack it on top of the stand of the same color while maintaining the balance. | |
| However, the stacking should be done in a specific color sequence - blue at the bottom, followed by green, | |
| and finally red at the top. The spheres and stands are each placed within separate zones.""" | |
| def __init__(self): | |
| super().__init__() | |
| self.max_steps = 10 | |
| self.lang_template = "stack the {color} sphere on the {color} stand" | |
| self.task_completed_desc = "done stacking spheres." | |
| self.additional_reset() | |
| def reset(self, env): | |
| super().reset(env) | |
| # Add zones. | |
| zone_size = (0.12, 0.12, 0) | |
| zone_urdf = 'zone/zone.urdf' | |
| zone_poses = [] | |
| for _ in range(3): | |
| zone_pose = self.get_random_pose(env, zone_size) | |
| env.add_object(zone_urdf, zone_pose, 'fixed') | |
| zone_poses.append(zone_pose) | |
| # Sphere and stand colors. | |
| colors = [ | |
| utils.COLORS['blue'], utils.COLORS['green'], utils.COLORS['red'] | |
| ] | |
| # Add spheres and stands. | |
| sphere_size = (0.04, 0.04, 0.04) | |
| sphere_urdf = 'sphere/sphere.urdf' | |
| stand_size = (0.05, 0.05, 0.05) | |
| stand_urdf = 'stacking/stand.urdf' | |
| spheres = [] | |
| stands = [] | |
| for i in range(3): | |
| sphere_pose = self.get_random_pose(env, sphere_size) | |
| sphere_id = env.add_object(sphere_urdf, sphere_pose, color=colors[i]) | |
| spheres.append(sphere_id) | |
| stand_pose = self.get_random_pose(env, stand_size) | |
| stand_id = env.add_object(stand_urdf, stand_pose, color=colors[i]) | |
| stands.append(stand_id) | |
| # Goal: each sphere is stacked on the stand of the same color in the correct zone. | |
| for i in range(3): | |
| self.add_goal(objs=[spheres[i]], matches=np.ones((1, 1)), targ_poses=[zone_poses[i]], replace=False, | |
| rotations=True, metric='pose', params=None, step_max_reward=1 / 3) | |
| self.lang_goals.append(self.lang_template.format(color=colors[i])) | |
| # Add distractors. | |
| n_distractors = 0 | |
| while n_distractors < 6: | |
| is_sphere = np.random.rand() > 0.5 | |
| urdf = sphere_urdf if is_sphere else stand_urdf | |
| size = sphere_size if is_sphere else stand_size | |
| pose = self.get_random_pose(env, obj_size=size) | |
| color = colors[n_distractors % len(colors)] | |
| obj_id = env.add_object(urdf, pose, color=color) | |
| n_distractors += 1 |