def multi_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
  ACTION_SPACE refers to the list of numbers each representing a number of
  actions.

  Args:
    behaviour_policy_logits: A list with length of ACTION_SPACE of float32
      tensors of shapes
      [T, B, ACTION_SPACE[0]],
      ...,
      [T, B, ACTION_SPACE[-1]]
      with un-normalized log-probabilities parameterizing the softmax behaviour
      policy.
    target_policy_logits: A list with length of ACTION_SPACE of float32
      tensors of shapes
      [T, B, ACTION_SPACE[0]],
      ...,
      [T, B, ACTION_SPACE[-1]]
      with un-normalized log-probabilities parameterizing the softmax target
      policy.
    actions: A list with length of ACTION_SPACE of int32
      tensors of shapes
      [T, B],
      ...,
      [T, B]
      with 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)).
  """

    for i in range(len(behaviour_policy_logits)):
        behaviour_policy_logits[i] = tf.convert_to_tensor(
            behaviour_policy_logits[i], dtype=tf.float32)
        target_policy_logits[i] = tf.convert_to_tensor(
            target_policy_logits[i], dtype=tf.float32)
        actions[i] = tf.convert_to_tensor(actions[i], dtype=tf.int32)

        # Make sure tensor ranks are as expected.
        # The rest will be checked by from_action_log_probs.
        behaviour_policy_logits[i].shape.assert_has_rank(3)
        target_policy_logits[i].shape.assert_has_rank(3)
        actions[i].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 = multi_log_probs_from_logits_and_actions(
            target_policy_logits, actions)
        behaviour_action_log_probs = multi_log_probs_from_logits_and_actions(
            behaviour_policy_logits, actions)

        log_rhos = get_log_rhos(target_action_log_probs,
                                behaviour_action_log_probs)

        vtrace_returns = 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())
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
  ACTION_SPACE refers to the list of numbers each representing a number of
  actions.

  Args:
    behaviour_policy_logits: A list with length of ACTION_SPACE of float32
      tensors of shapes
      [T, B, ACTION_SPACE[0]],
      ...,
      [T, B, ACTION_SPACE[-1]]
      with un-normalized log-probabilities parameterizing the softmax behaviour
      policy.
    target_policy_logits: A list with length of ACTION_SPACE of float32
      tensors of shapes
      [T, B, ACTION_SPACE[0]],
      ...,
      [T, B, ACTION_SPACE[-1]]
      with un-normalized log-probabilities parameterizing the softmax target
      policy.
    actions: A list with length of ACTION_SPACE of int32
      tensors of shapes
      [T, B],
      ...,
      [T, B]
      with 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)).