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 training the reward and constraint networks.
  
  Args:
    experience: A batch of experience data in the form of a `Trajectory` or
      `Transition`.
    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.
  Raises:
    ValueError:
      if the number of actions is greater than 1.
docstr_tok: ["Computes","loss","for","training","the","reward","and","constraint","networks",".","Args",":","experience",":","A","batch","of","experience","data","in","the","form","of","a","`","Trajectory","`","or","`","Transition","`",".","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",".","Raises",":","ValueError",":","if","the","number","of","actions","is","greater","than","1","."]
code: def _loss(
      self,
      experience: types.NestedTensor,
      weights: Optional[types.Float] = None,
      training: bool = False,
  ) -> tf_agent.LossInfo:
      """Computes loss for training the reward and constraint networks.
  
      Args:
        experience: A batch of experience data in the form of a `Trajectory` or
          `Transition`.
        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.
      Raises:
        ValueError:
          if the number of actions is greater than 1.
      """
      (
          observations,
          actions,
          rewards,
      ) = bandit_utils.process_experience_for_neural_agents(
          experience,
          self._accepts_per_arm_features,
          self.training_data_spec,
      )
      if self._observation_and_action_constraint_splitter is not None:
          (
              observations,
              _,
          ) = self._observation_and_action_constraint_splitter(
              observations
          )
  
      if self._constraints:
          rewards_tensor = rewards[bandit_spec_utils.REWARD_SPEC_KEY]
      else:
          rewards_tensor = rewards
      reward_loss = self.reward_loss(
          observations, actions, rewards_tensor, weights, training
      )
  
      constraint_loss = tf.constant(0.0)
      for i, c in enumerate(self._constraints, 0):
          constraint_targets = rewards[
              bandit_spec_utils.CONSTRAINTS_SPEC_KEY
          ][:, i]
          constraint_targets = tf.reshape(constraint_targets, [-1])
          constraint_loss += c.compute_loss(
              observations,
              actions,
              constraint_targets,
              weights,
              training,
          )
  
      self.compute_summaries(
          reward_loss,
          constraint_loss=(
              constraint_loss if self._constraints else None
          ),
      )
  
      total_loss = reward_loss
      if self._constraints:
          total_loss += constraint_loss
      return tf_agent.LossInfo(total_loss, extra=())
  
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 7 for calling the GreedyRewardPredictionAgent obj's underlying member method code with example usage: obj._loss(experience, weights, training) and returns: tf_agent