class_docstr: A neural reward network based bandit agent.
  
  This agent receives a neural network that it trains to predict rewards. The
  action is chosen greedily with respect to the prediction.
class_docstr_tok: ["A","neural","reward","network","based","bandit","agent",".","This","agent","receives","a","neural","network","that","it","trains","to","predict","rewards",".","The","action","is","chosen","greedily","with","respect","to","the","prediction","."]
code: def _train(self, experience, weights):
      experience = self._as_trajectory(experience)
  
      with tf.GradientTape() as tape:
          loss_info = self._loss(
              experience, weights=weights, training=True
          )
  
      variables_to_train = self._variables_to_train()
      if not variables_to_train:
          logging.info("No variable to train in the agent.")
          return loss_info
  
      grads = tape.gradient(loss_info.loss, variables_to_train)
      # Tuple is used for py3, where zip is a generator producing values once.
      grads_and_vars = tuple(zip(grads, variables_to_train))
      if self._gradient_clipping is not None:
          grads_and_vars = eager_utils.clip_gradient_norms(
              grads_and_vars, self._gradient_clipping
          )
  
      if self._summarize_grads_and_vars:
          eager_utils.add_variables_summaries(
              grads_and_vars, self.train_step_counter
          )
          eager_utils.add_gradients_summaries(
              grads_and_vars, self.train_step_counter
          )
  
      self._optimizer.apply_gradients(grads_and_vars)
      self.train_step_counter.assign_add(1)
      if self._num_samples_list:
          actions_flattened = tf.reshape(experience.action, [-1])
          if self._accepts_per_arm_features:
              num_samples_per_action_current = [
                  tf.cast(
                      tf.shape(actions_flattened)[0], dtype=tf.int64
                  )
              ]
          else:
              # Compute the number of samples for each action in the current batch.
              num_samples_per_action_current = [
                  tf.reduce_sum(
                      tf.cast(tf.equal(actions_flattened, k), tf.int64)
                  )
                  for k in range(self._num_actions)
              ]
          # Update the number of samples for each action.
          for a, b in zip(
              self._num_samples_list, num_samples_per_action_current
          ):
              tf.compat.v1.assign_add(a, b)
  
      return loss_info
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Iterable, Optional, Sequence, Text, Tuple
  from absl import logging
  import gin
  import tensorflow
  from tf_agents.agents import data_converter
  from tf_agents.agents import tf_agent
  from tf_agents.bandits.agents import utils
  from tf_agents.bandits.networks import heteroscedastic_q_network
  from tf_agents.bandits.policies import constraints
  from tf_agents.bandits.policies import greedy_reward_prediction_policy
  from tf_agents.bandits.specs import utils
  from tf_agents.policies import utils
  from tf_agents.typing import types
  from tf_agents.utils import common
  from tf_agents.utils import eager_utils
label_desc: Use image node_id 5 for calling the GreedyRewardPredictionAgent obj's underlying member method code with example usage: obj._train(experience, weights) and returns: loss_info, loss_info