def testCheckpoint(self):
spec = self.default_specs()
replay_table = episodic_table.EpisodicTable(spec, capacity=3)
empty_values = self.np_values(spec, 0)
test_values = self.np_values(spec, 5)
self.evaluate(tf.compat.v1.global_variables_initializer())
self.evaluate(replay_table.append(1, test_values))
checkpoint = tf.train.Checkpoint(table=replay_table)
tmpdir = self.get_temp_dir()
if tf.executing_eagerly():
location = checkpoint.save(os.path.join(tmpdir, 'ckpt'))
else:
with self.cached_session() as sess:
location = checkpoint.save(os.path.join(tmpdir, 'ckpt'),
session=sess)
reload_replay_table = episodic_table.EpisodicTable(spec, capacity=3)
reload_checkpoint = tf.train.Checkpoint(table=reload_replay_table)
status = reload_checkpoint.restore(location)
status.assert_consumed()
with self.cached_session() as sess:
status.initialize_or_restore(session=sess)
self._test_episode_lists(replay_table, empty_values, test_values)
self._test_episode_lists(reload_replay_table, empty_values,
test_values)
def testGetAddMultiple(self):
spec = self.default_specs()
replay_table = episodic_table.EpisodicTable(spec, capacity=4)
batch_size = 2
input_values = self.np_values(spec, batch_size)
expected_values = self.np_values(spec)
empty_values = self.np_values(spec, 0)
tensors = tf.nest.map_structure(
lambda x: tf.convert_to_tensor(value=x, dtype=tf.float32),
input_values)
write_op = replay_table.add([0, 1], tensors)
values_0 = replay_table.get_episode_values(0)
values_1 = replay_table.get_episode_values(1)
# This should be empty
values_2 = replay_table.get_episode_values(2)
self.evaluate(tf.compat.v1.global_variables_initializer())
self.evaluate(write_op)
tf.nest.map_structure(self.assertAllClose, expected_values,
self.evaluate(values_0))
tf.nest.map_structure(self.assertAllClose, expected_values,
self.evaluate(values_1))
tf.nest.map_structure(self.assertAllClose, empty_values,
self.evaluate(values_2))
def testGetEpisodeListsSomeRows(self):
spec = self.default_specs()
replay_table = episodic_table.EpisodicTable(spec, capacity=3)
empty_values = self.np_values(spec, 0)
test_values = self.np_values(spec, 5)
self.evaluate(tf.compat.v1.global_variables_initializer())
self.evaluate(replay_table.append(1, test_values))
# Get episode lists for rows 1 and 2. The returned values should look like:
# {'action': [<tensorlist for row 1>, <tensorlist for row 2>],
# 'camera': [<tensorlist for row 1>, <tensorlist for row 2>],
# 'lidar': [<tensorlist for row 1>, <tensorlist for row 2>]}
episode_lists = replay_table.get_episode_lists([1, 2])
for episode_list_slot in tf.nest.flatten(episode_lists):
self.assertEqual(episode_list_slot.shape.rank, 1)
self.assertEqual(self.evaluate(tf.size(input=episode_list_slot)),
2)
# Stack episode tensors for row 1, i.e.:
# {'action': <five items>, 'camera': <five items>, 'lidar': <five items>}
episode_tensors_1 = tf.nest.map_structure(
lambda slot, lists: replay_table._stack_tensor_list(
slot, lists[0]), replay_table.slots, episode_lists)
# Should be equivalent to
# Stack episode tensors for row 2, which is empty.
episode_tensors_2 = tf.nest.map_structure(
lambda slot, lists: replay_table._stack_tensor_list(
slot, lists[1]), replay_table.slots, episode_lists)
self.assertAllClose(self.evaluate(episode_tensors_1), test_values)
self.assertAllClose(self.evaluate(episode_tensors_2), empty_values)
def testGetEmpty(self):
spec = self.default_specs()
replay_table = episodic_table.EpisodicTable(spec, capacity=3)
empty_values = self.np_values(spec, 0)
values = replay_table.get_episode_values(0)
self.evaluate(tf.compat.v1.global_variables_initializer())
values_ = self.evaluate(values)
tf.nest.map_structure(self.assertAllClose, empty_values, values_)
def testGetEpisodeListsAllRows(self):
spec = self.default_specs()
replay_table = episodic_table.EpisodicTable(spec, capacity=3)
empty_values = self.np_values(spec, 0)
test_values = self.np_values(spec, 5)
self.evaluate(tf.compat.v1.global_variables_initializer())
self.evaluate(replay_table.append(1, test_values))
# Get episode lists for all rows. The returned values should look like:
# {'action': [<tensorlist for row 0>, <tensorlist for row 1>,
# <tensorlist for row 2>],
# ... }
self._test_episode_lists(replay_table, empty_values, test_values)
def testGetEpisodeListsOneRow(self):
spec = self.default_specs()
replay_table = episodic_table.EpisodicTable(spec, capacity=3)
test_values = self.np_values(spec, 5)
self.evaluate(tf.compat.v1.global_variables_initializer())
self.evaluate(replay_table.append(1, test_values))
# Get the episode lists we just added. The returned value should look like:
# {'action': <tensorlist>, 'camera': <tensorlist>, 'lidar': <tensorlist>}
episode_lists = replay_table.get_episode_lists(1)
for episode_list_slot in tf.nest.flatten(episode_lists):
self.assertEqual(episode_list_slot.shape.rank, 0)
episode_tensors = tf.nest.map_structure(
replay_table._stack_tensor_list, replay_table.slots, episode_lists)
self.assertAllClose(self.evaluate(episode_tensors), test_values)
def testExtend(self):
spec = self.default_specs()
replay_table = episodic_table.EpisodicTable(spec, capacity=3)
test_values = self.np_values(spec, 5)
self.evaluate(tf.compat.v1.global_variables_initializer())
self.evaluate(replay_table.append(1, test_values))
# Extend row 2 by the contents of row 1.
self.evaluate(replay_table.extend(2,
replay_table.get_episode_lists(1)))
# Extend rows 0 and 1 by the contents of rows 1 and 2.
self.evaluate(
replay_table.extend([0, 1], replay_table.get_episode_lists([1,
2])))
episode_0, episode_1, episode_2 = self.evaluate(
[replay_table.get_episode_values(r) for r in range(3)])
self.assertAllClose(episode_0, self.np_values(spec, 5))
self.assertAllClose(episode_1, self.np_values(spec, 10))
self.assertAllClose(episode_2, self.np_values(spec, 5))
def testGetAddSingle(self):
spec = self.default_specs()
replay_table = episodic_table.EpisodicTable(spec, capacity=3)
expected_values = self.np_values(spec)
tensors = tf.nest.map_structure(
lambda x: tf.convert_to_tensor(value=x, dtype=tf.float32),
expected_values)
add_op = replay_table.add([0], tensors)
values = replay_table.get_episode_values(0)
# Check static shape
assert_same_shape = lambda s, v: self.assertEqual(s.shape[1:], v.shape)
tf.nest.map_structure(assert_same_shape, values, spec)
self.evaluate(tf.compat.v1.global_variables_initializer())
self.evaluate(add_op)
values_ = self.evaluate(values)
tf.nest.map_structure(self.assertAllClose, expected_values, values_)
def testGetAddAppendMultiple(self):
spec = self.default_specs()
replay_table = episodic_table.EpisodicTable(spec, capacity=4)
batch_size = 2
input_values = self.np_values(spec, batch_size)
expected_values = self.np_values(spec)
empty_values = self.np_values(spec, 0)
tensors = tf.nest.map_structure(
lambda x: tf.convert_to_tensor(value=x, dtype=tf.float32),
input_values)
# Pull out the first entry in the batch, and add an outer
# dimension to represent a single time step that we'll append.
tensors_batch0 = tf.nest.map_structure(
lambda x: tf.expand_dims(x[0, ...], 0), tensors)
# We will append tensors_batch0 to row 0, which contains x[0].
expected_appended_values = tf.nest.map_structure(
lambda x: np.stack((x[0], x[0])), input_values)
# batch_size == 2, so add [0, 1]
write_op = replay_table.add([0, 1], tensors)
append_op_0 = lambda: replay_table.append(0, tensors_batch0)
values_0 = lambda: replay_table.get_episode_values(0)
values_1 = lambda: replay_table.get_episode_values(1)
values_2 = lambda: replay_table.get_episode_values(2)
self.evaluate(tf.compat.v1.global_variables_initializer())
self.evaluate(write_op)
tf.nest.map_structure(self.assertAllClose, expected_values,
self.evaluate(values_0()))
tf.nest.map_structure(self.assertAllClose, expected_values,
self.evaluate(values_1()))
tf.nest.map_structure(self.assertAllClose, empty_values,
self.evaluate(values_2()))
self.evaluate(append_op_0())
tf.nest.map_structure(self.assertAllClose, expected_appended_values,
self.evaluate(values_0()))
