code: def _build_module(self, input_layer):
      self.actions = tf.placeholder(tf.int32, [None, 2], name="actions")
      self.quantile_midpoints = tf.placeholder(
          tf.float32, [None, self.num_atoms], name="quantile_midpoints"
      )
      self.input = [self.actions, self.quantile_midpoints]
  
      # the output of the head is the N unordered quantile locations {theta_1, ..., theta_N}
      quantiles_locations = self.dense_layer(
          self.num_actions * self.num_atoms
      )(
          input_layer,
          name="output",
          bias_initializer=self.output_bias_initializer,
      )
      quantiles_locations = tf.reshape(
          quantiles_locations,
          (
              tf.shape(quantiles_locations)[0],
              self.num_actions,
              self.num_atoms,
          ),
      )
      self.output = quantiles_locations
  
      self.quantiles = tf.placeholder(
          tf.float32, shape=(None, self.num_atoms), name="quantiles"
      )
      self.target = self.quantiles
  
      # only the quantiles of the taken action are taken into account
      quantiles_for_used_actions = tf.gather_nd(
          quantiles_locations, self.actions
      )
  
      # reorder the output quantiles and the target quantiles as a preparation step for calculating the loss
      # the output quantiles vector and the quantile midpoints are tiled as rows of a NxN matrix (N = num quantiles)
      # the target quantiles vector is tiled as column of a NxN matrix
      theta_i = tf.tile(
          tf.expand_dims(quantiles_for_used_actions, -1),
          [1, 1, self.num_atoms],
      )
      T_theta_j = tf.tile(
          tf.expand_dims(self.target, -2), [1, self.num_atoms, 1]
      )
      tau_i = tf.tile(
          tf.expand_dims(self.quantile_midpoints, -1),
          [1, 1, self.num_atoms],
      )
  
      # Huber loss of T(theta_j) - theta_i
      error = T_theta_j - theta_i
      abs_error = tf.abs(error)
      quadratic = tf.minimum(abs_error, self.huber_loss_interval)
      huber_loss = (
          self.huber_loss_interval * (abs_error - quadratic)
          + 0.5 * quadratic**2
      )
  
      # Quantile Huber loss
      quantile_huber_loss = (
          tf.abs(tau_i - tf.cast(error < 0, dtype=tf.float32))
          * huber_loss
      )
  
      # Quantile regression loss (the probability for each quantile is 1/num_quantiles)
      quantile_regression_loss = tf.reduce_sum(
          quantile_huber_loss
      ) / float(self.num_atoms)
      self.loss = quantile_regression_loss
      tf.losses.add_loss(self.loss)
  
      self.q_values = tf.tensordot(
          tf.cast(self.output, tf.float64),
          self.quantile_probabilities,
          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/quantile_regression_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 QuantileRegressionQHead obj's underlying member method code with example usage: obj._build_module(input_layer) without return types