code: def main(unused_argv):
      tf.compat.v1.enable_v2_behavior()  # The trainer only runs with V2 enabled.
  
      means = [0.1, 0.2, 0.3, 0.45, 0.5]
      env = bern_env.BernoulliPyEnvironment(
          means=means, batch_size=BATCH_SIZE
      )
      environment = tf_py_environment.TFPyEnvironment(env)
  
      def optimal_reward_fn(unused_observation):
          return np.max(means)
  
      def optimal_action_fn(unused_observation):
          return np.int32(np.argmax(means))
  
      if FLAGS.agent == "BernTS":
          agent = bern_ts_agent.BernoulliThompsonSamplingAgent(
              time_step_spec=environment.time_step_spec(),
              action_spec=environment.action_spec(),
              dtype=tf.float64,
              batch_size=BATCH_SIZE,
          )
      else:
          raise ValueError("Only BernoulliTS is supported for now.")
  
      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],
          save_policy=False,
      )
  
file_path: agents/tf_agents/bandits/agents/examples/v2/train_eval_bernoulli.py
imports: import os
  from absl import app
  from absl import flags
  import numpy
  import tensorflow
  from tf_agents.bandits.agents import bernoulli_thompson_sampling_agent
  from tf_agents.bandits.agents.examples.v2 import trainer
  from tf_agents.bandits.environments import bernoulli_py_environment
  from tf_agents.bandits.metrics import tf_metrics
  from tf_agents.environments import tf_py_environment
label_desc: Use image node_id 1 for calling a global function with example usage: main(unused_argv) and returns: np, np