class_docstr: A CQL-SAC Agent based on the SAC Agent.
docstr: Computes CQL loss for SAC training in continuous action spaces.
  
  Extends the standard critic loss to minimize Q-values sampled from a policy
  and maximize values of the dataset actions.
  
  Based on the `CQL(H)` equation (4) in (Kumar et al., 2020):
  
  ```
  log_sum_exp(Q(s, a')) - Q(s, a)
  ```
  
  Other variants of CQL-SAC, such as `CQL(R)`, can override this function.
  
  Args:
    time_steps: A batch of timesteps.
    actions: A batch of actions.
    training: Whether training should be applied.
  
  Returns:
    cql_loss: A scalar CQL loss.
docstr_tok: ["Computes","CQL","loss","for","SAC","training","in","continuous","action","spaces",".","Extends","the","standard","critic","loss","to","minimize","Q-values","sampled","from","a","policy","and","maximize","values","of","the","dataset","actions",".","Based","on","the","`","CQL","(","H",")","`","equation","(","4",")","in","(","Kumar","et","al.",",","2020",")",":","``","`","log_sum_exp","(","Q","(","s",",","a","'",")",")","-","Q","(","s",",","a",")","``","`","Other","variants","of","CQL-SAC",",","such","as","`","CQL","(","R",")","`",",","can","override","this","function",".","Args",":","time_steps",":","A","batch","of","timesteps",".","actions",":","A","batch","of","actions",".","training",":","Whether","training","should","be","applied",".","Returns",":","cql_loss",":","A","scalar","CQL","loss","."]
code: def _cql_loss(
      self,
      time_steps: ts.TimeStep,
      actions: types.Tensor,
      training: Optional[bool] = False,
  ) -> types.Tensor:
      """Computes CQL loss for SAC training in continuous action spaces.
  
      Extends the standard critic loss to minimize Q-values sampled from a policy
      and maximize values of the dataset actions.
  
      Based on the `CQL(H)` equation (4) in (Kumar et al., 2020):
  
      ```
      log_sum_exp(Q(s, a')) - Q(s, a)
      ```
  
      Other variants of CQL-SAC, such as `CQL(R)`, can override this function.
  
      Args:
        time_steps: A batch of timesteps.
        actions: A batch of actions.
        training: Whether training should be applied.
  
      Returns:
        cql_loss: A scalar CQL loss.
      """
      # Tile tensors to match CQL action sampling.
      # Shape [batch_size * self._num_cql_samples, ...]
      observation = nest_utils.tile_batch(
          time_steps.observation, self._num_cql_samples
      )
      step_type = nest_utils.tile_batch(
          time_steps.step_type, self._num_cql_samples
      )
  
      # Sample self._num_cql_samples from the policy distribution.
      # We do not update actor during CQL loss.
      (
          sampled_actions,
          sampled_actions_log_probs,
      ) = self._sample_and_transpose_actions_and_log_probs(
          time_steps, self._num_cql_samples, training=False
      )
      # Shape [batch_size * self._num_cql_samples, ...]
      sampled_actions = self._flattened_multibatch_tensor(
          sampled_actions
      )
      target_input = (observation, sampled_actions)
      batch_size = nest_utils.get_outer_shape(
          time_steps, self._time_step_spec
      )[0]
      q_estimates1, q_estimates2 = self._get_q_values(
          target_input,
          step_type,
          reshape_batch_size=batch_size,
          training=False,
      )
      debug_summaries_dict = {}
      debug_summaries_dict["q_estimates1"] = tf.reduce_mean(
          q_estimates1
      )
      debug_summaries_dict["q_estimates2"] = tf.reduce_mean(
          q_estimates2
      )
  
      # We're supposed to be taking an unweighted sum of Q-values of actions
      # from the policy and uniform distributions. We correct for the fact that
      # we're drawing samples from either distribution by subtracting the
      # corresponding log probability.
      # Based on Appendix F:
      # ```
      # q_value_unif = \sum_{a_i~\Unif(a)}^{N} [exp(Q(s, a_i)) / \Unif(a)]
      # q_value_pi =   \sum_{a_i~\pi(a|s)}^{N} [exp(Q(s, a_i)) / \pi(a|s)]
      # log_sum_exp(Q(s, a')) = log((1/2N * q_value_unif) + (1/2N  * q_value_pi))
      # ```
      policy_log_probs1 = (
          q_estimates1 * self._softmax_temperature
      ) - sampled_actions_log_probs[..., None]
      policy_log_probs2 = (
          q_estimates2 * self._softmax_temperature
      ) - sampled_actions_log_probs[..., None]
  
      # Sample self._num_cql_samples from the uniform-at-random distribution.
      flattened_action_spec = tf.nest.flatten(self.action_spec)[0]
      uniform_actions = tf.random.uniform(
          tf.shape(sampled_actions),
          minval=flattened_action_spec.minimum,
          maxval=flattened_action_spec.maximum,
          seed=self._action_seed_stream(),
      )
      target_input = (observation, uniform_actions)
      q_uniform1, q_uniform2 = self._get_q_values(
          target_input, step_type, batch_size, training=False
      )
      debug_summaries_dict["q_uniform1"] = tf.reduce_mean(q_uniform1)
      debug_summaries_dict["q_uniform2"] = tf.reduce_mean(q_uniform2)
  
      # Uniform density is `(1/range)^dimension`, so the log probability of the
      # uniform distribution is `-log(range)*dimension`.
      # Once again, we subtract this from the Q-value to correct for drawing
      # from this distribution and contribute to an unweighted sum of Q-values.
      uniform_actions_log_probs = (
          tf.reduce_sum(
              -tf.math.log(
                  flattened_action_spec.maximum
                  - flattened_action_spec.minimum
              )
          )
          * flattened_action_spec.shape[0]
      )
      uniform_log_probs1 = (
          q_uniform1 * self._softmax_temperature
      ) - uniform_actions_log_probs
      uniform_log_probs2 = (
          q_uniform2 * self._softmax_temperature
      ) - uniform_actions_log_probs
  
      # Importance sampled estimate of the Q-value sum. We do this since we
      # can't tractably compute the exact Q-value in a continuous action space.
      # Based on the first part of equation (4):
      # ```log_sum_exp(Q(s, a'))```
      # Concatenates and collapses along the self._num_cql_samples dimension.
      combined_log_probs1 = tf.concat(
          [policy_log_probs1, uniform_log_probs1], axis=1
      )
      combined_log_probs2 = tf.concat(
          [policy_log_probs2, uniform_log_probs2], axis=1
      )
  
      logsumexp1 = (
          tf.math.reduce_logsumexp(combined_log_probs1, axis=1)
          * 1.0
          / self._softmax_temperature
      )
      logsumexp2 = (
          tf.math.reduce_logsumexp(combined_log_probs2, axis=1)
          * 1.0
          / self._softmax_temperature
      )
  
      target_input = (time_steps.observation, actions)
      q_original1, q_original2 = self._get_q_values(
          target_input,
          time_steps.step_type,
          reshape_batch_size=None,
          training=training,
      )
      debug_summaries_dict["q_original1"] = tf.reduce_mean(q_original1)
      debug_summaries_dict["q_original2"] = tf.reduce_mean(q_original2)
  
      cql_loss = (
          tf.reduce_mean(
              (logsumexp1 - q_original1) + (logsumexp2 - q_original2)
          )
          / 2.0
      )
  
      self._cql_loss_debug_summaries(debug_summaries_dict)
  
      return cql_loss
  
imports: from typing import Callable, Dict, NamedTuple, Optional, Text, Tuple, Union
  import gin
  import tensorflow
  import tensorflow_probability
  from tf_agents.agents import tf_agent
  from tf_agents.agents.sac import sac_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 nest_utils
  from tf_agents.utils import object_identity
label_desc: Use image node_id 10 for calling the CqlSacAgent obj's underlying member method code with example usage: obj._cql_loss(time_steps, actions, training) and returns: cql_loss