def testSequentialNetwork(self): output_spec = tensor_spec.BoundedTensorSpec([2], tf.float32, 0, 1) network = tanh_normal_projection_network.TanhNormalProjectionNetwork( output_spec) inputs = tf.random.stateless_uniform(shape=[3, 5], seed=[0, 0]) output, _ = network(inputs, outer_rank=1) # Create a squashed distribution. def create_dist(loc_and_scale): ndims = output_spec.shape.num_elements() loc = loc_and_scale[..., :ndims] scale = tf.exp(loc_and_scale[..., ndims:]) distribution = tfp.distributions.MultivariateNormalDiag( loc=loc, scale_diag=scale, validate_args=True, ) return distribution_utils.scale_distribution_to_spec( distribution, output_spec) # Create a sequential network. sequential_network = sequential.Sequential( [network._projection_layer] + [tf.keras.layers.Lambda(create_dist)]) sequential_output, _ = sequential_network(inputs) # Check that mode and standard deviation are the same. self.evaluate(tf.compat.v1.global_variables_initializer()) self.assertAllClose(self.evaluate(output.mode()), self.evaluate(sequential_output.mode())) self.assertAllClose(self.evaluate(output.stddev()), self.evaluate(output.stddev()))
def testTrainableVariables(self): output_spec = tensor_spec.BoundedTensorSpec([2], tf.float32, 0, 1) network = tanh_normal_projection_network.TanhNormalProjectionNetwork( output_spec) inputs = _get_inputs(batch_size=3, num_input_dims=5) network(inputs, outer_rank=1) self.evaluate(tf.compat.v1.global_variables_initializer()) # Dense kernel and bias. self.assertEqual(2, len(network.trainable_variables)) self.assertEqual((5, 4), network.trainable_variables[0].shape) self.assertEqual((4, ), network.trainable_variables[1].shape)
def testBuild(self): output_spec = tensor_spec.BoundedTensorSpec([2], tf.float32, 0, 1) network = tanh_normal_projection_network.TanhNormalProjectionNetwork( output_spec) inputs = _get_inputs(batch_size=3, num_input_dims=5) distribution, _ = network(inputs, outer_rank=1) self.evaluate(tf.compat.v1.global_variables_initializer()) self.assertEqual(tfp.distributions.MultivariateNormalDiag, type(distribution.input_distribution)) means = distribution.input_distribution.loc stds = distribution.input_distribution.scale self.assertAllEqual(means.shape.as_list(), [3] + output_spec.shape.as_list()) self.assertAllEqual(stds.shape.as_list(), [3] + output_spec.shape.as_list() * 2)
def __init__(self, input_tensor_spec, output_tensor_spec, image_encoder=None, fc_encoder_layers=(50, ), fc_layers=(1024, 1024), log_std_bounds=(-10, 2), name='Actor'): """Initializes an Actor network. Args: input_tensor_spec: Tensor spec matching the environment's observation_spec. output_tensor_spec: Tensor spec matching the environment's action_spec. image_encoder: Optional ImageEncoder used on the input observation. If None the network assumes the observation is a vector and this network simply applies fc layers. fc_encoder_layers: Dense layers followed by LayerNorm and a tanh. fc_layers: Iterable (list or tuple) with number of units for a dense layer stack. log_std_bounds: Bounds for scaling the log_std in the TanhNormalProjectionNetwork. name: Name for the network. """ def scale_and_exp(log_std): scale_min, scale_max = log_std_bounds log_std = tf.keras.activations.tanh(log_std) log_std = scale_min + 0.5 * (scale_max - scale_min) * (log_std + 1) return tf.exp(log_std) distribution_projection_network = ( tanh_normal_projection_network.TanhNormalProjectionNetwork( output_tensor_spec, std_transform=scale_and_exp)) super(Actor, self).__init__( input_tensor_spec=input_tensor_spec, state_spec=(), output_spec=distribution_projection_network.output_spec, name=name) # Because of Keras we need to assign to self AFTER we call super. self._image_encoder = image_encoder self._fc_encoder = tf.keras.Sequential() for fc_layer_units in fc_encoder_layers: self._fc_encoder.add( tf.keras.layers.Dense( fc_layer_units, # Default gain of 1.0 matches Pytorch. kernel_initializer=tf.keras.initializers.Orthogonal( gain=1.0), )) # 1e-5 is default epsilon in Pytorch. self._fc_encoder.add(tf.keras.layers.LayerNormalization(epsilon=1e-5)) self._fc_encoder.add( tf.keras.layers.Activation(tf.keras.activations.tanh)) self._dense_layers = tf.keras.Sequential() for units in fc_layers: self._dense_layers.add( tf.keras.layers.Dense( units, activation=tf.keras.activations.relu, kernel_initializer=tf.keras.initializers.Orthogonal( gain=1.0))) self._distribution_projection_network = distribution_projection_network