class_docstr: An extension of Rainbow to perform implicit quantile regression.
docstr: Builds a training op.
  
  Returns:
    train_op: An op performing one step of training from replay data.
docstr_tok: ["Builds","a","training","op",".","Returns",":","train_op",":","An","op","performing","one","step","of","training","from","replay","data","."]
code: def _build_train_op(self):
      """Builds a training op.
  
      Returns:
        train_op: An op performing one step of training from replay data.
      """
      batch_size = tf.shape(self._replay.rewards)[0]
  
      target_quantile_values = tf.stop_gradient(
          self._build_target_quantile_values_op()
      )
      # Reshape to self.num_tau_prime_samples x batch_size x 1 since this is
      # the manner in which the target_quantile_values are tiled.
      target_quantile_values = tf.reshape(
          target_quantile_values,
          [self.num_tau_prime_samples, batch_size, 1],
      )
      # Transpose dimensions so that the dimensionality is batch_size x
      # self.num_tau_prime_samples x 1 to prepare for computation of
      # Bellman errors.
      # Final shape of target_quantile_values:
      # batch_size x num_tau_prime_samples x 1.
      target_quantile_values = tf.transpose(
          target_quantile_values, [1, 0, 2]
      )
  
      # Shape of indices: (num_tau_samples x batch_size) x 1.
      # Expand dimension by one so that it can be used to index into all the
      # quantiles when using the tf.gather_nd function (see below).
      indices = tf.range(self.num_tau_samples * batch_size)[:, None]
  
      # Expand the dimension by one so that it can be used to index into all the
      # quantiles when using the tf.gather_nd function (see below).
      reshaped_actions = self._replay.actions[:, None]
      reshaped_actions = tf.tile(
          reshaped_actions, [self.num_tau_samples, 1]
      )
      # Shape of reshaped_actions: (num_tau_samples x batch_size) x 2.
      reshaped_actions = tf.concat([indices, reshaped_actions], axis=1)
  
      chosen_action_quantile_values = tf.gather_nd(
          self._replay_net_quantile_values, reshaped_actions
      )
      # Reshape to self.num_tau_samples x batch_size x 1 since this is the manner
      # in which the quantile values are tiled.
      chosen_action_quantile_values = tf.reshape(
          chosen_action_quantile_values,
          [self.num_tau_samples, batch_size, 1],
      )
      # Transpose dimensions so that the dimensionality is batch_size x
      # self.num_tau_samples x 1 to prepare for computation of
      # Bellman errors.
      # Final shape of chosen_action_quantile_values:
      # batch_size x num_tau_samples x 1.
      chosen_action_quantile_values = tf.transpose(
          chosen_action_quantile_values, [1, 0, 2]
      )
  
      # Shape of bellman_erors and huber_loss:
      # batch_size x num_tau_prime_samples x num_tau_samples x 1.
      bellman_errors = (
          target_quantile_values[:, :, None, :]
          - chosen_action_quantile_values[:, None, :, :]
      )
      # The huber loss (see Section 2.3 of the paper) is defined via two cases:
      # case_one: |bellman_errors| <= kappa
      # case_two: |bellman_errors| > kappa
      huber_loss_case_one = (
          tf.cast(tf.abs(bellman_errors) <= self.kappa, tf.float32)
          * 0.5
          * bellman_errors**2
      )
      huber_loss_case_two = (
          tf.cast(tf.abs(bellman_errors) > self.kappa, tf.float32)
          * self.kappa
          * (tf.abs(bellman_errors) - 0.5 * self.kappa)
      )
      huber_loss = huber_loss_case_one + huber_loss_case_two
  
      # Reshape replay_quantiles to batch_size x num_tau_samples x 1
      replay_quantiles = tf.reshape(
          self._replay_net_quantiles,
          [self.num_tau_samples, batch_size, 1],
      )
      replay_quantiles = tf.transpose(replay_quantiles, [1, 0, 2])
  
      # Tile by num_tau_prime_samples along a new dimension. Shape is now
      # batch_size x num_tau_prime_samples x num_tau_samples x 1.
      # These quantiles will be used for computation of the quantile huber loss
      # below (see section 2.3 of the paper).
      replay_quantiles = tf.cast(
          tf.tile(
              replay_quantiles[:, None, :, :],
              [1, self.num_tau_prime_samples, 1, 1],
          ),
          tf.float32,
      )
      # Shape: batch_size x num_tau_prime_samples x num_tau_samples x 1.
      quantile_huber_loss = (
          tf.abs(
              replay_quantiles
              - tf.stop_gradient(
                  tf.cast(bellman_errors < 0, tf.float32)
              )
          )
          * huber_loss
      ) / self.kappa
      # Sum over current quantile value (num_tau_samples) dimension,
      # average over target quantile value (num_tau_prime_samples) dimension.
      # Shape: batch_size x num_tau_prime_samples x 1.
      loss = tf.reduce_sum(quantile_huber_loss, axis=2)
      # Shape: batch_size x 1.
      loss = tf.reduce_mean(loss, axis=1)
  
      # TODO(kumasaurabh): Add prioritized replay functionality here.
      update_priorities_op = tf.no_op()
      with tf.control_dependencies([update_priorities_op]):
          if self.summary_writer is not None:
              with tf.compat.v1.variable_scope("Losses"):
                  tf.compat.v1.summary.scalar(
                      "QuantileLoss", tf.reduce_mean(loss)
                  )
          return self.optimizer.minimize(
              tf.reduce_mean(loss)
          ), tf.reduce_mean(loss)
  
file_path: dopamine/dopamine/agents/implicit_quantile/implicit_quantile_agent.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from dopamine.agents.rainbow import rainbow_agent
  from dopamine.discrete_domains import atari_lib
  import gin.tf
  import tensorflow
label_desc: Use image node_id 5 for calling the ImplicitQuantileAgent obj's underlying member method code with example usage: obj._build_train_op() and returns: self,  tf