class_docstr: A compact implementation of a simplified Rainbow agent.
class_docstr_tok: ["A","compact","implementation","of","a","simplified","Rainbow","agent","."]
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.
      """
      target_distribution = tf.stop_gradient(
          self._build_target_distribution()
      )
  
      # size of indices: batch_size x 1.
      indices = tf.range(tf.shape(self._replay_net_outputs.logits)[0])[
          :, None
      ]
      # size of reshaped_actions: batch_size x 2.
      reshaped_actions = tf.concat(
          [indices, self._replay.actions[:, None]], 1
      )
      # For each element of the batch, fetch the logits for its selected action.
      chosen_action_logits = tf.gather_nd(
          self._replay_net_outputs.logits, reshaped_actions
      )
  
      loss = tf.nn.softmax_cross_entropy_with_logits(
          labels=target_distribution, logits=chosen_action_logits
      )
  
      if self._replay_scheme == "prioritized":
          # The original prioritized experience replay uses a linear exponent
          # schedule 0.4 -> 1.0. Comparing the schedule to a fixed exponent of 0.5
          # on 5 games (Asterix, Pong, Q*Bert, Seaquest, Space Invaders) suggested
          # a fixed exponent actually performs better, except on Pong.
          probs = self._replay.transition["sampling_probabilities"]
          loss_weights = 1.0 / tf.sqrt(probs + 1e-10)
          loss_weights /= tf.reduce_max(loss_weights)
  
          # Rainbow and prioritized replay are parametrized by an exponent alpha,
          # but in both cases it is set to 0.5 - for simplicity's sake we leave it
          # as is here, using the more direct tf.sqrt(). Taking the square root
          # "makes sense", as we are dealing with a squared loss.
          # Add a small nonzero value to the loss to avoid 0 priority items. While
          # technically this may be okay, setting all items to 0 priority will cause
          # troubles, and also result in 1.0 / 0.0 = NaN correction terms.
          update_priorities_op = self._replay.tf_set_priority(
              self._replay.indices, tf.sqrt(loss + 1e-10)
          )
  
          # Weight the loss by the inverse priorities.
          loss = loss_weights * loss
      else:
          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(
                      "CrossEntropyLoss", tf.reduce_mean(loss)
                  )
          # Schaul et al. reports a slightly different rule, where 1/N is also
          # exponentiated by beta. Not doing so seems more reasonable, and did not
          # impact performance in our experiments.
          return self.optimizer.minimize(tf.reduce_mean(loss)), loss
  
file_path: dopamine/dopamine/agents/rainbow/rainbow_agent.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from dopamine.agents.dqn import dqn_agent
  from dopamine.discrete_domains import atari_lib
  from dopamine.replay_memory import prioritized_replay_buffer
  import gin.tf
  import tensorflow
label_desc: Use image node_id 5 for calling the RainbowAgent obj's underlying member method code with example usage: obj._build_train_op() and returns: self,  loss