code: def log_probs_from_logits_and_actions_and_spaces(
      policy_logits: List[TorchActionType],
      actions: List[TorchActionType],
      distribution_mapper: DistributionMapper,
  ) -> Tuple[List[TorchActionType], List[DictProbabilityDistribution]]:
      """Computes action log-probs from policy logits, actions and acton_spaces.
  
      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.
  
      :param policy_logits: A list (w.r.t. the substeps of the env) of dicts (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) of tensors
          (list[dict[str,[T, B]]])
      :param distribution_mapper: A distribution mapper providing a mapping of action heads to distributions.
  
      :return: A list (w.r.t. the substeps of the env) of dicts (w.r.t. the actions) of tensors of shape [T, B]
          corresponding to the sampling log probability of the chosen action w.r.t. the policy.
          And a list (w.r.t. the substeps of the env) of DictProbability distributions corresponding to the step-action-
          distributions.
      """
      log_probs = list()
      step_action_dists = list()
      for step_policy_logits, step_actions in zip(
          policy_logits, actions
      ):
          step_action_dist = (
              distribution_mapper.logits_dict_to_distribution(
                  logits_dict=step_policy_logits, temperature=1.0
              )
          )
          log_probs.append(step_action_dist.log_prob(step_actions))
          step_action_dists.append(step_action_dist)
      return log_probs, 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 1 for calling a global function with example usage: log_probs_from_logits_and_actions_and_spaces(policy_logits, actions, distribution_mapper) and returns: log_probs,  step_action_dists