class_docstr: A Duelling MLP Q-network.
class_docstr_tok: ["A","Duelling","MLP","Q-network","."]
docstr: Forward pass of the duelling network.
  
  Args:
    inputs: 2-D tensor of shape [batch_size, embedding_size].
  
  Returns:
    q_values: 2-D tensor of action values of shape [batch_size, num_actions]
docstr_tok: ["Forward","pass","of","the","duelling","network",".","Args",":","inputs",":","2-D","tensor","of","shape","[","batch_size",",","embedding_size","]",".","Returns",":","q_values",":","2-D","tensor","of","action","values","of","shape","[","batch_size",",","num_actions","]"]
code: def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
      """Forward pass of the duelling network.
  
      Args:
        inputs: 2-D tensor of shape [batch_size, embedding_size].
  
      Returns:
        q_values: 2-D tensor of action values of shape [batch_size, num_actions]
      """
  
      # Compute value & advantage for duelling.
      value = self._value_mlp(inputs)  # [B, 1]
      advantages = self._advantage_mlp(inputs)  # [B, A]
  
      # Advantages have zero mean.
      advantages -= jnp.mean(
          advantages, axis=-1, keepdims=True
      )  # [B, A]
  
      q_values = value + advantages  # [B, A]
  
      return q_values
  
imports: from typing import Sequence, Optional
  import haiku
  import jax.numpy
label_desc: Use image node_id 2 for calling the DuellingMLP obj's underlying member method code with example usage: obj.__call__(inputs) and returns: q_values