class_docstr: QMix impl. Assumes homogeneous agents for now.
  
  You must use MultiAgentEnv.with_agent_groups() to group agents
  together for QMix. This creates the proper Tuple obs/action spaces and
  populates the '_group_rewards' info field.
  
  Action masking: to specify an action mask for individual agents, use a
  dict space with an action_mask key, e.g. {"obs": ob, "action_mask": mask}.
  The mask space must be `Box(0, 1, (n_actions,))`.
class_docstr_tok: ["QMix","impl",".","Assumes","homogeneous","agents","for","now",".","You","must","use","MultiAgentEnv.with_agent_groups","(",")","to","group","agents","together","for","QMix",".","This","creates","the","proper","Tuple","obs\/action","spaces","and","populates","the","'_group_rewards","'","info","field",".","Action","masking",":","to","specify","an","action","mask","for","individual","agents",",","use","a","dict","space","with","an","action_mask","key",",","e.g",".","{","``","obs","''",":","ob",",","``","action_mask","''",":","mask","}",".","The","mask","space","must","be","`","Box","(","0",",","1",",","(","n_actions",",",")",")","`","."]
code: def compute_actions_from_input_dict(
      self,
      input_dict: Dict[str, TensorType],
      explore: bool = None,
      timestep: Optional[int] = None,
      **kwargs,
  ) -> Tuple[TensorType, List[TensorType], Dict[str, TensorType]]:
      obs_batch = input_dict[SampleBatch.OBS]
      state_batches = []
      i = 0
      while f"state_in_{i}" in input_dict:
          state_batches.append(input_dict[f"state_in_{i}"])
          i += 1
  
      explore = (
          explore if explore is not None else self.config["explore"]
      )
      obs_batch, action_mask, _ = self._unpack_observation(obs_batch)
      # We need to ensure we do not use the env global state
      # to compute actions
  
      # Compute actions
      with torch.no_grad():
          q_values, hiddens = _mac(
              self.model,
              torch.as_tensor(
                  obs_batch, dtype=torch.float, device=self.device
              ),
              [
                  torch.as_tensor(
                      np.array(s), dtype=torch.float, device=self.device
                  )
                  for s in state_batches
              ],
          )
          avail = torch.as_tensor(
              action_mask, dtype=torch.float, device=self.device
          )
          masked_q_values = q_values.clone()
          masked_q_values[avail == 0.0] = -float("inf")
          masked_q_values_folded = torch.reshape(
              masked_q_values, [-1] + list(masked_q_values.shape)[2:]
          )
          actions, _ = self.exploration.get_exploration_action(
              action_distribution=TorchCategorical(
                  masked_q_values_folded
              ),
              timestep=timestep,
              explore=explore,
          )
          actions = (
              torch.reshape(actions, list(masked_q_values.shape)[:-1])
              .cpu()
              .numpy()
          )
          hiddens = [s.cpu().numpy() for s in hiddens]
  
      return tuple(actions.transpose([1, 0])), hiddens, {}
  
imports: import logging
  from typing import Dict, List, Optional, Tuple
  import gymnasium
  import numpy
  import tree
  from rllib_qmix.qmix.mixers import QMixer, VDNMixer
  from rllib_qmix.qmix.model import RNNModel, _get_size
  from ray.rllib.env.multi_agent_env import ENV_STATE
  from ray.rllib.env.wrappers.group_agents_wrapper import GROUP_REWARDS
  from ray.rllib.models.catalog import ModelCatalog
  from ray.rllib.models.modelv2 import _unpack_obs
  from ray.rllib.models.torch.torch_action_dist import TorchCategorical
  from ray.rllib.policy.rnn_sequencing import chop_into_sequences
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.policy.torch_policy import TorchPolicy
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.metrics.learner_info import LEARNER_STATS_KEY
  from ray.rllib.utils.torch_utils import apply_grad_clipping
  from ray.rllib.utils.typing import TensorType
label_desc: Use image node_id 2 for calling the QMixTorchPolicy obj's underlying member method code with example usage: obj.compute_actions_from_input_dict(input_dict, explore, timestep) and returns: tuple,  hiddens,  dict