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
          covertype_dataset = (
              dataset_utilities.convert_covertype_dataset(
                  FLAGS.covertype_csv
              )
          )
          covertype_reward_distribution = tfd.Independent(
              tfd.Deterministic(tf.eye(7)), reinterpreted_batch_ndims=2
          )
          environment = ce.ClassificationBanditEnvironment(
              covertype_dataset,
              covertype_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,
                  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,
                  dtype=tf.float32,
              )
          elif FLAGS.agent == "epsGreedy":
              network = q_network.QNetwork(
                  input_tensor_spec=environment.time_step_spec().observation,
                  action_spec=environment.action_spec(),
                  fc_layer_params=LAYERS,
              )
              agent = eps_greedy_agent.NeuralEpsilonGreedyAgent(
                  time_step_spec=environment.time_step_spec(),
                  action_spec=environment.action_spec(),
                  reward_network=network,
                  optimizer=tf.compat.v1.train.AdamOptimizer(
                      learning_rate=LR
                  ),
                  epsilon=EPSILON,
              )
  
          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_covertype.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 import neural_epsilon_greedy_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
  from tf_agents.networks import q_network
label_desc: Use image node_id 1 for calling a global function with example usage: main(unused_argv) without return types