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: Perform update.
  
  That is collect the actor trajectories, compute the learner action logits, compute the loss
  and backprob it thought the networks.
docstr_tok: ["Perform","update",".","That","is","collect","the","actor","trajectories",",","compute","the","learner","action","logits",",","compute","the","loss","and","backprob","it","thought","the","networks","."]
code: def _update(self) -> None:
      """Perform update.
  
      That is collect the actor trajectories, compute the learner action logits, compute the loss
      and backprob it thought the networks.
      """
  
      start_update_time = time.time()
  
      # Collect batch from buffer ====================================================================================
      worker_output = self.distributed_workers.sample_batch(
          learner_device=self.learner_model.device
      )
      assert isinstance(worker_output, StructuredSpacesRecord)
      after_collection_time = time.time()
  
      # ==============================================================================================================
      # Compute critic losses ========================================================================================
      # ==============================================================================================================
  
      q_losses, q_values_mean = self._compute_critic_loss(worker_output)
  
      # NOTE: SPINNING UP IS ONLY USING ONE OPTIMIZER FOR BOTH Q NETWORKS!!!!!
      for q_optimizer, q_loss, q_params in zip(
          self.q_critic_optimizer,
          q_losses,
          self.learner_model.critic.per_critic_parameters(),
      ):
          loss_per_critic = torch.stack(list(q_loss.values())).sum(0)
          q_optimizer.zero_grad()
          loss_per_critic.backward(retain_graph=False)
          if self.algorithm_config.max_grad_norm > 0.0:
              nn.utils.clip_grad_norm_(
                  q_params, self.algorithm_config.max_grad_norm
              )
          q_optimizer.step()
  
      # ==============================================================================================================
      # Compute policy losses ========================================================================================
      # ==============================================================================================================
  
      # Freeze Q-networks so you don't waste computational effort
      # computing gradients for them during the policy learning step.
      for q_params in self.learner_model.critic.per_critic_parameters():
          for q_param in q_params:
              q_param.requires_grad = False
  
      (
          policy_losses,
          probs_policies,
          logp_policies,
          learner_entropies,
      ) = self._compute_policy_loss(worker_output)
  
      policy_loss = sum(policy_losses.values())
      self.policy_optimizer.zero_grad()
      policy_loss.backward()
      if self.algorithm_config.max_grad_norm > 0.0:
          nn.utils.clip_grad_norm_(
              self.learner_model.policy.parameters(),
              self.algorithm_config.max_grad_norm,
          )
      self.policy_optimizer.step()
  
      # Unfreeze Q-networks so you can optimize it at next step.
      for q_params in self.learner_model.critic.per_critic_parameters():
          for q_param in q_params:
              q_param.requires_grad = True
  
      # ==============================================================================================================
      # Compute entropy loss =========================================================================================
      # ==============================================================================================================
      entropy_losses = 0
      if self.algorithm_config.entropy_tuning:
          entropy_losses = self._compute_entropy_loss(
              probs_policies, logp_policies
          )
  
          entropy_loss = sum(entropy_losses.values())
          self.entropy_optimizer.zero_grad()
          entropy_loss.backward()
          self.entropy_optimizer.step()
          self.curr_entropy_coef = {
              step_key: torch.exp(log_alpha.detach())
              for step_key, log_alpha in self.curr_log_entropy_coef.items()
          }
  
      # ==============================================================================================================
      # Collect stats and log ========================================================================================
      # ==============================================================================================================
  
      # Policy data
      # for step_key in self.sub_step_keys:
      for step_key in self.sub_step_keys:
          policy_grad_norm = compute_gradient_norm(
              self.learner_model.policy.networks[step_key].parameters()
          )
          # log grad norms
          self.events.policy_grad_norm(
              step_key=step_key, value=policy_grad_norm
          )
  
          self.events.policy_loss(
              step_key=step_key,
              value=policy_losses[step_key].detach().item(),
          )
          self.events.policy_entropy(
              step_key=step_key,
              value=learner_entropies[step_key].mean().detach().item(),
          )
          if self.learner_model.critic.only_discrete_spaces[step_key]:
              value = (
                  torch.stack(list(logp_policies[step_key].values()))
                  .sum(0)
                  .mean()
                  .detach()
                  .item()
              )
          else:
              value = logp_policies[step_key].mean().detach().item()
          self.events.policy_mean_logp(step_key=step_key, value=value)
  
      # Critic data
      for critic_key in self.learner_model.critic.step_critic_keys:
          for idx, critic_step_key in enumerate(
              self.learner_model.critic.critic_key_mapping[critic_key]
          ):
              critic_grad_norm = compute_gradient_norm(
                  self.learner_model.critic.networks[
                      critic_step_key
                  ].parameters()
              )
              self.events.critic_value(
                  critic_key=critic_step_key,
                  value=q_values_mean[idx][critic_key],
              )
              self.events.critic_value_loss(
                  critic_key=critic_step_key,
                  value=q_losses[idx][critic_key].detach().item(),
              )
              self.events.critic_grad_norm(
                  critic_key=critic_step_key, value=critic_grad_norm
              )
  
          # self.events.errors_between_critics(critic_key=critic_key,
          #                                    value=error_between_critics[critic_key].mean().item())
  
      # Entropy coef data
      for step_key in self.sub_step_keys:
          if self.algorithm_config.entropy_tuning:
              self.events.entropy_loss(
                  step_key=step_key,
                  value=entropy_losses[step_key].detach().cpu().item(),
              )
          self.events.entropy_coef(
              step_key=step_key,
              value=self.curr_entropy_coef[step_key].cpu().item(),
          )
  
      # Timing events
      total_update_time = time.time() - start_update_time
      time_backprob = time.time() - after_collection_time
      self.events.time_backprob(
          time=time_backprob, percent=time_backprob / total_update_time
      )
      time_collecting_actors_total = (
          after_collection_time - start_update_time
      )
      self.events.time_sampling_from_buffer(
          time=time_collecting_actors_total,
          percent=time_collecting_actors_total / total_update_time,
      )
  
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 8 for calling the SAC obj's underlying member method code with example usage: obj._update() without return types