class_docstr: PyTorch policy class used with MAML.
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 PPO 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","PPO","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 PPO loss tensor given the input batch.
      """
      logits, state = model(train_batch)
      self.cur_lr = self.config["lr"]
  
      if self.config["worker_index"]:
          self.loss_obj = WorkerLoss(
              model=model,
              dist_class=dist_class,
              actions=train_batch[SampleBatch.ACTIONS],
              curr_logits=logits,
              behaviour_logits=train_batch[
                  SampleBatch.ACTION_DIST_INPUTS
              ],
              advantages=train_batch[Postprocessing.ADVANTAGES],
              value_fn=model.value_function(),
              value_targets=train_batch[Postprocessing.VALUE_TARGETS],
              vf_preds=train_batch[SampleBatch.VF_PREDS],
              cur_kl_coeff=0.0,
              entropy_coeff=self.config["entropy_coeff"],
              clip_param=self.config["clip_param"],
              vf_clip_param=self.config["vf_clip_param"],
              vf_loss_coeff=self.config["vf_loss_coeff"],
              clip_loss=False,
          )
      else:
          self.var_list = model.named_parameters()
  
          # `split` may not exist yet (during test-loss call), use a dummy value.
          # Cannot use get here due to train_batch being a TrackingDict.
          if "split" in train_batch:
              split = train_batch["split"]
          else:
              split_shape = (
                  self.config["inner_adaptation_steps"],
                  self.config["num_workers"],
              )
              split_const = int(
                  train_batch["obs"].shape[0]
                  // (split_shape[0] * split_shape[1])
              )
              split = torch.ones(split_shape, dtype=int) * split_const
          self.loss_obj = MAMLLoss(
              model=model,
              dist_class=dist_class,
              value_targets=train_batch[Postprocessing.VALUE_TARGETS],
              advantages=train_batch[Postprocessing.ADVANTAGES],
              actions=train_batch[SampleBatch.ACTIONS],
              behaviour_logits=train_batch[
                  SampleBatch.ACTION_DIST_INPUTS
              ],
              vf_preds=train_batch[SampleBatch.VF_PREDS],
              cur_kl_coeff=self.kl_coeff_val,
              policy_vars=self.var_list,
              obs=train_batch[SampleBatch.CUR_OBS],
              num_tasks=self.config["num_workers"],
              split=split,
              config=self.config,
              inner_adaptation_steps=self.config[
                  "inner_adaptation_steps"
              ],
              entropy_coeff=self.config["entropy_coeff"],
              clip_param=self.config["clip_param"],
              vf_clip_param=self.config["vf_clip_param"],
              vf_loss_coeff=self.config["vf_loss_coeff"],
              use_gae=self.config["use_gae"],
              meta_opt=self.meta_opt,
          )
  
      return self.loss_obj.loss
  
imports: import logging
  from typing import Dict, List, Type, Union
  import ray
  from ray.rllib.algorithms.ppo.ppo_tf_policy import validate_config
  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 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
  from ray.rllib.utils.typing import TensorType
label_desc: Use image node_id 2 for calling the MAMLTorchPolicy obj's underlying member method code with example usage: obj.loss(model, dist_class, train_batch) and returns: self