text: """
  PyTorch policy class used for APPO.
  
  Adapted from VTraceTFPolicy to use the PPO surrogate loss.
  Keep in sync with changes to VTraceTFPolicy.
  """
  
  import gym
  import numpy as np
  import logging
  from typing import Any, Dict, List, Optional, Type, Union
  
  import ray
  from ray.rllib.algorithms.appo.utils import make_appo_models
  import ray.rllib.algorithms.impala.vtrace_torch as vtrace
  from ray.rllib.algorithms.impala.impala_torch_policy import (
      make_time_major,
      VTraceOptimizer,
  )
  from ray.rllib.evaluation.episode import Episode
  from ray.rllib.evaluation.postprocessing import (
      compute_gae_for_sample_batch,
      Postprocessing,
  )
  from ray.rllib.models.action_dist import ActionDistribution
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.torch.torch_action_dist import (
      TorchDistributionWrapper,
      TorchCategorical,
  )
  from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.policy.torch_mixins import (
      EntropyCoeffSchedule,
      LearningRateSchedule,
      KLCoeffMixin,
      ValueNetworkMixin,
      TargetNetworkMixin,
  )
  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
  
  torch, nn = try_import_torch()
  
  logger = logging.getLogger(__name__)
  
  
  # We need this builder function because we want to share the same
  # custom logics between TF1 dynamic and TF2 eager policies.
  class APPOTorchPolicy(
      VTraceOptimizer,
      LearningRateSchedule,
      EntropyCoeffSchedule,
      KLCoeffMixin,
      ValueNetworkMixin,
      TargetNetworkMixin,
      TorchPolicyV2,
  ):
      """PyTorch policy class used with APPO."""
  
      def __init__(self, observation_space, action_space, config):
          config = dict(ray.rllib.algorithms.appo.appo.APPOConfig().to_dict(), **config)
  
          # Although this is a no-op, we call __init__ here to make it clear
          # that base.__init__ will use the make_model() call.
          VTraceOptimizer.__init__(self)
          LearningRateSchedule.__init__(self, config["lr"], config["lr_schedule"])
  
          TorchPolicyV2.__init__(
              self,
              observation_space,
              action_space,
              config,
              max_seq_len=config["model"]["max_seq_len"],
          )
  
          EntropyCoeffSchedule.__init__(
              self, config["entropy_coeff"], config["entropy_coeff_schedule"]
          )
          ValueNetworkMixin.__init__(self, config)
          KLCoeffMixin.__init__(self, config)
  
          # TODO: Don't require users to call this manually.
          self._initialize_loss_from_dummy_batch()
  
          # Initiate TargetNetwork ops after loss initialization.
          TargetNetworkMixin.__init__(self)
  
      @override(TorchPolicyV2)
      def init_view_requirements(self):
          self.view_requirements = self._get_default_view_requirements()
  
      @override(TorchPolicyV2)
      def make_model(self) -> ModelV2:
          return make_appo_models(self)
  
      @override(TorchPolicyV2)
      def loss(
          self,
          model: ModelV2,
          dist_class: Type[ActionDistribution],
          train_batch: SampleBatch,
      ) -> Union[TensorType, List[TensorType]]:
          """Constructs the loss for APPO.
  
          With IS modifications and V-trace for Advantage Estimation.
  
          Args:
              model (ModelV2): The Model to calculate the loss for.
              dist_class (Type[ActionDistribution]): The action distr. class.
              train_batch: The training data.
  
          Returns:
              Union[TensorType, List[TensorType]]: A single loss tensor or a list
                  of loss tensors.
          """
          target_model = self.target_models[model]
  
          model_out, _ = model(train_batch)
          action_dist = dist_class(model_out, model)
  
          if isinstance(self.action_space, gym.spaces.Discrete):
              is_multidiscrete = False
              output_hidden_shape = [self.action_space.n]
          elif isinstance(self.action_space, gym.spaces.multi_discrete.MultiDiscrete):
              is_multidiscrete = True
              output_hidden_shape = self.action_space.nvec.astype(np.int32)
          else:
              is_multidiscrete = False
              output_hidden_shape = 1
  
          def _make_time_major(*args, **kwargs):
              return make_time_major(
                  self, train_batch.get(SampleBatch.SEQ_LENS), *args, **kwargs
              )
  
          actions = train_batch[SampleBatch.ACTIONS]
          dones = train_batch[SampleBatch.DONES]
          rewards = train_batch[SampleBatch.REWARDS]
          behaviour_logits = train_batch[SampleBatch.ACTION_DIST_INPUTS]
  
          target_model_out, _ = target_model(train_batch)
  
          prev_action_dist = dist_class(behaviour_logits, model)
          values = model.value_function()
          values_time_major = _make_time_major(values)
  
          drop_last = self.config["vtrace"] and self.config["vtrace_drop_last_ts"]
  
          if self.is_recurrent():
              max_seq_len = torch.max(train_batch[SampleBatch.SEQ_LENS])
              mask = sequence_mask(train_batch[SampleBatch.SEQ_LENS], max_seq_len)
              mask = torch.reshape(mask, [-1])
              mask = _make_time_major(mask, drop_last=drop_last)
              num_valid = torch.sum(mask)
  
              def reduce_mean_valid(t):
                  return torch.sum(t[mask]) / num_valid
  
          else:
              reduce_mean_valid = torch.mean
  
          if self.config["vtrace"]:
              logger.debug(
                  "Using V-Trace surrogate loss (vtrace=True; " f"drop_last={drop_last})"
              )
  
              old_policy_behaviour_logits = target_model_out.detach()
              old_policy_action_dist = dist_class(old_policy_behaviour_logits, model)
  
              if isinstance(output_hidden_shape, (list, tuple, np.ndarray)):
                  unpacked_behaviour_logits = torch.split(
                      behaviour_logits, list(output_hidden_shape), dim=1
                  )
                  unpacked_old_policy_behaviour_logits = torch.split(
                      old_policy_behaviour_logits, list(output_hidden_shape), dim=1
                  )
              else:
                  unpacked_behaviour_logits = torch.chunk(
                      behaviour_logits, output_hidden_shape, dim=1
                  )
                  unpacked_old_policy_behaviour_logits = torch.chunk(
                      old_policy_behaviour_logits, output_hidden_shape, dim=1
                  )
  
              # Prepare actions for loss.
              loss_actions = (
                  actions if is_multidiscrete else torch.unsqueeze(actions, dim=1)
              )
  
              # Prepare KL for loss.
              action_kl = _make_time_major(
                  old_policy_action_dist.kl(action_dist), drop_last=drop_last
              )
  
              # Compute vtrace on the CPU for better perf.
              vtrace_returns = vtrace.multi_from_logits(
                  behaviour_policy_logits=_make_time_major(
                      unpacked_behaviour_logits, drop_last=drop_last
                  ),
                  target_policy_logits=_make_time_major(
                      unpacked_old_policy_behaviour_logits, drop_last=drop_last
                  ),
                  actions=torch.unbind(
                      _make_time_major(loss_actions, drop_last=drop_last), dim=2
                  ),
                  discounts=(1.0 - _make_time_major(dones, drop_last=drop_last).float())
                  * self.config["gamma"],
                  rewards=_make_time_major(rewards, drop_last=drop_last),
                  values=values_time_major[:-1] if drop_last else values_time_major,
                  bootstrap_value=values_time_major[-1],
                  dist_class=TorchCategorical if is_multidiscrete else dist_class,
                  model=model,
                  clip_rho_threshold=self.config["vtrace_clip_rho_threshold"],
                  clip_pg_rho_threshold=self.config["vtrace_clip_pg_rho_threshold"],
              )
  
              actions_logp = _make_time_major(
                  action_dist.logp(actions), drop_last=drop_last
              )
              prev_actions_logp = _make_time_major(
                  prev_action_dist.logp(actions), drop_last=drop_last
              )
              old_policy_actions_logp = _make_time_major(
                  old_policy_action_dist.logp(actions), drop_last=drop_last
              )
              is_ratio = torch.clamp(
                  torch.exp(prev_actions_logp - old_policy_actions_logp), 0.0, 2.0
              )
              logp_ratio = is_ratio * torch.exp(actions_logp - prev_actions_logp)
              self._is_ratio = is_ratio
  
              advantages = vtrace_returns.pg_advantages.to(logp_ratio.device)
              surrogate_loss = torch.min(
                  advantages * logp_ratio,
                  advantages
                  * torch.clamp(
                      logp_ratio,
                      1 - self.config["clip_param"],
                      1 + self.config["clip_param"],
                  ),
              )
  
              mean_kl_loss = reduce_mean_valid(action_kl)
              mean_policy_loss = -reduce_mean_valid(surrogate_loss)
  
              # The value function loss.
              value_targets = vtrace_returns.vs.to(values_time_major.device)
              if drop_last:
                  delta = values_time_major[:-1] - value_targets
              else:
                  delta = values_time_major - value_targets
              mean_vf_loss = 0.5 * reduce_mean_valid(torch.pow(delta, 2.0))
  
              # The entropy loss.
              mean_entropy = reduce_mean_valid(
                  _make_time_major(action_dist.entropy(), drop_last=drop_last)
              )
  
          else:
              logger.debug("Using PPO surrogate loss (vtrace=False)")
  
              # Prepare KL for Loss
              action_kl = _make_time_major(prev_action_dist.kl(action_dist))
  
              actions_logp = _make_time_major(action_dist.logp(actions))
              prev_actions_logp = _make_time_major(prev_action_dist.logp(actions))
              logp_ratio = torch.exp(actions_logp - prev_actions_logp)
  
              advantages = _make_time_major(train_batch[Postprocessing.ADVANTAGES])
              surrogate_loss = torch.min(
                  advantages * logp_ratio,
                  advantages
                  * torch.clamp(
                      logp_ratio,
                      1 - self.config["clip_param"],
                      1 + self.config["clip_param"],
                  ),
              )
  
              mean_kl_loss = reduce_mean_valid(action_kl)
              mean_policy_loss = -reduce_mean_valid(surrogate_loss)
  
              # The value function loss.
              value_targets = _make_time_major(train_batch[Postprocessing.VALUE_TARGETS])
              delta = values_time_major - value_targets
              mean_vf_loss = 0.5 * reduce_mean_valid(torch.pow(delta, 2.0))
  
              # The entropy loss.
              mean_entropy = reduce_mean_valid(_make_time_major(action_dist.entropy()))
  
          # The summed weighted loss
          total_loss = (
              mean_policy_loss
              + mean_vf_loss * self.config["vf_loss_coeff"]
              - mean_entropy * self.entropy_coeff
          )
  
          # Optional additional KL Loss
          if self.config["use_kl_loss"]:
              total_loss += self.kl_coeff * mean_kl_loss
  
          # 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["total_loss"] = total_loss
          model.tower_stats["mean_policy_loss"] = mean_policy_loss
          model.tower_stats["mean_kl_loss"] = mean_kl_loss
          model.tower_stats["mean_vf_loss"] = mean_vf_loss
          model.tower_stats["mean_entropy"] = mean_entropy
          model.tower_stats["value_targets"] = value_targets
          model.tower_stats["vf_explained_var"] = explained_variance(
              torch.reshape(value_targets, [-1]),
              torch.reshape(
                  values_time_major[:-1] if drop_last else values_time_major, [-1]
              ),
          )
  
          return total_loss
  
      @override(TorchPolicyV2)
      def stats_fn(self, train_batch: SampleBatch) -> Dict[str, TensorType]:
          """Stats function for APPO. Returns a dict with important loss stats.
  
          Args:
              policy: The Policy to generate stats for.
              train_batch: The SampleBatch (already) used for training.
  
          Returns:
              Dict[str, TensorType]: The stats dict.
          """
          stats_dict = {
              "cur_lr": self.cur_lr,
              "total_loss": torch.mean(torch.stack(self.get_tower_stats("total_loss"))),
              "policy_loss": torch.mean(
                  torch.stack(self.get_tower_stats("mean_policy_loss"))
              ),
              "entropy": torch.mean(torch.stack(self.get_tower_stats("mean_entropy"))),
              "entropy_coeff": self.entropy_coeff,
              "var_gnorm": global_norm(self.model.trainable_variables()),
              "vf_loss": torch.mean(torch.stack(self.get_tower_stats("mean_vf_loss"))),
              "vf_explained_var": torch.mean(
                  torch.stack(self.get_tower_stats("vf_explained_var"))
              ),
          }
  
          if self.config["vtrace"]:
              is_stat_mean = torch.mean(self._is_ratio, [0, 1])
              is_stat_var = torch.var(self._is_ratio, [0, 1])
              stats_dict["mean_IS"] = is_stat_mean
              stats_dict["var_IS"] = is_stat_var
  
          if self.config["use_kl_loss"]:
              stats_dict["kl"] = torch.mean(
                  torch.stack(self.get_tower_stats("mean_kl_loss"))
              )
              stats_dict["KL_Coeff"] = self.kl_coeff
  
          return convert_to_numpy(stats_dict)
  
      @override(TorchPolicyV2)
      def extra_action_out(
          self,
          input_dict: Dict[str, TensorType],
          state_batches: List[TensorType],
          model: TorchModelV2,
          action_dist: TorchDistributionWrapper,
      ) -> Dict[str, TensorType]:
          out = {}
          if not self.config["vtrace"]:
              out[SampleBatch.VF_PREDS] = model.value_function()
          return out
  
      @override(TorchPolicyV2)
      def postprocess_trajectory(
          self,
          sample_batch: SampleBatch,
          other_agent_batches: Optional[Dict[Any, SampleBatch]] = None,
          episode: Optional["Episode"] = None,
      ):
          # Call super's postprocess_trajectory first.
          sample_batch = super().postprocess_trajectory(
              sample_batch, other_agent_batches, episode
          )
          if not self.config["vtrace"]:
              # Do all post-processing always with no_grad().
              # Not using this here will introduce a memory leak
              # in torch (issue #6962).
              with torch.no_grad():
                  sample_batch = compute_gae_for_sample_batch(
                      self, sample_batch, other_agent_batches, episode
                  )
          return sample_batch
  
      @override(TorchPolicyV2)
      def extra_grad_process(
          self, optimizer: "torch.optim.Optimizer", loss: TensorType
      ) -> Dict[str, TensorType]:
          return apply_grad_clipping(self, optimizer, loss)
  
      @override(TorchPolicyV2)
      def get_batch_divisibility_req(self) -> int:
          return self.config["rollout_fragment_length"]
  