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)