class_docstr: An agent implementing the LinUCB algorithm on top of a neural network.
docstr: Computes the loss using LinUCB.
  
  Args:
    observation: A batch of observations.
    action: A batch of actions.
    reward: A batch of rewards.
    weights: unused weights.
    training: Whether the loss is being used to train.
  
  Returns:
    loss: A `LossInfo` containing the loss for the training step.
docstr_tok: ["Computes","the","loss","using","LinUCB",".","Args",":","observation",":","A","batch","of","observations",".","action",":","A","batch","of","actions",".","reward",":","A","batch","of","rewards",".","weights",":","unused","weights",".","training",":","Whether","the","loss","is","being","used","to","train",".","Returns",":","loss",":","A","`","LossInfo","`","containing","the","loss","for","the","training","step","."]
code: def compute_loss_using_linucb(
      self,
      observation: types.NestedTensor,
      action: types.Tensor,
      reward: types.Tensor,
      weights: Optional[types.Float] = None,
      training: bool = False,
  ) -> tf_agent.LossInfo:
      """Computes the loss using LinUCB.
  
      Args:
        observation: A batch of observations.
        action: A batch of actions.
        reward: A batch of rewards.
        weights: unused weights.
        training: Whether the loss is being used to train.
  
      Returns:
        loss: A `LossInfo` containing the loss for the training step.
      """
      del weights  # unused
  
      # The network is trained now. Update the covariance matrix.
      encoded_observation, _ = self._encoding_network(
          observation, training=training
      )
      encoded_observation = tf.cast(
          encoded_observation, dtype=self._dtype
      )
      encoded_observation = tf.reshape(
          encoded_observation, shape=[-1, self._encoding_dim]
      )
      for k in range(self._num_models):
          diag_mask = tf.linalg.tensor_diag(
              tf.cast(tf.equal(action, k), self._dtype)
          )
          observations_for_arm = tf.matmul(
              diag_mask, encoded_observation
          )
          rewards_for_arm = tf.matmul(
              diag_mask, tf.reshape(reward, [-1, 1])
          )
  
          num_samples_for_arm_current = tf.reduce_sum(diag_mask)
          tf.compat.v1.assign_add(
              self.num_samples[k], num_samples_for_arm_current
          )
          num_samples_for_arm_total = self.num_samples[k].read_value()
  
          # Update the matrix A and b.
          # pylint: disable=cell-var-from-loop
          def update(cov_matrix, data_vector):
              (
                  a_new,
                  b_new,
              ) = linear_agent.update_a_and_b_with_forgetting(
                  cov_matrix,
                  data_vector,
                  rewards_for_arm,
                  observations_for_arm,
                  self._gamma,
              )
              return a_new, b_new
  
          a_new, b_new = tf.cond(
              tf.squeeze(num_samples_for_arm_total) > 0,
              lambda: update(self.cov_matrix[k], self.data_vector[k]),
              lambda: (self.cov_matrix[k], self.data_vector[k]),
          )
          tf.compat.v1.assign(self.cov_matrix[k], a_new)
          tf.compat.v1.assign(self.data_vector[k], b_new)
  
      loss_tensor = tf.cast(
          -1.0 * tf.reduce_sum(reward), dtype=tf.float32
      )
      loss_info = tf_agent.LossInfo(loss=loss_tensor, extra=())
      tf.compat.v2.summary.scalar(
          name="using_reward_layer",
          data=0,
          step=self.train_step_counter,
      )
      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 14 for calling the NeuralLinUCBAgent obj's underlying member method code with example usage: obj.compute_loss_using_linucb(observation, action, reward, weights, training) and returns: loss_info, a_new,  b_new