code: def from_importance_weights(
      log_rhos: torch.Tensor,
      discounts: torch.Tensor,
      rewards: torch.Tensor,
      values: torch.Tensor,
      bootstrap_value: torch.Tensor,
      clip_rho_threshold: Union[float, None],
      clip_pg_rho_threshold: Union[float, None],
  ) -> VTraceReturns:
      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].
  
      :param 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.
      :param discounts: A float32 tensor of shape [T, B] with discounts encountered when following the behaviour policy.
      :param rewards: A float32 tensor of shape [T, B] containing rewards generated by following the behaviour policy.
      :param values: A float32 tensor of shape [T, B] with the value function estimates wrt. the target policy.
      :param bootstrap_value: A float32 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. If None, no clipping is
                                 applied.
      :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_sfrom_importance_weights)).
                                    If None, no clipping is applied.
  
      :return: 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 = convert_to_torch(
          log_rhos, device="cpu", cast=torch.float32, in_place="try"
      )
      discounts = convert_to_torch(
          discounts, device="cpu", cast=torch.float32, in_place="try"
      )
      rewards = convert_to_torch(
          rewards, device="cpu", cast=torch.float32, in_place="try"
      )
      values = convert_to_torch(
          values, device="cpu", cast=torch.float32, in_place="try"
      )
      bootstrap_value = convert_to_torch(
          bootstrap_value,
          device="cpu",
          cast=torch.float32,
          in_place="try",
      )
  
      # Make sure tensor ranks are consistent.
      rho_rank = len(log_rhos.size())  # Usually 2.
      assert rho_rank == len(values.size())
      assert rho_rank - 1 == len(
          bootstrap_value.size()
      ), "must have rank {}".format(rho_rank - 1)
      assert rho_rank == len(discounts.size())
      assert rho_rank == len(rewards.size())
  
      rhos = torch.exp(log_rhos)
      if clip_rho_threshold is not None:
          # Basically the min operation rho := min(rho_bar, pi(a_t|x_t)/mu(at|x_t))
          clipped_rhos = torch.clamp_max(rhos, clip_rho_threshold)
      else:
          clipped_rhos = rhos
  
      cs = torch.clamp_max(rhos, 1.0)
      # Append bootstrapped value to get [v1, ..., v_t+1]
      values_t_plus_1 = torch.cat(
          [values[1:], torch.unsqueeze(bootstrap_value, 0)], dim=0
      )
      deltas = clipped_rhos * (
          rewards + discounts * values_t_plus_1 - values
      )
  
      vs_minus_v_xs: List = [torch.zeros_like(bootstrap_value)]
      for i in reversed(range(len(discounts))):
          discount_t, c_t, delta_t = discounts[i], cs[i], deltas[i]
          vs_minus_v_xs.append(
              delta_t + discount_t * c_t * vs_minus_v_xs[-1]
          )
      vs_minus_v_xs: torch.Tensor = torch.stack(vs_minus_v_xs[1:])
      # Reverse the results back to original order.
      vs_minus_v_xs: torch.Tensor = torch.flip(vs_minus_v_xs, dims=[0])
      # Add V(x_s) to get v_s.
      vs = vs_minus_v_xs + values
  
      # Advantage for policy gradient.
      vs_t_plus_1 = torch.cat(
          [vs[1:], torch.unsqueeze(bootstrap_value, 0)], dim=0
      )
      if clip_pg_rho_threshold is not None:
          clipped_pg_rhos = torch.clamp_max(rhos, clip_pg_rho_threshold)
      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=vs.detach(), pg_advantages=pg_advantages.detach()
      )
  
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 3 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) and returns: VTraceReturns