class PutKnifeOnChoppingBoard(Task):
def init_task(self) -> None:
knife = Shape('knife')
success = ProximitySensor('success')
self.knife_block = Shape('knife_block')
self.chopping_board = Shape('chopping_board')
self.register_graspable_objects([knife])
cond_set = ConditionSet([DetectedCondition(knife, success),
NothingGrasped(self.robot.gripper)],
order_matters=True)
self.register_success_conditions([cond_set])
self.boundary = SpawnBoundary([Shape('boundary')])
def init_episode(self, index: int) -> List[str]:
self.boundary.clear()
self.boundary.sample(self.knife_block)
return ['put the knife on the chopping board',
'slide the knife out of the knife block and put it down on the '
'chopping board',
'place the knife on the chopping board',
'pick up the knife and leave it on the chopping board',
'move the knife from the holder to the chopping board']
def variation_count(self) -> int:
return 1
class TvOn(Task):
def init_task(self) -> None:
self._remote = Shape('tv_remote')
self._spawn_boundary = SpawnBoundary([Shape('spawn_boundary')])
self.register_graspable_objects([self._remote])
self.register_success_conditions(
[JointCondition(Joint('target_button_joint0'), 0.003)])
def init_episode(self, index: int) -> List[str]:
self._spawn_boundary.clear()
self._spawn_boundary.sample(self._remote)
return [
'turn on the TV',
'point the remote control at the television and turn on the '
'television',
'pick up the remote and rotate it such that the front of the '
'remote is pointed straight at the television, then set the '
'remote down and press the power button down in order to switch'
' on the TV',
'find the power button at the top of the remote, ensure the '
'remote is pointed at the tv, then turn the tv on'
]
def variation_count(self) -> int:
return 1
def base_rotation_bounds(
self
) -> Tuple[Tuple[float, float, float], Tuple[float, float, float]]:
return (0.0, 0.0, -0.5 * np.pi), (0.0, 0.0, +0.5 * np.pi)
class PlaceHangerOnRack(Task):
def init_task(self) -> None:
self.hanger_holder = Shape('hanger_holder')
hanger = Shape('clothes_hanger')
success_detector = ProximitySensor('success_detector0')
self.register_success_conditions([
NothingGrasped(self.robot.gripper),
DetectedCondition(hanger, success_detector)
])
self.register_graspable_objects([hanger])
self.workspace_boundary = SpawnBoundary([Shape('workspace')])
def init_episode(self, index: int) -> List[str]:
self.workspace_boundary.clear()
self.workspace_boundary.sample(self.hanger_holder)
return [
'pick up the hanger and place in on the rack'
'put the hanger on the rack', 'hang the hanger up',
'grasp the hanger and hang it up'
]
def variation_count(self) -> int:
return 1
def is_static_workspace(self) -> bool:
return True
class PutToiletRollOnStand(Task):
def init_task(self) -> None:
roll = Shape('toilet_roll')
self.boundary_roll = SpawnBoundary([Shape('boundary_roll')])
self.boundary_stand = SpawnBoundary([Shape('boundary_stand')])
self.roll_box = Shape('toilet_roll_box')
self.stand_base = Shape('stand_base')
success_sensor = ProximitySensor('success')
self.register_graspable_objects([roll])
self.register_success_conditions(
[DetectedCondition(roll, success_sensor),
NothingGrasped(self.robot.gripper)])
def init_episode(self, index: int) -> List[str]:
self.boundary_roll.clear()
self.boundary_stand.clear()
self.boundary_roll.sample(self.roll_box)
self.boundary_stand.sample(self.stand_base)
return ['put toilet roll on stand',
'place the toilet roll on the stand',
'slide the roll onto its holder']
def variation_count(self) -> int:
return 1
class PutTrayInOven(Task):
def init_task(self) -> None:
success_detector = ProximitySensor('success')
self.tray = Shape('tray')
self.register_graspable_objects([self.tray])
self.register_success_conditions(
[DetectedCondition(self.tray, success_detector),
NothingGrasped(self.robot.gripper)])
self.boundary = SpawnBoundary([Shape('oven_tray_boundary')])
def init_episode(self, index: int) -> List[str]:
self.boundary.clear()
self.boundary.sample(
self.tray, min_rotation=(0, 0, 0), max_rotation=(0, 0, 0))
return ['put tray in oven',
'place the tray in the oven',
'open the oven, then slide the tray in',
'open the oven door, pick up the tray, and put it down on the '
'oven shelf']
def variation_count(self) -> int:
return 1
def base_rotation_bounds(self) -> Tuple[List[float], List[float]]:
return [0, 0, -3.14 / 4.], [0, 0, 3.14 / 4.]
def boundary_root(self) -> Object:
return Shape('oven_boundary_root')
class PutPlateInColoredDishRack(Task):
def init_task(self) -> None:
plate = Shape('plate')
self.dish_rack = Shape('dish_rack')
self.plate_stand = Shape('plate_stand')
self.success_sensor = ProximitySensor('success')
self.success_poses = [Dummy('success_pos%d' % i) for i in range(3)]
self.pillars = [Shape('dish_rack_pillar%d' % i) for i in range(6)]
self.boundary = SpawnBoundary([Shape('boundary')])
self.register_graspable_objects([plate])
cond_set = ConditionSet([
DetectedCondition(plate, self.success_sensor),
NothingGrasped(self.robot.gripper)
],
order_matters=True)
self.register_success_conditions([cond_set])
def init_episode(self, index: int) -> List[str]:
color_name, _ = colors[index]
shuffled_pillar_indexes = list(range(3))
np.random.shuffle(shuffled_pillar_indexes)
# Color the other spokes
color_choices = [index] + list(
np.random.choice(
list(range(index)) + list(range(index + 1, len(colors))),
size=2,
replace=False))
for pillar_i, color_i in zip(shuffled_pillar_indexes, color_choices):
_, rgb = colors[color_i]
self.pillars[pillar_i * 2].set_color(rgb)
self.pillars[1 + pillar_i * 2].set_color(rgb)
# Move the success detector to the correct spot
success_pos = self.success_poses[shuffled_pillar_indexes[0]]
self.success_sensor.set_position(success_pos.get_position())
self.success_sensor.set_orientation(success_pos.get_orientation())
self.boundary.clear()
self.boundary.sample(self.plate_stand,
min_rotation=(0, 0, 0.5),
max_rotation=(0, 0, 0.5))
self.boundary.sample(self.dish_rack, min_distance=0.2)
return [
'put the plate between the %s pillars of the dish rack' #
% color_name,
'place the plate in the the %s section of the drying rack' %
color_name,
'pick up the plate and leave it between the %s spokes on the '
'dish rack' % color_name
]
def variation_count(self) -> int:
return len(colors)
def base_rotation_bounds(self) -> Tuple[List[float], List[float]]:
return [0, 0, 0], [0, 0, 0]
class PutGroceriesInCupboard(Task):
def init_task(self) -> None:
self.groceries = [Shape(name.replace(' ', '_'))
for name in GROCERY_NAMES]
self.grasp_points = [Dummy('%s_grasp_point' % name.replace(' ', '_'))
for name in GROCERY_NAMES]
self.waypoint1 = Dummy('waypoint1')
self.register_graspable_objects(self.groceries)
self.boundary = SpawnBoundary([Shape('workspace')])
def init_episode(self, index: int) -> List[str]:
self.boundary.clear()
[self.boundary.sample(g, min_distance=0.1) for g in self.groceries]
self.waypoint1.set_pose(self.grasp_points[index].get_pose())
self.register_success_conditions(
[DetectedCondition(self.groceries[index],
ProximitySensor('success')),
NothingGrasped(self.robot.gripper)])
return ['put the %s in the cupboard' % GROCERY_NAMES[index],
'pick up the %s and place it in the cupboard'
% GROCERY_NAMES[index],
'move the %s to the bottom shelf' % GROCERY_NAMES[index],
'put away the %s in the cupboard' % GROCERY_NAMES[index]]
def variation_count(self) -> int:
return len(GROCERY_NAMES)
def boundary_root(self) -> Object:
return Shape('boundary_root')
def base_rotation_bounds(self) -> Tuple[Tuple[float, float, float],
Tuple[float, float, float]]:
return (0.0, 0.0, -1.), (0.0, 0.0, 1.)
class PlugChargerInPowerSupply(Task):
def init_task(self) -> None:
sensor_right_hole = ProximitySensor('sensor1')
sensor_left_hole = ProximitySensor('sensor2')
charger = Shape('charger')
self.charger_base = Dummy('charger_base')
charger_right_stick = Shape('detector1')
charger_left_stick = Shape('detector2')
self.boundary = SpawnBoundary([Shape('charger_boundary')])
self.register_graspable_objects([charger])
self.register_success_conditions([
DetectedCondition(charger_right_stick, sensor_right_hole),
DetectedCondition(charger_left_stick, sensor_left_hole),
NothingGrasped(self.robot.gripper)
])
def init_episode(self, index: int) -> List[str]:
self.boundary.clear()
self.boundary.sample(self.charger_base)
return [
'plug charger in power supply',
'pick up the charger and plug in it',
'plug the charger into the mains',
'lift the charger up to the wall and plug it in',
'plug the charger into the wall'
]
def variation_count(self) -> int:
return 1
def base_rotation_bounds(self) -> Tuple[List[float], List[float]]:
return [0, 0, -3.14 / 4.], [0, 0, 3.14 / 4.]
class WipeDesk(Task):
def init_task(self) -> None:
self.dirt_spots = []
self.sponge = Shape('sponge')
self.sensor = ProximitySensor('sponge_sensor')
self.register_graspable_objects([self.sponge])
boundaries = [Shape('dirt_boundary')]
_, _, self.z_boundary = boundaries[0].get_position()
self.b = SpawnBoundary(boundaries)
def init_episode(self, index: int) -> List[str]:
self._place_dirt()
self.register_success_conditions([EmptyCondition(self.dirt_spots)])
return [
'wipe dirt off the desk', 'use the sponge to clean up the desk',
'remove the dirt from the desk',
'grip the sponge and wipe it back and forth over any dirt you '
'see', 'clean up the mess', 'wipe the dirt up'
]
def variation_count(self) -> int:
return 1
def step(self) -> None:
for d in self.dirt_spots:
if self.sensor.is_detected(d):
self.dirt_spots.remove(d)
d.remove()
def cleanup(self) -> None:
for d in self.dirt_spots:
d.remove()
self.dirt_spots = []
def _place_dirt(self):
for i in range(DIRT_POINTS):
spot = Shape.create(type=PrimitiveShape.CUBOID,
size=[.005, .005, .001],
mass=0,
static=True,
respondable=False,
renderable=True,
color=[0.58, 0.29, 0.0])
spot.set_parent(self.get_base())
spot.set_position([-1, -1, self.z_boundary + 0.001])
self.b.sample(spot,
min_distance=0.00,
min_rotation=(0.00, 0.00, 0.00),
max_rotation=(0.00, 0.00, 0.00))
self.dirt_spots.append(spot)
self.b.clear()
class BlockPyramid(Task):
def init_task(self) -> None:
self.blocks = [Shape('block_pyramid_block%d' % i) for i in range(6)]
self.distractors = [
Shape('block_pyramid_distractor_block%d' % i) for i in range(6)
]
success_detectors = [
ProximitySensor('block_pyramid_success_block%d' % i)
for i in range(3)
]
cond_set = ConditionSet([
DetectedSeveralCondition(self.blocks, success_detectors[0], 3),
DetectedSeveralCondition(self.blocks, success_detectors[1], 2),
DetectedSeveralCondition(self.blocks, success_detectors[2], 1),
NothingGrasped(self.robot.gripper)
])
self.register_success_conditions([cond_set])
self.register_graspable_objects(self.blocks + self.distractors)
self.spawn_boundary = SpawnBoundary(
[Shape('block_pyramid_boundary%d' % i) for i in range(4)])
def init_episode(self, index: int) -> List[str]:
color_name, color_rgb = colors[index]
for obj in self.blocks:
obj.set_color(color_rgb)
color_choice = np.random.choice(list(range(index)) +
list(range(index + 1, len(colors))),
size=1,
replace=False)[0]
name, rgb = colors[color_choice]
for obj in self.distractors:
obj.set_color(rgb)
self.spawn_boundary.clear()
for ob in self.blocks + self.distractors:
self.spawn_boundary.sample(ob, min_distance=0.08)
return [
'stack %s blocks in a pyramid' % color_name,
'create a pyramid with the %s objects' % color_name,
'make a pyramid out of %s cubes' % color_name,
'position the %s blocks in the shape of a pyramid' % color_name,
'use the %s blocks to build a pyramid' % color_name
]
def variation_count(self) -> int:
return len(colors)
class ChangeChannel(Task):
def init_task(self) -> None:
self._remote = Shape('tv_remote')
self._joint_conditions = [
JointCondition(Joint('target_button_joint1'), 0.003),
JointCondition(Joint('target_button_joint2'), 0.003)
]
self._w6 = Dummy('waypoint6')
self._w6z = self._w6.get_position()[2]
self.register_graspable_objects([self._remote])
self._spawn_boundary = SpawnBoundary([Shape('spawn_boundary')])
self._target_buttons = [
Shape('target_button_wrap1'), Shape('target_button_wrap2')]
def init_episode(self, index: int) -> List[str]:
detectors = [ProximitySensor('success%d' % i) for i in range(2)]
self.register_success_conditions([
self._joint_conditions[index],
DetectedCondition(self._remote, detectors[0]),
DetectedCondition(self._remote, detectors[1])
])
x, y, _ = self._target_buttons[index % 2].get_position()
self._w6.set_position([x, y, self._w6z])
self._spawn_boundary.clear()
self._spawn_boundary.sample(self._remote)
btn = ['plus', 'minus']
chnl = ['up', 'minus']
return ['turn the channel %s' % chnl[index],
'change the television channel %s' % chnl[index],
'point the remote at the tv and press the %s button to turn '
'the channel %s' % (btn[index], chnl[index]),
'using the tv remote, ensure it is facing the television and '
'press the %s button to increment the channel %s by one'
% (btn[index], chnl[index]),
'find the %s button on the remote, rotate the remote such that'
' it is pointed at the tv, then press the button to change '
'the channel %s' % (chnl[index], btn[index])]
def variation_count(self) -> int:
return 2
def base_rotation_bounds(self) -> Tuple[Tuple[float, float, float],
Tuple[float, float, float]]:
return (0.0, 0.0, -0.5*np.pi), (0.0, 0.0, +0.5 * np.pi)
class PlaceShapeInShapeSorter(Task):
def init_task(self) -> None:
self.shape_sorter = Shape('shape_sorter')
self.success_sensor = ProximitySensor('success')
self.shapes = [Shape(ob.replace(' ', '_')) for ob in SHAPE_NAMES]
self.drop_points = [
Dummy('%s_drop_point' % ob.replace(' ', '_'))
for ob in SHAPE_NAMES]
self.grasp_points = [
Dummy('%s_grasp_point' % ob.replace(' ', '_'))
for ob in SHAPE_NAMES]
self.waypoint1 = Dummy('waypoint1')
self.waypoint4 = Dummy('waypoint4')
self.register_graspable_objects(self.shapes)
self.register_waypoint_ability_start(0, self._set_grasp)
self.register_waypoint_ability_start(3, self._set_drop)
self.boundary = SpawnBoundary([Shape('boundary')])
def init_episode(self, index) -> List[str]:
self.variation_index = index
shape = SHAPE_NAMES[index]
self.register_success_conditions(
[DetectedCondition(self.shapes[index], self.success_sensor)])
self.boundary.clear()
[self.boundary.sample(s, min_distance=0.05) for s in self.shapes]
return ['put the %s in the shape sorter' % shape,
'pick up the %s and put it in the sorter' % shape,
'place the %s into its slot in the shape sorter' % shape,
'slot the %s into the shape sorter' % shape]
def variation_count(self) -> int:
return len(SHAPE_NAMES)
def _set_grasp(self, _):
gp = self.grasp_points[self.variation_index]
self.waypoint1.set_pose(gp.get_pose())
def _set_drop(self, _):
dp = self.drop_points[self.variation_index]
self.waypoint4.set_pose(dp.get_pose())
class PickUpCup(Task):
def init_task(self) -> None:
self.cup1 = Shape('cup1')
self.cup2 = Shape('cup2')
self.cup1_visual = Shape('cup1_visual')
self.cup2_visual = Shape('cup2_visual')
self.boundary = SpawnBoundary([Shape('boundary')])
self.success_sensor = ProximitySensor('success')
self.register_graspable_objects([self.cup1, self.cup2])
self.register_success_conditions([
DetectedCondition(self.cup1, self.success_sensor, negated=True),
GraspedCondition(self.robot.gripper, self.cup1),
])
def init_episode(self, index: int) -> List[str]:
self.variation_index = index
target_color_name, target_rgb = colors[index]
random_idx = np.random.choice(len(colors))
while random_idx == index:
random_idx = np.random.choice(len(colors))
_, other1_rgb = colors[random_idx]
self.cup1_visual.set_color(target_rgb)
self.cup2_visual.set_color(other1_rgb)
self.boundary.clear()
self.boundary.sample(self.cup2, min_distance=0.1)
self.boundary.sample(self.success_sensor, min_distance=0.1)
return [
'pick up the %s cup' % target_color_name,
'grasp the %s cup and lift it' % target_color_name,
'lift the %s cup' % target_color_name
]
def variation_count(self) -> int:
return len(colors)
class PutKnifeInKnifeBlock(Task):
def init_task(self) -> None:
knife = Shape('knife')
self.knife_base = Dummy('knife_base')
success = ProximitySensor('success')
self.knife_block = Shape('knife_block')
self.chopping_board = Shape('chopping_board')
self.register_graspable_objects([knife])
set = ConditionSet([
DetectedCondition(knife, success),
NothingGrasped(self.robot.gripper)
],
order_matters=True)
self.register_success_conditions([set])
self.boundary = SpawnBoundary([Shape('boundary')])
self.knife_bound = SpawnBoundary([Shape('knife_boundary')])
def init_episode(self, index: int) -> List[str]:
self.knife_bound.clear()
self.boundary.clear()
self.boundary.sample(self.knife_block)
self.boundary.sample(self.chopping_board)
self.knife_bound.sample(self.knife_base)
return [
'put the knife in the knife block',
'slide the knife into its slot in the knife block',
'place the knife in the knife block',
'pick up the knife and leave it in its holder',
'move the knife from the chopping board to the holder'
]
def variation_count(self) -> int:
return 1
def base_rotation_bounds(self) -> Tuple[List[float], List[float]]:
return [0, 0, 0], [0, 0, 0]
class EmptyContainer(Task):
def init_task(self) -> None:
self.large_container = Shape('large_container')
self.target_container0 = Shape('small_container0')
self.target_container1 = Shape('small_container1')
self.success_detector0 = ProximitySensor('success0')
self.success_detector1 = ProximitySensor('success1')
self.target_waypoint = Dummy('waypoint3')
self.spawn_boundary = SpawnBoundary([Shape('spawn_boundary')])
self.register_waypoint_ability_start(1, self._move_above_object)
self.register_waypoints_should_repeat(self._repeat)
self.bin_objects = []
def init_episode(self, index: int) -> List[str]:
self._variation_index = index
self.bin_objects = sample_procedural_objects(self.get_base(), 3)
self.bin_objects_not_done = list(self.bin_objects)
self.register_graspable_objects(self.bin_objects)
self.spawn_boundary.clear()
for ob in self.bin_objects:
ob.set_position([0.0, 0.0, 0.2],
relative_to=self.large_container,
reset_dynamics=False)
self.spawn_boundary.sample(ob,
ignore_collisions=True,
min_distance=0.05)
target_pos = [-5.9605e-8, -2.5005e-1, +1.7e-1]
conditions = []
target_color_name, target_color_rgb = colors[index]
color_choice = np.random.choice(list(range(index)) +
list(range(index + 1, len(colors))),
size=1,
replace=False)[0]
_, distractor_color_rgb = colors[color_choice]
if index % 2 == 0:
self.target_container0.set_color(target_color_rgb)
self.target_container1.set_color(distractor_color_rgb)
for ob in self.bin_objects:
conditions.append(DetectedCondition(ob,
self.success_detector0))
else:
self.target_container1.set_color(target_color_rgb)
self.target_container0.set_color(distractor_color_rgb)
for ob in self.bin_objects:
conditions.append(DetectedCondition(ob,
self.success_detector1))
target_pos[1] = -target_pos[1]
self.target_waypoint.set_position(target_pos,
relative_to=self.large_container,
reset_dynamics=True)
self.register_success_conditions(
[ConditionSet(conditions, simultaneously_met=True)])
return [
'empty the container in the to %s container' % target_color_name,
'clear all items from the large tray and put them in the %s '
'tray' % target_color_name,
'move all objects from the large container and drop them into '
'the smaller %s one' % target_color_name,
'remove whatever you find in the big box in the middle and '
'leave them in the %s one' % target_color_name,
'grasp and move all objects into the %s container' %
target_color_name
]
def variation_count(self) -> int:
return len(colors)
def cleanup(self) -> None:
[ob.remove() for ob in self.bin_objects if ob.still_exists()]
self.bin_objects = []
def step(self) -> None:
for ob in self.bin_objects_not_done:
if self._variation_index % 2 == 0:
if self.success_detector0.is_detected(ob):
self.bin_objects_not_done.remove(ob)
else:
if self.success_detector1.is_detected(ob):
self.bin_objects_not_done.remove(ob)
def _move_above_object(self, waypoint):
if len(self.bin_objects_not_done) <= 0:
raise RuntimeError('Should not be here.')
bin_obj = self.bin_objects_not_done[0]
way_obj = waypoint.get_waypoint_object()
way_obj.set_position(bin_obj.get_position())
x, y, _ = way_obj.get_orientation()
_, _, z = bin_obj.get_orientation(relative_to=way_obj)
way_obj.set_orientation([x, y, z])
def _repeat(self):
return len(self.bin_objects_not_done) > 0
class StackCups(Task):
def init_task(self) -> None:
success_sensor = ProximitySensor('success')
self.cup1 = Shape('cup1')
self.cup2 = Shape('cup2')
self.cup3 = Shape('cup3')
self.cup1_visual = Shape('cup1_visual')
self.cup2_visual = Shape('cup2_visual')
self.cup3_visaul = Shape('cup3_visual')
self.boundary = SpawnBoundary([Shape('boundary')])
self.register_graspable_objects([self.cup1, self.cup2, self.cup3])
self.register_success_conditions([
DetectedCondition(self.cup1, success_sensor),
DetectedCondition(self.cup3, success_sensor),
NothingGrasped(self.robot.gripper)
])
def init_episode(self, index: int) -> List[str]:
self.variation_index = index
target_color_name, target_rgb = colors[index]
random_idx = np.random.choice(len(colors))
while random_idx == index:
random_idx = np.random.choice(len(colors))
_, other1_rgb = colors[random_idx]
random_idx = np.random.choice(len(colors))
while random_idx == index:
random_idx = np.random.choice(len(colors))
_, other2_rgb = colors[random_idx]
self.cup2_visual.set_color(target_rgb)
self.cup1_visual.set_color(other1_rgb)
self.cup3_visaul.set_color(other2_rgb)
self.boundary.clear()
self.boundary.sample(self.cup2,
min_distance=0.05,
min_rotation=(0, 0, 0),
max_rotation=(0, 0, 0))
self.boundary.sample(self.cup1,
min_distance=0.05,
min_rotation=(0, 0, 0),
max_rotation=(0, 0, 0))
self.boundary.sample(self.cup3,
min_distance=0.05,
min_rotation=(0, 0, 0),
max_rotation=(0, 0, 0))
return [
'stack the other cups on top of the %s cup' % target_color_name,
'place two of the cups onto the odd cup out',
'put the remaining two cups on top of the %s cup' %
target_color_name,
'pick up and set the cups down into the %s cup' %
target_color_name,
'create a stack of cups with the %s cup as its base' %
target_color_name,
'keeping the %s cup on the table, stack the other two onto it' %
target_color_name
]
def variation_count(self) -> int:
return len(colors)
class PutAllGroceriesInCupboard(Task):
def init_task(self) -> None:
self.groceries = [
Shape(name.replace(' ', '_')) for name in GROCERY_NAMES
]
self.grasp_points = [
Dummy('%s_grasp_point' % name.replace(' ', '_'))
for name in GROCERY_NAMES
]
self.goals = [
Dummy('goal_%s' % name.replace(' ', '_')) for name in GROCERY_NAMES
]
self.waypoint1 = Dummy('waypoint1')
self.waypoint3 = Dummy('waypoint3')
self.register_graspable_objects(self.groceries)
self.boundary = SpawnBoundary([Shape('groceries_boundary')])
self.register_waypoint_ability_start(0, self._move_to_next_target)
self.register_waypoint_ability_start(3, self._move_to_drop_zone)
self.register_waypoints_should_repeat(self._repeat)
self.groceries_to_place = len(GROCERY_NAMES)
self.groceries_placed = 0
def init_episode(self, index: int) -> List[str]:
self.groceries_placed = 0
self.boundary.clear()
[self.boundary.sample(g, min_distance=0.15) for g in self.groceries]
self.register_success_conditions([
DetectedSeveralCondition(self.groceries[:self.groceries_to_place],
ProximitySensor('success'),
self.groceries_to_place),
NothingGrasped(self.robot.gripper)
])
return [
'put all of the groceries in the cupboard',
'pick up all of the groceries and place them in the cupboard',
'move the groceries to the shelves',
'put the groceries on the table into the cupboard',
'put away the groceries in the cupboard'
]
def variation_count(self) -> int:
return 1
def boundary_root(self) -> Object:
return Shape('boundary_root')
def is_static_workspace(self) -> bool:
return True
def _move_to_next_target(self, _):
if self.groceries_placed >= self.groceries_to_place:
raise RuntimeError('Should not be here.')
self.waypoint1.set_pose(
self.grasp_points[self.groceries_placed].get_pose())
def _move_to_drop_zone(self, _):
self.waypoint3.set_pose(self.goals[self.groceries_placed].get_pose())
def _repeat(self):
self.groceries_placed += 1
return self.groceries_placed < self.groceries_to_place
class PickAndLift(Task):
def init_task(self) -> None:
self.block = Shape('pick_and_lift_target')
self.register_graspable_objects([self.block])
self.boundary = SpawnBoundary([Shape('pick_and_lift_boundary')])
self.success_detector = ProximitySensor('pick_and_lift_success')
cond_set = ConditionSet([
GraspedCondition(self.robot.gripper, self.block),
DetectedCondition(self.block, self.success_detector)
])
self.register_success_conditions([cond_set])
self.target = Shape('target')
def base_rotation_bounds(self):
return [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]
def init_episode(self, index: int) -> List[str]:
block_color_name, block_rgb = colors[index]
self.block.set_color(block_rgb)
self.boundary.clear()
self.boundary.sample(self.success_detector,
min_rotation=(0.0, 0.0, 0.0),
max_rotation=(0.0, 0.0, 0.0))
if random.random() < self.ground_p:
pose = self.target.get_pose()
pose[2] = 7.75003076e-01
self.target.set_pose(pose)
for block in [self.block]: #+ self.distractors:
self.boundary.sample(block,
min_distance=0.1,
min_rotation=(0, 0, 0),
max_rotation=(0, 0, 0))
def _path_action(pos):
x, y, z, qx, qy, qz, qw = self.robot.arm.get_tip().get_pose()
try:
path = self.robot.arm.get_path(pos[:3],
quaternion=[qx, qy, qz, qw],
ignore_collisions=True)
done = False
observations = []
while not done:
done = path.step()
self.pyrep.step()
except Exception as e:
print(e)
print(
"***!!!! Failed to start pick and lift from demonstration !!!!***"
)
print("pos", pos)
print("current xyz", x, y, z)
if random.random() < self.not_special_p:
joint_pos = [
-1.8593069398775697e-05, -0.09871861338615417,
4.5094158849678934e-05, -2.7951412200927734,
-3.730739263119176e-05, 2.8503623008728027, 0.7854341268539429
]
gripper_pos = [0.04000139236450195, 0.04000053554773331]
self.robot.arm.set_joint_positions(joint_pos)
self.robot.gripper.set_joint_positions(gripper_pos)
self.robot.gripper.actuate(1, 1)
return
if not self.special_is_grip:
_path_action(self.block.get_position() + np.array([0, 0, 0.04]))
else:
_path_action(self.block.get_position())
done = False
while not done:
done = self.robot.gripper.actuate(0, velocity=1)
self.pyrep.step()
self.step()
for g_obj in self.get_graspable_objects():
self.robot.gripper.grasp(g_obj)
return ""
def get_reward(self):
dist = np.linalg.norm(
np.array(self.target.get_position()) -
np.array(self.block.get_position()))
thresh = 0.04
if sparse:
return -(dist > thresh).astype(np.float32), dist < thresh
#epsilon reward
reward = -dist
if dist < thresh:
reward += 0.1
return reward
def variation_count(self) -> int:
return len(colors)
def get_reward_and_done(self, sparse=False):
dist = np.linalg.norm(
np.array(self.target.get_position()) -
np.array(self.block.get_position()))
thresh = 0.04
if sparse:
return -(dist > thresh).astype(np.float32), dist < thresh
reward = -dist
if dist < thresh:
reward += 0.1
return reward, dist < thresh
#For HER
def compute_reward(self, achieved_goal, desired_goal, info):
dist = np.linalg.norm(np.array(desired_goal) - np.array(achieved_goal))
thresh = 0.04
return -(dist > thresh).astype(np.float32)
def compute_reward_and_done(self, achieved_goal, desired_goal, info):
dist = np.linalg.norm(np.array(desired_goal) - np.array(achieved_goal))
thresh = 0.04
return -(dist > thresh).astype(np.float32), dist < thresh
def variation_count(self) -> int:
return 1
def get_desired_goal(self):
return self.target.get_position()
def get_achieved_goal(self):
return self.block.get_position()
def get_state_obs(self):
joint_positions = self.robot.arm.get_joint_positions()
gripper_position = self.robot.arm.get_tip().get_position()
gripper_open = [self.robot.gripper.get_open_amount()[0] > 0.9]
gripper_vel = np.array(
self.robot.arm.get_tip().get_velocity()).flatten()
joint_vel = self.robot.arm.get_joint_velocities()
block = self.block.get_position()
return np.concatenate([
joint_positions, gripper_position, gripper_open, block,
gripper_vel, joint_vel
])
def get_save_state(self):
joint_positions = self.robot.arm.get_joint_positions()
gripper_positions = self.robot.gripper.get_joint_positions()
gripper_open = self.robot.gripper.get_open_amount()[0] > 0.9
block_position = self.block.get_position()
target_position = self.target.get_position()
return [
joint_positions, gripper_positions, gripper_open, block_position,
target_position
]
def restore_save_state(self, state):
joint_positions, gripper_positions, gripper_open, block_position, target_position = state
self.robot.arm.set_joint_positions(joint_positions)
self.robot.gripper.set_joint_positions(gripper_positions)
self.robot.gripper.actuate(gripper_open, 1)
self.block.set_position(block_position)
self.target.set_position(target_position)
class CustomPickAndLiftNoRotation(Task):
def init_task(self) -> None:
self.target_block = Shape('pick_and_lift_target')
self.target = Shape("success_visual")
self.target.set_renderable(False)
self.register_graspable_objects([self.target_block])
self.boundary = SpawnBoundary([Shape('pick_and_lift_boundary')])
self.success_detector = ProximitySensor('pick_and_lift_success')
self.front_camera_exists = Object.exists('cam_front')
if self.front_camera_exists:
self.front_camera = VisionSensor('cam_front')
self.init_front_camera_position = self.front_camera.get_position()
self.init_front_camera_orientation = self.front_camera.get_orientation(
)
self.panda_base = Shape("Panda_link0_visual")
cond_set = ConditionSet([
GraspedCondition(self.robot.gripper, self.target_block),
DetectedCondition(self.target_block, self.success_detector)
])
self.register_success_conditions([cond_set])
def init_episode(self, index: int) -> List[str]:
block_color_name, block_rgb = colors[index]
self.target_block.set_color(block_rgb)
self.boundary.clear()
self.boundary.sample(self.success_detector,
min_rotation=(0.0, 0.0, 0.0),
max_rotation=(0.0, 0.0, 0.0))
for block in [self.target_block]:
self.boundary.sample(block,
min_distance=0.1,
min_rotation=(0, 0, 0),
max_rotation=(0, 0, 0))
if self.front_camera_exists:
# apply new position to front camera for better generalization
self.front_camera_new_position()
return [
'pick up the %s block and lift it up to the target' %
block_color_name,
'grasp the %s block to the target' % block_color_name,
'lift the %s block up to the target' % block_color_name
]
def base_rotation_bounds(self):
# do not allow base rotation
return (0.0, 0.0, 0), (0.0, 0.0, 0)
def variation_count(self) -> int:
return len(colors)
def step(self) -> None:
if self.front_camera_exists:
# apply new position to front camera for better generalization
self.front_camera_new_position()
def get_low_dim_state(self) -> np.ndarray:
# get x, y, z from target block position
block_position = self.target_block.get_position()
# get x, y, z from target position
target_position = self.target.get_position()
return np.concatenate((block_position, target_position))
def reward(self):
""" the custom reward consists of two parts:
1. distance between the gripper and the target_block
2. distance between the target_block and the target (lift_target)
the total reward is the sum of both parts"""
# success conditions
object_grasped = self._success_conditions[0].condition_met()[0]
max_precision = 0.01 # 1cm
max_reward = 1 / max_precision
scale = 0.1
# first part
gripper_position = self.robot.arm.get_tip().get_position()
target_block_position = self.target_block.get_position()
dist = np.sqrt(
np.sum(np.square(
np.subtract(target_block_position, gripper_position)),
axis=0)) # euclidean norm
reward1 = min((1 / (dist + 0.00001)), max_reward)
reward1 = scale * reward1
# second part
if object_grasped:
print("Object Grasped!")
target_position = self.target.get_position()
dist = np.sqrt(
np.sum(np.square(
np.subtract(target_position, target_block_position)),
axis=0)) # euclidean norm
reward2 = min((1 / (dist + 0.00001)), max_reward)
reward2 = scale * reward2 + 5.0
else:
reward2 = 0
return reward1 + reward2
def front_camera_new_position(self):
""" changes the position of the front camera to a new one, and rotates the camera in a way that it sees the
target-block and the Panda-base. """
# reset camera position
self.front_camera.set_position(self.init_front_camera_position)
# --- move front camera to new position ---
# calculate the new (random) position
new_rel_front_camera_pos = np.random.uniform(low=-0.5,
high=0.5,
size=(3, ))
# move camera
self.front_camera.set_position(new_rel_front_camera_pos,
relative_to=self.front_camera)
# --- turn the camera ---
# -> camera should be turned to look at the new position -> mean of Panda_base and target_block
# get the coordinates
tar_block_obj_pos = self.target_block.get_position(
relative_to=self.front_camera)
panda_base_pos = self.panda_base.get_position(
relative_to=self.front_camera)
# calculate the mean to get the position to look at
look_at = np.mean(np.array([tar_block_obj_pos, panda_base_pos]),
axis=0)
# rotation arround y
rot_y = math.atan2(look_at[0], look_at[2])
# rotation arround x
rot_x = math.atan2(
look_at[1],
math.sqrt(math.pow(look_at[0], 2) + math.pow(look_at[2], 2)))
# calculate random rotation around z
rot_z = random.uniform(-math.pi, math.pi)
# apply rotation
old_orientation = self.front_camera.get_orientation(
relative_to=self.front_camera)
self.front_camera.set_orientation(
[old_orientation[0] - rot_x, old_orientation[1] + rot_y, rot_z],
relative_to=self.front_camera)
class EmptyContainer(Task):
def init_task(self) -> None:
self.spawn_boundary = SpawnBoundary([Shape('spawn_boundary')])
self.target = Shape('target')
self.bin_objects = []
self.block, self.block_visual = get_one_procedural_object_by_index(
self.get_base(), self.procedural_ind)
self.block_pose = self.block.get_pose()
self.block_visual.set_color(self.COLOR_TUPLE)
self.register_graspable_objects([self.block])
self.ind = 0
def init_episode(self, index: int) -> List[str]:
if self.procedural_mode == "increase":
if self.procedural_set:
self.block.remove()
self.block, self.block_visual = get_one_procedural_object_by_index(
self.get_base(), self.procedural_set[self.ind])
self.block_pose = self.block.get_pose()
self.block_visual.set_color(self.COLOR_TUPLE)
self.register_graspable_objects([self.block])
self.ind = (self.ind + 1) % len(self.procedural_set)
else:
self.block.remove()
self.block, self.block_visual = get_one_procedural_object_by_index(
self.get_base(), self.procedural_ind)
self.block_pose = self.block.get_pose()
self.block_visual.set_color(self.COLOR_TUPLE)
self.register_graspable_objects([self.block])
self.procedural_ind += 1
elif self.procedural_mode == "random":
assert self.procedural_set
self.block.remove()
self.block, self.block_visual = get_one_procedural_object_by_index(
self.get_base(), random.choice(self.procedural_set))
self.block_pose = self.block.get_pose()
self.block_visual.set_color(self.COLOR_TUPLE)
self.register_graspable_objects([self.block])
self.spawn_boundary.clear()
#Fix the target to spawn on a certain height.
self.spawn_boundary.sample(self.target)
pose = self.target.get_pose()
if random.random() < self.ground_p:
pose[2] = 7.75003076e-01
else:
pose[2] = 8.83000314e-01
self.target.set_pose(pose)
self.block.set_pose(self.block_pose)
self.block.set_position([0, 0, 0.76], reset_dynamics=False)
self.spawn_boundary.sample(self.block,
ignore_collisions=True,
min_distance=0.05,
min_rotation=(0, 0, 0),
max_rotation=(0, 0, 0))
def _path_action(pos):
x, y, z, qx, qy, qz, qw = self.robot.arm.get_tip().get_pose()
try:
path = self.robot.arm.get_path(pos[:3],
quaternion=[qx, qy, qz, qw],
ignore_collisions=True)
done = False
observations = []
while not done:
done = path.step()
self.pyrep.step()
except Exception as e:
print("***!!!! Failed to start from demonstration !!!!***")
if random.random() < self.not_special_p:
joint_pos = [
-1.8593069398775697e-05, -0.09871861338615417,
4.5094158849678934e-05, -2.7951412200927734,
-3.730739263119176e-05, 2.8503623008728027, 0.7854341268539429
]
gripper_pos = [0.04000139236450195, 0.04000053554773331]
self.robot.arm.set_joint_positions(joint_pos)
self.robot.gripper.set_joint_positions(gripper_pos)
self.robot.gripper.actuate(1, 1)
return
if not self.special_is_grip:
_path_action(self.block.get_position() + np.array([0, 0, 0.04]))
else:
block_pos = self.block.get_position()
_path_action(block_pos)
done = False
while not done:
done = self.robot.gripper.actuate(0, velocity=1)
self.pyrep.step()
self.step()
for g_obj in self.get_graspable_objects():
self.robot.gripper.grasp(g_obj)
return ''
def variation_count(self) -> int:
return len(colors)
def get_reward(self):
dist = np.linalg.norm(
np.array(self.target.get_position()) -
np.array(self.block.get_position()))
thresh = 0.04
if sparse:
return -(dist > thresh).astype(np.float32), dist < thresh
#epsilon reward
reward = -dist
if dist < thresh:
reward += 0.1
return reward
def variation_count(self) -> int:
return len(colors)
def get_reward_and_done(self, sparse=False):
dist = np.linalg.norm(
np.array(self.target.get_position()) -
np.array(self.block.get_position()))
thresh = 0.04
if sparse:
return -(dist > thresh).astype(np.float32), dist < thresh
#epsilon reward
reward = -dist
if dist < thresh:
reward += 0.1
return reward, dist < thresh
#For HER
def compute_reward(self, achieved_goal, desired_goal, info):
dist = np.linalg.norm(np.array(desired_goal) - np.array(achieved_goal))
thresh = 0.04
return -(dist > thresh).astype(np.float32)
def compute_reward_and_done(self, achieved_goal, desired_goal, info):
dist = np.linalg.norm(np.array(desired_goal) - np.array(achieved_goal))
thresh = 0.04
return -(dist > thresh).astype(np.float32), dist < thresh
def variation_count(self) -> int:
return 1
def get_desired_goal(self):
return self.target.get_position()
def get_achieved_goal(self):
return self.block.get_position()
# return self.block.get_position()
def get_state_obs(self):
joint_positions = self.robot.arm.get_joint_positions()
gripper_position = self.robot.arm.get_tip().get_position()
gripper_open = [self.robot.gripper.get_open_amount()[0] > 0.9]
gripper_vel = np.array(
self.robot.arm.get_tip().get_velocity()).flatten()
joint_vel = self.robot.arm.get_joint_velocities()
if hasattr(self, "is_mesh_obs") and self.is_mesh_obs:
vertices = np.zeros((600, 3))
v = self.block.get_mesh_data()[0]
vertices[:v.shape[0]] = v
vertices -= self.block.get_position()
block = vertices.flatten()
else:
block = self.block.get_pose()
return np.concatenate([
joint_positions, gripper_position, gripper_open, block,
gripper_vel, joint_vel
])
def get_save_state(self):
joint_positions = self.robot.arm.get_joint_positions()
gripper_positions = self.robot.gripper.get_joint_positions()
gripper_open = self.robot.gripper.get_open_amount()[0] > 0.9
block_pose = self.block.get_pose()
target_position = self.target.get_position()
return [
joint_positions, gripper_positions, gripper_open, block_pose,
target_position
]
def restore_save_state(self, state):
joint_positions, gripper_positions, gripper_open, block_pose, target_position = state
self.robot.arm.set_joint_positions(joint_positions)
self.robot.gripper.set_joint_positions(gripper_positions)
self.robot.gripper.actuate(gripper_open, 1)
self.block.set_pose(block_pose)
self.target.set_position(target_position)
class Scene(object):
"""Controls what is currently in the vrep scene. This is used for making
sure that the tasks are easily reachable. This may be just replaced by
environment. Responsible for moving all the objects. """
def __init__(self,
pyrep: PyRep,
robot: Robot,
obs_config=ObservationConfig()):
self._pyrep = pyrep
self._robot = robot
self._obs_config = obs_config
self._active_task = None
self._inital_task_state = None
self._start_arm_joint_pos = robot.arm.get_joint_positions()
self._starting_gripper_joint_pos = robot.gripper.get_joint_positions()
self._workspace = Shape('workspace')
self._workspace_boundary = SpawnBoundary([self._workspace])
self._cam_over_shoulder_left = VisionSensor('cam_over_shoulder_left')
self._cam_over_shoulder_right = VisionSensor('cam_over_shoulder_right')
self._cam_wrist = VisionSensor('cam_wrist')
self._cam_over_shoulder_left_mask = VisionSensor(
'cam_over_shoulder_left_mask')
self._cam_over_shoulder_right_mask = VisionSensor(
'cam_over_shoulder_right_mask')
self._cam_wrist_mask = VisionSensor('cam_wrist_mask')
self._has_init_task = self._has_init_episode = False
self._variation_index = 0
self._initial_robot_state = (robot.arm.get_configuration_tree(),
robot.gripper.get_configuration_tree())
# Set camera properties from observation config
self._set_camera_properties()
def load(self, task: Task) -> None:
"""Loads the task and positions at the centre of the workspace.
:param task: The task to load in the scene.
"""
task.load() # Load the task in to the scene
# Set at the centre of the workspace
task.get_base().set_position(self._workspace.get_position())
self._inital_task_state = task.get_state()
self._active_task = task
self._initial_task_pose = task.boundary_root().get_orientation()
self._has_init_task = self._has_init_episode = False
self._variation_index = 0
def unload(self) -> None:
"""Clears the scene. i.e. removes all tasks. """
if self._active_task is not None:
self._robot.gripper.release()
if self._has_init_task:
self._active_task.cleanup_()
self._active_task.unload()
self._active_task = None
self._variation_index = 0
def init_task(self) -> None:
self._active_task.init_task()
self._has_init_task = True
self._variation_index = 0
def init_episode(self,
index: int,
randomly_place: bool = True,
max_attempts: int = 5) -> List[str]:
"""Calls the task init_episode and puts randomly in the workspace.
"""
self._variation_index = index
if not self._has_init_task:
self.init_task()
# Try a few times to init and place in the workspace
attempts = 0
descriptions = None
while attempts < max_attempts:
descriptions = self._active_task.init_episode(index)
try:
if (randomly_place
and not self._active_task.is_static_workspace()):
self._place_task()
self._active_task.validate()
break
except (BoundaryError, WaypointError) as e:
self._active_task.cleanup_()
self._active_task.restore_state(self._inital_task_state)
attempts += 1
if attempts >= max_attempts:
raise e
# Let objects come to rest
[self._pyrep.step() for _ in range(STEPS_BEFORE_EPISODE_START)]
self._has_init_episode = True
return descriptions
def reset(self) -> None:
"""Resets the joint angles. """
self._robot.gripper.release()
arm, gripper = self._initial_robot_state
self._pyrep.set_configuration_tree(arm)
self._pyrep.set_configuration_tree(gripper)
self._robot.arm.set_joint_positions(self._start_arm_joint_pos)
self._robot.arm.set_joint_target_velocities(
[0] * len(self._robot.arm.joints))
self._robot.gripper.set_joint_positions(
self._starting_gripper_joint_pos)
self._robot.gripper.set_joint_target_velocities(
[0] * len(self._robot.gripper.joints))
if self._active_task is not None and self._has_init_task:
self._active_task.cleanup_()
self._active_task.restore_state(self._inital_task_state)
[self._pyrep.step_ui() for _ in range(20)]
self._active_task.set_initial_objects_in_scene()
def get_observation(self) -> Observation:
tip = self._robot.arm.get_tip()
joint_forces = None
if self._obs_config.joint_forces:
fs = self._robot.arm.get_joint_forces()
vels = self._robot.arm.get_joint_target_velocities()
joint_forces = self._obs_config.joint_forces_noise.apply(
np.array([-f if v < 0 else f for f, v in zip(fs, vels)]))
ee_forces_flat = None
if self._obs_config.gripper_touch_forces:
ee_forces = self._robot.gripper.get_touch_sensor_forces()
ee_forces_flat = []
for eef in ee_forces:
ee_forces_flat.extend(eef)
ee_forces_flat = np.array(ee_forces_flat)
lsc_ob = self._obs_config.left_shoulder_camera
rsc_ob = self._obs_config.right_shoulder_camera
wc_ob = self._obs_config.wrist_camera
obs = Observation(
left_shoulder_rgb=(lsc_ob.rgb_noise.apply(
self._cam_over_shoulder_left.capture_rgb())
if lsc_ob.rgb else None),
left_shoulder_depth=(lsc_ob.depth_noise.apply(
self._cam_over_shoulder_left.capture_depth())
if lsc_ob.depth else None),
right_shoulder_rgb=(rsc_ob.rgb_noise.apply(
self._cam_over_shoulder_right.capture_rgb())
if rsc_ob.rgb else None),
right_shoulder_depth=(rsc_ob.depth_noise.apply(
self._cam_over_shoulder_right.capture_depth())
if rsc_ob.depth else None),
wrist_rgb=(wc_ob.rgb_noise.apply(self._cam_wrist.capture_rgb())
if wc_ob.rgb else None),
wrist_depth=(wc_ob.depth_noise.apply(
self._cam_wrist.capture_depth()) if wc_ob.depth else None),
left_shoulder_mask=(
self._cam_over_shoulder_left_mask.capture_rgb()
if lsc_ob.mask else None),
right_shoulder_mask=(
self._cam_over_shoulder_right_mask.capture_rgb()
if rsc_ob.mask else None),
wrist_mask=(self._cam_wrist_mask.capture_rgb()
if wc_ob.mask else None),
joint_velocities=(self._obs_config.joint_velocities_noise.apply(
np.array(self._robot.arm.get_joint_velocities()))
if self._obs_config.joint_velocities else None),
joint_positions=(self._obs_config.joint_positions_noise.apply(
np.array(self._robot.arm.get_joint_positions()))
if self._obs_config.joint_positions else None),
joint_forces=(joint_forces
if self._obs_config.joint_forces else None),
gripper_open_amount=((
1.0 if self._robot.gripper.get_open_amount()[0] > 0.9 else 0.0)
if self._obs_config.gripper_open_amount else
None),
gripper_pose=(np.array(tip.get_pose())
if self._obs_config.gripper_pose else None),
gripper_touch_forces=(ee_forces_flat
if self._obs_config.gripper_touch_forces else
None),
gripper_joint_positions=(
np.array(self._robot.gripper.get_joint_positions())
if self._obs_config.gripper_joint_positions else None),
task_low_dim_state=(self._active_task.get_low_dim_state() if
self._obs_config.task_low_dim_state else None))
obs = self._active_task.decorate_observation(obs)
return obs
def step(self):
self._pyrep.step()
self._active_task.step()
def get_demo(self,
record: bool = True,
func=None,
randomly_place: bool = True) -> List[Observation]:
"""Returns a demo (list of observations)"""
if not self._has_init_task:
self.init_task()
if not self._has_init_episode:
self.init_episode(self._variation_index,
randomly_place=randomly_place)
self._has_init_episode = False
waypoints = self._active_task.get_waypoints()
if len(waypoints) == 0:
raise NoWaypointsError('No waypoints were found.',
self._active_task)
demo = []
if record:
self._pyrep.step() # Need this here or get_force doesn't work...
demo.append(self.get_observation())
while True:
success = False
for i, point in enumerate(waypoints):
point.start_of_path()
try:
path = point.get_path()
except ConfigurationPathError as e:
raise DemoError(
'Could not get a path for waypoint %d.' % i,
self._active_task) from e
ext = point.get_ext()
path.visualize()
done = False
success = False
while not done:
done = path.step()
self.step()
self._demo_record_step(demo, record, func)
success, term = self._active_task.success()
if success:
break
point.end_of_path()
path.clear_visualization()
if success:
# We can quit early because we have finished the task
break
# TODO: need to decide how I do the gripper actions
if len(ext) > 0:
contains_param = False
start_of_bracket = -1
gripper = self._robot.gripper
if 'open_gripper(' in ext:
gripper.release()
start_of_bracket = ext.index('open_gripper(') + 13
contains_param = ext[start_of_bracket] != ')'
if not contains_param:
done = False
while not done:
done = gripper.actuate(1.0, 0.04)
self._pyrep.step()
self._active_task.step()
if self._obs_config.record_gripper_closing:
self._demo_record_step(demo, record, func)
elif 'close_gripper(' in ext:
start_of_bracket = ext.index('close_gripper(') + 14
contains_param = ext[start_of_bracket] != ')'
if not contains_param:
done = False
while not done:
done = gripper.actuate(0.0, 0.04)
self._pyrep.step()
self._active_task.step()
if self._obs_config.record_gripper_closing:
self._demo_record_step(demo, record, func)
if contains_param:
rest = ext[start_of_bracket:]
num = float(rest[:rest.index(')')])
done = False
while not done:
done = gripper.actuate(num, 0.04)
self._pyrep.step()
self._active_task.step()
if self._obs_config.record_gripper_closing:
self._demo_record_step(demo, record, func)
if 'close_gripper(' in ext:
for g_obj in self._active_task.get_graspable_objects():
gripper.grasp(g_obj)
self._demo_record_step(demo, record, func)
if not self._active_task.should_repeat_waypoints() or success:
break
# Some tasks may need additional physics steps
# (e.g. ball rowling to goal)
if not success:
for _ in range(10):
self._pyrep.step()
self._active_task.step()
self._demo_record_step(demo, record, func)
success, term = self._active_task.success()
if success:
break
success, term = self._active_task.success()
if not success:
raise DemoError('Demo was completed, but was not successful.',
self._active_task)
return demo
def get_observation_config(self) -> ObservationConfig:
return self._obs_config
def _demo_record_step(self, demo_list, record, func):
if record:
demo_list.append(self.get_observation())
if func is not None:
func()
def _set_camera_properties(self) -> None:
def _set_rgb_props(rgb_cam: VisionSensor, rgb: bool, depth: bool,
conf: CameraConfig):
if not (rgb or depth):
rgb_cam.remove()
else:
rgb_cam.set_resolution(conf.image_size)
rgb_cam.set_render_mode(conf.render_mode)
def _set_mask_props(mask_cam: VisionSensor, mask: bool,
conf: CameraConfig):
if not mask:
mask_cam.remove()
else:
mask_cam.set_resolution(conf.image_size)
_set_rgb_props(self._cam_over_shoulder_left,
self._obs_config.left_shoulder_camera.rgb,
self._obs_config.left_shoulder_camera.depth,
self._obs_config.left_shoulder_camera)
_set_rgb_props(self._cam_over_shoulder_right,
self._obs_config.right_shoulder_camera.rgb,
self._obs_config.right_shoulder_camera.depth,
self._obs_config.right_shoulder_camera)
_set_rgb_props(self._cam_wrist, self._obs_config.wrist_camera.rgb,
self._obs_config.wrist_camera.depth,
self._obs_config.wrist_camera)
_set_mask_props(self._cam_over_shoulder_left_mask,
self._obs_config.left_shoulder_camera.mask,
self._obs_config.left_shoulder_camera)
_set_mask_props(self._cam_over_shoulder_right_mask,
self._obs_config.right_shoulder_camera.mask,
self._obs_config.right_shoulder_camera)
_set_mask_props(self._cam_wrist_mask,
self._obs_config.wrist_camera.mask,
self._obs_config.wrist_camera)
def _place_task(self) -> None:
self._workspace_boundary.clear()
# Find a place in the robot workspace for task
self._active_task.boundary_root().set_orientation(
self._initial_task_pose)
min_rot, max_rot = self._active_task.base_rotation_bounds()
self._workspace_boundary.sample(self._active_task.boundary_root(),
min_rotation=min_rot,
max_rotation=max_rot)
class Scene(object):
"""Controls what is currently in the vrep scene. This is used for making
sure that the tasks are easily reachable. This may be just replaced by
environment. Responsible for moving all the objects. """
def __init__(self,
pyrep: PyRep,
robot: Robot,
obs_config=ObservationConfig()):
self._pyrep = pyrep
self._robot = robot
self._obs_config = obs_config
self._active_task = None
self._inital_task_state = None
self._start_arm_joint_pos = robot.arm.get_joint_positions()
if robot.is_grip():
self._starting_gripper_joint_pos = robot.gripper.get_joint_positions(
)
else:
self._starting_gripper_joint_pos = None
self._workspace = Shape('workspace')
self._workspace_boundary = SpawnBoundary([self._workspace])
self._cam_over_shoulder_left = VisionSensor('cam_over_shoulder_left')
self._cam_over_shoulder_right = VisionSensor('cam_over_shoulder_right')
self._cam_wrist = VisionSensor('cam_wrist')
self._cam_front = VisionSensor('cam_front')
self._cam_over_shoulder_left_mask = VisionSensor(
'cam_over_shoulder_left_mask')
self._cam_over_shoulder_right_mask = VisionSensor(
'cam_over_shoulder_right_mask')
self._cam_wrist_mask = VisionSensor('cam_wrist_mask')
self._cam_front_mask = VisionSensor('cam_front_mask')
self._has_init_task = self._has_init_episode = False
self._variation_index = 0
self._initial_robot_state = (robot.arm.get_configuration_tree(),
robot.gripper.get_configuration_tree())
# Set camera properties from observation config
self._set_camera_properties()
x, y, z = self._workspace.get_position()
minx, maxx, miny, maxy, _, _ = self._workspace.get_bounding_box()
self._workspace_minx = x - np.fabs(minx)
self._workspace_maxx = x + maxx
self._workspace_miny = y - np.fabs(miny)
self._workspace_maxy = y + maxy
self._workspace_minz = z
self._workspace_maxz = z + 1.0 # 1M above workspace
def load(self, task: Task) -> None:
"""Loads the task and positions at the centre of the workspace.
:param task: The task to load in the scene.
"""
task.load() # Load the task in to the scene
# Set at the centre of the workspace
task.get_base().set_position(self._workspace.get_position())
self._inital_task_state = task.get_state()
self._active_task = task
self._initial_task_pose = task.boundary_root().get_orientation()
self._has_init_task = self._has_init_episode = False
self._variation_index = 0
def unload(self) -> None:
"""Clears the scene. i.e. removes all tasks. """
if self._active_task is not None:
self._robot.gripper.release()
if self._has_init_task:
self._active_task.cleanup_()
self._active_task.unload()
self._active_task = None
self._variation_index = 0
def init_task(self) -> None:
self._active_task.init_task()
self._inital_task_state = self._active_task.get_state()
self._has_init_task = True
self._variation_index = 0
def init_episode(self,
index: int,
randomly_place: bool = True,
max_attempts: int = 5) -> List[str]:
"""Calls the task init_episode and puts randomly in the workspace.
"""
self._variation_index = index
if not self._has_init_task:
self.init_task()
# Try a few times to init and place in the workspace
attempts = 0
descriptions = None
while attempts < max_attempts:
descriptions = self._active_task.init_episode(index)
try:
if (randomly_place
and not self._active_task.is_static_workspace()):
self._place_task()
self._active_task.validate()
break
except (BoundaryError, WaypointError) as e:
self._active_task.cleanup_()
self._active_task.restore_state(self._inital_task_state)
attempts += 1
if attempts >= max_attempts:
raise e
# Let objects come to rest
[self._pyrep.step() for _ in range(STEPS_BEFORE_EPISODE_START)]
self._has_init_episode = True
return descriptions
def reset(self) -> None:
"""Resets the joint angles. """
self._robot.gripper.release()
arm, gripper = self._initial_robot_state
self._pyrep.set_configuration_tree(arm)
self._pyrep.set_configuration_tree(gripper)
self._robot.arm.set_joint_positions(self._start_arm_joint_pos,
disable_dynamics=True)
self._robot.arm.set_joint_target_velocities(
[0] * len(self._robot.arm.joints))
if self._robot.is_grip():
self._robot.gripper.set_joint_positions(
self._starting_gripper_joint_pos, disable_dynamics=True)
self._robot.gripper.set_joint_target_velocities(
[0] * len(self._robot.gripper.joints))
else:
self._robot.gripper.release()
if self._active_task is not None and self._has_init_task:
self._active_task.cleanup_()
self._active_task.restore_state(self._inital_task_state)
self._active_task.set_initial_objects_in_scene()
def get_observation(self) -> Observation:
tip = self._robot.arm.get_tip()
joint_forces = None
if self._obs_config.joint_forces:
fs = self._robot.arm.get_joint_forces()
vels = self._robot.arm.get_joint_target_velocities()
joint_forces = self._obs_config.joint_forces_noise.apply(
np.array([-f if v < 0 else f for f, v in zip(fs, vels)]))
ee_forces_flat = None
if self._obs_config.gripper_touch_forces:
ee_forces = self._robot.gripper.get_touch_sensor_forces()
ee_forces_flat = []
for eef in ee_forces:
ee_forces_flat.extend(eef)
ee_forces_flat = np.array(ee_forces_flat)
lsc_ob = self._obs_config.left_shoulder_camera
rsc_ob = self._obs_config.right_shoulder_camera
wc_ob = self._obs_config.wrist_camera
fc_ob = self._obs_config.front_camera
lsc_mask_fn = (rgb_handles_to_mask
if lsc_ob.masks_as_one_channel else lambda x: x)
rsc_mask_fn = (rgb_handles_to_mask
if rsc_ob.masks_as_one_channel else lambda x: x)
wc_mask_fn = (rgb_handles_to_mask
if wc_ob.masks_as_one_channel else lambda x: x)
fc_mask_fn = (rgb_handles_to_mask
if fc_ob.masks_as_one_channel else lambda x: x)
def get_rgb_depth(sensor: VisionSensor, get_rgb: bool, get_depth: bool,
rgb_noise: NoiseModel, depth_noise: NoiseModel,
depth_in_meters: bool):
rgb = depth = None
if sensor is not None and (get_rgb or get_depth):
sensor.handle_explicitly()
if get_rgb:
rgb = sensor.capture_rgb()
if rgb_noise is not None:
rgb = rgb_noise.apply(rgb)
if get_depth:
depth = sensor.capture_depth(depth_in_meters)
if depth_noise is not None:
depth = depth_noise.apply(depth)
return rgb, depth
def get_mask(sensor: VisionSensor, mask_fn):
mask = None
if sensor is not None:
sensor.handle_explicitly()
mask = mask_fn(sensor.capture_rgb())
return mask
left_shoulder_rgb, left_shoulder_depth = get_rgb_depth(
self._cam_over_shoulder_left, lsc_ob.rgb, lsc_ob.depth,
lsc_ob.rgb_noise, lsc_ob.depth_noise, lsc_ob.depth_in_meters)
right_shoulder_rgb, right_shoulder_depth = get_rgb_depth(
self._cam_over_shoulder_right, rsc_ob.rgb, rsc_ob.depth,
rsc_ob.rgb_noise, rsc_ob.depth_noise, rsc_ob.depth_in_meters)
wrist_rgb, wrist_depth = get_rgb_depth(self._cam_wrist, wc_ob.rgb,
wc_ob.depth, wc_ob.rgb_noise,
wc_ob.depth_noise,
wc_ob.depth_in_meters)
front_rgb, front_depth = get_rgb_depth(self._cam_front, fc_ob.rgb,
fc_ob.depth, fc_ob.rgb_noise,
fc_ob.depth_noise,
fc_ob.depth_in_meters)
left_shoulder_mask = get_mask(self._cam_over_shoulder_left_mask,
lsc_mask_fn) if lsc_ob.mask else None
right_shoulder_mask = get_mask(self._cam_over_shoulder_right_mask,
rsc_mask_fn) if rsc_ob.mask else None
wrist_mask = get_mask(self._cam_wrist_mask,
wc_mask_fn) if wc_ob.mask else None
front_mask = get_mask(self._cam_front_mask,
fc_mask_fn) if fc_ob.mask else None
obs = Observation(
left_shoulder_rgb=left_shoulder_rgb,
left_shoulder_depth=left_shoulder_depth,
right_shoulder_rgb=right_shoulder_rgb,
right_shoulder_depth=right_shoulder_depth,
wrist_rgb=wrist_rgb,
wrist_depth=wrist_depth,
front_rgb=front_rgb,
front_depth=front_depth,
left_shoulder_mask=left_shoulder_mask,
right_shoulder_mask=right_shoulder_mask,
wrist_mask=wrist_mask,
front_mask=front_mask,
joint_velocities=(self._obs_config.joint_velocities_noise.apply(
np.array(self._robot.arm.get_joint_velocities()))
if self._obs_config.joint_velocities else None),
joint_positions=(self._obs_config.joint_positions_noise.apply(
np.array(self._robot.arm.get_joint_positions()))
if self._obs_config.joint_positions else None),
joint_forces=(joint_forces
if self._obs_config.joint_forces else None),
gripper_open=(self._robot.gripper.if_open()
if self._obs_config.gripper_open else None),
gripper_pose=(np.array(tip.get_pose())
if self._obs_config.gripper_pose else None),
gripper_matrix=(np.reshape(tip.get_matrix(), (3, 4))
if self._obs_config.gripper_matrix else None),
gripper_touch_forces=(ee_forces_flat
if self._obs_config.gripper_touch_forces else
None),
gripper_joint_positions=(
np.array(self._robot.gripper.get_joint_positions())
if self._obs_config.gripper_joint_positions else None),
wrist_camera_matrix=(np.reshape(self._cam_wrist.get_matrix(),
(3, 4))
if self._cam_wrist.still_exists() else None),
task_low_dim_state=(self._active_task.get_low_dim_state() if
self._obs_config.task_low_dim_state else None))
obs = self._active_task.decorate_observation(obs)
return obs
def step(self):
self._pyrep.step()
self._active_task.step()
def get_demo(self,
record: bool = True,
callable_each_step: Callable[[Observation], None] = None,
randomly_place: bool = True) -> Demo:
"""Returns a demo (list of observations)"""
if not self._has_init_task:
self.init_task()
if not self._has_init_episode:
self.init_episode(self._variation_index,
randomly_place=randomly_place)
self._has_init_episode = False
waypoints = self._active_task.get_waypoints()
if len(waypoints) == 0:
raise NoWaypointsError('No waypoints were found.',
self._active_task)
demo = []
if record:
self._pyrep.step() # Need this here or get_force doesn't work...
demo.append(self.get_observation())
while True:
success = False
for i, point in enumerate(waypoints):
point.start_of_path()
try:
path = point.get_path()
except ConfigurationPathError as e:
raise DemoError(
'Could not get a path for waypoint %d.' % i,
self._active_task) from e
ext = point.get_ext()
path.visualize()
done = False
success = False
while not done:
done = path.step()
self.step()
self._demo_record_step(demo, record, callable_each_step)
success, term = self._active_task.success()
point.end_of_path()
path.clear_visualization()
if len(ext) > 0:
contains_param = False
start_of_bracket = -1
gripper = self._robot.gripper
if 'open_gripper(' in ext:
gripper.release()
start_of_bracket = ext.index('open_gripper(') + 13
contains_param = ext[start_of_bracket] != ')'
if not contains_param:
done = False
while not done:
done = gripper.actuate(1.0, 0.04)
self._pyrep.step()
self._active_task.step()
if self._obs_config.record_gripper_closing:
self._demo_record_step(
demo, record, callable_each_step)
elif 'close_gripper(' in ext:
start_of_bracket = ext.index('close_gripper(') + 14
contains_param = ext[start_of_bracket] != ')'
if not contains_param:
done = False
while not done:
done = gripper.actuate(0.0, 0.04)
self._pyrep.step()
self._active_task.step()
if self._obs_config.record_gripper_closing:
self._demo_record_step(
demo, record, callable_each_step)
if contains_param:
rest = ext[start_of_bracket:]
num = float(rest[:rest.index(')')])
done = False
while not done:
done = gripper.actuate(num, 0.04)
self._pyrep.step()
self._active_task.step()
if self._obs_config.record_gripper_closing:
self._demo_record_step(demo, record,
callable_each_step)
if 'close_gripper(' in ext:
for g_obj in self._active_task.get_graspable_objects():
gripper.grasp(g_obj)
self._demo_record_step(demo, record, callable_each_step)
if not self._active_task.should_repeat_waypoints() or success:
break
# Some tasks may need additional physics steps
# (e.g. ball rowling to goal)
if not success:
for _ in range(10):
self._pyrep.step()
self._active_task.step()
self._demo_record_step(demo, record, callable_each_step)
success, term = self._active_task.success()
if success:
break
success, term = self._active_task.success()
if not success:
raise DemoError('Demo was completed, but was not successful.',
self._active_task)
return Demo(demo)
def get_observation_config(self) -> ObservationConfig:
return self._obs_config
def check_target_in_workspace(self, target_pos: np.ndarray) -> bool:
x, y, z = target_pos
return (self._workspace_maxx > x > self._workspace_minx
and self._workspace_maxy > y > self._workspace_miny
and self._workspace_maxz > z > self._workspace_minz)
def _demo_record_step(self, demo_list, record, func):
if record:
demo_list.append(self.get_observation())
if func is not None:
func(self.get_observation())
def _set_camera_properties(self) -> None:
def _set_rgb_props(rgb_cam: VisionSensor, rgb: bool, depth: bool,
conf: CameraConfig):
if not (rgb or depth):
rgb_cam.remove()
else:
rgb_cam.set_explicit_handling(1)
rgb_cam.set_resolution(conf.image_size)
rgb_cam.set_render_mode(conf.render_mode)
def _set_mask_props(mask_cam: VisionSensor, mask: bool,
conf: CameraConfig):
if not mask:
mask_cam.remove()
else:
mask_cam.set_explicit_handling(1)
mask_cam.set_resolution(conf.image_size)
_set_rgb_props(self._cam_over_shoulder_left,
self._obs_config.left_shoulder_camera.rgb,
self._obs_config.left_shoulder_camera.depth,
self._obs_config.left_shoulder_camera)
_set_rgb_props(self._cam_over_shoulder_right,
self._obs_config.right_shoulder_camera.rgb,
self._obs_config.right_shoulder_camera.depth,
self._obs_config.right_shoulder_camera)
_set_rgb_props(self._cam_wrist, self._obs_config.wrist_camera.rgb,
self._obs_config.wrist_camera.depth,
self._obs_config.wrist_camera)
_set_rgb_props(self._cam_front, self._obs_config.front_camera.rgb,
self._obs_config.front_camera.depth,
self._obs_config.front_camera)
_set_mask_props(self._cam_over_shoulder_left_mask,
self._obs_config.left_shoulder_camera.mask,
self._obs_config.left_shoulder_camera)
_set_mask_props(self._cam_over_shoulder_right_mask,
self._obs_config.right_shoulder_camera.mask,
self._obs_config.right_shoulder_camera)
_set_mask_props(self._cam_wrist_mask,
self._obs_config.wrist_camera.mask,
self._obs_config.wrist_camera)
_set_mask_props(self._cam_front_mask,
self._obs_config.front_camera.mask,
self._obs_config.front_camera)
def _place_task(self) -> None:
self._workspace_boundary.clear()
# Find a place in the robot workspace for task
self._active_task.boundary_root().set_orientation(
self._initial_task_pose)
min_rot, max_rot = self._active_task.base_rotation_bounds()
self._workspace_boundary.sample(self._active_task.boundary_root(),
min_rotation=min_rot,
max_rotation=max_rot)
class SlideBlockToTarget(Task):
def init_task(self) -> None:
self.block = Shape('block')
self.success_sensor = ProximitySensor('success')
self.register_success_conditions([
DetectedCondition(self.block, self.success_sensor)])
self.target = Shape('target')
self.workspace = Shape('workspace')
self.target_boundary = SpawnBoundary([Shape('target_boundary')])
def init_episode(self, index: int) -> List[str]:
self._variation_index = index
self.target_boundary.clear()
self.target_boundary.sample(self.block, change_position=False, min_rotation=(0,0,0), max_rotation=(0,0,0)) #just add the block to the boundary.
notdoneyet = True
while notdoneyet:
self.target_boundary.sample(self.target, min_distance=0.05, min_rotation=(0,0,0), max_rotation=(0,0,0))
notdoneyet = self.get_reward_and_done()[1]
joint_pos = [-5.435585626401007e-05, 0.39650705456733704, 3.4842159948311746e-05, -2.453382730484009, 0.00016987835988402367, 3.0057222843170166, 0.785248875617981]
gripper_pos = [0.0004903227090835571, 0.0007378421723842621]
self.robot.arm.set_joint_positions(joint_pos)
self.robot.gripper.set_joint_positions(gripper_pos)
self.robot.gripper.actuate(0, 1)
return ['slide the block to target',
'slide the block onto the target',
'push the block until it is sitting on top of the target',
'slide the block towards the green target',
'cover the target with the block by pushing the block in its'
' direction']
def get_reward_and_done(self, sparse=False):
dist = np.linalg.norm(np.array(self.target.get_position()) - np.array(self.block.get_position()))
thresh = 0.04
if sparse:
return -(dist > thresh).astype(np.float32), dist < thresh
#epsilon reward
reward = -dist
if dist < thresh:
reward += 0.1
return reward, dist < thresh
#For HER
def compute_reward(self, achieved_goal, desired_goal, info):
dist = np.linalg.norm(np.array(desired_goal) - np.array(achieved_goal))
thresh = 0.04
return -(dist > thresh).astype(np.float32)
def compute_reward_and_done(self, achieved_goal, desired_goal, info):
dist = np.linalg.norm(np.array(desired_goal) - np.array(achieved_goal))
thresh = 0.04
return -(dist > thresh).astype(np.float32), dist < thresh
def variation_count(self) -> int:
return 1
def get_desired_goal(self):
return self.target.get_position()
def get_achieved_goal(self):
return self.block.get_position()
def get_state_obs(self):
joint_positions = self.robot.arm.get_joint_positions()
joint_vel = self.robot.arm.get_joint_velocities()
gripper_position = self.robot.arm.get_tip().get_position()
gripper_vel = np.array(self.robot.arm.get_tip().get_velocity()).flatten()
block = self.block.get_pose()
return np.concatenate([joint_positions, gripper_position, block, joint_vel, gripper_vel])
def get_save_state(self):
joint_positions = self.robot.arm.get_joint_positions()
gripper_positions = self.robot.gripper.get_joint_positions()
block_position = self.block.get_position()
target_position = self.target.get_position()
return [joint_positions, gripper_positions, block_position, target_position]
def restore_save_state(self, state):
joint_positions, gripper_positions , block_position, target_position = state
self.robot.arm.set_joint_positions(joint_positions)
self.robot.gripper.set_joint_positions(gripper_positions)
self.block.set_position(block_position)
self.target.set_position(target_position)
