def init_task(self) -> None: self.square_ring = Shape('hannoi_square_ring') self.success_centre = Dummy('success_centre') success_detectors = [ ProximitySensor('success_detector%d' % i) for i in range(4) ] self.register_graspable_objects([self.square_ring]) success_condition = ConditionSet([ DetectedCondition(self.square_ring, sd) for sd in success_detectors ]) self.register_success_conditions([success_condition])
class HannoiSquare(Task): def init_task(self) -> None: self.square_ring = Shape('hannoi_square_ring') self.success_centre = Dummy('success_centre') success_detectors = [ ProximitySensor('success_detector%d' % i) for i in range(4) ] self.register_graspable_objects([self.square_ring]) success_condition = ConditionSet([ DetectedCondition(self.square_ring, sd) for sd in success_detectors ]) self.register_success_conditions([success_condition]) def init_episode(self, index: int) -> List[str]: pillar0 = Shape('hannoi_square_pillar0') pillar1 = Shape('hannoi_square_pillar1') pillar2 = Shape('hannoi_square_pillar2') spokes = [pillar0, pillar1, pillar2] color_name, color_rgb = colors[index] chosen_pillar = np.random.choice(spokes) chosen_pillar.set_color(color_rgb) _, _, z = self.success_centre.get_position() x, y, _ = chosen_pillar.get_position() self.success_centre.set_position([x, y, z]) color_choices = np.random.choice(list(range(index)) + list(range(index + 1, len(colors))), size=2, replace=False) spokes.remove(chosen_pillar) for spoke, i in zip(spokes, color_choices): name, rgb = colors[i] spoke.set_color(rgb) boundary1 = Shape('hannoi_square_boundary0') square_ring = Shape('hannoi_square_ring') b = SpawnBoundary([boundary1]) b.sample(square_ring) return [ 'put the ring on the %s spoke' % color_name, 'slide the ring onto the %s colored spoke' % color_name, 'place the ring onto the %s spoke' % color_name ] def variation_count(self) -> int: return len(colors)
def __init__(self, pyrep: PyRep, robot: Robot, scene: Scene, task: Task, action_mode: ActionMode, dataset_root: str, obs_config: ObservationConfig, static_positions: bool = False, attach_grasped_objects: bool = True): self._pyrep = pyrep self._robot = robot self._scene = scene self._task = task self._variation_number = 0 self._action_mode = action_mode self._dataset_root = dataset_root self._obs_config = obs_config self._static_positions = static_positions self._attach_grasped_objects = attach_grasped_objects self._reset_called = False self._prev_ee_velocity = None self._enable_path_observations = False self._scene.load(self._task) self._pyrep.start() self._target_workspace_check = Dummy.create() self._last_e = None
def test_action_mode_ee_position_plan_ee_frame(self): task = self.get_task( ReachTarget, ArmActionMode.EE_POSE_PLAN_EE_FRAME) _, obs = task.reset() dummy = Dummy.create() dummy.set_position([0, 0, 0.05], relative_to=task._robot.arm.get_tip()) action = [0, 0, 0.05, 0, 0, 0, 1, 1] obs, reward, term = task.step(action) [self.assertAlmostEqual(a, p, delta=0.01) for a, p in zip(dummy.get_position(), obs.gripper_pose[:3])]