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","."]
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 `Tensor` containing the loss for the training step.
  Raises:
    ValueError:
      if the number of actions is greater than 1.
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","`","Tensor","`","containing","the","loss","for","the","training","step",".","Raises",":","ValueError",":","if","the","number","of","actions","is","greater","than","1","."]
code: def reward_loss(
      self,
      observations: types.NestedTensor,
      actions: types.Tensor,
      rewards: types.Tensor,
      weights: Optional[types.Float] = None,
      training: bool = False,
  ) -> types.Tensor:
      """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 `Tensor` containing the loss for the training step.
      Raises:
        ValueError:
          if the number of actions is greater than 1.
      """
      with tf.name_scope("loss"):
          sample_weights = weights if weights is not None else 1
          if self._heteroscedastic:
              predictions, _ = self._reward_network(
                  observations, training=training
              )
              predicted_values = predictions.q_value_logits
              predicted_log_variance = predictions.log_variance
              action_predicted_log_variance = common.index_with_actions(
                  predicted_log_variance,
                  tf.cast(actions, dtype=tf.int32),
              )
              sample_weights = (
                  sample_weights
                  * 0.5
                  * tf.exp(-action_predicted_log_variance)
              )
  
              loss = 0.5 * tf.reduce_mean(action_predicted_log_variance)
              # loss = 1/(2 * var(x)) * (y - f(x))^2 + 1/2 * log var(x)
              # Kendall, Alex, and Yarin Gal. "What Uncertainties Do We Need in
              # Bayesian Deep Learning for Computer Vision?." Advances in Neural
              # Information Processing Systems. 2017. https://arxiv.org/abs/1703.04977
          else:
              predicted_values, _ = self._reward_network(
                  observations, training=training
              )
              loss = tf.constant(0.0)
  
          action_predicted_values = common.index_with_actions(
              predicted_values, tf.cast(actions, dtype=tf.int32)
          )
  
          # Apply Laplacian smoothing on the estimated rewards, if applicable.
          if self._laplacian_matrix is not None:
              smoothness_batched = tf.matmul(
                  predicted_values,
                  tf.matmul(
                      self._laplacian_matrix,
                      predicted_values,
                      transpose_b=True,
                  ),
              )
              loss += self._laplacian_smoothing_weight * tf.reduce_mean(
                  tf.linalg.tensor_diag_part(smoothness_batched)
                  * sample_weights
              )
  
          # Reduction is done outside of the loss function because non-scalar
          # weights with unknown shapes may trigger shape validation that fails
          # XLA compilation.
          loss += tf.reduce_mean(
              tf.multiply(
                  self._error_loss_fn(
                      rewards,
                      action_predicted_values,
                      reduction=tf.compat.v1.losses.Reduction.NONE,
                  ),
                  sample_weights,
              )
          )
  
      return loss
  
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 6 for calling the GreedyRewardPredictionAgent obj's underlying member method code with example usage: obj.reward_loss(observations, actions, rewards, weights, training) and returns: loss