docstr: Forward pass for the mixer.
  
  Args:
      agent_qs: Tensor of shape [B, T, n_agents, n_actions]
      states: Tensor of shape [B, T, state_dim]
docstr_tok: ["Forward","pass","for","the","mixer",".","Args",":","agent_qs",":","Tensor","of","shape","[","B",",","T",",","n_agents",",","n_actions","]","states",":","Tensor","of","shape","[","B",",","T",",","state_dim","]"]
code: def forward(self, agent_qs, states):
      """Forward pass for the mixer.
  
      Args:
          agent_qs: Tensor of shape [B, T, n_agents, n_actions]
          states: Tensor of shape [B, T, state_dim]
      """
      bs = agent_qs.size(0)
      states = states.reshape(-1, self.state_dim)
      agent_qs = agent_qs.view(-1, 1, self.n_agents)
      # First layer
      w1 = torch.abs(self.hyper_w_1(states))
      b1 = self.hyper_b_1(states)
      w1 = w1.view(-1, self.n_agents, self.embed_dim)
      b1 = b1.view(-1, 1, self.embed_dim)
      hidden = nn.functional.elu(torch.bmm(agent_qs, w1) + b1)
      # Second layer
      w_final = torch.abs(self.hyper_w_final(states))
      w_final = w_final.view(-1, self.embed_dim, 1)
      # State-dependent bias
      v = self.V(states).view(-1, 1, 1)
      # Compute final output
      y = torch.bmm(hidden, w_final) + v
      # Reshape and return
      q_tot = y.view(bs, -1, 1)
      return q_tot
  
file_path: ray/rllib_contrib/qmix/src/rllib_qmix/qmix/mixers.py
imports: import numpy
  from ray.rllib.utils.framework import try_import_torch
label_desc: Use image node_id 2 for calling the QMixer obj's underlying member method code with example usage: obj.forward(agent_qs, states) and returns: q_tot