def __init__(self, max_len, dtype, shape=(), name='TFDeque'):
self._max_len = tf.convert_to_tensor(max_len, dtype=tf.int32)
self._spec = tf.TensorSpec(shape, dtype, name='Buffer')
self._buffer = table.Table(self._spec, capacity=max_len)
self._head = common.create_variable(
initial_value=0, dtype=tf.int32, shape=(), name=name + 'Head')
def testReadWriteSingle(self):
spec = [
specs.TensorSpec([3], tf.float32, 'action'),
[
specs.TensorSpec([5], tf.float32, 'camera'),
specs.TensorSpec([3, 2], tf.float32, 'lidar')
]
]
replay_table = table.Table(spec, capacity=3)
variables = replay_table.variables()
self.assertEqual(3, len(variables))
self.assertAllEqual(
['Table/action:0', 'Table/camera:0', 'Table/lidar:0'],
[v.name for v in variables])
expected_values = [
1 * np.ones(spec[0].shape.as_list()),
[
2 * np.ones(spec[1][0].shape.as_list()),
3 * np.ones(spec[1][1].shape.as_list())
]
]
tensors = tf.nest.map_structure(
lambda x: tf.convert_to_tensor(value=x, dtype=tf.float32),
expected_values)
write_op = replay_table.write(0, tensors)
read_op = replay_table.read(0)
self.evaluate(tf.compat.v1.global_variables_initializer())
self.evaluate(write_op)
read_value_ = self.evaluate(read_op)
tf.nest.map_structure(self.assertAllClose, read_value_,
expected_values)
def testReadWriteBatch(self):
spec = [
specs.TensorSpec([3], tf.float32, 'action'),
[
specs.TensorSpec([5], tf.float32, 'camera'),
specs.TensorSpec([3, 2], tf.float32, 'lidar')
]
]
replay_table = table.Table(spec, capacity=4)
batch_size = 2
expected_values = [
1 * np.ones([batch_size] + spec[0].shape.as_list()),
[
2 * np.ones([batch_size] + spec[1][0].shape.as_list()),
3 * np.ones([batch_size] + spec[1][1].shape.as_list())
]
]
tensors = tf.nest.map_structure(
lambda x: tf.convert_to_tensor(value=x, dtype=tf.float32),
expected_values)
write_op = replay_table.write(list(range(batch_size)), tensors)
read_op = replay_table.read(list(range(batch_size)))
self.evaluate(tf.compat.v1.global_variables_initializer())
self.evaluate(write_op)
read_value_ = self.evaluate(read_op)
tf.nest.map_structure(self.assertAllClose, read_value_,
expected_values)
def testReadWriteNamedTuple(self):
# pylint: disable=invalid-name
Observation = collections.namedtuple('Observation',
['action', 'camera', 'lidar'])
# pylint: enable=invalid-name
spec = Observation(action=specs.TensorSpec([3], tf.float32, 'action'),
camera=specs.TensorSpec([5], tf.float32, 'camera'),
lidar=specs.TensorSpec([3, 2], tf.float32, 'lidar'))
replay_table = table.Table(spec, capacity=3)
variables = replay_table.variables()
self.assertEqual(3, len(variables))
self.assertAllEqual(
['Table/action:0', 'Table/camera:0', 'Table/lidar:0'],
[v.name for v in variables])
expected_values = Observation(
action=1 * np.ones(spec.action.shape.as_list()),
camera=2 * np.ones(spec.camera.shape.as_list()),
lidar=3 * np.ones(spec.lidar.shape.as_list()))
tensors = tf.nest.map_structure(
lambda x: tf.convert_to_tensor(value=x, dtype=tf.float32),
expected_values)
write_op = replay_table.write(0, tensors)
read_op = replay_table.read(0)
self.evaluate(tf.compat.v1.global_variables_initializer())
self.evaluate(write_op)
read_value_ = self.evaluate(read_op)
tf.nest.map_structure(self.assertAllClose, read_value_,
expected_values)
def testWritePartialSlots(self):
spec = [
specs.TensorSpec([3], tf.float32, 'action'),
[
specs.TensorSpec([5], tf.float32, 'camera'),
specs.TensorSpec([3, 2], tf.float32, 'lidar')
]
]
replay_table = table.Table(spec, capacity=4)
batch_size = 2
action1 = 1 * np.ones([batch_size] + spec[0].shape.as_list())
camera1 = 2 * np.ones([batch_size] + spec[1][0].shape.as_list())
lidar1 = 3 * np.ones([batch_size] + spec[1][1].shape.as_list())
write_op1 = replay_table.write(list(range(batch_size)),
[action1, [camera1, lidar1]])
lidar2 = 10 * np.ones([batch_size] + spec[1][1].shape.as_list())
action2 = 20 * np.ones([batch_size] + spec[0].shape.as_list())
write_op2 = replay_table.write(list(range(batch_size)),
[lidar2, [action2]],
['lidar', ['action']])
read_op = replay_table.read(list(range(batch_size)))
self.evaluate(tf.compat.v1.global_variables_initializer())
self.evaluate(write_op1)
self.evaluate(write_op2)
read_value_ = self.evaluate(read_op)
expected_values = [action2, [camera1, lidar2]]
tf.nest.map_structure(self.assertAllClose, read_value_,
expected_values)
def testReadPartialSlots(self):
spec = [
specs.TensorSpec([3], tf.float32, 'action'),
[
specs.TensorSpec([5], tf.float32, 'camera'),
specs.TensorSpec([3, 2], tf.float32, 'lidar')
]
]
replay_table = table.Table(spec, capacity=4)
batch_size = 2
action = 1 * np.ones([batch_size] + spec[0].shape.as_list())
camera = 2 * np.ones([batch_size] + spec[1][0].shape.as_list())
lidar = 3 * np.ones([batch_size] + spec[1][1].shape.as_list())
values = [action, [camera, lidar]]
tensors = nest.map_structure(
lambda x: tf.convert_to_tensor(x, dtype=tf.float32), values)
write_op = replay_table.write(list(range(batch_size)), tensors)
read_op = replay_table.read(list(range(batch_size)),
slots=['lidar', ['action']])
self.evaluate(tf.global_variables_initializer())
self.evaluate(write_op)
read_value_ = self.evaluate(read_op)
expected_values = [lidar, [action]]
nest.map_structure(self.assertAllClose, read_value_, expected_values)
def testReadWriteString(self):
spec = [
specs.TensorSpec([3], tf.float32, 'action'),
[
specs.TensorSpec([], tf.string, 'camera'),
specs.TensorSpec([3, 2], tf.float32, 'lidar')
]
]
replay_table = table.Table(spec, capacity=3)
variables = replay_table.variables()
self.assertEqual(3, len(variables))
self.assertAllEqual(
['Table/action:0', 'Table/camera:0', 'Table/lidar:0'],
[v.name for v in variables])
expected_values = [
1 * np.ones(spec[0].shape.as_list()),
[b'foo', 3 * np.ones(spec[1][1].shape.as_list())]
]
tensors = nest.map_structure(
lambda x, dtype: tf.convert_to_tensor(x, dtype=dtype),
expected_values, [tf.float32, [tf.string, tf.float32]])
write_op = replay_table.write(0, tensors)
read_op = replay_table.read(0)
self.evaluate(tf.global_variables_initializer())
self.evaluate(write_op)
read_value_ = self.evaluate(read_op)
self.assertAllClose(read_value_[0], expected_values[0])
self.assertEqual(read_value_[1][0], expected_values[1][0])
self.assertAllClose(read_value_[1][1], expected_values[1][1])
def testReadWriteDict(self):
spec = {
'action': specs.TensorSpec([3], tf.float32, 'action'),
'camera': specs.TensorSpec([5], tf.float32, 'camera'),
'lidar': specs.TensorSpec([3, 2], tf.float32, 'lidar')
}
replay_table = table.Table(spec, capacity=3)
variables = replay_table.variables()
self.assertEqual(3, len(variables))
self.assertAllEqual(
['Table/action:0', 'Table/camera:0', 'Table/lidar:0'],
[v.name for v in variables])
expected_values = {
'action': 1 * np.ones(spec['action'].shape.as_list()),
'camera': 2 * np.ones(spec['camera'].shape.as_list()),
'lidar': 3 * np.ones(spec['lidar'].shape.as_list())
}
tensors = nest.map_structure(
lambda x: tf.convert_to_tensor(x, dtype=tf.float32),
expected_values)
write_op = replay_table.write(0, tensors)
read_op = replay_table.read(0)
self.evaluate(tf.global_variables_initializer())
self.evaluate(write_op)
read_value_ = self.evaluate(read_op)
nest.map_structure(self.assertAllClose, read_value_, expected_values)
def testSaveRestore(self):
spec = [
specs.TensorSpec([3], tf.float32),
specs.TensorSpec([5], tf.float32, 'lidar'),
specs.TensorSpec([3, 2], tf.float32, 'lidar')
]
replay_table = table.Table(spec, capacity=3)
self.evaluate(tf.compat.v1.global_variables_initializer())
directory = self.get_temp_dir()
prefix = os.path.join(directory, 'table')
root = tf.train.Checkpoint(table=replay_table)
save_path = root.save(prefix)
root.restore(save_path).assert_consumed().run_restore_ops()
def testDuplicateSpecNames(self):
spec = [
specs.TensorSpec([3], tf.float32, 'lidar'),
specs.TensorSpec([5], tf.float32, 'lidar'),
specs.TensorSpec([3, 2], tf.float32, 'lidar')
]
replay_table = table.Table(spec, capacity=3)
variables = replay_table.variables()
self.assertEqual(3, len(variables))
self.assertAllEqual(
['Table/lidar:0', 'Table/lidar_1:0', 'Table/lidar_2:0'],
[v.name for v in variables])
expected_slots = ['lidar', 'lidar_1', 'lidar_2']
self.assertAllEqual(replay_table.slots, expected_slots)
tensors = replay_table.read(0, expected_slots)
tf.nest.map_structure(lambda x, y: self.assertEqual(x.shape, y.shape),
spec, tensors)
