class_docstr: PyTorch Policy class used with A3C.
class_docstr_tok: ["PyTorch","Policy","class","used","with","A3C","."]
docstr: Constructs the loss function.
  
  Args:
      model: The Model to calculate the loss for.
      dist_class: The action distr. class.
      train_batch: The training data.
  
  Returns:
      The A3C loss tensor given the input batch.
docstr_tok: ["Constructs","the","loss","function",".","Args",":","model",":","The","Model","to","calculate","the","loss","for",".","dist_class",":","The","action","distr",".","class",".","train_batch",":","The","training","data",".","Returns",":","The","A3C","loss","tensor","given","the","input","batch","."]
code: def loss(
      self,
      model: ModelV2,
      dist_class: Type[TorchDistributionWrapper],
      train_batch: SampleBatch,
  ) -> Union[TensorType, List[TensorType]]:
      """Constructs the loss function.
  
      Args:
          model: The Model to calculate the loss for.
          dist_class: The action distr. class.
          train_batch: The training data.
  
      Returns:
          The A3C loss tensor given the input batch.
      """
      logits, _ = model(train_batch)
      values = model.value_function()
  
      if self.is_recurrent():
          B = len(train_batch[SampleBatch.SEQ_LENS])
          max_seq_len = logits.shape[0] // B
          mask_orig = sequence_mask(
              train_batch[SampleBatch.SEQ_LENS], max_seq_len
          )
          valid_mask = torch.reshape(mask_orig, [-1])
      else:
          valid_mask = torch.ones_like(values, dtype=torch.bool)
  
      dist = dist_class(logits, model)
      log_probs = dist.logp(train_batch[SampleBatch.ACTIONS]).reshape(
          -1
      )
      pi_err = -torch.sum(
          torch.masked_select(
              log_probs * train_batch[Postprocessing.ADVANTAGES],
              valid_mask,
          )
      )
  
      # Compute a value function loss.
      if self.config["use_critic"]:
          value_err = 0.5 * torch.sum(
              torch.pow(
                  torch.masked_select(
                      values.reshape(-1)
                      - train_batch[Postprocessing.VALUE_TARGETS],
                      valid_mask,
                  ),
                  2.0,
              )
          )
      # Ignore the value function.
      else:
          value_err = 0.0
  
      entropy = torch.sum(
          torch.masked_select(dist.entropy(), valid_mask)
      )
  
      total_loss = (
          pi_err
          + value_err * self.config["vf_loss_coeff"]
          - entropy * self.entropy_coeff
      )
  
      # 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["entropy"] = entropy
      model.tower_stats["pi_err"] = pi_err
      model.tower_stats["value_err"] = value_err
  
      return total_loss
  
imports: from typing import Dict, List, Optional, Type, Union
  from ray.rllib.evaluation.episode import Episode
  from ray.rllib.evaluation.postprocessing import Postprocessing, compute_gae_for_sample_batch
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.torch.torch_action_dist import TorchDistributionWrapper
  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, sequence_mask
  from ray.rllib.utils.typing import AgentID, TensorType
label_desc: Use image node_id 2 for calling the A3CTorchPolicy obj's underlying member method code with example usage: obj.loss(model, dist_class, train_batch) and returns: total_loss