class_docstr: An agent implementing the LinUCB algorithm on top of a neural network.
docstr: Computes loss using the reward layer.
  
  Args:
    observation: A batch of observations.
    action: A batch of actions.
    reward: 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","using","the","reward","layer",".","Args",":","observation",":","A","batch","of","observations",".","action",":","A","batch","of","actions",".","reward",":","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 compute_loss_using_reward_layer(
      self,
      observation: types.NestedTensor,
      action: types.Tensor,
      reward: types.Tensor,
      weights: Optional[types.Float] = None,
      training: bool = False,
  ) -> tf_agent.LossInfo:
      """Computes loss using the reward layer.
  
      Args:
        observation: A batch of observations.
        action: A batch of actions.
        reward: 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.
      """
      # Update the neural network params.
      with tf.GradientTape() as tape:
          loss_info = self._loss_using_reward_layer(
              observation, action, reward, weights, training=training
          )
      tf.debugging.check_numerics(loss_info[0], "Loss is inf or nan")
      tf.compat.v2.summary.scalar(
          name="using_reward_layer",
          data=1,
          step=self.train_step_counter,
      )
      if self._summarize_grads_and_vars:
          self.compute_summaries(loss_info.loss)
      variables_to_train = (
          self._encoding_network.trainable_weights
          + self._reward_layer.trainable_weights
      )
      if not variables_to_train:
          raise ValueError("No variable to train in the agent.")
  
      grads = tape.gradient(loss_info.loss, variables_to_train)
      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:
          with tf.name_scope("Reward_network/"):
              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)
  
      return loss_info
  
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 13 for calling the NeuralLinUCBAgent obj's underlying member method code with example usage: obj.compute_loss_using_reward_layer(observation, action, reward, weights, training) and returns: loss_info