docstr: Forward pass of the loss.
  
  Args:
      rewards: Tensor of shape [B, T, n_agents]
      actions: Tensor of shape [B, T, n_agents]
      terminated: Tensor of shape [B, T, n_agents]
      mask: Tensor of shape [B, T, n_agents]
      obs: Tensor of shape [B, T, n_agents, obs_size]
      next_obs: Tensor of shape [B, T, n_agents, obs_size]
      action_mask: Tensor of shape [B, T, n_agents, n_actions]
      next_action_mask: Tensor of shape [B, T, n_agents, n_actions]
      state: Tensor of shape [B, T, state_dim] (optional)
      next_state: Tensor of shape [B, T, state_dim] (optional)
docstr_tok: ["Forward","pass","of","the","loss",".","Args",":","rewards",":","Tensor","of","shape","[","B",",","T",",","n_agents","]","actions",":","Tensor","of","shape","[","B",",","T",",","n_agents","]","terminated",":","Tensor","of","shape","[","B",",","T",",","n_agents","]","mask",":","Tensor","of","shape","[","B",",","T",",","n_agents","]","obs",":","Tensor","of","shape","[","B",",","T",",","n_agents",",","obs_size","]","next_obs",":","Tensor","of","shape","[","B",",","T",",","n_agents",",","obs_size","]","action_mask",":","Tensor","of","shape","[","B",",","T",",","n_agents",",","n_actions","]","next_action_mask",":","Tensor","of","shape","[","B",",","T",",","n_agents",",","n_actions","]","state",":","Tensor","of","shape","[","B",",","T",",","state_dim","]","(","optional",")","next_state",":","Tensor","of","shape","[","B",",","T",",","state_dim","]","(","optional",")"]
code: def forward(
      self,
      rewards,
      actions,
      terminated,
      mask,
      obs,
      next_obs,
      action_mask,
      next_action_mask,
      state=None,
      next_state=None,
  ):
      """Forward pass of the loss.
  
      Args:
          rewards: Tensor of shape [B, T, n_agents]
          actions: Tensor of shape [B, T, n_agents]
          terminated: Tensor of shape [B, T, n_agents]
          mask: Tensor of shape [B, T, n_agents]
          obs: Tensor of shape [B, T, n_agents, obs_size]
          next_obs: Tensor of shape [B, T, n_agents, obs_size]
          action_mask: Tensor of shape [B, T, n_agents, n_actions]
          next_action_mask: Tensor of shape [B, T, n_agents, n_actions]
          state: Tensor of shape [B, T, state_dim] (optional)
          next_state: Tensor of shape [B, T, state_dim] (optional)
      """
  
      # Assert either none or both of state and next_state are given
      if state is None and next_state is None:
          state = obs  # default to state being all agents' observations
          next_state = next_obs
      elif (state is None) != (next_state is None):
          raise ValueError(
              "Expected either neither or both of `state` and "
              "`next_state` to be given. Got: "
              "\n`state` = {}\n`next_state` = {}".format(
                  state, next_state
              )
          )
  
      # Calculate estimated Q-Values
      mac_out = _unroll_mac(self.model, obs)
  
      # Pick the Q-Values for the actions taken -> [B * n_agents, T]
      chosen_action_qvals = torch.gather(
          mac_out, dim=3, index=actions.unsqueeze(3)
      ).squeeze(3)
  
      # Calculate the Q-Values necessary for the target
      target_mac_out = _unroll_mac(self.target_model, next_obs)
  
      # Mask out unavailable actions for the t+1 step
      ignore_action_tp1 = (next_action_mask == 0) & (
          mask == 1
      ).unsqueeze(-1)
      target_mac_out[ignore_action_tp1] = -np.inf
  
      # Max over target Q-Values
      if self.double_q:
          # Double Q learning computes the target Q values by selecting the
          # t+1 timestep action according to the "policy" neural network and
          # then estimating the Q-value of that action with the "target"
          # neural network
  
          # Compute the t+1 Q-values to be used in action selection
          # using next_obs
          mac_out_tp1 = _unroll_mac(self.model, next_obs)
  
          # mask out unallowed actions
          mac_out_tp1[ignore_action_tp1] = -np.inf
  
          # obtain best actions at t+1 according to policy NN
          cur_max_actions = mac_out_tp1.argmax(dim=3, keepdim=True)
  
          # use the target network to estimate the Q-values of policy
          # network's selected actions
          target_max_qvals = torch.gather(
              target_mac_out, 3, cur_max_actions
          ).squeeze(3)
      else:
          target_max_qvals = target_mac_out.max(dim=3)[0]
  
      assert (
          target_max_qvals.min().item() != -np.inf
      ), "target_max_qvals contains a masked action; \
              there may be a state with no valid actions."
  
      # Mix
      if self.mixer is not None:
          chosen_action_qvals = self.mixer(chosen_action_qvals, state)
          target_max_qvals = self.target_mixer(
              target_max_qvals, next_state
          )
  
      # Calculate 1-step Q-Learning targets
      targets = (
          rewards + self.gamma * (1 - terminated) * target_max_qvals
      )
  
      # Td-error
      td_error = chosen_action_qvals - targets.detach()
  
      mask = mask.expand_as(td_error)
  
      # 0-out the targets that came from padded data
      masked_td_error = td_error * mask
  
      # Normal L2 loss, take mean over actual data
      loss = (masked_td_error**2).sum() / mask.sum()
      return loss, mask, masked_td_error, chosen_action_qvals, targets
  
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 QMixLoss obj's underlying member method code with example usage: obj.forward(rewards, actions, terminated, mask, obs, next_obs, action_mask, next_action_mask, state, next_state) and returns: loss,  mask,  masked_td_error,  chosen_action_qvals,  targets