code: def from_logits(
      behaviour_policy_logits: List[TorchActionType],
      target_policy_logits: List[TorchActionType],
      actions: List[TorchActionType],
      distribution_mapper: DistributionMapper,
      discounts: torch.Tensor,
      rewards: torch.Tensor,
      values: List[torch.Tensor],
      bootstrap_value: List[torch.Tensor],
      clip_rho_threshold: Union[float, None],
      clip_pg_rho_threshold: Union[float, None],
      device: Union[str, None],
  ) -> VTraceFromLogitsReturns:
      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.
  
      :param behaviour_policy_logits: A list (w.r.t. the substeps of the env) of dict (w.r.t. the actions) of tensors
                                      of un-normalized log-probabilities (shape list[dict[str,[T, B, NUM_ACTIONS]]])
      :param target_policy_logits: A list (w.r.t. the substeps of the env) of dict (w.r.t. the actions) of tensors
                                   of un-normalized log-probabilities (shape list[dict[str,[T, B, NUM_ACTIONS]]])
      :param actions: An list (w.r.t. the substeps of the env) of dicts (w.r.t. the actions) with actions sampled from
                      the behavior policy. (list[dict[str,[T, B]]])
      :param distribution_mapper: A distribution mapper providing a mapping of action heads to distributions.
      :param discounts: A float32 tensor of shape [T, B] with the discount encountered when following the behavior policy.
      :param rewards: A float32 tensor of shape [T, B] with the rewards generated by following the behavior policy.
      :param values: A list (w.r.t. the substeps of the env) of float32 tensors of shape [T, B] with the value function
                     estimates wrt. the target policy.
      :param bootstrap_value: A list (w.r.t. the substeps of the env) of float32 tensors of shape [B] with the value
                              function estimate at time T.
      :param 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.
      :param 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)).
      :param device: the device the results should be sent to before returning it
  
      :return: A `VTraceFromLogitsReturns` namedtuple with the following fields:
               vs: A list (w.r.t. the substeps of the env) of float32 tensors of shape [T, B]. Can be used as target to
               train a baseline (V(x_t) - vs_t)^2.
               pg_advantages: A list (w.r.t. the substeps of the env) of float32 tensors of shape [T, B]. Can be used as
               an estimate of the advantage in the calculation of policy gradients.
               log_rhos: A list (w.r.t. the substeps of the env) of float32 tensors of shape [T, B] containing the log
               importance sampling weights (log rhos).
               behaviour_action_log_probs: A list (w.r.t. the substeps of the env) of float32 tensors of shape [T, B]
               containing the behaviour policy action log probabilities (log \mu(a_t)).
               target_action_log_probs: A list (w.r.t. the substeps of the env) of float32 tensors of shape [T, B]
               containing target policy action probabilities (log \pi(a_t)).
               target_step_action_dists: A list (w.r.t. the substeps of the env) of the action probability distributions
               w.r.t. to the target policy
      """
  
      (
          behaviour_action_log_probs,
          _,
      ) = log_probs_from_logits_and_actions_and_spaces(
          behaviour_policy_logits, actions, distribution_mapper
      )
      (
          target_action_log_probs,
          target_step_action_dists,
      ) = log_probs_from_logits_and_actions_and_spaces(
          target_policy_logits, actions, distribution_mapper
      )
  
      log_rhos = get_log_rhos(
          target_action_log_probs=target_action_log_probs,
          behaviour_action_log_probs=behaviour_action_log_probs,
      )
      vss = list()
      pg_advantagess = list()
      for step_log_rhos, step_values, step_bootstrap_values in zip(
          log_rhos, values, bootstrap_value
      ):
          vs, pg_advantages = from_importance_weights(
              log_rhos=step_log_rhos,
              discounts=discounts,
              rewards=rewards,
              values=step_values,
              bootstrap_value=step_bootstrap_values,
              clip_rho_threshold=clip_rho_threshold,
              clip_pg_rho_threshold=clip_pg_rho_threshold,
          )
          vss.append(vs.to(device) if device is not None else vs)
          pg_advantagess.append(
              pg_advantages.to(device)
              if device is not None
              else pg_advantages
          )
  
      return VTraceFromLogitsReturns(
          log_rhos=log_rhos,
          behaviour_action_log_probs=behaviour_action_log_probs,
          target_action_log_probs=target_action_log_probs,
          vs=vss,
          pg_advantages=pg_advantagess,
          target_step_action_dists=target_step_action_dists,
      )
  
file_path: maze/maze/train/trainers/impala/impala_vtrace.py
imports: import collections
  from typing import List, Tuple, Union
  import torch
  from maze.core.env.action_conversion import TorchActionType
  from maze.distributions.dict import DictProbabilityDistribution
  from maze.distributions.distribution_mapper import DistributionMapper
  from maze.perception.perception_utils import convert_to_torch
label_desc: Use image node_id 2 for calling a global function with example usage: from_logits(behaviour_policy_logits, target_policy_logits, actions, distribution_mapper, discounts, rewards, values, bootstrap_value, clip_rho_threshold, clip_pg_rho_threshold, device) and returns: VTraceFromLogitsReturns