class_docstr: One critic is shared across all sub-steps or actors (default to use for standard gym-style environments).
  Can be instantiated via the
  :class:`~maze.perception.models.critics.shared_state_action_critics_composer.SharedStateActionCriticComposer`.
class_docstr_tok: ["One","critic","is","shared","across","all","sub-steps","or","actors","(","default","to","use","for","standard","gym-style","environments",")",".","Can","be","instantiated","via","the",":","class",":","`","~maze.perception.models.critics.shared_state_action_critics_composer.SharedStateActionCriticComposer","`","."]
docstr: implementation of
  :class:`~maze.core.agent.torch_state_action_critic.TorchStateActionCritic`
docstr_tok: ["implementation","of",":","class",":","`","~maze.core.agent.torch_state_action_critic.TorchStateActionCritic","`"]
code: def predict_next_q_values(
      self,
      next_observations: Dict[Union[str, int], Dict[str, torch.Tensor]],
      next_actions: Dict[Union[str, int], Dict[str, torch.Tensor]],
      next_actions_logits: Dict[
          Union[str, int], Dict[str, torch.Tensor]
      ],
      next_actions_log_probs: Dict[
          Union[str, int], Dict[str, torch.Tensor]
      ],
      alpha: Dict[Union[str, int], torch.Tensor],
  ) -> Dict[
      Union[str, int], Union[torch.Tensor, Dict[str, torch.Tensor]]
  ]:
      """implementation of
      :class:`~maze.core.agent.torch_state_action_critic.TorchStateActionCritic`
      """
  
      flattened_next_observations = flatten_spaces(
          next_observations.values()
      )
      flattened_next_actions = flatten_spaces(next_actions.values())
      flattened_next_actions_logits = flatten_spaces(
          next_actions_logits.values()
      )
      flattened_next_action_log_probs = flatten_spaces(
          next_actions_log_probs.values()
      )
  
      assert len(self.step_critic_keys) == 1
      step_id = self.step_critic_keys[0]
      alpha = sum(alpha.values())
  
      if all(self.only_discrete_spaces.values()):
          next_q_values = self.compute_state_action_values_step(
              flattened_next_observations,
              critic_id=(step_id, self.target_key),
          )
          transpose_next_q_value = {
              k: [dic[k] for dic in next_q_values]
              for k in next_q_values[0]
          }
          next_q_value = dict()
          for q_action_head, q_values in transpose_next_q_value.items():
              action_key = q_action_head.replace("_q_values", "")
              tmp_q_value = torch.stack(q_values).min(dim=0).values
              next_action_probs = logits_to_probs(
                  flattened_next_actions_logits[action_key]
              )
              next_action_log_probs = torch.log(
                  next_action_probs
                  + (next_action_probs == 0.0).float() * 1e-8
              )
  
              # output shape of V(st) is (rollout_length, batch_dim)
              next_q_value[action_key] = (
                  torch.matmul(
                      next_action_probs.unsqueeze(-2),
                      (
                          tmp_q_value - alpha * next_action_log_probs
                      ).unsqueeze(-1),
                  )
                  .squeeze(-1)
                  .squeeze(-1)
              )
  
      else:
          next_q_value = self.compute_state_action_value_step(
              flattened_next_observations,
              flattened_next_actions,
              (step_id, self.target_key),
          )
          next_q_value = torch.stack(next_q_value).min(
              dim=0
          ).values - alpha * torch.stack(
              list(flattened_next_action_log_probs.values())
          ).mean(
              dim=0
          )
  
      return {step_id: next_q_value}
  
imports: import copy
  import itertools
  from abc import abstractmethod
  from typing import Mapping, Union, List, Dict
  import torch
  from gym import spaces
  from torch import nn
  from torch.distributions.utils import logits_to_probs
  from maze.core.agent.state_action_critic import StateActionCritic
  from maze.core.agent.torch_model import TorchModel
  from maze.core.annotations import override
  from maze.core.wrappers.observation_preprocessing.preprocessors.one_hot import OneHotPreProcessor
  from maze.perception.blocks.inference import InferenceBlock
  from maze.perception.perception_utils import flatten_spaces, convert_to_numpy, convert_to_torch
  from maze.perception.weight_init import make_module_init_normc
  from maze.utils.bcolors import BColors
label_desc: Use image node_id 2 for calling the TorchSharedStateActionCritic obj's underlying member method code with example usage: obj.predict_next_q_values(next_observations, next_actions, next_actions_logits, next_actions_log_probs, alpha) and returns: dict