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 __init__(self, obs_space, action_space, config):
      # We want to error out on instantiation and not on import, because tune
      # imports all RLlib algorithms when registering them
      # TODO (Artur): Find a way to only import algorithms when needed
      if not torch:
          raise ImportError(
              "Could not import PyTorch, which QMix requires."
          )
  
      _validate(obs_space, action_space)
      self.framework = "torch"
  
      self.n_agents = len(obs_space.original_space.spaces)
      config["model"]["n_agents"] = self.n_agents
      self.n_actions = action_space.spaces[0].n
      self.h_size = config["model"]["lstm_cell_size"]
      self.has_env_global_state = False
      self.has_action_mask = False
  
      agent_obs_space = obs_space.original_space.spaces[0]
      if isinstance(agent_obs_space, gym.spaces.Dict):
          space_keys = set(agent_obs_space.spaces.keys())
          if "obs" not in space_keys:
              raise ValueError(
                  "Dict obs space must have subspace labeled `obs`"
              )
          self.obs_size = _get_size(agent_obs_space.spaces["obs"])
          if "action_mask" in space_keys:
              mask_shape = tuple(
                  agent_obs_space.spaces["action_mask"].shape
              )
              if mask_shape != (self.n_actions,):
                  raise ValueError(
                      "Action mask shape must be {}, got {}".format(
                          (self.n_actions,), mask_shape
                      )
                  )
              self.has_action_mask = True
          if ENV_STATE in space_keys:
              self.env_global_state_shape = _get_size(
                  agent_obs_space.spaces[ENV_STATE]
              )
              self.has_env_global_state = True
          else:
              self.env_global_state_shape = (
                  self.obs_size,
                  self.n_agents,
              )
          # The real agent obs space is nested inside the dict
          config["model"]["full_obs_space"] = agent_obs_space
          agent_obs_space = agent_obs_space.spaces["obs"]
      else:
          self.obs_size = _get_size(agent_obs_space)
          self.env_global_state_shape = (self.obs_size, self.n_agents)
  
      self.model = ModelCatalog.get_model_v2(
          agent_obs_space,
          action_space.spaces[0],
          self.n_actions,
          config["model"],
          framework="torch",
          name="model",
          default_model=RNNModel,
      )
  
      super().__init__(
          obs_space, action_space, config, model=self.model
      )
  
      self.target_model = ModelCatalog.get_model_v2(
          agent_obs_space,
          action_space.spaces[0],
          self.n_actions,
          config["model"],
          framework="torch",
          name="target_model",
          default_model=RNNModel,
      ).to(self.device)
  
      self.exploration = self._create_exploration()
  
      # Setup the mixer network.
      if config["mixer"] is None:
          self.mixer = None
          self.target_mixer = None
      elif config["mixer"] == "qmix":
          self.mixer = QMixer(
              self.n_agents,
              self.env_global_state_shape,
              config["mixing_embed_dim"],
          ).to(self.device)
          self.target_mixer = QMixer(
              self.n_agents,
              self.env_global_state_shape,
              config["mixing_embed_dim"],
          ).to(self.device)
      elif config["mixer"] == "vdn":
          self.mixer = VDNMixer().to(self.device)
          self.target_mixer = VDNMixer().to(self.device)
      else:
          raise ValueError(
              "Unknown mixer type {}".format(config["mixer"])
          )
  
      self.cur_epsilon = 1.0
      self.update_target()  # initial sync
  
      # Setup optimizer
      self.params = list(self.model.parameters())
      if self.mixer:
          self.params += list(self.mixer.parameters())
      self.loss = QMixLoss(
          self.model,
          self.target_model,
          self.mixer,
          self.target_mixer,
          self.n_agents,
          self.n_actions,
          self.config["double_q"],
          self.config["gamma"],
      )
      from torch.optim import RMSprop
  
      self.rmsprop_optimizer = RMSprop(
          params=self.params,
          lr=config["lr"],
          alpha=config["optim_alpha"],
          eps=config["optim_eps"],
      )
  
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 1 to create a new QMixTorchPolicy object from inherited base classes: TorchPolicy with example: obj = QMixTorchPolicy(obs_space, action_space, config)