class_docstr: PyTorch policy class used with PG.
class_docstr_tok: ["PyTorch","policy","class","used","with","PG","."]
docstr: Returns the calculated loss in a stats dict.
  
  Args:
      policy: The Policy object.
      train_batch: The data used for training.
  
  Returns:
      Dict[str, TensorType]: The stats dict.
docstr_tok: ["Returns","the","calculated","loss","in","a","stats","dict",".","Args",":","policy",":","The","Policy","object",".","train_batch",":","The","data","used","for","training",".","Returns",":","Dict","[","str",",","TensorType","]",":","The","stats","dict","."]
code: def stats_fn(self, train_batch: SampleBatch) -> Dict[str, TensorType]:
      """Returns the calculated loss in a stats dict.
  
      Args:
          policy: The Policy object.
          train_batch: The data used for training.
  
      Returns:
          Dict[str, TensorType]: The stats dict.
      """
  
      return convert_to_numpy(
          {
              "policy_loss": torch.mean(
                  torch.stack(self.get_tower_stats("policy_loss"))
              ),
              "cur_lr": self.cur_lr,
          }
      )
  
imports: import logging
  from typing import Dict, List, Type, Union, Optional, Tuple
  from ray.rllib.evaluation.episode import Episode
  from ray.rllib.utils.typing import AgentID
  from ray.rllib.policy.torch_policy_v2 import TorchPolicyV2
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils.numpy import convert_to_numpy
  from ray.rllib.algorithms.pg.pg import PGConfig
  from ray.rllib.algorithms.pg.utils import post_process_advantages
  from ray.rllib.evaluation.postprocessing import Postprocessing
  from ray.rllib.models.torch.torch_action_dist import TorchDistributionWrapper
  from ray.rllib.models.modelv2 import ModelV2
  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.utils.framework import try_import_torch
  from ray.rllib.utils.typing import TensorType
label_desc: Use image node_id 3 for calling the PGTorchPolicy obj's underlying member method code with example usage: obj.stats_fn(train_batch) and returns: convert_to_numpy