text: """Adapted from A3CTFPolicy to add V-trace.
  
  Keep in sync with changes to A3CTFPolicy and VtraceSurrogatePolicy."""
  
  import numpy as np
  import logging
  import gym
  from typing import Dict, List, Type, Union
  
  import ray
  from ray.rllib.algorithms.impala import vtrace_tf as vtrace
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.tf.tf_action_dist import Categorical, TFActionDistribution
  from ray.rllib.policy.dynamic_tf_policy_v2 import DynamicTFPolicyV2
  from ray.rllib.policy.eager_tf_policy_v2 import EagerTFPolicyV2
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.policy.tf_mixins import LearningRateSchedule, EntropyCoeffSchedule
  from ray.rllib.utils import force_list
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils.framework import try_import_tf
  from ray.rllib.utils.tf_utils import explained_variance
  from ray.rllib.utils.typing import (
      LocalOptimizer,
      ModelGradients,
      TensorType,
      TFPolicyV2Type,
  )
  
  tf1, tf, tfv = try_import_tf()
  
  logger = logging.getLogger(__name__)
  
  
  class VTraceLoss:
      def __init__(
          self,
          actions,
          actions_logp,
          actions_entropy,
          dones,
          behaviour_action_logp,
          behaviour_logits,
          target_logits,
          discount,
          rewards,
          values,
          bootstrap_value,
          dist_class,
          model,
          valid_mask,
          config,
          vf_loss_coeff=0.5,
          entropy_coeff=0.01,
          clip_rho_threshold=1.0,
          clip_pg_rho_threshold=1.0,
      ):
          """Policy gradient loss with vtrace importance weighting.
  
          VTraceLoss takes tensors of shape [T, B, ...], where `B` is the
          batch_size. The reason we need to know `B` is for V-trace to properly
          handle episode cut boundaries.
  
          Args:
              actions: An int|float32 tensor of shape [T, B, ACTION_SPACE].
              actions_logp: A float32 tensor of shape [T, B].
              actions_entropy: A float32 tensor of shape [T, B].
              dones: A bool tensor of shape [T, B].
              behaviour_action_logp: Tensor of shape [T, B].
              behaviour_logits: A list with length of ACTION_SPACE of float32
                  tensors of shapes
                  [T, B, ACTION_SPACE[0]],
                  ...,
                  [T, B, ACTION_SPACE[-1]]
              target_logits: A list with length of ACTION_SPACE of float32
                  tensors of shapes
                  [T, B, ACTION_SPACE[0]],
                  ...,
                  [T, B, ACTION_SPACE[-1]]
              discount: A float32 scalar.
              rewards: A float32 tensor of shape [T, B].
              values: A float32 tensor of shape [T, B].
              bootstrap_value: A float32 tensor of shape [B].
              dist_class: action distribution class for logits.
              valid_mask: A bool tensor of valid RNN input elements (#2992).
              config: Algorithm config dict.
          """
  
          # Compute vtrace on the CPU for better perf.
          with tf.device("/cpu:0"):
              self.vtrace_returns = vtrace.multi_from_logits(
                  behaviour_action_log_probs=behaviour_action_logp,
                  behaviour_policy_logits=behaviour_logits,
                  target_policy_logits=target_logits,
                  actions=tf.unstack(actions, axis=2),
                  discounts=tf.cast(~tf.cast(dones, tf.bool), tf.float32) * discount,
                  rewards=rewards,
                  values=values,
                  bootstrap_value=bootstrap_value,
                  dist_class=dist_class,
                  model=model,
                  clip_rho_threshold=tf.cast(clip_rho_threshold, tf.float32),
                  clip_pg_rho_threshold=tf.cast(clip_pg_rho_threshold, tf.float32),
              )
              self.value_targets = self.vtrace_returns.vs
  
          # The policy gradients loss.
          masked_pi_loss = tf.boolean_mask(
              actions_logp * self.vtrace_returns.pg_advantages, valid_mask
          )
          self.pi_loss = -tf.reduce_sum(masked_pi_loss)
          self.mean_pi_loss = -tf.reduce_mean(masked_pi_loss)
  
          # The baseline loss.
          delta = tf.boolean_mask(values - self.vtrace_returns.vs, valid_mask)
          delta_squarred = tf.math.square(delta)
          self.vf_loss = 0.5 * tf.reduce_sum(delta_squarred)
          self.mean_vf_loss = 0.5 * tf.reduce_mean(delta_squarred)
  
          # The entropy loss.
          masked_entropy = tf.boolean_mask(actions_entropy, valid_mask)
          self.entropy = tf.reduce_sum(masked_entropy)
          self.mean_entropy = tf.reduce_mean(masked_entropy)
  
          # The summed weighted loss.
          self.total_loss = self.pi_loss - self.entropy * entropy_coeff
  
          # Optional vf loss (or in a separate term due to separate
          # optimizers/networks).
          self.loss_wo_vf = self.total_loss
          if not config["_separate_vf_optimizer"]:
              self.total_loss += self.vf_loss * vf_loss_coeff
  
  
  def _make_time_major(policy, seq_lens, tensor, drop_last=False):
      """Swaps batch and trajectory axis.
  
      Args:
          policy: Policy reference
          seq_lens: Sequence lengths if recurrent or None
          tensor: A tensor or list of tensors to reshape.
          drop_last: A bool indicating whether to drop the last
          trajectory item.
  
      Returns:
          res: A tensor with swapped axes or a list of tensors with
          swapped axes.
      """
      if isinstance(tensor, list):
          return [_make_time_major(policy, seq_lens, t, drop_last) for t in tensor]
  
      if policy.is_recurrent():
          B = tf.shape(seq_lens)[0]
          T = tf.shape(tensor)[0] // B
      else:
          # Important: chop the tensor into batches at known episode cut
          # boundaries.
          # TODO: (sven) this is kind of a hack and won't work for
          #  batch_mode=complete_episodes.
          T = policy.config["rollout_fragment_length"]
          B = tf.shape(tensor)[0] // T
      rs = tf.reshape(tensor, tf.concat([[B, T], tf.shape(tensor)[1:]], axis=0))
  
      # swap B and T axes
      res = tf.transpose(rs, [1, 0] + list(range(2, 1 + int(tf.shape(tensor).shape[0]))))
  
      if drop_last:
          return res[:-1]
      return res
  
  
  class VTraceClipGradients:
      """VTrace version of gradient computation logic."""
  
      def __init__(self):
          """No special initialization required."""
          pass
  
      def compute_gradients_fn(
          self, optimizer: LocalOptimizer, loss: TensorType
      ) -> ModelGradients:
          # Supporting more than one loss/optimizer.
          if self.config["_tf_policy_handles_more_than_one_loss"]:
              optimizers = force_list(optimizer)
              losses = force_list(loss)
              assert len(optimizers) == len(losses)
              clipped_grads_and_vars = []
              for optim, loss_ in zip(optimizers, losses):
                  grads_and_vars = optim.compute_gradients(
                      loss_, self.model.trainable_variables()
                  )
                  clipped_g_and_v = []
                  for g, v in grads_and_vars:
                      if g is not None:
                          clipped_g, _ = tf.clip_by_global_norm(
                              [g], self.config["grad_clip"]
                          )
                          clipped_g_and_v.append((clipped_g[0], v))
                  clipped_grads_and_vars.append(clipped_g_and_v)
  
              self.grads = [g for g_and_v in clipped_grads_and_vars for (g, v) in g_and_v]
          # Only one optimizer and and loss term.
          else:
              grads_and_vars = optimizer.compute_gradients(
                  loss, self.model.trainable_variables()
              )
              grads = [g for (g, v) in grads_and_vars]
              self.grads, _ = tf.clip_by_global_norm(grads, self.config["grad_clip"])
              clipped_grads_and_vars = list(
                  zip(self.grads, self.model.trainable_variables())
              )
  
          return clipped_grads_and_vars
  
  
  class VTraceOptimizer:
      """Optimizer function for VTrace policies."""
  
      def __init__(self):
          pass
  
      # TODO: maybe standardize this function, so the choice of optimizers are more
      #  predictable for common algorithms.
      def optimizer(
          self,
      ) -> Union["tf.keras.optimizers.Optimizer", List["tf.keras.optimizers.Optimizer"]]:
          config = self.config
          if config["opt_type"] == "adam":
              if config["framework"] in ["tf2", "tfe"]:
                  optim = tf.keras.optimizers.Adam(self.cur_lr)
                  if config["_separate_vf_optimizer"]:
                      return optim, tf.keras.optimizers.Adam(config["_lr_vf"])
              else:
                  optim = tf1.train.AdamOptimizer(self.cur_lr)
                  if config["_separate_vf_optimizer"]:
                      return optim, tf1.train.AdamOptimizer(config["_lr_vf"])
          else:
              if config["_separate_vf_optimizer"]:
                  raise ValueError(
                      "RMSProp optimizer not supported for separate"
                      "vf- and policy losses yet! Set `opt_type=adam`"
                  )
  
              if tfv == 2:
                  optim = tf.keras.optimizers.RMSprop(
                      self.cur_lr, config["decay"], config["momentum"], config["epsilon"]
                  )
              else:
                  optim = tf1.train.RMSPropOptimizer(
                      self.cur_lr, config["decay"], config["momentum"], config["epsilon"]
                  )
  
          return optim
  
  
  # We need this builder function because we want to share the same
  # custom logics between TF1 dynamic and TF2 eager policies.
  def get_impala_tf_policy(base: TFPolicyV2Type) -> TFPolicyV2Type:
      """Construct an ImpalaTFPolicy inheriting either dynamic or eager base policies.
  
      Args:
          base: Base class for this policy. DynamicTFPolicyV2 or EagerTFPolicyV2.
  
      Returns:
          A TF Policy to be used with Impala.
      """
      # VTrace mixins are placed in front of more general mixins to make sure
      # their functions like optimizer() overrides all the other implementations
      # (e.g., LearningRateSchedule.optimizer())
      class ImpalaTFPolicy(
          VTraceClipGradients,
          VTraceOptimizer,
          LearningRateSchedule,
          EntropyCoeffSchedule,
          base,
      ):
          def __init__(
              self,
              obs_space,
              action_space,
              config,
              existing_model=None,
              existing_inputs=None,
          ):
              # First thing first, enable eager execution if necessary.
              base.enable_eager_execution_if_necessary()
  
              config = dict(
                  ray.rllib.algorithms.impala.impala.ImpalaConfig().to_dict(), **config
              )
  
              # Initialize base class.
              base.__init__(
                  self,
                  obs_space,
                  action_space,
                  config,
                  existing_inputs=existing_inputs,
                  existing_model=existing_model,
              )
  
              VTraceClipGradients.__init__(self)
              VTraceOptimizer.__init__(self)
              LearningRateSchedule.__init__(self, config["lr"], config["lr_schedule"])
              EntropyCoeffSchedule.__init__(
                  self, config["entropy_coeff"], config["entropy_coeff_schedule"]
              )
  
              # Note: this is a bit ugly, but loss and optimizer initialization must
              # happen after all the MixIns are initialized.
              self.maybe_initialize_optimizer_and_loss()
  
          @override(base)
          def loss(
              self,
              model: Union[ModelV2, "tf.keras.Model"],
              dist_class: Type[TFActionDistribution],
              train_batch: SampleBatch,
          ) -> Union[TensorType, List[TensorType]]:
              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.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, **kw):
                  return _make_time_major(
                      self, train_batch.get(SampleBatch.SEQ_LENS), *args, **kw
                  )
  
              actions = train_batch[SampleBatch.ACTIONS]
              dones = train_batch[SampleBatch.DONES]
              rewards = train_batch[SampleBatch.REWARDS]
              behaviour_action_logp = train_batch[SampleBatch.ACTION_LOGP]
              behaviour_logits = train_batch[SampleBatch.ACTION_DIST_INPUTS]
              unpacked_behaviour_logits = tf.split(
                  behaviour_logits, output_hidden_shape, axis=1
              )
              unpacked_outputs = tf.split(model_out, output_hidden_shape, axis=1)
              values = model.value_function()
  
              if self.is_recurrent():
                  max_seq_len = tf.reduce_max(train_batch[SampleBatch.SEQ_LENS])
                  mask = tf.sequence_mask(train_batch[SampleBatch.SEQ_LENS], max_seq_len)
                  mask = tf.reshape(mask, [-1])
              else:
                  mask = tf.ones_like(rewards)
  
              # Prepare actions for loss
              loss_actions = (
                  actions if is_multidiscrete else tf.expand_dims(actions, axis=1)
              )
  
              # Inputs are reshaped from [B * T] => [(T|T-1), B] for V-trace calc.
              drop_last = self.config["vtrace_drop_last_ts"]
              self.vtrace_loss = VTraceLoss(
                  actions=make_time_major(loss_actions, drop_last=drop_last),
                  actions_logp=make_time_major(
                      action_dist.logp(actions), drop_last=drop_last
                  ),
                  actions_entropy=make_time_major(
                      action_dist.multi_entropy(), drop_last=drop_last
                  ),
                  dones=make_time_major(dones, drop_last=drop_last),
                  behaviour_action_logp=make_time_major(
                      behaviour_action_logp, drop_last=drop_last
                  ),
                  behaviour_logits=make_time_major(
                      unpacked_behaviour_logits, drop_last=drop_last
                  ),
                  target_logits=make_time_major(unpacked_outputs, drop_last=drop_last),
                  discount=self.config["gamma"],
                  rewards=make_time_major(rewards, drop_last=drop_last),
                  values=make_time_major(values, drop_last=drop_last),
                  bootstrap_value=make_time_major(values)[-1],
                  dist_class=Categorical if is_multidiscrete else dist_class,
                  model=model,
                  valid_mask=make_time_major(mask, drop_last=drop_last),
                  config=self.config,
                  vf_loss_coeff=self.config["vf_loss_coeff"],
                  entropy_coeff=self.entropy_coeff,
                  clip_rho_threshold=self.config["vtrace_clip_rho_threshold"],
                  clip_pg_rho_threshold=self.config["vtrace_clip_pg_rho_threshold"],
              )
  
              if self.config.get("_separate_vf_optimizer"):
                  return self.vtrace_loss.loss_wo_vf, self.vtrace_loss.vf_loss
              else:
                  return self.vtrace_loss.total_loss
  
          @override(base)
          def stats_fn(self, train_batch: SampleBatch) -> Dict[str, TensorType]:
              drop_last = self.config["vtrace"] and self.config["vtrace_drop_last_ts"]
              values_batched = _make_time_major(
                  self,
                  train_batch.get(SampleBatch.SEQ_LENS),
                  self.model.value_function(),
                  drop_last=drop_last,
              )
  
              return {
                  "cur_lr": tf.cast(self.cur_lr, tf.float64),
                  "policy_loss": self.vtrace_loss.mean_pi_loss,
                  "entropy": self.vtrace_loss.mean_entropy,
                  "entropy_coeff": tf.cast(self.entropy_coeff, tf.float64),
                  "var_gnorm": tf.linalg.global_norm(self.model.trainable_variables()),
                  "vf_loss": self.vtrace_loss.mean_vf_loss,
                  "vf_explained_var": explained_variance(
                      tf.reshape(self.vtrace_loss.value_targets, [-1]),
                      tf.reshape(values_batched, [-1]),
                  ),
              }
  
          @override(base)
          def grad_stats_fn(
              self, train_batch: SampleBatch, grads: ModelGradients
          ) -> Dict[str, TensorType]:
              # We have support for more than one loss (list of lists of grads).
              if self.config.get("_tf_policy_handles_more_than_one_loss"):
                  grad_gnorm = [tf.linalg.global_norm(g) for g in grads]
              # Old case: We have a single list of grads (only one loss term and
              # optimizer).
              else:
                  grad_gnorm = tf.linalg.global_norm(grads)
  
              return {
                  "grad_gnorm": grad_gnorm,
              }
  
          @override(base)
          def get_batch_divisibility_req(self) -> int:
              return self.config["rollout_fragment_length"]
  
      return ImpalaTFPolicy
  
  
  ImpalaTF1Policy = get_impala_tf_policy(DynamicTFPolicyV2)
  ImpalaTF2Policy = get_impala_tf_policy(EagerTFPolicyV2)
  