class_docstr: Each sub-step or actor gets its individual critic.
  Can be instantiated via the
  :class:`~maze.perception.models.critics.step_state_action_critic_composer.StepStateActionCriticComposer`.
class_docstr_tok: ["Each","sub-step","or","actor","gets","its","individual","critic",".","Can","be","instantiated","via","the",":","class",":","`","~maze.perception.models.critics.step_state_action_critic_composer.StepStateActionCriticComposer","`","."]
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`
      """
  
      next_q_values = dict()
      for step_id in next_observations.keys():
          if self.only_discrete_spaces[step_id]:
              next_q_value = self.compute_state_action_values_step(
                  next_observations[step_id],
                  critic_id=(step_id, self.target_key),
              )
              transpose_next_q_value = {
                  k: [dic[k] for dic in next_q_value]
                  for k in next_q_value[0]
              }
              next_q_values[step_id] = 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(
                      next_actions_logits[step_id][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_values[step_id][action_key] = (
                      torch.matmul(
                          next_action_probs.unsqueeze(-2),
                          (
                              tmp_q_value
                              - alpha[step_id] * next_action_log_probs
                          ).unsqueeze(-1),
                      )
                      .squeeze(-1)
                      .squeeze(-1)
                  )
          else:
              next_q_value = self.compute_state_action_value_step(
                  next_observations[step_id],
                  next_actions[step_id],
                  (step_id, self.target_key),
              )
              # output shape of V(st) is (rollout_length, batch_size)
              next_q_values[step_id] = torch.stack(next_q_value).min(
                  dim=0
              ).values - alpha[step_id] * torch.stack(
                  list(next_actions_log_probs[step_id].values())
              ).mean(
                  dim=0
              )
  
      return next_q_values
  
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 TorchStepStateActionCritic 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: next_q_values