docstr: Policy gradient loss with vtrace importance weighting.
  
  VTraceLoss takes tensors of shape [T, B, ...], where `B` is the
  batch_size. The reason we need to know `B` is for V-trace to properly
  handle episode cut boundaries.
  
  Args:
      actions: An int|float32 tensor of shape [T, B, ACTION_SPACE].
      actions_logp: A float32 tensor of shape [T, B].
      actions_entropy: A float32 tensor of shape [T, B].
      dones: A bool tensor of shape [T, B].
      behaviour_action_logp: Tensor of shape [T, B].
      behaviour_logits: A list with length of ACTION_SPACE of float32
          tensors of shapes
          [T, B, ACTION_SPACE[0]],
          ...,
          [T, B, ACTION_SPACE[-1]]
      target_logits: A list with length of ACTION_SPACE of float32
          tensors of shapes
          [T, B, ACTION_SPACE[0]],
          ...,
          [T, B, ACTION_SPACE[-1]]
      discount: A float32 scalar.
      rewards: A float32 tensor of shape [T, B].
      values: A float32 tensor of shape [T, B].
      bootstrap_value: A float32 tensor of shape [B].
      dist_class: action distribution class for logits.
      valid_mask: A bool tensor of valid RNN input elements (#2992).
      config: Algorithm config dict.
docstr_tok: ["Policy","gradient","loss","with","vtrace","importance","weighting",".","VTraceLoss","takes","tensors","of","shape","[","T",",","B",",","...","]",",","where","`","B","`","is","the","batch_size",".","The","reason","we","need","to","know","`","B","`","is","for","V-trace","to","properly","handle","episode","cut","boundaries",".","Args",":","actions",":","An","int|float32","tensor","of","shape","[","T",",","B",",","ACTION_SPACE","]",".","actions_logp",":","A","float32","tensor","of","shape","[","T",",","B","]",".","actions_entropy",":","A","float32","tensor","of","shape","[","T",",","B","]",".","dones",":","A","bool","tensor","of","shape","[","T",",","B","]",".","behaviour_action_logp",":","Tensor","of","shape","[","T",",","B","]",".","behaviour_logits",":","A","list","with","length","of","ACTION_SPACE","of","float32","tensors","of","shapes","[","T",",","B",",","ACTION_SPACE","[","0","]","]",",","...",",","[","T",",","B",",","ACTION_SPACE","[","-1","]","]","target_logits",":","A","list","with","length","of","ACTION_SPACE","of","float32","tensors","of","shapes","[","T",",","B",",","ACTION_SPACE","[","0","]","]",",","...",",","[","T",",","B",",","ACTION_SPACE","[","-1","]","]","discount",":","A","float32","scalar",".","rewards",":","A","float32","tensor","of","shape","[","T",",","B","]",".","values",":","A","float32","tensor","of","shape","[","T",",","B","]",".","bootstrap_value",":","A","float32","tensor","of","shape","[","B","]",".","dist_class",":","action","distribution","class","for","logits",".","valid_mask",":","A","bool","tensor","of","valid","RNN","input","elements","(","#","2992",")",".","config",":","Algorithm","config","dict","."]
code: def __init__(
      self,
      actions,
      actions_logp,
      actions_entropy,
      dones,
      behaviour_action_logp,
      behaviour_logits,
      target_logits,
      discount,
      rewards,
      values,
      bootstrap_value,
      dist_class,
      model,
      valid_mask,
      config,
      vf_loss_coeff=0.5,
      entropy_coeff=0.01,
      clip_rho_threshold=1.0,
      clip_pg_rho_threshold=1.0,
  ):
      """Policy gradient loss with vtrace importance weighting.
  
      VTraceLoss takes tensors of shape [T, B, ...], where `B` is the
      batch_size. The reason we need to know `B` is for V-trace to properly
      handle episode cut boundaries.
  
      Args:
          actions: An int|float32 tensor of shape [T, B, ACTION_SPACE].
          actions_logp: A float32 tensor of shape [T, B].
          actions_entropy: A float32 tensor of shape [T, B].
          dones: A bool tensor of shape [T, B].
          behaviour_action_logp: Tensor of shape [T, B].
          behaviour_logits: A list with length of ACTION_SPACE of float32
              tensors of shapes
              [T, B, ACTION_SPACE[0]],
              ...,
              [T, B, ACTION_SPACE[-1]]
          target_logits: A list with length of ACTION_SPACE of float32
              tensors of shapes
              [T, B, ACTION_SPACE[0]],
              ...,
              [T, B, ACTION_SPACE[-1]]
          discount: A float32 scalar.
          rewards: A float32 tensor of shape [T, B].
          values: A float32 tensor of shape [T, B].
          bootstrap_value: A float32 tensor of shape [B].
          dist_class: action distribution class for logits.
          valid_mask: A bool tensor of valid RNN input elements (#2992).
          config: Algorithm config dict.
      """
      import ray.rllib.algorithms.impala.vtrace_torch as vtrace
  
      if valid_mask is None:
          valid_mask = torch.ones_like(actions_logp)
  
      # Compute vtrace on the CPU for better perf
      # (devices handled inside `vtrace.multi_from_logits`).
      device = behaviour_action_logp[0].device
      self.vtrace_returns = vtrace.multi_from_logits(
          behaviour_action_log_probs=behaviour_action_logp,
          behaviour_policy_logits=behaviour_logits,
          target_policy_logits=target_logits,
          actions=torch.unbind(actions, dim=2),
          discounts=(1.0 - dones.float()) * discount,
          rewards=rewards,
          values=values,
          bootstrap_value=bootstrap_value,
          dist_class=dist_class,
          model=model,
          clip_rho_threshold=clip_rho_threshold,
          clip_pg_rho_threshold=clip_pg_rho_threshold,
      )
      # Move v-trace results back to GPU for actual loss computing.
      self.value_targets = self.vtrace_returns.vs.to(device)
  
      # The policy gradients loss.
      self.pi_loss = -torch.sum(
          actions_logp
          * self.vtrace_returns.pg_advantages.to(device)
          * valid_mask
      )
  
      # The baseline loss.
      delta = (values - self.value_targets) * valid_mask
      self.vf_loss = 0.5 * torch.sum(torch.pow(delta, 2.0))
  
      # The entropy loss.
      self.entropy = torch.sum(actions_entropy * valid_mask)
      self.mean_entropy = self.entropy / torch.sum(valid_mask)
  
      # The summed weighted loss.
      self.total_loss = self.pi_loss - self.entropy * entropy_coeff
  
      # Optional vf loss (or in a separate term due to separate
      # optimizers/networks).
      self.loss_wo_vf = self.total_loss
      if not config["_separate_vf_optimizer"]:
          self.total_loss += self.vf_loss * vf_loss_coeff
  
imports: import gymnasium
  import logging
  import numpy
  from typing import Dict, List, Optional, Type, Union
  import ray
  from ray.rllib.evaluation.episode import Episode
  from ray.rllib.evaluation.postprocessing import compute_bootstrap_value
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.action_dist import ActionDistribution
  from ray.rllib.models.torch.torch_action_dist import TorchCategorical
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.policy.torch_mixins import EntropyCoeffSchedule, LearningRateSchedule, ValueNetworkMixin
  from ray.rllib.policy.torch_policy_v2 import TorchPolicyV2
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.numpy import convert_to_numpy
  from ray.rllib.utils.torch_utils import apply_grad_clipping, explained_variance, global_norm, sequence_mask
  from ray.rllib.utils.typing import TensorType
label_desc: Use image node_id 1 to create a new VTraceLoss object  with example: obj = VTraceLoss(actions, actions_logp, actions_entropy, dones, behaviour_action_logp, behaviour_logits, target_logits, discount, rewards, values, bootstrap_value, dist_class, model, valid_mask, config, vf_loss_coeff, entropy_coeff, clip_rho_threshold, clip_pg_rho_threshold)