def testHandlePreprocessingLayers(self, lstm_size, rnn_construction_fn):
observation_spec = (tensor_spec.TensorSpec([1], tf.float32),
tensor_spec.TensorSpec([], tf.float32))
time_step_spec = ts.time_step_spec(observation_spec)
time_step = tensor_spec.sample_spec_nest(time_step_spec, outer_dims=(3, 4))
action_spec = [
tensor_spec.BoundedTensorSpec((2,), tf.float32, 2, 3),
tensor_spec.BoundedTensorSpec((3,), tf.int32, 0, 3)
]
preprocessing_layers = (tf.keras.layers.Dense(4),
sequential_layer.SequentialLayer([
tf.keras.layers.Reshape((1,)),
tf.keras.layers.Dense(4)
]))
net = actor_distribution_rnn_network.ActorDistributionRnnNetwork(
observation_spec,
action_spec,
preprocessing_layers=preprocessing_layers,
lstm_size=lstm_size,
preprocessing_combiner=tf.keras.layers.Add(),
rnn_construction_fn=rnn_construction_fn,
rnn_construction_kwargs={'lstm_size': 3})
initial_state = actor_policy.ActorPolicy(time_step_spec, action_spec,
net).get_initial_state(3)
action_distributions, _ = net(time_step.observation, time_step.step_type,
initial_state)
self.evaluate(tf.compat.v1.global_variables_initializer())
self.assertEqual([3, 4, 2], action_distributions[0].mode().shape.as_list())
self.assertEqual([3, 4, 3], action_distributions[1].mode().shape.as_list())
self.assertGreater(len(net.trainable_variables), 4)
def test_dict_spec_and_pre_processing(self):
input_spec = {
'a': tensor_spec.TensorSpec((32, 32, 3), tf.float32),
'b': tensor_spec.TensorSpec((32, 32, 3), tf.float32)
}
network = encoding_network.EncodingNetwork(
input_spec,
preprocessing_layers={
'a':
sequential_layer.SequentialLayer([
tf.keras.layers.Dense(4, activation='tanh'),
tf.keras.layers.Flatten()
]),
'b':
tf.keras.layers.Flatten()
},
fc_layer_params=(),
preprocessing_combiner=tf.keras.layers.Concatenate(axis=-1),
activation_fn=tf.keras.activations.tanh,
)
sample_input = tensor_spec.sample_spec_nest(input_spec)
output, _ = network(sample_input)
# 6144 is the shape from a concat of flat (32, 32, 3) x2.
self.assertEqual((7168,), output.shape)
def testTrainableVariablesNestedNetwork(self):
sequential_inner = sequential_layer.SequentialLayer(
[tf.keras.layers.Dense(3),
tf.keras.layers.Dense(4)])
sequential = sequential_layer.SequentialLayer(
[tf.keras.layers.Dense(3), sequential_inner])
sequential.build((3, 2))
self.evaluate(tf.compat.v1.global_variables_initializer())
variables = self.evaluate(sequential.trainable_variables)
self.assertLen(variables, 6)
self.assertLen(sequential.variables, 6)
self.assertLen(sequential_inner.variables, 4)
self.assertTrue(sequential.trainable)
sequential.trainable = False
self.assertFalse(sequential.trainable)
self.assertEmpty(sequential.trainable_variables)
self.assertLen(sequential.variables, 6)
def __init__(self, input_tensor_spec, output_tensor_spec):
super(ActorNetwork, self).__init__(input_tensor_spec=input_tensor_spec,
state_spec=(),
name='TestActorNetwork')
num_actions = output_tensor_spec.shape.num_elements()
self._sequential_layer = sequential_layer.SequentialLayer([
tf.keras.layers.Dense(50),
tf.keras.layers.Dense(10),
tf.keras.layers.Dense(num_actions)
])
def testCopy(self):
sequential = sequential_layer.SequentialLayer(
[tf.keras.layers.Dense(3),
tf.keras.layers.Dense(4, use_bias=False)])
clone = type(sequential).from_config(sequential.get_config())
self.assertLen(clone.layers, 2)
for l1, l2 in zip(sequential.layers, clone.layers):
self.assertEqual(l1.dtype, l2.dtype)
self.assertEqual(l1.units, l2.units)
self.assertEqual(l1.use_bias, l2.use_bias)
def testBuild(self):
sequential = sequential_layer.SequentialLayer(
[tf.keras.layers.Dense(4, use_bias=False),
tf.keras.layers.ReLU()])
inputs = np.ones((2, 3))
out = sequential(inputs)
self.evaluate(tf.compat.v1.global_variables_initializer())
out = self.evaluate(out)
weights = self.evaluate(sequential.layers[0].weights[0])
expected = np.dot(inputs, weights)
expected[expected < 0] = 0
self.assertAllClose(expected, out)
def test_layers_buildable(self):
input_spec = {
'a': tensor_spec.TensorSpec((32, 32, 3), tf.float32),
'b': tensor_spec.TensorSpec((32, 32, 3), tf.float32)
}
network = encoding_network.EncodingNetwork(
input_spec,
preprocessing_layers={
'a':
sequential_layer.SequentialLayer([
tf.keras.layers.Dense(4, activation='tanh'),
tf.keras.layers.Flatten()
]),
'b':
tf.keras.layers.Flatten()
},
fc_layer_params=(),
preprocessing_combiner=tf.keras.layers.Concatenate(axis=-1),
activation_fn=tf.keras.activations.tanh,
)
network.create_variables()
self.assertNotEmpty(network.variables)