code: def main(unused_argv):
      tf.compat.v1.enable_v2_behavior()  # The trainer only runs with V2 enabled.
  
      class LinearNormalReward(object):
          def __init__(self, theta):
              self.theta = theta
  
          def __call__(self, x):
              mu = np.dot(x, self.theta)
              return np.random.normal(mu, 1)
  
      def _global_context_sampling_fn():
          return np.random.randint(-10, 10, [4]).astype(np.float32)
  
      def _arm_context_sampling_fn():
          return np.random.randint(-2, 3, [5]).astype(np.float32)
  
      reward_fn = LinearNormalReward(HIDDEN_PARAM)
  
      observation_and_action_constraint_splitter = None
      num_actions_fn = None
      if FLAGS.add_num_actions_feature:
          num_actions_fn = lambda: NUM_ACTIONS
  
      env = sspe.StationaryStochasticPerArmPyEnvironment(
          _global_context_sampling_fn,
          _arm_context_sampling_fn,
          NUM_ACTIONS,
          reward_fn,
          num_actions_fn,
          batch_size=BATCH_SIZE,
      )
      environment = tf_py_environment.TFPyEnvironment(env)
  
      if FLAGS.agent == "LinUCB":
          agent = lin_ucb_agent.LinearUCBAgent(
              time_step_spec=environment.time_step_spec(),
              action_spec=environment.action_spec(),
              alpha=AGENT_ALPHA,
              accepts_per_arm_features=True,
              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,
              observation_and_action_constraint_splitter=(
                  observation_and_action_constraint_splitter
              ),
              accepts_per_arm_features=True,
              dtype=tf.float32,
          )
      elif FLAGS.agent == "epsGreedy":
          obs_spec = environment.observation_spec()
          if FLAGS.network == "commontower":
              network = global_and_arm_feature_network.create_feed_forward_common_tower_network(
                  obs_spec, (40, 30), (30, 40), (40, 20)
              )
          elif FLAGS.network == "dotproduct":
              network = global_and_arm_feature_network.create_feed_forward_dot_product_network(
                  obs_spec, (4, 3, 6), (3, 4, 6)
              )
          agent = neural_epsilon_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,
              observation_and_action_constraint_splitter=(
                  observation_and_action_constraint_splitter
              ),
              accepts_per_arm_features=True,
              emit_policy_info=policy_utilities.InfoFields.PREDICTED_REWARDS_MEAN,
          )
      elif FLAGS.agent == "NeuralLinUCB":
          obs_spec = environment.observation_spec()
          network = global_and_arm_feature_network.create_feed_forward_common_tower_network(
              obs_spec, (40, 30), (30, 40), (40, 20), ENCODING_DIM
          )
          agent = neural_linucb_agent.NeuralLinUCBAgent(
              time_step_spec=environment.time_step_spec(),
              action_spec=environment.action_spec(),
              encoding_network=network,
              encoding_network_num_train_steps=EPS_PHASE_STEPS,
              encoding_dim=ENCODING_DIM,
              optimizer=tf.compat.v1.train.AdamOptimizer(
                  learning_rate=LR
              ),
              alpha=1.0,
              gamma=1.0,
              epsilon_greedy=EPSILON,
              accepts_per_arm_features=True,
              debug_summaries=True,
              summarize_grads_and_vars=True,
              emit_policy_info=policy_utilities.InfoFields.PREDICTED_REWARDS_MEAN,
          )
  
      def _all_rewards(observation, hidden_param):
          """Outputs rewards for all actions, given an observation."""
          hidden_param = tf.cast(hidden_param, dtype=tf.float32)
          global_obs = observation[bandit_spec_utils.GLOBAL_FEATURE_KEY]
          per_arm_obs = observation[
              bandit_spec_utils.PER_ARM_FEATURE_KEY
          ]
          num_actions = tf.shape(per_arm_obs)[1]
          tiled_global = tf.tile(
              tf.expand_dims(global_obs, axis=1), [1, num_actions, 1]
          )
          concatenated = tf.concat([tiled_global, per_arm_obs], axis=-1)
          rewards = tf.linalg.matvec(concatenated, hidden_param)
          return rewards
  
      def optimal_reward(observation, hidden_param):
          return tf.reduce_max(
              _all_rewards(observation, hidden_param), axis=1
          )
  
      def optimal_action(observation, hidden_param):
          return tf.argmax(
              _all_rewards(observation, hidden_param),
              axis=1,
              output_type=tf.int32,
          )
  
      optimal_reward_fn = functools.partial(
          optimal_reward, hidden_param=HIDDEN_PARAM
      )
      optimal_action_fn = functools.partial(
          optimal_action, hidden_param=HIDDEN_PARAM
      )
      regret_metric = tf_bandit_metrics.RegretMetric(optimal_reward_fn)
      suboptimal_arms_metric = tf_bandit_metrics.SuboptimalArmsMetric(
          optimal_action_fn
      )
  
      if FLAGS.drop_arm_obs:
          drop_arm_feature_fn = functools.partial(
              bandit_spec_utils.drop_arm_observation
          )
      else:
          drop_arm_feature_fn = None
      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],
          training_data_spec_transformation_fn=drop_arm_feature_fn,
      )
  
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 numpy
  import tensorflow
  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 import neural_linucb_agent
  from tf_agents.bandits.agents.examples.v2 import trainer
  from tf_agents.bandits.environments import stationary_stochastic_per_arm_py_environment
  from tf_agents.bandits.metrics import tf_metrics
  from tf_agents.bandits.networks import global_and_arm_feature_network
  from tf_agents.bandits.specs import utils
  from tf_agents.environments import tf_py_environment
  from tf_agents.policies import utils
label_desc: Use image node_id 1 for calling a global function with example usage: main(unused_argv) and returns: np, np, rewards, tf, tf, np