code: def main(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,
              ],
          )
  
file_path: agents/tf_agents/bandits/agents/examples/v2/train_eval_mushroom.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import os
  from absl import app
  from absl import flags
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.agents import lin_ucb_agent
  from tf_agents.bandits.agents import linear_thompson_sampling_agent
  from tf_agents.bandits.agents.examples.v2 import trainer
  from tf_agents.bandits.environments import classification_environment
  from tf_agents.bandits.environments import dataset_utilities
  from tf_agents.bandits.environments import environment_utilities
  from tf_agents.bandits.metrics import tf_metrics
label_desc: Use image node_id 1 for calling a global function with example usage: main(unused_argv) without return types