code: def main(unused_argv):
      tf.compat.v1.enable_v2_behavior()  # The trainer only runs with V2 enabled.
      if MU_BASE[0] > MU_HIGH:
          raise ValueError(
              "The reward for action1 (MU_BASE[0]) should always be "
              "less than the reward for the optimal outer action "
              "(MU_HIGH)."
          )
  
      with tf.device("/CPU:0"):  # due to b/128333994
          env = wheel_py_environment.WheelPyEnvironment(
              DELTA, MU_BASE, STD_BASE, MU_HIGH, STD_HIGH, BATCH_SIZE
          )
          environment = tf_py_environment.TFPyEnvironment(env)
  
          optimal_reward_fn = functools.partial(
              environment_utilities.tf_wheel_bandit_compute_optimal_reward,
              delta=DELTA,
              mu_inside=MU_BASE[0],
              mu_high=MU_HIGH,
          )
          optimal_action_fn = functools.partial(
              environment_utilities.tf_wheel_bandit_compute_optimal_action,
              delta=DELTA,
          )
          network = q_network.QNetwork(
              input_tensor_spec=environment.time_step_spec().observation,
              action_spec=environment.action_spec(),
              fc_layer_params=(LAYERS),
          )
  
          if FLAGS.agent == "LinUCB":
              agent = lin_ucb_agent.LinearUCBAgent(
                  time_step_spec=environment.time_step_spec(),
                  action_spec=environment.action_spec(),
                  alpha=AGENT_ALPHA,
                  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":
              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,
              )
          elif FLAGS.agent == "random":
              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=1.0,
              )
          elif FLAGS.agent == "Mix":
              emit_policy_info = (
                  policy_utilities.InfoFields.PREDICTED_REWARDS_MEAN,
              )
              agent_epsgreedy = (
                  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
                      ),
                      emit_policy_info=emit_policy_info,
                      epsilon=EPSILON,
                  )
              )
              agent_linucb = lin_ucb_agent.LinearUCBAgent(
                  time_step_spec=environment.time_step_spec(),
                  action_spec=environment.action_spec(),
                  alpha=AGENT_ALPHA,
                  emit_policy_info=emit_policy_info,
                  dtype=tf.float32,
              )
              agent_random = 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
                  ),
                  emit_policy_info=emit_policy_info,
                  epsilon=1.0,
              )
              agent_halfrandom = (
                  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
                      ),
                      emit_policy_info=emit_policy_info,
                      epsilon=0.5,
                  )
              )
              agent = exp3_mixture_agent.Exp3MixtureAgent(
                  (
                      agent_epsgreedy,
                      agent_linucb,
                      agent_random,
                      agent_halfrandom,
                  )
              )
  
          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_wheel.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
  from tf_agents.bandits.agents import exp3_mixture_agent
  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 environment_utilities
  from tf_agents.bandits.environments import wheel_py_environment
  from tf_agents.bandits.metrics import tf_metrics
  from tf_agents.environments import tf_py_environment
  from tf_agents.networks import q_network
  from tf_agents.policies import utils
label_desc: Use image node_id 1 for calling a global function with example usage: main(unused_argv) without return types