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
          observation_shape = [CONTEXT_DIM]
          overall_shape = [BATCH_SIZE] + observation_shape
          observation_distribution = tfd.Normal(
              loc=tf.zeros(overall_shape), scale=tf.ones(overall_shape)
          )
          interval_distribution = tfd.Deterministic(STATIONARITY_LENGTH)
          action_shape = [NUM_ACTIONS]
          observation_to_reward_shape = observation_shape + action_shape
          observation_to_reward_distribution = tfd.Normal(
              loc=tf.zeros(observation_to_reward_shape),
              scale=tf.ones(observation_to_reward_shape),
          )
          additive_reward_distribution = tfd.Normal(
              loc=tf.zeros(action_shape),
              scale=(REWARD_NOISE_VARIANCE * tf.ones(action_shape)),
          )
          environment_dynamics = pse.PiecewiseStationaryDynamics(
              observation_distribution,
              interval_distribution,
              observation_to_reward_distribution,
              additive_reward_distribution,
          )
          environment = nse.NonStationaryStochasticEnvironment(
              environment_dynamics
          )
  
          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(
              environment.environment_dynamics.compute_optimal_reward
          )
          suboptimal_arms_metric = tf_bandit_metrics.SuboptimalArmsMetric(
              environment.environment_dynamics.compute_optimal_action
          )
  
          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_piecewise_linear.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  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 non_stationary_stochastic_environment
  from tf_agents.bandits.environments import piecewise_stochastic_environment
  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