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
          action_reward_fns = (
              environment_utilities.sliding_linear_reward_fn_generator(
                  CONTEXT_DIM, NUM_ACTIONS, REWARD_NOISE_VARIANCE
              )
          )
  
          env = sspe.StationaryStochasticPyEnvironment(
              functools.partial(
                  environment_utilities.context_sampling_fn,
                  batch_size=BATCH_SIZE,
                  context_dim=CONTEXT_DIM,
              ),
              action_reward_fns,
              batch_size=BATCH_SIZE,
          )
          environment = tf_py_environment.TFPyEnvironment(env)
  
          optimal_reward_fn = functools.partial(
              environment_utilities.tf_compute_optimal_reward,
              per_action_reward_fns=action_reward_fns,
          )
  
          optimal_action_fn = functools.partial(
              environment_utilities.tf_compute_optimal_action,
              per_action_reward_fns=action_reward_fns,
          )
  
          q_net = q_network.QNetwork(
              environment.observation_spec(),
              environment.action_spec(),
              fc_layer_params=(50, 50),
          )
  
          agent = dqn_agent.DqnAgent(
              environment.time_step_spec(),
              environment.action_spec(),
              q_network=q_net,
              epsilon_greedy=0.1,
              target_update_tau=0.05,
              target_update_period=5,
              optimizer=tf.compat.v1.train.AdamOptimizer(
                  learning_rate=1e-2
              ),
              td_errors_loss_fn=common.element_wise_squared_loss,
          )
  
          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_dqn.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.agents.dqn import dqn_agent
  from tf_agents.bandits.agents.examples.v2 import trainer
  from tf_agents.bandits.environments import environment_utilities
  from tf_agents.bandits.environments import stationary_stochastic_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.utils import common
label_desc: Use image node_id 1 for calling a global function with example usage: main(unused_argv) without return types