class PlacingDataset:
def __init__(self, episodes, seed=None):
self.episodes = episodes
self.episodes_place_success_index = list(map(lambda e: e[0], filter(lambda e: e[1]['actions'][1]['reward'] > 0, enumerate(episodes))))
self.episodes_different_objects_ids = {
'wooden': ('2019-12-16-17-01-18-409', '2020-01-22-09-34-02-952'),
'baby': ('2020-01-22-09-34-02-953', '2020-01-29-23-17-15-032'),
}
# Get indexes of episodes between the ids (from above) which have a positive place action
self.episodes_different_objects_index = {
k: list(map(lambda e: e[0], filter(lambda e: v[0] <= e[1]['id'] <= v[1] and e[1]['actions'][1]['reward'] > 0, enumerate(episodes))))
for k, v in self.episodes_different_objects_ids.items()
}
self.size_input = (752, 480)
self.size_memory_scale = 4
self.size_cropped = (200, 200)
self.size_result = (32, 32)
self.size_cropped_area = (self.size_cropped[0] // self.size_memory_scale, self.size_cropped[1] // self.size_memory_scale)
self.use_hindsight = True
self.use_further_hindsight = False
self.use_negative_foresight = True
self.use_own_goal = True
self.use_different_episodes_as_goals = True
self.jittered_hindsight_images = 1
self.jittered_hindsight_x_images = 2 # Only if place reward > 0
self.jittered_goal_images = 1
self.different_episodes_images = 1
self.different_episodes_images_success = 4 # Only if place reward > 0
self.different_object_images = 4 # Only if place reward > 0
self.different_jittered_object_images = 0 # Only if place reward > 0
self.box_distance = 0.281 # [m]
self.indexer = GraspIndexer([0.05, 0.07, 0.086]) # [m]
# self.indexer = GraspIndexer([0.025, 0.05, 0.07, 0.086]) # [m]
self.cameras = ('ed',)
# self.cameras = ('ed', 'rd', 'rc')
self.seed = seed
self.random_gen = np.random.RandomState(seed)
@lru_cache(maxsize=None)
def load_image(self, collection, episode_id, action_id, suffix):
image = Loader.get_image(collection, episode_id, action_id, suffix, as_float=True)
draw_around_box(image, box=Config.box)
image.mat = cv2.resize(image.mat, (self.size_input[0] // self.size_memory_scale, self.size_input[1] // self.size_memory_scale))
image.pixel_size /= self.size_memory_scale
return image
def area_of_interest(self, image, pose):
area = get_area_of_interest_new(
image,
RobotPose(all_data=pose),
size_cropped=self.size_cropped_area,
size_result=self.size_result,
border_color=image.value_from_depth(self.box_distance) / (255 * 255),
)
if len(area.mat.shape) == 2:
return np.expand_dims(area.mat, 2)
return area.mat
def jitter_pose(self, pose, scale_x=0.05, scale_y=0.05, scale_a=1.5, around=True):
new_pose = copy.deepcopy(pose)
if around:
low = [np.minimum(0.001, scale_x), np.minimum(0.001, scale_y), np.minimum(0.06, scale_a)]
mode = [np.minimum(0.006, scale_x), np.minimum(0.006, scale_y), np.minimum(0.32, scale_a)]
high = [scale_x + 1e-6, scale_y + 1e-6, scale_a + 1e-6]
dx, dy, da = self.random_gen.choice([-1, 1], size=3) * self.random_gen.triangular(low, mode, high, size=3)
else:
low = [-scale_x - 1e-6, -scale_y - 1e-6, -scale_a - 1e-6]
mode = [0.0, 0.0, 0.0]
high = [scale_x + 1e-6, scale_y + 1e-6, scale_a + 1e-6]
dx, dy, da = self.random_gen.triangular(low, mode, high, size=3)
new_pose['x'] += np.cos(pose['a']) * dx - np.sin(pose['a']) * dy
new_pose['y'] += np.sin(pose['a']) * dx + np.cos(pose['a']) * dy
new_pose['a'] += da
return new_pose
def generator(self, index):
e = self.episodes[index]
result = []
collection = e['collection']
episode_id = e['id']
grasp = e['actions'][0]
grasp_before = self.load_image(collection, episode_id, 0, 'ed-v')
grasp_before_area = self.area_of_interest(grasp_before, grasp['pose'])
grasp_index = self.indexer.from_pose(grasp['pose'])
# Only single grasp
if len(e['actions']) == 1:
pass
place = e['actions'][1]
place_before = self.load_image(collection, episode_id, 1, 'ed-v')
place_after = self.load_image(collection, episode_id, 1, 'ed-after')
# Generate goal has no action_id
def generate_goal(g_collection, g_episode_id, g_suffix, g_pose, g_suffix_before='v', g_reward=0, g_index=None, g_place_weight=1.0, g_merge_weight=1.0, jitter=None):
if g_collection == collection and g_episode_id == episode_id and g_suffix == 'v' and g_suffix_before == 'v':
place_goal_before = place_before
place_goal = place_before
elif g_collection == collection and g_episode_id == episode_id and g_suffix == 'v' and g_suffix_before == 'after':
place_goal_before = place_after
place_goal = place_before
elif g_collection == collection and g_episode_id == episode_id and g_suffix == 'after' and g_suffix_before == 'v':
place_goal_before = place_before
place_goal = place_after
elif g_collection == collection and g_episode_id == episode_id and g_suffix == 'after' and g_suffix_before == 'after':
place_goal_before = place_after
place_goal = place_after
else:
goal_e = self.episodes[g_index]
g_collection = g_collection if g_collection else goal_e['collection']
g_episode_id = g_episode_id if g_episode_id else goal_e['id']
g_pose = g_pose if g_pose else goal_e['actions'][1]['pose']
place_goal_before = self.load_image(g_collection, g_episode_id, 1, 'ed-' + g_suffix_before)
place_goal = self.load_image(g_collection, g_episode_id, 1, 'ed-' + g_suffix)
if isinstance(jitter, dict):
g_pose = self.jitter_pose(g_pose, **jitter)
place_before_area = self.area_of_interest(place_goal_before, g_pose)
place_goal_area = self.area_of_interest(place_goal, g_pose)
reward_grasp = grasp['reward']
reward_place = g_reward * grasp['reward'] * place['reward']
reward_merge = reward_place
grasp_weight = g_reward
place_weight = (1.0 + 3.0 * reward_place) * reward_grasp * g_place_weight
merge_weight = (1.0 + 3.0 * reward_merge) * reward_grasp * g_merge_weight
return (
grasp_before_area,
place_before_area,
place_goal_area,
(reward_grasp, grasp_index, grasp_weight),
(reward_place, 0, place_weight),
(reward_merge, 0, merge_weight),
)
if self.use_hindsight:
result.append(generate_goal(collection, episode_id, 'after', place['pose'], g_reward=1))
result += [
generate_goal(collection, episode_id, 'after', place['pose'], jitter={})
for _ in range(self.jittered_hindsight_images)
]
if place['reward'] > 0:
result += [
generate_goal(collection, episode_id, 'after', place['pose'], jitter={'scale_x': 0.02, 'scale_y': 0.01, 'scale_a': 0.2})
for _ in range(self.jittered_hindsight_x_images)
]
if self.use_further_hindsight and 'bin_episode' in place:
for i in range(index + 1, len(self.episodes)):
place_later = self.episodes[i]['actions'][1]
if place_later['bin_episode'] != place['bin_episode']:
break
if place_later['reward'] > 0:
result.append(generate_goal(None, None, 'after', place['pose'], g_index=i, g_reward=1))
if self.use_negative_foresight:
g_suffix, g_suffix_before = random.choice([('v', 'v'), ('after', 'after'), ('v', 'after')])
result.append(generate_goal(collection, episode_id, g_suffix, place['pose'], g_suffix_before=g_suffix_before, jitter={'around': False}))
if self.use_own_goal and 'ed-goal' in place['images']:
result.append(generate_goal(collection, episode_id, 'goal', place['pose'], g_place_weight=0.2, g_merge_weight=0.7, g_index=index))
result += [
generate_goal(collection, episode_id, 'goal', place['pose'], g_index=index, jitter={})
for _ in range(self.jittered_goal_images)
]
if self.use_different_episodes_as_goals:
result += [
generate_goal(None, None, 'after', None, g_index=goal_index, g_place_weight=0.0)
for goal_index in self.random_gen.choice(self.episodes_place_success_index, size=self.different_episodes_images)
]
if place['reward'] > 0:
result += [
generate_goal(None, None, 'after', None, g_index=goal_index, g_place_weight=0.0)
for goal_index in self.random_gen.choice(self.episodes_place_success_index, size=self.different_episodes_images_success)
]
for k, v in self.episodes_different_objects_ids.items():
if v[0] <= e['id'] <= v[1]:
result += [
generate_goal(None, None, 'after', None, g_index=goal_index, g_place_weight=0.0)
for goal_index in self.random_gen.choice(self.episodes_different_objects_index[k], size=self.different_object_images)
]
# result += [
# generate_goal(None, None, 'after', None, g_index=goal_index, jitter={})
# for goal_index in self.random_gen.choice(self.episodes_different_objects_index[k], size=self.different_jittered_object_images)
# ]
return [np.array(t, dtype=np.float32) for t in zip(*result)]
def tf_generator(self, index):
r = tf.py_function(
self.generator,
[index],
(tf.float32,) * 6,
)
r[0].set_shape((None, 32, 32, 1))
r[1].set_shape((None, 32, 32, 1))
r[2].set_shape((None, 32, 32, 1))
r[3].set_shape((None, 3))
r[4].set_shape((None, 3))
r[5].set_shape((None, 3))
return (r[0], r[1], r[2]), (r[3], r[4], r[5])
def get_data(self, shuffle=None):
data = tf.data.Dataset.range(0, len(self.episodes))
if shuffle:
shuffle_number = len(self.episodes) if shuffle == 'all' else int(shuffle)
data = data.shuffle(shuffle_number, seed=self.seed)
data = data.map(self.tf_generator, num_parallel_calls=tf.data.experimental.AUTOTUNE)
# data = data.map(self.tf_generator)
return data.interleave(lambda *x: tf.data.Dataset.from_tensor_slices(x), cycle_length=1)
class PlacingDataset:
def __init__(self, episodes, seed=None):
self.episodes = episodes
self.episodes_place_success_index = list(
map(
lambda e: e[0],
filter(
lambda e: len(e[1]['actions']) > 1 and e[1]['actions'][1][
'reward'] > 0, enumerate(episodes))))
# self.episodes_different_objects_index = list(map(lambda e: e[0], filter(lambda e: '2019-12-16-17-01-18-409' <= e[1]['id'] <= '2020-01-16-17-09-47-989' and e[1]['actions'][1]['reward'] > 0, enumerate(episodes))))
self.size_input = (752, 480)
self.size_memory_scale = 4
self.size_cropped = (200, 200)
self.size_result = (32, 32)
self.size_cropped_area = (self.size_cropped[0] //
self.size_memory_scale,
self.size_cropped[1] //
self.size_memory_scale)
self.use_hindsight = True
self.use_further_hindsight = False
self.use_negative_foresight = True
self.use_own_goal = True
self.use_different_episodes_as_goals = True
self.jittered_hindsight_images = 3
self.jittered_hindsight_x_images = 3
self.jittered_goal_images = 2
self.different_episodes_images = 2
self.different_object_images = 5
# self.different_jittered_object_images = 2
self.box_distance = 0.281 # [m]
# self.indexer = GraspIndexer([0.05, 0.07, 0.086]) # [m]
self.indexer = GraspIndexer([0.025, 0.05, 0.07, 0.086]) # [m]
self.cameras = ('ed', 'rd', 'rc')
self.random_gen = np.random.RandomState(seed)
@lru_cache(maxsize=None)
def load_image(self, collection, episode_id, action_id, suffix):
image = Loader.get_image(collection,
episode_id,
action_id,
suffix,
as_float=True)
draw_around_box(image, box=Config.box)
image.mat = cv2.resize(image.mat,
(self.size_input[0] // self.size_memory_scale,
self.size_input[1] // self.size_memory_scale))
image.pixel_size /= self.size_memory_scale
return image
def area_of_interest(self, image, pose):
area = get_area_of_interest_new(
image,
RobotPose(all_data=pose),
size_cropped=self.size_cropped_area,
size_result=self.size_result,
border_color=image.value_from_depth(self.box_distance) /
(255 * 255),
)
if len(area.mat.shape) == 2:
return np.expand_dims(area.mat, 2)
return area.mat
def jitter_pose(self, pose, scale_x=0.05, scale_y=0.05, scale_a=1.5):
new_pose = copy.deepcopy(pose)
low = [
np.minimum(0.002, scale_x),
np.minimum(0.002, scale_y),
np.minimum(0.05, scale_a)
]
high = [scale_x, scale_y, scale_a]
dx, dy, da = self.random_gen.choice(
[-1, 1], size=3) * self.random_gen.uniform(low, high, size=3)
new_pose['x'] += np.cos(pose['a']) * dx - np.sin(pose['a']) * dy
new_pose['y'] += np.sin(pose['a']) * dx + np.cos(pose['a']) * dy
new_pose['a'] += da
return new_pose
def generator(self, index):
e = self.episodes[index]
result = []
collection = e['collection']
episode_id = e['id']
grasp = e['actions'][0]
grasp_before = tuple(
self.load_image(collection, episode_id, 0, camera + '-v')
for camera in self.cameras)
grasp_before_area = tuple(
self.area_of_interest(image, grasp['pose'])
for image in grasp_before)
grasp_index = self.indexer.from_pose(grasp['pose'])
# Only grasp
if len(e['actions']) == 1:
zeros = (np.zeros(self.size_result + (1, )),
np.zeros(self.size_result + (1, )),
np.zeros(self.size_result + (3, )))
result = [
grasp_before_area + zeros + zeros + (
(grasp['reward'], grasp_index, 0.4),
(0, 0, 0),
(0, 0, 0),
)
]
return [np.array(t, dtype=np.float32) for t in zip(*result)]
place = e['actions'][1]
place_before = tuple(
self.load_image(collection, episode_id, 1, camera + '-v')
for camera in self.cameras)
place_after = tuple(
self.load_image(collection, episode_id, 1, camera + '-after')
for camera in self.cameras)
# Generate goal has no action_id
def generate_goal(g_collection,
g_episode_id,
g_suffix,
g_pose,
g_reward=0,
g_index=None,
g_merge_weight=1.0,
jitter=None):
if g_collection == collection and g_episode_id == episode_id and g_suffix == 'v':
place_goal_before = place_before
place_goal = place_before
elif g_collection == collection and g_episode_id == episode_id and g_suffix == 'after':
place_goal_before = place_before
place_goal = place_after
else:
goal_e = self.episodes[g_index]
g_collection = g_collection if g_collection else goal_e[
'collection']
g_episode_id = g_episode_id if g_episode_id else goal_e['id']
g_pose = g_pose if g_pose else goal_e['actions'][1]['pose']
place_goal_before = tuple(
self.load_image(g_collection, g_episode_id, 1, camera +
'-v') for camera in self.cameras)
place_goal = tuple(
self.load_image(g_collection, g_episode_id, 1, camera +
'-' + g_suffix) for camera in self.cameras)
if isinstance(jitter, dict):
g_pose = self.jitter_pose(g_pose, **jitter)
place_before_area = tuple(
self.area_of_interest(image, g_pose)
for image in place_goal_before)
place_goal_area = tuple(
self.area_of_interest(image, g_pose) for image in place_goal)
goal_reward = g_reward
reward_grasp = grasp['reward']
reward_place = goal_reward * place['reward']
reward_merge = reward_place
grasp_weight = g_reward
place_weight = (1.0 + 5.0 * reward_place) * reward_grasp
merge_weight = (1.0 + 5.0 * reward_merge) * g_merge_weight
return grasp_before_area + place_before_area + place_goal_area + (
(reward_grasp, grasp_index, grasp_weight),
(reward_place, 0, place_weight),
(reward_merge, 0, merge_weight),
)
if self.use_hindsight:
result.append(
generate_goal(collection,
episode_id,
'after',
place['pose'],
g_reward=1))
result += [
generate_goal(collection,
episode_id,
'after',
place['pose'],
jitter={})
for _ in range(self.jittered_hindsight_images)
]
result += [
generate_goal(collection,
episode_id,
'after',
place['pose'],
jitter={
'scale_x': 0.025,
'scale_y': 0,
'scale_a': 0
})
for _ in range(self.jittered_hindsight_x_images)
]
if self.use_further_hindsight and 'bin_episode' in place:
for i in range(index + 1, len(self.episodes)):
place_later = self.episodes[i]['actions'][1]
if place_later['bin_episode'] != place['bin_episode']:
break
if place_later['reward'] > 0:
result.append(
generate_goal(None,
None,
'after',
place['pose'],
g_index=i,
g_reward=1))
if self.use_negative_foresight:
result.append(
generate_goal(collection, episode_id, 'v', place['pose']))
if self.use_own_goal and 'ed-goal' in place['images']:
result.append(
generate_goal(collection,
episode_id,
'goal',
place['pose'],
g_index=index))
result += [
generate_goal(collection,
episode_id,
'goal',
place['pose'],
g_index=index,
g_merge_weight=0.5,
jitter={})
for _ in range(self.jittered_goal_images)
]
if self.use_different_episodes_as_goals:
result += [
generate_goal(None,
None,
'after',
None,
g_index=goal_index,
g_merge_weight=0.3) for goal_index in
self.random_gen.choice(self.episodes_place_success_index,
size=self.different_episodes_images)
]
# result += [
# generate_goal(None, None, 'after', None, g_index=goal_index, g_merge_weight=0.3)
# for goal_index in self.random_gen.choice(self.episodes_different_objects_index, size=self.different_object_images)
# ]
# result += [
# generate_goal(None, None, 'after', None, g_index=goal_index, g_merge_weight=0.3, jitter={})
# for goal_index in self.random_gen.choice(self.episodes_different_objects_index, size=self.different_jittered_object_images)
# ]
return [np.array(t, dtype=np.float32) for t in zip(*result)]
def tf_generator(self, index):
r = tf.py_function(
self.generator,
[index],
(tf.float32, ) * (3 * len(self.cameras) + 3),
)
r[0].set_shape((None, 32, 32, 1))
r[1].set_shape((None, 32, 32, 1))
r[2].set_shape((None, 32, 32, 3))
r[3].set_shape((None, 32, 32, 1))
r[4].set_shape((None, 32, 32, 1))
r[5].set_shape((None, 32, 32, 3))
r[6].set_shape((None, 32, 32, 1))
r[7].set_shape((None, 32, 32, 1))
r[8].set_shape((None, 32, 32, 3))
r[9].set_shape((None, 3))
r[10].set_shape((None, 3))
r[11].set_shape((None, 3))
return (r[0], r[1], r[2], r[3], r[4], r[5], r[6], r[7],
r[8]), (r[9], r[10], r[11])
def get_data(self, shuffle=None):
data = tf.data.Dataset.range(0, len(self.episodes))
if shuffle:
shuffle_number = len(
self.episodes) if shuffle == 'all' else int(shuffle)
data = data.shuffle(shuffle_number)
data = data.map(self.tf_generator,
num_parallel_calls=tf.data.experimental.AUTOTUNE)
return data.interleave(
lambda *x: tf.data.Dataset.from_tensor_slices(x), cycle_length=1)
