code: def vtrace_tf2(
      *,
      target_action_log_probs: "tf.Tensor",
      behaviour_action_log_probs: "tf.Tensor",
      discounts: "tf.Tensor",
      rewards: "tf.Tensor",
      values: "tf.Tensor",
      bootstrap_value: "tf.Tensor",
      clip_rho_threshold: Union[float, "tf.Tensor"] = 1.0,
      clip_pg_rho_threshold: Union[float, "tf.Tensor"] = 1.0,
  ):
      r"""V-trace for softmax policies implemented with tensorflow.
  
      Calculates V-trace actor critic targets for softmax polices as described in
      "IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner
      Architectures" by Espeholt, Soyer, Munos et al. (https://arxiv.org/abs/1802.01561)
  
      The V-trace implementation used here closely resembles the one found in the
      scalable-agent repository by Google DeepMind, available at
      https://github.com/deepmind/scalable_agent. This version has been optimized to
      minimize the number of floating-point operations required per V-Trace
      calculation, achieved through the use of dynamic programming techniques. It's
      important to note that the mathematical expressions used in this implementation
      may appear quite different from those presented in the IMPALA paper.
  
      The following terminology applies:
          - `target policy` refers to the policy we are interested in improving.
          - `behaviour policy` refers to the policy that generated the given
              rewards and actions.
          - `T` refers to the time dimension. This is usually either the length of the
              trajectory or the length of the sequence if recurrent.
          - `B` refers to the batch size.
  
      Args:
          target_action_log_probs: Action log probs from the target policy. A float32
              tensor of shape [T, B].
          behaviour_action_log_probs: Action log probs from the behaviour policy. A
              float32 tensor of shape [T, B].
          discounts: A float32 tensor of shape [T, B] with the discount encountered when
              following the behaviour policy. This will be 0 for terminal timesteps
              (done=True) and gamma (the discount factor) otherwise.
          rewards: A float32 tensor of shape [T, B] with the rewards generated by
              following the behaviour policy.
          values: A float32 tensor of shape [T, B] with the value function estimates
              wrt. the target policy.
          bootstrap_value: A float32 of shape [B] with the value function estimate at
              time T.
          clip_rho_threshold: A scalar float32 tensor with the clipping threshold for
              importance weights (rho) when calculating the baseline targets (vs).
              rho^bar in the paper.
          clip_pg_rho_threshold: A scalar float32 tensor with the clipping threshold
              on rho_s in \rho_s \delta log \pi(a|x) (r + \gamma v_{s+1} - V(x_s)).
      """
      log_rhos = target_action_log_probs - behaviour_action_log_probs
  
      rhos = tf.math.exp(log_rhos)
      if clip_rho_threshold is not None:
          clipped_rhos = tf.minimum(
              clip_rho_threshold, rhos, name="clipped_rhos"
          )
      else:
          clipped_rhos = rhos
  
      cs = tf.minimum(1.0, rhos, name="cs")
      # Append bootstrapped value to get [v1, ..., v_t+1]
      values_t_plus_1 = tf.concat(
          [values[1:], tf.expand_dims(bootstrap_value, 0)], axis=0
      )
  
      deltas = clipped_rhos * (
          rewards + discounts * values_t_plus_1 - values
      )
  
      # All sequences are reversed, computation starts from the back.
      sequences = (
          tf.reverse(discounts, axis=[0]),
          tf.reverse(cs, axis=[0]),
          tf.reverse(deltas, axis=[0]),
      )
  
      # V-trace vs are calculated through a scan from the back to the
      # beginning of the given trajectory.
      def scanfunc(acc, sequence_item):
          discount_t, c_t, delta_t = sequence_item
          return delta_t + discount_t * c_t * acc
  
      with tf.device("/cpu:0"):
          initial_values = tf.zeros_like(bootstrap_value)
          vs_minus_v_xs = tf.nest.map_structure(
              tf.stop_gradient,
              tf.scan(
                  fn=scanfunc,
                  elems=sequences,
                  initializer=initial_values,
                  parallel_iterations=1,
                  name="scan",
              ),
          )
      # Reverse the results back to original order.
      vs_minus_v_xs = tf.reverse(vs_minus_v_xs, [0])
  
      # Add V(x_s) to get v_s.
      vs = tf.add(vs_minus_v_xs, values)
  
      # Advantage for policy gradient.
      vs_t_plus_1 = tf.concat(
          [vs[1:], tf.expand_dims(bootstrap_value, 0)], axis=0
      )
      if clip_pg_rho_threshold is not None:
          clipped_pg_rhos = tf.minimum(clip_pg_rho_threshold, rhos)
      else:
          clipped_pg_rhos = rhos
      pg_advantages = clipped_pg_rhos * (
          rewards + discounts * vs_t_plus_1 - values
      )
  
      # Make sure no gradients backpropagated through the returned values.
      return tf.stop_gradient(vs), tf.stop_gradient(pg_advantages)
  
label_desc: Use image node_id 2 for calling a global function with example usage: vtrace_tf2() and returns: tf,  tf, unknown