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_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`
      """
  
      flattened_observations = flatten_spaces(observations.values())
      flattened_actions = flatten_spaces(actions.values())
  
      assert len(self.step_critic_keys) == 1
      step_id = self.step_critic_keys[0]
      if all(self.only_discrete_spaces.values()):
          out = self.compute_state_action_values_step(
              flattened_observations, step_id
          )
          # output shape List[Dict[str, (rollout_length, batch_dim)]]
          if gather_output:
              out = [
                  {
                      action_key: action_value.gather(
                          -1,
                          flattened_actions[
                              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_value = out
      else:
          q_value = self.compute_state_action_value_step(
              flattened_observations, flattened_actions, step_id
          )
  
      q_values = {step_id: q_value}
      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 TorchSharedStateActionCritic obj's underlying member method code with example usage: obj.predict_q_values(observations, actions, gather_output) and returns: q_values