//  [python]
def _loss(self, experience, weights=None, training=False):
    """Computes loss for QtOpt training.

    Args:
      experience: A batch of experience data in the form of a `Trajectory` or
        `Transition`. The structure of `experience` must match that of
        `self.collect_policy.step_spec`.  If a `Trajectory`, all tensors in
        `experience` must be shaped `[B, T, ...]` where `T` must be equal to
        `self.train_sequence_length` if that property is not `None`.
      weights: Optional scalar or elementwise (per-batch-entry) importance
        weights.  The output td_loss will be scaled by these weights, and the
        final scalar loss is the mean of these values.
      training: Whether this loss is being used for training.

    Returns:
      loss: An instance of `QtOptLossInfo`.
    Raises:
      ValueError:
        if the number of actions is greater than 1.
    """
    transition = self._as_transition(experience)
    time_steps, policy_steps, next_time_steps = transition
    actions = policy_steps.action

    with tf.name_scope('loss'):
      q_values = self._compute_q_values(
          time_steps, actions, policy_steps.state, training=training
      )

      next_q_values = self._compute_next_q_values(
          next_time_steps, policy_steps.info, policy_steps.state
      )

      # This applies to any value of n_step_update and also in RNN-QtOpt.
      # In the RNN-QtOpt case, inputs and outputs contain a time dimension.
      td_targets = compute_td_targets(
          next_q_values,
          rewards=self._reward_scale_factor * next_time_steps.reward,
          discounts=self._gamma * next_time_steps.discount,
      )

      valid_mask = tf.cast(~time_steps.is_last(), tf.float32)
      td_error = valid_mask * (td_targets - q_values)

      td_loss = valid_mask * self._td_errors_loss_fn(td_targets, q_values)

      if nest_utils.is_batched_nested_tensors(
          time_steps, self.time_step_spec, num_outer_dims=2
      ):
        # Do a sum over the time dimension.
        td_loss = tf.reduce_sum(input_tensor=td_loss, axis=1)

      # Aggregate across the elements of the batch and add regularization loss.
      # Note: We use an element wise loss above to ensure each element is always
      #   weighted by 1/N where N is the batch size, even when some of the
      #   weights are zero due to boundary transitions. Weighting by 1/K where K
      #   is the actual number of non-zero weight would artificially increase
      #   their contribution in the loss. Think about what would happen as
      #   the number of boundary samples increases.

      agg_loss = common.aggregate_losses(
          per_example_loss=td_loss,
          sample_weight=weights,
          regularization_loss=self._q_network.losses,
      )
      total_loss = agg_loss.total_loss

      losses_dict = {
          'td_loss': agg_loss.weighted,
          'reg_loss': agg_loss.regularization,
      }
      total_auxiliary_loss = self._add_auxiliary_losses(
          transition, weights, losses_dict
      )
      total_loss += total_auxiliary_loss

      losses_dict['total_loss'] = total_loss

      common.summarize_scalar_dict(
          losses_dict, step=self.train_step_counter, name_scope='Losses/'
      )

      if self._summarize_grads_and_vars:
        with tf.name_scope('Variables/'):
          for var in self._q_network.trainable_weights:
            tf.compat.v2.summary.histogram(
                name=var.name.replace(':', '_'),
                data=var,
                step=self.train_step_counter,
            )
      if self._debug_summaries:
        diff_q_values = q_values - next_q_values
        common.generate_tensor_summaries(
            'td_error', td_error, self.train_step_counter
        )
        common.generate_tensor_summaries(
            'q_values', q_values, self.train_step_counter
        )
        common.generate_tensor_summaries(
            'next_q_values', next_q_values, self.train_step_counter
        )
        common.generate_tensor_summaries(
            'diff_q_values', diff_q_values, self.train_step_counter
        )
        common.generate_tensor_summaries(
            'reward', next_time_steps.reward, self.train_step_counter
        )

      return tf_agent.LossInfo(
          total_loss, QtOptLossInfo(td_loss=td_loss, td_error=td_error)
      )