code: def from_importance_weights(
      log_rhos,
      discounts,
      rewards,
      values,
      bootstrap_value,
      clip_rho_threshold=1.0,
      clip_pg_rho_threshold=1.0,
      name="vtrace_from_importance_weights",
  ):
      r"""V-trace from log importance weights.
  
      Calculates V-trace actor critic targets as described in
  
      "IMPALA: Scalable Distributed Deep-RL with
      Importance Weighted Actor-Learner Architectures"
      by Espeholt, Soyer, Munos et al.
  
      In the notation used throughout documentation and comments, T refers to the
      time dimension ranging from 0 to T-1. B refers to the batch size. This code
      also supports the case where all tensors have the same number of additional
      dimensions, e.g., `rewards` is [T, B, C], `values` is [T, B, C],
      `bootstrap_value` is [B, C].
  
      Args:
        log_rhos: A float32 tensor of shape [T, B] representing the
          log importance sampling weights, i.e.
          log(target_policy(a) / behaviour_policy(a)). V-trace performs operations
          on rhos in log-space for numerical stability.
        discounts: A float32 tensor of shape [T, B] with discounts encountered when
          following the behaviour policy.
        rewards: A float32 tensor of shape [T, B] containing 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. If None, no clipping is applied.
        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)). If
          None, no clipping is applied.
        name: The name scope that all V-trace operations will be created in.
  
      Returns:
        A VTraceReturns namedtuple (vs, pg_advantages) where:
          vs: A float32 tensor of shape [T, B]. Can be used as target to
            train a baseline (V(x_t) - vs_t)^2.
          pg_advantages: A float32 tensor of shape [T, B]. Can be used as the
            advantage in the calculation of policy gradients.
      """
      log_rhos = tf.convert_to_tensor(log_rhos, dtype=tf.float32)
      discounts = tf.convert_to_tensor(discounts, dtype=tf.float32)
      rewards = tf.convert_to_tensor(rewards, dtype=tf.float32)
      values = tf.convert_to_tensor(values, dtype=tf.float32)
      bootstrap_value = tf.convert_to_tensor(
          bootstrap_value, dtype=tf.float32
      )
      if clip_rho_threshold is not None:
          clip_rho_threshold = tf.convert_to_tensor(
              clip_rho_threshold, dtype=tf.float32
          )
      if clip_pg_rho_threshold is not None:
          clip_pg_rho_threshold = tf.convert_to_tensor(
              clip_pg_rho_threshold, dtype=tf.float32
          )
  
      # Make sure tensor ranks are consistent.
      rho_rank = log_rhos.shape.ndims  # Usually 2.
      values.shape.assert_has_rank(rho_rank)
      bootstrap_value.shape.assert_has_rank(rho_rank - 1)
      discounts.shape.assert_has_rank(rho_rank)
      rewards.shape.assert_has_rank(rho_rank)
      if clip_rho_threshold is not None:
          clip_rho_threshold.shape.assert_has_rank(0)
      if clip_pg_rho_threshold is not None:
          clip_pg_rho_threshold.shape.assert_has_rank(0)
  
      with tf1.name_scope(
          name,
          values=[
              log_rhos,
              discounts,
              rewards,
              values,
              bootstrap_value,
          ],
      ):
          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
  
          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], name="vs_minus_v_xs"
          )
  
          # Add V(x_s) to get v_s.
          vs = tf.add(vs_minus_v_xs, values, name="vs")
  
          # 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, name="clipped_pg_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 VTraceReturns(
              vs=tf.stop_gradient(vs),
              pg_advantages=tf.stop_gradient(pg_advantages),
          )
  
file_path: ray/rllib/algorithms/impala/vtrace_tf.py
imports: import collections
  from ray.rllib.models.tf.tf_action_dist import Categorical
  from ray.rllib.utils.framework import try_import_tf
label_desc: Use image node_id 5 for calling a global function with example usage: from_importance_weights(log_rhos, discounts, rewards, values, bootstrap_value, clip_rho_threshold, clip_pg_rho_threshold, name) and returns: VTraceReturns, unknown