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",",",")",")","`","."]
docstr: Unpacks the observation, action mask, and state (if present)
  from agent grouping.
  
  Returns:
      obs (np.ndarray): obs tensor of shape [B, n_agents, obs_size]
      mask (np.ndarray): action mask, if any
      state (np.ndarray or None): state tensor of shape [B, state_size]
          or None if it is not in the batch
docstr_tok: ["Unpacks","the","observation",",","action","mask",",","and","state","(","if","present",")","from","agent","grouping",".","Returns",":","obs","(","np.ndarray",")",":","obs","tensor","of","shape","[","B",",","n_agents",",","obs_size","]","mask","(","np.ndarray",")",":","action","mask",",","if","any","state","(","np.ndarray","or","None",")",":","state","tensor","of","shape","[","B",",","state_size","]","or","None","if","it","is","not","in","the","batch"]
code: def _unpack_observation(self, obs_batch):
      """Unpacks the observation, action mask, and state (if present)
      from agent grouping.
  
      Returns:
          obs (np.ndarray): obs tensor of shape [B, n_agents, obs_size]
          mask (np.ndarray): action mask, if any
          state (np.ndarray or None): state tensor of shape [B, state_size]
              or None if it is not in the batch
      """
  
      unpacked = _unpack_obs(
          np.array(obs_batch, dtype=np.float32),
          self.observation_space.original_space,
          tensorlib=np,
      )
  
      if isinstance(unpacked[0], dict):
          assert "obs" in unpacked[0]
          unpacked_obs = [
              np.concatenate(tree.flatten(u["obs"]), 1)
              for u in unpacked
          ]
      else:
          unpacked_obs = unpacked
  
      obs = np.concatenate(unpacked_obs, axis=1).reshape(
          [len(obs_batch), self.n_agents, self.obs_size]
      )
  
      if self.has_action_mask:
          action_mask = np.concatenate(
              [o["action_mask"] for o in unpacked], axis=1
          ).reshape([len(obs_batch), self.n_agents, self.n_actions])
      else:
          action_mask = np.ones(
              [len(obs_batch), self.n_agents, self.n_actions],
              dtype=np.float32,
          )
  
      if self.has_env_global_state:
          state = np.concatenate(
              tree.flatten(unpacked[0][ENV_STATE]), 1
          )
      else:
          state = None
      return obs, action_mask, state
  
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 16 for calling the QMixTorchPolicy obj's underlying member method code with example usage: obj._unpack_observation(obs_batch) and returns: obs,  action_mask,  state