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 learn_on_batch(self, samples):
      (
          obs_batch,
          action_mask,
          env_global_state,
      ) = self._unpack_observation(samples[SampleBatch.CUR_OBS])
      (
          next_obs_batch,
          next_action_mask,
          next_env_global_state,
      ) = self._unpack_observation(samples[SampleBatch.NEXT_OBS])
      group_rewards = self._get_group_rewards(
          samples[SampleBatch.INFOS]
      )
  
      input_list = [
          group_rewards,
          action_mask,
          next_action_mask,
          samples[SampleBatch.ACTIONS],
          samples[SampleBatch.TERMINATEDS],
          obs_batch,
          next_obs_batch,
      ]
      if self.has_env_global_state:
          input_list.extend([env_global_state, next_env_global_state])
  
      output_list, _, seq_lens = chop_into_sequences(
          episode_ids=samples[SampleBatch.EPS_ID],
          unroll_ids=samples[SampleBatch.UNROLL_ID],
          agent_indices=samples[SampleBatch.AGENT_INDEX],
          feature_columns=input_list,
          state_columns=[],  # RNN states not used here
          max_seq_len=self.config["model"]["max_seq_len"],
          dynamic_max=True,
      )
      # These will be padded to shape [B * T, ...]
      if self.has_env_global_state:
          (
              rew,
              action_mask,
              next_action_mask,
              act,
              terminateds,
              obs,
              next_obs,
              env_global_state,
              next_env_global_state,
          ) = output_list
      else:
          (
              rew,
              action_mask,
              next_action_mask,
              act,
              terminateds,
              obs,
              next_obs,
          ) = output_list
      B, T = len(seq_lens), max(seq_lens)
  
      def to_batches(arr, dtype):
          new_shape = [B, T] + list(arr.shape[1:])
          return torch.as_tensor(
              np.reshape(arr, new_shape),
              dtype=dtype,
              device=self.device,
          )
  
      rewards = to_batches(rew, torch.float)
      actions = to_batches(act, torch.long)
      obs = to_batches(obs, torch.float).reshape(
          [B, T, self.n_agents, self.obs_size]
      )
      action_mask = to_batches(action_mask, torch.float)
      next_obs = to_batches(next_obs, torch.float).reshape(
          [B, T, self.n_agents, self.obs_size]
      )
      next_action_mask = to_batches(next_action_mask, torch.float)
      if self.has_env_global_state:
          env_global_state = to_batches(env_global_state, torch.float)
          next_env_global_state = to_batches(
              next_env_global_state, torch.float
          )
  
      # TODO(ekl) this treats group termination as individual termination
      terminated = (
          to_batches(terminateds, torch.float)
          .unsqueeze(2)
          .expand(B, T, self.n_agents)
      )
  
      # Create mask for where index is < unpadded sequence length
      filled = np.reshape(
          np.tile(np.arange(T, dtype=np.float32), B), [B, T]
      ) < np.expand_dims(seq_lens, 1)
      mask = (
          torch.as_tensor(filled, dtype=torch.float, device=self.device)
          .unsqueeze(2)
          .expand(B, T, self.n_agents)
      )
  
      # Compute loss
      (
          loss_out,
          mask,
          masked_td_error,
          chosen_action_qvals,
          targets,
      ) = self.loss(
          rewards,
          actions,
          terminated,
          mask,
          obs,
          next_obs,
          action_mask,
          next_action_mask,
          env_global_state,
          next_env_global_state,
      )
  
      # Optimise
      self.rmsprop_optimizer.zero_grad()
      loss_out.backward()
      grad_norm_info = apply_grad_clipping(
          self, self.rmsprop_optimizer, loss_out
      )
      self.rmsprop_optimizer.step()
  
      mask_elems = mask.sum().item()
      stats = {
          "loss": loss_out.item(),
          "td_error_abs": masked_td_error.abs().sum().item()
          / mask_elems,
          "q_taken_mean": (chosen_action_qvals * mask).sum().item()
          / mask_elems,
          "target_mean": (targets * mask).sum().item() / mask_elems,
      }
      stats.update(grad_norm_info)
  
      return {LEARNER_STATS_KEY: stats}
  
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 5 for calling the QMixTorchPolicy obj's underlying member method code with example usage: obj.learn_on_batch(samples) and returns: dict, torch