class_docstr: Natural Policy Optimization with Task Embeddings.
  
  See https://karolhausman.github.io/pdf/hausman17nips-ws2.pdf for algorithm
  reference.
  
  Args:
      env_spec (EnvSpec): Environment specification.
      policy (garage.tf.policies.TaskEmbeddingPolicy): Policy.
      baseline (garage.tf.baselines.Baseline): The baseline.
      sampler (garage.sampler.Sampler): Sampler.
      scope (str): Scope for identifying the algorithm.
          Must be specified if running multiple algorithms
          simultaneously, each using different environments
          and policies.
      discount (float): Discount.
      gae_lambda (float): Lambda used for generalized advantage
          estimation.
      center_adv (bool): Whether to rescale the advantages
          so that they have mean 0 and standard deviation 1.
      positive_adv (bool): Whether to shift the advantages
          so that they are always positive. When used in
          conjunction with center_adv the advantages will be
          standardized before shifting.
      fixed_horizon (bool): Whether to fix horizon.
      lr_clip_range (float): The limit on the likelihood ratio between
          policies, as in PPO.
      max_kl_step (float): The maximum KL divergence between old and new
          policies, as in TRPO.
      optimizer (object): The optimizer of the algorithm. Should be the
          optimizers in garage.tf.optimizers.
      optimizer_args (dict): The arguments of the optimizer.
      policy_ent_coeff (float): The coefficient of the policy entropy.
          Setting it to zero would mean no entropy regularization.
      encoder_ent_coeff (float): The coefficient of the policy encoder
          entropy. Setting it to zero would mean no entropy regularization.
      use_softplus_entropy (bool): Whether to estimate the softmax
          distribution of the entropy to prevent the entropy from being
          negative.
      stop_ce_gradient (bool): Whether to stop the cross entropy gradient.
      inference (garage.tf.embeddings.StochasticEncoder): A encoder
          that infers the task embedding from state trajectory.
      inference_optimizer (object): The optimizer of the inference. Should be
          an optimizer in garage.tf.optimizers.
      inference_optimizer_args (dict): The arguments of the inference
          optimizer.
      inference_ce_coeff (float): The coefficient of the cross entropy of
          task embeddings inferred from task one-hot and state trajectory.
          This is effectively the coefficient of log-prob of inference.
      name (str): The name of the algorithm.
docstr: Build policy loss and other output tensors.
  
  Args:
      i (namedtuple): Collection of variables to compute policy loss.
  
  Returns:
      tf.Tensor: Policy loss.
      tf.Tensor: Mean policy KL divergence.
docstr_tok: ["Build","policy","loss","and","other","output","tensors",".","Args",":","i","(","namedtuple",")",":","Collection","of","variables","to","compute","policy","loss",".","Returns",":","tf.Tensor",":","Policy","loss",".","tf.Tensor",":","Mean","policy","KL","divergence","."]
code: def _build_policy_loss(self, i):
      """Build policy loss and other output tensors.
  
      Args:
          i (namedtuple): Collection of variables to compute policy loss.
  
      Returns:
          tf.Tensor: Policy loss.
          tf.Tensor: Mean policy KL divergence.
  
      """
      # pylint: disable=too-many-statements
      self._policy_network, self._encoder_network = self.policy.build(
          i.augmented_obs_var, i.task_var, name="loss_policy"
      )
      (
          self._old_policy_network,
          self._old_encoder_network,
      ) = self._old_policy.build(
          i.augmented_obs_var, i.task_var, name="loss_old_policy"
      )
      self._infer_network = self._inference.build(
          i.augmented_traj_var, name="loss_infer"
      )
      self._old_infer_network = self._old_inference.build(
          i.augmented_traj_var, name="loss_old_infer"
      )
  
      pol_dist = self._policy_network.dist
      old_pol_dist = self._old_policy_network.dist
  
      # Entropy terms
      (
          encoder_entropy,
          inference_ce,
          policy_entropy,
      ) = self._build_entropy_terms(i)
  
      # Augment the path rewards with entropy terms
      with tf.name_scope("augmented_rewards"):
          rewards = (
              i.reward_var
              - (self.inference_ce_coeff * inference_ce)
              + (self._policy_ent_coeff * policy_entropy)
          )
  
      with tf.name_scope("policy_loss"):
          with tf.name_scope("advantages"):
              adv = compute_advantages(
                  self._discount,
                  self._gae_lambda,
                  self.max_episode_length,
                  i.baseline_var,
                  rewards,
                  name="advantages",
              )
              adv = tf.reshape(adv, [-1, self.max_episode_length])
  
          # Optionally normalize advantages
          eps = tf.constant(1e-8, dtype=tf.float32)
          if self._center_adv:
              adv = center_advs(adv, axes=[0], eps=eps)
  
          if self._positive_adv:
              adv = positive_advs(adv, eps)
  
          # Calculate loss function and KL divergence
          with tf.name_scope("kl"):
              kl = old_pol_dist.kl_divergence(pol_dist)
              pol_mean_kl = tf.reduce_mean(kl)
  
          ll = pol_dist.log_prob(i.action_var, name="log_likelihood")
  
          # Calculate surrogate loss
          with tf.name_scope("surr_loss"):
              old_ll = old_pol_dist.log_prob(i.action_var)
              old_ll = tf.stop_gradient(old_ll)
              # Clip early to avoid overflow
              lr = tf.exp(
                  tf.minimum(
                      ll - old_ll, np.log(1 + self._lr_clip_range)
                  )
              )
  
              surrogate = lr * adv
  
              surrogate = tf.debugging.check_numerics(
                  surrogate, message="surrogate"
              )
  
          # Finalize objective function
          with tf.name_scope("loss"):
              lr_clip = tf.clip_by_value(
                  lr,
                  1 - self._lr_clip_range,
                  1 + self._lr_clip_range,
                  name="lr_clip",
              )
              surr_clip = lr_clip * adv
              obj = tf.minimum(surrogate, surr_clip, name="surr_obj")
              obj = tf.boolean_mask(obj, i.valid_var)
              # Maximize E[surrogate objective] by minimizing
              # -E_t[surrogate objective]
              loss = -tf.reduce_mean(obj)
  
              # Encoder entropy bonus
              loss -= self.encoder_ent_coeff * encoder_entropy
  
          encoder_mean_kl = self._build_encoder_kl()
  
          # Diagnostic functions
          self._f_policy_kl = (
              tf.compat.v1.get_default_session().make_callable(
                  pol_mean_kl,
                  feed_list=flatten_inputs(self._policy_opt_inputs),
              )
          )
  
          self._f_rewards = (
              tf.compat.v1.get_default_session().make_callable(
                  rewards,
                  feed_list=flatten_inputs(self._policy_opt_inputs),
              )
          )
  
          returns = discounted_returns(
              self._discount,
              self.max_episode_length,
              rewards,
              name="returns",
          )
          self._f_returns = (
              tf.compat.v1.get_default_session().make_callable(
                  returns,
                  feed_list=flatten_inputs(self._policy_opt_inputs),
              )
          )
  
      return loss, pol_mean_kl, encoder_mean_kl
  
imports: import akro
  from dowel import Histogram, logger, tabular
  import numpy
  import scipy.stats
  import tensorflow
  from garage import InOutSpec, log_performance
  from garage.experiment import deterministic
  from garage.np import discount_cumsum, explained_variance_1d, pad_batch_array, rrse, sliding_window
  from garage.np.algos import RLAlgorithm
  from garage.tf import center_advs, compile_function, compute_advantages, concat_tensor_list, discounted_returns, flatten_inputs, graph_inputs, pad_tensor_dict, positive_advs, stack_tensor_dict_list
  from garage.tf.embeddings import StochasticEncoder
  from garage.tf.optimizers import LBFGSOptimizer
  from garage.tf.policies import TaskEmbeddingPolicy
label_desc: Use image node_id 8 for calling the TENPO obj's underlying member method code with example usage: obj._build_policy_loss(i) and returns: loss,  pol_mean_kl,  encoder_mean_kl