def convert_episode_to_tf_records(base_directory, new_directory, dataloader,
path):
episode = frame.load_episode(path)
features, labels = dataloader.load_features_and_labels_episode(episode)
assert path.rfind(base_directory) > -1
new_path = path[path.rfind(base_directory) + len(base_directory) + 1:]
new_path = os.path.splitext(new_path)[0]
new_path = os.path.splitext(new_path)[0]
new_path = os.path.join(new_directory, new_path + ".tfrecord")
options = tf.python_io.TFRecordOptions(
compression_type=tf.python_io.TFRecordCompressionType.GZIP)
os.makedirs(new_path, exist_ok=True)
for i, f in enumerate(episode.frames):
writer = tf.python_io.TFRecordWriter(os.path.join(
new_path, "{}.tfrecord".format(i)),
options=options)
example = tf.train.Example(features=tf.train.Features(
feature={
'action': _int64_feature([f.action]),
'label': _int64_feature(
[f.label] if f.label is not None else []),
'observation': _float_feature(f.observation.reshape(-1)),
'observation_shape': _int64_feature(f.observation.shape),
'image': _bytes_feature([f.image.tobytes()]),
'image_shape': _int64_feature(f.image.shape),
}))
writer.write(example.SerializeToString())
writer.close()
return new_path
def load_features_episode(self, episode_path):
features = {}
episode = frame.load_episode(episode_path)
observations = [frame.observation for frame in episode.frames]
actions = [frame.action for frame in episode.frames]
features['observation'] = np.concatenate([
np.expand_dims(observation, axis=0) for observation in observations
],
axis=0)
labels = np.array([
pong_catastrophe.is_catastrophe(observation, location="bottom")
for observation in observations
])
return features, labels
def load_features_episode(self, episode_path):
features = {}
episode = frame.load_episode(episode_path)
observations = [frame.observation for frame in episode.frames]
actions = [frame.action for frame in episode.frames]
features['observation'] = np.concatenate([
np.expand_dims(observation, axis=0) for observation in observations
],
axis=0)
features['action'] = np.expand_dims(np.array(actions), axis=1)
is_catastrophe = np.array([
pong_catastrophe.is_catastrophe(observation, location="bottom")
for observation in observations
])
is_catastrophe = is_catastrophe[1:]
for key, value in features.items():
features[key] = features[key][:-1]
labels = is_catastrophe
for i in range(1, self.block_radius):
labels = np.logical_or(labels, (np.pad(is_catastrophe[i:],
(0, i), 'constant')))
# Filter out instances where a catastrophe is already in progress
not_already_catastrophe = np.logical_not(is_catastrophe[0:-1])
labels = labels[1:][not_already_catastrophe]
for key, value in features.items():
features[key] = features[key][1:][not_already_catastrophe]
### USE HUMAN LABELS
new_labels = np.array([bool(frame.label) for frame in episode.frames])
new_labels = new_labels[:len(labels)]
labels = new_labels
print([frame.label for frame in episode.frames])
print('total cats', np.sum(labels))
return features, labels
def __iter__(self):
for episode_path in self.episode_paths:
episode = frame_module.load_episode(episode_path)
yield episode_path, episode, 0 # reset the frame_index
obs, reward, done, info = self.env.step(action)
if should_block(self.last_obs, action):
info["frame/should_have_blocked"] = True
self.last_obs = obs
return obs, reward, done, info
if __name__ == "__main__":
from IPython import get_ipython # isort:ignore
from matplotlib import pyplot as plt # isort:ignore
ipython = get_ipython()
ipython.magic("matplotlib inline")
eps = frame.episode_paths("logs/SpaceInvadersRandom")
ep = frame.load_episode(eps[0])
s = set()
f = np.copy(ep.frames[50].image)
barrier_damage(f)
for i in range(0, len(ep.frames)):
if should_block(ep.frames[i].image, ep.frames[i].action):
print(i, is_ship_below_barrier(ep.frames[i].image))
plt.imshow(ep.frames[i].image[155:195, :, :])
plt.show()
plt.imshow(ep.frames[i + 1].image[155:195, :, :])
plt.show()
plt.imshow(ep.frames[i + 2].image[155:195, :, :])
plt.show()
# for i in range(f.shape[0]):
# for j in range(f.shape[1]):