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_q_values(
      self,
      observations: Dict[Union[str, int], Dict[str, torch.Tensor]],
      actions: Dict[Union[str, int], Dict[str, torch.Tensor]],
      gather_output: bool,
  ) -> Dict[
      Union[str, int],
      List[Union[torch.Tensor, Dict[str, torch.Tensor]]],
  ]:
      """implementation of
      :class:`~maze.core.agent.torch_state_action_critic.TorchStateActionCritic`
      """
  
      q_values = dict()
      for step_id in observations.keys():
          if self.only_discrete_spaces[step_id]:
              out = self.compute_state_action_values_step(
                  observations[step_id], step_id
              )
              # output shape List[Dict[str, (rollout_length, batch_dim)]]
              if gather_output:
                  out = [
                      {
                          action_key: action_value.gather(
                              -1,
                              actions[step_id][
                                  action_key.replace("_q_values", "")
                              ]
                              .long()
                              .unsqueeze(-1),
                          ).squeeze(-1)
                          for action_key, action_value in critic_out.items()
                      }
                      for critic_out in out
                  ]
              q_values[step_id] = out
          else:
              q_values[step_id] = self.compute_state_action_value_step(
                  observations[step_id], actions[step_id], step_id
              )
  
      return 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 1 for calling the TorchStepStateActionCritic obj's underlying member method code with example usage: obj.predict_q_values(observations, actions, gather_output) and returns: q_values