import tensorflow as tf



@registry.register_model
class FeedForwardCnnSmallCategoricalPolicy(PolicyBase):
  """Small cnn network with categorical output."""

  def body(self, features):
    observations = features["inputs_raw"]
    # Axis 0    - Batch.
    # Axis 1    - Input Frames, 4 frames.
    # Axis 2, 3 - Height & Width.
    # Axis 4    - Channels RGB, 3 colours.
    x = tf.transpose(observations, [0, 2, 3, 1, 4])
    x_shape = common_layers.shape_list(x)
    x = tf.reshape(x, x_shape[:-2] + [-1])
    dropout = getattr(self.hparams, "dropout_ppo", 0.0)
    with tf.variable_scope("feed_forward_cnn_small"):
      x = tf.cast(x, tf.float32) / 255.0
      x = tf.layers.conv2d(x, 32, (5, 5), strides=(2, 2),
                           activation=tf.nn.relu, padding="same")
      x = tf.layers.conv2d(x, 32, (5, 5), strides=(2, 2),
                           activation=tf.nn.relu, padding="same")

      flat_x = tf.layers.flatten(x)
      if self.use_epochs:
        epoch = features["epoch"] + tf.zeros([x_shape[0]], dtype=tf.int32)
        # Randomly set epoch to 0 in some cases as that's the inference value.
        rand = tf.random.uniform([x_shape[0]])
