class_docstr: Multi step soft actor critic.
  
  :param algorithm_config: Algorithm options.
  :param learner_model: Structured torch actor critic to train.
  :param distributed_actors: Distributed actors for collection of training rollouts.
  :param model_selection: Optional model selection class, receives model evaluation results.
  :param evaluator: The evaluator to use.
class_docstr_tok: ["Multi","step","soft","actor","critic",".",":","param","algorithm_config",":","Algorithm","options",".",":","param","learner_model",":","Structured","torch","actor","critic","to","train",".",":","param","distributed_actors",":","Distributed","actors","for","collection","of","training","rollouts",".",":","param","model_selection",":","Optional","model","selection","class",",","receives","model","evaluation","results",".",":","param","evaluator",":","The","evaluator","to","use","."]
docstr: Compute the critic losses.
  
  :param worker_output: The batched output of the workers.
  :return: Return the critic losses as a list w.r.t. the number of critics used + mean values for stats.
docstr_tok: ["Compute","the","critic","losses",".",":","param","worker_output",":","The","batched","output","of","the","workers",".",":","return",":","Return","the","critic","losses","as","a","list","w.r.t",".","the","number","of","critics","used","+","mean","values","for","stats","."]
code: def _compute_critic_loss(
      self, worker_output: StructuredSpacesRecord
  ) -> Tuple[
      List[Dict[Union[str, int], torch.Tensor]],
      List[Dict[Union[str, int], torch.Tensor]],
  ]:
      """Compute the critic losses.
  
      :param worker_output: The batched output of the workers.
      :return: Return the critic losses as a list w.r.t. the number of critics used + mean values for stats.
      """
      next_actions = dict()
      next_actions_logits = dict()
      next_action_log_probs = dict()
  
      q_values_selected = self.learner_model.critic.predict_q_values(
          worker_output.observations_dict,
          worker_output.actions_dict,
          gather_output=True,
      )
      q_values_mean = {
          step_key: [
              curr_q.detach().mean().item()
              if isinstance(curr_q, torch.Tensor)
              else torch.stack(list(curr_q.values()))
              .mean(dim=1)
              .detach()
              .mean()
              .item()
              for curr_q in q_values_list
          ]
          for step_key, q_values_list in q_values_selected.items()
      }
  
      with torch.no_grad():
          for step_key in self.sub_step_keys:
              next_policy_output = self.learner_model.policy.compute_substep_policy_output(
                  worker_output.next_observations_dict[step_key],
                  actor_id=ActorID(step_key, 0),
              )
              next_action = next_policy_output.prob_dist.sample()
  
              next_action_log_probs[
                  step_key
              ] = next_policy_output.prob_dist.log_prob(next_action)
              next_actions_logits[
                  step_key
              ] = next_policy_output.action_logits
              next_actions[step_key] = next_action
  
          next_q_values = (
              self.learner_model.critic.predict_next_q_values(
                  worker_output.next_observations_dict,
                  next_actions,
                  next_actions_logits,
                  next_action_log_probs,
                  self.curr_entropy_coef,
              )
          )
  
          target_q_values = dict()
  
          # TODO: Take into account all rewards, not just from the last sub-step
          last_rewards = list(worker_output.rewards_dict.values())[-1]
          last_dones = list(worker_output.dones_dict.values())[-1]
  
          for step_key, next_q_value_per_step in next_q_values.items():
              if self.learner_model.critic.only_discrete_spaces[
                  step_key
              ]:
                  assert isinstance(next_q_value_per_step, dict)
                  target_q_values[step_key] = {
                      action_key: (
                          last_rewards
                          + (~last_dones).float()
                          * self.algorithm_config.gamma
                          * next_action_q_value
                      )
                      for action_key, next_action_q_value in next_q_value_per_step.items()
                  }
              else:
                  assert isinstance(next_q_value_per_step, torch.Tensor)
                  target_q_values[step_key] = (
                      last_rewards
                      + (~last_dones).float()
                      * self.algorithm_config.gamma
                      * next_q_value_per_step
                  )
  
      q_losses = dict()
      for step_key in q_values_selected:
          per_critic_values = list()
          for q_values_per_sub_critic in q_values_selected[step_key]:
              target_q_values_per_step = target_q_values[step_key]
              if self.learner_model.critic.only_discrete_spaces[
                  step_key
              ]:
                  assert isinstance(q_values_per_sub_critic, dict)
                  per_action_per_critic_loss = list()
                  for (
                      action_key,
                      q_values_per_action,
                  ) in q_values_per_sub_critic.items():
                      org_action_key = action_key.replace(
                          "_q_values", ""
                      )
                      per_action_loss = (
                          (
                              q_values_per_action
                              - target_q_values_per_step[org_action_key]
                          )
                          .pow(2)
                          .mean()
                      )
                      per_action_per_critic_loss.append(per_action_loss)
                  # Sum the q_value of individual action in one step together
                  per_critic_values.append(
                      torch.stack(per_action_per_critic_loss).sum(dim=0)
                  )
              else:
                  assert isinstance(
                      q_values_per_sub_critic, torch.Tensor
                  )
                  per_critic_values.append(
                      (
                          q_values_per_sub_critic
                          - target_q_values_per_step
                      )
                      .pow(2)
                      .mean()
                  )
          q_losses[step_key] = per_critic_values
  
      # Transpose list of lists to get into right format and sum values from different steps together (but keep them
      #   separate w.r.t. the q network
      q_losses = [
          dict(zip(q_losses, t)) for t in zip(*q_losses.values())
      ]
      q_values_mean = [
          dict(zip(q_values_mean, t))
          for t in zip(*q_values_mean.values())
      ]
      return q_losses, q_values_mean
  
imports: import sys
  import time
  from typing import Union, Optional, BinaryIO, Dict, Tuple, List
  import numpy
  import torch
  import torch.nn
  from gym import spaces
  from maze.core.agent.torch_actor_critic import TorchActorCritic
  from maze.core.annotations import override
  from maze.core.env.base_env_events import BaseEnvEvents
  from maze.core.env.structured_env import ActorID
  from maze.core.log_stats.log_stats import increment_log_step, LogStatsLevel
  from maze.core.trajectory_recording.records.structured_spaces_record import StructuredSpacesRecord
  from maze.train.trainers.common.evaluators.rollout_evaluator import RolloutEvaluator
  from maze.train.trainers.common.model_selection.best_model_selection import BestModelSelection
  from maze.train.trainers.common.trainer import Trainer
  from maze.train.utils.train_utils import compute_gradient_norm
  from maze.utils.bcolors import BColors
  from maze.train.parallelization.distributed_actors.base_distributed_workers_with_buffer import BaseDistributedWorkersWithBuffer
  from maze.train.trainers.sac.sac_algorithm_config import SACAlgorithmConfig
  from maze.train.trainers.sac.sac_events import SACEvents
  from torch.distributions.utils import logits_to_probs
label_desc: Use image node_id 9 for calling the SAC obj's underlying member method code with example usage: obj._compute_critic_loss(worker_output) and returns: q_losses,  q_values_mean