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: The policy losses as well a few other metrics needed for the entropy loss computation and stats.
docstr_tok: ["Compute","the","critic","losses",".",":","param","worker_output",":","The","batched","output","of","the","workers",".",":","return",":","The","policy","losses","as","well","a","few","other","metrics","needed","for","the","entropy","loss","computation","and","stats","."]
code: def _compute_policy_loss(
      self, worker_output: StructuredSpacesRecord
  ) -> Tuple[
      Dict[Union[str, int], torch.Tensor],
      Dict[
          Union[str, int], Union[torch.Tensor, Dict[str, torch.Tensor]]
      ],
      Dict[
          Union[str, int], Union[torch.Tensor, Dict[str, torch.Tensor]]
      ],
      Dict[Union[str, int], torch.Tensor],
  ]:
      """Compute the critic losses.
  
      :param worker_output: The batched output of the workers.
      :return: The policy losses as well a few other metrics needed for the entropy loss computation and stats.
      """
  
      # Sample actions and compute action log probabilities (continuous steps)/ action probabilities (discrete steps)
      (
          policy_losses,
          action_entropies,
          action_log_probs,
          actions_sampled,
      ) = (dict(), dict(), dict(), dict())
      action_probs = dict()
  
      for step_key in self.sub_step_keys:
          step_obs = worker_output.observations_dict[step_key]
          learner_policy_out = (
              self.learner_model.policy.compute_substep_policy_output(
                  step_obs, ActorID(step_key, 0)
              )
          )
          learner_action = learner_policy_out.prob_dist.sample()
  
          # Average the logp_policy of all actions in this step (all steps if shared critic)
          if self.learner_model.critic.only_discrete_spaces[step_key]:
              probs_policy = {
                  action_key: logits_to_probs(x)
                  for action_key, x in learner_policy_out.action_logits.items()
              }
              logp_policy = {
                  action_key: torch.log(x + (x == 0.0).float() * 1e-8)
                  for action_key, x in probs_policy.items()
              }
          else:
              probs_policy = None
              logp_policy = torch.stack(
                  list(
                      learner_policy_out.prob_dist.log_prob(
                          learner_action
                      ).values()
                  )
              ).mean(dim=0)
  
          action_probs[step_key] = probs_policy
          action_log_probs[step_key] = logp_policy
          actions_sampled[step_key] = learner_action
          action_entropies[step_key] = learner_policy_out.entropy
  
      # Predict Q values
      q_values = self.learner_model.critic.predict_q_values(
          worker_output.observations_dict,
          actions_sampled,
          gather_output=False,
      )
      if len(q_values) < len(self.sub_step_keys):
          assert len(q_values) == 1
          critic_key = list(q_values.keys())[0]
          q_values = {
              step_key: q_values[critic_key]
              for step_key in self.sub_step_keys
          }
  
      # Compute loss
      for step_key in self.sub_step_keys:
          action_log_probs_step = action_log_probs[step_key]
          q_values_step = q_values[step_key]
  
          if self.learner_model.critic.only_discrete_spaces[step_key]:
              action_probs_step = action_probs[step_key]
  
              policy_losses_per_action = list()
              # Compute the policy loss for each individual action
              for action_key in action_log_probs_step.keys():
                  q_action_key = action_key + "_q_values"
                  action_q_values = (
                      torch.stack(
                          [
                              q_values_sub_critic[q_action_key]
                              for q_values_sub_critic in q_values_step
                          ]
                      )
                      .min(dim=0)
                      .values
                  )
                  q_term = (
                      self.curr_entropy_coef[step_key]
                      * action_log_probs_step[action_key]
                      - action_q_values
                  )
                  action_policy_loss = (
                      torch.matmul(
                          action_probs_step[action_key].unsqueeze(-2),
                          q_term.unsqueeze(-1),
                      )
                      .squeeze(-1)
                      .squeeze(-1)
                  )
                  policy_losses_per_action.append(action_policy_loss)
              # Sum the losses of all action together
              policy_losses_per_step = torch.stack(
                  policy_losses_per_action
              ).sum(dim=0)
              # Average the losses w.r.t. to the batch
              policy_losses[step_key] = policy_losses_per_step.mean()
          else:
              # Do not detach q_values in discrete setting
              q_value_per_step = (
                  torch.stack(q_values_step).min(dim=0).values
              )
              # Average the losses w.r.t. to the batch
              policy_losses[step_key] = torch.mean(
                  (
                      self.curr_entropy_coef[step_key]
                      * action_log_probs_step
                      - q_value_per_step
                  )
              )
  
      return (
          policy_losses,
          action_probs,
          action_log_probs,
          action_entropies,
      )
  
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 10 for calling the SAC obj's underlying member method code with example usage: obj._compute_policy_loss(worker_output) and returns: policy_losses,  action_probs,  action_log_probs,  action_entropies