code: def project_distribution(
      supports, weights, target_support, validate_args=False
  ):
      """Projects a batch of (support, weights) onto target_support.
  
      Based on equation (7) in (Bellemare et al., 2017):
        https://arxiv.org/abs/1707.06887
      In the rest of the comments we will refer to this equation simply as Eq7.
  
      This code is not easy to digest, so we will use a running example to clarify
      what is going on, with the following sample inputs:
  
        * supports =       [[0, 2, 4, 6, 8],
                            [1, 3, 4, 5, 6]]
        * weights =        [[0.1, 0.6, 0.1, 0.1, 0.1],
                            [0.1, 0.2, 0.5, 0.1, 0.1]]
        * target_support = [4, 5, 6, 7, 8]
  
      In the code below, comments preceded with 'Ex:' will be referencing the above
      values.
  
      Args:
        supports: Tensor of shape (batch_size, num_dims) defining supports for the
          distribution.
        weights: Tensor of shape (batch_size, num_dims) defining weights on the
          original support points. Although for the CategoricalDQN agent these
          weights are probabilities, it is not required that they are.
        target_support: Tensor of shape (num_dims) defining support of the projected
          distribution. The values must be monotonically increasing. Vmin and Vmax
          will be inferred from the first and last elements of this tensor,
          respectively. The values in this tensor must be equally spaced.
        validate_args: Whether we will verify the contents of the
          target_support parameter.
  
      Returns:
        A Tensor of shape (batch_size, num_dims) with the projection of a batch of
        (support, weights) onto target_support.
  
      Raises:
        ValueError: If target_support has no dimensions, or if shapes of supports,
          weights, and target_support are incompatible.
      """
      target_support_deltas = target_support[1:] - target_support[:-1]
      # delta_z = `\Delta z` in Eq7.
      delta_z = target_support_deltas[0]
      validate_deps = []
      supports.shape.assert_is_compatible_with(weights.shape)
      supports[0].shape.assert_is_compatible_with(target_support.shape)
      target_support.shape.assert_has_rank(1)
      if validate_args:
          # Assert that supports and weights have the same shapes.
          validate_deps.append(
              tf.Assert(
                  tf.reduce_all(
                      tf.equal(tf.shape(supports), tf.shape(weights))
                  ),
                  [supports, weights],
              )
          )
          # Assert that elements of supports and target_support have the same shape.
          validate_deps.append(
              tf.Assert(
                  tf.reduce_all(
                      tf.equal(
                          tf.shape(supports)[1],
                          tf.shape(target_support),
                      )
                  ),
                  [supports, target_support],
              )
          )
          # Assert that target_support has a single dimension.
          validate_deps.append(
              tf.Assert(
                  tf.equal(tf.size(tf.shape(target_support)), 1),
                  [target_support],
              )
          )
          # Assert that the target_support is monotonically increasing.
          validate_deps.append(
              tf.Assert(
                  tf.reduce_all(target_support_deltas > 0),
                  [target_support],
              )
          )
          # Assert that the values in target_support are equally spaced.
          validate_deps.append(
              tf.Assert(
                  tf.reduce_all(
                      tf.equal(target_support_deltas, delta_z)
                  ),
                  [target_support],
              )
          )
  
      with tf.control_dependencies(validate_deps):
          # Ex: `v_min, v_max = 4, 8`.
          v_min, v_max = target_support[0], target_support[-1]
          # Ex: `batch_size = 2`.
          batch_size = tf.shape(supports)[0]
          # `N` in Eq7.
          # Ex: `num_dims = 5`.
          num_dims = tf.shape(target_support)[0]
          # clipped_support = `[\hat{T}_{z_j}]^{V_max}_{V_min}` in Eq7.
          # Ex: `clipped_support = [[[ 4.  4.  4.  6.  8.]]
          #                         [[ 4.  4.  4.  5.  6.]]]`.
          clipped_support = tf.clip_by_value(supports, v_min, v_max)[
              :, None, :
          ]
          # Ex: `tiled_support = [[[[ 4.  4.  4.  6.  8.]
          #                         [ 4.  4.  4.  6.  8.]
          #                         [ 4.  4.  4.  6.  8.]
          #                         [ 4.  4.  4.  6.  8.]
          #                         [ 4.  4.  4.  6.  8.]]
          #                        [[ 4.  4.  4.  5.  6.]
          #                         [ 4.  4.  4.  5.  6.]
          #                         [ 4.  4.  4.  5.  6.]
          #                         [ 4.  4.  4.  5.  6.]
          #                         [ 4.  4.  4.  5.  6.]]]]`.
          tiled_support = tf.tile(
              [clipped_support], [1, 1, num_dims, 1]
          )
          # Ex: `reshaped_target_support = [[[ 4.]
          #                                  [ 5.]
          #                                  [ 6.]
          #                                  [ 7.]
          #                                  [ 8.]]
          #                                 [[ 4.]
          #                                  [ 5.]
          #                                  [ 6.]
          #                                  [ 7.]
          #                                  [ 8.]]]`.
          reshaped_target_support = tf.tile(
              target_support[:, None], [batch_size, 1]
          )
          reshaped_target_support = tf.reshape(
              reshaped_target_support, [batch_size, num_dims, 1]
          )
          # numerator = `|clipped_support - z_i|` in Eq7.
          # Ex: `numerator = [[[[ 0.  0.  0.  2.  4.]
          #                     [ 1.  1.  1.  1.  3.]
          #                     [ 2.  2.  2.  0.  2.]
          #                     [ 3.  3.  3.  1.  1.]
          #                     [ 4.  4.  4.  2.  0.]]
          #                    [[ 0.  0.  0.  1.  2.]
          #                     [ 1.  1.  1.  0.  1.]
          #                     [ 2.  2.  2.  1.  0.]
          #                     [ 3.  3.  3.  2.  1.]
          #                     [ 4.  4.  4.  3.  2.]]]]`.
          numerator = tf.abs(tiled_support - reshaped_target_support)
          quotient = 1 - (numerator / delta_z)
          # clipped_quotient = `[1 - numerator / (\Delta z)]_0^1` in Eq7.
          # Ex: `clipped_quotient = [[[[ 1.  1.  1.  0.  0.]
          #                            [ 0.  0.  0.  0.  0.]
          #                            [ 0.  0.  0.  1.  0.]
          #                            [ 0.  0.  0.  0.  0.]
          #                            [ 0.  0.  0.  0.  1.]]
          #                           [[ 1.  1.  1.  0.  0.]
          #                            [ 0.  0.  0.  1.  0.]
          #                            [ 0.  0.  0.  0.  1.]
          #                            [ 0.  0.  0.  0.  0.]
          #                            [ 0.  0.  0.  0.  0.]]]]`.
          clipped_quotient = tf.clip_by_value(quotient, 0, 1)
          # Ex: `weights = [[ 0.1  0.6  0.1  0.1  0.1]
          #                 [ 0.1  0.2  0.5  0.1  0.1]]`.
          weights = weights[:, None, :]
          # inner_prod = `\sum_{j=0}^{N-1} clipped_quotient * p_j(x', \pi(x'))`
          # in Eq7.
          # Ex: `inner_prod = [[[[ 0.1  0.6  0.1  0.  0. ]
          #                      [ 0.   0.   0.   0.  0. ]
          #                      [ 0.   0.   0.   0.1 0. ]
          #                      [ 0.   0.   0.   0.  0. ]
          #                      [ 0.   0.   0.   0.  0.1]]
          #                     [[ 0.1  0.2  0.5  0.  0. ]
          #                      [ 0.   0.   0.   0.1 0. ]
          #                      [ 0.   0.   0.   0.  0.1]
          #                      [ 0.   0.   0.   0.  0. ]
          #                      [ 0.   0.   0.   0.  0. ]]]]`.
          inner_prod = clipped_quotient * weights
          # Ex: `projection = [[ 0.8 0.0 0.1 0.0 0.1]
          #                    [ 0.8 0.1 0.1 0.0 0.0]]`.
          projection = tf.reduce_sum(inner_prod, 3)
          projection = tf.reshape(projection, [batch_size, num_dims])
          return projection
  
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 1 for calling a global function with example usage: project_distribution(supports, weights, target_support, validate_args) and returns: projection