class_docstr: A compact implementation of a simplified Rainbow agent.
class_docstr_tok: ["A","compact","implementation","of","a","simplified","Rainbow","agent","."]
docstr: Builds the C51 target distribution as per Bellemare et al. (2017).
  
  First, we compute the support of the Bellman target, r + gamma Z'. Where Z'
  is the support of the next state distribution:
  
    * Evenly spaced in [-vmax, vmax] if the current state is nonterminal;
    * 0 otherwise (duplicated num_atoms times).
  
  Second, we compute the next-state probabilities, corresponding to the action
  with highest expected value.
  
  Finally we project the Bellman target (support + probabilities) onto the
  original support.
  
  Returns:
    target_distribution: tf.tensor, the target distribution from the replay.
docstr_tok: ["Builds","the","C51","target","distribution","as","per","Bellemare","et","al",".","(","2017",")",".","First",",","we","compute","the","support","of","the","Bellman","target",",","r","+","gamma","Z","'",".","Where","Z'","is","the","support","of","the","next","state","distribution",":","*","Evenly","spaced","in","[","-vmax",",","vmax","]","if","the","current","state","is","nonterminal",";","*","0","otherwise","(","duplicated","num_atoms","times",")",".","Second",",","we","compute","the","next-state","probabilities",",","corresponding","to","the","action","with","highest","expected","value",".","Finally","we","project","the","Bellman","target","(","support","+","probabilities",")","onto","the","original","support",".","Returns",":","target_distribution",":","tf.tensor",",","the","target","distribution","from","the","replay","."]
code: def _build_target_distribution(self):
      """Builds the C51 target distribution as per Bellemare et al. (2017).
  
      First, we compute the support of the Bellman target, r + gamma Z'. Where Z'
      is the support of the next state distribution:
  
        * Evenly spaced in [-vmax, vmax] if the current state is nonterminal;
        * 0 otherwise (duplicated num_atoms times).
  
      Second, we compute the next-state probabilities, corresponding to the action
      with highest expected value.
  
      Finally we project the Bellman target (support + probabilities) onto the
      original support.
  
      Returns:
        target_distribution: tf.tensor, the target distribution from the replay.
      """
      batch_size = self._replay.batch_size
  
      # size of rewards: batch_size x 1
      rewards = self._replay.rewards[:, None]
  
      # size of tiled_support: batch_size x num_atoms
      tiled_support = tf.tile(self._support, [batch_size])
      tiled_support = tf.reshape(
          tiled_support, [batch_size, self._num_atoms]
      )
  
      # size of target_support: batch_size x num_atoms
  
      is_terminal_multiplier = 1.0 - tf.cast(
          self._replay.terminals, tf.float32
      )
      # Incorporate terminal state to discount factor.
      # size of gamma_with_terminal: batch_size x 1
      gamma_with_terminal = (
          self.cumulative_gamma * is_terminal_multiplier
      )
      gamma_with_terminal = gamma_with_terminal[:, None]
  
      target_support = rewards + gamma_with_terminal * tiled_support
  
      # size of next_qt_argmax: 1 x batch_size
      next_qt_argmax = tf.argmax(
          self._replay_next_target_net_outputs.q_values, axis=1
      )[:, None]
      batch_indices = tf.range(tf.cast(batch_size, tf.int64))[:, None]
      # size of next_qt_argmax: batch_size x 2
      batch_indexed_next_qt_argmax = tf.concat(
          [batch_indices, next_qt_argmax], axis=1
      )
  
      # size of next_probabilities: batch_size x num_atoms
      next_probabilities = tf.gather_nd(
          self._replay_next_target_net_outputs.probabilities,
          batch_indexed_next_qt_argmax,
      )
  
      return project_distribution(
          target_support, next_probabilities, self._support
      )
  
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 4 for calling the RainbowAgent obj's underlying member method code with example usage: obj._build_target_distribution() and returns: project_distribution