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])]