code: def _build_module(self, input_layer):
      # state value tower - V
      with tf.variable_scope("state_value"):
          state_value = self.dense_layer(512)(
              input_layer,
              activation=self.activation_function,
              name="fc1",
          )
          state_value = self.dense_layer(self.num_atoms)(
              state_value, name="fc2"
          )
          state_value = tf.expand_dims(state_value, axis=1)
  
      # action advantage tower - A
      with tf.variable_scope("action_advantage"):
          action_advantage = self.dense_layer(512)(
              input_layer,
              activation=self.activation_function,
              name="fc1",
          )
          action_advantage = self.dense_layer(
              self.num_actions * self.num_atoms
          )(action_advantage, name="fc2")
          action_advantage = tf.reshape(
              action_advantage,
              (
                  tf.shape(input_layer)[0],
                  self.num_actions,
                  self.num_atoms,
              ),
          )
          action_mean = tf.reduce_mean(
              action_advantage, axis=1, keepdims=True
          )
          action_advantage = action_advantage - action_mean
  
      # merge to state-action value function Q
      values_distribution = tf.add(
          state_value, action_advantage, name="output"
      )
  
      # softmax on atoms dimension
      self.output = tf.nn.softmax(values_distribution)
  
      # calculate cross entropy loss
      self.distributions = tf.placeholder(
          tf.float32,
          shape=(None, self.num_actions, self.num_atoms),
          name="distributions",
      )
      self.target = self.distributions
      self.loss = tf.nn.softmax_cross_entropy_with_logits(
          labels=self.target, logits=values_distribution
      )
      tf.losses.add_loss(self.loss)
  
      self.q_values = tf.tensordot(
          tf.cast(self.output, tf.float64), self.z_values, 1
      )
  
      # used in batch-rl to estimate a probablity distribution over actions
      self.softmax = self.add_softmax_with_temperature()
  
file_path: coach/rl_coach/architectures/tensorflow_components/heads/rainbow_q_head.py
imports: import tensorflow
  import numpy
  from rl_coach.architectures.tensorflow_components.heads import QHead
  from rl_coach.architectures.tensorflow_components.layers import Dense
  from rl_coach.base_parameters import AgentParameters
  from rl_coach.spaces import SpacesDefinition
label_desc: Use image node_id 2 for calling the RainbowQHead obj's underlying member method code with example usage: obj._build_module(input_layer) without return types