import tensorflow as tf


@registry.register_model
class FeedForwardCategoricalPolicy(PolicyBase):
  """Feed-forward categorical."""

  def body(self, features):
    observations = features["inputs_raw"]
    observations = tf.cast(observations, tf.float32)
    flat_observations = tf.layers.flatten(observations)
    with tf.variable_scope("policy"):
      x = flat_observations
      for size in self.hparams.policy_layers:
        x = tf.layers.dense(x, size, activation=tf.nn.relu)
      logits = tf.layers.dense(x, self.hparams.problem.num_actions)
      logits = tf.expand_dims(logits, axis=1)
    with tf.variable_scope("value"):
      x = flat_observations
      for size in self.hparams.value_layers:
        x = tf.layers.dense(x, size, activation=tf.nn.relu)
      value = tf.layers.dense(x, 1)
    logits = clip_logits(logits, self.hparams)
    return {"target_policy": logits, "target_value": value}


@registry.register_model
class FeedForwardCnnSmallCategoricalPolicy(PolicyBase):
