text: # Copyright (c) 2019 Horizon Robotics. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #      http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  import tensorflow as tf
  from tf_agents.networks.encoding_network import EncodingNetwork as TFAEncodingNetwork
  from alf.layers import NestConcatenate
  
  
  class EncodingNetwork(TFAEncodingNetwork):
      """Feed Forward network with CNN and FNN layers."""
  
      def __init__(self,
                   input_tensor_spec,
                   last_layer_size,
                   last_activation_fn=None,
                   dtype=tf.float32,
                   last_kernel_initializer=None,
                   last_bias_initializer=tf.initializers.Zeros(),
                   preprocessing_combiner=NestConcatenate(axis=-1),
                   **xargs):
          """Create an EncodingNetwork
  
          This EncodingNetwork allows the last layer to have different setting
          from the other layers.
  
          Args:
              last_layer_size (int): size of the last layer
              last_activation_fn: Activation function of the last layer.
              last_kernel_initializer: Initializer for the kernel of the last
                  layer. If none is provided a default
                  tf.initializers.VarianceScaling is used.
              last_bias_initializer: initializer for the bias of the last layer.
              preprocessing_combiner: (Optional.) A keras layer that takes a flat
                  list of tensors and combines them. Good options include
                  `tf.keras.layers.Add` and `tf.keras.layers.Concatenate(axis=-1)`.
                  This layer must not be already built. For more details see
                  the documentation of `networks.EncodingNetwork`. If there is only
                  one input, this will be ignored.
              xargs (dict): See tf_agents.networks.encoding_network.EncodingNetwork
                for detail
          """
          if len(tf.nest.flatten(input_tensor_spec)) == 1:
              preprocessing_combiner = None
          super(EncodingNetwork, self).__init__(
              input_tensor_spec=input_tensor_spec,
              preprocessing_combiner=preprocessing_combiner,
              dtype=dtype,
              **xargs)
  
          if not last_kernel_initializer:
              last_kernel_initializer = tf.initializers.VarianceScaling(
                  scale=2.0, mode='fan_in', distribution='truncated_normal')
          self._last_layer = tf.keras.layers.Dense(
              last_layer_size,
              activation=last_activation_fn,
              kernel_initializer=last_kernel_initializer,
              bias_initializer=last_bias_initializer,
              dtype=dtype)
  
      def call(self, observation, step_type=None, network_state=()):
          state, network_state = super(EncodingNetwork, self).call(
              observation, step_type=step_type, network_state=network_state)
          return self._last_layer(state), network_state
  