import tensorflow as tf


      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
