class_docstr: An agent implementing the LinUCB algorithm on top of a neural network.
docstr: Computes loss for reward prediction training.
  
  Args:
    observations: A batch of observations.
    actions: A batch of actions.
    rewards: A batch of rewards.
    weights: Optional scalar or elementwise (per-batch-entry) importance
      weights.  The output batch loss will be scaled by these weights, and the
      final scalar loss is the mean of these values.
    training: Whether the loss is being used for training.
  
  Returns:
    loss: A `LossInfo` containing the loss for the training step.
docstr_tok: ["Computes","loss","for","reward","prediction","training",".","Args",":","observations",":","A","batch","of","observations",".","actions",":","A","batch","of","actions",".","rewards",":","A","batch","of","rewards",".","weights",":","Optional","scalar","or","elementwise","(","per-batch-entry",")","importance","weights",".","The","output","batch","loss","will","be","scaled","by","these","weights",",","and","the","final","scalar","loss","is","the","mean","of","these","values",".","training",":","Whether","the","loss","is","being","used","for","training",".","Returns",":","loss",":","A","`","LossInfo","`","containing","the","loss","for","the","training","step","."]
code: def _loss_using_reward_layer(
      self,
      observations: types.NestedTensor,
      actions: types.Tensor,
      rewards: types.Tensor,
      weights: Optional[types.Float] = None,
      training: bool = False,
  ) -> tf_agent.LossInfo:
      """Computes loss for reward prediction training.
  
      Args:
        observations: A batch of observations.
        actions: A batch of actions.
        rewards: A batch of rewards.
        weights: Optional scalar or elementwise (per-batch-entry) importance
          weights.  The output batch loss will be scaled by these weights, and the
          final scalar loss is the mean of these values.
        training: Whether the loss is being used for training.
  
      Returns:
        loss: A `LossInfo` containing the loss for the training step.
      """
      with tf.name_scope("loss"):
          encoded_observation, _ = self._encoding_network(
              observations, training=training
          )
          encoded_observation = tf.reshape(
              encoded_observation, shape=[-1, self._encoding_dim]
          )
          predicted_rewards = self._reward_layer(
              encoded_observation, training=training
          )
          chosen_actions_predicted_rewards = common.index_with_actions(
              predicted_rewards, tf.cast(actions, dtype=tf.int32)
          )
  
          loss = self._error_loss_fn(
              rewards,
              chosen_actions_predicted_rewards,
              1 if weights is None else weights,
          )
          if self._summarize_grads_and_vars:
              with tf.name_scope("Per_arm_loss/"):
                  for k in range(self._num_models):
                      loss_mask_for_arm = tf.cast(
                          tf.equal(actions, k), tf.float32
                      )
                      loss_for_arm = self._error_loss_fn(
                          rewards,
                          chosen_actions_predicted_rewards,
                          weights=loss_mask_for_arm,
                      )
                      tf.compat.v2.summary.scalar(
                          name="loss_arm_" + str(k),
                          data=loss_for_arm,
                          step=self.train_step_counter,
                      )
  
      return tf_agent.LossInfo(loss, extra=())
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Optional, Sequence, Text
  import gin
  import tensorflow
  from tf_agents.agents import data_converter
  from tf_agents.agents import tf_agent
  from tf_agents.bandits.agents import linear_bandit_agent
  from tf_agents.bandits.agents import utils
  from tf_agents.bandits.policies import neural_linucb_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 12 for calling the NeuralLinUCBAgent obj's underlying member method code with example usage: obj._loss_using_reward_layer(observations, actions, rewards, weights, training) and returns: tf_agent