import tensorflow as tf

  return (dist_entropy, sample_masks, sample_log_prob)


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

  run_config_args = {
      'model_dir': FLAGS.model_dir,
