code: def vtrace_from_logits(
      behaviour_policy_logits,
      target_policy_logits,
      actions,
      discounts,
      rewards,
      values,
      bootstrap_value,
      clip_rho_threshold=1.0,
      clip_pg_rho_threshold=1.0,
      name="vtrace_from_logits",
  ):
      r"""V-trace for softmax policies.
  
      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.
  
      Target policy refers to the policy we are interested in improving and
      behaviour policy refers to the policy that generated the given
      rewards and actions.
  
      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 and
      `NUM_ACTIONS` refers to the number of actions.
  
      Args:
        behaviour_policy_logits: A float32 tensor of shape `[T, B, NUM_ACTIONS]`
          with un-normalized log-probabilities parametrizing the softmax behaviour
          policy.
        target_policy_logits: A float32 tensor of shape `[T, B, NUM_ACTIONS]` with
          un-normalized log-probabilities parametrizing the softmax target policy.
        actions: An int32 tensor of shape `[T, B]` of actions sampled from the
          behaviour policy.
        discounts: A float32 tensor of shape `[T, B]` with the discount encountered
          when following the behaviour policy.
        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)).
        name: The name scope that all V-trace operations will be created in.
  
      Returns:
        A `VTraceFromLogitsReturns` namedtuple with the following fields:
          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 float 32 tensor of shape `[T, B]`. Can be used as an
            estimate of the advantage in the calculation of policy gradients.
          log_rhos: A float32 tensor of shape `[T, B]` containing the log importance
            sampling weights (log rhos).
          behaviour_action_log_probs: A float32 tensor of shape `[T, B]` containing
            behaviour policy action log probabilities (log \mu(a_t)).
          target_action_log_probs: A float32 tensor of shape `[T, B]` containing
            target policy action probabilities (log \pi(a_t)).
      """
      behaviour_policy_logits = tf.convert_to_tensor(
          behaviour_policy_logits, dtype=tf.float32
      )
      target_policy_logits = tf.convert_to_tensor(
          target_policy_logits, dtype=tf.float32
      )
      actions = tf.convert_to_tensor(actions, dtype=tf.int32)
  
      # Make sure tensor ranks are as expected.
      # The rest will be checked by from_action_log_probs.
      behaviour_policy_logits.shape.assert_has_rank(3)
      target_policy_logits.shape.assert_has_rank(3)
      actions.shape.assert_has_rank(2)
  
      with tf.name_scope(
          name,
          values=[
              behaviour_policy_logits,
              target_policy_logits,
              actions,
              discounts,
              rewards,
              values,
              bootstrap_value,
          ],
      ):
          target_action_log_probs = log_probs_from_logits_and_actions(
              target_policy_logits, actions
          )
          behaviour_action_log_probs = (
              log_probs_from_logits_and_actions(
                  behaviour_policy_logits, actions
              )
          )
          log_rhos = (
              target_action_log_probs - behaviour_action_log_probs
          )
          vtrace_returns = vtrace_from_importance_weights(
              log_rhos=log_rhos,
              discounts=discounts,
              rewards=rewards,
              values=values,
              bootstrap_value=bootstrap_value,
              clip_rho_threshold=clip_rho_threshold,
              clip_pg_rho_threshold=clip_pg_rho_threshold,
          )
          return VTraceFromLogitsReturns(
              log_rhos=log_rhos,
              behaviour_action_log_probs=behaviour_action_log_probs,
              target_action_log_probs=target_action_log_probs,
              **vtrace_returns._asdict(),
          )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import tensorflow.compat.v1
label_desc: Use image node_id 2 for calling a global function with example usage: vtrace_from_logits(behaviour_policy_logits, target_policy_logits, actions, discounts, rewards, values, bootstrap_value, clip_rho_threshold, clip_pg_rho_threshold, name) and returns: VTraceFromLogitsReturns