import tensorflow as tf

@registry.register_model
class DenseBitwiseCategoricalPolicy(PolicyBase):
  """Dense network with bitwise input and categorical output."""

  def body(self, features):
    observations = features["inputs"]
    flat_x = tf.layers.flatten(observations)
    with tf.variable_scope("dense_bitwise"):
      flat_x = discretization.int_to_bit_embed(flat_x, 8, 32)

      x = tf.layers.dense(flat_x, 256, activation=tf.nn.relu)
      x = tf.layers.dense(flat_x, 128, activation=tf.nn.relu)

      logits = tf.layers.dense(x, self.hparams.problem.num_actions)

      value = tf.layers.dense(x, 1)[..., 0]

    return {"target_policy": logits, "target_value": value}


@registry.register_model
class RandomPolicy(PolicyBase):
