class_docstr: A SAC Agent.
class_docstr_tok: ["A","SAC","Agent","."]
docstr: Computes the critic loss for SAC training.
  
  Args:
    time_steps: A batch of timesteps.
    actions: A batch of actions.
    next_time_steps: A batch of next timesteps.
    td_errors_loss_fn: A function(td_targets, predictions) to compute
      elementwise (per-batch-entry) loss.
    gamma: Discount for future rewards.
    reward_scale_factor: Multiplicative factor to scale rewards.
    weights: Optional scalar or elementwise (per-batch-entry) importance
      weights.
    training: Whether this loss is being used for training.
  
  Returns:
    critic_loss: A scalar critic loss.
docstr_tok: ["Computes","the","critic","loss","for","SAC","training",".","Args",":","time_steps",":","A","batch","of","timesteps",".","actions",":","A","batch","of","actions",".","next_time_steps",":","A","batch","of","next","timesteps",".","td_errors_loss_fn",":","A","function","(","td_targets",",","predictions",")","to","compute","elementwise","(","per-batch-entry",")","loss",".","gamma",":","Discount","for","future","rewards",".","reward_scale_factor",":","Multiplicative","factor","to","scale","rewards",".","weights",":","Optional","scalar","or","elementwise","(","per-batch-entry",")","importance","weights",".","training",":","Whether","this","loss","is","being","used","for","training",".","Returns",":","critic_loss",":","A","scalar","critic","loss","."]
code: def critic_loss(
      self,
      time_steps: ts.TimeStep,
      actions: types.Tensor,
      next_time_steps: ts.TimeStep,
      td_errors_loss_fn: types.LossFn,
      gamma: types.Float = 1.0,
      reward_scale_factor: types.Float = 1.0,
      weights: Optional[types.Tensor] = None,
      training: bool = False,
  ) -> types.Tensor:
      """Computes the critic loss for SAC training.
  
      Args:
        time_steps: A batch of timesteps.
        actions: A batch of actions.
        next_time_steps: A batch of next timesteps.
        td_errors_loss_fn: A function(td_targets, predictions) to compute
          elementwise (per-batch-entry) loss.
        gamma: Discount for future rewards.
        reward_scale_factor: Multiplicative factor to scale rewards.
        weights: Optional scalar or elementwise (per-batch-entry) importance
          weights.
        training: Whether this loss is being used for training.
  
      Returns:
        critic_loss: A scalar critic loss.
      """
      with tf.name_scope("critic_loss"):
          nest_utils.assert_same_structure(actions, self.action_spec)
          nest_utils.assert_same_structure(
              time_steps, self.time_step_spec
          )
          nest_utils.assert_same_structure(
              next_time_steps, self.time_step_spec
          )
  
          # We do not update actor or target networks in critic loss.
          next_actions, next_log_pis = self._actions_and_log_probs(
              next_time_steps, training=False
          )
          target_input = (next_time_steps.observation, next_actions)
          (
              target_q_values1,
              unused_network_state1,
          ) = self._target_critic_network_1(
              target_input,
              step_type=next_time_steps.step_type,
              training=False,
          )
          (
              target_q_values2,
              unused_network_state2,
          ) = self._target_critic_network_2(
              target_input,
              step_type=next_time_steps.step_type,
              training=False,
          )
          target_q_values = (
              tf.minimum(target_q_values1, target_q_values2)
              - tf.exp(self._log_alpha) * next_log_pis
          )
  
          td_targets = tf.stop_gradient(
              reward_scale_factor * next_time_steps.reward
              + gamma * next_time_steps.discount * target_q_values
          )
  
          pred_input = (time_steps.observation, actions)
          pred_td_targets1, _ = self._critic_network_1(
              pred_input,
              step_type=time_steps.step_type,
              training=training,
          )
          pred_td_targets2, _ = self._critic_network_2(
              pred_input,
              step_type=time_steps.step_type,
              training=training,
          )
          critic_loss1 = td_errors_loss_fn(td_targets, pred_td_targets1)
          critic_loss2 = td_errors_loss_fn(td_targets, pred_td_targets2)
          critic_loss = critic_loss1 + critic_loss2
  
          if critic_loss.shape.rank > 1:
              # Sum over the time dimension.
              critic_loss = tf.reduce_sum(
                  critic_loss, axis=range(1, critic_loss.shape.rank)
              )
  
          agg_loss = common.aggregate_losses(
              per_example_loss=critic_loss,
              sample_weight=weights,
              regularization_loss=(
                  self._critic_network_1.losses
                  + self._critic_network_2.losses
              ),
          )
          critic_loss = agg_loss.total_loss
  
          self._critic_loss_debug_summaries(
              td_targets, pred_td_targets1, pred_td_targets2
          )
  
          return critic_loss
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  from typing import Callable, Optional, Text
  import gin
  import numpy
  from six.moves import zip
  import tensorflow
  import tensorflow_probability
  from tf_agents.agents import data_converter
  from tf_agents.agents import tf_agent
  from tf_agents.networks import network
  from tf_agents.policies import actor_policy
  from tf_agents.policies import tf_policy
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
  from tf_agents.utils import common
  from tf_agents.utils import eager_utils
  from tf_agents.utils import nest_utils
  from tf_agents.utils import object_identity
label_desc: Use image node_id 11 for calling the SacAgent obj's underlying member method code with example usage: obj.critic_loss(time_steps, actions, next_time_steps, td_errors_loss_fn, gamma, reward_scale_factor, weights, training) and returns: critic_loss