def METHOD_NAME(unused_argv):
  tf.compat.v1.enable_v2_behavior()  # The trainer only runs with V2 enabled.
  with tf.device('/CPU:0'):  # due to b/128333994
    mushroom_reward_distribution = (
        dataset_utilities.mushroom_reward_distribution(
            r_noeat=0.0,
            r_eat_safe=5.0,
            r_eat_poison_bad=-35.0,
            r_eat_poison_good=5.0,
            prob_poison_bad=0.5,
        )
    )
    mushroom_dataset = dataset_utilities.convert_mushroom_csv_to_tf_dataset(
        FLAGS.mushroom_csv
    )
    environment = ce.ClassificationBanditEnvironment(
        mushroom_dataset, mushroom_reward_distribution, BATCH_SIZE
    )
    optimal_reward_fn = functools.partial(
        env_util.compute_optimal_reward_with_classification_environment,
        environment=environment,
    )
    optimal_action_fn = functools.partial(
        env_util.compute_optimal_action_with_classification_environment,
        environment=environment,
    )
    if FLAGS.agent == 'LinUCB':
      agent = lin_ucb_agent.LinearUCBAgent(
          time_step_spec=environment.time_step_spec(),
          action_spec=environment.action_spec(),
          alpha=AGENT_ALPHA,
          gamma=0.95,
          emit_log_probability=False,
          dtype=tf.float32,
      )
    elif FLAGS.agent == 'LinTS':
      agent = lin_ts_agent.LinearThompsonSamplingAgent(
          time_step_spec=environment.time_step_spec(),
          action_spec=environment.action_spec(),
          alpha=AGENT_ALPHA,
          gamma=0.95,
          dtype=tf.float32,
      )
    regret_metric = tf_bandit_metrics.RegretMetric(optimal_reward_fn)
    suboptimal_arms_metric = tf_bandit_metrics.SuboptimalArmsMetric(
        optimal_action_fn
    )
    trainer.train(
        root_dir=FLAGS.root_dir,
        agent=agent,
        environment=environment,
        training_loops=TRAINING_LOOPS,
        steps_per_loop=STEPS_PER_LOOP,
        additional_metrics=[regret_metric, suboptimal_arms_metric],
    )