class_docstr: PyTorch policy class used with PGTrainer.
docstr: The basic policy gradients loss function.
  
  Calculates the vanilla policy gradient loss based on:
  L = -E[ log(pi(a|s)) * A]
  
  Args:
      model: The Model to calculate the loss for.
      dist_class: The action distr. class.
      train_batch: The training data.
  
  Returns:
      Union[TensorType, List[TensorType]]: A single loss tensor or a list
          of loss tensors.
docstr_tok: ["The","basic","policy","gradients","loss","function",".","Calculates","the","vanilla","policy","gradient","loss","based","on",":","L","=","-E","[","log","(","pi","(","a|s",")",")","*","A","]","Args",":","model",":","The","Model","to","calculate","the","loss","for",".","dist_class",":","The","action","distr",".","class",".","train_batch",":","The","training","data",".","Returns",":","Union","[","TensorType",",","List","[","TensorType","]","]",":","A","single","loss","tensor","or","a","list","of","loss","tensors","."]
code: def loss(
      self,
      model: ModelV2,
      dist_class: Type[TorchDistributionWrapper],
      train_batch: SampleBatch,
  ) -> Union[TensorType, List[TensorType]]:
      """The basic policy gradients loss function.
  
      Calculates the vanilla policy gradient loss based on:
      L = -E[ log(pi(a|s)) * A]
  
      Args:
          model: The Model to calculate the loss for.
          dist_class: The action distr. class.
          train_batch: The training data.
  
      Returns:
          Union[TensorType, List[TensorType]]: A single loss tensor or a list
              of loss tensors.
      """
      # Pass the training data through our model to get distribution parameters.
      dist_inputs, _ = model(train_batch)
  
      # Create an action distribution object.
      action_dist = dist_class(dist_inputs, model)
  
      # Calculate the vanilla PG loss based on:
      # L = -E[ log(pi(a|s)) * A]
      log_probs = action_dist.logp(train_batch[SampleBatch.ACTIONS])
  
      # Final policy loss.
      policy_loss = -torch.mean(
          log_probs * train_batch[Postprocessing.ADVANTAGES]
      )
  
      # Store values for stats function in model (tower), such that for
      # multi-GPU, we do not override them during the parallel loss phase.
      model.tower_stats["policy_loss"] = policy_loss
  
      return policy_loss
  
imports: import logging
  from typing import Dict, List, Optional, Tuple, Type, Union
  from rllib_pg.pg.pg import PGConfig
  from rllib_pg.pg.utils import post_process_advantages
  from ray.rllib.evaluation.episode import Episode
  from ray.rllib.evaluation.postprocessing import Postprocessing
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.torch.torch_action_dist import TorchDistributionWrapper
  from ray.rllib.policy import Policy
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.policy.torch_mixins import LearningRateSchedule
  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.typing import AgentID, TensorType
label_desc: Use image node_id 2 for calling the PGTorchPolicy obj's underlying member method code with example usage: obj.loss(model, dist_class, train_batch) and returns: policy_loss