import tensorflow as tf

        name='loss_logits_rl_w',
        initializer=tf.initializers.zeros(),
        shape=[
            FLAGS.num_choices,
        ],
        dtype=tf.float32)
    dist = tfp.distributions.Categorical(logits=logits)
    dist_entropy = tf.reduce_sum(dist.entropy())

    sample = dist.sample()
    sample_masks = 1. * tf.cast(sample, tf.float32) / FLAGS.loss_weight_scale
    sample_log_prob = tf.reduce_mean(dist.log_prob(sample))

  return (dist_entropy, sample_masks, sample_log_prob)


def main(unused_argv):
  tf.set_random_seed(FLAGS.random_seed)

  run_config_args = {
